studyhub/graph/models.py

from __future__ import annotations

from dataclasses import asdict, dataclass, field


@dataclass
class Module:
    """One course module = one syllabus (taught by one lecturer)."""
    id: str
    title: str
    lecturer: str
    hours: str
    objective: str
    language: str
    source_file: str                       # the syllabus path
    topic_ids: list[str] = field(default_factory=list)
    source_ids: list[str] = field(default_factory=list)   # content files in this module


@dataclass
class Topic:
    """A 'Main topics' bullet under a module; gathers passages across sources."""
    id: str
    module_id: str
    title: str
    chunk_ids: list[str] = field(default_factory=list)


@dataclass
class Graph:
    modules: list[Module] = field(default_factory=list)
    topics: list[Topic] = field(default_factory=list)

    def module_by_id(self, module_id: str):
        return next((m for m in self.modules if m.id == module_id), None)

    def topic_by_id(self, topic_id: str):
        return next((t for t in self.topics if t.id == topic_id), None)

    def topics_for_module(self, module_id: str) -> list[Topic]:
        return [t for t in self.topics if t.module_id == module_id]


def graph_to_dict(g: Graph) -> dict:
    return {"modules": [asdict(m) for m in g.modules], "topics": [asdict(t) for t in g.topics]}


def graph_from_dict(d: dict) -> Graph:
    return Graph(
        modules=[Module(**m) for m in d.get("modules", [])],
        topics=[Topic(**t) for t in d.get("topics", [])],
    )