From ceedd02353c74f063188fc103500fb848909a810 Mon Sep 17 00:00:00 2001 From: Anish Sarkar <104695310+AnishSarkar22@users.noreply.github.com> Date: Wed, 20 May 2026 02:01:36 +0530 Subject: [PATCH 01/15] refactor: extract shared memory service --- .../builtins/memory/tools/update_memory.py | 351 ++------------- .../app/agents/new_chat/memory_extraction.py | 196 +-------- .../agents/new_chat/tools/update_memory.py | 414 ++---------------- .../app/services/memory/__init__.py | 29 ++ .../app/services/memory/prompts.py | 110 +++++ .../app/services/memory/rewrite.py | 35 ++ .../app/services/memory/schemas.py | 23 + .../app/services/memory/service.py | 300 +++++++++++++ .../app/services/memory/validation.py | 158 +++++++ .../unit/services/test_memory_service.py | 204 +++++++++ 10 files changed, 946 insertions(+), 874 deletions(-) create mode 100644 surfsense_backend/app/services/memory/__init__.py create mode 100644 surfsense_backend/app/services/memory/prompts.py create mode 100644 surfsense_backend/app/services/memory/rewrite.py create mode 100644 surfsense_backend/app/services/memory/schemas.py create mode 100644 surfsense_backend/app/services/memory/service.py create mode 100644 surfsense_backend/app/services/memory/validation.py create mode 100644 surfsense_backend/tests/unit/services/test_memory_service.py diff --git a/surfsense_backend/app/agents/multi_agent_chat/subagents/builtins/memory/tools/update_memory.py b/surfsense_backend/app/agents/multi_agent_chat/subagents/builtins/memory/tools/update_memory.py index 23375a081..67bcc3e06 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/subagents/builtins/memory/tools/update_memory.py +++ b/surfsense_backend/app/agents/multi_agent_chat/subagents/builtins/memory/tools/update_memory.py @@ -1,280 +1,23 @@ -"""Overwrite one markdown memory document per user or team, with size and shrink guards.""" +"""Memory update tools backed by the canonical memory service.""" from __future__ import annotations import logging -import re -from typing import Any, Literal +from typing import Any from uuid import UUID -from langchain_core.messages import HumanMessage from langchain_core.tools import tool -from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession -from app.db import SearchSpace, User +from app.services.memory import ( + MEMORY_HARD_LIMIT, + MEMORY_SOFT_LIMIT, + MemoryScope, + save_memory, +) logger = logging.getLogger(__name__) -MEMORY_SOFT_LIMIT = 18_000 -MEMORY_HARD_LIMIT = 25_000 - -_SECTION_HEADING_RE = re.compile(r"^##\s+(.+)$", re.MULTILINE) -_HEADING_NORMALIZE_RE = re.compile(r"\s+") - -_MARKER_RE = re.compile(r"\[(fact|pref|instr)\]") -_BULLET_FORMAT_RE = re.compile(r"^- \(\d{4}-\d{2}-\d{2}\) \[(fact|pref|instr)\] .+$") -_PERSONAL_ONLY_MARKERS = {"pref", "instr"} - - -# --------------------------------------------------------------------------- -# Diff validation -# --------------------------------------------------------------------------- - - -def _extract_headings(memory: str) -> set[str]: - """Return all ``## …`` heading texts (without the ``## `` prefix).""" - return set(_SECTION_HEADING_RE.findall(memory)) - - -def _normalize_heading(heading: str) -> str: - """Normalize heading text for robust scope checks.""" - return _HEADING_NORMALIZE_RE.sub(" ", heading.strip().lower()) - - -def _validate_memory_scope( - content: str, scope: Literal["user", "team"] -) -> dict[str, Any] | None: - """Reject personal-only markers ([pref], [instr]) in team memory.""" - if scope != "team": - return None - - markers = set(_MARKER_RE.findall(content)) - leaked = sorted(markers & _PERSONAL_ONLY_MARKERS) - if leaked: - tags = ", ".join(f"[{m}]" for m in leaked) - return { - "status": "error", - "message": ( - f"Team memory cannot include personal markers: {tags}. " - "Use [fact] only in team memory." - ), - } - return None - - -def _validate_bullet_format(content: str) -> list[str]: - """Return warnings for bullet lines that don't match the required format. - - Expected: ``- (YYYY-MM-DD) [fact|pref|instr] text`` - """ - warnings: list[str] = [] - for line in content.splitlines(): - stripped = line.strip() - if not stripped.startswith("- "): - continue - if not _BULLET_FORMAT_RE.match(stripped): - short = stripped[:80] + ("..." if len(stripped) > 80 else "") - warnings.append(f"Malformed bullet: {short}") - return warnings - - -def _validate_diff(old_memory: str | None, new_memory: str) -> list[str]: - """Return a list of warning strings about suspicious changes.""" - if not old_memory: - return [] - - warnings: list[str] = [] - old_headings = _extract_headings(old_memory) - new_headings = _extract_headings(new_memory) - dropped = old_headings - new_headings - if dropped: - names = ", ".join(sorted(dropped)) - warnings.append( - f"Sections removed: {names}. " - "If unintentional, the user can restore from the settings page." - ) - - old_len = len(old_memory) - new_len = len(new_memory) - if old_len > 0 and new_len < old_len * 0.4: - warnings.append( - f"Memory shrank significantly ({old_len:,} -> {new_len:,} chars). " - "Possible data loss." - ) - return warnings - - -# --------------------------------------------------------------------------- -# Size validation & soft warning -# --------------------------------------------------------------------------- - - -def _validate_memory_size(content: str) -> dict[str, Any] | None: - """Return an error/warning dict if *content* is too large, else None.""" - length = len(content) - if length > MEMORY_HARD_LIMIT: - return { - "status": "error", - "message": ( - f"Memory exceeds {MEMORY_HARD_LIMIT:,} character limit " - f"({length:,} chars). Consolidate by merging related items, " - "removing outdated entries, and shortening descriptions. " - "Then call update_memory again." - ), - } - return None - - -def _soft_warning(content: str) -> str | None: - """Return a warning string if content exceeds the soft limit.""" - length = len(content) - if length > MEMORY_SOFT_LIMIT: - return ( - f"Memory is at {length:,}/{MEMORY_HARD_LIMIT:,} characters. " - "Consolidate by merging related items and removing less important " - "entries on your next update." - ) - return None - - -# --------------------------------------------------------------------------- -# Forced rewrite when memory exceeds the hard limit -# --------------------------------------------------------------------------- - -_FORCED_REWRITE_PROMPT = """\ -You are a memory curator. The following memory document exceeds the character \ -limit and must be shortened. - -RULES: -1. Rewrite the document to be under {target} characters. -2. Preserve existing ## headings. Every entry must remain under a heading. You may merge - or rename headings to consolidate, but keep names personal and descriptive. -3. Priority for keeping content: [instr] > [pref] > [fact]. -4. Merge duplicate entries, remove outdated entries, shorten verbose descriptions. -5. Every bullet MUST have format: - (YYYY-MM-DD) [fact|pref|instr] text -6. Preserve the user's first name in entries — do not replace it with "the user". -7. Output ONLY the consolidated markdown — no explanations, no wrapping. - - -{content} -""" - - -async def _forced_rewrite(content: str, llm: Any) -> str | None: - """Use a focused LLM call to compress *content* under the hard limit. - - Returns the rewritten string, or ``None`` if the call fails. - """ - try: - prompt = _FORCED_REWRITE_PROMPT.format( - target=MEMORY_HARD_LIMIT, content=content - ) - response = await llm.ainvoke( - [HumanMessage(content=prompt)], - config={"tags": ["surfsense:internal"]}, - ) - text = ( - response.content - if isinstance(response.content, str) - else str(response.content) - ) - return text.strip() - except Exception: - logger.exception("Forced rewrite LLM call failed") - return None - - -# --------------------------------------------------------------------------- -# Shared save-and-respond logic -# --------------------------------------------------------------------------- - - -async def _save_memory( - *, - updated_memory: str, - old_memory: str | None, - llm: Any | None, - apply_fn, - commit_fn, - rollback_fn, - label: str, - scope: Literal["user", "team"], -) -> dict[str, Any]: - """Validate, optionally force-rewrite if over the hard limit, save, and - return a response dict. - - Parameters - ---------- - updated_memory : str - The new document the agent submitted. - old_memory : str | None - The previously persisted document (for diff checks). - llm : Any | None - LLM instance for forced rewrite (may be ``None``). - apply_fn : callable(str) -> None - Callback that sets the new memory on the ORM object. - commit_fn : coroutine - ``session.commit``. - rollback_fn : coroutine - ``session.rollback``. - label : str - Human label for log messages (e.g. "user memory", "team memory"). - """ - content = updated_memory - - # --- forced rewrite if over the hard limit --- - if len(content) > MEMORY_HARD_LIMIT and llm is not None: - rewritten = await _forced_rewrite(content, llm) - if rewritten is not None and len(rewritten) < len(content): - content = rewritten - - # --- hard-limit gate (reject if still too large after rewrite) --- - size_err = _validate_memory_size(content) - if size_err: - return size_err - - scope_err = _validate_memory_scope(content, scope) - if scope_err: - return scope_err - - # --- persist --- - try: - apply_fn(content) - await commit_fn() - except Exception as e: - logger.exception("Failed to update %s: %s", label, e) - await rollback_fn() - return {"status": "error", "message": f"Failed to update {label}: {e}"} - - # --- build response --- - resp: dict[str, Any] = { - "status": "saved", - "message": f"{label.capitalize()} updated.", - } - - if content is not updated_memory: - resp["notice"] = "Memory was automatically rewritten to fit within limits." - - diff_warnings = _validate_diff(old_memory, content) - if diff_warnings: - resp["diff_warnings"] = diff_warnings - - format_warnings = _validate_bullet_format(content) - if format_warnings: - resp["format_warnings"] = format_warnings - - warning = _soft_warning(content) - if warning: - resp["warning"] = warning - - return resp - - -# --------------------------------------------------------------------------- -# Tool factories -# --------------------------------------------------------------------------- - def create_update_memory_tool( user_id: str | UUID, @@ -287,40 +30,22 @@ def create_update_memory_tool( async def update_memory(updated_memory: str) -> dict[str, Any]: """Update the user's personal memory document. - Your current memory is shown in in the system prompt. - When the user shares important long-term information (preferences, - facts, instructions, context), rewrite the memory document to include - the new information. Merge new facts with existing ones, update - contradictions, remove outdated entries, and keep it concise. - - Args: - updated_memory: The FULL updated markdown document (not a diff). + The current memory is shown in . Pass the FULL updated + markdown document, not a diff. """ try: - result = await db_session.execute(select(User).where(User.id == uid)) - user = result.scalars().first() - if not user: - return {"status": "error", "message": "User not found."} - - old_memory = user.memory_md - - return await _save_memory( - updated_memory=updated_memory, - old_memory=old_memory, + result = await save_memory( + scope=MemoryScope.USER, + target_id=uid, + content=updated_memory, + session=db_session, llm=llm, - apply_fn=lambda content: setattr(user, "memory_md", content), - commit_fn=db_session.commit, - rollback_fn=db_session.rollback, - label="memory", - scope="user", ) + return result.to_dict() except Exception as e: logger.exception("Failed to update user memory: %s", e) await db_session.rollback() - return { - "status": "error", - "message": f"Failed to update memory: {e}", - } + return {"status": "error", "message": f"Failed to update memory: {e}"} return update_memory @@ -334,36 +59,18 @@ def create_update_team_memory_tool( async def update_memory(updated_memory: str) -> dict[str, Any]: """Update the team's shared memory document for this search space. - Your current team memory is shown in in the system - prompt. When the team shares important long-term information - (decisions, conventions, key facts, priorities), rewrite the memory - document to include the new information. Merge new facts with - existing ones, update contradictions, remove outdated entries, and - keep it concise. - - Args: - updated_memory: The FULL updated markdown document (not a diff). + The current team memory is shown in . Pass the FULL updated + markdown document, not a diff. """ try: - result = await db_session.execute( - select(SearchSpace).where(SearchSpace.id == search_space_id) - ) - space = result.scalars().first() - if not space: - return {"status": "error", "message": "Search space not found."} - - old_memory = space.shared_memory_md - - return await _save_memory( - updated_memory=updated_memory, - old_memory=old_memory, + result = await save_memory( + scope=MemoryScope.TEAM, + target_id=search_space_id, + content=updated_memory, + session=db_session, llm=llm, - apply_fn=lambda content: setattr(space, "shared_memory_md", content), - commit_fn=db_session.commit, - rollback_fn=db_session.rollback, - label="team memory", - scope="team", ) + return result.to_dict() except Exception as e: logger.exception("Failed to update team memory: %s", e) await db_session.rollback() @@ -373,3 +80,11 @@ def create_update_team_memory_tool( } return update_memory + + +__all__ = [ + "MEMORY_HARD_LIMIT", + "MEMORY_SOFT_LIMIT", + "create_update_memory_tool", + "create_update_team_memory_tool", +] diff --git a/surfsense_backend/app/agents/new_chat/memory_extraction.py b/surfsense_backend/app/agents/new_chat/memory_extraction.py index e31774a7c..d44b58f7b 100644 --- a/surfsense_backend/app/agents/new_chat/memory_extraction.py +++ b/surfsense_backend/app/agents/new_chat/memory_extraction.py @@ -1,9 +1,4 @@ -"""Background memory extraction for the SurfSense agent. - -After each agent response, if the agent did not call ``update_memory`` during -the turn, this module can run a lightweight LLM call to decide whether the -latest message contains long-term information worth persisting. -""" +"""Background memory extraction for the SurfSense agent.""" from __future__ import annotations @@ -11,102 +6,11 @@ import logging from typing import Any from uuid import UUID -from langchain_core.messages import HumanMessage -from sqlalchemy import select - -from app.agents.new_chat.tools.update_memory import _save_memory -from app.db import SearchSpace, User, shielded_async_session -from app.utils.content_utils import extract_text_content +from app.db import User, shielded_async_session +from app.services.memory import MemoryScope, extract_and_save logger = logging.getLogger(__name__) -_MEMORY_EXTRACT_PROMPT = """\ -You are a memory extraction assistant. Analyze the user's message and decide \ -if it contains any long-term information worth persisting to memory. - -Worth remembering: preferences, background/identity, goals, projects, \ -instructions, tools/languages they use, decisions, expertise, workplace — \ -durable facts that will matter in future conversations. - -NOT worth remembering: greetings, one-off factual questions, session \ -logistics, ephemeral requests, follow-up clarifications with no new personal \ -info, things that only matter for the current task. - -If the message contains memorizable information, output the FULL updated \ -memory document with the new facts merged into the existing content. Follow \ -these rules: -- Every entry MUST be under a ## heading. Preserve existing headings; create new ones - freely. Keep heading names short (2-3 words) and natural. Do NOT include the user's - name in headings. -- Keep entries as single bullet points. Be descriptive but concise — include relevant - details and context rather than just a few words. -- Every bullet MUST use format: - (YYYY-MM-DD) [fact|pref|instr] text - [fact] = durable facts, [pref] = preferences, [instr] = standing instructions. -- Use the user's first name (from ) in entry text, not "the user". -- If a new fact contradicts an existing entry, update the existing entry. -- Do not duplicate information that is already present. - -If nothing is worth remembering, output exactly: NO_UPDATE - -{user_name} - - -{current_memory} - - - -{user_message} -""" - -_TEAM_MEMORY_EXTRACT_PROMPT = """\ -You are a team-memory extraction assistant. Analyze the latest message and \ -decide if it contains durable TEAM-level information worth persisting. - -Decision policy: -- Prioritize recall for durable team context, while avoiding personal-only facts. -- Do NOT require explicit consensus language. A direct team-level statement can - be stored if it is stable and broadly useful for future team chats. -- If evidence is weak or clearly tentative, output NO_UPDATE. - -Worth remembering (team-level only): -- Decisions and defaults that guide future team work -- Team conventions/standards (naming, review policy, coding norms) -- Stable org/project facts (locations, ownership, constraints) -- Long-lived architecture/process facts -- Ongoing priorities that are likely relevant beyond this turn - -NOT worth remembering: -- Personal preferences or biography of one person -- Questions, brainstorming, tentative ideas, or speculation -- One-off requests, status updates, TODOs, logistics for this session -- Information scoped only to a single ephemeral task - -If the message contains memorizable team information, output the FULL updated \ -team memory document with new facts merged into existing content. Follow rules: -- Every entry MUST be under a ## heading. Preserve existing headings; create new ones - freely. Keep heading names short (2-3 words) and natural. -- Keep entries as single bullet points. Be descriptive but concise — include relevant - details and context rather than just a few words. -- Every bullet MUST use format: - (YYYY-MM-DD) [fact] text - Team memory uses ONLY the [fact] marker. Never use [pref] or [instr]. -- If a new fact contradicts an existing entry, update the existing entry. -- Do not duplicate existing information. -- Preserve neutral team phrasing; avoid person-specific memory unless role-anchored. - -If nothing is worth remembering, output exactly: NO_UPDATE - - -{current_memory} - - - -{author} - - - -{user_message} -""" - async def extract_and_save_memory( *, @@ -114,57 +18,31 @@ async def extract_and_save_memory( user_id: str | None, llm: Any, ) -> None: - """Background task: extract memorizable info and persist it. + """Fire-and-forget personal memory extraction. - Designed to be fire-and-forget — catches all exceptions internally. + The service uses structured output, so free-form ``NO_UPDATE`` text can no + longer be accidentally persisted as memory. """ if not user_id: return try: uid = UUID(user_id) if isinstance(user_id, str) else user_id - async with shielded_async_session() as session: - result = await session.execute(select(User).where(User.id == uid)) - user = result.scalars().first() - if not user: - return - - old_memory = user.memory_md - first_name = ( - user.display_name.strip().split()[0] - if user.display_name and user.display_name.strip() - else "The user" - ) - prompt = _MEMORY_EXTRACT_PROMPT.format( - current_memory=old_memory or "(empty)", + user = await session.get(User, uid) + actor_display_name = user.display_name if user else None + result = await extract_and_save( + scope=MemoryScope.USER, + target_id=uid, user_message=user_message, - user_name=first_name, - ) - response = await llm.ainvoke( - [HumanMessage(content=prompt)], - config={"tags": ["surfsense:internal", "memory-extraction"]}, - ) - text = extract_text_content(response.content).strip() - - if text == "NO_UPDATE" or not text: - logger.debug("Memory extraction: no update needed (user %s)", uid) - return - - save_result = await _save_memory( - updated_memory=text, - old_memory=old_memory, + actor_display_name=actor_display_name, + session=session, llm=llm, - apply_fn=lambda content: setattr(user, "memory_md", content), - commit_fn=session.commit, - rollback_fn=session.rollback, - label="memory", - scope="user", ) logger.info( "Background memory extraction for user %s: %s", uid, - save_result.get("status"), + result.status, ) except Exception: logger.exception("Background user memory extraction failed") @@ -177,56 +55,24 @@ async def extract_and_save_team_memory( llm: Any, author_display_name: str | None = None, ) -> None: - """Background task: extract team-level memory and persist it. - - Runs only for shared threads. Designed to be fire-and-forget and catches - exceptions internally. - """ + """Fire-and-forget team-level memory extraction.""" if not search_space_id: return try: async with shielded_async_session() as session: - result = await session.execute( - select(SearchSpace).where(SearchSpace.id == search_space_id) - ) - space = result.scalars().first() - if not space: - return - - old_memory = space.shared_memory_md - prompt = _TEAM_MEMORY_EXTRACT_PROMPT.format( - current_memory=old_memory or "(empty)", - author=author_display_name or "Unknown team member", + result = await extract_and_save( + scope=MemoryScope.TEAM, + target_id=search_space_id, user_message=user_message, - ) - response = await llm.ainvoke( - [HumanMessage(content=prompt)], - config={"tags": ["surfsense:internal", "team-memory-extraction"]}, - ) - text = extract_text_content(response.content).strip() - - if text == "NO_UPDATE" or not text: - logger.debug( - "Team memory extraction: no update needed (space %s)", - search_space_id, - ) - return - - save_result = await _save_memory( - updated_memory=text, - old_memory=old_memory, + actor_display_name=author_display_name, + session=session, llm=llm, - apply_fn=lambda content: setattr(space, "shared_memory_md", content), - commit_fn=session.commit, - rollback_fn=session.rollback, - label="team memory", - scope="team", ) logger.info( "Background team memory extraction for space %s: %s", search_space_id, - save_result.get("status"), + result.status, ) except Exception: logger.exception("Background team memory extraction failed") diff --git a/surfsense_backend/app/agents/new_chat/tools/update_memory.py b/surfsense_backend/app/agents/new_chat/tools/update_memory.py index 062668aac..78a65201b 100644 --- a/surfsense_backend/app/agents/new_chat/tools/update_memory.py +++ b/surfsense_backend/app/agents/new_chat/tools/update_memory.py @@ -1,369 +1,53 @@ -"""Markdown-document memory tool for the SurfSense agent. - -Replaces the old row-per-fact save_memory / recall_memory tools with a single -update_memory tool that overwrites a freeform markdown TEXT column. The LLM -always sees the current memory in / tags injected -by MemoryInjectionMiddleware, so it passes the FULL updated document each time. - -Overflow handling: - - Soft limit (18K chars): a warning is returned telling the agent to - consolidate on the next update. - - Hard limit (25K chars): a forced LLM-driven rewrite compresses the document. - If it still exceeds the limit after rewriting, the save is rejected. - - Diff validation: warns when entire ``##`` sections are dropped or when the - document shrinks by more than 60%. -""" +"""Memory update tools backed by the canonical memory service.""" from __future__ import annotations import logging -import re -from typing import Any, Literal +from typing import Any from uuid import UUID -from langchain_core.messages import HumanMessage from langchain_core.tools import tool -from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession -from app.db import SearchSpace, User, async_session_maker -from app.utils.content_utils import extract_text_content +from app.db import async_session_maker +from app.services.memory import MemoryScope, save_memory logger = logging.getLogger(__name__) -MEMORY_SOFT_LIMIT = 18_000 -MEMORY_HARD_LIMIT = 25_000 - -_SECTION_HEADING_RE = re.compile(r"^##\s+(.+)$", re.MULTILINE) -_HEADING_NORMALIZE_RE = re.compile(r"\s+") - -_MARKER_RE = re.compile(r"\[(fact|pref|instr)\]") -_BULLET_FORMAT_RE = re.compile(r"^- \(\d{4}-\d{2}-\d{2}\) \[(fact|pref|instr)\] .+$") -_PERSONAL_ONLY_MARKERS = {"pref", "instr"} - - -# --------------------------------------------------------------------------- -# Diff validation -# --------------------------------------------------------------------------- - - -def _extract_headings(memory: str) -> set[str]: - """Return all ``## …`` heading texts (without the ``## `` prefix).""" - return set(_SECTION_HEADING_RE.findall(memory)) - - -def _normalize_heading(heading: str) -> str: - """Normalize heading text for robust scope checks.""" - return _HEADING_NORMALIZE_RE.sub(" ", heading.strip().lower()) - - -def _validate_memory_scope( - content: str, scope: Literal["user", "team"] -) -> dict[str, Any] | None: - """Reject personal-only markers ([pref], [instr]) in team memory.""" - if scope != "team": - return None - - markers = set(_MARKER_RE.findall(content)) - leaked = sorted(markers & _PERSONAL_ONLY_MARKERS) - if leaked: - tags = ", ".join(f"[{m}]" for m in leaked) - return { - "status": "error", - "message": ( - f"Team memory cannot include personal markers: {tags}. " - "Use [fact] only in team memory." - ), - } - return None - - -def _validate_bullet_format(content: str) -> list[str]: - """Return warnings for bullet lines that don't match the required format. - - Expected: ``- (YYYY-MM-DD) [fact|pref|instr] text`` - """ - warnings: list[str] = [] - for line in content.splitlines(): - stripped = line.strip() - if not stripped.startswith("- "): - continue - if not _BULLET_FORMAT_RE.match(stripped): - short = stripped[:80] + ("..." if len(stripped) > 80 else "") - warnings.append(f"Malformed bullet: {short}") - return warnings - - -def _validate_diff(old_memory: str | None, new_memory: str) -> list[str]: - """Return a list of warning strings about suspicious changes.""" - if not old_memory: - return [] - - warnings: list[str] = [] - old_headings = _extract_headings(old_memory) - new_headings = _extract_headings(new_memory) - dropped = old_headings - new_headings - if dropped: - names = ", ".join(sorted(dropped)) - warnings.append( - f"Sections removed: {names}. " - "If unintentional, the user can restore from the settings page." - ) - - old_len = len(old_memory) - new_len = len(new_memory) - if old_len > 0 and new_len < old_len * 0.4: - warnings.append( - f"Memory shrank significantly ({old_len:,} -> {new_len:,} chars). " - "Possible data loss." - ) - return warnings - - -# --------------------------------------------------------------------------- -# Size validation & soft warning -# --------------------------------------------------------------------------- - - -def _validate_memory_size(content: str) -> dict[str, Any] | None: - """Return an error/warning dict if *content* is too large, else None.""" - length = len(content) - if length > MEMORY_HARD_LIMIT: - return { - "status": "error", - "message": ( - f"Memory exceeds {MEMORY_HARD_LIMIT:,} character limit " - f"({length:,} chars). Consolidate by merging related items, " - "removing outdated entries, and shortening descriptions. " - "Then call update_memory again." - ), - } - return None - - -def _soft_warning(content: str) -> str | None: - """Return a warning string if content exceeds the soft limit.""" - length = len(content) - if length > MEMORY_SOFT_LIMIT: - return ( - f"Memory is at {length:,}/{MEMORY_HARD_LIMIT:,} characters. " - "Consolidate by merging related items and removing less important " - "entries on your next update." - ) - return None - - -# --------------------------------------------------------------------------- -# Forced rewrite when memory exceeds the hard limit -# --------------------------------------------------------------------------- - -_FORCED_REWRITE_PROMPT = """\ -You are a memory curator. The following memory document exceeds the character \ -limit and must be shortened. - -RULES: -1. Rewrite the document to be under {target} characters. -2. Preserve existing ## headings. Every entry must remain under a heading. You may merge - or rename headings to consolidate, but keep names personal and descriptive. -3. Priority for keeping content: [instr] > [pref] > [fact]. -4. Merge duplicate entries, remove outdated entries, shorten verbose descriptions. -5. Every bullet MUST have format: - (YYYY-MM-DD) [fact|pref|instr] text -6. Preserve the user's first name in entries — do not replace it with "the user". -7. Output ONLY the consolidated markdown — no explanations, no wrapping. - - -{content} -""" - - -async def _forced_rewrite(content: str, llm: Any) -> str | None: - """Use a focused LLM call to compress *content* under the hard limit. - - Returns the rewritten string, or ``None`` if the call fails. - """ - try: - prompt = _FORCED_REWRITE_PROMPT.format( - target=MEMORY_HARD_LIMIT, content=content - ) - response = await llm.ainvoke( - [HumanMessage(content=prompt)], - config={"tags": ["surfsense:internal"]}, - ) - text = extract_text_content(response.content).strip() - if not text: - logger.warning("Forced rewrite returned empty text; aborting rewrite") - return None - return text - except Exception: - logger.exception("Forced rewrite LLM call failed") - return None - - -# --------------------------------------------------------------------------- -# Shared save-and-respond logic -# --------------------------------------------------------------------------- - - -async def _save_memory( - *, - updated_memory: str, - old_memory: str | None, - llm: Any | None, - apply_fn, - commit_fn, - rollback_fn, - label: str, - scope: Literal["user", "team"], -) -> dict[str, Any]: - """Validate, optionally force-rewrite if over the hard limit, save, and - return a response dict. - - Parameters - ---------- - updated_memory : str - The new document the agent submitted. - old_memory : str | None - The previously persisted document (for diff checks). - llm : Any | None - LLM instance for forced rewrite (may be ``None``). - apply_fn : callable(str) -> None - Callback that sets the new memory on the ORM object. - commit_fn : coroutine - ``session.commit``. - rollback_fn : coroutine - ``session.rollback``. - label : str - Human label for log messages (e.g. "user memory", "team memory"). - """ - if not isinstance(updated_memory, str): - logger.warning( - "Refusing non-string memory payload (type=%s)", - type(updated_memory).__name__, - ) - return { - "status": "error", - "message": "Internal error: memory payload must be a string.", - } - - content = updated_memory - - # --- forced rewrite if over the hard limit --- - if len(content) > MEMORY_HARD_LIMIT and llm is not None: - rewritten = await _forced_rewrite(content, llm) - if rewritten is not None and len(rewritten) < len(content): - content = rewritten - - # --- hard-limit gate (reject if still too large after rewrite) --- - size_err = _validate_memory_size(content) - if size_err: - return size_err - - scope_err = _validate_memory_scope(content, scope) - if scope_err: - return scope_err - - # --- persist --- - try: - apply_fn(content) - await commit_fn() - except Exception as e: - logger.exception("Failed to update %s: %s", label, e) - await rollback_fn() - return {"status": "error", "message": f"Failed to update {label}: {e}"} - - # --- build response --- - resp: dict[str, Any] = { - "status": "saved", - "message": f"{label.capitalize()} updated.", - } - - if content is not updated_memory: - resp["notice"] = "Memory was automatically rewritten to fit within limits." - - diff_warnings = _validate_diff(old_memory, content) - if diff_warnings: - resp["diff_warnings"] = diff_warnings - - format_warnings = _validate_bullet_format(content) - if format_warnings: - resp["format_warnings"] = format_warnings - - warning = _soft_warning(content) - if warning: - resp["warning"] = warning - - return resp - - -# --------------------------------------------------------------------------- -# Tool factories -# --------------------------------------------------------------------------- - def create_update_memory_tool( user_id: str | UUID, db_session: AsyncSession, llm: Any | None = None, ): - """Factory function to create the user-memory update tool. + """Factory for the user-memory update tool. - The tool acquires its own short-lived ``AsyncSession`` per call via - :data:`async_session_maker` so the closure is safe to share across - HTTP requests by the compiled-agent cache. Capturing a per-request - session here would surface stale/closed sessions on cache hits. - The session's bound ``commit``/``rollback`` methods are captured at - call time, after ``async with`` has bound ``db_session`` locally. - - Args: - user_id: ID of the user whose memory document is being updated. - db_session: Reserved for registry compatibility. Per-call sessions - are opened via :data:`async_session_maker` inside the tool body. - llm: Optional LLM for the forced-rewrite path. - - Returns: - Configured update_memory tool for the user-memory scope. + Uses a fresh short-lived session per call so compiled-agent caches never + retain a stale request-scoped session. """ - del db_session # per-call session — see docstring + del db_session uid = UUID(user_id) if isinstance(user_id, str) else user_id @tool async def update_memory(updated_memory: str) -> dict[str, Any]: """Update the user's personal memory document. - Your current memory is shown in in the system prompt. - When the user shares important long-term information (preferences, - facts, instructions, context), rewrite the memory document to include - the new information. Merge new facts with existing ones, update - contradictions, remove outdated entries, and keep it concise. - - Args: - updated_memory: The FULL updated markdown document (not a diff). + The current memory is shown in . Pass the FULL updated + markdown document, not a diff. """ try: async with async_session_maker() as db_session: - result = await db_session.execute(select(User).where(User.id == uid)) - user = result.scalars().first() - if not user: - return {"status": "error", "message": "User not found."} - - old_memory = user.memory_md - - return await _save_memory( - updated_memory=updated_memory, - old_memory=old_memory, + result = await save_memory( + scope=MemoryScope.USER, + target_id=uid, + content=updated_memory, + session=db_session, llm=llm, - apply_fn=lambda content: setattr(user, "memory_md", content), - commit_fn=db_session.commit, - rollback_fn=db_session.rollback, - label="memory", - scope="user", ) + return result.to_dict() except Exception as e: logger.exception("Failed to update user memory: %s", e) - return { - "status": "error", - "message": f"Failed to update memory: {e}", - } + return {"status": "error", "message": f"Failed to update memory: {e}"} return update_memory @@ -373,64 +57,26 @@ def create_update_team_memory_tool( db_session: AsyncSession, llm: Any | None = None, ): - """Factory function to create the team-memory update tool. - - The tool acquires its own short-lived ``AsyncSession`` per call via - :data:`async_session_maker` so the closure is safe to share across - HTTP requests by the compiled-agent cache. Capturing a per-request - session here would surface stale/closed sessions on cache hits. - The session's bound ``commit``/``rollback`` methods are captured at - call time, after ``async with`` has bound ``db_session`` locally. - - Args: - search_space_id: ID of the search space whose team memory is being - updated. - db_session: Reserved for registry compatibility. Per-call sessions - are opened via :data:`async_session_maker` inside the tool body. - llm: Optional LLM for the forced-rewrite path. - - Returns: - Configured update_memory tool for the team-memory scope. - """ - del db_session # per-call session — see docstring + """Factory for the team-memory update tool.""" + del db_session @tool async def update_memory(updated_memory: str) -> dict[str, Any]: """Update the team's shared memory document for this search space. - Your current team memory is shown in in the system - prompt. When the team shares important long-term information - (decisions, conventions, key facts, priorities), rewrite the memory - document to include the new information. Merge new facts with - existing ones, update contradictions, remove outdated entries, and - keep it concise. - - Args: - updated_memory: The FULL updated markdown document (not a diff). + The current team memory is shown in . Pass the FULL updated + markdown document, not a diff. """ try: async with async_session_maker() as db_session: - result = await db_session.execute( - select(SearchSpace).where(SearchSpace.id == search_space_id) - ) - space = result.scalars().first() - if not space: - return {"status": "error", "message": "Search space not found."} - - old_memory = space.shared_memory_md - - return await _save_memory( - updated_memory=updated_memory, - old_memory=old_memory, + result = await save_memory( + scope=MemoryScope.TEAM, + target_id=search_space_id, + content=updated_memory, + session=db_session, llm=llm, - apply_fn=lambda content: setattr( - space, "shared_memory_md", content - ), - commit_fn=db_session.commit, - rollback_fn=db_session.rollback, - label="team memory", - scope="team", ) + return result.to_dict() except Exception as e: logger.exception("Failed to update team memory: %s", e) return { @@ -439,3 +85,9 @@ def create_update_team_memory_tool( } return update_memory + + +__all__ = [ + "create_update_memory_tool", + "create_update_team_memory_tool", +] diff --git a/surfsense_backend/app/services/memory/__init__.py b/surfsense_backend/app/services/memory/__init__.py new file mode 100644 index 000000000..d72f45e1f --- /dev/null +++ b/surfsense_backend/app/services/memory/__init__.py @@ -0,0 +1,29 @@ +"""First-class memory service for user and team markdown memory.""" + +from .service import ( + MemoryScope, + SaveResult, + extract_and_save, + read_memory, + reset_memory, + save_memory, +) +from .validation import ( + MEMORY_HARD_LIMIT, + MEMORY_SOFT_LIMIT, + validate_bullet_format, + validate_memory_scope, +) + +__all__ = [ + "MEMORY_HARD_LIMIT", + "MEMORY_SOFT_LIMIT", + "MemoryScope", + "SaveResult", + "extract_and_save", + "read_memory", + "reset_memory", + "save_memory", + "validate_bullet_format", + "validate_memory_scope", +] diff --git a/surfsense_backend/app/services/memory/prompts.py b/surfsense_backend/app/services/memory/prompts.py new file mode 100644 index 000000000..fbf27fd08 --- /dev/null +++ b/surfsense_backend/app/services/memory/prompts.py @@ -0,0 +1,110 @@ +"""Prompts used by the memory service.""" + +FORCED_REWRITE_PROMPT = """\ +You are a memory curator. The following memory document exceeds the character \ +limit and must be shortened. + +RULES: +1. Rewrite the document to be under {target} characters. +2. Output Markdown only. Use clear `##` headings and concise bullet points. +3. New-format bullets should look like: `- YYYY-MM-DD: memory text`. +4. If the input contains legacy markers like `(YYYY-MM-DD) [fact]`, preserve the + information but remove the inline marker in the output. +5. Preserve durable instructions and preferences before generic facts when + compressing personal memory. +6. Preserve existing headings when useful; merge duplicate headings and bullets. +7. Output ONLY the consolidated markdown — no explanations, no wrapping. + + +{content} +""" + +USER_MEMORY_EXTRACT_PROMPT = """\ +You are a memory extraction assistant. Analyze the user's message and decide \ +if it contains any long-term information worth persisting to personal memory. + +Worth remembering: preferences, background/identity, goals, projects, \ +instructions, tools/languages they use, decisions, expertise, workplace — \ +durable facts that will matter in future conversations. + +NOT worth remembering: greetings, one-off factual questions, session \ +logistics, ephemeral requests, follow-up clarifications with no new personal \ +info, things that only matter for the current task. + +If there is nothing durable to remember, choose `action = no_update`. + +If the message contains memorizable information, choose `action = save` and \ +return the FULL updated memory document with the new information merged into \ +existing content. + +FORMAT RULES FOR `updated_memory`: +- Markdown only. +- Every entry should be under a `##` heading. +- Recommended headings: `## Facts`, `## Preferences`, `## Instructions`. +- New bullets should use: `- YYYY-MM-DD: memory text`. +- If current memory uses legacy `(YYYY-MM-DD) [fact|pref|instr]` markers, + preserve the information but write the updated document in the new + heading-based format. +- Use the user's first name from `` when helpful, not "the user". +- Do not duplicate existing information. + +{user_name} + + +{current_memory} + + + +{user_message} +""" + +TEAM_MEMORY_EXTRACT_PROMPT = """\ +You are a team-memory extraction assistant. Analyze the latest message and \ +decide if it contains durable TEAM-level information worth persisting. + +Decision policy: +- Prioritize recall for durable team context, while avoiding personal-only facts. +- Do NOT require explicit consensus language. A direct team-level statement can + be stored if it is stable and broadly useful for future team chats. +- If evidence is weak or clearly tentative, choose `action = no_update`. + +Worth remembering (team-level only): +- Decisions and defaults that guide future team work +- Team conventions/standards (naming, review policy, coding norms) +- Stable org/project facts (locations, ownership, constraints) +- Long-lived architecture/process facts +- Ongoing priorities that are likely relevant beyond this turn + +NOT worth remembering: +- Personal preferences or biography of one person +- Questions, brainstorming, tentative ideas, or speculation +- One-off requests, status updates, TODOs, logistics for this session +- Information scoped only to a single ephemeral task + +If the message contains memorizable team information, choose `action = save` \ +and return the FULL updated team memory document with new facts merged into \ +existing content. + +FORMAT RULES FOR `updated_memory`: +- Markdown only. +- Every entry should be under a `##` heading. +- Recommended headings: `## Product Decisions`, `## Engineering Conventions`, + `## Project Facts`, `## Open Questions`. +- New bullets should use: `- YYYY-MM-DD: memory text`. +- If current memory uses legacy `(YYYY-MM-DD) [fact]` markers, preserve the + information but write the updated document in the new heading-based format. +- Do not create personal headings such as `## Preferences`, `## Instructions`, + or `## Personal Notes`. +- Preserve neutral team phrasing; avoid person-specific memory unless role-anchored. + + +{current_memory} + + + +{author} + + + +{user_message} +""" diff --git a/surfsense_backend/app/services/memory/rewrite.py b/surfsense_backend/app/services/memory/rewrite.py new file mode 100644 index 000000000..270904ce7 --- /dev/null +++ b/surfsense_backend/app/services/memory/rewrite.py @@ -0,0 +1,35 @@ +"""LLM-backed memory rewrite helpers.""" + +from __future__ import annotations + +import logging +from typing import Any + +from langchain_core.messages import HumanMessage + +from app.services.memory.prompts import FORCED_REWRITE_PROMPT +from app.services.memory.validation import MEMORY_HARD_LIMIT +from app.utils.content_utils import extract_text_content + +logger = logging.getLogger(__name__) + + +async def forced_rewrite(content: str, llm: Any) -> str | None: + """Use a focused LLM call to compress memory under the hard limit.""" + try: + prompt = FORCED_REWRITE_PROMPT.format( + target=MEMORY_HARD_LIMIT, + content=content, + ) + response = await llm.ainvoke( + [HumanMessage(content=prompt)], + config={"tags": ["surfsense:internal", "memory-rewrite"]}, + ) + text = extract_text_content(response.content).strip() + if not text: + logger.warning("Forced memory rewrite returned empty text") + return None + return text + except Exception: + logger.exception("Forced memory rewrite LLM call failed") + return None diff --git a/surfsense_backend/app/services/memory/schemas.py b/surfsense_backend/app/services/memory/schemas.py new file mode 100644 index 000000000..9b40ee5b1 --- /dev/null +++ b/surfsense_backend/app/services/memory/schemas.py @@ -0,0 +1,23 @@ +"""Structured output schemas for memory extraction.""" + +from __future__ import annotations + +from typing import Literal + +from pydantic import BaseModel, Field + + +class MemoryExtractionDecision(BaseModel): + """Structured extraction result; avoids string sentinel parsing.""" + + action: Literal["no_update", "save"] = Field( + description="Choose no_update when nothing durable should be saved; choose save otherwise." + ) + reason: str | None = Field( + default=None, + description="Short reason for no_update, or brief summary of the memory update.", + ) + updated_memory: str | None = Field( + default=None, + description="The full updated markdown memory document when action is save.", + ) diff --git a/surfsense_backend/app/services/memory/service.py b/surfsense_backend/app/services/memory/service.py new file mode 100644 index 000000000..85459c28c --- /dev/null +++ b/surfsense_backend/app/services/memory/service.py @@ -0,0 +1,300 @@ +"""Canonical read/write/reset/extract service for markdown memory.""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass, field +from enum import StrEnum +from typing import Any, Literal +from uuid import UUID + +from langchain_core.messages import HumanMessage +from pydantic import BaseModel +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from app.db import SearchSpace, User +from app.services.memory.prompts import ( + TEAM_MEMORY_EXTRACT_PROMPT, + USER_MEMORY_EXTRACT_PROMPT, +) +from app.services.memory.rewrite import forced_rewrite +from app.services.memory.schemas import MemoryExtractionDecision +from app.services.memory.validation import ( + MEMORY_HARD_LIMIT, + soft_limit_warning, + strip_preamble_to_first_heading, + validate_bullet_format, + validate_diff, + validate_heading_sanity, + validate_memory_scope, + validate_memory_size, +) + +logger = logging.getLogger(__name__) + + +class MemoryScope(StrEnum): + USER = "user" + TEAM = "team" + + +@dataclass(frozen=True) +class SaveResult: + status: Literal["saved", "error", "no_op"] + message: str + memory_md: str = "" + warnings: list[str] = field(default_factory=list) + diff_warnings: list[str] = field(default_factory=list) + format_warnings: list[str] = field(default_factory=list) + notice: str | None = None + + def to_dict(self) -> dict[str, Any]: + data: dict[str, Any] = { + "status": self.status, + "message": self.message, + "memory_md": self.memory_md, + } + if self.notice: + data["notice"] = self.notice + if self.warnings: + data["warnings"] = self.warnings + if len(self.warnings) == 1: + data["warning"] = self.warnings[0] + if self.diff_warnings: + data["diff_warnings"] = self.diff_warnings + if self.format_warnings: + data["format_warnings"] = self.format_warnings + return data + + +class MemoryRead(BaseModel): + memory_md: str + + +def _normalize_scope(scope: MemoryScope | str) -> MemoryScope: + return scope if isinstance(scope, MemoryScope) else MemoryScope(scope) + + +def _normalize_user_id(target_id: str | UUID) -> UUID: + return UUID(target_id) if isinstance(target_id, str) else target_id + + +async def _load_target( + *, + scope: MemoryScope | str, + target_id: str | int | UUID, + session: AsyncSession, +) -> User | SearchSpace | None: + normalized = _normalize_scope(scope) + if normalized is MemoryScope.USER: + result = await session.execute( + select(User).where(User.id == _normalize_user_id(target_id)) # type: ignore[arg-type] + ) + return result.scalars().first() + result = await session.execute(select(SearchSpace).where(SearchSpace.id == int(target_id))) + return result.scalars().first() + + +def _get_memory(target: User | SearchSpace, scope: MemoryScope) -> str: + if scope is MemoryScope.USER: + return getattr(target, "memory_md", None) or "" + return getattr(target, "shared_memory_md", None) or "" + + +def _set_memory(target: User | SearchSpace, scope: MemoryScope, content: str) -> None: + if scope is MemoryScope.USER: + target.memory_md = content + else: + target.shared_memory_md = content + + +async def read_memory( + *, + scope: MemoryScope | str, + target_id: str | int | UUID, + session: AsyncSession, +) -> str: + normalized = _normalize_scope(scope) + target = await _load_target(scope=normalized, target_id=target_id, session=session) + if target is None: + return "" + return _get_memory(target, normalized) + + +async def save_memory( + *, + scope: MemoryScope | str, + target_id: str | int | UUID, + content: str, + session: AsyncSession, + llm: Any | None = None, +) -> SaveResult: + normalized = _normalize_scope(scope) + if not isinstance(content, str): + return SaveResult( + status="error", + message="Internal error: memory payload must be a string.", + ) + + target = await _load_target(scope=normalized, target_id=target_id, session=session) + if target is None: + return SaveResult( + status="error", + message="User not found." if normalized is MemoryScope.USER else "Search space not found.", + ) + + old_memory = _get_memory(target, normalized) + next_content = strip_preamble_to_first_heading(content.strip()) + notice: str | None = None + warnings: list[str] = [] + + if len(next_content) > MEMORY_HARD_LIMIT and llm is not None: + rewritten = await forced_rewrite(next_content, llm) + if rewritten is not None and len(rewritten) < len(next_content): + next_content = strip_preamble_to_first_heading(rewritten) + notice = "Memory was automatically rewritten to fit within limits." + + for validation in ( + validate_memory_size(next_content), + validate_heading_sanity(next_content), + ): + if validation: + return SaveResult( + status="error", + message=validation["message"], + memory_md=old_memory, + ) + + scope_error, scope_warnings = validate_memory_scope( + next_content, + normalized.value, + old_memory=old_memory, + ) + warnings.extend(scope_warnings) + if scope_error: + return SaveResult( + status="error", + message=scope_error["message"], + memory_md=old_memory, + warnings=warnings, + ) + + try: + _set_memory(target, normalized, next_content) + session.add(target) + await session.commit() + except Exception as e: + logger.exception("Failed to update %s memory: %s", normalized.value, e) + await session.rollback() + return SaveResult( + status="error", + message=f"Failed to update {normalized.value} memory: {e}", + memory_md=old_memory, + ) + + diff_warnings = validate_diff(old_memory, next_content) + format_warnings = validate_bullet_format(next_content) + warning = soft_limit_warning(next_content) + if warning: + warnings.append(warning) + + return SaveResult( + status="saved", + message=( + "Memory updated." + if normalized is MemoryScope.USER + else "Team memory updated." + ), + memory_md=next_content, + warnings=warnings, + diff_warnings=diff_warnings, + format_warnings=format_warnings, + notice=notice, + ) + + +async def reset_memory( + *, + scope: MemoryScope | str, + target_id: str | int | UUID, + session: AsyncSession, +) -> SaveResult: + return await save_memory( + scope=scope, + target_id=target_id, + content="", + session=session, + llm=None, + ) + + +async def extract_and_save( + *, + scope: MemoryScope | str, + target_id: str | int | UUID, + user_message: str, + actor_display_name: str | None, + session: AsyncSession, + llm: Any, +) -> SaveResult: + normalized = _normalize_scope(scope) + current_memory = await read_memory( + scope=normalized, + target_id=target_id, + session=session, + ) + + if normalized is MemoryScope.USER: + first_name = ( + actor_display_name.strip().split()[0] + if actor_display_name and actor_display_name.strip() + else "The user" + ) + prompt = USER_MEMORY_EXTRACT_PROMPT.format( + current_memory=current_memory or "(empty)", + user_message=user_message, + user_name=first_name, + ) + else: + prompt = TEAM_MEMORY_EXTRACT_PROMPT.format( + current_memory=current_memory or "(empty)", + author=actor_display_name or "Unknown team member", + user_message=user_message, + ) + + try: + structured = llm.with_structured_output(MemoryExtractionDecision) + decision = await structured.ainvoke( + [HumanMessage(content=prompt)], + config={"tags": ["surfsense:internal", "memory-extraction"]}, + ) + except Exception: + logger.exception("Structured memory extraction failed") + return SaveResult( + status="error", + message="Structured memory extraction failed.", + memory_md=current_memory, + ) + + if decision.action == "no_update": + return SaveResult( + status="no_op", + message=decision.reason or "No durable memory to persist.", + memory_md=current_memory, + ) + + if not decision.updated_memory: + return SaveResult( + status="error", + message="Structured memory extraction chose save without updated_memory.", + memory_md=current_memory, + ) + + return await save_memory( + scope=normalized, + target_id=target_id, + content=decision.updated_memory, + session=session, + llm=llm, + ) diff --git a/surfsense_backend/app/services/memory/validation.py b/surfsense_backend/app/services/memory/validation.py new file mode 100644 index 000000000..0e856943b --- /dev/null +++ b/surfsense_backend/app/services/memory/validation.py @@ -0,0 +1,158 @@ +"""Validation helpers for markdown-backed memory.""" + +from __future__ import annotations + +import re +from typing import Literal + +MEMORY_SOFT_LIMIT = 18_000 +MEMORY_HARD_LIMIT = 25_000 + +_SECTION_HEADING_RE = re.compile(r"^##\s+(.+)$", re.MULTILINE) +_HEADING_LINE_RE = re.compile(r"^##\s+\S+", re.MULTILINE) +_HEADING_NORMALIZE_RE = re.compile(r"[^a-z0-9]+") +_LEGACY_BULLET_RE = re.compile(r"^-\s+\(\d{4}-\d{2}-\d{2}\)\s+\[(fact|pref|instr)\]\s+.+$") +_NEW_BULLET_RE = re.compile(r"^-\s+\d{4}-\d{2}-\d{2}:\s+.+$") + +_FORBIDDEN_TEAM_HEADINGS = { + "preferences", + "instructions", + "personal notes", + "personal instructions", +} + + +def has_markdown_heading(content: str) -> bool: + return bool(_HEADING_LINE_RE.search(content)) + + +def strip_preamble_to_first_heading(content: str) -> str: + """Drop model preamble before the first ``##`` heading, if one exists.""" + match = _HEADING_LINE_RE.search(content) + if not match: + return content.strip() + return content[match.start() :].strip() + + +def extract_headings(memory: str | None) -> set[str]: + if not memory: + return set() + return {_normalize_heading(h) for h in _SECTION_HEADING_RE.findall(memory)} + + +def _normalize_heading(heading: str) -> str: + return _HEADING_NORMALIZE_RE.sub(" ", heading.strip().lower()).strip() + + +def validate_memory_size(content: str) -> dict[str, str] | None: + length = len(content) + if length > MEMORY_HARD_LIMIT: + return { + "status": "error", + "message": ( + f"Memory exceeds {MEMORY_HARD_LIMIT:,} character limit " + f"({length:,} chars). Consolidate by merging related items, " + "removing outdated entries, and shortening descriptions." + ), + } + return None + + +def validate_heading_sanity(content: str) -> dict[str, str] | None: + """Block long prose blobs without headings unless they are legacy bullets.""" + stripped = content.strip() + if not stripped: + return None + if has_markdown_heading(stripped): + return None + if len(stripped) <= 40: + return None + if any(_LEGACY_BULLET_RE.match(line.strip()) for line in stripped.splitlines()): + return None + return { + "status": "error", + "message": "Memory must be markdown with at least one ## heading.", + } + + +def validate_memory_scope( + content: str, + scope: Literal["user", "team"], + *, + old_memory: str | None = None, +) -> tuple[dict[str, str] | None, list[str]]: + """Reject new personal headings in team memory, grandfather existing ones.""" + if scope != "team": + return None, [] + + old_forbidden = extract_headings(old_memory) & _FORBIDDEN_TEAM_HEADINGS + new_forbidden = extract_headings(content) & _FORBIDDEN_TEAM_HEADINGS + introduced = sorted(new_forbidden - old_forbidden) + grandfathered = sorted(new_forbidden & old_forbidden) + + warnings: list[str] = [] + if grandfathered: + warnings.append( + "Team memory contains legacy personal headings: " + + ", ".join(grandfathered) + + ". Please consolidate them into team-safe headings." + ) + if introduced: + return ( + { + "status": "error", + "message": ( + "Team memory cannot introduce personal headings: " + + ", ".join(introduced) + + ". Use team-safe headings instead." + ), + }, + warnings, + ) + return None, warnings + + +def validate_bullet_format(content: str) -> list[str]: + warnings: list[str] = [] + for line in content.splitlines(): + stripped = line.strip() + if not stripped.startswith("- "): + continue + if _NEW_BULLET_RE.match(stripped) or _LEGACY_BULLET_RE.match(stripped): + continue + short = stripped[:80] + ("..." if len(stripped) > 80 else "") + warnings.append(f"Non-standard memory bullet: {short}") + return warnings + + +def validate_diff(old_memory: str | None, new_memory: str) -> list[str]: + if not old_memory: + return [] + + warnings: list[str] = [] + old_headings = extract_headings(old_memory) + new_headings = extract_headings(new_memory) + dropped = old_headings - new_headings + if dropped: + names = ", ".join(sorted(dropped)) + warnings.append( + f"Sections removed: {names}. If unintentional, restore from the settings page." + ) + + old_len = len(old_memory) + new_len = len(new_memory) + if old_len > 0 and new_len < old_len * 0.4: + warnings.append( + f"Memory shrank significantly ({old_len:,} -> {new_len:,} chars). Possible data loss." + ) + return warnings + + +def soft_limit_warning(content: str) -> str | None: + length = len(content) + if length > MEMORY_SOFT_LIMIT: + return ( + f"Memory is at {length:,}/{MEMORY_HARD_LIMIT:,} characters. " + "Consolidate by merging related items and removing less important entries." + ) + return None diff --git a/surfsense_backend/tests/unit/services/test_memory_service.py b/surfsense_backend/tests/unit/services/test_memory_service.py new file mode 100644 index 000000000..c16e34062 --- /dev/null +++ b/surfsense_backend/tests/unit/services/test_memory_service.py @@ -0,0 +1,204 @@ +"""Unit tests for the first-class memory service.""" + +from types import SimpleNamespace + +import pytest + +from app.services.memory import ( + MemoryScope, + extract_and_save, + reset_memory, + save_memory, +) +from app.services.memory.schemas import MemoryExtractionDecision + +pytestmark = pytest.mark.unit + + +class _FakeSession: + def __init__(self) -> None: + self.commit_calls = 0 + self.rollback_calls = 0 + self.added = [] + + def add(self, obj) -> None: + self.added.append(obj) + + async def commit(self) -> None: + self.commit_calls += 1 + + async def rollback(self) -> None: + self.rollback_calls += 1 + + +class _StructuredLLM: + def __init__(self, decision: MemoryExtractionDecision) -> None: + self.decision = decision + + def with_structured_output(self, _schema): + return self + + async def ainvoke(self, *_args, **_kwargs): + return self.decision + + +@pytest.mark.asyncio +async def test_save_memory_saves_heading_based_memory(monkeypatch) -> None: + target = SimpleNamespace(memory_md="") + session = _FakeSession() + + async def fake_load_target(**_kwargs): + return target + + monkeypatch.setattr("app.services.memory.service._load_target", fake_load_target) + + result = await save_memory( + scope=MemoryScope.USER, + target_id="00000000-0000-0000-0000-000000000000", + content="## Facts\n- 2026-05-19: Anish works on SurfSense\n", + session=session, + ) + + assert result.status == "saved" + assert target.memory_md.startswith("## Facts") + assert session.commit_calls == 1 + + +@pytest.mark.asyncio +async def test_save_memory_accepts_legacy_marker_payload(monkeypatch) -> None: + target = SimpleNamespace(memory_md="") + session = _FakeSession() + + async def fake_load_target(**_kwargs): + return target + + monkeypatch.setattr("app.services.memory.service._load_target", fake_load_target) + + result = await save_memory( + scope=MemoryScope.USER, + target_id="00000000-0000-0000-0000-000000000000", + content="- (2026-05-19) [fact] Legacy marker memory\n", + session=session, + ) + + assert result.status == "saved" + assert "[fact]" in target.memory_md + + +@pytest.mark.asyncio +async def test_save_memory_rejects_long_no_heading_payload(monkeypatch) -> None: + target = SimpleNamespace(memory_md="## Facts\n- 2026-05-19: Existing\n") + session = _FakeSession() + + async def fake_load_target(**_kwargs): + return target + + monkeypatch.setattr("app.services.memory.service._load_target", fake_load_target) + + result = await save_memory( + scope=MemoryScope.USER, + target_id="00000000-0000-0000-0000-000000000000", + content="reasoning text before NO_UPDATE should not become saved memory", + session=session, + ) + + assert result.status == "error" + assert session.commit_calls == 0 + assert target.memory_md.startswith("## Facts") + + +@pytest.mark.asyncio +async def test_save_memory_grandfathers_existing_team_personal_heading(monkeypatch) -> None: + content = "## Preferences\n- 2026-05-19: Existing legacy heading\n" + target = SimpleNamespace(shared_memory_md=content) + session = _FakeSession() + + async def fake_load_target(**_kwargs): + return target + + monkeypatch.setattr("app.services.memory.service._load_target", fake_load_target) + + result = await save_memory( + scope=MemoryScope.TEAM, + target_id=1, + content=content, + session=session, + ) + + assert result.status == "saved" + assert result.warnings + assert session.commit_calls == 1 + + +@pytest.mark.asyncio +async def test_reset_memory_clears_memory(monkeypatch) -> None: + target = SimpleNamespace(memory_md="## Facts\n- 2026-05-19: Existing\n") + session = _FakeSession() + + async def fake_load_target(**_kwargs): + return target + + monkeypatch.setattr("app.services.memory.service._load_target", fake_load_target) + + result = await reset_memory( + scope=MemoryScope.USER, + target_id="00000000-0000-0000-0000-000000000000", + session=session, + ) + + assert result.status == "saved" + assert target.memory_md == "" + + +@pytest.mark.asyncio +async def test_extract_and_save_no_update_does_not_commit(monkeypatch) -> None: + target = SimpleNamespace(memory_md="## Facts\n- 2026-05-19: Existing\n") + session = _FakeSession() + + async def fake_load_target(**_kwargs): + return target + + monkeypatch.setattr("app.services.memory.service._load_target", fake_load_target) + + result = await extract_and_save( + scope=MemoryScope.USER, + target_id="00000000-0000-0000-0000-000000000000", + user_message="hello", + actor_display_name="Anish", + session=session, + llm=_StructuredLLM( + MemoryExtractionDecision(action="no_update", reason="Greeting only") + ), + ) + + assert result.status == "no_op" + assert session.commit_calls == 0 + + +@pytest.mark.asyncio +async def test_extract_and_save_persists_structured_update(monkeypatch) -> None: + target = SimpleNamespace(memory_md="") + session = _FakeSession() + + async def fake_load_target(**_kwargs): + return target + + monkeypatch.setattr("app.services.memory.service._load_target", fake_load_target) + + result = await extract_and_save( + scope=MemoryScope.USER, + target_id="00000000-0000-0000-0000-000000000000", + user_message="I work on SurfSense", + actor_display_name="Anish", + session=session, + llm=_StructuredLLM( + MemoryExtractionDecision( + action="save", + updated_memory="## Facts\n- 2026-05-19: Anish works on SurfSense\n", + ) + ), + ) + + assert result.status == "saved" + assert "SurfSense" in target.memory_md + assert session.commit_calls == 1 From 5247dc709708556a5be7a1a61bef9444b43248e6 Mon Sep 17 00:00:00 2001 From: Anish Sarkar <104695310+AnishSarkar22@users.noreply.github.com> Date: Wed, 20 May 2026 02:02:10 +0530 Subject: [PATCH 02/15] feat: refine private and team memory protocols --- .../prompts/memory_protocol/private.md | 6 + .../prompts/memory_protocol/team.md | 8 + .../update_memory/private/description.md | 10 +- .../tools/update_memory/private/example.md | 10 +- .../tools/update_memory/team/description.md | 16 +- .../tools/update_memory/team/example.md | 4 +- .../builtins/memory/system_prompt.md | 7 + .../new_chat/middleware/memory_injection.py | 2 +- .../prompts/base/memory_protocol_private.md | 6 + .../prompts/base/memory_protocol_team.md | 8 + .../prompts/examples/update_memory_private.md | 14 +- .../prompts/examples/update_memory_team.md | 4 +- .../prompts/tools/update_memory_private.md | 51 ++-- .../prompts/tools/update_memory_team.md | 48 ++-- .../new_chat/test_memory_response_content.py | 41 +-- .../tools/test_update_memory_scope.py | 261 +++++++----------- 16 files changed, 232 insertions(+), 264 deletions(-) diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/system_prompt/prompts/memory_protocol/private.md b/surfsense_backend/app/agents/multi_agent_chat/main_agent/system_prompt/prompts/memory_protocol/private.md index 4dd511014..bcb80f0f4 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/main_agent/system_prompt/prompts/memory_protocol/private.md +++ b/surfsense_backend/app/agents/multi_agent_chat/main_agent/system_prompt/prompts/memory_protocol/private.md @@ -6,4 +6,10 @@ standing instructions? If yes, call `update_memory` **alongside** your normal response — don't defer it to a later turn. Skip ephemeral chat noise (one-off Q/A, greetings, session logistics). Stay within the budget shown in ``. + +Memory is heading-based markdown. New entries should be under `##` headings +such as `## Facts`, `## Preferences`, or `## Instructions`, with bullets like +`- YYYY-MM-DD: text`. If existing memory contains legacy +`(YYYY-MM-DD) [fact|pref|instr]` markers, preserve the information but write +new saves in the heading-based format. diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/system_prompt/prompts/memory_protocol/team.md b/surfsense_backend/app/agents/multi_agent_chat/main_agent/system_prompt/prompts/memory_protocol/team.md index decd23c4d..14d9a6793 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/main_agent/system_prompt/prompts/memory_protocol/team.md +++ b/surfsense_backend/app/agents/multi_agent_chat/main_agent/system_prompt/prompts/memory_protocol/team.md @@ -6,4 +6,12 @@ key facts? If yes, call `update_memory` **alongside** your normal response — don't defer it to a later turn. Skip ephemeral chat noise (one-off Q/A, greetings, session logistics). Stay within the budget shown in ``. + +Team memory is heading-based markdown. New entries should be under `##` +headings such as `## Product Decisions`, `## Engineering Conventions`, +`## Project Facts`, or `## Open Questions`, with bullets like +`- YYYY-MM-DD: text`. If existing memory contains legacy `(YYYY-MM-DD) [fact]` +markers, preserve the information but write new saves in the heading-based +format. Do not create personal headings such as `## Preferences` or +`## Instructions`. diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/system_prompt/prompts/tools/update_memory/private/description.md b/surfsense_backend/app/agents/multi_agent_chat/main_agent/system_prompt/prompts/tools/update_memory/private/description.md index e7fa842b1..01169ff60 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/main_agent/system_prompt/prompts/tools/update_memory/private/description.md +++ b/surfsense_backend/app/agents/multi_agent_chat/main_agent/system_prompt/prompts/tools/update_memory/private/description.md @@ -9,7 +9,9 @@ - Skip ephemeral chat noise (one-off Q/A, greetings, session logistics). - Args: `updated_memory` — FULL replacement markdown (merge and curate, don't only append). - - Formatting: bullets `- (YYYY-MM-DD) [marker] text` with markers `[fact]`, - `[pref]`, `[instr]` (priority when trimming: `instr > pref > fact`). - Group bullets under short `##` headings; stay under the limit shown in - ``. + - Formatting: heading-based markdown with entries under `##` headings. + Recommended headings are `## Facts`, `## Preferences`, `## Instructions`, + though clearer natural headings are allowed. New bullets should look like + `- YYYY-MM-DD: text`; stay under the limit shown in ``. + - If existing memory uses legacy `(YYYY-MM-DD) [fact|pref|instr]` markers, + preserve the information but write the updated document in the new format. diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/system_prompt/prompts/tools/update_memory/private/example.md b/surfsense_backend/app/agents/multi_agent_chat/main_agent/system_prompt/prompts/tools/update_memory/private/example.md index 2505bdf87..9afadb02c 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/main_agent/system_prompt/prompts/tools/update_memory/private/example.md +++ b/surfsense_backend/app/agents/multi_agent_chat/main_agent/system_prompt/prompts/tools/update_memory/private/example.md @@ -1,28 +1,28 @@ Alex, is empty. user: "I'm a space enthusiast, explain astrophage to me" -→ update_memory(updated_memory="## Interests & background\n- (2025-03-15) [fact] Alex is a space enthusiast\n") +→ update_memory(updated_memory="## Facts\n- 2025-03-15: Alex is a space enthusiast\n") (Casual durable fact; use first name, neutral heading.) user: "Remember that I prefer concise answers over detailed explanations" -→ update_memory(updated_memory="## Interests & background\n- (2025-03-15) [fact] Alex is a space enthusiast\n\n## Response style\n- (2025-03-15) [pref] Alex prefers concise answers over detailed explanations\n") +→ update_memory(updated_memory="## Facts\n- 2025-03-15: Alex is a space enthusiast\n\n## Preferences\n- 2025-03-15: Alex prefers concise answers over detailed explanations\n") (Durable preference; merge with existing memory.) user: "I actually moved to Tokyo last month" -→ update_memory(updated_memory="...\n\n## Personal context\n- (2025-03-15) [fact] Alex lives in Tokyo (previously London)\n...") +→ update_memory(updated_memory="...\n\n## Facts\n- 2025-03-15: Alex lives in Tokyo (previously London)\n...") (Updated fact; date reflects when recorded.) user: "I'm a freelance photographer working on a nature documentary" -→ update_memory(updated_memory="...\n\n## Current focus\n- (2025-03-15) [fact] Alex is a freelance photographer\n- (2025-03-15) [fact] Alex is working on a nature documentary\n") +→ update_memory(updated_memory="...\n\n## Current Focus\n- 2025-03-15: Alex is a freelance photographer\n- 2025-03-15: Alex is working on a nature documentary\n") user: "Always respond in bullet points" -→ update_memory(updated_memory="...\n\n## Response style\n- (2025-03-15) [instr] Always respond to Alex in bullet points\n") +→ update_memory(updated_memory="...\n\n## Instructions\n- 2025-03-15: Always respond to Alex in bullet points\n") diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/system_prompt/prompts/tools/update_memory/team/description.md b/surfsense_backend/app/agents/multi_agent_chat/main_agent/system_prompt/prompts/tools/update_memory/team/description.md index 13341a910..8459f9e7a 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/main_agent/system_prompt/prompts/tools/update_memory/team/description.md +++ b/surfsense_backend/app/agents/multi_agent_chat/main_agent/system_prompt/prompts/tools/update_memory/team/description.md @@ -9,8 +9,14 @@ - Skip ephemeral chat noise (one-off Q/A, greetings, session logistics). - Args: `updated_memory` — FULL replacement markdown (merge and curate, don't only append). - - Formatting: bullets `- (YYYY-MM-DD) [fact] text`. Team memory uses ONLY - the `[fact]` marker (never `[pref]` or `[instr]`). Group bullets under - short `##` headings (2-3 words each); stay under the limit shown in - ``. When trimming, prioritise: decisions/conventions > key - facts > current priorities. + - Formatting: heading-based markdown with entries under `##` headings. + Recommended headings are `## Product Decisions`, + `## Engineering Conventions`, `## Project Facts`, and `## Open Questions`. + New bullets should look like `- YYYY-MM-DD: text`; stay under the limit + shown in ``. + - If existing memory uses legacy `(YYYY-MM-DD) [fact]` markers, preserve the + information but write the updated document in the new format. + - Do not create personal headings such as `## Preferences`, + `## Instructions`, `## Personal Notes`, or `## Personal Instructions`. + When trimming, prioritise: decisions/conventions > key facts > current + priorities. diff --git a/surfsense_backend/app/agents/multi_agent_chat/main_agent/system_prompt/prompts/tools/update_memory/team/example.md b/surfsense_backend/app/agents/multi_agent_chat/main_agent/system_prompt/prompts/tools/update_memory/team/example.md index 8bd8fcfe4..5d06d9a0c 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/main_agent/system_prompt/prompts/tools/update_memory/team/example.md +++ b/surfsense_backend/app/agents/multi_agent_chat/main_agent/system_prompt/prompts/tools/update_memory/team/example.md @@ -1,9 +1,9 @@ user: "Let's remember that we decided to do weekly standup meetings on Mondays" -→ update_memory(updated_memory="...\n\n## Team rituals\n- (2025-03-15) [fact] Weekly standup meetings on Mondays\n...") +→ update_memory(updated_memory="...\n\n## Product Decisions\n- 2025-03-15: Weekly standup meetings happen on Mondays\n...") user: "Our office is in downtown Seattle, 5th floor" -→ update_memory(updated_memory="...\n\n## Workspace\n- (2025-03-15) [fact] Office location: downtown Seattle, 5th floor\n...") +→ update_memory(updated_memory="...\n\n## Project Facts\n- 2025-03-15: Office location is downtown Seattle, 5th floor\n...") diff --git a/surfsense_backend/app/agents/multi_agent_chat/subagents/builtins/memory/system_prompt.md b/surfsense_backend/app/agents/multi_agent_chat/subagents/builtins/memory/system_prompt.md index 32becf233..13f7b68a5 100644 --- a/surfsense_backend/app/agents/multi_agent_chat/subagents/builtins/memory/system_prompt.md +++ b/surfsense_backend/app/agents/multi_agent_chat/subagents/builtins/memory/system_prompt.md @@ -18,6 +18,10 @@ Persist durable preferences/facts/instructions with `update_memory` while avoidi - Do not store transient chatter. - Do not store secrets unless explicitly instructed. - If memory intent is unclear, return `status=blocked` with the missing intent signal. +- Persisted memory is heading-based markdown. New saved bullets should look like + `- YYYY-MM-DD: text` under `##` headings. If existing memory has legacy + `(YYYY-MM-DD) [fact|pref|instr]` markers, preserve the information but write + the updated document in the heading-based format. @@ -53,4 +57,7 @@ Rules: - `status=success` -> `next_step=null`, `missing_fields=null`. - `status=partial|blocked|error` -> `next_step` must be non-null. - `status=blocked` due to missing required inputs -> `missing_fields` must be non-null. +- `evidence.memory_category` is a semantic classification for supervisor logs + only. It is not the persisted storage format and must not force inline + `[fact|preference|instruction]` markers into saved memory. diff --git a/surfsense_backend/app/agents/new_chat/middleware/memory_injection.py b/surfsense_backend/app/agents/new_chat/middleware/memory_injection.py index 6179adccd..1d447aa28 100644 --- a/surfsense_backend/app/agents/new_chat/middleware/memory_injection.py +++ b/surfsense_backend/app/agents/new_chat/middleware/memory_injection.py @@ -17,8 +17,8 @@ from langgraph.runtime import Runtime from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession -from app.agents.new_chat.tools.update_memory import MEMORY_HARD_LIMIT, MEMORY_SOFT_LIMIT from app.db import ChatVisibility, SearchSpace, User, shielded_async_session +from app.services.memory import MEMORY_HARD_LIMIT, MEMORY_SOFT_LIMIT logger = logging.getLogger(__name__) diff --git a/surfsense_backend/app/agents/new_chat/prompts/base/memory_protocol_private.md b/surfsense_backend/app/agents/new_chat/prompts/base/memory_protocol_private.md index 8f7da14f8..22fed418a 100644 --- a/surfsense_backend/app/agents/new_chat/prompts/base/memory_protocol_private.md +++ b/surfsense_backend/app/agents/new_chat/prompts/base/memory_protocol_private.md @@ -3,4 +3,10 @@ IMPORTANT — After understanding each user message, ALWAYS check: does this mes reveal durable facts about the user (role, interests, preferences, projects, background, or standing instructions)? If yes, you MUST call update_memory alongside your normal response — do not defer this to a later turn. + +Memory is stored as a heading-based markdown document. New entries should be +under `##` headings such as `## Facts`, `## Preferences`, or `## Instructions` +with bullets like `- YYYY-MM-DD: text`. If existing memory contains legacy +`(YYYY-MM-DD) [fact|pref|instr]` markers, preserve the information but write +new saves in the heading-based format. diff --git a/surfsense_backend/app/agents/new_chat/prompts/base/memory_protocol_team.md b/surfsense_backend/app/agents/new_chat/prompts/base/memory_protocol_team.md index 61d89cc5d..38ec798c0 100644 --- a/surfsense_backend/app/agents/new_chat/prompts/base/memory_protocol_team.md +++ b/surfsense_backend/app/agents/new_chat/prompts/base/memory_protocol_team.md @@ -3,4 +3,12 @@ IMPORTANT — After understanding each user message, ALWAYS check: does this mes reveal durable facts about the team (decisions, conventions, architecture, processes, or key facts)? If yes, you MUST call update_memory alongside your normal response — do not defer this to a later turn. + +Team memory is stored as a heading-based markdown document. New entries should +be under `##` headings such as `## Product Decisions`, +`## Engineering Conventions`, `## Project Facts`, or `## Open Questions` with +bullets like `- YYYY-MM-DD: text`. If existing memory contains legacy +`(YYYY-MM-DD) [fact]` markers, preserve the information but write new saves in +the heading-based format. Do not create personal headings such as +`## Preferences` or `## Instructions`. diff --git a/surfsense_backend/app/agents/new_chat/prompts/examples/update_memory_private.md b/surfsense_backend/app/agents/new_chat/prompts/examples/update_memory_private.md index f83fe40b4..496bdcae3 100644 --- a/surfsense_backend/app/agents/new_chat/prompts/examples/update_memory_private.md +++ b/surfsense_backend/app/agents/new_chat/prompts/examples/update_memory_private.md @@ -1,16 +1,16 @@ - Alex, is empty. User: "I'm a space enthusiast, explain astrophage to me" - - The user casually shared a durable fact. Use their first name in the entry, short neutral heading: - update_memory(updated_memory="## Interests & background\n- (2025-03-15) [fact] Alex is a space enthusiast\n") + - The user casually shared a durable fact: + update_memory(updated_memory="## Facts\n- 2025-03-15: Alex is a space enthusiast\n") - User: "Remember that I prefer concise answers over detailed explanations" - - Durable preference. Merge with existing memory, add a new heading: - update_memory(updated_memory="## Interests & background\n- (2025-03-15) [fact] Alex is a space enthusiast\n\n## Response style\n- (2025-03-15) [pref] Alex prefers concise answers over detailed explanations\n") + - Durable preference. Merge with existing memory: + update_memory(updated_memory="## Facts\n- 2025-03-15: Alex is a space enthusiast\n\n## Preferences\n- 2025-03-15: Alex prefers concise answers over detailed explanations\n") - User: "I actually moved to Tokyo last month" - Updated fact, date prefix reflects when recorded: - update_memory(updated_memory="## Interests & background\n...\n\n## Personal context\n- (2025-03-15) [fact] Alex lives in Tokyo (previously London)\n...") + update_memory(updated_memory="## Facts\n- 2025-03-15: Alex lives in Tokyo (previously London)\n...") - User: "I'm a freelance photographer working on a nature documentary" - Durable background info under a fitting heading: - update_memory(updated_memory="...\n\n## Current focus\n- (2025-03-15) [fact] Alex is a freelance photographer\n- (2025-03-15) [fact] Alex is working on a nature documentary\n") + update_memory(updated_memory="...\n\n## Current Focus\n- 2025-03-15: Alex is a freelance photographer\n- 2025-03-15: Alex is working on a nature documentary\n") - User: "Always respond in bullet points" - Standing instruction: - update_memory(updated_memory="...\n\n## Response style\n- (2025-03-15) [instr] Always respond to Alex in bullet points\n") + update_memory(updated_memory="...\n\n## Instructions\n- 2025-03-15: Always respond to Alex in bullet points\n") diff --git a/surfsense_backend/app/agents/new_chat/prompts/examples/update_memory_team.md b/surfsense_backend/app/agents/new_chat/prompts/examples/update_memory_team.md index 1c74fdf6e..16b90babf 100644 --- a/surfsense_backend/app/agents/new_chat/prompts/examples/update_memory_team.md +++ b/surfsense_backend/app/agents/new_chat/prompts/examples/update_memory_team.md @@ -1,7 +1,7 @@ - User: "Let's remember that we decided to do weekly standup meetings on Mondays" - Durable team decision: - update_memory(updated_memory="- (2025-03-15) [fact] Weekly standup meetings on Mondays\n...") + update_memory(updated_memory="## Product Decisions\n- 2025-03-15: Weekly standup meetings happen on Mondays\n...") - User: "Our office is in downtown Seattle, 5th floor" - Durable team fact: - update_memory(updated_memory="- (2025-03-15) [fact] Office location: downtown Seattle, 5th floor\n...") + update_memory(updated_memory="## Project Facts\n- 2025-03-15: Office location is downtown Seattle, 5th floor\n...") diff --git a/surfsense_backend/app/agents/new_chat/prompts/tools/update_memory_private.md b/surfsense_backend/app/agents/new_chat/prompts/tools/update_memory_private.md index 184013804..65de785e9 100644 --- a/surfsense_backend/app/agents/new_chat/prompts/tools/update_memory_private.md +++ b/surfsense_backend/app/agents/new_chat/prompts/tools/update_memory_private.md @@ -1,31 +1,26 @@ - update_memory: Update your personal memory document about the user. - - Your current memory is already in in your context. The `chars` and - `limit` attributes show your current usage and the maximum allowed size. - - This is your curated long-term memory — the distilled essence of what you know about - the user, not raw conversation logs. - - Call update_memory when: - * The user explicitly asks to remember or forget something - * The user shares durable facts or preferences that will matter in future conversations - - The user's first name is provided in . Use it in memory entries - instead of "the user" (e.g. "{name} works at..." not "The user works at..."). - Do not store the name itself as a separate memory entry. - - Do not store short-lived or ephemeral info: one-off questions, greetings, - session logistics, or things that only matter for the current task. + - Your current memory is already in in your context. The `chars` + and `limit` attributes show current usage and the maximum allowed size. + - This is curated long-term memory, not raw conversation logs. + - Call update_memory when the user explicitly asks to remember/forget + something or shares durable facts, preferences, or standing instructions. + - The user's first name is provided in . Use it in entries instead + of "the user" when helpful. Do not store the name alone as a memory entry. + - Do not store short-lived info: one-off questions, greetings, session + logistics, or things that only matter for the current task. - Args: - - updated_memory: The FULL updated markdown document (not a diff). - Merge new facts with existing ones, update contradictions, remove outdated entries. - Treat every update as a curation pass — consolidate, don't just append. - - Every bullet MUST use this format: - (YYYY-MM-DD) [marker] text - Markers: - [fact] — durable facts (role, background, projects, tools, expertise) - [pref] — preferences (response style, languages, formats, tools) - [instr] — standing instructions (always/never do, response rules) - - Keep it concise and well under the character limit shown in . - - Every entry MUST be under a `##` heading. Keep heading names short (2-3 words) and - natural. Do NOT include the user's name in headings. Organize by context — e.g. - who they are, what they're focused on, how they prefer things. Create, split, or - merge headings freely as the memory grows. - - Each entry MUST be a single bullet point. Be descriptive but concise — include relevant - details and context rather than just a few words. - - During consolidation, prioritize keeping: [instr] > [pref] > [fact]. + - updated_memory: The FULL updated markdown document, not a diff. Merge new + facts with existing ones, update contradictions, remove outdated entries, + and consolidate instead of only appending. + - Use heading-based Markdown: + * Every entry must be under a `##` heading. + * Recommended headings: `## Facts`, `## Preferences`, `## Instructions`. + Specific natural headings are allowed when clearer. + * New bullets should use `- YYYY-MM-DD: text`. + * Each entry should be one concise but descriptive bullet. + - If existing memory uses legacy `(YYYY-MM-DD) [fact|pref|instr]` markers, + preserve the information but write the updated document in the new + heading-based format. + - During consolidation, prioritize durable instructions and preferences before + generic facts. diff --git a/surfsense_backend/app/agents/new_chat/prompts/tools/update_memory_team.md b/surfsense_backend/app/agents/new_chat/prompts/tools/update_memory_team.md index 7eaca8818..79d4ead3a 100644 --- a/surfsense_backend/app/agents/new_chat/prompts/tools/update_memory_team.md +++ b/surfsense_backend/app/agents/new_chat/prompts/tools/update_memory_team.md @@ -1,26 +1,28 @@ - update_memory: Update the team's shared memory document for this search space. - - Your current team memory is already in in your context. The `chars` - and `limit` attributes show current usage and the maximum allowed size. - - This is the team's curated long-term memory — decisions, conventions, key facts. - - NEVER store personal memory in team memory (e.g. personal bio, individual - preferences, or user-only standing instructions). - - Call update_memory when: - * A team member explicitly asks to remember or forget something - * The conversation surfaces durable team decisions, conventions, or facts - that will matter in future conversations - - Do not store short-lived or ephemeral info: one-off questions, greetings, - session logistics, or things that only matter for the current task. + - Your current team memory is already in in your context. The + `chars` and `limit` attributes show current usage and the maximum allowed size. + - This is curated long-term team memory: decisions, conventions, architecture, + processes, and key shared facts. + - NEVER store personal memory in team memory: individual bios, personal + preferences, or user-only standing instructions. + - Call update_memory when a team member asks to remember/forget something, or + when the conversation surfaces durable team context that matters later. + - Do not store short-lived info: one-off questions, greetings, session + logistics, or things that only matter for the current task. - Args: - - updated_memory: The FULL updated markdown document (not a diff). - Merge new facts with existing ones, update contradictions, remove outdated entries. - Treat every update as a curation pass — consolidate, don't just append. - - Every bullet MUST use this format: - (YYYY-MM-DD) [fact] text - Team memory uses ONLY the [fact] marker. Never use [pref] or [instr] in team memory. - - Keep it concise and well under the character limit shown in . - - Every entry MUST be under a `##` heading. Keep heading names short (2-3 words) and - natural. Organize by context — e.g. what the team decided, current architecture, - active processes. Create, split, or merge headings freely as the memory grows. - - Each entry MUST be a single bullet point. Be descriptive but concise — include relevant - details and context rather than just a few words. - - During consolidation, prioritize keeping: decisions/conventions > key facts > current priorities. + - updated_memory: The FULL updated markdown document, not a diff. Merge new + facts with existing ones, update contradictions, remove outdated entries, + and consolidate instead of only appending. + - Use heading-based Markdown: + * Every entry must be under a `##` heading. + * Recommended headings: `## Product Decisions`, `## Engineering Conventions`, + `## Project Facts`, `## Open Questions`. + * New bullets should use `- YYYY-MM-DD: text`. + * Each entry should be one concise but descriptive bullet. + - If existing memory uses legacy `(YYYY-MM-DD) [fact]` markers, preserve the + information but write the updated document in the new heading-based format. + - Do not create personal headings such as `## Preferences`, `## Instructions`, + `## Personal Notes`, or `## Personal Instructions`. + - During consolidation, prioritize decisions/conventions, then key facts, then + current priorities. diff --git a/surfsense_backend/tests/unit/agents/new_chat/test_memory_response_content.py b/surfsense_backend/tests/unit/agents/new_chat/test_memory_response_content.py index 1f338ee3e..c2f52659c 100644 --- a/surfsense_backend/tests/unit/agents/new_chat/test_memory_response_content.py +++ b/surfsense_backend/tests/unit/agents/new_chat/test_memory_response_content.py @@ -2,28 +2,12 @@ import pytest -from app.agents.new_chat.tools.update_memory import _save_memory +from app.services.memory import MemoryScope, save_memory from app.utils.content_utils import extract_text_content pytestmark = pytest.mark.unit -class _Recorder: - def __init__(self) -> None: - self.applied_content: str | None = None - self.commit_calls = 0 - self.rollback_calls = 0 - - def apply(self, content: str) -> None: - self.applied_content = content - - async def commit(self) -> None: - self.commit_calls += 1 - - async def rollback(self) -> None: - self.rollback_calls += 1 - - def test_extract_text_content_keeps_no_update_bare_string_from_content_blocks() -> None: content = [ {"type": "thinking", "thinking": "No"}, @@ -69,21 +53,12 @@ def test_extract_text_content_preserves_plain_string_responses() -> None: @pytest.mark.asyncio async def test_save_memory_rejects_non_string_payload_before_commit() -> None: - recorder = _Recorder() - - result = await _save_memory( - updated_memory=["NO_UPDATE"], # type: ignore[arg-type] - old_memory=None, - llm=None, - apply_fn=recorder.apply, - commit_fn=recorder.commit, - rollback_fn=recorder.rollback, - label="memory", - scope="user", + result = await save_memory( + scope=MemoryScope.USER, + target_id="00000000-0000-0000-0000-000000000000", + content=["NO_UPDATE"], # type: ignore[arg-type] + session=None, # type: ignore[arg-type] ) - assert result["status"] == "error" - assert "must be a string" in result["message"] - assert recorder.applied_content is None - assert recorder.commit_calls == 0 - assert recorder.rollback_calls == 0 + assert result.status == "error" + assert "must be a string" in result.message diff --git a/surfsense_backend/tests/unit/agents/new_chat/tools/test_update_memory_scope.py b/surfsense_backend/tests/unit/agents/new_chat/tools/test_update_memory_scope.py index f7fbacf50..60310d907 100644 --- a/surfsense_backend/tests/unit/agents/new_chat/tools/test_update_memory_scope.py +++ b/surfsense_backend/tests/unit/agents/new_chat/tools/test_update_memory_scope.py @@ -1,24 +1,24 @@ -"""Unit tests for memory scope validation and bullet format validation.""" +"""Unit tests for heading-based memory validation.""" import pytest -from app.agents.new_chat.tools.update_memory import ( - _save_memory, - _validate_bullet_format, - _validate_memory_scope, +from app.services.memory import MemoryScope, save_memory +from app.services.memory.validation import ( + validate_bullet_format, + validate_memory_scope, ) pytestmark = pytest.mark.unit -class _Recorder: +class _FakeSession: def __init__(self) -> None: - self.applied_content: str | None = None + self.added = [] self.commit_calls = 0 self.rollback_calls = 0 - def apply(self, content: str) -> None: - self.applied_content = content + def add(self, obj) -> None: + self.added.append(obj) async def commit(self) -> None: self.commit_calls += 1 @@ -27,172 +27,125 @@ class _Recorder: self.rollback_calls += 1 -# --------------------------------------------------------------------------- -# _validate_memory_scope — marker-based -# --------------------------------------------------------------------------- - - -def test_validate_memory_scope_rejects_pref_marker_in_team_scope() -> None: - content = "- (2026-04-10) [pref] Prefers dark mode\n" - result = _validate_memory_scope(content, "team") +def test_validate_memory_scope_rejects_new_personal_heading_in_team() -> None: + content = "## Preferences\n- 2026-04-10: Prefers dark mode\n" + result, _warnings = validate_memory_scope(content, "team") assert result is not None assert result["status"] == "error" - assert "[pref]" in result["message"] + assert "preferences" in result["message"] -def test_validate_memory_scope_rejects_instr_marker_in_team_scope() -> None: - content = "- (2026-04-10) [instr] Always respond in Spanish\n" - result = _validate_memory_scope(content, "team") - assert result is not None - assert result["status"] == "error" - assert "[instr]" in result["message"] +def test_validate_memory_scope_allows_old_marker_payload_in_team_scope() -> None: + content = "- (2026-04-10) [pref] Legacy personal marker remains readable\n" + result, _warnings = validate_memory_scope(content, "team") + assert result is None -def test_validate_memory_scope_rejects_both_personal_markers_in_team() -> None: +def test_validate_memory_scope_allows_team_headings() -> None: + content = "## Engineering Conventions\n- 2026-04-10: Uses PostgreSQL\n" + result, _warnings = validate_memory_scope(content, "team") + assert result is None + + +def test_validate_bullet_format_accepts_new_and_legacy_bullets() -> None: content = ( - "- (2026-04-10) [pref] Prefers dark mode\n" - "- (2026-04-10) [instr] Always respond in Spanish\n" + "## Facts\n" + "- 2026-04-10: Senior Python developer\n" + "- (2026-04-10) [fact] Legacy fact is preserved\n" ) - result = _validate_memory_scope(content, "team") - assert result is not None - assert result["status"] == "error" - assert "[instr]" in result["message"] - assert "[pref]" in result["message"] - - -def test_validate_memory_scope_allows_fact_in_team_scope() -> None: - content = "- (2026-04-10) [fact] Office is in downtown Seattle\n" - result = _validate_memory_scope(content, "team") - assert result is None - - -def test_validate_memory_scope_allows_all_markers_in_user_scope() -> None: - content = ( - "- (2026-04-10) [fact] Python developer\n" - "- (2026-04-10) [pref] Prefers concise answers\n" - "- (2026-04-10) [instr] Always use bullet points\n" - ) - result = _validate_memory_scope(content, "user") - assert result is None - - -def test_validate_memory_scope_allows_any_heading_in_team() -> None: - content = "## Architecture\n- (2026-04-10) [fact] Uses PostgreSQL for persistence\n" - result = _validate_memory_scope(content, "team") - assert result is None - - -def test_validate_memory_scope_allows_any_heading_in_user() -> None: - content = "## My Projects\n- (2026-04-10) [fact] Working on SurfSense\n" - result = _validate_memory_scope(content, "user") - assert result is None - - -# --------------------------------------------------------------------------- -# _validate_bullet_format -# --------------------------------------------------------------------------- - - -def test_validate_bullet_format_passes_valid_bullets() -> None: - content = ( - "## Work\n" - "- (2026-04-10) [fact] Senior Python developer\n" - "- (2026-04-10) [pref] Prefers dark mode\n" - "- (2026-04-10) [instr] Always respond in bullet points\n" - ) - warnings = _validate_bullet_format(content) + warnings = validate_bullet_format(content) assert warnings == [] -def test_validate_bullet_format_warns_on_missing_marker() -> None: - content = "- (2026-04-10) Senior Python developer\n" - warnings = _validate_bullet_format(content) +def test_validate_bullet_format_warns_on_nonstandard_bullet() -> None: + content = "## Facts\n- Senior Python developer\n" + warnings = validate_bullet_format(content) assert len(warnings) == 1 - assert "Malformed bullet" in warnings[0] - - -def test_validate_bullet_format_warns_on_missing_date() -> None: - content = "- [fact] Senior Python developer\n" - warnings = _validate_bullet_format(content) - assert len(warnings) == 1 - assert "Malformed bullet" in warnings[0] - - -def test_validate_bullet_format_warns_on_unknown_marker() -> None: - content = "- (2026-04-10) [context] Working on project X\n" - warnings = _validate_bullet_format(content) - assert len(warnings) == 1 - assert "Malformed bullet" in warnings[0] - - -def test_validate_bullet_format_ignores_non_bullet_lines() -> None: - content = "## Some Heading\nSome paragraph text\n" - warnings = _validate_bullet_format(content) - assert warnings == [] - - -def test_validate_bullet_format_warns_on_old_format_without_marker() -> None: - content = "## About the user\n- (2026-04-10) Likes cats\n" - warnings = _validate_bullet_format(content) - assert len(warnings) == 1 - - -# --------------------------------------------------------------------------- -# _save_memory — end-to-end with marker scope check -# --------------------------------------------------------------------------- + assert "Non-standard memory bullet" in warnings[0] @pytest.mark.asyncio -async def test_save_memory_blocks_pref_in_team_before_commit() -> None: - recorder = _Recorder() - result = await _save_memory( - updated_memory="- (2026-04-10) [pref] Prefers dark mode\n", - old_memory=None, - llm=None, - apply_fn=recorder.apply, - commit_fn=recorder.commit, - rollback_fn=recorder.rollback, - label="team memory", - scope="team", +async def test_save_memory_blocks_new_personal_heading_in_team_before_commit( + monkeypatch, +) -> None: + target = type("Target", (), {"shared_memory_md": ""})() + session = _FakeSession() + + async def fake_load_target(**_kwargs): + return target + + monkeypatch.setattr("app.services.memory.service._load_target", fake_load_target) + + result = await save_memory( + scope=MemoryScope.TEAM, + target_id=1, + content="## Preferences\n- 2026-04-10: Prefers dark mode\n", + session=session, ) - assert result["status"] == "error" - assert recorder.commit_calls == 0 - assert recorder.applied_content is None + assert result.status == "error" + assert session.commit_calls == 0 + assert target.shared_memory_md == "" @pytest.mark.asyncio -async def test_save_memory_allows_fact_in_team_and_commits() -> None: - recorder = _Recorder() - content = "- (2026-04-10) [fact] Weekly standup on Mondays\n" - result = await _save_memory( - updated_memory=content, - old_memory=None, - llm=None, - apply_fn=recorder.apply, - commit_fn=recorder.commit, - rollback_fn=recorder.rollback, - label="team memory", - scope="team", +async def test_save_memory_allows_grandfathered_personal_heading_in_team(monkeypatch) -> None: + content = "## Preferences\n- 2026-04-10: Prefers dark mode\n" + target = type("Target", (), {"shared_memory_md": content})() + session = _FakeSession() + + async def fake_load_target(**_kwargs): + return target + + monkeypatch.setattr("app.services.memory.service._load_target", fake_load_target) + + result = await save_memory( + scope=MemoryScope.TEAM, + target_id=1, + content=content, + session=session, ) - assert result["status"] == "saved" - assert recorder.commit_calls == 1 - assert recorder.applied_content == content + assert result.status == "saved" + assert session.commit_calls == 1 + assert target.shared_memory_md == content.strip() + assert result.warnings @pytest.mark.asyncio -async def test_save_memory_includes_format_warnings() -> None: - recorder = _Recorder() - content = "- (2026-04-10) Missing marker text\n" - result = await _save_memory( - updated_memory=content, - old_memory=None, - llm=None, - apply_fn=recorder.apply, - commit_fn=recorder.commit, - rollback_fn=recorder.rollback, - label="memory", - scope="user", +async def test_save_memory_strips_preamble_before_heading(monkeypatch) -> None: + target = type("Target", (), {"memory_md": ""})() + session = _FakeSession() + + async def fake_load_target(**_kwargs): + return target + + monkeypatch.setattr("app.services.memory.service._load_target", fake_load_target) + + result = await save_memory( + scope=MemoryScope.USER, + target_id="00000000-0000-0000-0000-000000000000", + content="Sure, here is the update:\n\n## Facts\n- 2026-04-10: Likes cats\n", + session=session, ) - assert result["status"] == "saved" - assert "format_warnings" in result - assert len(result["format_warnings"]) == 1 + assert result.status == "saved" + assert target.memory_md == "## Facts\n- 2026-04-10: Likes cats" + + +@pytest.mark.asyncio +async def test_save_memory_rejects_long_no_heading_payload(monkeypatch) -> None: + target = type("Target", (), {"memory_md": ""})() + session = _FakeSession() + + async def fake_load_target(**_kwargs): + return target + + monkeypatch.setattr("app.services.memory.service._load_target", fake_load_target) + + result = await save_memory( + scope=MemoryScope.USER, + target_id="00000000-0000-0000-0000-000000000000", + content="NO_UPDATE because there is nothing durable to remember.", + session=session, + ) + assert result.status == "error" + assert "## heading" in result.message + assert session.commit_calls == 0 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 03/15] 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 From 89a8438864b5a52d9927b8af3e37df8b1116d900 Mon Sep 17 00:00:00 2001 From: Anish Sarkar <104695310+AnishSarkar22@users.noreply.github.com> Date: Wed, 20 May 2026 02:02:42 +0530 Subject: [PATCH 04/15] feat: wire memory settings to memory API --- .../components/MemoryContent.tsx | 151 +---------------- .../settings/team-memory-manager.tsx | 152 +----------------- .../contracts/types/search-space.types.ts | 1 - surfsense_web/hooks/use-memory.ts | 109 +++++++++++++ 4 files changed, 121 insertions(+), 292 deletions(-) create mode 100644 surfsense_web/hooks/use-memory.ts diff --git a/surfsense_web/app/dashboard/[search_space_id]/user-settings/components/MemoryContent.tsx b/surfsense_web/app/dashboard/[search_space_id]/user-settings/components/MemoryContent.tsx index 3542f0925..dc002244f 100644 --- a/surfsense_web/app/dashboard/[search_space_id]/user-settings/components/MemoryContent.tsx +++ b/surfsense_web/app/dashboard/[search_space_id]/user-settings/components/MemoryContent.tsx @@ -1,10 +1,8 @@ "use client"; import { useAtomValue } from "jotai"; -import { ArrowUp, ChevronDown, ClipboardCopy, Download, Info, Pencil } from "lucide-react"; -import { useCallback, useEffect, useRef, useState } from "react"; +import { ChevronDown, ClipboardCopy, Download, Info } from "lucide-react"; import { toast } from "sonner"; -import { z } from "zod"; import { activeSearchSpaceIdAtom } from "@/atoms/search-spaces/search-space-query.atoms"; import { PlateEditor } from "@/components/editor/plate-editor"; import { Alert, AlertDescription } from "@/components/ui/alert"; @@ -16,102 +14,23 @@ import { DropdownMenuTrigger, } from "@/components/ui/dropdown-menu"; import { Spinner } from "@/components/ui/spinner"; - -import { baseApiService } from "@/lib/apis/base-api.service"; - -const MEMORY_HARD_LIMIT = 25_000; - -const MemoryReadSchema = z.object({ - memory_md: z.string(), -}); +import { MEMORY_HARD_LIMIT, useUserMemory } from "@/hooks/use-memory"; export function MemoryContent() { const activeSearchSpaceId = useAtomValue(activeSearchSpaceIdAtom); - const [memory, setMemory] = useState(""); - const [loading, setLoading] = useState(true); - const [saving, setSaving] = useState(false); - const [editQuery, setEditQuery] = useState(""); - const [editing, setEditing] = useState(false); - const [showInput, setShowInput] = useState(false); - const textareaRef = useRef(null); - const inputContainerRef = useRef(null); - - const fetchMemory = useCallback(async () => { - try { - setLoading(true); - const data = await baseApiService.get("/api/v1/users/me/memory", MemoryReadSchema); - setMemory(data.memory_md); - } catch { - toast.error("Failed to load memory"); - } finally { - setLoading(false); - } - }, []); - - useEffect(() => { - fetchMemory(); - }, [fetchMemory]); - - useEffect(() => { - if (!showInput) return; - - const handlePointerDownOutside = (event: MouseEvent | TouchEvent) => { - const target = event.target; - if (!(target instanceof Node)) return; - if (inputContainerRef.current?.contains(target)) return; - - setShowInput(false); - }; - - document.addEventListener("mousedown", handlePointerDownOutside); - document.addEventListener("touchstart", handlePointerDownOutside, { passive: true }); - - return () => { - document.removeEventListener("mousedown", handlePointerDownOutside); - document.removeEventListener("touchstart", handlePointerDownOutside); - }; - }, [showInput]); + const { memory, displayMemory, loading, saving, reset } = useUserMemory( + Number(activeSearchSpaceId) + ); const handleClear = async () => { try { - setSaving(true); - const data = await baseApiService.put("/api/v1/users/me/memory", MemoryReadSchema, { - body: { memory_md: "" }, - }); - setMemory(data.memory_md); + await reset(); toast.success("Memory cleared"); } catch { toast.error("Failed to clear memory"); - } finally { - setSaving(false); } }; - const handleEdit = async () => { - const query = editQuery.trim(); - if (!query) return; - - try { - setEditing(true); - const data = await baseApiService.post("/api/v1/users/me/memory/edit", MemoryReadSchema, { - body: { query, search_space_id: Number(activeSearchSpaceId) }, - }); - setMemory(data.memory_md); - setEditQuery(""); - setShowInput(false); - toast.success("Memory updated"); - } catch { - toast.error("Failed to edit memory"); - } finally { - setEditing(false); - } - }; - - const openInput = () => { - setShowInput(true); - requestAnimationFrame(() => textareaRef.current?.focus()); - }; - const handleDownload = () => { if (!memory) return; try { @@ -139,14 +58,6 @@ export function MemoryContent() { } }; - const handleKeyDown = (e: React.KeyboardEvent) => { - if (e.key === "Enter" && !e.shiftKey) { - e.preventDefault(); - handleEdit(); - } - }; - - const displayMemory = memory.replace(/\(\d{4}-\d{2}-\d{2}\)\s*\[(fact|pref|instr)\]\s*/g, ""); const charCount = memory.length; const getCounterColor = () => { @@ -198,54 +109,6 @@ export function MemoryContent() { className="px-5 py-4 text-sm min-h-full" /> - - {showInput ? ( -
-
- setEditQuery(e.target.value)} - onKeyDown={handleKeyDown} - placeholder="Tell SurfSense what to remember or forget" - disabled={editing} - className="flex-1 bg-transparent text-sm outline-none placeholder:text-muted-foreground/70" - /> - -
-
- ) : ( - - )}
@@ -263,7 +126,7 @@ export function MemoryContent() { size="sm" className="text-xs sm:text-sm" onClick={handleClear} - disabled={saving || editing || !memory} + disabled={saving || !memory} > Reset Memory Reset diff --git a/surfsense_web/components/settings/team-memory-manager.tsx b/surfsense_web/components/settings/team-memory-manager.tsx index 9d3a40e46..6a2cbf52f 100644 --- a/surfsense_web/components/settings/team-memory-manager.tsx +++ b/surfsense_web/components/settings/team-memory-manager.tsx @@ -1,12 +1,7 @@ "use client"; -import { useQuery, useQueryClient } from "@tanstack/react-query"; -import { useAtomValue } from "jotai"; -import { ArrowUp, ChevronDown, ClipboardCopy, Download, Info, Pencil } from "lucide-react"; -import { useEffect, useRef, useState } from "react"; +import { ChevronDown, ClipboardCopy, Download, Info } from "lucide-react"; import { toast } from "sonner"; -import { z } from "zod"; -import { updateSearchSpaceMutationAtom } from "@/atoms/search-spaces/search-space-mutation.atoms"; import { PlateEditor } from "@/components/editor/plate-editor"; import { Alert, AlertDescription } from "@/components/ui/alert"; import { Button } from "@/components/ui/button"; @@ -17,105 +12,24 @@ import { DropdownMenuTrigger, } from "@/components/ui/dropdown-menu"; import { Spinner } from "@/components/ui/spinner"; -import { baseApiService } from "@/lib/apis/base-api.service"; -import { searchSpacesApiService } from "@/lib/apis/search-spaces-api.service"; -import { cacheKeys } from "@/lib/query-client/cache-keys"; - -const MEMORY_HARD_LIMIT = 25_000; - -const SearchSpaceSchema = z - .object({ - shared_memory_md: z.string().optional().default(""), - }) - .passthrough(); +import { MEMORY_HARD_LIMIT, useTeamMemory } from "@/hooks/use-memory"; interface TeamMemoryManagerProps { searchSpaceId: number; } export function TeamMemoryManager({ searchSpaceId }: TeamMemoryManagerProps) { - const queryClient = useQueryClient(); - const { data: searchSpace, isLoading: loading } = useQuery({ - queryKey: cacheKeys.searchSpaces.detail(searchSpaceId.toString()), - queryFn: () => searchSpacesApiService.getSearchSpace({ id: searchSpaceId }), - enabled: !!searchSpaceId, - }); - - const { mutateAsync: updateSearchSpace } = useAtomValue(updateSearchSpaceMutationAtom); - - const [saving, setSaving] = useState(false); - const [editQuery, setEditQuery] = useState(""); - const [editing, setEditing] = useState(false); - const [showInput, setShowInput] = useState(false); - const textareaRef = useRef(null); - const inputContainerRef = useRef(null); - - const memory = searchSpace?.shared_memory_md || ""; - - useEffect(() => { - if (!showInput) return; - - const handlePointerDownOutside = (event: MouseEvent | TouchEvent) => { - const target = event.target; - if (!(target instanceof Node)) return; - if (inputContainerRef.current?.contains(target)) return; - - setShowInput(false); - }; - - document.addEventListener("mousedown", handlePointerDownOutside); - document.addEventListener("touchstart", handlePointerDownOutside, { passive: true }); - - return () => { - document.removeEventListener("mousedown", handlePointerDownOutside); - document.removeEventListener("touchstart", handlePointerDownOutside); - }; - }, [showInput]); + const { memory, displayMemory, loading, saving, reset } = useTeamMemory(searchSpaceId); const handleClear = async () => { try { - setSaving(true); - await updateSearchSpace({ - id: searchSpaceId, - data: { shared_memory_md: "" }, - }); + await reset(); toast.success("Team memory cleared"); } catch { toast.error("Failed to clear team memory"); - } finally { - setSaving(false); } }; - const handleEdit = async () => { - const query = editQuery.trim(); - if (!query) return; - - try { - setEditing(true); - await baseApiService.post( - `/api/v1/searchspaces/${searchSpaceId}/memory/edit`, - SearchSpaceSchema, - { body: { query } } - ); - setEditQuery(""); - setShowInput(false); - await queryClient.invalidateQueries({ - queryKey: cacheKeys.searchSpaces.detail(searchSpaceId.toString()), - }); - toast.success("Team memory updated"); - } catch { - toast.error("Failed to edit team memory"); - } finally { - setEditing(false); - } - }; - - const openInput = () => { - setShowInput(true); - requestAnimationFrame(() => textareaRef.current?.focus()); - }; - const handleDownload = () => { if (!memory) return; try { @@ -143,14 +57,6 @@ export function TeamMemoryManager({ searchSpaceId }: TeamMemoryManagerProps) { } }; - const handleKeyDown = (e: React.KeyboardEvent) => { - if (e.key === "Enter" && !e.shiftKey) { - e.preventDefault(); - handleEdit(); - } - }; - - const displayMemory = memory.replace(/\(\d{4}-\d{2}-\d{2}\)\s*\[(fact|pref|instr)\]\s*/g, ""); const charCount = memory.length; const getCounterColor = () => { @@ -204,54 +110,6 @@ export function TeamMemoryManager({ searchSpaceId }: TeamMemoryManagerProps) { className="px-5 py-4 text-sm min-h-full" />
- - {showInput ? ( -
-
- setEditQuery(e.target.value)} - onKeyDown={handleKeyDown} - placeholder="Tell SurfSense what to remember or forget about your team" - disabled={editing} - className="flex-1 bg-transparent text-sm outline-none placeholder:text-muted-foreground/70" - /> - -
-
- ) : ( - - )}
@@ -269,7 +127,7 @@ export function TeamMemoryManager({ searchSpaceId }: TeamMemoryManagerProps) { size="sm" className="text-xs sm:text-sm" onClick={handleClear} - disabled={saving || editing || !memory} + disabled={saving || !memory} > Reset Memory Reset diff --git a/surfsense_web/contracts/types/search-space.types.ts b/surfsense_web/contracts/types/search-space.types.ts index 7449f82b1..08918e2af 100644 --- a/surfsense_web/contracts/types/search-space.types.ts +++ b/surfsense_web/contracts/types/search-space.types.ts @@ -56,7 +56,6 @@ export const updateSearchSpaceRequest = z.object({ description: true, citations_enabled: true, qna_custom_instructions: true, - shared_memory_md: true, ai_file_sort_enabled: true, }) .partial(), diff --git a/surfsense_web/hooks/use-memory.ts b/surfsense_web/hooks/use-memory.ts new file mode 100644 index 000000000..1f7a51790 --- /dev/null +++ b/surfsense_web/hooks/use-memory.ts @@ -0,0 +1,109 @@ +"use client"; + +import { useCallback, useEffect, useState } from "react"; +import { z } from "zod"; +import { baseApiService } from "@/lib/apis/base-api.service"; + +export const MEMORY_HARD_LIMIT = 25_000; + +const MemoryReadSchema = z.object({ + memory_md: z.string(), +}); + +type MemoryScope = "user" | "team"; + +interface UseMemoryOptions { + scope: MemoryScope; + searchSpaceId?: number | null; + autoLoad?: boolean; +} + +function getMemoryPath(scope: MemoryScope, searchSpaceId?: number | null) { + if (scope === "user") return "/api/v1/users/me/memory"; + if (!searchSpaceId) throw new Error("searchSpaceId is required for team memory"); + return `/api/v1/searchspaces/${searchSpaceId}/memory`; +} + +export function stripMemoryDisplayPrefixes(memory: string) { + return memory.replace( + /^\s*-\s+(?:\(\d{4}-\d{2}-\d{2}\)\s*\[(?:fact|pref|instr)\]\s*|\d{4}-\d{2}-\d{2}:\s*)/gim, + "- " + ); +} + +export function useMemory({ scope, searchSpaceId, autoLoad = true }: UseMemoryOptions) { + const [memory, setMemory] = useState(""); + const [loading, setLoading] = useState(autoLoad); + const [saving, setSaving] = useState(false); + + const load = useCallback(async () => { + setLoading(true); + try { + const data = await baseApiService.get(getMemoryPath(scope, searchSpaceId), MemoryReadSchema); + setMemory(data.memory_md); + return data.memory_md; + } finally { + setLoading(false); + } + }, [scope, searchSpaceId]); + + useEffect(() => { + if (!autoLoad) return; + load().catch(() => { + setLoading(false); + }); + }, [autoLoad, load]); + + const save = useCallback( + async (memoryMd: string) => { + setSaving(true); + try { + const data = await baseApiService.put( + getMemoryPath(scope, searchSpaceId), + MemoryReadSchema, + { + body: { memory_md: memoryMd }, + } + ); + setMemory(data.memory_md); + return data.memory_md; + } finally { + setSaving(false); + } + }, + [scope, searchSpaceId] + ); + + const reset = useCallback(async () => { + setSaving(true); + try { + const data = await baseApiService.post( + `${getMemoryPath(scope, searchSpaceId)}/reset`, + MemoryReadSchema + ); + setMemory(data.memory_md); + return data.memory_md; + } finally { + setSaving(false); + } + }, [scope, searchSpaceId]); + + return { + memory, + setMemory, + displayMemory: stripMemoryDisplayPrefixes(memory), + loading, + saving, + load, + save, + reset, + }; +} + +export function useUserMemory(searchSpaceId?: number | null) { + return useMemory({ scope: "user", searchSpaceId }); +} + +export function useTeamMemory(searchSpaceId?: number | null) { + return useMemory({ scope: "team", searchSpaceId }); +} From cb1cf26ef3436c233fb22b3d0b791f5241552c15 Mon Sep 17 00:00:00 2001 From: Anish Sarkar <104695310+AnishSarkar22@users.noreply.github.com> Date: Wed, 20 May 2026 02:02:59 +0530 Subject: [PATCH 05/15] feat: improve document editor panel behavior --- .../atoms/editor/editor-panel.atom.ts | 26 +++- .../components/documents/DocumentNode.tsx | 132 ++++++++++++++---- .../components/documents/FolderTreeView.tsx | 79 ++++++----- .../components/editor-panel/editor-panel.tsx | 108 ++++++++++++-- .../layout/ui/right-panel/RightPanel.tsx | 19 ++- .../layout/ui/sidebar/DocumentsSidebar.tsx | 91 +++++++++++- .../contracts/enums/connectorIcons.tsx | 4 + .../contracts/types/document.types.ts | 2 + 8 files changed, 380 insertions(+), 81 deletions(-) diff --git a/surfsense_web/atoms/editor/editor-panel.atom.ts b/surfsense_web/atoms/editor/editor-panel.atom.ts index 28563e7d3..c302c66ee 100644 --- a/surfsense_web/atoms/editor/editor-panel.atom.ts +++ b/surfsense_web/atoms/editor/editor-panel.atom.ts @@ -3,10 +3,11 @@ import { rightPanelCollapsedAtom, rightPanelTabAtom } from "@/atoms/layout/right interface EditorPanelState { isOpen: boolean; - kind: "document" | "local_file"; + kind: "document" | "local_file" | "memory"; documentId: number | null; localFilePath: string | null; searchSpaceId: number | null; + memoryScope: "user" | "team" | null; title: string | null; } @@ -16,6 +17,7 @@ const initialState: EditorPanelState = { documentId: null, localFilePath: null, searchSpaceId: null, + memoryScope: null, title: null, }; @@ -38,6 +40,12 @@ export const openEditorPanelAtom = atom( title?: string; searchSpaceId?: number; } + | { + kind: "memory"; + memoryScope: "user" | "team"; + title?: string; + searchSpaceId?: number; + } ) => { if (!get(editorPanelAtom).isOpen) { set(preEditorCollapsedAtom, get(rightPanelCollapsedAtom)); @@ -49,6 +57,21 @@ export const openEditorPanelAtom = atom( documentId: null, localFilePath: payload.localFilePath, searchSpaceId: payload.searchSpaceId ?? null, + memoryScope: null, + title: payload.title ?? null, + }); + set(rightPanelTabAtom, "editor"); + set(rightPanelCollapsedAtom, false); + return; + } + if (payload.kind === "memory") { + set(editorPanelAtom, { + isOpen: true, + kind: "memory", + documentId: null, + localFilePath: null, + searchSpaceId: payload.searchSpaceId ?? null, + memoryScope: payload.memoryScope, title: payload.title ?? null, }); set(rightPanelTabAtom, "editor"); @@ -61,6 +84,7 @@ export const openEditorPanelAtom = atom( documentId: payload.documentId, localFilePath: null, searchSpaceId: payload.searchSpaceId, + memoryScope: null, title: payload.title ?? null, }); set(rightPanelTabAtom, "editor"); diff --git a/surfsense_web/components/documents/DocumentNode.tsx b/surfsense_web/components/documents/DocumentNode.tsx index 0f3cd4a19..bef2c6ba2 100644 --- a/surfsense_web/components/documents/DocumentNode.tsx +++ b/surfsense_web/components/documents/DocumentNode.tsx @@ -9,6 +9,7 @@ import { MoreHorizontal, Move, Pencil, + RotateCcw, Trash2, } from "lucide-react"; import React, { useCallback, useRef, useState } from "react"; @@ -61,8 +62,13 @@ interface DocumentNodeProps { onEdit: (doc: DocumentNodeDoc) => void; onDelete: (doc: DocumentNodeDoc) => void; onMove: (doc: DocumentNodeDoc) => void; + onReset?: (doc: DocumentNodeDoc) => void; onExport?: (doc: DocumentNodeDoc, format: string) => void; onVersionHistory?: (doc: DocumentNodeDoc) => void; + canDelete?: boolean; + canMove?: boolean; + canMention?: boolean; + canEdit?: boolean; contextMenuOpen?: boolean; onContextMenuOpenChange?: (open: boolean) => void; } @@ -76,8 +82,13 @@ export const DocumentNode = React.memo(function DocumentNode({ onEdit, onDelete, onMove, + onReset, onExport, onVersionHistory, + canDelete = true, + canMove = true, + canMention = true, + canEdit = true, contextMenuOpen, onContextMenuOpenChange, }: DocumentNodeProps) { @@ -85,8 +96,13 @@ export const DocumentNode = React.memo(function DocumentNode({ const isFailed = statusState === "failed"; const isProcessing = statusState === "pending" || statusState === "processing"; const isUnavailable = isProcessing || isFailed; - const isSelectable = !isUnavailable; - const isEditable = EDITABLE_DOCUMENT_TYPES.has(doc.document_type) && !isUnavailable; + const isMemoryDocument = + doc.document_type === "USER_MEMORY" || doc.document_type === "TEAM_MEMORY"; + const isSelectable = canMention && !isUnavailable; + const isEditable = + canEdit && + (isMemoryDocument || EDITABLE_DOCUMENT_TYPES.has(doc.document_type)) && + !isUnavailable; const handleCheckChange = useCallback(() => { if (isSelectable) { @@ -94,13 +110,22 @@ export const DocumentNode = React.memo(function DocumentNode({ } }, [doc, isMentioned, isSelectable, onToggleChatMention]); + const handlePrimaryClick = useCallback(() => { + if (canMention) { + handleCheckChange(); + return; + } + onPreview(doc); + }, [canMention, doc, handleCheckChange, onPreview]); + const [{ isDragging }, drag] = useDrag( () => ({ type: DND_TYPES.DOCUMENT, item: { id: doc.id }, + canDrag: canMove, collect: (monitor) => ({ isDragging: monitor.isDragging() }), }), - [doc.id] + [canMove, doc.id] ); const [dropdownOpen, setDropdownOpen] = useState(false); @@ -130,9 +155,11 @@ export const DocumentNode = React.memo(function DocumentNode({ const attachRef = useCallback( (node: HTMLDivElement | null) => { (rowRef as React.MutableRefObject).current = node; - drag(node); + if (canMove) { + drag(node); + } }, - [drag] + [canMove, drag] ); return ( @@ -187,12 +214,39 @@ export const DocumentNode = React.memo(function DocumentNode({ ); } return ( - e.stopPropagation()} - className="h-3.5 w-3.5 shrink-0" - /> + <> + {isMemoryDocument ? ( + + ) : canMention ? ( + e.stopPropagation()} + className="h-3.5 w-3.5 shrink-0" + /> + ) : ( + + {getDocumentTypeIcon( + doc.document_type as DocumentTypeEnum, + "h-3.5 w-3.5 text-muted-foreground" + )} + + )} + ); })()} @@ -205,8 +259,8 @@ export const DocumentNode = React.memo(function DocumentNode({
@@ -708,7 +788,9 @@ function DesktopEditorPanel() { const hasTarget = panelState.kind === "document" ? !!panelState.documentId && !!panelState.searchSpaceId - : !!panelState.localFilePath; + : panelState.kind === "local_file" + ? !!panelState.localFilePath + : !!panelState.memoryScope; if (!panelState.isOpen || !hasTarget) return null; return ( @@ -717,6 +799,7 @@ function DesktopEditorPanel() { kind={panelState.kind} documentId={panelState.documentId ?? undefined} localFilePath={panelState.localFilePath ?? undefined} + memoryScope={panelState.memoryScope ?? undefined} searchSpaceId={panelState.searchSpaceId ?? undefined} title={panelState.title} onClose={closePanel} @@ -734,7 +817,7 @@ function MobileEditorDrawer() { const hasTarget = panelState.kind === "document" ? !!panelState.documentId && !!panelState.searchSpaceId - : !!panelState.localFilePath; + : !!panelState.memoryScope; if (!hasTarget) return null; return ( @@ -756,6 +839,7 @@ function MobileEditorDrawer() { kind={panelState.kind} documentId={panelState.documentId ?? undefined} localFilePath={panelState.localFilePath ?? undefined} + memoryScope={panelState.memoryScope ?? undefined} searchSpaceId={panelState.searchSpaceId ?? undefined} title={panelState.title} /> @@ -771,7 +855,9 @@ export function EditorPanel() { const hasTarget = panelState.kind === "document" ? !!panelState.documentId && !!panelState.searchSpaceId - : !!panelState.localFilePath; + : panelState.kind === "local_file" + ? !!panelState.localFilePath + : !!panelState.memoryScope; if (!panelState.isOpen || !hasTarget) return null; if (!isDesktop && panelState.kind === "local_file") return null; @@ -789,7 +875,9 @@ export function MobileEditorPanel() { const hasTarget = panelState.kind === "document" ? !!panelState.documentId && !!panelState.searchSpaceId - : !!panelState.localFilePath; + : panelState.kind === "local_file" + ? !!panelState.localFilePath + : !!panelState.memoryScope; if (isDesktop || !panelState.isOpen || !hasTarget || panelState.kind === "local_file") return null; diff --git a/surfsense_web/components/layout/ui/right-panel/RightPanel.tsx b/surfsense_web/components/layout/ui/right-panel/RightPanel.tsx index b379e58e3..5a7588979 100644 --- a/surfsense_web/components/layout/ui/right-panel/RightPanel.tsx +++ b/surfsense_web/components/layout/ui/right-panel/RightPanel.tsx @@ -103,7 +103,11 @@ export function RightPanelToggleButton({ const reportOpen = reportState.isOpen && !!reportState.reportId; const editorOpen = editorState.isOpen && - (editorState.kind === "document" ? !!editorState.documentId : !!editorState.localFilePath); + (editorState.kind === "document" + ? !!editorState.documentId + : editorState.kind === "memory" + ? !!editorState.memoryScope + : !!editorState.localFilePath); const hitlEditOpen = hitlEditState.isOpen && !!hitlEditState.onSave; const citationOpen = citationState.isOpen && citationState.chunkId != null; const hasContent = documentsOpen || reportOpen || editorOpen || hitlEditOpen || citationOpen; @@ -151,7 +155,11 @@ export function RightPanelExpandButton() { const reportOpen = reportState.isOpen && !!reportState.reportId; const editorOpen = editorState.isOpen && - (editorState.kind === "document" ? !!editorState.documentId : !!editorState.localFilePath); + (editorState.kind === "document" + ? !!editorState.documentId + : editorState.kind === "memory" + ? !!editorState.memoryScope + : !!editorState.localFilePath); const hitlEditOpen = hitlEditState.isOpen && !!hitlEditState.onSave; const citationOpen = citationState.isOpen && citationState.chunkId != null; const hasContent = documentsOpen || reportOpen || editorOpen || hitlEditOpen || citationOpen; @@ -193,7 +201,11 @@ export function RightPanel({ const reportOpen = reportState.isOpen && !!reportState.reportId; const editorOpen = editorState.isOpen && - (editorState.kind === "document" ? !!editorState.documentId : !!editorState.localFilePath); + (editorState.kind === "document" + ? !!editorState.documentId + : editorState.kind === "memory" + ? !!editorState.memoryScope + : !!editorState.localFilePath); const hitlEditOpen = hitlEditState.isOpen && !!hitlEditState.onSave; const citationOpen = citationState.isOpen && citationState.chunkId != null; @@ -292,6 +304,7 @@ export function RightPanel({ kind={editorState.kind} documentId={editorState.documentId ?? undefined} localFilePath={editorState.localFilePath ?? undefined} + memoryScope={editorState.memoryScope ?? undefined} searchSpaceId={editorState.searchSpaceId ?? undefined} title={editorState.title} onClose={closeEditor} diff --git a/surfsense_web/components/layout/ui/sidebar/DocumentsSidebar.tsx b/surfsense_web/components/layout/ui/sidebar/DocumentsSidebar.tsx index cdb757cb2..0c37d003c 100644 --- a/surfsense_web/components/layout/ui/sidebar/DocumentsSidebar.tsx +++ b/surfsense_web/components/layout/ui/sidebar/DocumentsSidebar.tsx @@ -88,7 +88,31 @@ const DesktopLocalTabContent = dynamic( { ssr: false } ); -const NON_DELETABLE_DOCUMENT_TYPES: readonly string[] = ["SURFSENSE_DOCS"]; +const NON_DELETABLE_DOCUMENT_TYPES: readonly string[] = [ + "SURFSENSE_DOCS", + "USER_MEMORY", + "TEAM_MEMORY", +]; +const MEMORY_DOCUMENTS: DocumentNodeDoc[] = [ + { + id: -1001, + title: "MEMORY.md", + document_type: "USER_MEMORY", + folderId: null, + status: { state: "ready" }, + }, + { + id: -1002, + title: "TEAM_MEMORY.md", + document_type: "TEAM_MEMORY", + folderId: null, + status: { state: "ready" }, + }, +]; + +function isMemoryDocument(doc: { document_type: string }) { + return doc.document_type === "USER_MEMORY" || doc.document_type === "TEAM_MEMORY"; +} const LOCAL_FILESYSTEM_TRUST_KEY = "surfsense.local-filesystem-trust.v1"; const MAX_LOCAL_FILESYSTEM_ROOTS = 10; @@ -879,6 +903,7 @@ function AuthenticatedDocumentsSidebarBase({ const handleToggleChatMention = useCallback( (doc: { id: number; title: string; document_type: string }, isMentioned: boolean) => { + if (isMemoryDocument(doc)) return; const key = getMentionDocKey({ ...doc, kind: "doc" }); if (isMentioned) { setSidebarDocs((prev) => prev.filter((d) => getMentionDocKey(d) !== key)); @@ -927,11 +952,66 @@ function AuthenticatedDocumentsSidebarBase({ [treeFolders, setSidebarDocs] ); + const treeDocumentsWithMemory = useMemo( + () => [...MEMORY_DOCUMENTS, ...treeDocuments], + [treeDocuments] + ); + const searchFilteredDocuments = useMemo(() => { const query = debouncedSearch.trim().toLowerCase(); - if (!query) return treeDocuments; - return treeDocuments.filter((d) => d.title.toLowerCase().includes(query)); - }, [treeDocuments, debouncedSearch]); + if (!query) return treeDocumentsWithMemory; + return treeDocumentsWithMemory.filter((d) => d.title.toLowerCase().includes(query)); + }, [treeDocumentsWithMemory, debouncedSearch]); + + const openMemoryDocument = useCallback( + (doc: DocumentNodeDoc) => { + if (doc.document_type === "USER_MEMORY") { + openEditorPanel({ + kind: "memory", + memoryScope: "user", + searchSpaceId, + title: doc.title, + }); + return true; + } + if (doc.document_type === "TEAM_MEMORY") { + openEditorPanel({ + kind: "memory", + memoryScope: "team", + searchSpaceId, + title: doc.title, + }); + return true; + } + return false; + }, + [openEditorPanel, searchSpaceId] + ); + + const handleResetMemoryDocument = useCallback( + async (doc: DocumentNodeDoc) => { + if (!isMemoryDocument(doc)) return; + if (!window.confirm(`Reset ${doc.title.toLowerCase()}? This clears the memory document.`)) { + return; + } + const endpoint = + doc.document_type === "USER_MEMORY" + ? `${process.env.NEXT_PUBLIC_FASTAPI_BACKEND_URL}/api/v1/users/me/memory/reset` + : `${process.env.NEXT_PUBLIC_FASTAPI_BACKEND_URL}/api/v1/searchspaces/${searchSpaceId}/memory/reset`; + try { + const response = await authenticatedFetch(endpoint, { method: "POST" }); + if (!response.ok) { + const errorData = await response.json().catch(() => ({ detail: "Reset failed" })); + throw new Error(errorData.detail || "Reset failed"); + } + toast.success(`${doc.title} reset`); + openMemoryDocument(doc); + } catch (error) { + toast.error((error as Error)?.message || `Failed to reset ${doc.title.toLowerCase()}`); + } + }, + [openMemoryDocument, searchSpaceId] + ); const typeCounts = useMemo(() => { const counts: Partial> = {}; @@ -1169,6 +1249,7 @@ function AuthenticatedDocumentsSidebarBase({ onCreateFolder={handleCreateFolder} searchQuery={debouncedSearch.trim() || undefined} onPreviewDocument={(doc) => { + if (openMemoryDocument(doc)) return; openEditorPanel({ documentId: doc.id, searchSpaceId, @@ -1176,6 +1257,7 @@ function AuthenticatedDocumentsSidebarBase({ }); }} onEditDocument={(doc) => { + if (openMemoryDocument(doc)) return; openEditorPanel({ documentId: doc.id, searchSpaceId, @@ -1184,6 +1266,7 @@ function AuthenticatedDocumentsSidebarBase({ }} onDeleteDocument={(doc) => handleDeleteDocument(doc.id)} onMoveDocument={handleMoveDocument} + onResetDocument={handleResetMemoryDocument} onExportDocument={handleExportDocument} onVersionHistory={(doc) => setVersionDocId(doc.id)} activeTypes={activeTypes} diff --git a/surfsense_web/contracts/enums/connectorIcons.tsx b/surfsense_web/contracts/enums/connectorIcons.tsx index 1c6745db5..610bd508b 100644 --- a/surfsense_web/contracts/enums/connectorIcons.tsx +++ b/surfsense_web/contracts/enums/connectorIcons.tsx @@ -1,6 +1,7 @@ import { IconUsersGroup } from "@tabler/icons-react"; import { BookOpen, + Brain, File, FileText, Globe, @@ -120,6 +121,9 @@ export const getConnectorIcon = (connectorType: EnumConnectorName | string, clas return ; case "SURFSENSE_DOCS": return ; + case "USER_MEMORY": + case "TEAM_MEMORY": + return ; case "DEEP": return ; case "DEEPER": diff --git a/surfsense_web/contracts/types/document.types.ts b/surfsense_web/contracts/types/document.types.ts index 7b8784568..ccc15fa62 100644 --- a/surfsense_web/contracts/types/document.types.ts +++ b/surfsense_web/contracts/types/document.types.ts @@ -29,6 +29,8 @@ export const documentTypeEnum = z.enum([ "LOCAL_FOLDER_FILE", "SURFSENSE_DOCS", "NOTE", + "USER_MEMORY", + "TEAM_MEMORY", "COMPOSIO_GOOGLE_DRIVE_CONNECTOR", "COMPOSIO_GMAIL_CONNECTOR", "COMPOSIO_GOOGLE_CALENDAR_CONNECTOR", From 73043a07567a1516a518eff4f6f8dd33875d1a88 Mon Sep 17 00:00:00 2001 From: Anish Sarkar <104695310+AnishSarkar22@users.noreply.github.com> Date: Wed, 20 May 2026 03:17:05 +0530 Subject: [PATCH 06/15] feat: enhance memory API responses with limits and update UI components for memory limit handling --- surfsense_backend/app/routes/memory_routes.py | 12 ++--- .../app/routes/team_memory_routes.py | 18 +++---- .../app/services/memory/__init__.py | 5 ++ .../app/services/memory/schemas.py | 16 +++++- .../app/services/memory/service.py | 8 +-- .../components/MemoryContent.tsx | 18 +++---- .../components/editor-panel/editor-panel.tsx | 54 +++++++++++++++++-- .../settings/team-memory-manager.tsx | 18 +++---- surfsense_web/hooks/use-memory.ts | 34 +++++++++++- 9 files changed, 132 insertions(+), 51 deletions(-) diff --git a/surfsense_backend/app/routes/memory_routes.py b/surfsense_backend/app/routes/memory_routes.py index 7b674a584..8e73a277c 100644 --- a/surfsense_backend/app/routes/memory_routes.py +++ b/surfsense_backend/app/routes/memory_routes.py @@ -8,7 +8,9 @@ from sqlalchemy.ext.asyncio import AsyncSession from app.db import User, get_async_session from app.services.memory import ( + MemoryRead, MemoryScope, + memory_limits, read_memory, reset_memory, save_memory, @@ -18,10 +20,6 @@ from app.users import current_active_user router = APIRouter() -class MemoryRead(BaseModel): - memory_md: str - - class MemoryUpdate(BaseModel): memory_md: str @@ -36,7 +34,7 @@ async def get_user_memory( target_id=user.id, session=session, ) - return MemoryRead(memory_md=memory_md) + return MemoryRead(memory_md=memory_md, limits=memory_limits()) @router.put("/users/me/memory", response_model=MemoryRead) @@ -53,7 +51,7 @@ async def update_user_memory( ) if result.status == "error": raise HTTPException(status_code=400, detail=result.message) - return MemoryRead(memory_md=result.memory_md) + return MemoryRead(memory_md=result.memory_md, limits=memory_limits()) @router.post("/users/me/memory/reset", response_model=MemoryRead) @@ -68,4 +66,4 @@ async def reset_user_memory( ) if result.status == "error": raise HTTPException(status_code=400, detail=result.message) - return MemoryRead(memory_md=result.memory_md) + return MemoryRead(memory_md=result.memory_md, limits=memory_limits()) diff --git a/surfsense_backend/app/routes/team_memory_routes.py b/surfsense_backend/app/routes/team_memory_routes.py index 3e552ce32..b37a99b03 100644 --- a/surfsense_backend/app/routes/team_memory_routes.py +++ b/surfsense_backend/app/routes/team_memory_routes.py @@ -8,7 +8,9 @@ from sqlalchemy.ext.asyncio import AsyncSession from app.db import User, get_async_session from app.services.memory import ( + MemoryRead, MemoryScope, + memory_limits, read_memory, reset_memory, save_memory, @@ -19,15 +21,11 @@ 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) +@router.get("/searchspaces/{search_space_id}/memory", response_model=MemoryRead) async def get_team_memory( search_space_id: int, session: AsyncSession = Depends(get_async_session), @@ -39,10 +37,10 @@ async def get_team_memory( target_id=search_space_id, session=session, ) - return TeamMemoryRead(memory_md=memory_md) + return MemoryRead(memory_md=memory_md, limits=memory_limits()) -@router.put("/searchspaces/{search_space_id}/memory", response_model=TeamMemoryRead) +@router.put("/searchspaces/{search_space_id}/memory", response_model=MemoryRead) async def update_team_memory( search_space_id: int, body: TeamMemoryUpdate, @@ -58,10 +56,10 @@ async def update_team_memory( ) if result.status == "error": raise HTTPException(status_code=400, detail=result.message) - return TeamMemoryRead(memory_md=result.memory_md) + return MemoryRead(memory_md=result.memory_md, limits=memory_limits()) -@router.post("/searchspaces/{search_space_id}/memory/reset", response_model=TeamMemoryRead) +@router.post("/searchspaces/{search_space_id}/memory/reset", response_model=MemoryRead) async def reset_team_memory( search_space_id: int, session: AsyncSession = Depends(get_async_session), @@ -75,4 +73,4 @@ async def reset_team_memory( ) if result.status == "error": raise HTTPException(status_code=400, detail=result.message) - return TeamMemoryRead(memory_md=result.memory_md) + return MemoryRead(memory_md=result.memory_md, limits=memory_limits()) diff --git a/surfsense_backend/app/services/memory/__init__.py b/surfsense_backend/app/services/memory/__init__.py index d72f45e1f..27d0592fd 100644 --- a/surfsense_backend/app/services/memory/__init__.py +++ b/surfsense_backend/app/services/memory/__init__.py @@ -1,9 +1,11 @@ """First-class memory service for user and team markdown memory.""" +from .schemas import MemoryLimits, MemoryRead from .service import ( MemoryScope, SaveResult, extract_and_save, + memory_limits, read_memory, reset_memory, save_memory, @@ -18,9 +20,12 @@ from .validation import ( __all__ = [ "MEMORY_HARD_LIMIT", "MEMORY_SOFT_LIMIT", + "MemoryLimits", + "MemoryRead", "MemoryScope", "SaveResult", "extract_and_save", + "memory_limits", "read_memory", "reset_memory", "save_memory", diff --git a/surfsense_backend/app/services/memory/schemas.py b/surfsense_backend/app/services/memory/schemas.py index 9b40ee5b1..623e4aa93 100644 --- a/surfsense_backend/app/services/memory/schemas.py +++ b/surfsense_backend/app/services/memory/schemas.py @@ -1,4 +1,4 @@ -"""Structured output schemas for memory extraction.""" +"""Schemas for memory API responses and structured extraction.""" from __future__ import annotations @@ -7,6 +7,20 @@ from typing import Literal from pydantic import BaseModel, Field +class MemoryLimits(BaseModel): + """Canonical memory size limits exposed to clients.""" + + soft: int + hard: int + + +class MemoryRead(BaseModel): + """Memory document payload returned by user and team memory APIs.""" + + memory_md: str + limits: MemoryLimits + + class MemoryExtractionDecision(BaseModel): """Structured extraction result; avoids string sentinel parsing.""" diff --git a/surfsense_backend/app/services/memory/service.py b/surfsense_backend/app/services/memory/service.py index 85459c28c..8159977a7 100644 --- a/surfsense_backend/app/services/memory/service.py +++ b/surfsense_backend/app/services/memory/service.py @@ -9,7 +9,6 @@ from typing import Any, Literal from uuid import UUID from langchain_core.messages import HumanMessage -from pydantic import BaseModel from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession @@ -19,9 +18,10 @@ from app.services.memory.prompts import ( USER_MEMORY_EXTRACT_PROMPT, ) from app.services.memory.rewrite import forced_rewrite -from app.services.memory.schemas import MemoryExtractionDecision +from app.services.memory.schemas import MemoryExtractionDecision, MemoryLimits from app.services.memory.validation import ( MEMORY_HARD_LIMIT, + MEMORY_SOFT_LIMIT, soft_limit_warning, strip_preamble_to_first_heading, validate_bullet_format, @@ -68,8 +68,8 @@ class SaveResult: return data -class MemoryRead(BaseModel): - memory_md: str +def memory_limits() -> MemoryLimits: + return MemoryLimits(soft=MEMORY_SOFT_LIMIT, hard=MEMORY_HARD_LIMIT) def _normalize_scope(scope: MemoryScope | str) -> MemoryScope: diff --git a/surfsense_web/app/dashboard/[search_space_id]/user-settings/components/MemoryContent.tsx b/surfsense_web/app/dashboard/[search_space_id]/user-settings/components/MemoryContent.tsx index dc002244f..c7cb3d1d4 100644 --- a/surfsense_web/app/dashboard/[search_space_id]/user-settings/components/MemoryContent.tsx +++ b/surfsense_web/app/dashboard/[search_space_id]/user-settings/components/MemoryContent.tsx @@ -14,11 +14,11 @@ import { DropdownMenuTrigger, } from "@/components/ui/dropdown-menu"; import { Spinner } from "@/components/ui/spinner"; -import { MEMORY_HARD_LIMIT, useUserMemory } from "@/hooks/use-memory"; +import { getMemoryLimitState, useUserMemory } from "@/hooks/use-memory"; export function MemoryContent() { const activeSearchSpaceId = useAtomValue(activeSearchSpaceIdAtom); - const { memory, displayMemory, loading, saving, reset } = useUserMemory( + const { memory, displayMemory, limits, loading, saving, reset } = useUserMemory( Number(activeSearchSpaceId) ); @@ -59,11 +59,11 @@ export function MemoryContent() { }; const charCount = memory.length; + const limitState = getMemoryLimitState(charCount, limits); const getCounterColor = () => { - if (charCount > MEMORY_HARD_LIMIT) return "text-red-500"; - if (charCount > 15_000) return "text-orange-500"; - if (charCount > 10_000) return "text-yellow-500"; + if (limitState.level === "error") return "text-red-500"; + if (limitState.level === "warning") return "text-orange-500"; return "text-muted-foreground"; }; @@ -112,13 +112,7 @@ export function MemoryContent() {
- - {charCount.toLocaleString()} / {MEMORY_HARD_LIMIT.toLocaleString()} - characters - chars - {charCount > 15_000 && charCount <= MEMORY_HARD_LIMIT && " - Approaching limit"} - {charCount > MEMORY_HARD_LIMIT && " - Exceeds limit"} - + {limitState.label}
- - - - - - - - Copy as Markdown - - - - Download as Markdown - - - -
-
- - ); -} diff --git a/surfsense_web/app/dashboard/[search_space_id]/user-settings/layout-shell.tsx b/surfsense_web/app/dashboard/[search_space_id]/user-settings/layout-shell.tsx index 820021622..037568db3 100644 --- a/surfsense_web/app/dashboard/[search_space_id]/user-settings/layout-shell.tsx +++ b/surfsense_web/app/dashboard/[search_space_id]/user-settings/layout-shell.tsx @@ -1,7 +1,6 @@ "use client"; import { - Brain, CircleUser, Keyboard, KeyRound, @@ -26,7 +25,6 @@ export type UserSettingsTab = | "api-key" | "prompts" | "community-prompts" - | "memory" | "agent-permissions" | "agent-status" | "purchases" @@ -75,11 +73,6 @@ export function UserSettingsLayoutShell({ searchSpaceId, children }: UserSetting label: "Community Prompts", icon: , }, - { - value: "memory" as const, - label: "Memory", - icon: , - }, { value: "agent-permissions" as const, label: "Agent Permissions", diff --git a/surfsense_web/app/dashboard/[search_space_id]/user-settings/memory/page.tsx b/surfsense_web/app/dashboard/[search_space_id]/user-settings/memory/page.tsx deleted file mode 100644 index b10c5bce5..000000000 --- a/surfsense_web/app/dashboard/[search_space_id]/user-settings/memory/page.tsx +++ /dev/null @@ -1,5 +0,0 @@ -import { MemoryContent } from "../components/MemoryContent"; - -export default function Page() { - return ; -} diff --git a/surfsense_web/components/settings/team-memory-manager.tsx b/surfsense_web/components/settings/team-memory-manager.tsx deleted file mode 100644 index 4a730d45f..000000000 --- a/surfsense_web/components/settings/team-memory-manager.tsx +++ /dev/null @@ -1,151 +0,0 @@ -"use client"; - -import { ChevronDown, ClipboardCopy, Download, Info } from "lucide-react"; -import { toast } from "sonner"; -import { PlateEditor } from "@/components/editor/plate-editor"; -import { Alert, AlertDescription } from "@/components/ui/alert"; -import { Button } from "@/components/ui/button"; -import { - DropdownMenu, - DropdownMenuContent, - DropdownMenuItem, - DropdownMenuTrigger, -} from "@/components/ui/dropdown-menu"; -import { Spinner } from "@/components/ui/spinner"; -import { getMemoryLimitState, useTeamMemory } from "@/hooks/use-memory"; - -interface TeamMemoryManagerProps { - searchSpaceId: number; -} - -export function TeamMemoryManager({ searchSpaceId }: TeamMemoryManagerProps) { - const { memory, displayMemory, limits, loading, saving, reset } = useTeamMemory(searchSpaceId); - - const handleClear = async () => { - try { - await reset(); - toast.success("Team memory cleared"); - } catch { - toast.error("Failed to clear team memory"); - } - }; - - const handleDownload = () => { - if (!memory) return; - try { - const blob = new Blob([memory], { type: "text/markdown;charset=utf-8" }); - const url = URL.createObjectURL(blob); - const a = document.createElement("a"); - a.href = url; - a.download = "team-memory.md"; - document.body.appendChild(a); - a.click(); - document.body.removeChild(a); - URL.revokeObjectURL(url); - } catch { - toast.error("Failed to download team memory"); - } - }; - - const handleCopyMarkdown = async () => { - if (!memory) return; - try { - await navigator.clipboard.writeText(memory); - toast.success("Copied to clipboard"); - } catch { - toast.error("Failed to copy team memory"); - } - }; - - const charCount = memory.length; - const limitState = getMemoryLimitState(charCount, limits); - - const getCounterColor = () => { - if (limitState.level === "error") return "text-red-500"; - if (limitState.level === "warning") return "text-orange-500"; - return "text-muted-foreground"; - }; - - if (loading) { - return ( -
- -
- ); - } - - if (!memory) { - return ( -
-

- What does SurfSense remember about your team? -

-

- Nothing yet. SurfSense picks up on team decisions and conventions as your team chats. -

-
- ); - } - - return ( -
- - - -

- SurfSense uses this shared memory to provide team-wide context across all conversations - in this search space. -

- - - -
-
- -
-
- -
- {limitState.label} -
- - - - - - - - - Copy as Markdown - - - - Download as Markdown - - - -
-
-
- ); -} From 78a3c71bb59fb6f876b62f2c88c68a046b9e0644 Mon Sep 17 00:00:00 2001 From: Anish Sarkar <104695310+AnishSarkar22@users.noreply.github.com> Date: Wed, 20 May 2026 12:50:15 +0530 Subject: [PATCH 10/15] feat: implement memory document fetching and saving functionality in the editor panel, and remove deprecated memory hook --- .../components/documents/DocumentNode.tsx | 8 +- .../components/editor-panel/editor-panel.tsx | 78 +++------- .../components/editor-panel/memory.ts | 116 ++++++++++++++ surfsense_web/hooks/use-memory.ts | 141 ------------------ 4 files changed, 140 insertions(+), 203 deletions(-) create mode 100644 surfsense_web/components/editor-panel/memory.ts delete mode 100644 surfsense_web/hooks/use-memory.ts diff --git a/surfsense_web/components/documents/DocumentNode.tsx b/surfsense_web/components/documents/DocumentNode.tsx index a5b02cbb3..86c9b899a 100644 --- a/surfsense_web/components/documents/DocumentNode.tsx +++ b/surfsense_web/components/documents/DocumentNode.tsx @@ -216,13 +216,9 @@ export const DocumentNode = React.memo(function DocumentNode({ return ( <> {isMemoryDocument ? ( - + ) : canMention ? ( ({ detail: "Failed to fetch memory" })); - throw new Error(errorData.detail || "Failed to fetch memory"); - } - const data = (await response.json()) as { - memory_md?: string; - limits?: MemoryLimits; - }; - setMemoryLimits(data.limits ?? null); - const content: EditorContent = { - document_id: memoryScope === "team" ? -1002 : -1001, - title: title || (memoryScope === "team" ? "Team Memory" : "Personal Memory"), - document_type: memoryScope === "team" ? "TEAM_MEMORY" : "USER_MEMORY", - source_markdown: data.memory_md ?? "", - }; + setMemoryLimits(limits); + const content: EditorContent = document; markdownRef.current = content.source_markdown; setDisplayTitle(content.title); setEditorDoc(content); @@ -370,34 +356,14 @@ export function EditorPanelContent({ return true; } if (isMemoryMode) { - if (memoryScope === "team" && !searchSpaceId) { - throw new Error("Missing search space context"); - } - const response = await authenticatedFetch( - `${process.env.NEXT_PUBLIC_FASTAPI_BACKEND_URL}${ - memoryScope === "team" - ? `/api/v1/searchspaces/${searchSpaceId}/memory` - : "/api/v1/users/me/memory" - }`, - { - method: "PUT", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ memory_md: markdownRef.current }), - } - ); - if (!response.ok) { - const errorData = await response - .json() - .catch(() => ({ detail: "Failed to save memory" })); - throw new Error(errorData.detail || "Failed to save memory"); - } - const data = (await response.json()) as { - memory_md?: string; - limits?: MemoryLimits; - }; - const savedContent = data.memory_md ?? markdownRef.current; + if (!memoryScope) throw new Error("Missing memory context"); + const { markdown: savedContent, limits } = await saveMemoryMarkdown({ + scope: memoryScope, + searchSpaceId, + markdown: markdownRef.current, + }); markdownRef.current = savedContent; - setMemoryLimits(data.limits ?? memoryLimits); + setMemoryLimits(limits ?? memoryLimits); setEditorDoc((prev) => (prev ? { ...prev, source_markdown: savedContent } : prev)); setEditedMarkdown(null); if (!options?.silent) { diff --git a/surfsense_web/components/editor-panel/memory.ts b/surfsense_web/components/editor-panel/memory.ts new file mode 100644 index 000000000..aa5b1f68d --- /dev/null +++ b/surfsense_web/components/editor-panel/memory.ts @@ -0,0 +1,116 @@ +"use client"; + +import { authenticatedFetch } from "@/lib/auth-utils"; + +export type MemoryScope = "user" | "team"; + +export interface MemoryLimits { + soft: number; + hard: number; +} + +export type MemoryLimitLevel = "ok" | "warning" | "error"; + +export interface MemoryEditorDocument { + document_id: number; + title: string; + document_type: "USER_MEMORY" | "TEAM_MEMORY"; + source_markdown: string; +} + +interface MemoryReadResponse { + memory_md?: string; + limits?: MemoryLimits; +} + +function getMemoryPath(scope: MemoryScope, searchSpaceId?: number | null) { + if (scope === "user") return "/api/v1/users/me/memory"; + if (!searchSpaceId) throw new Error("Missing search space context"); + return `/api/v1/searchspaces/${searchSpaceId}/memory`; +} + +function getBackendUrl(path: string) { + return `${process.env.NEXT_PUBLIC_FASTAPI_BACKEND_URL}${path}`; +} + +export function getMemoryLimitState(length: number, limits?: MemoryLimits | null) { + if (!limits) { + return { + level: "ok" as MemoryLimitLevel, + label: `${length.toLocaleString()} chars`, + isOverLimit: false, + }; + } + + const isOverLimit = length > limits.hard; + const isNearLimit = length > limits.soft; + const level: MemoryLimitLevel = isOverLimit ? "error" : isNearLimit ? "warning" : "ok"; + const suffix = isOverLimit ? " - Exceeds limit" : isNearLimit ? " - Approaching limit" : ""; + + return { + level, + label: `${length.toLocaleString()}/${limits.hard.toLocaleString()} chars${suffix}`, + isOverLimit, + }; +} + +export async function fetchMemoryEditorDocument({ + scope, + searchSpaceId, + title, + signal, +}: { + scope: MemoryScope; + searchSpaceId?: number | null; + title?: string | null; + signal?: AbortSignal; +}) { + const response = await authenticatedFetch(getBackendUrl(getMemoryPath(scope, searchSpaceId)), { + method: "GET", + signal, + }); + if (!response.ok) { + const errorData = await response.json().catch(() => ({ detail: "Failed to fetch memory" })); + throw new Error(errorData.detail || "Failed to fetch memory"); + } + + const data = (await response.json()) as MemoryReadResponse; + const isTeamMemory = scope === "team"; + + return { + limits: data.limits ?? null, + document: { + document_id: isTeamMemory ? -1002 : -1001, + title: title || (isTeamMemory ? "Team Memory" : "Personal Memory"), + document_type: isTeamMemory ? "TEAM_MEMORY" : "USER_MEMORY", + source_markdown: data.memory_md ?? "", + } satisfies MemoryEditorDocument, + }; +} + +export async function saveMemoryMarkdown({ + scope, + searchSpaceId, + markdown, +}: { + scope: MemoryScope; + searchSpaceId?: number | null; + markdown: string; +}) { + const response = await authenticatedFetch(getBackendUrl(getMemoryPath(scope, searchSpaceId)), { + method: "PUT", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ memory_md: markdown }), + }); + if (!response.ok) { + const errorData = await response.json().catch(() => ({ detail: "Failed to save memory" })); + throw new Error(errorData.detail || "Failed to save memory"); + } + + const data = (await response.json()) as MemoryReadResponse; + + return { + markdown: data.memory_md ?? markdown, + limits: data.limits, + }; +} diff --git a/surfsense_web/hooks/use-memory.ts b/surfsense_web/hooks/use-memory.ts deleted file mode 100644 index 609aad537..000000000 --- a/surfsense_web/hooks/use-memory.ts +++ /dev/null @@ -1,141 +0,0 @@ -"use client"; - -import { useCallback, useEffect, useState } from "react"; -import { z } from "zod"; -import { baseApiService } from "@/lib/apis/base-api.service"; - -const MemoryLimitsSchema = z.object({ - soft: z.number(), - hard: z.number(), -}); - -const MemoryReadSchema = z.object({ - memory_md: z.string(), - limits: MemoryLimitsSchema, -}); - -type MemoryScope = "user" | "team"; -export type MemoryLimits = z.infer; -export type MemoryLimitLevel = "ok" | "warning" | "error"; - -interface UseMemoryOptions { - scope: MemoryScope; - searchSpaceId?: number | null; - autoLoad?: boolean; -} - -function getMemoryPath(scope: MemoryScope, searchSpaceId?: number | null) { - if (scope === "user") return "/api/v1/users/me/memory"; - if (!searchSpaceId) throw new Error("searchSpaceId is required for team memory"); - return `/api/v1/searchspaces/${searchSpaceId}/memory`; -} - -export function stripMemoryDisplayPrefixes(memory: string) { - return memory.replace( - /^\s*-\s+(?:\(\d{4}-\d{2}-\d{2}\)\s*\[(?:fact|pref|instr)\]\s*|\d{4}-\d{2}-\d{2}:\s*)/gim, - "- " - ); -} - -export function getMemoryLimitState(length: number, limits?: MemoryLimits | null) { - if (!limits) { - return { - level: "ok" as MemoryLimitLevel, - label: `${length.toLocaleString()} chars`, - isOverLimit: false, - }; - } - - const isOverLimit = length > limits.hard; - const isNearLimit = length > limits.soft; - const level: MemoryLimitLevel = isOverLimit ? "error" : isNearLimit ? "warning" : "ok"; - const suffix = isOverLimit ? " - Exceeds limit" : isNearLimit ? " - Approaching limit" : ""; - - return { - level, - label: `${length.toLocaleString()}/${limits.hard.toLocaleString()} chars${suffix}`, - isOverLimit, - }; -} - -export function useMemory({ scope, searchSpaceId, autoLoad = true }: UseMemoryOptions) { - const [memory, setMemory] = useState(""); - const [limits, setLimits] = useState(null); - const [loading, setLoading] = useState(autoLoad); - const [saving, setSaving] = useState(false); - - const load = useCallback(async () => { - setLoading(true); - try { - const data = await baseApiService.get(getMemoryPath(scope, searchSpaceId), MemoryReadSchema); - setMemory(data.memory_md); - setLimits(data.limits); - return data.memory_md; - } finally { - setLoading(false); - } - }, [scope, searchSpaceId]); - - useEffect(() => { - if (!autoLoad) return; - load().catch(() => { - setLoading(false); - }); - }, [autoLoad, load]); - - const save = useCallback( - async (memoryMd: string) => { - setSaving(true); - try { - const data = await baseApiService.put( - getMemoryPath(scope, searchSpaceId), - MemoryReadSchema, - { - body: { memory_md: memoryMd }, - } - ); - setMemory(data.memory_md); - setLimits(data.limits); - return data.memory_md; - } finally { - setSaving(false); - } - }, - [scope, searchSpaceId] - ); - - const reset = useCallback(async () => { - setSaving(true); - try { - const data = await baseApiService.post( - `${getMemoryPath(scope, searchSpaceId)}/reset`, - MemoryReadSchema - ); - setMemory(data.memory_md); - setLimits(data.limits); - return data.memory_md; - } finally { - setSaving(false); - } - }, [scope, searchSpaceId]); - - return { - memory, - setMemory, - limits, - displayMemory: stripMemoryDisplayPrefixes(memory), - loading, - saving, - load, - save, - reset, - }; -} - -export function useUserMemory(searchSpaceId?: number | null) { - return useMemory({ scope: "user", searchSpaceId }); -} - -export function useTeamMemory(searchSpaceId?: number | null) { - return useMemory({ scope: "team", searchSpaceId }); -} From fe07de3f9c027cad75493ce8e0670d347d497fc3 Mon Sep 17 00:00:00 2001 From: Anish Sarkar <104695310+AnishSarkar22@users.noreply.github.com> Date: Wed, 20 May 2026 12:55:10 +0530 Subject: [PATCH 11/15] chore: ran linting --- surfsense_backend/app/services/memory/service.py | 8 ++++++-- surfsense_backend/app/services/memory/validation.py | 4 +++- .../agents/new_chat/tools/test_update_memory_scope.py | 4 +++- .../tests/unit/services/test_memory_service.py | 4 +++- surfsense_web/components/documents/DocumentNode.tsx | 5 +---- surfsense_web/components/editor-panel/editor-panel.tsx | 2 +- .../components/layout/ui/sidebar/DocumentsSidebar.tsx | 5 ++++- 7 files changed, 21 insertions(+), 11 deletions(-) diff --git a/surfsense_backend/app/services/memory/service.py b/surfsense_backend/app/services/memory/service.py index 8159977a7..d4a7d0974 100644 --- a/surfsense_backend/app/services/memory/service.py +++ b/surfsense_backend/app/services/memory/service.py @@ -92,7 +92,9 @@ async def _load_target( select(User).where(User.id == _normalize_user_id(target_id)) # type: ignore[arg-type] ) return result.scalars().first() - result = await session.execute(select(SearchSpace).where(SearchSpace.id == int(target_id))) + result = await session.execute( + select(SearchSpace).where(SearchSpace.id == int(target_id)) + ) return result.scalars().first() @@ -141,7 +143,9 @@ async def save_memory( if target is None: return SaveResult( status="error", - message="User not found." if normalized is MemoryScope.USER else "Search space not found.", + message="User not found." + if normalized is MemoryScope.USER + else "Search space not found.", ) old_memory = _get_memory(target, normalized) diff --git a/surfsense_backend/app/services/memory/validation.py b/surfsense_backend/app/services/memory/validation.py index 0e856943b..f9c5007d9 100644 --- a/surfsense_backend/app/services/memory/validation.py +++ b/surfsense_backend/app/services/memory/validation.py @@ -11,7 +11,9 @@ MEMORY_HARD_LIMIT = 25_000 _SECTION_HEADING_RE = re.compile(r"^##\s+(.+)$", re.MULTILINE) _HEADING_LINE_RE = re.compile(r"^##\s+\S+", re.MULTILINE) _HEADING_NORMALIZE_RE = re.compile(r"[^a-z0-9]+") -_LEGACY_BULLET_RE = re.compile(r"^-\s+\(\d{4}-\d{2}-\d{2}\)\s+\[(fact|pref|instr)\]\s+.+$") +_LEGACY_BULLET_RE = re.compile( + r"^-\s+\(\d{4}-\d{2}-\d{2}\)\s+\[(fact|pref|instr)\]\s+.+$" +) _NEW_BULLET_RE = re.compile(r"^-\s+\d{4}-\d{2}-\d{2}:\s+.+$") _FORBIDDEN_TEAM_HEADINGS = { diff --git a/surfsense_backend/tests/unit/agents/new_chat/tools/test_update_memory_scope.py b/surfsense_backend/tests/unit/agents/new_chat/tools/test_update_memory_scope.py index 60310d907..f1a0f97f0 100644 --- a/surfsense_backend/tests/unit/agents/new_chat/tools/test_update_memory_scope.py +++ b/surfsense_backend/tests/unit/agents/new_chat/tools/test_update_memory_scope.py @@ -88,7 +88,9 @@ async def test_save_memory_blocks_new_personal_heading_in_team_before_commit( @pytest.mark.asyncio -async def test_save_memory_allows_grandfathered_personal_heading_in_team(monkeypatch) -> None: +async def test_save_memory_allows_grandfathered_personal_heading_in_team( + monkeypatch, +) -> None: content = "## Preferences\n- 2026-04-10: Prefers dark mode\n" target = type("Target", (), {"shared_memory_md": content})() session = _FakeSession() diff --git a/surfsense_backend/tests/unit/services/test_memory_service.py b/surfsense_backend/tests/unit/services/test_memory_service.py index c16e34062..e7fef2cac 100644 --- a/surfsense_backend/tests/unit/services/test_memory_service.py +++ b/surfsense_backend/tests/unit/services/test_memory_service.py @@ -108,7 +108,9 @@ async def test_save_memory_rejects_long_no_heading_payload(monkeypatch) -> None: @pytest.mark.asyncio -async def test_save_memory_grandfathers_existing_team_personal_heading(monkeypatch) -> None: +async def test_save_memory_grandfathers_existing_team_personal_heading( + monkeypatch, +) -> None: content = "## Preferences\n- 2026-05-19: Existing legacy heading\n" target = SimpleNamespace(shared_memory_md=content) session = _FakeSession() diff --git a/surfsense_web/components/documents/DocumentNode.tsx b/surfsense_web/components/documents/DocumentNode.tsx index 86c9b899a..a13bd0079 100644 --- a/surfsense_web/components/documents/DocumentNode.tsx +++ b/surfsense_web/components/documents/DocumentNode.tsx @@ -216,10 +216,7 @@ export const DocumentNode = React.memo(function DocumentNode({ return ( <> {isMemoryDocument ? ( - + Date: Wed, 20 May 2026 13:20:05 +0530 Subject: [PATCH 12/15] feat: add memory document model and parsing functionality for markdown handling --- .../app/services/memory/document.py | 200 ++++++++++++++++++ .../app/services/memory/service.py | 3 + .../app/services/memory/validation.py | 52 ++--- .../tools/test_update_memory_scope.py | 21 ++ .../unit/services/test_memory_service.py | 2 +- 5 files changed, 241 insertions(+), 37 deletions(-) create mode 100644 surfsense_backend/app/services/memory/document.py diff --git a/surfsense_backend/app/services/memory/document.py b/surfsense_backend/app/services/memory/document.py new file mode 100644 index 000000000..498195e25 --- /dev/null +++ b/surfsense_backend/app/services/memory/document.py @@ -0,0 +1,200 @@ +"""Memory-specific markdown document model and canonical renderer. + +This intentionally parses only SurfSense memory's small markdown contract: +``##`` sections with dated bullet items. Unknown lines are preserved so user +edits are not lost, while legacy marker bullets are normalized on render. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from datetime import date + +DEFAULT_LEGACY_SECTION = "Memory" +LEGACY_MARKERS = frozenset({"fact", "pref", "instr"}) + + +@dataclass(frozen=True) +class MemoryBullet: + entry_date: date + text: str + + +@dataclass(frozen=True) +class MemoryRawLine: + text: str + + +MemoryLine = MemoryBullet | MemoryRawLine + + +@dataclass(frozen=True) +class MemorySection: + heading: str + lines: list[MemoryLine] = field(default_factory=list) + explicit_heading: bool = True + + +@dataclass(frozen=True) +class MemoryDocument: + sections: list[MemorySection] = field(default_factory=list) + + @property + def has_explicit_heading(self) -> bool: + return any(section.explicit_heading for section in self.sections) + + +def is_section_heading(line: str) -> bool: + return line.startswith("## ") and bool(line[3:].strip()) + + +def heading_text(line: str) -> str: + return line[3:].strip() + + +def normalize_heading(heading: str) -> str: + chars: list[str] = [] + previous_was_space = True + for char in heading.strip().lower(): + if char.isalnum(): + chars.append(char) + previous_was_space = False + elif not previous_was_space: + chars.append(" ") + previous_was_space = True + return "".join(chars).strip() + + +def parse_bullet_line(line: str) -> MemoryBullet | None: + stripped = line.strip() + if not stripped.startswith("- "): + return None + + body = stripped[2:] + parsed = _parse_canonical_bullet(body) + if parsed is not None: + return parsed + return _parse_legacy_bullet(body) + + +def _parse_canonical_bullet(body: str) -> MemoryBullet | None: + if len(body) < 13 or body[10:12] != ": ": + return None + try: + entry_date = date.fromisoformat(body[:10]) + except ValueError: + return None + text = body[12:].strip() + if not text: + return None + return MemoryBullet(entry_date=entry_date, text=text) + + +def _parse_legacy_bullet(body: str) -> MemoryBullet | None: + if len(body) < 20 or not body.startswith("("): + return None + if len(body) < 14 or body[11:14] != ") [": + return None + try: + entry_date = date.fromisoformat(body[1:11]) + except ValueError: + return None + + marker_end = body.find("] ", 14) + if marker_end == -1: + return None + marker = body[14:marker_end] + if marker not in LEGACY_MARKERS: + return None + + text = body[marker_end + 2 :].strip() + if not text: + return None + return MemoryBullet(entry_date=entry_date, text=text) + + +def parse_memory_document(content: str | None) -> MemoryDocument: + if not content: + return MemoryDocument() + + sections: list[MemorySection] = [] + current_heading: str | None = None + current_explicit = True + current_lines: list[MemoryLine] = [] + + def flush_current() -> None: + nonlocal current_heading, current_explicit, current_lines + if current_heading is None: + return + sections.append( + MemorySection( + heading=current_heading, + lines=current_lines, + explicit_heading=current_explicit, + ) + ) + current_heading = None + current_explicit = True + current_lines = [] + + for raw_line in content.strip().splitlines(): + line = raw_line.rstrip() + if is_section_heading(line): + flush_current() + current_heading = heading_text(line) + current_explicit = True + current_lines = [] + continue + + bullet = parse_bullet_line(line) + if current_heading is None: + if bullet is None: + continue + current_heading = DEFAULT_LEGACY_SECTION + current_explicit = False + current_lines = [bullet] + continue + + current_lines.append(bullet if bullet is not None else MemoryRawLine(text=line)) + + flush_current() + return MemoryDocument(sections=sections) + + +def render_memory_document(document: MemoryDocument) -> str: + rendered_sections: list[str] = [] + for section in document.sections: + section_lines = [f"## {section.heading}"] + for line in section.lines: + if isinstance(line, MemoryBullet): + section_lines.append(f"- {line.entry_date.isoformat()}: {line.text}") + else: + section_lines.append(line.text) + rendered_sections.append("\n".join(section_lines).strip()) + return "\n\n".join(section for section in rendered_sections if section).strip() + + +def extract_headings(memory: str | None) -> set[str]: + document = parse_memory_document(memory) + return { + normalize_heading(section.heading) + for section in document.sections + if section.explicit_heading + } + + +def has_explicit_heading(content: str) -> bool: + return parse_memory_document(content).has_explicit_heading + + +def nonstandard_bullets(content: str) -> list[str]: + warnings: list[str] = [] + for line in content.splitlines(): + stripped = line.strip() + if not stripped.startswith("- "): + continue + if parse_bullet_line(stripped) is not None: + continue + short = stripped[:80] + ("..." if len(stripped) > 80 else "") + warnings.append(f"Non-standard memory bullet: {short}") + return warnings diff --git a/surfsense_backend/app/services/memory/service.py b/surfsense_backend/app/services/memory/service.py index d4a7d0974..dd4459e77 100644 --- a/surfsense_backend/app/services/memory/service.py +++ b/surfsense_backend/app/services/memory/service.py @@ -13,6 +13,7 @@ from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from app.db import SearchSpace, User +from app.services.memory.document import parse_memory_document, render_memory_document from app.services.memory.prompts import ( TEAM_MEMORY_EXTRACT_PROMPT, USER_MEMORY_EXTRACT_PROMPT, @@ -184,6 +185,8 @@ async def save_memory( warnings=warnings, ) + next_content = render_memory_document(parse_memory_document(next_content)) + try: _set_memory(target, normalized, next_content) session.add(target) diff --git a/surfsense_backend/app/services/memory/validation.py b/surfsense_backend/app/services/memory/validation.py index f9c5007d9..6565f39c7 100644 --- a/surfsense_backend/app/services/memory/validation.py +++ b/surfsense_backend/app/services/memory/validation.py @@ -2,20 +2,18 @@ from __future__ import annotations -import re from typing import Literal +from app.services.memory.document import ( + extract_headings, + has_explicit_heading, + nonstandard_bullets, + parse_memory_document, +) + MEMORY_SOFT_LIMIT = 18_000 MEMORY_HARD_LIMIT = 25_000 -_SECTION_HEADING_RE = re.compile(r"^##\s+(.+)$", re.MULTILINE) -_HEADING_LINE_RE = re.compile(r"^##\s+\S+", re.MULTILINE) -_HEADING_NORMALIZE_RE = re.compile(r"[^a-z0-9]+") -_LEGACY_BULLET_RE = re.compile( - r"^-\s+\(\d{4}-\d{2}-\d{2}\)\s+\[(fact|pref|instr)\]\s+.+$" -) -_NEW_BULLET_RE = re.compile(r"^-\s+\d{4}-\d{2}-\d{2}:\s+.+$") - _FORBIDDEN_TEAM_HEADINGS = { "preferences", "instructions", @@ -25,25 +23,16 @@ _FORBIDDEN_TEAM_HEADINGS = { def has_markdown_heading(content: str) -> bool: - return bool(_HEADING_LINE_RE.search(content)) + return has_explicit_heading(content) def strip_preamble_to_first_heading(content: str) -> str: """Drop model preamble before the first ``##`` heading, if one exists.""" - match = _HEADING_LINE_RE.search(content) - if not match: - return content.strip() - return content[match.start() :].strip() - - -def extract_headings(memory: str | None) -> set[str]: - if not memory: - return set() - return {_normalize_heading(h) for h in _SECTION_HEADING_RE.findall(memory)} - - -def _normalize_heading(heading: str) -> str: - return _HEADING_NORMALIZE_RE.sub(" ", heading.strip().lower()).strip() + lines = content.splitlines() + for index, line in enumerate(lines): + if line.startswith("## ") and line[3:].strip(): + return "\n".join(lines[index:]).strip() + return content.strip() def validate_memory_size(content: str) -> dict[str, str] | None: @@ -69,7 +58,7 @@ def validate_heading_sanity(content: str) -> dict[str, str] | None: return None if len(stripped) <= 40: return None - if any(_LEGACY_BULLET_RE.match(line.strip()) for line in stripped.splitlines()): + if parse_memory_document(stripped).sections: return None return { "status": "error", @@ -115,16 +104,7 @@ def validate_memory_scope( def validate_bullet_format(content: str) -> list[str]: - warnings: list[str] = [] - for line in content.splitlines(): - stripped = line.strip() - if not stripped.startswith("- "): - continue - if _NEW_BULLET_RE.match(stripped) or _LEGACY_BULLET_RE.match(stripped): - continue - short = stripped[:80] + ("..." if len(stripped) > 80 else "") - warnings.append(f"Non-standard memory bullet: {short}") - return warnings + return nonstandard_bullets(content) def validate_diff(old_memory: str | None, new_memory: str) -> list[str]: @@ -138,7 +118,7 @@ def validate_diff(old_memory: str | None, new_memory: str) -> list[str]: if dropped: names = ", ".join(sorted(dropped)) warnings.append( - f"Sections removed: {names}. If unintentional, restore from the settings page." + f"Sections removed: {names}. If unintentional, restore them from the memory document." ) old_len = len(old_memory) diff --git a/surfsense_backend/tests/unit/agents/new_chat/tools/test_update_memory_scope.py b/surfsense_backend/tests/unit/agents/new_chat/tools/test_update_memory_scope.py index f1a0f97f0..c941d7d65 100644 --- a/surfsense_backend/tests/unit/agents/new_chat/tools/test_update_memory_scope.py +++ b/surfsense_backend/tests/unit/agents/new_chat/tools/test_update_memory_scope.py @@ -64,6 +64,27 @@ def test_validate_bullet_format_warns_on_nonstandard_bullet() -> None: assert "Non-standard memory bullet" in warnings[0] +@pytest.mark.asyncio +async def test_save_memory_normalizes_legacy_marker_bullets(monkeypatch) -> None: + target = type("Target", (), {"memory_md": ""})() + session = _FakeSession() + + async def fake_load_target(**_kwargs): + return target + + monkeypatch.setattr("app.services.memory.service._load_target", fake_load_target) + + result = await save_memory( + scope=MemoryScope.USER, + target_id="00000000-0000-0000-0000-000000000000", + content="- (2026-04-10) [fact] Legacy fact is preserved\n", + session=session, + ) + + assert result.status == "saved" + assert target.memory_md == "## Memory\n- 2026-04-10: Legacy fact is preserved" + + @pytest.mark.asyncio async def test_save_memory_blocks_new_personal_heading_in_team_before_commit( monkeypatch, diff --git a/surfsense_backend/tests/unit/services/test_memory_service.py b/surfsense_backend/tests/unit/services/test_memory_service.py index e7fef2cac..0a45bf3aa 100644 --- a/surfsense_backend/tests/unit/services/test_memory_service.py +++ b/surfsense_backend/tests/unit/services/test_memory_service.py @@ -82,7 +82,7 @@ async def test_save_memory_accepts_legacy_marker_payload(monkeypatch) -> None: ) assert result.status == "saved" - assert "[fact]" in target.memory_md + assert target.memory_md == "## Memory\n- 2026-05-19: Legacy marker memory" @pytest.mark.asyncio From 132e7b3c44488cba85831f0f6c49e284036cade8 Mon Sep 17 00:00:00 2001 From: Anish Sarkar <104695310+AnishSarkar22@users.noreply.github.com> Date: Wed, 20 May 2026 14:03:28 +0530 Subject: [PATCH 13/15] refactor: remove memory extraction functions and related components from the new chat agent --- .../app/agents/new_chat/memory_extraction.py | 78 ---------------- .../app/services/memory/__init__.py | 2 - .../app/services/memory/prompts.py | 90 ------------------- .../app/services/memory/schemas.py | 20 +---- .../app/services/memory/service.py | 78 +--------------- .../app/tasks/chat/stream_new_chat.py | 35 -------- .../streaming/graph_stream/event_stream.py | 1 - .../chat/streaming/graph_stream/result.py | 1 - .../tasks/chat/streaming/handlers/tool_end.py | 3 - .../app/tasks/chat/streaming/relay/state.py | 1 - .../unit/services/test_memory_service.py | 67 -------------- .../chat/streaming/test_stream_output.py | 1 - 12 files changed, 2 insertions(+), 375 deletions(-) delete mode 100644 surfsense_backend/app/agents/new_chat/memory_extraction.py diff --git a/surfsense_backend/app/agents/new_chat/memory_extraction.py b/surfsense_backend/app/agents/new_chat/memory_extraction.py deleted file mode 100644 index d44b58f7b..000000000 --- a/surfsense_backend/app/agents/new_chat/memory_extraction.py +++ /dev/null @@ -1,78 +0,0 @@ -"""Background memory extraction for the SurfSense agent.""" - -from __future__ import annotations - -import logging -from typing import Any -from uuid import UUID - -from app.db import User, shielded_async_session -from app.services.memory import MemoryScope, extract_and_save - -logger = logging.getLogger(__name__) - - -async def extract_and_save_memory( - *, - user_message: str, - user_id: str | None, - llm: Any, -) -> None: - """Fire-and-forget personal memory extraction. - - The service uses structured output, so free-form ``NO_UPDATE`` text can no - longer be accidentally persisted as memory. - """ - if not user_id: - return - - try: - uid = UUID(user_id) if isinstance(user_id, str) else user_id - async with shielded_async_session() as session: - user = await session.get(User, uid) - actor_display_name = user.display_name if user else None - result = await extract_and_save( - scope=MemoryScope.USER, - target_id=uid, - user_message=user_message, - actor_display_name=actor_display_name, - session=session, - llm=llm, - ) - logger.info( - "Background memory extraction for user %s: %s", - uid, - result.status, - ) - except Exception: - logger.exception("Background user memory extraction failed") - - -async def extract_and_save_team_memory( - *, - user_message: str, - search_space_id: int | None, - llm: Any, - author_display_name: str | None = None, -) -> None: - """Fire-and-forget team-level memory extraction.""" - if not search_space_id: - return - - try: - async with shielded_async_session() as session: - result = await extract_and_save( - scope=MemoryScope.TEAM, - target_id=search_space_id, - user_message=user_message, - actor_display_name=author_display_name, - session=session, - llm=llm, - ) - logger.info( - "Background team memory extraction for space %s: %s", - search_space_id, - result.status, - ) - except Exception: - logger.exception("Background team memory extraction failed") diff --git a/surfsense_backend/app/services/memory/__init__.py b/surfsense_backend/app/services/memory/__init__.py index 27d0592fd..eef6559c2 100644 --- a/surfsense_backend/app/services/memory/__init__.py +++ b/surfsense_backend/app/services/memory/__init__.py @@ -4,7 +4,6 @@ from .schemas import MemoryLimits, MemoryRead from .service import ( MemoryScope, SaveResult, - extract_and_save, memory_limits, read_memory, reset_memory, @@ -24,7 +23,6 @@ __all__ = [ "MemoryRead", "MemoryScope", "SaveResult", - "extract_and_save", "memory_limits", "read_memory", "reset_memory", diff --git a/surfsense_backend/app/services/memory/prompts.py b/surfsense_backend/app/services/memory/prompts.py index fbf27fd08..25c09e9c5 100644 --- a/surfsense_backend/app/services/memory/prompts.py +++ b/surfsense_backend/app/services/memory/prompts.py @@ -18,93 +18,3 @@ RULES: {content} """ - -USER_MEMORY_EXTRACT_PROMPT = """\ -You are a memory extraction assistant. Analyze the user's message and decide \ -if it contains any long-term information worth persisting to personal memory. - -Worth remembering: preferences, background/identity, goals, projects, \ -instructions, tools/languages they use, decisions, expertise, workplace — \ -durable facts that will matter in future conversations. - -NOT worth remembering: greetings, one-off factual questions, session \ -logistics, ephemeral requests, follow-up clarifications with no new personal \ -info, things that only matter for the current task. - -If there is nothing durable to remember, choose `action = no_update`. - -If the message contains memorizable information, choose `action = save` and \ -return the FULL updated memory document with the new information merged into \ -existing content. - -FORMAT RULES FOR `updated_memory`: -- Markdown only. -- Every entry should be under a `##` heading. -- Recommended headings: `## Facts`, `## Preferences`, `## Instructions`. -- New bullets should use: `- YYYY-MM-DD: memory text`. -- If current memory uses legacy `(YYYY-MM-DD) [fact|pref|instr]` markers, - preserve the information but write the updated document in the new - heading-based format. -- Use the user's first name from `` when helpful, not "the user". -- Do not duplicate existing information. - -{user_name} - - -{current_memory} - - - -{user_message} -""" - -TEAM_MEMORY_EXTRACT_PROMPT = """\ -You are a team-memory extraction assistant. Analyze the latest message and \ -decide if it contains durable TEAM-level information worth persisting. - -Decision policy: -- Prioritize recall for durable team context, while avoiding personal-only facts. -- Do NOT require explicit consensus language. A direct team-level statement can - be stored if it is stable and broadly useful for future team chats. -- If evidence is weak or clearly tentative, choose `action = no_update`. - -Worth remembering (team-level only): -- Decisions and defaults that guide future team work -- Team conventions/standards (naming, review policy, coding norms) -- Stable org/project facts (locations, ownership, constraints) -- Long-lived architecture/process facts -- Ongoing priorities that are likely relevant beyond this turn - -NOT worth remembering: -- Personal preferences or biography of one person -- Questions, brainstorming, tentative ideas, or speculation -- One-off requests, status updates, TODOs, logistics for this session -- Information scoped only to a single ephemeral task - -If the message contains memorizable team information, choose `action = save` \ -and return the FULL updated team memory document with new facts merged into \ -existing content. - -FORMAT RULES FOR `updated_memory`: -- Markdown only. -- Every entry should be under a `##` heading. -- Recommended headings: `## Product Decisions`, `## Engineering Conventions`, - `## Project Facts`, `## Open Questions`. -- New bullets should use: `- YYYY-MM-DD: memory text`. -- If current memory uses legacy `(YYYY-MM-DD) [fact]` markers, preserve the - information but write the updated document in the new heading-based format. -- Do not create personal headings such as `## Preferences`, `## Instructions`, - or `## Personal Notes`. -- Preserve neutral team phrasing; avoid person-specific memory unless role-anchored. - - -{current_memory} - - - -{author} - - - -{user_message} -""" diff --git a/surfsense_backend/app/services/memory/schemas.py b/surfsense_backend/app/services/memory/schemas.py index 623e4aa93..78c69d800 100644 --- a/surfsense_backend/app/services/memory/schemas.py +++ b/surfsense_backend/app/services/memory/schemas.py @@ -2,9 +2,7 @@ from __future__ import annotations -from typing import Literal - -from pydantic import BaseModel, Field +from pydantic import BaseModel class MemoryLimits(BaseModel): @@ -19,19 +17,3 @@ class MemoryRead(BaseModel): memory_md: str limits: MemoryLimits - - -class MemoryExtractionDecision(BaseModel): - """Structured extraction result; avoids string sentinel parsing.""" - - action: Literal["no_update", "save"] = Field( - description="Choose no_update when nothing durable should be saved; choose save otherwise." - ) - reason: str | None = Field( - default=None, - description="Short reason for no_update, or brief summary of the memory update.", - ) - updated_memory: str | None = Field( - default=None, - description="The full updated markdown memory document when action is save.", - ) diff --git a/surfsense_backend/app/services/memory/service.py b/surfsense_backend/app/services/memory/service.py index dd4459e77..c33b91679 100644 --- a/surfsense_backend/app/services/memory/service.py +++ b/surfsense_backend/app/services/memory/service.py @@ -8,18 +8,13 @@ from enum import StrEnum from typing import Any, Literal from uuid import UUID -from langchain_core.messages import HumanMessage from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from app.db import SearchSpace, User from app.services.memory.document import parse_memory_document, render_memory_document -from app.services.memory.prompts import ( - TEAM_MEMORY_EXTRACT_PROMPT, - USER_MEMORY_EXTRACT_PROMPT, -) from app.services.memory.rewrite import forced_rewrite -from app.services.memory.schemas import MemoryExtractionDecision, MemoryLimits +from app.services.memory.schemas import MemoryLimits from app.services.memory.validation import ( MEMORY_HARD_LIMIT, MEMORY_SOFT_LIMIT, @@ -234,74 +229,3 @@ async def reset_memory( session=session, llm=None, ) - - -async def extract_and_save( - *, - scope: MemoryScope | str, - target_id: str | int | UUID, - user_message: str, - actor_display_name: str | None, - session: AsyncSession, - llm: Any, -) -> SaveResult: - normalized = _normalize_scope(scope) - current_memory = await read_memory( - scope=normalized, - target_id=target_id, - session=session, - ) - - if normalized is MemoryScope.USER: - first_name = ( - actor_display_name.strip().split()[0] - if actor_display_name and actor_display_name.strip() - else "The user" - ) - prompt = USER_MEMORY_EXTRACT_PROMPT.format( - current_memory=current_memory or "(empty)", - user_message=user_message, - user_name=first_name, - ) - else: - prompt = TEAM_MEMORY_EXTRACT_PROMPT.format( - current_memory=current_memory or "(empty)", - author=actor_display_name or "Unknown team member", - user_message=user_message, - ) - - try: - structured = llm.with_structured_output(MemoryExtractionDecision) - decision = await structured.ainvoke( - [HumanMessage(content=prompt)], - config={"tags": ["surfsense:internal", "memory-extraction"]}, - ) - except Exception: - logger.exception("Structured memory extraction failed") - return SaveResult( - status="error", - message="Structured memory extraction failed.", - memory_md=current_memory, - ) - - if decision.action == "no_update": - return SaveResult( - status="no_op", - message=decision.reason or "No durable memory to persist.", - memory_md=current_memory, - ) - - if not decision.updated_memory: - return SaveResult( - status="error", - message="Structured memory extraction chose save without updated_memory.", - memory_md=current_memory, - ) - - return await save_memory( - scope=normalized, - target_id=target_id, - content=decision.updated_memory, - session=session, - llm=llm, - ) diff --git a/surfsense_backend/app/tasks/chat/stream_new_chat.py b/surfsense_backend/app/tasks/chat/stream_new_chat.py index 9a69b6164..564fd81de 100644 --- a/surfsense_backend/app/tasks/chat/stream_new_chat.py +++ b/surfsense_backend/app/tasks/chat/stream_new_chat.py @@ -39,10 +39,6 @@ from app.agents.new_chat.llm_config import ( load_agent_config, load_global_llm_config_by_id, ) -from app.agents.new_chat.memory_extraction import ( - extract_and_save_memory, - extract_and_save_team_memory, -) from app.agents.new_chat.mention_resolver import resolve_mentions, substitute_in_text from app.agents.new_chat.middleware.busy_mutex import ( end_turn, @@ -283,7 +279,6 @@ class StreamResult: accumulated_text: str = "" is_interrupted: bool = False sandbox_files: list[str] = field(default_factory=list) - agent_called_update_memory: bool = False request_id: str | None = None turn_id: str = "" filesystem_mode: str = "cloud" @@ -2208,36 +2203,6 @@ async def stream_new_chat( }, ) - # Fire background memory extraction if the agent didn't handle it. - # Shared threads write to team memory; private threads write to user memory. - if not stream_result.agent_called_update_memory: - memory_seed = user_query.strip() or ( - f"[{len(user_image_data_urls or [])} image(s)]" - if user_image_data_urls - else "(message)" - ) - if visibility == ChatVisibility.SEARCH_SPACE: - task = asyncio.create_task( - extract_and_save_team_memory( - user_message=memory_seed, - search_space_id=search_space_id, - llm=llm, - author_display_name=current_user_display_name, - ) - ) - _background_tasks.add(task) - task.add_done_callback(_background_tasks.discard) - elif user_id: - task = asyncio.create_task( - extract_and_save_memory( - user_message=memory_seed, - user_id=user_id, - llm=llm, - ) - ) - _background_tasks.add(task) - task.add_done_callback(_background_tasks.discard) - # Finish the step and message yield streaming_service.format_data("turn-status", {"status": "idle"}) yield streaming_service.format_finish_step() diff --git a/surfsense_backend/app/tasks/chat/streaming/graph_stream/event_stream.py b/surfsense_backend/app/tasks/chat/streaming/graph_stream/event_stream.py index 9a309f9d7..50e7a1360 100644 --- a/surfsense_backend/app/tasks/chat/streaming/graph_stream/event_stream.py +++ b/surfsense_backend/app/tasks/chat/streaming/graph_stream/event_stream.py @@ -48,4 +48,3 @@ async def stream_output( yield frame result.accumulated_text = state.accumulated_text - result.agent_called_update_memory = state.called_update_memory diff --git a/surfsense_backend/app/tasks/chat/streaming/graph_stream/result.py b/surfsense_backend/app/tasks/chat/streaming/graph_stream/result.py index 391f14f24..1d3f1e88a 100644 --- a/surfsense_backend/app/tasks/chat/streaming/graph_stream/result.py +++ b/surfsense_backend/app/tasks/chat/streaming/graph_stream/result.py @@ -11,7 +11,6 @@ class StreamingResult: accumulated_text: str = "" is_interrupted: bool = False sandbox_files: list[str] = field(default_factory=list) - agent_called_update_memory: bool = False request_id: str | None = None turn_id: str = "" filesystem_mode: str = "cloud" diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tool_end.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tool_end.py index 57ab617c5..ad4a17d08 100644 --- a/surfsense_backend/app/tasks/chat/streaming/handlers/tool_end.py +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tool_end.py @@ -36,9 +36,6 @@ def iter_tool_end_frames( raw_output = event.get("data", {}).get("output", "") staged_file_path = state.file_path_by_run.pop(run_id, None) if run_id else None - if tool_name == "update_memory": - state.called_update_memory = True - if hasattr(raw_output, "content"): content = raw_output.content if isinstance(content, str): diff --git a/surfsense_backend/app/tasks/chat/streaming/relay/state.py b/surfsense_backend/app/tasks/chat/streaming/relay/state.py index 27898403d..f99fc8edb 100644 --- a/surfsense_backend/app/tasks/chat/streaming/relay/state.py +++ b/surfsense_backend/app/tasks/chat/streaming/relay/state.py @@ -32,7 +32,6 @@ class AgentEventRelayState: last_active_step_items: list[str] = field(default_factory=list) just_finished_tool: bool = False active_tool_depth: int = 0 - called_update_memory: bool = False current_reasoning_id: str | None = None pending_tool_call_chunks: list[dict[str, Any]] = field(default_factory=list) lc_tool_call_id_by_run: dict[str, str] = field(default_factory=dict) diff --git a/surfsense_backend/tests/unit/services/test_memory_service.py b/surfsense_backend/tests/unit/services/test_memory_service.py index 0a45bf3aa..94918d25b 100644 --- a/surfsense_backend/tests/unit/services/test_memory_service.py +++ b/surfsense_backend/tests/unit/services/test_memory_service.py @@ -6,11 +6,9 @@ import pytest from app.services.memory import ( MemoryScope, - extract_and_save, reset_memory, save_memory, ) -from app.services.memory.schemas import MemoryExtractionDecision pytestmark = pytest.mark.unit @@ -31,17 +29,6 @@ class _FakeSession: self.rollback_calls += 1 -class _StructuredLLM: - def __init__(self, decision: MemoryExtractionDecision) -> None: - self.decision = decision - - def with_structured_output(self, _schema): - return self - - async def ainvoke(self, *_args, **_kwargs): - return self.decision - - @pytest.mark.asyncio async def test_save_memory_saves_heading_based_memory(monkeypatch) -> None: target = SimpleNamespace(memory_md="") @@ -150,57 +137,3 @@ async def test_reset_memory_clears_memory(monkeypatch) -> None: assert result.status == "saved" assert target.memory_md == "" - - -@pytest.mark.asyncio -async def test_extract_and_save_no_update_does_not_commit(monkeypatch) -> None: - target = SimpleNamespace(memory_md="## Facts\n- 2026-05-19: Existing\n") - session = _FakeSession() - - async def fake_load_target(**_kwargs): - return target - - monkeypatch.setattr("app.services.memory.service._load_target", fake_load_target) - - result = await extract_and_save( - scope=MemoryScope.USER, - target_id="00000000-0000-0000-0000-000000000000", - user_message="hello", - actor_display_name="Anish", - session=session, - llm=_StructuredLLM( - MemoryExtractionDecision(action="no_update", reason="Greeting only") - ), - ) - - assert result.status == "no_op" - assert session.commit_calls == 0 - - -@pytest.mark.asyncio -async def test_extract_and_save_persists_structured_update(monkeypatch) -> None: - target = SimpleNamespace(memory_md="") - session = _FakeSession() - - async def fake_load_target(**_kwargs): - return target - - monkeypatch.setattr("app.services.memory.service._load_target", fake_load_target) - - result = await extract_and_save( - scope=MemoryScope.USER, - target_id="00000000-0000-0000-0000-000000000000", - user_message="I work on SurfSense", - actor_display_name="Anish", - session=session, - llm=_StructuredLLM( - MemoryExtractionDecision( - action="save", - updated_memory="## Facts\n- 2026-05-19: Anish works on SurfSense\n", - ) - ), - ) - - assert result.status == "saved" - assert "SurfSense" in target.memory_md - assert session.commit_calls == 1 diff --git a/surfsense_backend/tests/unit/tasks/chat/streaming/test_stream_output.py b/surfsense_backend/tests/unit/tasks/chat/streaming/test_stream_output.py index c0123b76d..c53dad5fb 100644 --- a/surfsense_backend/tests/unit/tasks/chat/streaming/test_stream_output.py +++ b/surfsense_backend/tests/unit/tasks/chat/streaming/test_stream_output.py @@ -89,7 +89,6 @@ async def test_stream_output_emits_text_lifecycle_and_updates_result() -> None: "text_end:text-1", ] assert result.accumulated_text == "Hello world" - assert result.agent_called_update_memory is False async def test_stream_output_passes_runtime_context_to_agent() -> None: From 8c9be9796a93a88022d9c8edbe6eed7d1ac57a4f Mon Sep 17 00:00:00 2001 From: Anish Sarkar <104695310+AnishSarkar22@users.noreply.github.com> Date: Wed, 20 May 2026 15:03:35 +0530 Subject: [PATCH 14/15] feat: add no-update sentinel handling to save_memory function and corresponding unit tests --- .../app/services/memory/service.py | 16 +++++++ .../unit/services/test_memory_service.py | 48 +++++++++++++++++++ 2 files changed, 64 insertions(+) diff --git a/surfsense_backend/app/services/memory/service.py b/surfsense_backend/app/services/memory/service.py index c33b91679..feca000c9 100644 --- a/surfsense_backend/app/services/memory/service.py +++ b/surfsense_backend/app/services/memory/service.py @@ -29,6 +29,15 @@ from app.services.memory.validation import ( logger = logging.getLogger(__name__) +_NO_UPDATE_SENTINELS = frozenset( + { + "NO_UPDATE", + "NO UPDATE", + "NO_CHANGE", + "NO CHANGE", + } +) + class MemoryScope(StrEnum): USER = "user" @@ -149,6 +158,13 @@ async def save_memory( notice: str | None = None warnings: list[str] = [] + if next_content.upper() in _NO_UPDATE_SENTINELS: + return SaveResult( + status="no_op", + message="No memory update requested.", + memory_md=old_memory, + ) + if len(next_content) > MEMORY_HARD_LIMIT and llm is not None: rewritten = await forced_rewrite(next_content, llm) if rewritten is not None and len(rewritten) < len(next_content): diff --git a/surfsense_backend/tests/unit/services/test_memory_service.py b/surfsense_backend/tests/unit/services/test_memory_service.py index 94918d25b..820e6aa28 100644 --- a/surfsense_backend/tests/unit/services/test_memory_service.py +++ b/surfsense_backend/tests/unit/services/test_memory_service.py @@ -94,6 +94,54 @@ async def test_save_memory_rejects_long_no_heading_payload(monkeypatch) -> None: assert target.memory_md.startswith("## Facts") +@pytest.mark.asyncio +async def test_save_memory_no_update_sentinel_is_no_op(monkeypatch) -> None: + existing = "## Preferences\n- 2026-05-20: Existing preference\n" + target = SimpleNamespace(memory_md=existing) + session = _FakeSession() + + async def fake_load_target(**_kwargs): + return target + + monkeypatch.setattr("app.services.memory.service._load_target", fake_load_target) + + result = await save_memory( + scope=MemoryScope.USER, + target_id="00000000-0000-0000-0000-000000000000", + content="NO_UPDATE", + session=session, + ) + + assert result.status == "no_op" + assert result.memory_md == existing + assert target.memory_md == existing + assert session.commit_calls == 0 + + +@pytest.mark.asyncio +async def test_save_memory_no_update_sentinel_is_case_insensitive(monkeypatch) -> None: + existing = "## Preferences\n- 2026-05-20: Existing preference\n" + target = SimpleNamespace(memory_md=existing) + session = _FakeSession() + + async def fake_load_target(**_kwargs): + return target + + monkeypatch.setattr("app.services.memory.service._load_target", fake_load_target) + + result = await save_memory( + scope=MemoryScope.USER, + target_id="00000000-0000-0000-0000-000000000000", + content=" no update ", + session=session, + ) + + assert result.status == "no_op" + assert result.memory_md == existing + assert target.memory_md == existing + assert session.commit_calls == 0 + + @pytest.mark.asyncio async def test_save_memory_grandfathers_existing_team_personal_heading( monkeypatch, From 39c29d651f50dd958ef9fe20459f0b2de892d820 Mon Sep 17 00:00:00 2001 From: Anish Sarkar <104695310+AnishSarkar22@users.noreply.github.com> Date: Wed, 20 May 2026 15:29:41 +0530 Subject: [PATCH 15/15] feat: enhance token display in MessageInfoDropdown with improved visual separation --- .../assistant-ui/assistant-message.tsx | 25 +++++++++++++------ 1 file changed, 17 insertions(+), 8 deletions(-) diff --git a/surfsense_web/components/assistant-ui/assistant-message.tsx b/surfsense_web/components/assistant-ui/assistant-message.tsx index ac1732441..f6c91e8bf 100644 --- a/surfsense_web/components/assistant-ui/assistant-message.tsx +++ b/surfsense_web/components/assistant-ui/assistant-message.tsx @@ -14,6 +14,7 @@ import { ClipboardPaste, CopyIcon, DownloadIcon, + Dot, ExternalLink, Globe, MessageCircleReply, @@ -330,9 +331,14 @@ const MessageInfoDropdown: FC<{ chatTurnId: string | null | undefined }> = ({ ch {icon} {name} - - {counts.total_tokens.toLocaleString()} tokens - {costMicros && costMicros > 0 ? ` · ${formatTurnCost(costMicros)}` : ""} + + {counts.total_tokens.toLocaleString()} tokens + {costMicros && costMicros > 0 ? ( + <> + ); @@ -342,11 +348,14 @@ const MessageInfoDropdown: FC<{ chatTurnId: string | null | undefined }> = ({ ch className="focus:bg-accent focus:text-accent-foreground relative flex cursor-default flex-col items-start gap-0.5 rounded-sm px-2 py-1.5 text-sm outline-hidden select-none" onSelect={(e) => e.preventDefault()} > - - {usage.total_tokens.toLocaleString()} tokens - {usage.cost_micros && usage.cost_micros > 0 - ? ` · ${formatTurnCost(usage.cost_micros)}` - : ""} + + {usage.total_tokens.toLocaleString()} tokens + {usage.cost_micros && usage.cost_micros > 0 ? ( + <> + )}