open_notebook/graphs/mindmap.py

import json
from typing import TypedDict, Optional, List, Dict, Any
from langgraph.graph import StateGraph, END
from loguru import logger
from open_notebook.ai.provision import provision_langchain_model
from open_notebook.database.repository import repo_query, repo_upsert, ensure_record_id
from open_notebook.domain.notebook import Notebook, Source


class MindMapState(TypedDict):
    notebook_id: str
    raw_content: str
    concepts: List[Dict[str, Any]]
    relationships: List[Dict[str, Any]]
    graph_data: Optional[Dict[str, Any]]
    error: Optional[str]


CONCEPT_EXTRACTION_PROMPT = """
You are an expert knowledge mapper. Analyze the following research content and extract
the most important concepts, entities, themes, and ideas.

CONTENT:
{content}

Return ONLY valid JSON in this exact schema (no markdown, no explanation):
{{
  "concepts": [
    {{
      "id": "unique_snake_case_id",
      "label": "Human Readable Label",
      "category": "one of: core_concept | entity | theme | finding | method | tool | place | person",
      "importance": 1-10,
      "description": "one sentence description"
    }}
  ]
}}

Rules:
- Extract 25 to 40 concepts maximum
- Core concepts (importance 8-10) should be the central nodes
- Avoid duplicates — merge similar concepts
- Keep labels concise (2-4 words max)
"""

RELATIONSHIP_EXTRACTION_PROMPT = """
Given these concepts from a research notebook, identify meaningful relationships between them.

CONCEPTS:
{concepts_json}

Return ONLY valid JSON (no markdown, no explanation):
{{
  "relationships": [
    {{
      "source": "concept_id_1",
      "target": "concept_id_2",
      "label": "relationship verb (2-3 words)",
      "strength": 1-5
    }}
  ]
}}

Rules:
- Only create relationships that are explicitly supported by the content
- Strength 5 = very strong direct relationship, 1 = weak tangential connection
- Aim for 30-60 relationships
- Every concept should have at least one connection
"""


async def load_content_node(state: MindMapState) -> MindMapState:
    try:
        notebook = await Notebook.get(state["notebook_id"])
        sources = await notebook.get_sources()
        notes = await notebook.get_notes()

        content_parts = []
        for source in sources:
            if source.content:
                content_parts.append(f"[SOURCE: {source.name}]\n{str(source.content)[:8000]}")
        for note in notes:
            if note.content:
                content_parts.append(f"[NOTE: {note.title}]\n{str(note.content)[:2000]}")

        state["raw_content"] = "\n\n---\n\n".join(content_parts)
        if not state["raw_content"].strip():
            state["error"] = "No content found in this notebook to generate a mind map."
        return state
    except Exception as e:
        logger.error(f"MindMap load_content_node error: {e}")
        state["error"] = str(e)
        return state


async def extract_concepts_node(state: MindMapState) -> MindMapState:
    if state.get("error"):
        return state
    try:
        content = state["raw_content"][:40000]
        model = await provision_langchain_model(content, None, "extraction")
        prompt = CONCEPT_EXTRACTION_PROMPT.format(content=content)
        response = await model.ainvoke(prompt)
        text = response.content if hasattr(response, "content") else str(response)
        text = text.strip().lstrip("```json").rstrip("```").strip()
        data = json.loads(text)
        state["concepts"] = data["concepts"]
        return state
    except Exception as e:
        logger.error(f"MindMap extract_concepts_node error: {e}")
        state["error"] = str(e)
        return state


async def build_relationships_node(state: MindMapState) -> MindMapState:
    if state.get("error"):
        return state
    try:
        concepts_json = json.dumps(
            [{"id": c["id"], "label": c["label"], "category": c["category"]}
             for c in state["concepts"]], indent=2
        )
        prompt = RELATIONSHIP_EXTRACTION_PROMPT.format(concepts_json=concepts_json)
        model = await provision_langchain_model(prompt, None, "extraction")
        response = await model.ainvoke(prompt)
        text = response.content if hasattr(response, "content") else str(response)
        text = text.strip().lstrip("```json").rstrip("```").strip()
        data = json.loads(text)
        state["relationships"] = data["relationships"]
        return state
    except Exception as e:
        logger.error(f"MindMap build_relationships_node error: {e}")
        state["error"] = str(e)
        return state


async def structure_graph_node(state: MindMapState) -> MindMapState:
    if state.get("error"):
        return state
    try:
        category_colors = {
            "core_concept": "#6366f1",
            "entity": "#0ea5e9",
            "theme": "#10b981",
            "finding": "#f59e0b",
            "method": "#8b5cf6",
            "tool": "#06b6d4",
            "place": "#14b8a6",
            "person": "#f97316",
        }
        nodes = [
            {
                "id": c["id"],
                "label": c["label"],
                "category": c["category"],
                "description": c.get("description", ""),
                "importance": c.get("importance", 5),
                "color": category_colors.get(c["category"], "#94a3b8"),
                "size": max(20, c.get("importance", 5) * 8),
            }
            for c in state["concepts"]
        ]
        valid_ids = {c["id"] for c in state["concepts"]}
        links = [
            {
                "source": r["source"],
                "target": r["target"],
                "label": r.get("label", ""),
                "strength": r.get("strength", 2),
            }
            for r in state["relationships"]
            if r["source"] in valid_ids and r["target"] in valid_ids
        ]
        state["graph_data"] = {
            "nodes": nodes,
            "links": links,
            "metadata": {
                "node_count": len(nodes),
                "link_count": len(links),
                "categories": list(category_colors.keys()),
            }
        }
        return state
    except Exception as e:
        logger.error(f"MindMap structure_graph_node error: {e}")
        state["error"] = str(e)
        return state


async def save_result_node(state: MindMapState) -> MindMapState:
    if state.get("error") or not state.get("graph_data"):
        return state
    try:
        await repo_upsert("mind_map", f"mindmap:{state['notebook_id']}", {
            "notebook_id": state["notebook_id"],
            "graph_data": state["graph_data"],
            "concept_count": len(state["concepts"]),
            "relationship_count": len(state["relationships"]),
        })
        return state
    except Exception as e:
        logger.error(f"MindMap save_result_node error: {e}")
        state["error"] = str(e)
        return state


workflow = StateGraph(MindMapState)
workflow.add_node("load_content", load_content_node)
workflow.add_node("extract_concepts", extract_concepts_node)
workflow.add_node("build_relationships", build_relationships_node)
workflow.add_node("structure_graph", structure_graph_node)
workflow.add_node("save_result", save_result_node)

workflow.set_entry_point("load_content")
workflow.add_edge("load_content", "extract_concepts")
workflow.add_edge("extract_concepts", "build_relationships")
workflow.add_edge("build_relationships", "structure_graph")
workflow.add_edge("structure_graph", "save_result")
workflow.add_edge("save_result", END)

mindmap_graph = workflow.compile()