diff --git a/surfsense_backend/app/agents/new_chat/chat_deepagent.py b/surfsense_backend/app/agents/new_chat/chat_deepagent.py index 9da6ea3c2..726f82285 100644 --- a/surfsense_backend/app/agents/new_chat/chat_deepagent.py +++ b/surfsense_backend/app/agents/new_chat/chat_deepagent.py @@ -22,6 +22,7 @@ from app.agents.new_chat.system_prompt import ( build_surfsense_system_prompt, ) from app.agents.new_chat.tools.registry import build_tools_async +from app.db import ChatVisibility from app.services.connector_service import ConnectorService # ============================================================================= @@ -126,6 +127,7 @@ async def create_surfsense_deep_agent( disabled_tools: list[str] | None = None, additional_tools: Sequence[BaseTool] | None = None, firecrawl_api_key: str | None = None, + thread_visibility: ChatVisibility | None = None, ): """ Create a SurfSense deep agent with configurable tools and prompts. @@ -227,15 +229,15 @@ async def create_surfsense_deep_agent( logging.warning(f"Failed to discover available connectors/document types: {e}") - # Build dependencies dict for the tools registry + visibility = thread_visibility or ChatVisibility.PRIVATE dependencies = { "search_space_id": search_space_id, "db_session": db_session, "connector_service": connector_service, "firecrawl_api_key": firecrawl_api_key, - "user_id": user_id, # Required for memory tools - "thread_id": thread_id, # For podcast tool - # Dynamic connector/document type discovery for knowledge base tool + "user_id": user_id, + "thread_id": thread_id, + "thread_visibility": visibility, "available_connectors": available_connectors, "available_document_types": available_document_types, } diff --git a/surfsense_backend/app/agents/new_chat/tools/registry.py b/surfsense_backend/app/agents/new_chat/tools/registry.py index 2cf43c973..30201e8df 100644 --- a/surfsense_backend/app/agents/new_chat/tools/registry.py +++ b/surfsense_backend/app/agents/new_chat/tools/registry.py @@ -51,8 +51,14 @@ from .mcp_tool import load_mcp_tools from .podcast import create_generate_podcast_tool from .scrape_webpage import create_scrape_webpage_tool from .search_surfsense_docs import create_search_surfsense_docs_tool +from .shared_memory import ( + create_recall_shared_memory_tool, + create_save_shared_memory_tool, +) from .user_memory import create_recall_memory_tool, create_save_memory_tool +from app.db import ChatVisibility + # ============================================================================= # Tool Definition # ============================================================================= @@ -156,29 +162,42 @@ BUILTIN_TOOLS: list[ToolDefinition] = [ requires=["db_session"], ), # ========================================================================= - # USER MEMORY TOOLS - Claude-like memory feature + # USER MEMORY TOOLS - private or team store by thread_visibility # ========================================================================= - # Save memory tool - stores facts/preferences about the user ToolDefinition( name="save_memory", - description="Save facts, preferences, or context about the user for personalized responses", - factory=lambda deps: create_save_memory_tool( - user_id=deps["user_id"], - search_space_id=deps["search_space_id"], - db_session=deps["db_session"], + description="Save facts, preferences, or context for personalized or team responses", + factory=lambda deps: ( + create_save_shared_memory_tool( + search_space_id=deps["search_space_id"], + created_by_id=deps["user_id"], + db_session=deps["db_session"], + ) + if deps["thread_visibility"] == ChatVisibility.SEARCH_SPACE + else create_save_memory_tool( + user_id=deps["user_id"], + search_space_id=deps["search_space_id"], + db_session=deps["db_session"], + ) ), - requires=["user_id", "search_space_id", "db_session"], + requires=["user_id", "search_space_id", "db_session", "thread_visibility"], ), - # Recall memory tool - retrieves relevant user memories ToolDefinition( name="recall_memory", - description="Recall user memories for personalized and contextual responses", - factory=lambda deps: create_recall_memory_tool( - user_id=deps["user_id"], - search_space_id=deps["search_space_id"], - db_session=deps["db_session"], + description="Recall relevant memories (personal or team) for context", + factory=lambda deps: ( + create_recall_shared_memory_tool( + search_space_id=deps["search_space_id"], + db_session=deps["db_session"], + ) + if deps["thread_visibility"] == ChatVisibility.SEARCH_SPACE + else create_recall_memory_tool( + user_id=deps["user_id"], + search_space_id=deps["search_space_id"], + db_session=deps["db_session"], + ) ), - requires=["user_id", "search_space_id", "db_session"], + requires=["user_id", "search_space_id", "db_session", "thread_visibility"], ), # ========================================================================= # ADD YOUR CUSTOM TOOLS BELOW diff --git a/surfsense_backend/app/agents/new_chat/tools/shared_memory.py b/surfsense_backend/app/agents/new_chat/tools/shared_memory.py index 57d158f16..aa4a738ce 100644 --- a/surfsense_backend/app/agents/new_chat/tools/shared_memory.py +++ b/surfsense_backend/app/agents/new_chat/tools/shared_memory.py @@ -4,6 +4,7 @@ import logging from typing import Any from uuid import UUID +from langchain_core.tools import tool from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession @@ -177,3 +178,101 @@ def format_shared_memories_for_context( ) parts.append("") return "\n".join(parts) + + +def create_save_shared_memory_tool( + search_space_id: int, + created_by_id: str | UUID, + db_session: AsyncSession, +): + """ + Factory function to create the save_memory tool for shared (team) chats. + + Args: + search_space_id: The search space ID + created_by_id: The user ID of the person adding the memory + db_session: Database session for executing queries + + Returns: + A configured tool function for saving team memories + """ + + @tool + async def save_memory( + content: str, + category: str = "fact", + ) -> dict[str, Any]: + """ + Save a fact, preference, or context to the team's shared memory for future reference. + + Use this tool when: + - User or a team member says "remember this", "keep this in mind", or similar in this shared chat + - The team agrees on something to remember (e.g., decisions, conventions, where things live) + - Someone shares a preference or fact that should be visible to the whole team + + The saved information will be available in future shared conversations in this space. + + Args: + content: The fact/preference/context to remember. + Phrase it clearly, e.g., "API keys are stored in Vault", + "The team prefers weekly demos on Fridays" + category: Type of memory. One of: + - "preference": Team or workspace preferences + - "fact": Facts the team agreed on (e.g., processes, locations) + - "instruction": Standing instructions for the team + - "context": Current context (e.g., ongoing projects, goals) + + Returns: + A dictionary with the save status and memory details + """ + return await save_shared_memory( + db_session, search_space_id, created_by_id, content, category + ) + + return save_memory + + +def create_recall_shared_memory_tool( + search_space_id: int, + db_session: AsyncSession, +): + """ + Factory function to create the recall_memory tool for shared (team) chats. + + Args: + search_space_id: The search space ID + db_session: Database session for executing queries + + Returns: + A configured tool function for recalling team memories + """ + + @tool + async def recall_memory( + query: str | None = None, + category: str | None = None, + top_k: int = DEFAULT_RECALL_TOP_K, + ) -> dict[str, Any]: + """ + Recall relevant team memories for this space to provide contextual responses. + + Use this tool when: + - You need team context to answer (e.g., "where do we store X?", "what did we decide about Y?") + - Someone asks about something the team agreed to remember + - Team preferences or conventions would improve the response + + Args: + query: Optional search query to find specific memories. + If not provided, returns the most recent memories. + category: Optional category filter. One of: + "preference", "fact", "instruction", "context" + top_k: Number of memories to retrieve (default: 5, max: 20) + + Returns: + A dictionary containing relevant memories and formatted context + """ + return await recall_shared_memory( + db_session, search_space_id, query, category, top_k + ) + + return recall_memory