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

"""Agent Memory plugin routes — CRUD for persistent agent memories."""

from __future__ import annotations

import uuid

from fastapi import APIRouter, Depends, HTTPException, Query, Response, status
from sqlalchemy import func, select
from sqlalchemy.ext.asyncio import AsyncSession

from app.core.db import get_db
from app.deps import get_current_user, require_permission
from app.plugins.builtins.agent_memory.models import AgentMemory
from app.plugins.builtins.agent_memory.schemas import (
    AgentMemoryCreate,
    AgentMemoryUpdate,
)
from app.plugins.builtins.agent_memory.services import (
    delete_memory,
    retrieve_relevant_memories,
    store_memory,
)

router = APIRouter(prefix="/api/v1/agent-memory", tags=["agent-memory"])


@router.post("", status_code=status.HTTP_201_CREATED, dependencies=[Depends(require_permission("agent_memory:write"))])
async def create_memory(
    body: AgentMemoryCreate,
    db: AsyncSession = Depends(get_db),
    current_user: dict = Depends(get_current_user),
):
    """Create a new agent memory with embedding."""
    tenant_id = uuid.UUID(current_user["tenant_id"])
    user_id = uuid.UUID(current_user["user_id"])
    agent_id = _parse_uuid(body.agent_id, "agent_id")

    result = await store_memory(
        db=db,
        tenant_id=tenant_id,
        agent_id=agent_id,
        content=body.content,
        memory_type=body.memory_type,
        owner_id=user_id,
    )
    return result


@router.get("", dependencies=[Depends(require_permission("agent_memory:read"))])
async def list_memories(
    agent_id: str = Query(..., description="Filter by agent UUID"),
    memory_type: str | None = Query(None, description="Filter by memory type"),
    page: int = Query(1, ge=1),
    page_size: int = Query(50, ge=1, le=200),
    db: AsyncSession = Depends(get_db),
    current_user: dict = Depends(get_current_user),
):
    """List memories for an agent, with optional type filter."""
    tenant_id = uuid.UUID(current_user["tenant_id"])
    aid = _parse_uuid(agent_id, "agent_id")

    base = select(AgentMemory).where(
        AgentMemory.tenant_id == tenant_id,
        AgentMemory.agent_id == aid,
        AgentMemory.deleted_at.is_(None),
    )
    if memory_type:
        base = base.where(AgentMemory.memory_type == memory_type)

    # Count total
    count_q = select(func.count()).select_from(base.subquery())
    total_result = await db.execute(count_q)
    total = total_result.scalar_one()

    # Paginated query
    offset = (page - 1) * page_size
    stmt = base.order_by(AgentMemory.created_at.desc()).offset(offset).limit(page_size)
    result = await db.execute(stmt)
    memories = result.scalars().all()

    return {
        "items": [
            {
                "id": str(m.id),
                "agent_id": str(m.agent_id),
                "memory_type": m.memory_type,
                "content": m.content,
                "owner_id": str(m.owner_id) if m.owner_id else None,
                "created_at": m.created_at.isoformat() if m.created_at else None,
                "updated_at": m.updated_at.isoformat() if m.updated_at else None,
                "score": 0.0,
            }
            for m in memories
        ],
        "total": total,
    }


@router.get("/search", dependencies=[Depends(require_permission("agent_memory:read"))])
async def search_memories(
    agent_id: str = Query(..., description="Agent UUID"),
    query: str = Query(..., min_length=1, description="Natural language query"),
    memory_type: str | None = Query(None, description="Optional type filter"),
    limit: int = Query(10, ge=1, le=100),
    min_score: float = Query(0.5, ge=0.0, le=1.0),
    db: AsyncSession = Depends(get_db),
    current_user: dict = Depends(get_current_user),
):
    """Semantic search over agent memories using pgvector."""
    tenant_id = uuid.UUID(current_user["tenant_id"])
    aid = _parse_uuid(agent_id, "agent_id")

    results = await retrieve_relevant_memories(
        db=db,
        tenant_id=tenant_id,
        agent_id=aid,
        query=query,
        limit=limit,
        memory_type=memory_type,
        min_score=min_score,
    )
    return {"items": results, "total": len(results)}


@router.patch("/{memory_id}", dependencies=[Depends(require_permission("agent_memory:write"))])
async def update_memory(
    memory_id: str,
    body: AgentMemoryUpdate,
    db: AsyncSession = Depends(get_db),
    current_user: dict = Depends(get_current_user),
):
    """Update a memory's content or type."""
    tenant_id = uuid.UUID(current_user["tenant_id"])
    mid = _parse_uuid(memory_id, "memory_id")

    result = await db.execute(
        select(AgentMemory).where(
            AgentMemory.id == mid,
            AgentMemory.tenant_id == tenant_id,
        )
    )
    memory = result.scalar_one_or_none()
    if memory is None:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail={"detail": "Memory not found", "code": "not_found"},
        )

    data = body.model_dump(exclude_unset=True)
    if "content" in data:
        memory.content = data["content"]
        # Regenerate embedding for updated content
        from app.plugins.builtins.unified_search.contracts import generate_embedding
        embedding = await generate_embedding(data["content"], db=db, tenant_id=tenant_id)
        if embedding:
            from sqlalchemy import text as sql_text
            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})
    if "memory_type" in data:
        memory.memory_type = data["memory_type"]

    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,
        "updated_at": memory.updated_at.isoformat() if memory.updated_at else None,
    }


@router.delete("/{memory_id}", status_code=status.HTTP_204_NO_CONTENT, dependencies=[Depends(require_permission("agent_memory:write"))])
async def delete_memory_route(
    memory_id: str,
    db: AsyncSession = Depends(get_db),
    current_user: dict = Depends(get_current_user),
):
    """Delete a memory by ID."""
    tenant_id = uuid.UUID(current_user["tenant_id"])
    mid = _parse_uuid(memory_id, "memory_id")

    success = await delete_memory(db, tenant_id, mid)
    if not success:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail={"detail": "Memory not found", "code": "not_found"},
        )
    return Response(status_code=status.HTTP_204_NO_CONTENT)


def _parse_uuid(val: str, field: str) -> uuid.UUID:
    try:
        return uuid.UUID(val)
    except (ValueError, TypeError):
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail={"detail": f"Invalid {field}", "code": "invalid_id"},
        ) from None