studyhub/graph/extract.py

from __future__ import annotations

import json
import re
from dataclasses import dataclass


@dataclass
class ModuleSpec:
    title: str
    lecturer: str
    hours: str
    objective: str
    language: str
    topics: list[str]


_SYSTEM = (
    "Extract the course module from this syllabus. Return ONLY a JSON object with keys: "
    "title, lecturer, hours, objective, language, and topics (a list of the 'Main topics' "
    "bullets, each a short string). No prose, no code fences."
)


def _parse_json(raw: str) -> dict:
    raw = re.sub(r"```(?:json)?|```", "", raw or "").strip()
    m = re.search(r"\{.*\}", raw, flags=re.DOTALL)
    if not m:
        return {}
    try:
        return json.loads(m.group(0))
    except Exception:
        return {}


def _heuristic_topics(text: str) -> list[str]:
    low = text.lower()
    i = low.find("main topics")
    if i < 0:
        return []
    out: list[str] = []
    for line in text[i:i + 1800].splitlines()[1:]:
        item = re.sub(r"^[\s•\-\*•●▪]+", "", line).strip()
        low_item = item.lower()
        if low_item.startswith(("prerequisit", "assessment", "text", "notes", "lecturer", "n. of", "objective")):
            break
        if 3 <= len(item) <= 180:
            out.append(item)
    return out[:12]


def extract_module(text: str, llm) -> ModuleSpec:
    """LLM-extract the module + topics; fall back to a heuristic 'Main topics' scan."""
    try:
        data = _parse_json(llm.complete(_SYSTEM, text[:6000]))
    except Exception:
        data = {}
    topics = data.get("topics") or _heuristic_topics(text)
    return ModuleSpec(
        title=str(data.get("title") or "").strip(),
        lecturer=str(data.get("lecturer") or "").strip(),
        hours=str(data.get("hours") or "").strip(),
        objective=str(data.get("objective") or "").strip(),
        language=str(data.get("language") or "").strip(),
        topics=[str(t).strip() for t in topics if str(t).strip()],
    )