File size: 4,435 Bytes
2e818da
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
"""GraphStateManager -> single source of truth for node state per session.

Numeric scoring/monotone clamping was removed -> apply_node_patch now only
handles status ("ongoing"/"completed"), description updates, and new children.
"""
import json
import os
from typing import Dict

from app.schemas.graph import NodeData, NodePatch

_GRAPH_DIR = os.path.expanduser("~/.studybuddy/graphs")

class GraphStateManager:
    def __init__(self) -> None:
        os.makedirs(_GRAPH_DIR, exist_ok=True)
        # { project_id: { node_id: NodeData } }
        self._graphs: Dict[str, Dict[str, NodeData]] = {}

    def _path(self, project_id: str) -> str:
        return os.path.join(_GRAPH_DIR, f"{project_id}.json")

    def _save(self, project_id: str) -> None:
        if project_id in self._graphs:
            with open(self._path(project_id), "w", encoding="utf-8") as f:
                json.dump([n.model_dump() for n in self._graphs[project_id].values()], f)

    def _load(self, project_id: str) -> bool:
        if project_id in self._graphs:
            return True
        p = self._path(project_id)
        if os.path.exists(p):
            with open(p, "r", encoding="utf-8") as f:
                nodes = [NodeData(**n) for n in json.load(f)]
            self._graphs[project_id] = {n.id: n for n in nodes}
            return True
        return False

    def add_node(self, project_id: str, node: NodeData) -> None:
        self._load(project_id)
        self._graphs.setdefault(project_id, {})[node.id] = node
        self._save(project_id)

    def set_graph(self, project_id: str, nodes: list[NodeData]) -> None:
        self._graphs[project_id] = {n.id: n for n in nodes}
        self._save(project_id)

    def get_node(self, project_id: str, node_id: str) -> NodeData:
        self._load(project_id)
        return self._graphs[project_id][node_id]

    def list_nodes(self, project_id: str) -> list[NodeData]:
        self._load(project_id)
        return list(self._graphs.get(project_id, {}).values())

    def apply_node_patch(self, project_id: str, patch: NodePatch) -> NodeData:
        self._load(project_id)
        node = self._graphs.setdefault(project_id, {}).get(patch.node_id)
        if node is None:
            node = NodeData(id=patch.node_id, label=patch.node_id)
            self._graphs[project_id][patch.node_id] = node

        if patch.status is not None:
            node.status = patch.status
        if patch.updated_label is not None:
            node.label = patch.updated_label
        if patch.updated_description is not None:
            node.description = patch.updated_description
        if patch.new_children:
            node.children_ids.extend(
                c for c in patch.new_children if c not in node.children_ids
            )


        self._save(project_id)
        return node

    def clear_session(self, project_id: str) -> None:
        self._graphs.pop(project_id, None)
        p = self._path(project_id)
        if os.path.exists(p):
            os.remove(p)

    # ------------------------------------------------------------------ #
    # Document-keyed cache -> reuse a built graph across sessions per PDF  #
    # ------------------------------------------------------------------ #

    def _doc_path(self, document_id: str) -> str:
        return os.path.join(_GRAPH_DIR, f"doc_{document_id}.json")

    def save_doc_graph(self, document_id: str, nodes: list[NodeData], edges: list) -> None:
        if not document_id:
            return
        with open(self._doc_path(document_id), "w", encoding="utf-8") as f:
            json.dump({"nodes": [n.model_dump() for n in nodes], "edges": edges}, f)

    def load_doc_graph(self, document_id: str):
        """Return (nodes, edges) for a previously built document, or None."""
        if not document_id:
            return None
        p = self._doc_path(document_id)
        if not os.path.exists(p):
            return None
        try:
            with open(p, "r", encoding="utf-8") as f:
                data = json.load(f)
            nodes = [NodeData(**n) for n in data.get("nodes", [])]
            return nodes, data.get("edges", [])
        except Exception:
            return None

    def clear_document_graph(self, document_id: str) -> None:
        if not document_id:
            return
        path = self._doc_path(document_id)
        if os.path.exists(path):
            os.remove(path)