studyhub/core/citations.py

from __future__ import annotations

import re
from dataclasses import dataclass

from ..ingest.models import Chunk

_CIT = re.compile(r"\[(S\d+)\]")


@dataclass
class Citation:
    sid: str
    label: str
    file: str
    kind: str
    locator: str
    is_notes: bool


def extract_citations(answer_text: str, id_map: dict[str, Chunk]) -> list[Citation]:
    out: list[Citation] = []
    seen: set[str] = set()
    for m in _CIT.finditer(answer_text):
        sid = m.group(1)
        if sid in id_map and sid not in seen:
            seen.add(sid)
            a = id_map[sid].anchor
            out.append(Citation(sid=sid, label=a.label, file=a.file, kind=a.kind,
                                locator=a.locator, is_notes=a.is_notes))
    return out


def render_sources(citations: list[Citation]) -> str:
    return "\n".join(f"[{c.sid}] 📓 {c.label}" for c in citations)