leocrm/app/plugins/builtins/agent_memory/services.py

"""Agent Memory services — store and retrieve memories with semantic search."""

from __future__ import annotations

import uuid
from typing import Any

from sqlalchemy import select
from sqlalchemy import text as sql_text
from sqlalchemy.ext.asyncio import AsyncSession

from app.plugins.builtins.agent_memory.models import AgentMemory
from app.plugins.builtins.unified_search.contracts import generate_embedding


async def store_memory(
    db: AsyncSession,
    tenant_id: uuid.UUID,
    agent_id: uuid.UUID,
    content: str,
    memory_type: str = "fact",
    owner_id: uuid.UUID | None = None,
) -> dict[str, Any]:
    """Store a new memory with embedding generation.

    Args:
        db: Database session.
        tenant_id: Tenant UUID.
        agent_id: Agent UUID.
        content: Memory content text.
        memory_type: Type of memory (fact, context, pattern, instruction).
        owner_id: Optional owner user UUID.

    Returns:
        Dict with the created memory data.
    """
    # Generate embedding for semantic search
    embedding = await generate_embedding(content, db=db, tenant_id=tenant_id)

    memory = AgentMemory(
        tenant_id=tenant_id,
        agent_id=agent_id,
        memory_type=memory_type,
        content=content,
        owner_id=owner_id,
    )
    db.add(memory)
    await db.flush()
    await db.refresh(memory)

    # Store embedding if generated successfully
    if embedding:
        sql = sql_text(
            "UPDATE agent_memories SET embedding = cast(:emb AS vector) WHERE id = :mid"
        )
        await db.execute(sql, {"emb": str(embedding), "mid": memory.id})
        await db.flush()

    return {
        "id": str(memory.id),
        "agent_id": str(memory.agent_id),
        "memory_type": memory.memory_type,
        "content": memory.content,
        "owner_id": str(memory.owner_id) if memory.owner_id else None,
        "created_at": memory.created_at.isoformat() if memory.created_at else None,
    }


async def retrieve_relevant_memories(
    db: AsyncSession,
    tenant_id: uuid.UUID,
    agent_id: uuid.UUID,
    query: str,
    limit: int = 10,
    memory_type: str | None = None,
    min_score: float = 0.5,
) -> list[dict[str, Any]]:
    """Retrieve semantically relevant memories for an agent.

    Uses pgvector cosine similarity search to find memories whose
    embedding is closest to the query embedding.

    Args:
        db: Database session.
        tenant_id: Tenant UUID.
        agent_id: Agent UUID.
        query: Natural language query to match against memories.
        limit: Maximum number of results.
        memory_type: Optional filter by memory type.
        min_score: Minimum similarity score threshold (0.0 to 1.0).

    Returns:
        List of memory dicts with similarity scores, sorted by relevance.
    """
    # Generate embedding for the query
    query_embedding = await generate_embedding(query, db=db, tenant_id=tenant_id)
    if not query_embedding:
        # Fallback: return recent memories if embedding fails
        stmt = select(AgentMemory).where(
            AgentMemory.tenant_id == tenant_id,
            AgentMemory.agent_id == agent_id,
            AgentMemory.deleted_at.is_(None),
        )
        if memory_type:
            stmt = stmt.where(AgentMemory.memory_type == memory_type)
        stmt = stmt.order_by(AgentMemory.created_at.desc()).limit(limit)
        result = await db.execute(stmt)
        memories = result.scalars().all()
        return [
            {
                "id": str(m.id),
                "agent_id": str(m.agent_id),
                "memory_type": m.memory_type,
                "content": m.content,
                "score": 0.0,
                "created_at": m.created_at.isoformat() if m.created_at else None,
            }
            for m in memories
        ]

    # Vector similarity search
    type_filter = "AND m.memory_type = :mtype" if memory_type else ""
    sql = sql_text(f"""
        SELECT m.*, 1 - (m.embedding <=> cast(:emb AS vector)) AS score
        FROM agent_memories m
        WHERE m.tenant_id = :tid
          AND m.agent_id = :aid
          AND m.deleted_at IS NULL
          AND m.embedding IS NOT NULL
          {type_filter}
        ORDER BY m.embedding <=> cast(:emb AS vector)
        LIMIT :lim
    """)

    params: dict[str, Any] = {
        "emb": str(query_embedding),
        "tid": tenant_id,
        "aid": agent_id,
        "lim": limit,
    }
    if memory_type:
        params["mtype"] = memory_type

    result = await db.execute(sql, params)
    rows = result.mappings().all()

    return [
        {
            "id": str(row["id"]),
            "agent_id": str(row["agent_id"]),
            "memory_type": row["memory_type"],
            "content": row["content"],
            "score": float(row["score"]),
            "created_at": row["created_at"].isoformat() if row.get("created_at") else None,
        }
        for row in rows
        if float(row["score"]) >= min_score
    ]


async def delete_memory(
    db: AsyncSession,
    tenant_id: uuid.UUID,
    memory_id: uuid.UUID,
) -> bool:
    """Delete a memory by ID."""
    result = await db.execute(
        select(AgentMemory).where(
            AgentMemory.id == memory_id,
            AgentMemory.tenant_id == tenant_id,
        )
    )
    memory = result.scalar_one_or_none()
    if memory is None:
        return False
    await db.delete(memory)
    return True