From 3178309e1ae21238bd09e2e0d301740b45f6308f Mon Sep 17 00:00:00 2001 From: Anish Sarkar <104695310+AnishSarkar22@users.noreply.github.com> Date: Wed, 20 May 2026 02:02:27 +0530 Subject: [PATCH] feat: add team memory routes --- surfsense_backend/app/routes/__init__.py | 2 + surfsense_backend/app/routes/memory_routes.py | 141 ++++-------------- .../app/routes/search_spaces_routes.py | 111 -------------- .../app/routes/team_memory_routes.py | 78 ++++++++++ surfsense_backend/app/schemas/search_space.py | 1 - 5 files changed, 111 insertions(+), 222 deletions(-) create mode 100644 surfsense_backend/app/routes/team_memory_routes.py diff --git a/surfsense_backend/app/routes/__init__.py b/surfsense_backend/app/routes/__init__.py index 5b6a74376..ec4d1650f 100644 --- a/surfsense_backend/app/routes/__init__.py +++ b/surfsense_backend/app/routes/__init__.py @@ -54,6 +54,7 @@ from .search_spaces_routes import router as search_spaces_router from .slack_add_connector_route import router as slack_add_connector_router from .stripe_routes import router as stripe_router from .surfsense_docs_routes import router as surfsense_docs_router +from .team_memory_routes import router as team_memory_router from .teams_add_connector_route import router as teams_add_connector_router from .video_presentations_routes import router as video_presentations_router from .vision_llm_routes import router as vision_llm_router @@ -117,3 +118,4 @@ router.include_router(stripe_router) # Stripe checkout for additional page pack router.include_router(youtube_router) # YouTube playlist resolution router.include_router(prompts_router) router.include_router(memory_router) # User personal memory (memory.md style) +router.include_router(team_memory_router) # Search-space team memory diff --git a/surfsense_backend/app/routes/memory_routes.py b/surfsense_backend/app/routes/memory_routes.py index e57ca4055..7b674a584 100644 --- a/surfsense_backend/app/routes/memory_routes.py +++ b/surfsense_backend/app/routes/memory_routes.py @@ -1,24 +1,19 @@ -"""Routes for user memory management (personal memory.md).""" +"""Routes for user memory management.""" from __future__ import annotations -import logging - from fastapi import APIRouter, Depends, HTTPException -from langchain_core.messages import HumanMessage from pydantic import BaseModel from sqlalchemy.ext.asyncio import AsyncSession -from app.agents.new_chat.llm_config import ( - create_chat_litellm_from_agent_config, - load_agent_llm_config_for_search_space, -) -from app.agents.new_chat.tools.update_memory import MEMORY_HARD_LIMIT, _save_memory from app.db import User, get_async_session +from app.services.memory import ( + MemoryScope, + read_memory, + reset_memory, + save_memory, +) from app.users import current_active_user -from app.utils.content_utils import extract_text_content - -logger = logging.getLogger(__name__) router = APIRouter() @@ -31,45 +26,17 @@ class MemoryUpdate(BaseModel): memory_md: str -class MemoryEditRequest(BaseModel): - query: str - search_space_id: int - - -_MEMORY_EDIT_PROMPT = """\ -You are a memory editor. The user wants to modify their memory document. \ -Apply the user's instruction to the existing memory document and output the \ -FULL updated document. - -RULES: -1. If the instruction asks to add something, add it with format: \ -- (YYYY-MM-DD) [fact|pref|instr] text, under an existing or new ## heading. \ -Heading names should be personal and descriptive, not generic categories. -2. If the instruction asks to remove something, remove the matching entry. -3. If the instruction asks to change something, update the matching entry. -4. Preserve existing ## headings and all other entries. -5. Every bullet must include a marker: [fact], [pref], or [instr]. -6. Use the user's first name (from ) in entries instead of "the user". -7. Output ONLY the updated markdown — no explanations, no wrapping. - -{user_name} - - -{current_memory} - - - -{instruction} -""" - - @router.get("/users/me/memory", response_model=MemoryRead) async def get_user_memory( user: User = Depends(current_active_user), session: AsyncSession = Depends(get_async_session), ): - await session.refresh(user, ["memory_md"]) - return MemoryRead(memory_md=user.memory_md or "") + memory_md = await read_memory( + scope=MemoryScope.USER, + target_id=user.id, + session=session, + ) + return MemoryRead(memory_md=memory_md) @router.put("/users/me/memory", response_model=MemoryRead) @@ -78,73 +45,27 @@ async def update_user_memory( user: User = Depends(current_active_user), session: AsyncSession = Depends(get_async_session), ): - if len(body.memory_md) > MEMORY_HARD_LIMIT: - raise HTTPException( - status_code=400, - detail=f"Memory exceeds {MEMORY_HARD_LIMIT:,} character limit ({len(body.memory_md):,} chars).", - ) - user.memory_md = body.memory_md - session.add(user) - await session.commit() - await session.refresh(user, ["memory_md"]) - return MemoryRead(memory_md=user.memory_md or "") + result = await save_memory( + scope=MemoryScope.USER, + target_id=user.id, + content=body.memory_md, + session=session, + ) + if result.status == "error": + raise HTTPException(status_code=400, detail=result.message) + return MemoryRead(memory_md=result.memory_md) -@router.post("/users/me/memory/edit", response_model=MemoryRead) -async def edit_user_memory( - body: MemoryEditRequest, +@router.post("/users/me/memory/reset", response_model=MemoryRead) +async def reset_user_memory( user: User = Depends(current_active_user), session: AsyncSession = Depends(get_async_session), ): - """Apply a natural language edit to the user's personal memory via LLM.""" - agent_config = await load_agent_llm_config_for_search_space( - session, body.search_space_id + result = await reset_memory( + scope=MemoryScope.USER, + target_id=user.id, + session=session, ) - if not agent_config: - raise HTTPException(status_code=500, detail="No LLM configuration available.") - llm = create_chat_litellm_from_agent_config(agent_config) - if not llm: - raise HTTPException(status_code=500, detail="Failed to create LLM instance.") - - await session.refresh(user, ["memory_md", "display_name"]) - current_memory = user.memory_md or "" - first_name = ( - user.display_name.strip().split()[0] - if user.display_name and user.display_name.strip() - else "The user" - ) - - prompt = _MEMORY_EDIT_PROMPT.format( - current_memory=current_memory or "(empty)", - instruction=body.query, - user_name=first_name, - ) - try: - response = await llm.ainvoke( - [HumanMessage(content=prompt)], - config={"tags": ["surfsense:internal", "memory-edit"]}, - ) - updated = extract_text_content(response.content).strip() - except Exception as e: - logger.exception("Memory edit LLM call failed: %s", e) - raise HTTPException(status_code=500, detail="Memory edit failed.") from e - - if not updated: - raise HTTPException(status_code=400, detail="LLM returned empty result.") - - result = await _save_memory( - updated_memory=updated, - old_memory=current_memory, - llm=llm, - apply_fn=lambda content: setattr(user, "memory_md", content), - commit_fn=session.commit, - rollback_fn=session.rollback, - label="memory", - scope="user", - ) - - if result.get("status") == "error": - raise HTTPException(status_code=400, detail=result["message"]) - - await session.refresh(user, ["memory_md"]) - return MemoryRead(memory_md=user.memory_md or "") + if result.status == "error": + raise HTTPException(status_code=400, detail=result.message) + return MemoryRead(memory_md=result.memory_md) diff --git a/surfsense_backend/app/routes/search_spaces_routes.py b/surfsense_backend/app/routes/search_spaces_routes.py index 0f0e43035..db230b0f5 100644 --- a/surfsense_backend/app/routes/search_spaces_routes.py +++ b/surfsense_backend/app/routes/search_spaces_routes.py @@ -1,17 +1,10 @@ import logging from fastapi import APIRouter, Depends, HTTPException -from langchain_core.messages import HumanMessage -from pydantic import BaseModel as PydanticBaseModel from sqlalchemy import func, update from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.future import select -from app.agents.new_chat.llm_config import ( - create_chat_litellm_from_agent_config, - load_agent_llm_config_for_search_space, -) -from app.agents.new_chat.tools.update_memory import MEMORY_HARD_LIMIT, _save_memory from app.config import config from app.db import ( ImageGenerationConfig, @@ -35,7 +28,6 @@ from app.schemas import ( SearchSpaceWithStats, ) from app.users import current_active_user -from app.utils.content_utils import extract_text_content from app.utils.rbac import check_permission, check_search_space_access logger = logging.getLogger(__name__) @@ -43,34 +35,6 @@ logger = logging.getLogger(__name__) router = APIRouter() -class _TeamMemoryEditRequest(PydanticBaseModel): - query: str - - -_TEAM_MEMORY_EDIT_PROMPT = """\ -You are a memory editor for a team workspace. The user wants to modify the \ -team's shared memory document. Apply the user's instruction to the existing \ -memory document and output the FULL updated document. - -RULES: -1. If the instruction asks to add something, add it with format: \ -- (YYYY-MM-DD) [fact] text, under an existing or new ## heading. \ -Heading names should be descriptive, not generic categories. -2. If the instruction asks to remove something, remove the matching entry. -3. If the instruction asks to change something, update the matching entry. -4. Preserve existing ## headings and all other entries. -5. NEVER use [pref] or [instr] markers. Team memory uses [fact] only. -6. Output ONLY the updated markdown — no explanations, no wrapping. - - -{current_memory} - - - -{instruction} -""" - - async def create_default_roles_and_membership( session: AsyncSession, search_space_id: int, @@ -294,15 +258,6 @@ async def update_search_space( update_data = search_space_update.model_dump(exclude_unset=True) - if ( - "shared_memory_md" in update_data - and len(update_data["shared_memory_md"] or "") > MEMORY_HARD_LIMIT - ): - raise HTTPException( - status_code=400, - detail=f"Team memory exceeds {MEMORY_HARD_LIMIT:,} character limit.", - ) - for key, value in update_data.items(): setattr(db_search_space, key, value) await session.commit() @@ -317,72 +272,6 @@ async def update_search_space( ) from e -@router.post( - "/searchspaces/{search_space_id}/memory/edit", - response_model=SearchSpaceRead, -) -async def edit_team_memory( - search_space_id: int, - body: _TeamMemoryEditRequest, - session: AsyncSession = Depends(get_async_session), - user: User = Depends(current_active_user), -): - """Apply a natural language edit to the team memory via LLM.""" - await check_search_space_access(session, user, search_space_id) - - agent_config = await load_agent_llm_config_for_search_space( - session, search_space_id - ) - if not agent_config: - raise HTTPException(status_code=500, detail="No LLM configuration available.") - llm = create_chat_litellm_from_agent_config(agent_config) - if not llm: - raise HTTPException(status_code=500, detail="Failed to create LLM instance.") - - result = await session.execute( - select(SearchSpace).filter(SearchSpace.id == search_space_id) - ) - db_search_space = result.scalars().first() - if not db_search_space: - raise HTTPException(status_code=404, detail="Search space not found") - - current_memory = db_search_space.shared_memory_md or "" - - prompt = _TEAM_MEMORY_EDIT_PROMPT.format( - current_memory=current_memory or "(empty)", - instruction=body.query, - ) - try: - response = await llm.ainvoke( - [HumanMessage(content=prompt)], - config={"tags": ["surfsense:internal", "memory-edit"]}, - ) - updated = extract_text_content(response.content).strip() - except Exception as e: - logger.exception("Team memory edit LLM call failed: %s", e) - raise HTTPException(status_code=500, detail="Team memory edit failed.") from e - - if not updated: - raise HTTPException(status_code=400, detail="LLM returned empty result.") - - save_result = await _save_memory( - updated_memory=updated, - old_memory=current_memory, - llm=llm, - apply_fn=lambda content: setattr(db_search_space, "shared_memory_md", content), - commit_fn=session.commit, - rollback_fn=session.rollback, - label="team memory", - scope="team", - ) - - if save_result.get("status") == "error": - raise HTTPException(status_code=400, detail=save_result["message"]) - - await session.refresh(db_search_space) - return db_search_space - - @router.post("/searchspaces/{search_space_id}/ai-sort") async def trigger_ai_sort( search_space_id: int, diff --git a/surfsense_backend/app/routes/team_memory_routes.py b/surfsense_backend/app/routes/team_memory_routes.py new file mode 100644 index 000000000..3e552ce32 --- /dev/null +++ b/surfsense_backend/app/routes/team_memory_routes.py @@ -0,0 +1,78 @@ +"""Routes for search-space team memory.""" + +from __future__ import annotations + +from fastapi import APIRouter, Depends, HTTPException +from pydantic import BaseModel +from sqlalchemy.ext.asyncio import AsyncSession + +from app.db import User, get_async_session +from app.services.memory import ( + MemoryScope, + read_memory, + reset_memory, + save_memory, +) +from app.users import current_active_user +from app.utils.rbac import check_search_space_access + +router = APIRouter() + + +class TeamMemoryRead(BaseModel): + memory_md: str + + +class TeamMemoryUpdate(BaseModel): + memory_md: str + + +@router.get("/searchspaces/{search_space_id}/memory", response_model=TeamMemoryRead) +async def get_team_memory( + search_space_id: int, + session: AsyncSession = Depends(get_async_session), + user: User = Depends(current_active_user), +): + await check_search_space_access(session, user, search_space_id) + memory_md = await read_memory( + scope=MemoryScope.TEAM, + target_id=search_space_id, + session=session, + ) + return TeamMemoryRead(memory_md=memory_md) + + +@router.put("/searchspaces/{search_space_id}/memory", response_model=TeamMemoryRead) +async def update_team_memory( + search_space_id: int, + body: TeamMemoryUpdate, + session: AsyncSession = Depends(get_async_session), + user: User = Depends(current_active_user), +): + await check_search_space_access(session, user, search_space_id) + result = await save_memory( + scope=MemoryScope.TEAM, + target_id=search_space_id, + content=body.memory_md, + session=session, + ) + if result.status == "error": + raise HTTPException(status_code=400, detail=result.message) + return TeamMemoryRead(memory_md=result.memory_md) + + +@router.post("/searchspaces/{search_space_id}/memory/reset", response_model=TeamMemoryRead) +async def reset_team_memory( + search_space_id: int, + session: AsyncSession = Depends(get_async_session), + user: User = Depends(current_active_user), +): + await check_search_space_access(session, user, search_space_id) + result = await reset_memory( + scope=MemoryScope.TEAM, + target_id=search_space_id, + session=session, + ) + if result.status == "error": + raise HTTPException(status_code=400, detail=result.message) + return TeamMemoryRead(memory_md=result.memory_md) diff --git a/surfsense_backend/app/schemas/search_space.py b/surfsense_backend/app/schemas/search_space.py index 77e34ea4b..70ed0004e 100644 --- a/surfsense_backend/app/schemas/search_space.py +++ b/surfsense_backend/app/schemas/search_space.py @@ -21,7 +21,6 @@ class SearchSpaceUpdate(BaseModel): description: str | None = None citations_enabled: bool | None = None qna_custom_instructions: str | None = None - shared_memory_md: str | None = None ai_file_sort_enabled: bool | None = None