studyhub/web/ingest_worker.py

"""In-process ingest worker: a user upload becomes searchable within minutes.

A daemon thread takes uploads off a queue and runs the offline pipeline's own
machinery per file (parser dispatch -> chunker -> embedder -> fresh FAISS), then
atomically swaps the live RetrievalIndex, links the file into the topic graph,
persists the pipeline-shaped artifact set to index_dir and (optionally) pushes
it to the dataset's index/. Manifest entries are format-identical to the
pipeline's, so the weekly Action rebuild treats worker-ingested files as
unchanged and re-embeds nothing. Hub failures only log: local state stays
authoritative and the Action is the consistency backstop.

Import-light: heavy deps (faiss, yake, PyMuPDF/pptx/docx) stay lazy inside the
pipeline helpers; nothing here pulls torch or FlagEmbedding at import time.
"""
from __future__ import annotations

import logging
import math
import queue
import threading
from pathlib import Path

import numpy as np

from ..core.loader import RetrievalIndex
from ..graph.build import _NAME_WEIGHT, _name_scores, _tokens
from ..graph.models import Module, Topic
from ..graph.store import save_graph
from ..ingest.chunk import chunk_blocks
from ..ingest.concepts import extract_concepts
from ..ingest.index import index_from_vectors, save_faiss
from ..ingest.models import Chunk
from ..ingest.pipeline import PARSERS
from ..ingest.store import file_hash, read_json, write_chunks, write_json

log = logging.getLogger(__name__)

_ACTIVE = ("pending", "processing", "indexed")


class IngestWorker:
    """Background ingestion of user uploads into the live service and index_dir.

    Status lifecycle per path: "pending" -> "processing" -> "indexed" | "failed: <reason>";
    "unknown" for paths never seen. Paths are store paths ("uploads/<user>/<file>"); on
    disk and in citations the file lives at "_uploads/<user>/<file>" under materials_dir,
    where /api/download already resolves it.
    """

    def __init__(self, service, *, index_dir, materials_dir, dataset_repo: str = "", hf_token: str = ""):
        self.service = service
        self.index_dir = Path(index_dir)
        self.materials_dir = Path(materials_dir)
        self.dataset_repo = dataset_repo
        self.hf_token = hf_token
        self._queue: queue.Queue[tuple[str, bytes] | None] = queue.Queue()
        self._status: dict[str, str] = {}
        self._lock = threading.Lock()
        self._stop = threading.Event()
        self._thread: threading.Thread | None = None
        self._module_routing: tuple | None = None        # lazy: (module_vecs, identities, idf)
        self._topic_vecs: dict[str, np.ndarray] = {}     # module_id -> topic embedding matrix

    # -- public API (request handlers; must never block) --

    def enqueue(self, path_in_repo: str, data: bytes) -> None:
        with self._lock:
            self._status[path_in_repo] = "pending"
        self._queue.put((path_in_repo, bytes(data)))

    def status(self, path_in_repo: str) -> str:
        with self._lock:
            s = self._status.get(path_in_repo)
        if s is not None:
            return s
        # not seen this boot, but maybe indexed by a previous run / the Action
        if self._rel(path_in_repo) in {c.file for c in self.service.index.chunks}:
            return "indexed"
        return "unknown"

    def statuses(self, prefix: str) -> dict[str, str]:
        with self._lock:
            return {p: s for p, s in self._status.items() if p.startswith(prefix)}

    def boot_reconcile(self, uploaded_paths: list[str]) -> int:
        """Enqueue uploads missing from the live index; mark present ones "indexed".

        No hub calls in here — the server passes the listing, and bytes come from the
        local dataset snapshot (base/uploads/**) or a prior merge into materials.
        Returns the number of files enqueued."""
        indexed = {c.file for c in self.service.index.chunks}
        n = 0
        for path in uploaded_paths:
            rel = self._rel(path)
            if rel in indexed:
                with self._lock:
                    self._status.setdefault(path, "indexed")
                continue
            with self._lock:
                if self._status.get(path) in _ACTIVE:
                    continue
            if not self._is_safe(path):
                with self._lock:
                    self._status[path] = "failed: unsafe path"
                continue
            data = self._local_bytes(path, rel)
            if data is None:
                with self._lock:
                    self._status[path] = "failed: upload bytes not found locally"
                continue
            self.enqueue(path, data)
            n += 1
        return n

    def start(self) -> None:
        if self._thread is not None and self._thread.is_alive():
            return
        self._stop.clear()
        self._thread = threading.Thread(target=self._run, name="ingest-worker", daemon=True)
        self._thread.start()

    def stop(self, timeout: float = 5.0) -> None:
        self._stop.set()
        self._queue.put(None)                      # wake the blocking get
        if self._thread is not None:
            self._thread.join(timeout)

    # -- worker thread --

    def _run(self) -> None:
        while True:
            item = self._queue.get()
            if item is None:
                if self._stop.is_set():
                    return
                continue                            # stale wake-up from a previous stop()
            path, data = item
            with self._lock:
                self._status[path] = "processing"
            try:
                self._process(path, data)
            except Exception as exc:                # one bad file must not kill the worker
                log.exception("ingest failed for %s", path)
                with self._lock:
                    self._status[path] = f"failed: {type(exc).__name__}: {exc}"[:300]
            else:
                with self._lock:
                    self._status[path] = "indexed"

    def _process(self, path: str, data: bytes) -> None:
        rel = self._rel(path)
        if not self._is_safe(rel):
            raise ValueError(f"unsafe path: {path}")
        parser = PARSERS.get(Path(rel).suffix.lower())
        if parser is None:
            raise ValueError(f"unsupported file type: {Path(rel).suffix or '?'}")

        dest = self.materials_dir / rel             # where /api/download resolves uploads
        dest.parent.mkdir(parents=True, exist_ok=True)
        dest.write_bytes(data)

        blocks = [b for b in parser(dest, rel) if b.has_text]
        chunks = chunk_blocks(blocks)               # ids identical to a pipeline run
        if not chunks:
            raise ValueError("no extractable text")
        for c in chunks:
            c.concepts = extract_concepts(c.text)
        vectors = self.service.embedder.encode([c.text for c in chunks]).astype("float32")

        new_index = self._swap_index(rel, chunks, vectors)
        self._link_graph(rel, chunks, vectors)
        self._persist(rel, dest, new_index)
        self._push_to_hub(rel)

    # -- index --

    def _swap_index(self, rel: str, new_chunks: list[Chunk], new_vecs: np.ndarray) -> RetrievalIndex:
        """Rebuild-and-swap: the live FAISS index is never mutated — concurrent searches
        keep their old snapshot; one attribute assignment publishes the new one. A
        re-upload of the same path replaces its previous chunks."""
        old = self.service.index
        dim = int(self.service.embedder.dim)
        kept_chunks: list[Chunk] = []
        kept_vecs = np.zeros((0, dim), dtype="float32")
        if old.chunks:
            if old.faiss_index is None or old.faiss_index.ntotal != len(old.chunks):
                raise RuntimeError("live index inconsistent; refusing to rebuild")
            all_vecs = np.asarray(old.faiss_index.reconstruct_n(0, old.faiss_index.ntotal), dtype="float32")
            keep = [i for i, c in enumerate(old.chunks) if c.file != rel]
            kept_chunks = [old.chunks[i] for i in keep]
            if keep:
                kept_vecs = all_vecs[keep]
        chunks = kept_chunks + list(new_chunks)
        vectors = np.vstack([kept_vecs, np.asarray(new_vecs, dtype="float32")])
        dropped = {c.id for c in old.chunks} - {c.id for c in kept_chunks}
        new_index = RetrievalIndex(
            chunks=chunks,
            faiss_index=index_from_vectors(vectors, dim),
            concept_index=self._merge_concepts(old.concept_index, dropped, new_chunks),
            by_id={c.id: c for c in chunks},
            row_by_id={c.id: i for i, c in enumerate(chunks)},
        )
        self.service.index = new_index
        return new_index

    @staticmethod
    def _merge_concepts(existing: dict, dropped: set[str], new_chunks: list[Chunk]) -> dict:
        """Pipeline-shaped merge: existing entries survive (minus replaced chunk ids); new
        concepts need >= 2 chunks (build_concept_index's min_chunks) unless already known."""
        inverted: dict[str, set[str]] = {}
        for concept, ids in existing.items():
            kept = {i for i in ids if i not in dropped}
            if kept:
                inverted[concept] = kept
        for c in new_chunks:
            for concept in {x.lower() for x in c.concepts}:
                inverted.setdefault(concept, set()).add(c.id)
        return {k: sorted(v) for k, v in inverted.items() if k in existing or len(v) >= 2}

    # -- topic graph --

    def _link_graph(self, rel: str, new_chunks: list[Chunk], new_vecs: np.ndarray) -> None:
        """Route the file to one module exactly like build_graph (filename affinity beats
        content cosine), then each chunk to its nearest topic within that module."""
        graph = getattr(self.service, "graph", None)
        if graph is None or not graph.modules or not graph.topics:
            return
        module = self._route_module(rel, new_vecs, graph.modules)
        prefix = f"{rel}#"
        for t in graph.topics:                      # a re-upload's old links are replaced
            if any(cid.startswith(prefix) for cid in t.chunk_ids):
                t.chunk_ids = [cid for cid in t.chunk_ids if not cid.startswith(prefix)]
        topics = graph.topics_for_module(module.id)
        if topics:
            tvecs = self._topic_vectors(module, topics)
            for c, v in zip(new_chunks, np.asarray(new_vecs, dtype="float32")):
                topics[int(np.argmax(tvecs @ v))].chunk_ids.append(c.id)
        if rel not in module.source_ids:
            module.source_ids.append(rel)

    def _route_module(self, rel: str, new_vecs: np.ndarray, modules: list[Module]) -> Module:
        routing = self._module_routing              # modules only change via Action + restart
        if routing is None or len(routing[1]) != len(modules):
            vecs = self.service.embedder.encode(
                [f"{m.title}. {m.lecturer}. {m.objective}" for m in modules]).astype("float32")
            vecs /= np.linalg.norm(vecs, axis=1, keepdims=True) + 1e-9
            identities = [_tokens(f"{m.title} {m.lecturer} {m.source_file}") for m in modules]
            df: dict[str, int] = {}
            for ident in identities:
                for tok in ident:
                    df[tok] = df.get(tok, 0) + 1
            idf = {tok: math.log(1 + len(modules) / n) for tok, n in df.items()}
            routing = self._module_routing = (vecs, identities, idf)
        module_vecs, identities, idf = routing
        file_vec = np.mean(np.asarray(new_vecs, dtype="float32"), axis=0)
        file_vec /= np.linalg.norm(file_vec) + 1e-9
        score = _NAME_WEIGHT * _name_scores(rel, identities, idf) + module_vecs @ file_vec
        return modules[int(np.argmax(score))]

    def _topic_vectors(self, module: Module, topics: list[Topic]) -> np.ndarray:
        vecs = self._topic_vecs.get(module.id)
        if vecs is None or vecs.shape[0] != len(topics):
            vecs = self.service.embedder.encode(
                [f"{module.title}. {t.title}" for t in topics]).astype("float32")
            self._topic_vecs[module.id] = vecs
        return vecs

    # -- persistence --

    def _persist(self, rel: str, dest: Path, index: RetrievalIndex) -> None:
        """Mirror the pipeline's artifact set so the next Action rebuild reuses this file
        (same files, same manifest format: {materials-relative path: sha256})."""
        self.index_dir.mkdir(parents=True, exist_ok=True)
        write_chunks(index.chunks, self.index_dir / "chunks.jsonl")
        save_faiss(index.faiss_index, self.index_dir / "embeddings.faiss")
        write_json(index.concept_index, self.index_dir / "concept_index.json")
        manifest_path = self.index_dir / "manifest.json"
        manifest = read_json(manifest_path) if manifest_path.exists() else {}
        manifest[rel] = file_hash(dest)
        write_json(manifest, manifest_path)
        graph = getattr(self.service, "graph", None)
        if graph is not None:
            save_graph(graph, self.index_dir / "graph.json")

    def _push_to_hub(self, rel: str) -> None:
        if not (self.dataset_repo and self.hf_token):
            return
        try:
            from huggingface_hub import HfApi

            HfApi(token=self.hf_token).upload_folder(
                folder_path=str(self.index_dir), path_in_repo="index",
                repo_id=self.dataset_repo, repo_type="dataset",
                commit_message=f"ingest-worker: index {rel}",
            )
        except Exception:   # local state stays authoritative; the weekly Action reconciles
            log.exception("hub push failed for %s (Action will reconcile)", rel)

    # -- helpers --

    @staticmethod
    def _rel(path_in_repo: str) -> str:
        """Store path 'uploads/<user>/<file>' -> materials-relative '_uploads/<user>/<file>'
        (the path ingest, citations and /api/download all agree on)."""
        return f"_uploads/{path_in_repo.removeprefix('uploads/')}"

    @staticmethod
    def _is_safe(path: str) -> bool:
        return ".." not in Path(path).parts and not Path(path).is_absolute()

    def _local_bytes(self, path: str, rel: str) -> bytes | None:
        """Upload bytes from the local dataset snapshot (base/uploads/**) or an earlier
        merge into materials (originals/_uploads/**); None if neither exists."""
        for cand in (self.materials_dir.parent / path, self.materials_dir / rel):
            try:
                if cand.is_file():
                    return cand.read_bytes()
            except OSError:
                continue
        return None