Wire thread_visibility to agent; memory tools use private or shared backend by visibility

This commit is contained in:
CREDO23 2026-02-06 17:23:46 +02:00
parent d71a2be66f
commit 0306597df7
3 changed files with 139 additions and 19 deletions

View file

@ -22,6 +22,7 @@ from app.agents.new_chat.system_prompt import (
build_surfsense_system_prompt, build_surfsense_system_prompt,
) )
from app.agents.new_chat.tools.registry import build_tools_async from app.agents.new_chat.tools.registry import build_tools_async
from app.db import ChatVisibility
from app.services.connector_service import ConnectorService from app.services.connector_service import ConnectorService
# ============================================================================= # =============================================================================
@ -126,6 +127,7 @@ async def create_surfsense_deep_agent(
disabled_tools: list[str] | None = None, disabled_tools: list[str] | None = None,
additional_tools: Sequence[BaseTool] | None = None, additional_tools: Sequence[BaseTool] | None = None,
firecrawl_api_key: str | None = None, firecrawl_api_key: str | None = None,
thread_visibility: ChatVisibility | None = None,
): ):
""" """
Create a SurfSense deep agent with configurable tools and prompts. 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}") 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 = { dependencies = {
"search_space_id": search_space_id, "search_space_id": search_space_id,
"db_session": db_session, "db_session": db_session,
"connector_service": connector_service, "connector_service": connector_service,
"firecrawl_api_key": firecrawl_api_key, "firecrawl_api_key": firecrawl_api_key,
"user_id": user_id, # Required for memory tools "user_id": user_id,
"thread_id": thread_id, # For podcast tool "thread_id": thread_id,
# Dynamic connector/document type discovery for knowledge base tool "thread_visibility": visibility,
"available_connectors": available_connectors, "available_connectors": available_connectors,
"available_document_types": available_document_types, "available_document_types": available_document_types,
} }

View file

@ -51,8 +51,14 @@ from .mcp_tool import load_mcp_tools
from .podcast import create_generate_podcast_tool from .podcast import create_generate_podcast_tool
from .scrape_webpage import create_scrape_webpage_tool from .scrape_webpage import create_scrape_webpage_tool
from .search_surfsense_docs import create_search_surfsense_docs_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 .user_memory import create_recall_memory_tool, create_save_memory_tool
from app.db import ChatVisibility
# ============================================================================= # =============================================================================
# Tool Definition # Tool Definition
# ============================================================================= # =============================================================================
@ -156,29 +162,42 @@ BUILTIN_TOOLS: list[ToolDefinition] = [
requires=["db_session"], 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( ToolDefinition(
name="save_memory", name="save_memory",
description="Save facts, preferences, or context about the user for personalized responses", description="Save facts, preferences, or context for personalized or team responses",
factory=lambda deps: create_save_memory_tool( 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"], user_id=deps["user_id"],
search_space_id=deps["search_space_id"], search_space_id=deps["search_space_id"],
db_session=deps["db_session"], 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( ToolDefinition(
name="recall_memory", name="recall_memory",
description="Recall user memories for personalized and contextual responses", description="Recall relevant memories (personal or team) for context",
factory=lambda deps: create_recall_memory_tool( 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"], user_id=deps["user_id"],
search_space_id=deps["search_space_id"], search_space_id=deps["search_space_id"],
db_session=deps["db_session"], 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 # ADD YOUR CUSTOM TOOLS BELOW

View file

@ -4,6 +4,7 @@ import logging
from typing import Any from typing import Any
from uuid import UUID from uuid import UUID
from langchain_core.tools import tool
from sqlalchemy import select from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.ext.asyncio import AsyncSession
@ -177,3 +178,101 @@ def format_shared_memories_for_context(
) )
parts.append("</team_memories>") parts.append("</team_memories>")
return "\n".join(parts) 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