study-buddy / app /services /graph_state.py
GitHub Actions
deploy d092bea3608b7a29952f16357fda39b7a29e399b
2e818da
Raw
History Blame Contribute Delete
4.44 kB
"""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)