mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-06-12 20:45:20 +02:00
feat: drop legacy memory tables and update memory management tools to streamline user and shared memory handling
This commit is contained in:
Anish Sarkar
parent
f46f3b4989
commit
b8091114b5
6 changed files with 93 additions and 758 deletions
|
|
@ -0,0 +1,91 @@
|
|||
"""Drop legacy user_memories and shared_memories tables
|
||||
|
||||
Revision ID: 122
|
||||
Revises: 121
|
||||
|
||||
The old row-per-fact memory system (user_memories, shared_memories tables and
|
||||
memorycategory enum) is replaced by memory_md / shared_memory_md TEXT columns
|
||||
added in migration 121.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import op
|
||||
from app.config import config
|
||||
|
||||
revision: str = "122"
|
||||
down_revision: str | None = "121"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
EMBEDDING_DIM = config.embedding_model_instance.dimension
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.execute("DROP TABLE IF EXISTS shared_memories CASCADE;")
|
||||
op.execute("DROP TABLE IF EXISTS user_memories CASCADE;")
|
||||
op.execute("DROP TYPE IF EXISTS memorycategory;")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.execute(
|
||||
"""
|
||||
DO $$
|
||||
BEGIN
|
||||
IF NOT EXISTS (SELECT 1 FROM pg_type WHERE typname = 'memorycategory') THEN
|
||||
CREATE TYPE memorycategory AS ENUM (
|
||||
'preference',
|
||||
'fact',
|
||||
'instruction',
|
||||
'context'
|
||||
);
|
||||
END IF;
|
||||
END$$;
|
||||
"""
|
||||
)
|
||||
|
||||
op.execute(
|
||||
f"""
|
||||
CREATE TABLE IF NOT EXISTS user_memories (
|
||||
id SERIAL PRIMARY KEY,
|
||||
created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
|
||||
user_id UUID NOT NULL REFERENCES "user"(id) ON DELETE CASCADE,
|
||||
search_space_id INTEGER REFERENCES searchspaces(id) ON DELETE CASCADE,
|
||||
memory_text TEXT NOT NULL,
|
||||
category memorycategory NOT NULL DEFAULT 'fact',
|
||||
embedding vector({EMBEDDING_DIM}),
|
||||
updated_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW()
|
||||
);
|
||||
"""
|
||||
)
|
||||
op.execute("CREATE INDEX IF NOT EXISTS ix_user_memories_user_id ON user_memories(user_id);")
|
||||
op.execute("CREATE INDEX IF NOT EXISTS ix_user_memories_search_space_id ON user_memories(search_space_id);")
|
||||
op.execute("CREATE INDEX IF NOT EXISTS ix_user_memories_updated_at ON user_memories(updated_at);")
|
||||
op.execute("CREATE INDEX IF NOT EXISTS ix_user_memories_category ON user_memories(category);")
|
||||
op.execute("CREATE INDEX IF NOT EXISTS ix_user_memories_user_search_space ON user_memories(user_id, search_space_id);")
|
||||
op.execute(
|
||||
"CREATE INDEX IF NOT EXISTS user_memories_vector_index ON user_memories USING hnsw (embedding public.vector_cosine_ops);"
|
||||
)
|
||||
|
||||
op.execute(
|
||||
f"""
|
||||
CREATE TABLE IF NOT EXISTS shared_memories (
|
||||
id SERIAL PRIMARY KEY,
|
||||
created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
|
||||
updated_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
|
||||
search_space_id INTEGER NOT NULL REFERENCES searchspaces(id) ON DELETE CASCADE,
|
||||
created_by_id UUID NOT NULL REFERENCES "user"(id) ON DELETE CASCADE,
|
||||
memory_text TEXT NOT NULL,
|
||||
category memorycategory NOT NULL DEFAULT 'fact',
|
||||
embedding vector({EMBEDDING_DIM})
|
||||
);
|
||||
"""
|
||||
)
|
||||
op.execute("CREATE INDEX IF NOT EXISTS ix_shared_memories_search_space_id ON shared_memories(search_space_id);")
|
||||
op.execute("CREATE INDEX IF NOT EXISTS ix_shared_memories_updated_at ON shared_memories(updated_at);")
|
||||
op.execute("CREATE INDEX IF NOT EXISTS ix_shared_memories_created_by_id ON shared_memories(created_by_id);")
|
||||
op.execute(
|
||||
"CREATE INDEX IF NOT EXISTS shared_memories_vector_index ON shared_memories USING hnsw (embedding public.vector_cosine_ops);"
|
||||
)
|
||||
|
|
@ -169,8 +169,7 @@ async def create_surfsense_deep_agent(
|
|||
- generate_podcast: Generate audio podcasts from content
|
||||
- generate_image: Generate images from text descriptions using AI models
|
||||
- scrape_webpage: Extract content from webpages
|
||||
- save_memory: Store facts/preferences about the user
|
||||
- recall_memory: Retrieve relevant user memories
|
||||
- update_memory: Update the user's personal or team memory document
|
||||
|
||||
The agent also includes TodoListMiddleware by default (via create_deep_agent) which provides:
|
||||
- write_todos: Create and update planning/todo lists for complex tasks
|
||||
|
|
|
|||
|
|
@ -248,7 +248,7 @@ _TOOL_INSTRUCTIONS["web_search"] = """
|
|||
"""
|
||||
|
||||
# Memory tool instructions have private and shared variants.
|
||||
# We store them keyed as "save_memory" / "recall_memory" with sub-keys.
|
||||
# We store them keyed as "update_memory" with sub-keys.
|
||||
_MEMORY_TOOL_INSTRUCTIONS: dict[str, dict[str, str]] = {
|
||||
"update_memory": {
|
||||
"private": """
|
||||
|
|
|
|||
|
|
@ -1,281 +0,0 @@
|
|||
"""Shared (team) memory backend for search-space-scoped AI context."""
|
||||
|
||||
import asyncio
|
||||
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
|
||||
|
||||
from app.db import MemoryCategory, SharedMemory, User
|
||||
from app.utils.document_converters import embed_text
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
DEFAULT_RECALL_TOP_K = 5
|
||||
MAX_MEMORIES_PER_SEARCH_SPACE = 250
|
||||
|
||||
|
||||
async def get_shared_memory_count(
|
||||
db_session: AsyncSession,
|
||||
search_space_id: int,
|
||||
) -> int:
|
||||
result = await db_session.execute(
|
||||
select(SharedMemory).where(SharedMemory.search_space_id == search_space_id)
|
||||
)
|
||||
return len(result.scalars().all())
|
||||
|
||||
|
||||
async def delete_oldest_shared_memory(
|
||||
db_session: AsyncSession,
|
||||
search_space_id: int,
|
||||
) -> None:
|
||||
result = await db_session.execute(
|
||||
select(SharedMemory)
|
||||
.where(SharedMemory.search_space_id == search_space_id)
|
||||
.order_by(SharedMemory.updated_at.asc())
|
||||
.limit(1)
|
||||
)
|
||||
oldest = result.scalars().first()
|
||||
if oldest:
|
||||
await db_session.delete(oldest)
|
||||
await db_session.commit()
|
||||
|
||||
|
||||
def _to_uuid(value: str | UUID) -> UUID:
|
||||
if isinstance(value, UUID):
|
||||
return value
|
||||
return UUID(value)
|
||||
|
||||
|
||||
async def save_shared_memory(
|
||||
db_session: AsyncSession,
|
||||
search_space_id: int,
|
||||
created_by_id: str | UUID,
|
||||
content: str,
|
||||
category: str = "fact",
|
||||
) -> dict[str, Any]:
|
||||
category = category.lower() if category else "fact"
|
||||
valid = ["preference", "fact", "instruction", "context"]
|
||||
if category not in valid:
|
||||
category = "fact"
|
||||
try:
|
||||
count = await get_shared_memory_count(db_session, search_space_id)
|
||||
if count >= MAX_MEMORIES_PER_SEARCH_SPACE:
|
||||
await delete_oldest_shared_memory(db_session, search_space_id)
|
||||
embedding = await asyncio.to_thread(embed_text, content)
|
||||
row = SharedMemory(
|
||||
search_space_id=search_space_id,
|
||||
created_by_id=_to_uuid(created_by_id),
|
||||
memory_text=content,
|
||||
category=MemoryCategory(category),
|
||||
embedding=embedding,
|
||||
)
|
||||
db_session.add(row)
|
||||
await db_session.commit()
|
||||
await db_session.refresh(row)
|
||||
return {
|
||||
"status": "saved",
|
||||
"memory_id": row.id,
|
||||
"memory_text": content,
|
||||
"category": category,
|
||||
"message": f"I'll remember: {content}",
|
||||
}
|
||||
except Exception as e:
|
||||
logger.exception("Failed to save shared memory: %s", e)
|
||||
await db_session.rollback()
|
||||
return {
|
||||
"status": "error",
|
||||
"error": str(e),
|
||||
"message": "Failed to save memory. Please try again.",
|
||||
}
|
||||
|
||||
|
||||
async def recall_shared_memory(
|
||||
db_session: AsyncSession,
|
||||
search_space_id: int,
|
||||
query: str | None = None,
|
||||
category: str | None = None,
|
||||
top_k: int = DEFAULT_RECALL_TOP_K,
|
||||
) -> dict[str, Any]:
|
||||
top_k = min(max(top_k, 1), 20)
|
||||
try:
|
||||
valid_categories = ["preference", "fact", "instruction", "context"]
|
||||
stmt = select(SharedMemory).where(
|
||||
SharedMemory.search_space_id == search_space_id
|
||||
)
|
||||
if category and category in valid_categories:
|
||||
stmt = stmt.where(SharedMemory.category == MemoryCategory(category))
|
||||
if query:
|
||||
query_embedding = await asyncio.to_thread(embed_text, query)
|
||||
stmt = stmt.order_by(
|
||||
SharedMemory.embedding.op("<=>")(query_embedding)
|
||||
).limit(top_k)
|
||||
else:
|
||||
stmt = stmt.order_by(SharedMemory.updated_at.desc()).limit(top_k)
|
||||
result = await db_session.execute(stmt)
|
||||
rows = result.scalars().all()
|
||||
memory_list = [
|
||||
{
|
||||
"id": m.id,
|
||||
"memory_text": m.memory_text,
|
||||
"category": m.category.value if m.category else "unknown",
|
||||
"updated_at": m.updated_at.isoformat() if m.updated_at else None,
|
||||
"created_by_id": str(m.created_by_id) if m.created_by_id else None,
|
||||
}
|
||||
for m in rows
|
||||
]
|
||||
created_by_ids = list(
|
||||
{m["created_by_id"] for m in memory_list if m["created_by_id"]}
|
||||
)
|
||||
created_by_map: dict[str, str] = {}
|
||||
if created_by_ids:
|
||||
uuids = [UUID(uid) for uid in created_by_ids]
|
||||
users_result = await db_session.execute(
|
||||
select(User).where(User.id.in_(uuids))
|
||||
)
|
||||
for u in users_result.scalars().all():
|
||||
created_by_map[str(u.id)] = u.display_name or "A team member"
|
||||
formatted_context = format_shared_memories_for_context(
|
||||
memory_list, created_by_map
|
||||
)
|
||||
return {
|
||||
"status": "success",
|
||||
"count": len(memory_list),
|
||||
"memories": memory_list,
|
||||
"formatted_context": formatted_context,
|
||||
}
|
||||
except Exception as e:
|
||||
logger.exception("Failed to recall shared memory: %s", e)
|
||||
await db_session.rollback()
|
||||
return {
|
||||
"status": "error",
|
||||
"error": str(e),
|
||||
"memories": [],
|
||||
"formatted_context": "Failed to recall memories.",
|
||||
}
|
||||
|
||||
|
||||
def format_shared_memories_for_context(
|
||||
memories: list[dict[str, Any]],
|
||||
created_by_map: dict[str, str] | None = None,
|
||||
) -> str:
|
||||
if not memories:
|
||||
return "No relevant team memories found."
|
||||
created_by_map = created_by_map or {}
|
||||
parts = ["<team_memories>"]
|
||||
for memory in memories:
|
||||
category = memory.get("category", "unknown")
|
||||
text = memory.get("memory_text", "")
|
||||
updated = memory.get("updated_at", "")
|
||||
created_by_id = memory.get("created_by_id")
|
||||
added_by = (
|
||||
created_by_map.get(str(created_by_id), "A team member")
|
||||
if created_by_id is not None
|
||||
else "A team member"
|
||||
)
|
||||
parts.append(
|
||||
f" <memory category='{category}' updated='{updated}' added_by='{added_by}'>{text}</memory>"
|
||||
)
|
||||
parts.append("</team_memories>")
|
||||
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
|
||||
|
|
@ -1,351 +0,0 @@
|
|||
"""
|
||||
User memory tools for the SurfSense agent.
|
||||
|
||||
This module provides tools for storing and retrieving user memories,
|
||||
enabling personalized AI responses similar to Claude's memory feature.
|
||||
|
||||
Features:
|
||||
- save_memory: Store facts, preferences, and context about the user
|
||||
- recall_memory: Retrieve relevant memories using semantic search
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
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
|
||||
|
||||
from app.db import MemoryCategory, UserMemory
|
||||
from app.utils.document_converters import embed_text
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Constants
|
||||
# =============================================================================
|
||||
|
||||
# Default number of memories to retrieve
|
||||
DEFAULT_RECALL_TOP_K = 5
|
||||
|
||||
# Maximum number of memories per user (to prevent unbounded growth)
|
||||
MAX_MEMORIES_PER_USER = 100
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Helper Functions
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def _to_uuid(user_id: str) -> UUID:
|
||||
"""Convert a string user_id to a UUID object."""
|
||||
if isinstance(user_id, UUID):
|
||||
return user_id
|
||||
return UUID(user_id)
|
||||
|
||||
|
||||
async def get_user_memory_count(
|
||||
db_session: AsyncSession,
|
||||
user_id: str,
|
||||
search_space_id: int | None = None,
|
||||
) -> int:
|
||||
"""Get the count of memories for a user."""
|
||||
uuid_user_id = _to_uuid(user_id)
|
||||
query = select(UserMemory).where(UserMemory.user_id == uuid_user_id)
|
||||
if search_space_id is not None:
|
||||
query = query.where(
|
||||
(UserMemory.search_space_id == search_space_id)
|
||||
| (UserMemory.search_space_id.is_(None))
|
||||
)
|
||||
result = await db_session.execute(query)
|
||||
return len(result.scalars().all())
|
||||
|
||||
|
||||
async def delete_oldest_memory(
|
||||
db_session: AsyncSession,
|
||||
user_id: str,
|
||||
search_space_id: int | None = None,
|
||||
) -> None:
|
||||
"""Delete the oldest memory for a user to make room for new ones."""
|
||||
uuid_user_id = _to_uuid(user_id)
|
||||
query = (
|
||||
select(UserMemory)
|
||||
.where(UserMemory.user_id == uuid_user_id)
|
||||
.order_by(UserMemory.updated_at.asc())
|
||||
.limit(1)
|
||||
)
|
||||
if search_space_id is not None:
|
||||
query = query.where(
|
||||
(UserMemory.search_space_id == search_space_id)
|
||||
| (UserMemory.search_space_id.is_(None))
|
||||
)
|
||||
result = await db_session.execute(query)
|
||||
oldest_memory = result.scalars().first()
|
||||
if oldest_memory:
|
||||
await db_session.delete(oldest_memory)
|
||||
await db_session.commit()
|
||||
|
||||
|
||||
def format_memories_for_context(memories: list[dict[str, Any]]) -> str:
|
||||
"""Format retrieved memories into a readable context string for the LLM."""
|
||||
if not memories:
|
||||
return "No relevant memories found for this user."
|
||||
|
||||
parts = ["<user_memories>"]
|
||||
for memory in memories:
|
||||
category = memory.get("category", "unknown")
|
||||
text = memory.get("memory_text", "")
|
||||
updated = memory.get("updated_at", "")
|
||||
parts.append(
|
||||
f" <memory category='{category}' updated='{updated}'>{text}</memory>"
|
||||
)
|
||||
parts.append("</user_memories>")
|
||||
|
||||
return "\n".join(parts)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Tool Factory Functions
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def create_save_memory_tool(
|
||||
user_id: str,
|
||||
search_space_id: int,
|
||||
db_session: AsyncSession,
|
||||
):
|
||||
"""
|
||||
Factory function to create the save_memory tool.
|
||||
|
||||
Args:
|
||||
user_id: The user's UUID
|
||||
search_space_id: The search space ID (for space-specific memories)
|
||||
db_session: Database session for executing queries
|
||||
|
||||
Returns:
|
||||
A configured tool function for saving user memories
|
||||
"""
|
||||
|
||||
@tool
|
||||
async def save_memory(
|
||||
content: str,
|
||||
category: str = "fact",
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Save a fact, preference, or context about the user for future reference.
|
||||
|
||||
Use this tool when:
|
||||
- User explicitly says "remember this", "keep this in mind", or similar
|
||||
- User shares personal preferences (e.g., "I prefer Python over JavaScript")
|
||||
- User shares important facts about themselves (name, role, interests, projects)
|
||||
- User gives standing instructions (e.g., "always respond in bullet points")
|
||||
- User shares relevant context (e.g., "I'm working on project X")
|
||||
|
||||
The saved information will be available in future conversations to provide
|
||||
more personalized and contextual responses.
|
||||
|
||||
Args:
|
||||
content: The fact/preference/context to remember.
|
||||
Phrase it clearly, e.g., "User prefers dark mode",
|
||||
"User is a senior Python developer", "User is working on an AI project"
|
||||
category: Type of memory. One of:
|
||||
- "preference": User preferences (e.g., coding style, tools, formats)
|
||||
- "fact": Facts about the user (e.g., name, role, expertise)
|
||||
- "instruction": Standing instructions (e.g., response format preferences)
|
||||
- "context": Current context (e.g., ongoing projects, goals)
|
||||
|
||||
Returns:
|
||||
A dictionary with the save status and memory details
|
||||
"""
|
||||
# Normalize and validate category (LLMs may send uppercase)
|
||||
category = category.lower() if category else "fact"
|
||||
valid_categories = ["preference", "fact", "instruction", "context"]
|
||||
if category not in valid_categories:
|
||||
category = "fact"
|
||||
|
||||
try:
|
||||
# Convert user_id to UUID
|
||||
uuid_user_id = _to_uuid(user_id)
|
||||
|
||||
# Check if we've hit the memory limit
|
||||
memory_count = await get_user_memory_count(
|
||||
db_session, user_id, search_space_id
|
||||
)
|
||||
if memory_count >= MAX_MEMORIES_PER_USER:
|
||||
# Delete oldest memory to make room
|
||||
await delete_oldest_memory(db_session, user_id, search_space_id)
|
||||
|
||||
embedding = await asyncio.to_thread(embed_text, content)
|
||||
|
||||
# Create new memory using ORM
|
||||
# The pgvector Vector column type handles embedding conversion automatically
|
||||
new_memory = UserMemory(
|
||||
user_id=uuid_user_id,
|
||||
search_space_id=search_space_id,
|
||||
memory_text=content,
|
||||
category=MemoryCategory(category), # Convert string to enum
|
||||
embedding=embedding, # Pass embedding directly (list or numpy array)
|
||||
)
|
||||
|
||||
db_session.add(new_memory)
|
||||
await db_session.commit()
|
||||
await db_session.refresh(new_memory)
|
||||
|
||||
return {
|
||||
"status": "saved",
|
||||
"memory_id": new_memory.id,
|
||||
"memory_text": content,
|
||||
"category": category,
|
||||
"message": f"I'll remember: {content}",
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f"Failed to save memory for user {user_id}: {e}")
|
||||
# Rollback the session to clear any failed transaction state
|
||||
await db_session.rollback()
|
||||
return {
|
||||
"status": "error",
|
||||
"error": str(e),
|
||||
"message": "Failed to save memory. Please try again.",
|
||||
}
|
||||
|
||||
return save_memory
|
||||
|
||||
|
||||
def create_recall_memory_tool(
|
||||
user_id: str,
|
||||
search_space_id: int,
|
||||
db_session: AsyncSession,
|
||||
):
|
||||
"""
|
||||
Factory function to create the recall_memory tool.
|
||||
|
||||
Args:
|
||||
user_id: The user's UUID
|
||||
search_space_id: The search space ID
|
||||
db_session: Database session for executing queries
|
||||
|
||||
Returns:
|
||||
A configured tool function for recalling user 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 memories about the user to provide personalized responses.
|
||||
|
||||
Use this tool when:
|
||||
- You need user context to give a better, more personalized answer
|
||||
- User asks about their preferences or past information they shared
|
||||
- User references something they told you before
|
||||
- Personalization would significantly improve the response quality
|
||||
- User asks "what do you know about me?" or similar
|
||||
|
||||
Args:
|
||||
query: Optional search query to find specific memories.
|
||||
If not provided, returns the most recent memories.
|
||||
Example: "programming preferences", "current projects"
|
||||
category: Optional category filter. One of:
|
||||
"preference", "fact", "instruction", "context"
|
||||
If not provided, searches all categories.
|
||||
top_k: Number of memories to retrieve (default: 5, max: 20)
|
||||
|
||||
Returns:
|
||||
A dictionary containing relevant memories and formatted context
|
||||
"""
|
||||
top_k = min(max(top_k, 1), 20) # Clamp between 1 and 20
|
||||
|
||||
try:
|
||||
# Convert user_id to UUID
|
||||
uuid_user_id = _to_uuid(user_id)
|
||||
|
||||
if query:
|
||||
query_embedding = await asyncio.to_thread(embed_text, query)
|
||||
|
||||
# Build query with vector similarity
|
||||
stmt = (
|
||||
select(UserMemory)
|
||||
.where(UserMemory.user_id == uuid_user_id)
|
||||
.where(
|
||||
(UserMemory.search_space_id == search_space_id)
|
||||
| (UserMemory.search_space_id.is_(None))
|
||||
)
|
||||
)
|
||||
|
||||
# Add category filter if specified
|
||||
if category and category in [
|
||||
"preference",
|
||||
"fact",
|
||||
"instruction",
|
||||
"context",
|
||||
]:
|
||||
stmt = stmt.where(UserMemory.category == MemoryCategory(category))
|
||||
|
||||
# Order by vector similarity
|
||||
stmt = stmt.order_by(
|
||||
UserMemory.embedding.op("<=>")(query_embedding)
|
||||
).limit(top_k)
|
||||
|
||||
else:
|
||||
# No query - return most recent memories
|
||||
stmt = (
|
||||
select(UserMemory)
|
||||
.where(UserMemory.user_id == uuid_user_id)
|
||||
.where(
|
||||
(UserMemory.search_space_id == search_space_id)
|
||||
| (UserMemory.search_space_id.is_(None))
|
||||
)
|
||||
)
|
||||
|
||||
# Add category filter if specified
|
||||
if category and category in [
|
||||
"preference",
|
||||
"fact",
|
||||
"instruction",
|
||||
"context",
|
||||
]:
|
||||
stmt = stmt.where(UserMemory.category == MemoryCategory(category))
|
||||
|
||||
stmt = stmt.order_by(UserMemory.updated_at.desc()).limit(top_k)
|
||||
|
||||
result = await db_session.execute(stmt)
|
||||
memories = result.scalars().all()
|
||||
|
||||
# Format memories for response
|
||||
memory_list = [
|
||||
{
|
||||
"id": m.id,
|
||||
"memory_text": m.memory_text,
|
||||
"category": m.category.value if m.category else "unknown",
|
||||
"updated_at": m.updated_at.isoformat() if m.updated_at else None,
|
||||
}
|
||||
for m in memories
|
||||
]
|
||||
|
||||
formatted_context = format_memories_for_context(memory_list)
|
||||
|
||||
return {
|
||||
"status": "success",
|
||||
"count": len(memory_list),
|
||||
"memories": memory_list,
|
||||
"formatted_context": formatted_context,
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f"Failed to recall memories for user {user_id}: {e}")
|
||||
await db_session.rollback()
|
||||
return {
|
||||
"status": "error",
|
||||
"error": str(e),
|
||||
"memories": [],
|
||||
"formatted_context": "Failed to recall memories.",
|
||||
}
|
||||
|
||||
return recall_memory
|
||||
|
|
@ -861,99 +861,6 @@ class ChatSessionState(BaseModel):
|
|||
ai_responding_to_user = relationship("User")
|
||||
|
||||
|
||||
class MemoryCategory(StrEnum):
|
||||
"""Categories for user memories."""
|
||||
|
||||
# Using lowercase keys to match PostgreSQL enum values
|
||||
preference = "preference" # User preferences (e.g., "prefers dark mode")
|
||||
fact = "fact" # Facts about the user (e.g., "is a Python developer")
|
||||
instruction = (
|
||||
"instruction" # Standing instructions (e.g., "always respond in bullet points")
|
||||
)
|
||||
context = "context" # Contextual information (e.g., "working on project X")
|
||||
|
||||
|
||||
class UserMemory(BaseModel, TimestampMixin):
|
||||
"""
|
||||
Private memory: facts, preferences, context per user per search space.
|
||||
Used only for private chats (not shared/team chats).
|
||||
"""
|
||||
|
||||
__tablename__ = "user_memories"
|
||||
|
||||
user_id = Column(
|
||||
UUID(as_uuid=True),
|
||||
ForeignKey("user.id", ondelete="CASCADE"),
|
||||
nullable=False,
|
||||
index=True,
|
||||
)
|
||||
# Optional association with a search space (if memory is space-specific)
|
||||
search_space_id = Column(
|
||||
Integer,
|
||||
ForeignKey("searchspaces.id", ondelete="CASCADE"),
|
||||
nullable=True,
|
||||
index=True,
|
||||
)
|
||||
|
||||
# The actual memory content
|
||||
memory_text = Column(Text, nullable=False)
|
||||
# Category for organization and filtering
|
||||
category = Column(
|
||||
SQLAlchemyEnum(MemoryCategory),
|
||||
nullable=False,
|
||||
default=MemoryCategory.fact,
|
||||
)
|
||||
# Vector embedding for semantic search
|
||||
embedding = Column(Vector(config.embedding_model_instance.dimension))
|
||||
|
||||
# Track when memory was last updated
|
||||
updated_at = Column(
|
||||
TIMESTAMP(timezone=True),
|
||||
nullable=False,
|
||||
default=lambda: datetime.now(UTC),
|
||||
onupdate=lambda: datetime.now(UTC),
|
||||
index=True,
|
||||
)
|
||||
|
||||
# Relationships
|
||||
user = relationship("User", back_populates="memories")
|
||||
search_space = relationship("SearchSpace", back_populates="user_memories")
|
||||
|
||||
|
||||
class SharedMemory(BaseModel, TimestampMixin):
|
||||
__tablename__ = "shared_memories"
|
||||
|
||||
search_space_id = Column(
|
||||
Integer,
|
||||
ForeignKey("searchspaces.id", ondelete="CASCADE"),
|
||||
nullable=False,
|
||||
index=True,
|
||||
)
|
||||
created_by_id = Column(
|
||||
UUID(as_uuid=True),
|
||||
ForeignKey("user.id", ondelete="CASCADE"),
|
||||
nullable=False,
|
||||
index=True,
|
||||
)
|
||||
memory_text = Column(Text, nullable=False)
|
||||
category = Column(
|
||||
SQLAlchemyEnum(MemoryCategory),
|
||||
nullable=False,
|
||||
default=MemoryCategory.fact,
|
||||
)
|
||||
embedding = Column(Vector(config.embedding_model_instance.dimension))
|
||||
updated_at = Column(
|
||||
TIMESTAMP(timezone=True),
|
||||
nullable=False,
|
||||
default=lambda: datetime.now(UTC),
|
||||
onupdate=lambda: datetime.now(UTC),
|
||||
index=True,
|
||||
)
|
||||
|
||||
search_space = relationship("SearchSpace", back_populates="shared_memories")
|
||||
created_by = relationship("User")
|
||||
|
||||
|
||||
class Folder(BaseModel, TimestampMixin):
|
||||
__tablename__ = "folders"
|
||||
|
||||
|
|
@ -1518,20 +1425,6 @@ class SearchSpace(BaseModel, TimestampMixin):
|
|||
cascade="all, delete-orphan",
|
||||
)
|
||||
|
||||
# User memories associated with this search space
|
||||
user_memories = relationship(
|
||||
"UserMemory",
|
||||
back_populates="search_space",
|
||||
order_by="UserMemory.updated_at.desc()",
|
||||
cascade="all, delete-orphan",
|
||||
)
|
||||
shared_memories = relationship(
|
||||
"SharedMemory",
|
||||
back_populates="search_space",
|
||||
order_by="SharedMemory.updated_at.desc()",
|
||||
cascade="all, delete-orphan",
|
||||
)
|
||||
|
||||
|
||||
class SearchSourceConnector(BaseModel, TimestampMixin):
|
||||
__tablename__ = "search_source_connectors"
|
||||
|
|
@ -2032,14 +1925,6 @@ if config.AUTH_TYPE == "GOOGLE":
|
|||
passive_deletes=True,
|
||||
)
|
||||
|
||||
# User memories for personalized AI responses
|
||||
memories = relationship(
|
||||
"UserMemory",
|
||||
back_populates="user",
|
||||
order_by="UserMemory.updated_at.desc()",
|
||||
cascade="all, delete-orphan",
|
||||
)
|
||||
|
||||
# Incentive tasks completed by this user
|
||||
incentive_tasks = relationship(
|
||||
"UserIncentiveTask",
|
||||
|
|
@ -2154,14 +2039,6 @@ else:
|
|||
passive_deletes=True,
|
||||
)
|
||||
|
||||
# User memories for personalized AI responses
|
||||
memories = relationship(
|
||||
"UserMemory",
|
||||
back_populates="user",
|
||||
order_by="UserMemory.updated_at.desc()",
|
||||
cascade="all, delete-orphan",
|
||||
)
|
||||
|
||||
# Incentive tasks completed by this user
|
||||
incentive_tasks = relationship(
|
||||
"UserIncentiveTask",
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue