Upload tools/retrieval/__init__.py with huggingface_hub

tools/retrieval/__init__.py

@@ -0,0 +1,198 @@
+"""Pluggable retrieval for serve.py (Phase 5b).
+
+Retriever(config_path) loads the sources listed in retrieval/config.yaml. Each
+source implements search(query, k) -> list of passages
+{"text", "title", "source", "score", "meta"}; the Retriever normalises scores
+per source, applies the source weight, and returns the merged top-k.
+Source types: kb (literature KB), wiki_dump (prebuilt BM25 over a Wikipedia
+dump), folder (local .txt/.md/.jsonl, indexed lazily), search_api (HTTP JSON).
+"""
+import glob, json, os, time, urllib.parse, urllib.request
+import yaml
+from .bm25 import BM25Index, build as build_bm25
+from .text import toks, chunk_words, nfc
+from . import kb as kbmod
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+DEFAULT_CONFIG = os.path.join(HERE, "config.yaml")
+
+class Passage(dict):
+    pass
+
+class KBSource:
+    type = "kb"
+    def __init__(self, cfg):
+        self.name = cfg["name"]; self.k = int(cfg.get("k", 3)); self.weight = float(cfg.get("weight", 1.0))
+        self.units = kbmod.load_units(cfg.get("glob", "data/kb/*.jsonl"))
+        d = cfg.get("index_dir", "data/index/literature_kb")
+        marker = os.path.join(d, "meta.json")
+        newest = max((os.path.getmtime(f) for f in glob.glob(cfg.get("glob", "data/kb/*.jsonl"))), default=0)
+        if not BM25Index.exists(d) or os.path.getmtime(marker) < newest:
+            build_bm25(d, (toks(kbmod.unit_text(u)) for u in self.units))
+        self.idx = BM25Index(d)
+    def search(self, query, k=None):
+        out = []
+        for i, s in self.idx.search(query, k or self.k):
+            u = self.units[i]
+            out.append(Passage(text=kbmod.format_unit(u), title=f"{u.get('work')} {u.get('number', '')}".strip(),
+                               source=self.name, score=s, meta={"unit": u}))
+        return out
+
+class WikiDumpSource:
+    type = "wiki_dump"
+    def __init__(self, cfg):
+        self.name = cfg["name"]; self.k = int(cfg.get("k", 4)); self.weight = float(cfg.get("weight", 1.0))
+        self.dir = cfg["index_dir"]
+        fs_dir = cfg.get("family_safe_index_dir")
+        if os.environ.get("FAMILY_SAFE", "1") != "0" and fs_dir and BM25Index.exists(fs_dir) and os.path.exists(os.path.join(fs_dir, "passages.jsonl")):
+            self.dir = fs_dir   # family-safe index (severe lexicon hits dropped at build time) when FAMILY_SAFE=1
+        self.ok = BM25Index.exists(self.dir) and os.path.exists(os.path.join(self.dir, "passages.jsonl"))
+        if not self.ok:
+            print(f"[retrieval] source {self.name} disabled: index missing at {self.dir} (run retrieval/build_wiki_index.py)")
+            return
+        self.idx = BM25Index(self.dir)
+        self.offsets = __import__("numpy").load(os.path.join(self.dir, "passage_offsets.npy"))
+        self.fh = open(os.path.join(self.dir, "passages.jsonl"), "rb")
+    def _passage(self, i):
+        self.fh.seek(int(self.offsets[i])); return json.loads(self.fh.readline())
+    def search(self, query, k=None):
+        if not self.ok:
+            return []
+        kk = k or self.k
+        qn = " ".join(t for t in toks(query) if t not in __import__("retrieval.text", fromlist=["STOPWORDS"]).STOPWORDS)
+        out = []
+        for i, s in self.idx.search(query, max(kk * 5, 20)):
+            p = self._passage(i)
+            tn = " ".join(toks(p["title"]))
+            if qn and tn == qn:
+                s *= 1.5          # exact title match: the article about the thing asked
+            elif qn and qn in tn:
+                s *= 1.15
+            out.append(Passage(text=p["text"], title=p["title"], source=self.name, score=s, meta={"title": p["title"]}))
+        out.sort(key=lambda h: -h["score"])
+        return out[:kk]
+
+class FolderSource:
+    type = "folder"
+    def __init__(self, cfg):
+        self.name = cfg["name"]; self.k = int(cfg.get("k", 4)); self.weight = float(cfg.get("weight", 1.0))
+        self.path = cfg["path"]; self.dir = cfg.get("index_dir", f"data/index/{self.name}")
+        self.ok = os.path.isdir(self.path)
+        if not self.ok:
+            print(f"[retrieval] source {self.name} disabled: folder {self.path} not found"); return
+        files = self._files()
+        newest = max((os.path.getmtime(f) for f in files), default=0)
+        marker = os.path.join(self.dir, "meta.json")
+        if not BM25Index.exists(self.dir) or os.path.getmtime(marker) < newest:
+            self._build(files)
+        self.idx = BM25Index(self.dir)
+        self.passages = [json.loads(l) for l in open(os.path.join(self.dir, "passages.jsonl"), encoding="utf-8")]
+    def _files(self):
+        out = []
+        for root, _, fs in os.walk(self.path):
+            out += [os.path.join(root, f) for f in fs if f.endswith((".txt", ".md", ".jsonl"))]
+        return sorted(out)
+    def _build(self, files):
+        os.makedirs(self.dir, exist_ok=True)
+        passages = []
+        for f in files:
+            rel = os.path.relpath(f, self.path)
+            if f.endswith(".jsonl"):
+                for l in open(f, encoding="utf-8"):
+                    try:
+                        d = json.loads(l)
+                    except Exception:
+                        continue
+                    for c in chunk_words(str(d.get("text", "")), target=200):
+                        passages.append({"title": d.get("title") or rel, "text": c, "file": rel})
+            else:
+                for c in chunk_words(open(f, encoding="utf-8", errors="replace").read(), target=200):
+                    passages.append({"title": rel, "text": c, "file": rel})
+        with open(os.path.join(self.dir, "passages.jsonl"), "w", encoding="utf-8") as w:
+            for p in passages:
+                w.write(json.dumps(p, ensure_ascii=False) + "\n")
+        build_bm25(self.dir, (toks(p["title"] + " " + p["text"]) for p in passages))
+        print(f"[retrieval] built folder index {self.name}: {len(files)} files, {len(passages)} passages")
+    def search(self, query, k=None):
+        if not self.ok:
+            return []
+        return [Passage(text=self.passages[i]["text"], title=self.passages[i]["title"], source=self.name, score=s,
+                        meta={"file": self.passages[i]["file"]}) for i, s in self.idx.search(query, k or self.k)]
+
+class SearchAPISource:
+    type = "search_api"
+    def __init__(self, cfg):
+        self.name = cfg["name"]; self.k = int(cfg.get("k", 4)); self.weight = float(cfg.get("weight", 1.0))
+        self.url = cfg["url"]; self.results_key = cfg.get("results_key", "results")
+        self.text_key = cfg.get("text_key", "text"); self.title_key = cfg.get("title_key", "title")
+        self.header = cfg.get("api_key_header", "Authorization"); self.bearer = bool(cfg.get("bearer", True))
+        env = cfg.get("api_key_env")
+        self.key = os.environ.get(env) if env else None
+        self.ok = bool(self.key) or not env
+        if not self.ok:
+            print(f"[retrieval] source {self.name} disabled: env var {env} not set")
+        self.timeout = float(cfg.get("timeout", 8))
+    def search(self, query, k=None):
+        if not self.ok:
+            return []
+        url = self.url.replace("{query}", urllib.parse.quote(query))
+        req = urllib.request.Request(url, headers={"User-Agent": "tamil-lm-serve/1.0"})
+        if self.key:
+            req.add_header(self.header, f"Bearer {self.key}" if self.bearer else self.key)
+        try:
+            with urllib.request.urlopen(req, timeout=self.timeout) as r:
+                data = json.loads(r.read().decode("utf-8"))
+        except Exception as e:
+            print(f"[retrieval] {self.name} request failed: {type(e).__name__}"); return []
+        node = data
+        for part in self.results_key.split("."):
+            node = node.get(part, []) if isinstance(node, dict) else []
+        out = []
+        for j, hit in enumerate(node[: k or self.k]):
+            if not isinstance(hit, dict):
+                continue
+            out.append(Passage(text=str(hit.get(self.text_key, "")), title=str(hit.get(self.title_key, "")),
+                               source=self.name, score=float(len(node) - j), meta={"url": hit.get("url")}))
+        return out
+
+SOURCE_TYPES = {c.type: c for c in (KBSource, WikiDumpSource, FolderSource, SearchAPISource)}
+
+class Retriever:
+    def __init__(self, config_path=DEFAULT_CONFIG, only=None):
+        cfg = yaml.safe_load(open(config_path))
+        self.top_k = int(cfg.get("top_k", 4))
+        self.sources = []
+        for s in cfg.get("sources", []):
+            if only and s["name"] not in only:
+                continue
+            cls = SOURCE_TYPES.get(s.get("type"))
+            if cls is None:
+                print(f"[retrieval] unknown source type {s.get('type')} for {s.get('name')}; skipped"); continue
+            t0 = time.time()
+            self.sources.append(cls(s))
+            print(f"[retrieval] loaded {s['name']} ({s['type']}) in {time.time() - t0:.1f}s")
+    def source(self, name):
+        return next((s for s in self.sources if s.name == name), None)
+    def search(self, query, k=None):
+        merged = []
+        for src in self.sources:
+            hits = src.search(query)
+            if not hits:
+                continue
+            # rank across sources by the raw BM25 score times the source weight; per-source normalisation
+            # made every source's best hit look equally good (a weak literature match outranked the
+            # right Wikipedia article for "rasam recipe", 2026-09-07)
+            for h in hits:
+                h["score_raw"] = h["score"]; h["score"] = src.weight * h["score"]
+                merged.append(h)
+        merged.sort(key=lambda h: -h["score"])
+        return merged[: k or self.top_k]
+
+def format_passages(hits):
+    out = []
+    for h in hits:
+        if h["source"] == "literature_kb":
+            out.append(h["text"])
+        else:
+            out.append(f"[{h['source']}: {h['title']}]\n{h['text']}")
+    return "\n\n".join(out)