feat: add team memory routes

This commit is contained in:
Anish Sarkar 2026-05-20 02:02:27 +05:30
parent 5247dc7097
commit 3178309e1a
5 changed files with 111 additions and 222 deletions

View file

@ -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 .slack_add_connector_route import router as slack_add_connector_router
from .stripe_routes import router as stripe_router from .stripe_routes import router as stripe_router
from .surfsense_docs_routes import router as surfsense_docs_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 .teams_add_connector_route import router as teams_add_connector_router
from .video_presentations_routes import router as video_presentations_router from .video_presentations_routes import router as video_presentations_router
from .vision_llm_routes import router as vision_llm_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(youtube_router) # YouTube playlist resolution
router.include_router(prompts_router) router.include_router(prompts_router)
router.include_router(memory_router) # User personal memory (memory.md style) router.include_router(memory_router) # User personal memory (memory.md style)
router.include_router(team_memory_router) # Search-space team memory

View file

@ -1,24 +1,19 @@
"""Routes for user memory management (personal memory.md).""" """Routes for user memory management."""
from __future__ import annotations from __future__ import annotations
import logging
from fastapi import APIRouter, Depends, HTTPException from fastapi import APIRouter, Depends, HTTPException
from langchain_core.messages import HumanMessage
from pydantic import BaseModel from pydantic import BaseModel
from sqlalchemy.ext.asyncio import AsyncSession 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.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.users import current_active_user
from app.utils.content_utils import extract_text_content
logger = logging.getLogger(__name__)
router = APIRouter() router = APIRouter()
@ -31,45 +26,17 @@ class MemoryUpdate(BaseModel):
memory_md: str 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 <user_name>) in entries instead of "the user".
7. Output ONLY the updated markdown no explanations, no wrapping.
<user_name>{user_name}</user_name>
<current_memory>
{current_memory}
</current_memory>
<user_instruction>
{instruction}
</user_instruction>"""
@router.get("/users/me/memory", response_model=MemoryRead) @router.get("/users/me/memory", response_model=MemoryRead)
async def get_user_memory( async def get_user_memory(
user: User = Depends(current_active_user), user: User = Depends(current_active_user),
session: AsyncSession = Depends(get_async_session), session: AsyncSession = Depends(get_async_session),
): ):
await session.refresh(user, ["memory_md"]) memory_md = await read_memory(
return MemoryRead(memory_md=user.memory_md or "") scope=MemoryScope.USER,
target_id=user.id,
session=session,
)
return MemoryRead(memory_md=memory_md)
@router.put("/users/me/memory", response_model=MemoryRead) @router.put("/users/me/memory", response_model=MemoryRead)
@ -78,73 +45,27 @@ async def update_user_memory(
user: User = Depends(current_active_user), user: User = Depends(current_active_user),
session: AsyncSession = Depends(get_async_session), session: AsyncSession = Depends(get_async_session),
): ):
if len(body.memory_md) > MEMORY_HARD_LIMIT: result = await save_memory(
raise HTTPException( scope=MemoryScope.USER,
status_code=400, target_id=user.id,
detail=f"Memory exceeds {MEMORY_HARD_LIMIT:,} character limit ({len(body.memory_md):,} chars).", content=body.memory_md,
) session=session,
user.memory_md = body.memory_md )
session.add(user) if result.status == "error":
await session.commit() raise HTTPException(status_code=400, detail=result.message)
await session.refresh(user, ["memory_md"]) return MemoryRead(memory_md=result.memory_md)
return MemoryRead(memory_md=user.memory_md or "")
@router.post("/users/me/memory/edit", response_model=MemoryRead) @router.post("/users/me/memory/reset", response_model=MemoryRead)
async def edit_user_memory( async def reset_user_memory(
body: MemoryEditRequest,
user: User = Depends(current_active_user), user: User = Depends(current_active_user),
session: AsyncSession = Depends(get_async_session), session: AsyncSession = Depends(get_async_session),
): ):
"""Apply a natural language edit to the user's personal memory via LLM.""" result = await reset_memory(
agent_config = await load_agent_llm_config_for_search_space( scope=MemoryScope.USER,
session, body.search_space_id target_id=user.id,
session=session,
) )
if not agent_config: if result.status == "error":
raise HTTPException(status_code=500, detail="No LLM configuration available.") raise HTTPException(status_code=400, detail=result.message)
llm = create_chat_litellm_from_agent_config(agent_config) return MemoryRead(memory_md=result.memory_md)
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 "")

View file

@ -1,17 +1,10 @@
import logging import logging
from fastapi import APIRouter, Depends, HTTPException 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 import func, update
from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.future import select 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.config import config
from app.db import ( from app.db import (
ImageGenerationConfig, ImageGenerationConfig,
@ -35,7 +28,6 @@ from app.schemas import (
SearchSpaceWithStats, SearchSpaceWithStats,
) )
from app.users import current_active_user 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 from app.utils.rbac import check_permission, check_search_space_access
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@ -43,34 +35,6 @@ logger = logging.getLogger(__name__)
router = APIRouter() 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>
{current_memory}
</current_memory>
<user_instruction>
{instruction}
</user_instruction>"""
async def create_default_roles_and_membership( async def create_default_roles_and_membership(
session: AsyncSession, session: AsyncSession,
search_space_id: int, search_space_id: int,
@ -294,15 +258,6 @@ async def update_search_space(
update_data = search_space_update.model_dump(exclude_unset=True) 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(): for key, value in update_data.items():
setattr(db_search_space, key, value) setattr(db_search_space, key, value)
await session.commit() await session.commit()
@ -317,72 +272,6 @@ async def update_search_space(
) from e ) 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") @router.post("/searchspaces/{search_space_id}/ai-sort")
async def trigger_ai_sort( async def trigger_ai_sort(
search_space_id: int, search_space_id: int,

View file

@ -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)

View file

@ -21,7 +21,6 @@ class SearchSpaceUpdate(BaseModel):
description: str | None = None description: str | None = None
citations_enabled: bool | None = None citations_enabled: bool | None = None
qna_custom_instructions: str | None = None qna_custom_instructions: str | None = None
shared_memory_md: str | None = None
ai_file_sort_enabled: bool | None = None ai_file_sort_enabled: bool | None = None