app/ai/tools/user_memory.py

# app/ai/tools/user_memory.py
"""
User Memory Service — stores and retrieves per-user memories in MongoDB.

Collection: user_memories
Schema: {user_id, key, value, category, created_at, updated_at}
Categories: "preference", "search_pattern", "feedback", "context"
"""

from datetime import datetime, timezone
from typing import Any, Dict, List, Optional

from structlog import get_logger

logger = get_logger(__name__)

_COLLECTION = "user_memories"


async def save_memory(
    user_id: str,
    key: str,
    value: Any,
    category: str = "preference",
) -> bool:
    """
    Upsert a memory for a user.
    If the key already exists, update it; otherwise create it.
    """
    from app.database import get_db

    try:
        db = await get_db()
        now = datetime.now(timezone.utc)

        result = await db[_COLLECTION].update_one(
            {"user_id": user_id, "key": key},
            {
                "$set": {
                    "value": value,
                    "category": category,
                    "updated_at": now,
                },
                "$setOnInsert": {
                    "user_id": user_id,
                    "key": key,
                    "created_at": now,
                },
            },
            upsert=True,
        )
        logger.info(
            "Memory saved",
            user_id=user_id,
            key=key,
            category=category,
            upserted=result.upserted_id is not None,
        )
        return True
    except Exception as exc:
        logger.error("Failed to save memory", user_id=user_id, key=key, error=str(exc))
        return False


async def get_memories(
    user_id: str,
    category: Optional[str] = None,
) -> List[Dict]:
    """
    Retrieve all memories for a user, optionally filtered by category.
    """
    from app.database import get_db

    try:
        db = await get_db()
        query: Dict[str, Any] = {"user_id": user_id}
        if category:
            query["category"] = category

        cursor = db[_COLLECTION].find(query).sort("updated_at", -1)
        docs = await cursor.to_list(length=100)

        # Convert ObjectId to string for serialization
        for doc in docs:
            doc["_id"] = str(doc["_id"])

        return docs
    except Exception as exc:
        logger.error("Failed to get memories", user_id=user_id, error=str(exc))
        return []


async def get_memory(user_id: str, key: str) -> Optional[Dict]:
    """Retrieve a single memory by key."""
    from app.database import get_db

    try:
        db = await get_db()
        doc = await db[_COLLECTION].find_one({"user_id": user_id, "key": key})
        if doc:
            doc["_id"] = str(doc["_id"])
        return doc
    except Exception as exc:
        logger.error("Failed to get memory", user_id=user_id, key=key, error=str(exc))
        return None


async def delete_memory(user_id: str, key: str) -> bool:
    """Delete a specific memory."""
    from app.database import get_db

    try:
        db = await get_db()
        result = await db[_COLLECTION].delete_one({"user_id": user_id, "key": key})
        deleted = result.deleted_count > 0
        logger.info("Memory deleted", user_id=user_id, key=key, deleted=deleted)
        return deleted
    except Exception as exc:
        logger.error("Failed to delete memory", user_id=user_id, key=key, error=str(exc))
        return False


async def get_user_context(user_id: str) -> str:
    """
    Build a formatted string of user memories for injection into the brain prompt.
    Returns empty string if no memories exist.
    """
    memories = await get_memories(user_id)
    if not memories:
        return ""

    lines = []
    by_category: Dict[str, List[Dict]] = {}
    for mem in memories:
        cat = mem.get("category", "other")
        by_category.setdefault(cat, []).append(mem)

    for cat, items in by_category.items():
        lines.append(f"[{cat.upper()}]")
        for item in items[:15]:  # Cap per category to save tokens
            lines.append(f"  - {item['key']}: {item['value']}")

    return "\n".join(lines)