api/mindmap_service.py

from datetime import datetime, timezone
from typing import Any, Dict, Optional
from loguru import logger
from open_notebook.database.repository import repo_query, repo_upsert
from open_notebook.graphs.mindmap import mindmap_graph


async def generate_mindmap(notebook_id: str) -> Dict[str, Any]:
    result = await mindmap_graph.ainvoke({"notebook_id": notebook_id})
    if result.get("error"):
        raise RuntimeError(result["error"])
    return {
        "graph_data": result["graph_data"],
        "cached": False,
        "generated_at": datetime.now(timezone.utc).isoformat(),
    }


async def get_cached_mindmap(notebook_id: str) -> Optional[Dict[str, Any]]:
    results = await repo_query(
        "SELECT * FROM mind_map WHERE notebook_id=$nid",
        {"nid": notebook_id},
    )
    if results:
        record = results[0]
        graph_data = record.get("graph_data")
        if graph_data:
            return {
                "graph_data": graph_data,
                "cached": True,
                "generated_at": str(record.get("updated_at", "")),
            }
    return None


async def invalidate_mindmap(notebook_id: str) -> None:
    await repo_upsert("mind_map", f"mindmap:{notebook_id}", {"graph_data": None})