mem0-python-server/mem0core/routes.py

from fastapi import APIRouter, Request
from fastapi.responses import HTMLResponse
from mem0 import Memory

from .config import COLLECTION_CONVERSATIONAL, COLLECTION_KNOWLEDGE, RERANKER_URL
from .handlers import extract_user_id, handle_add, handle_recent, handle_search
from .prompts import PROMPTS
from .reranker import rerank_results
from .responses import SafeJSONResponse
from .storage import chroma_get_all, sqlite_delete_ids


def build_router(memory_conv: Memory, memory_know: Memory) -> APIRouter:
    """Build and return all API routes bound to the provided memory instances."""
    router = APIRouter()

    with open("dashboard.html", "r", encoding="utf-8") as handle:
        dashboard_html = handle.read()

    @router.get("/dashboard", response_class=HTMLResponse, summary="Render dashboard", tags=["dashboard"])
    async def dashboard():
        """Serve the local dashboard HTML used for memory inspection and admin actions."""
        return HTMLResponse(content=dashboard_html)

    @router.get("/health", summary="Service health and prompt preview", tags=["health"])
    async def health():
        """Return runtime health plus prompt snippets for quick configuration verification."""
        return SafeJSONResponse(
            content={
                "status": "ok",
                "reranker_url": RERANKER_URL,
                "collections": {
                    "conversational": COLLECTION_CONVERSATIONAL,
                    "knowledge": COLLECTION_KNOWLEDGE,
                },
                "prompts": {
                    key: {p_key: text[:80] + "…" for p_key, text in prompt_set.items()}
                    for key, prompt_set in PROMPTS.items()
                },
            }
        )

    @router.post("/memories", summary="Add conversational memory", tags=["memories"])
    async def add_memory(req: Request):
        """Store conversational memory with LLM extraction and deduplication enabled."""
        return await handle_add(req, memory_conv, verbatim_allowed=False)

    @router.post("/memories/search", summary="Search conversational memory", tags=["memories"])
    async def search_memories(req: Request):
        """Search conversational memory and rerank candidates by relevance."""
        return await handle_search(req, memory_conv)

    @router.post("/memories/recent", summary="Recent conversational memory", tags=["memories"])
    async def recent_memories(req: Request):
        """Return newest conversational memories ordered by creation time."""
        return await handle_recent(req, memory_conv)

    @router.delete("/memories", summary="Delete conversational memory by filter", tags=["memories"])
    async def delete_memory(req: Request):
        """Delete conversational memories using a mem0 filter object from the request body."""
        data = await req.json()
        return SafeJSONResponse(content=memory_conv.delete(data.get("filter", {})))

    @router.post("/memories/all", summary="List all conversational memories", tags=["memories"])
    async def memories_all(req: Request):
        """Fetch full conversational history for a user by paging Chroma directly."""
        data = await req.json()
        user_id = extract_user_id(data) or "main"

        rows = chroma_get_all(
            memory_conv.vector_store.collection,
            user_id,
            include=["metadatas", "documents"],
        )
        items = []
        for row in rows:
            meta = row.get("metadata") or {}
            items.append(
                {
                    "id": row["id"],
                    "memory": row.get("document") or meta.get("data", ""),
                    "created_at": meta.get("created_at"),
                    "metadata": meta,
                    "user_id": user_id,
                }
            )

        items.sort(key=lambda r: r.get("created_at") or "", reverse=True)
        print(f"[memories/all] user={user_id} total={len(items)}")
        return SafeJSONResponse(content={"results": items})

    @router.post("/knowledge", summary="Add knowledge chunk", tags=["knowledge"])
    async def add_knowledge(req: Request):
        """Store knowledge verbatim without LLM extraction (ingestor already summarizes)."""
        return await handle_add(req, memory_know, verbatim_allowed=True)

    @router.post("/knowledge/search", summary="Search knowledge base", tags=["knowledge"])
    async def search_knowledge(req: Request):
        """Search knowledge memories and rerank candidates."""
        return await handle_search(req, memory_know)

    @router.post("/knowledge/recent", summary="Recent knowledge entries", tags=["knowledge"])
    async def recent_knowledge(req: Request):
        """Return newest knowledge entries for a user."""
        return await handle_recent(req, memory_know)

    @router.delete("/knowledge", summary="Delete knowledge by filter", tags=["knowledge"])
    async def delete_knowledge(req: Request):
        """Delete knowledge entries using a mem0 filter object from the request body."""
        data = await req.json()
        return SafeJSONResponse(content=memory_know.delete(data.get("filter", {})))

    @router.post("/knowledge/sources", summary="Knowledge source counts", tags=["knowledge"])
    async def knowledge_sources(req: Request):
        """List distinct source_file values with counts, bypassing mem0 get_all caps."""
        data = await req.json()
        user_id = extract_user_id(data) or "knowledge_base"

        rows = chroma_get_all(memory_know.vector_store.collection, user_id)
        counts = {}
        for row in rows:
            src = (row.get("metadata") or {}).get("source_file", "(no source)")
            counts[src] = counts.get(src, 0) + 1

        sources = [{"source_file": k, "count": v} for k, v in sorted(counts.items(), key=lambda x: -x[1])]
        print(f"[sources] user={user_id} total={len(rows)} books={len(sources)}")
        return SafeJSONResponse(content={"sources": sources, "total": len(rows)})

    @router.delete("/knowledge/by-source", summary="Delete all entries for one source file", tags=["knowledge"])
    async def delete_knowledge_by_source(req: Request):
        """Delete all knowledge entries for a source_file from Chroma and SQLite."""
        data = await req.json()
        source_file = data.get("source_file")
        user_id = extract_user_id(data) or "knowledge_base"

        if not source_file:
            return SafeJSONResponse(content={"error": "Missing source_file"}, status_code=400)

        rows = chroma_get_all(memory_know.vector_store.collection, user_id)
        to_delete = [
            row["id"]
            for row in rows
            if (row.get("metadata") or {}).get("source_file") == source_file
        ]

        if not to_delete:
            return SafeJSONResponse(content={"deleted": 0, "message": "no entries found for that source"})

        try:
            memory_know.vector_store.collection.delete(ids=to_delete)
        except Exception as exc:
            return SafeJSONResponse(content={"error": f"chroma delete failed: {exc}"}, status_code=500)

        sqlite_deleted = sqlite_delete_ids(to_delete)
        print(f"[delete by-source] source={source_file} chroma={len(to_delete)} sqlite={sqlite_deleted}")

        return SafeJSONResponse(
            content={
                "deleted": len(to_delete),
                "sqlite_deleted": sqlite_deleted,
                "source_file": source_file,
            }
        )

    @router.delete("/memory/{memory_id}", summary="Delete one memory by ID", tags=["memories", "knowledge"])
    async def delete_single_memory(memory_id: str, req: Request):
        """Delete a single memory from selected collection and mirror deletion in SQLite."""
        data = await req.json()
        collection = data.get("collection", "knowledge")
        mem = memory_know if collection == "knowledge" else memory_conv

        try:
            mem.vector_store.collection.delete(ids=[memory_id])
        except Exception as exc:
            return SafeJSONResponse(content={"error": f"chroma delete failed: {exc}"}, status_code=500)

        sqlite_delete_ids([memory_id])
        print(f"[delete single] id={memory_id} collection={collection}")
        return SafeJSONResponse(content={"deleted": memory_id})

    @router.post("/search", summary="Search both collections", tags=["search"])
    async def search_all(req: Request):
        """Search conversational + knowledge collections, then rerank merged results."""
        data = await req.json()
        query = (data.get("query") or "").strip()
        user_id = extract_user_id(data)
        limit = int(data.get("limit", 5))

        if not query:
            return SafeJSONResponse(content={"results": []})

        fetch_k = max(limit * 3, 15)

        def fetch(mem: Memory, tag: str):
            try:
                response = mem.search(query, user_id=user_id, limit=fetch_k)
                items = response.get("results", [])
            except Exception:
                items = []
            for item in items:
                item["_source"] = tag
            return items

        conv_items = fetch(memory_conv, "conversational")
        know_items = fetch(memory_know, "knowledge")
        merged = rerank_results(query, conv_items + know_items, top_k=limit)

        print(
            f"[search/all] user={user_id} query={query!r} "
            f"conv={len(conv_items)} know={len(know_items)} merged={len(merged)}"
        )
        return SafeJSONResponse(content={"results": merged})

    return router