ingest.py

"""Build the knowledge base for the lab assistant.

Run this whenever ``knowledge/`` changes:

    py -3.12 ingest.py

Steps:
1. Read the markdown knowledge files and ``publications.jsonl``.
2. Enrich publications that have an arXiv id with the official abstract.
3. Chunk everything into retrieval passages.
4. Embed the passages (backend chosen by ``EMBED_BACKEND``).
5. Build a lightweight paper graph (edges = shared authors or shared topics).
6. Write ``data/chunks.json``, ``data/embeddings.npy``, ``data/graph.json``.
"""

from __future__ import annotations

import json
import re
import time
from typing import Dict, List

import numpy as np

import config
import embeddings


# --------------------------------------------------------------------------- #
# Loading + chunking
# --------------------------------------------------------------------------- #
def _split_markdown(text: str, source: str, title: str, url: str) -> List[dict]:
    """Split a markdown doc into chunks at '## ' headings, then by paragraph
    if a section is very long."""
    chunks: List[dict] = []
    sections = re.split(r"\n(?=## )", text)
    for sec in sections:
        sec = sec.strip()
        if not sec:
            continue
        heading = sec.splitlines()[0].lstrip("# ").strip()
        # keep sections whole unless they are long, then split on blank lines
        parts = [sec] if len(sec) < 900 else [p for p in sec.split("\n\n") if p.strip()]
        for part in parts:
            chunks.append(
                {
                    "doc_id": source,
                    "title": f"{title} — {heading}" if heading else title,
                    "source": source,
                    "url": url,
                    "topics": [],
                    "authors": [],
                    "text": part.strip(),
                }
            )
    return chunks


def _fetch_arxiv_abstracts(arxiv_ids: List[str]) -> Dict[str, str]:
    """Fetch abstracts from the arXiv API. Returns {arxiv_id: abstract}.
    Fails soft: on any network error returns whatever it has (possibly empty)."""
    ids = [a for a in arxiv_ids if a]
    if not ids:
        return {}
    # arXiv asks for a descriptive User-Agent and a short delay between requests;
    # the default requests UA is sometimes rate-limited (HTTP 429).
    headers = {"User-Agent": "labbot-ingest/1.0 (research assistant; mailto:noreply@example.com)"}
    url = "https://export.arxiv.org/api/query"
    xml = ""
    try:
        import requests

        for attempt in range(3):
            resp = requests.get(
                url,
                params={"id_list": ",".join(ids), "max_results": len(ids)},
                headers=headers,
                timeout=30,
            )
            if resp.status_code == 429:
                wait = 3 * (attempt + 1)
                print(f"[ingest] arXiv rate-limited (429); retrying in {wait}s ...")
                time.sleep(wait)
                continue
            resp.raise_for_status()
            xml = resp.text
            break
    except Exception as exc:
        print(f"[ingest] arXiv fetch failed ({exc}); using manual abstracts only.")
        return {}
    if not xml:
        print("[ingest] arXiv still rate-limited; using manual abstracts only.")
        return {}

    out: Dict[str, str] = {}
    # crude but dependency-free Atom parsing
    for entry in re.findall(r"<entry>(.*?)</entry>", xml, re.DOTALL):
        id_match = re.search(r"<id>(.*?)</id>", entry, re.DOTALL)
        sum_match = re.search(r"<summary>(.*?)</summary>", entry, re.DOTALL)
        if not (id_match and sum_match):
            continue
        raw_id = id_match.group(1).strip()
        # normalize ".../abs/2502.18470v1" -> "2502.18470"
        aid = re.sub(r"v\d+$", "", raw_id.rsplit("/", 1)[-1])
        summary = re.sub(r"\s+", " ", sum_match.group(1)).strip()
        out[aid] = summary
    return out


def _load_publications() -> List[dict]:
    pubs = []
    with open(config.PUBLICATIONS_FILE, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                pubs.append(json.loads(line))

    abstracts = _fetch_arxiv_abstracts([p.get("arxiv", "") for p in pubs])
    fetched = sum(1 for p in pubs if abstracts.get(p.get("arxiv", "")))
    print(f"[ingest] fetched {fetched}/{len(pubs)} abstracts from arXiv.")

    chunks = []
    for p in pubs:
        official = abstracts.get(p.get("arxiv", ""), "")
        manual = p.get("abstract", "")
        authors = ", ".join(p.get("authors", []))
        body = official or manual
        text = (
            f"{p['title']} ({p.get('venue', '')} {p.get('year', '')}). "
            f"Authors: {authors}. {body}"
        ).strip()
        chunks.append(
            {
                "doc_id": p["id"],
                "title": p["title"],
                "source": "publication",
                "url": p.get("url", config.HOMEPAGE_URL),
                "topics": p.get("topics", []),
                "authors": p.get("authors", []),
                "text": text,
            }
        )
    return chunks


def build_chunks() -> List[dict]:
    chunks: List[dict] = []
    md_files = [
        ("profile.md", "profile", "Profile of Prof. Liang Zhao", config.HOMEPAGE_URL),
        ("faq.md", "faq", "Prospective Students & Collaborators FAQ", config.HOMEPAGE_URL),
        ("tools.md", "tools", "Lab Tools & Systems", config.HOMEPAGE_URL),
    ]
    for fname, source, title, url in md_files:
        path = config.KNOWLEDGE_DIR / fname
        if path.exists():
            chunks.extend(_split_markdown(path.read_text(encoding="utf-8"), source, title, url))
    chunks.extend(_load_publications())

    # assign stable ids
    for i, c in enumerate(chunks):
        c["chunk_id"] = i
    return chunks


# --------------------------------------------------------------------------- #
# Paper graph (shared authors or shared topics) — the GRAG/CG-RAG flavour
# --------------------------------------------------------------------------- #
def build_graph(chunks: List[dict]) -> Dict[str, List[str]]:
    pubs = [c for c in chunks if c["source"] == "publication"]
    adj: Dict[str, set] = {c["doc_id"]: set() for c in pubs}
    for i, a in enumerate(pubs):
        a_authors = {x for x in a["authors"] if x != config.PROFESSOR_NAME}
        a_topics = set(a["topics"])
        for b in pubs[i + 1 :]:
            b_authors = {x for x in b["authors"] if x != config.PROFESSOR_NAME}
            shared_author = bool(a_authors & b_authors)
            shared_topics = len(a_topics & set(b["topics"]))
            if shared_author or shared_topics >= 2:
                adj[a["doc_id"]].add(b["doc_id"])
                adj[b["doc_id"]].add(a["doc_id"])
    return {k: sorted(v) for k, v in adj.items()}


# --------------------------------------------------------------------------- #
# Main
# --------------------------------------------------------------------------- #
def main() -> None:
    t0 = time.time()
    config.DATA_DIR.mkdir(parents=True, exist_ok=True)

    chunks = build_chunks()
    print(f"[ingest] built {len(chunks)} chunks from knowledge/.")

    print(f"[ingest] embedding with backend='{embeddings.active_backend()}' ...")
    vectors = embeddings.embed_texts([c["text"] for c in chunks])
    print(f"[ingest] embeddings shape = {vectors.shape}")

    graph = build_graph(chunks)
    n_edges = sum(len(v) for v in graph.values()) // 2
    print(f"[ingest] paper graph: {len(graph)} nodes, {n_edges} edges.")

    np.save(config.EMBEDDINGS_FILE, vectors)
    config.CHUNKS_FILE.write_text(json.dumps(chunks, ensure_ascii=False, indent=1), encoding="utf-8")
    config.GRAPH_FILE.write_text(json.dumps(graph, ensure_ascii=False, indent=1), encoding="utf-8")

    print(
        f"[ingest] wrote {config.CHUNKS_FILE.name}, {config.EMBEDDINGS_FILE.name}, "
        f"{config.GRAPH_FILE.name} in {time.time() - t0:.1f}s."
    )


if __name__ == "__main__":
    main()