mem0 pydantic ai toolset

mem0_toolset.py

import pydantic
from pydantic import TypeAdapter
from pydantic_ai import RunContext
from pydantic_ai.toolsets import FunctionToolset

from llm.agents.core.agent_types import Deps

MEMORY_MAX_RETRIES = 3


class MemoryRecord(pydantic.BaseModel):
    id: str = pydantic.Field(description="Unique identifier")
    memory: str = pydantic.Field(description="Memory content")


MemoryAdapter = TypeAdapter(list[MemoryRecord])


async def get_all_memory(ctx: RunContext[Deps]) -> str:
    """Retrieve all long-term memory records for the user.

    This tool fetches the complete set of memories stored for the user, which can include
    preferences, past conversations, important facts, or any information the LLM has deemed
    worth remembering. Use this when you need to recall all stored information about the user.

    Args:
        ctx: Execution context with dependencies for the agent.

    Returns:
        str: Serialized XML string containing all memory records with their IDs and values,
             or "No memories found." if the user has no stored memories.
    """
    memories = await ctx.deps.memory.get_all(user_id=str(ctx.deps.db_user.id))

    return MemoryAdapter.validate_python(memories.get("results", []))


async def create_long_term_memory_record(ctx: RunContext[Deps], value: str) -> None:
    """Create a new long-term memory record for the user.

    This tool stores important information about the user that should persist across conversations.
    Use this to remember user preferences, important facts, past interactions, or any information
    that would be valuable to recall in future conversations.

    Args:
        ctx: Execution context with dependencies for the agent.
        value: The content to store in the long-term memory record. This should be clear,
               concise information that captures something important about the user.
    """
    await ctx.deps.memory.add(
        [
            {
                "role": "user",
                "content": value,
            }
        ],
        user_id=str(ctx.deps.db_user.id),
        infer=False,
    )


async def delete_long_term_memory_record(ctx: RunContext[Deps], memory_id: str) -> bool:
    """Delete a specific long-term memory record.

    This tool permanently removes a memory record from the user's long-term memory. Use this
    when information is no longer relevant, incorrect, or when the user requests to forget
    something. Be careful as deleted memories cannot be recovered.

    Args:
        ctx: Execution context with dependencies for the agent.
        memory_id: The ID of the memory record to delete. You can get memory IDs from
                   get_all_memory() or read_long_term_memory_records().

    Returns:
        bool: True if the memory record was successfully deleted, False if the memory
              was not found or couldn't be deleted.
    """
    return await ctx.deps.memory.delete(memory_id=memory_id)


async def update_long_term_memory_record(ctx: RunContext[Deps], memory_id: str, value: str) -> dict[str, str]:
    """Update an existing long-term memory record.

    This tool modifies the content of an existing memory record while preserving its ID.
    Use this when user information has changed or when you need to correct/update
    previously stored information.

    Args:
        ctx: Execution context with dependencies for the agent.
        memory_id: The ID of the memory record to update. You can get memory IDs from
                   get_all_memory() or read_long_term_memory_records().
        value: The new content for the memory record. This replaces the existing content.

    Returns:
        dict[str, str]: Updated memory record data containing the new information.
    """
    return await ctx.deps.memory.update(memory_id=memory_id, data=value)


async def read_long_term_memory_records(
    ctx: RunContext[Deps],
    query: str,
) -> str:
    """Search for long-term memory records based on semantic similarity to a query.

    This tool performs a semantic search through the user's long-term memories, finding
    records that are relevant to the provided query. Use this when you need to recall
    specific information about the user based on meaning rather than exact keywords.

    Args:
        ctx: Execution context with dependencies for the agent.
        query: The search query to find relevant memories. This should be a natural language
               description of what you're looking for in the user's memory.

    Returns:
        str: Serialized XML string containing matching memory records with their IDs and values,
             or empty string if no relevant memories are found.
    """
    memories = await ctx.deps.memory.search(query, user_id=str(ctx.deps.db_user.id))
    return MemoryAdapter.validate_python(memories.get("results", []))


memory_toolset = FunctionToolset(
    tools=[
        create_long_term_memory_record,
        read_long_term_memory_records,
        update_long_term_memory_record,
        delete_long_term_memory_record,
        get_all_memory,
    ],
    require_parameter_descriptions=True,
    max_retries=2,
)