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()