Heimgeist/backend/rag/corpus_enricher.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
RAG Corpus Enrichment (second pass)

What this does
- Reads a JSONL corpus (one record per line) from build_corpus.py and adds LLM-generated enrichment:
  * headline (<= 12 words, language = --summary-lang)
  * summary (2–4 sentences, language = --summary-lang)
  * keywords (5–12, normalized & deduped)
  * entities (name + canonical type: PERSON|ORG|PRODUCT|WORK|PLACE|EVENT|DATE|OTHER)
  * 2–4 likely Q/A pairs (language = --summary-lang)

- Writes two outputs:
  1) --out        : original record + enrichment fields + embedding_text_hint (what your indexer should embed)
  2) --shadow-out : compact “shadow” record for retrieval with normalized shadow_text and useful metadata
       (includes: parent_id, span, size metrics, quality_flags)

Design for speed & robustness
- Default local model: phi4:latest (good balance on Apple/M1 Max). Swap with --model if desired.
- Pooled HTTP via requests.Session (one per thread), bounded semaphore on Ollama calls.
- JSON mode with strict schema + robust repair if the model returns non-JSON.
- Head/Mid/Tail sampling for long texts to stay within context quickly.
- Caching:
    * Main enrichment cache keyed by (prompt_version + model + lang + sampled_text + record_id + record_type)
    * Translation cache keyed by (model + target_lang + field_text)
- Post-enforcement:
    * clamp headline to <=12 words, summary to 2–4 sentences
    * ensure 5–12 keywords, dedup & normalize
    * canonicalize entity types and dedup by name
    * top-up Q/A to required count with a tiny follow-up call (cheap)
    * verify/translate fields to --summary-lang if needed

CLI example
  python rag_enhance_corpus.py \
    --in corpus.jsonl \
    --out corpus.enhanced.jsonl \
    --shadow-out corpus.shadow.jsonl \
    --summary-lang en \
    --ollama http://localhost:11434 \
    --model phi4:latest \
    --concurrency 8 \
    --keep-alive 15m \
    --min-chars 120 \
    --max-text 12000 \
    --timeout 120

Requires: requests, tqdm, (optional) langid, (optional) orjson
"""

from __future__ import annotations
import argparse
import concurrent.futures as cf
import hashlib
import json
import os
import re
import sys
import threading
import time
from dataclasses import dataclass, asdict
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Tuple, Callable
import multiprocessing as mp_context

# Optional faster JSON if available
try:
    import orjson as _orjson
    def json_dumps(obj) -> str:
        return _orjson.dumps(obj, option=_orjson.OPT_NON_STR_KEYS).decode("utf-8")
    def json_loads(s: str) -> Any:
        return _orjson.loads(s)
except Exception:
    def json_dumps(obj) -> str:
        return json.dumps(obj, ensure_ascii=False)
    def json_loads(s: str) -> Any:
        return json.loads(s)

try:
    import langid  # optional language detection
except Exception:
    langid = None

try:
    from tqdm import tqdm
except Exception:
    tqdm = None

# -------------------------
# Constants & helpers
# -------------------------

PROMPT_VERSION = "v4.1-standard-deep"

ENTITY_CANON = {
    "PERSON": "PERSON",
    "ORG": "ORG", "ORGANIZATION": "ORG", "COMPANY": "ORG", "INSTITUTION": "ORG", "COUNTRY": "ORG",
    "PRODUCT": "PRODUCT", "TOOL": "PRODUCT", "LIBRARY": "PRODUCT",
    "WORK": "WORK", "BOOK": "WORK", "PAPER": "WORK", "ARTICLE": "WORK", "MOVIE": "WORK",
    "PLACE": "PLACE", "LOCATION": "PLACE", "CITY": "PLACE", "REGION": "PLACE", "ADDRESS": "PLACE",
    "EVENT": "EVENT", "CONFERENCE": "EVENT", "MEETING": "EVENT",
    "DATE": "DATE", "TIME": "DATE", "YEAR": "DATE",
    "OTHER": "OTHER",
}

QA_TARGET_DEFAULT = 3      # aim for 3 Q/A pairs for normal docs
QA_TARGET_SHORT = 2        # short docs can have 2

# thread-local session pool
_TLS = threading.local()

def get_session():
    import requests
    s = getattr(_TLS, "session", None)
    if s is None:
        s = requests.Session()
        setattr(_TLS, "session", s)
    return s

def log(msg: str, *, verbose: bool = True):
    if verbose:
        print(msg, flush=True)

def strip_think(s: str) -> str:
    return re.sub(r"<\s*think\s*>.*?<\s*/\s*think\s*>", "", s, flags=re.S | re.I)

def sanitize_text(s: str) -> str:
    if not s:
        return ""
    s = strip_think(s)
    s = re.sub(r"^\s*```(?:\w+)?\s*|\s*```\s*$", "", s, flags=re.M)  # strip stray code fences
    s = re.sub(r"[ \t]+", " ", s)
    s = re.sub(r"\n{3,}", "\n\n", s)
    return s.strip()

def detect_lang_quick(s: str) -> Optional[str]:
    s = (s or "").strip()
    if not s:
        return None
    try:
        if langid is not None:
            lang, _ = langid.classify(s[:4000])
            return lang
    except Exception:
        pass
    return None

def sentence_split(text: str) -> List[str]:
    # very light heuristic splitter
    parts = re.split(r"(?<=[.!?])\s+(?=[A-ZÄÖÜ0-9\"'])", text.strip())
    # fall back if we ended up with nothing
    if len(parts) == 1 and len(parts[0]) > 0:
        return [parts[0]]
    return [p.strip() for p in parts if p.strip()]

def clamp_sentences(text: str, min_s: int = 2, max_s: int = 4) -> str:
    sents = sentence_split(text)
    if not sents:
        return ""
    sents = sents[:max_s]
    # if after clamp we have < min_s and original had more, pad; else keep as is
    return " ".join(sents)

def clamp_words(text: str, max_words: int) -> str:
    words = text.strip().split()
    if len(words) <= max_words:
        return text.strip()
    return " ".join(words[:max_words])

def normalize_keywords(items: List[str]) -> List[str]:
    out: List[str] = []
    seen = set()
    for x in items or []:
        s = sanitize_text(str(x))
        s = re.sub(r"^[,;:.\-–—\s]+|[,;:.\-–—\s]+$", "", s)
        s = re.sub(r"\s+", " ", s)
        if not s:
            continue
        key = s.lower()
        if key in seen:
            continue
        seen.add(key)
        out.append(s)
    # enforce 5–12 by trimming if needed
    if len(out) > 12:
        out = out[:12]
    return out

def canonicalize_entities(ents: List[Dict[str, Any]]) -> List[Dict[str, str]]:
    out: List[Dict[str, str]] = []
    seen = set()
    for e in ents or []:
        if not isinstance(e, dict): continue
        name = sanitize_text(str(e.get("name", "")))
        if not name: continue
        typ_raw = sanitize_text(str(e.get("type", ""))).upper()
        typ = ENTITY_CANON.get(typ_raw, ENTITY_CANON.get(typ_raw.split("_")[0], "OTHER"))
        key = name.lower()
        if key in seen: continue
        seen.add(key)
        out.append({"name": name, "type": typ})
    return out

def text_size_metrics(text: str) -> Dict[str, int]:
    text = text or ""
    return {
        "char_count": len(text),
        "word_count": len(text.split()),
        "line_count": len([ln for ln in text.splitlines()]),
    }

def head_mid_tail_sample(s: str, max_chars: int) -> str:
    if len(s) <= max_chars:
        return s
    third = max_chars // 3
    head = s[:third]
    mid_start = max(0, len(s)//2 - third//2)
    mid = s[mid_start:mid_start + third]
    tail = s[-third:]
    return f"{head}\n\n[...] (sample)\n\n{mid}\n\n[...] (sample)\n\n{tail}"

def looks_like_ocr_noise(s: str) -> bool:
    s = s or ""
    if not s.strip():
        return False
    letters = sum(ch.isalpha() for ch in s)
    punct = sum(ch in "!@#$%^&*()[]{}<>/\\|~`" for ch in s)
    ratio_letters = letters / max(1, len(s))
    ratio_punct = punct / max(1, len(s))
    return ratio_letters < 0.45 and ratio_punct > 0.08

def build_doc_hint(rec: Dict[str, Any]) -> str:
    rt = rec.get("record_type") or ""
    mime = rec.get("mime") or ""
    title = rec.get("title") or ""
    if rt == "image":
        return "This record is derived from an IMAGE. If text exists, it may be OCR; otherwise it is an image description. Summaries should read like quality alt-text and include short visible text only if clearly legible."
    if rt == "av":
        return "This record is derived from an AUDIO/VIDEO transcript. Focus on the main points, speakers (if known), and concrete facts. Q&A should target answerable details from the transcript."
    if rt == "code-summary":
        return "This record summarizes a code file. Keywords should emphasize APIs, functions, modules, and side effects. Q&A should focus on how to use or extend the code."
    # PDFs/HTML/TXT/etc.
    if "pdf" in mime:
        return "This record is a PDF page or document content."
    if "html" in mime:
        return "This record is HTML/webpage content."
    if "text" in mime:
        return "This record is plain text content."
    return f"This record is of type '{rt}' with mime '{mime}'. Title (if any): {title}"

def pick_text(d: Dict[str, Any]) -> str:
    return d.get("text") or d.get("content") or d.get("body") or ""

def stable_hash(text: str, model: str, lang: str, rec_id: str, rec_type: str) -> str:
    h = hashlib.sha1()
    for part in (PROMPT_VERSION, model, lang, rec_id or "", rec_type or "", text):
        h.update(part.encode("utf-8", errors="ignore"))
        h.update(b"\x00")
    return h.hexdigest()

# -------------------------
# Ollama calls
# -------------------------

def ollama_generate_json(
    host: str,
    model: str,
    system_prompt: str,
    user_prompt: str,
    *,
    keep_alive: str = "15m",
    timeout: int = 120,
    options: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
    """
    Call Ollama /api/generate in JSON mode (format='json').
    Robust JSON repair if needed.
    """
    session = get_session()
    payload = {
        "model": model,
        "system": system_prompt,
        "prompt": user_prompt,
        "format": "json",
        "stream": False,
        "keep_alive": keep_alive,
    }
    if options:
        payload["options"] = options
    r = session.post(f"{host.rstrip('/')}/api/generate", json=payload, timeout=timeout)
    r.raise_for_status()
    data = r.json()
    raw = sanitize_text(data.get("response", ""))
    try:
        return json_loads(raw)
    except Exception:
        m = re.search(r"\{.*\}", raw, flags=re.S)
        if m:
            try:
                return json_loads(m.group(0))
            except Exception:
                pass
        raise ValueError("Ollama returned incomplete or invalid JSON.")

def ollama_generate_text(
    host: str,
    model: str,
    system_prompt: str,
    user_prompt: str,
    *,
    keep_alive: str = "15m",
    timeout: int = 120,
    options: Optional[Dict[str, Any]] = None,
) -> str:
    """
    Plain text response (no enforced JSON). Used for tiny follow-ups if desired.
    """
    session = get_session()
    payload = {
        "model": model,
        "system": system_prompt,
        "prompt": user_prompt,
        "stream": False,
        "keep_alive": keep_alive,
    }
    if options:
        payload["options"] = options
    r = session.post(f"{host.rstrip('/')}/api/generate", json=payload, timeout=timeout)
    r.raise_for_status()
    data = r.json()
    return sanitize_text(data.get("response", ""))

# -------------------------
# Prompts
# -------------------------

def build_system(summary_lang: str) -> str:
    return (
        "You are a precise, concise, multilingual document tagger for retrieval-augmented generation (RAG). "
        "Return ONLY JSON matching the schema. Avoid markdown. No extra commentary.\n"
        f"Output language for headline/summary/keywords/Q&A must be '{summary_lang}'."
    )

def build_user_main(text: str, summary_lang: str, doc_hint: str, want_qa: int, *, include_qa: bool) -> str:
    want_qa = max(2, min(4, int(want_qa))) if include_qa else 0
    qa_style = "Also produce 2–4 useful QA pairs." if include_qa else "Set qa to an empty list."
    qa_constraint = (
        "- Q&A must be answerable ONLY from the TEXT; keep questions <= 16 words; answers concise (<= ~80 words).\n"
        if include_qa else "- Return an empty qa list.\n"
    )
    # Fixed internal instruction for style/tone
    fixed_instruction = (
        "Produce concise headlines (≤12 words) and 2–4 sentence summaries; "
        f"5–12 normalized keywords (kebab-case) and named entities with types. {qa_style} "
        "Keep strictly grounded in the source."
    )
    return (
        f"{doc_hint}\n\n"
        "You will receive a document TEXT. Produce JSON matching this schema strictly:\n"
        "{\n"
        '  "headline": string,               # <= 12 words\n'
        '  "summary": string,                # 2-4 sentences, faithful and specific\n'
        '  "keywords": [string, ...],        # 5-12 salient terms; multi-word allowed; no hashtags\n'
        '  "entities": [                     # up to ~12 unique items\n'
        '     {"name": string, "type": "PERSON|ORG|PRODUCT|WORK|PLACE|EVENT|DATE|OTHER"}\n'
        "  ],\n"
        f'  "qa": [                           # exactly {want_qa} Q&A pairs\n'
        '     {"q": string, "a": string}\n'
        "  ]\n"
        "}\n\n"
        f"Style Instruction: {fixed_instruction}\n\n"
        f"Constraints:\n"
        f"- Headline and summary MUST be in {summary_lang}.\n"
        "- Extract proper nouns and salient terms as entities; deduplicate by name.\n"
        f"{qa_constraint}"
        "- Be terse and informative; no filler.\n\n"
        "TEXT:\n" + text
    )


def unwrap_enrichment_response(value: Any) -> Dict[str, Any]:
    if not isinstance(value, dict):
        raise ValueError("Enrichment response is not a JSON object.")
    nested = value.get("summary")
    if not value.get("headline") and isinstance(nested, str) and nested.lstrip().startswith("{"):
        try:
            decoded = json_loads(nested)
        except Exception as exc:
            raise ValueError("Enrichment response contains truncated nested JSON.") from exc
        if isinstance(decoded, dict):
            value = decoded
    if not isinstance(value.get("summary"), str) or not str(value.get("summary") or "").strip():
        raise ValueError("Enrichment response is missing a summary.")
    for field in ("keywords", "entities", "qa"):
        if field in value and not isinstance(value[field], list):
            raise ValueError(f"Enrichment response field '{field}' must be a list.")
    return value

def build_user_qa_topup(text: str, summary_lang: str, need: int) -> str:
    need = max(1, min(3, int(need)))
    return (
        "We have a document TEXT and need ONLY additional Q&A pairs for retrieval. "
        "Return STRICT JSON of the form: {\n"
        '  "qa": [ {"q": string, "a": string}, ... ]\n'
        "}\n"
        f"Output language: {summary_lang}. Provide exactly {need} pairs. "
        "Questions <= 16 words; answers concise (<= ~80 words).\n\n"
        "TEXT:\n" + text
    )

def build_system_translate(target_lang: str) -> str:
    return (
        "You are a translator. Return ONLY JSON of the form {\"text\": \"...\"}. "
        "Do not add commentary."
    )

def build_user_translate(text: str, target_lang: str) -> str:
    return (
        f"Translate into {target_lang} preserving meaning and tone.\n"
        "Return: {\"text\": \"...\"} only.\n\n"
        "TEXT:\n" + text
    )

# -------------------------
# Shadow rendering
# -------------------------

def render_shadow(rec: Dict[str, Any], enrichment: Dict[str, Any]) -> Dict[str, Any]:
    """
    Build a compact record for retrieval. 'shadow_text' concatenates fields in a stable order.
    Adds: parent_id, span, size metrics, quality_flags.
    """
    parts: List[str] = []
    h = enrichment.get("headline") or rec.get("title") or ""
    s = enrichment.get("summary") or ""
    kws = enrichment.get("keywords") or []
    ents = enrichment.get("entities") or []
    qas = enrichment.get("qa") or []

    if h: parts.append(f"headline: {h}")
    if s: parts.append(f"summary: {s}")
    if kws:
        kw_line = ", ".join(str(k).strip() for k in kws if str(k).strip())
        parts.append("keywords: " + kw_line)
        # tiny topical boost (helps small embedders)
        parts.append("keywords_boost: " + kw_line)
    if ents:
        uniq = {}
        for e in ents:
            name = (e.get("name") or "").strip()
            et = (e.get("type") or "OTHER").strip().upper()
            if name and name.lower() not in uniq:
                uniq[name.lower()] = (name, et)
        if uniq:
            parts.append("entities: " + "; ".join(f"{n} [{t}]" for n, t in uniq.values()))
    if qas:
        qas_strs = []
        for qa in qas[:4]:
            q = (qa.get("q") or "").strip()
            a = (qa.get("a") or "").strip()
            if q and a:
                qas_strs.append(f"Q: {q}\nA: {a}")
        if qas_strs:
            parts.append("qa:\n" + "\n".join(qas_strs))

    shadow_text = "\n".join(parts).strip()
    meta = {
        "prompt_version": PROMPT_VERSION,
        "size": text_size_metrics(shadow_text),
    }
    parent_id = rec.get("parent_id")
    span = rec.get("span") if isinstance(rec.get("span"), dict) else None

    out = {
        "id": rec.get("id"),
        "parent_id": parent_id,
        "source_path": rec.get("source_path"),
        "url": rec.get("url"),
        "title": rec.get("title"),
        "record_type": rec.get("record_type"),
        "mime": rec.get("mime"),
        "lang": rec.get("lang"),
        "span": span,
        "shadow_text": shadow_text,
        "shadow_meta": meta,
    }
    return out

# -------------------------
# Cache
# -------------------------

class Cache:
    def __init__(self, root: Path, prefix: str = ""):
        self.root = root
        self.root.mkdir(parents=True, exist_ok=True)
        self.lock = threading.Lock()
        self.prefix = prefix

    def _path(self, key: str) -> Path:
        k = (self.prefix + key)
        sub = self.root / k[:2] / (k + ".json")
        sub.parent.mkdir(parents=True, exist_ok=True)
        return sub

    def get(self, key: str) -> Optional[Dict[str, Any]]:
        p = self._path(key)
        if not p.exists():
            return None
        try:
            return json_loads(p.read_text(encoding="utf-8"))
        except Exception:
            return None

    def put(self, key: str, value: Dict[str, Any]):
        p = self._path(key)
        with self.lock:
            p.write_text(json_dumps(value), encoding="utf-8")

# -------------------------
# Post-process + translation guards
# -------------------------

def enforce_schema_and_language(
    out: Dict[str, Any],
    *,
    target_lang: str,
    rec_text_sample: str,
    rec_is_short: bool,
    include_qa: bool,
    perform_translate,
    stats: Dict[str, int],
) -> Dict[str, Any]:
    quality_flags: List[str] = []

    # headline
    headline = sanitize_text(str(out.get("headline", "")))
    if headline:
        hd = clamp_words(headline, 12)
        if hd != headline:
            quality_flags.append("headline_clamped")
        headline = hd

    # summary
    summary = sanitize_text(str(out.get("summary", "")))
    if summary:
        sm = clamp_sentences(summary, 2, 4)
        if sm != summary:
            quality_flags.append("summary_clamped")
        summary = sm

    # keywords
    kws = out.get("keywords", [])
    if isinstance(kws, list):
        kws = normalize_keywords(kws)
        if len(kws) < 5 and headline:
            # augment from headline tokens if we’re short
            extra = [w for w in re.split(r"[,\s]+", headline) if len(w) > 3]
            kws = normalize_keywords((kws or []) + extra)
        if len(kws) < 5 and summary:
            extra = [w for w in re.findall(r"[A-Za-zÀ-ÖØ-öø-ÿ][A-Za-zÀ-ÖØ-öø-ÿ\-]{2,}", summary)]
            kws = normalize_keywords((kws or []) + extra)
        if len(kws) > 12:
            kws = kws[:12]
    else:
        kws = []

    # entities
    ents = canonicalize_entities(out.get("entities", []))

    # QA
    qas = []
    if include_qa:
        for qa in out.get("qa", []) or []:
            if not isinstance(qa, dict): continue
            q = clamp_words(sanitize_text(str(qa.get("q", ""))), 16)
            a = sanitize_text(str(qa.get("a", "")))
            if q and a and len(a) >= 30:
                qas.append({"q": q, "a": a})
        target = QA_TARGET_SHORT if rec_is_short else QA_TARGET_DEFAULT
        if len(qas) < target:
            need = target - len(qas)
            add = perform_translate("__QATOPUP__", rec_text_sample, need)
            extra = []
            if isinstance(add, dict):
                for qa in add.get("qa", []) or []:
                    if not isinstance(qa, dict): continue
                    q = clamp_words(sanitize_text(str(qa.get("q", ""))), 16)
                    a = sanitize_text(str(qa.get("a", "")))
                    if q and a and len(a) >= 30:
                        extra.append({"q": q, "a": a})
            if extra:
                qas.extend(extra[:need])
                quality_flags.append("qa_topped_up")
                stats["qa_topped_up"] += 1

    # Language guard (per-field)
    def _guard_lang(field_value: str, field_name: str) -> str:
        if not field_value:
            return field_value
        detected = detect_lang_quick(field_value)
        if detected and target_lang and detected != target_lang:
            tr = perform_translate(field_name, field_value, 0)  # 0 = translate exactly this string
            if isinstance(tr, dict):
                txt = sanitize_text(str(tr.get("text", "")))
            else:
                txt = sanitize_text(str(tr) if tr else "")
            if txt:
                quality_flags.append(f"{field_name}_translated")
                stats["translated_fields"] += 1
                return txt
        return field_value

    headline = _guard_lang(headline, "headline")
    summary  = _guard_lang(summary, "summary")
    # translate Q&A fields if needed
    fixed_qas = []
    for qa in qas:
        q = _guard_lang(qa["q"], "qa_q")
        a = _guard_lang(qa["a"], "qa_a")
        fixed_qas.append({"q": q, "a": a})
    qas = fixed_qas

    return {
        "headline": headline,
        "summary": summary,
        "keywords": kws,
        "entities": ents,
        "qa": qas,
        "quality_flags": quality_flags,
    }

# -------------------------
# Worker
# -------------------------

@dataclass
class Args:
    inp: str
    out: str
    shadow_out: str
    ollama: str
    model: str
    summary_lang: str
    concurrency: int
    keep_alive: str
    timeout: int
    min_chars: int
    max_text: int
    force: bool
    cache_dir: str
    verbose: bool
    deep_source_paths: frozenset[str]

def enrich_one(
    rec: Dict[str, Any],
    *,
    args: Args,
    cache_main: Cache,
    cache_tr: Cache,
    sem: threading.BoundedSemaphore,
    stats: Dict[str, int],
) -> Tuple[Dict[str, Any], Dict[str, Any]]:
    """
    Return (enriched_record, shadow_record)
    """
    base_text = sanitize_text(pick_text(rec))
    rec_id = str(rec.get("id") or "")
    rec_type = str(rec.get("record_type") or "")
    doc_hint = build_doc_hint(rec)
    source_path = str(rec.get("source_path") or "")
    try:
        source_key = str(Path(source_path).expanduser().resolve())
    except Exception:
        source_key = source_path
    include_qa = source_key in args.deep_source_paths
    level = "deep" if include_qa else "standard"

    is_short = len(base_text) < args.min_chars and not include_qa
    sampled = base_text if len(base_text) <= args.max_text else head_mid_tail_sample(base_text, args.max_text)
    target_lang = args.summary_lang
    if target_lang == "auto":
        target_lang = str(rec.get("lang") or "").strip().lower() or (detect_lang_quick(sampled) or "en")
    qa_target = QA_TARGET_SHORT if is_short else QA_TARGET_DEFAULT

    # short fast-path (no LLM)
    if is_short:
        enrichment = {
            "headline": (rec.get("title") or base_text[:80]).strip(),
            "summary": clamp_sentences(base_text[:400], 1, 3),
            "keywords": [],
            "entities": [],
            "qa": [],
            "model": None,
            "prompt_version": PROMPT_VERSION,
            "cached": True,
            "strategy": "short-fastpath",
            "level": level,
        }
        enr = dict(rec)
        enr.update({
            "headline": enrichment["headline"],
            "summary": enrichment["summary"],
            "keywords": enrichment["keywords"],
            "entities": enrichment["entities"],
            "qa": enrichment["qa"],
            "enrichment_meta": {
                "model": None,
                "prompt_version": PROMPT_VERSION,
                "cached": True,
                "strategy": "short-fastpath",
                "level": level,
                "ok": True,
                "error": None,
            }
        })
        shadow = render_shadow(rec, enrichment)
        # embedding hint prefers shadow_text
        enr["embedding_text_hint"] = shadow["shadow_text"]
        return enr, shadow

    # OCR noise guard: nudge the prompt to produce a descriptive summary
    if looks_like_ocr_noise(sampled):
        doc_hint += " The TEXT appears noisy/garbled (possibly OCR). Summarize what the document likely conveys and any clearly legible details; avoid copying garbled strings."

    # caching
    key = stable_hash(sampled, args.model, target_lang, rec_id, f"{rec_type}:{level}")
    if not args.force:
        hit = cache_main.get(key)
        if hit is not None:
            enriched = dict(rec)
            enriched.update({
                "headline": hit.get("headline"),
                "summary": hit.get("summary"),
                "keywords": hit.get("keywords"),
                "entities": hit.get("entities"),
                "qa": hit.get("qa"),
                "enrichment_meta": {
                    "model": hit.get("model"),
                    "prompt_version": hit.get("prompt_version"),
                    "cached": True,
                    "strategy": hit.get("strategy"),
                    "level": hit.get("level") or level,
                    "ok": True,
                    "error": None,
                }
            })
            shadow = render_shadow(rec, hit)
            enriched["embedding_text_hint"] = shadow["shadow_text"]
            stats["cache_hits"] += 1
            return enriched, shadow

    # tiny helper: translation or QA top-up calls (cached for translations)
    def perform_translate(kind: str, payload: str, need_pairs: int) -> Dict[str, Any] | str:
        if kind == "__QATOPUP__":
            # request exactly need_pairs additional pairs
            sys_prompt = build_system(target_lang)
            usr_prompt = build_user_qa_topup(sampled, target_lang, need_pairs)
            opts = {"temperature": 0.2, "repeat_penalty": 1.1, "top_p": 0.9, "num_predict": 280}
            with sem:
                tries, backoff, last = 2, 1.5, None
                for i in range(tries):
                    try:
                        return ollama_generate_json(args.ollama, args.model, sys_prompt, usr_prompt,
                                                    keep_alive=args.keep_alive, timeout=args.timeout, options=opts)
                    except Exception as e:
                        last = e
                        time.sleep(backoff ** (i+1))
                # failure → empty result
                return {"qa": []}
        else:
            # per-field translation caching
            tr_key = stable_hash(payload, args.model, target_lang, kind, "translate")
            if not args.force:
                tr_hit = cache_tr.get(tr_key)
                if tr_hit is not None:
                    return tr_hit
            sys_prompt = build_system_translate(target_lang)
            usr_prompt = build_user_translate(payload, target_lang)
            opts = {"temperature": 0.2, "repeat_penalty": 1.05, "top_p": 0.9, "num_predict": 200}
            with sem:
                tries, backoff, last = 2, 1.5, None
                for i in range(tries):
                    try:
                        out = ollama_generate_json(args.ollama, args.model, sys_prompt, usr_prompt,
                                                   keep_alive=args.keep_alive, timeout=args.timeout, options=opts)
                        # normalize
                        if not isinstance(out, dict):
                            out = {"text": sanitize_text(str(out))}
                        else:
                            out["text"] = sanitize_text(str(out.get("text", "")))
                        cache_tr.put(tr_key, out)
                        return out
                    except Exception as e:
                        last = e
                        time.sleep(backoff ** (i+1))
                return {"text": payload}  # give up: return original

    # main call
    system = build_system(target_lang)
    user = build_user_main(sampled, target_lang, doc_hint, qa_target, include_qa=include_qa)
    options = {
        "temperature": 0.2,
        "repeat_penalty": 1.1,
        "top_p": 0.9,
        "num_predict": 1000 if include_qa else 700,
    }

    with sem:
        tries, backoff, last_exc = 3, 1.5, None
        for i in range(tries):
            try:
                out = ollama_generate_json(args.ollama, args.model, system, user,
                                           keep_alive=args.keep_alive, timeout=args.timeout, options=options)
                out = unwrap_enrichment_response(out)
                # sanitize + normalize structure
                for k in ("headline", "summary"):
                    if k in out and isinstance(out[k], str):
                        out[k] = sanitize_text(out[k])

                # normalize arrays to expected types
                out["keywords"] = [sanitize_text(str(x)) for x in out.get("keywords", []) if str(x).strip()]
                ents = []
                for e in out.get("entities", []) or []:
                    if isinstance(e, dict):
                        name = sanitize_text(str(e.get("name", "")))
                        typ = sanitize_text(str(e.get("type", "OTHER")))
                        if name:
                            ents.append({"name": name, "type": typ})
                out["entities"] = ents

                qas = []
                for qa in out.get("qa", []) or []:
                    if isinstance(qa, dict):
                        q = sanitize_text(str(qa.get("q", "")))
                        a = sanitize_text(str(qa.get("a", "")))
                        if q and a:
                            qas.append({"q": q, "a": a})
                out["qa"] = qas

                # post-enforce schema + language
                fixed = enforce_schema_and_language(
                    out,
                    target_lang=target_lang,
                    rec_text_sample=sampled,
                    rec_is_short=is_short,
                    include_qa=include_qa,
                    perform_translate=perform_translate,
                    stats=stats,
                )

                result = {
                    "headline": fixed["headline"],
                    "summary": fixed["summary"],
                    "keywords": fixed["keywords"],
                    "entities": fixed["entities"],
                    "qa": fixed["qa"],
                    "quality_flags": fixed["quality_flags"],
                    "model": args.model,
                    "prompt_version": PROMPT_VERSION,
                    "cached": False,
                    "strategy": "sampled" if len(base_text) > args.max_text else "full",
                    "level": level,
                }

                # save to cache
                cache_main.put(key, result)

                enriched = dict(rec)
                enriched.update({
                    "headline": result["headline"],
                    "summary": result["summary"],
                    "keywords": result["keywords"],
                    "entities": result["entities"],
                    "qa": result["qa"],
                    "enrichment_meta": {
                        "model": args.model,
                        "prompt_version": PROMPT_VERSION,
                        "cached": False,
                        "strategy": result["strategy"],
                        "level": level,
                        "ok": True,
                        "error": None,
                        "quality_flags": result["quality_flags"],
                    }
                })
                shadow = render_shadow(rec, result)
                enriched["embedding_text_hint"] = shadow["shadow_text"]
                return enriched, shadow

            except Exception as e:
                last_exc = e
                time.sleep(backoff ** (i+1))

    # fallback if everything failed
    stats["fallbacks"] += 1
    fallback_summary = clamp_sentences(sampled[:1000], 2, 4)
    fallback = {
        "headline": (rec.get("title") or sampled.split("\n", 1)[0][:80]).strip(),
        "summary": fallback_summary,
        "keywords": [],
        "entities": [],
        "qa": [],
        "model": None,
        "prompt_version": PROMPT_VERSION,
        "cached": False,
        "strategy": f"fallback:{type(last_exc).__name__ if last_exc else 'error'}",
        "level": level,
        "quality_flags": ["fallback"],
    }
    enriched = dict(rec)
    enriched.update({
        "headline": fallback["headline"],
        "summary": fallback["summary"],
        "keywords": [],
        "entities": [],
        "qa": [],
        "enrichment_meta": {
            "model": None,
            "prompt_version": PROMPT_VERSION,
            "cached": False,
            "strategy": fallback["strategy"],
            "level": level,
            "ok": False,
            "error": str(last_exc) if last_exc else "unknown",
            "quality_flags": ["fallback"],
        }
    })
    shadow = render_shadow(rec, fallback)
    enriched["embedding_text_hint"] = shadow["shadow_text"]
    return enriched, shadow

# -------------------------
# IO
# -------------------------

def iter_jsonl(path: Path) -> Iterable[Dict[str, Any]]:
    with path.open("r", encoding="utf-8") as f:
        for line in f:
            if not line.strip():
                continue
            try:
                yield json_loads(line)
            except Exception:
                continue

def write_line(path: Path, obj: Dict[str, Any], lock: threading.Lock, *, dry_run: bool = False):
    if dry_run:
        return
    line = json_dumps(obj) + "\n"
    with lock:
        with path.open("a", encoding="utf-8") as fh:
            fh.write(line)
            fh.flush()

# -------------------------
# CLI
# -------------------------

def parse_args():
    p = argparse.ArgumentParser(description="Enrich a JSONL corpus with small-LLM generated summaries, keywords, entities and Q&A.")
    p.add_argument("--in", dest="inp", required=True, help="Input JSONL (from build_corpus.py)")
    p.add_argument("--out", required=True, help="Output JSONL with enrichment fields merged into each record")
    p.add_argument("--shadow-out", required=True, help="Output JSONL of compact 'shadow' records for retrieval")
    p.add_argument("--ollama", default="http://localhost:11434", help="Ollama base URL")
    p.add_argument("--model", default="phi4:latest", help="Local model (e.g., 'phi4:latest' or 'llama3.2:3b')")
    p.add_argument("--summary-lang", default="en", help="Language of headline/summary/keywords/Q&A")
    p.add_argument("--concurrency", type=int, default=max(2, (os.cpu_count() or 4)//2), help="Parallel HTTP workers")
    p.add_argument("--keep-alive", default="15m", help="Ollama keep_alive value (e.g., '15m', '-1' for forever)")
    p.add_argument("--timeout", type=int, default=120, help="HTTP timeout per request (seconds)")
    p.add_argument("--min-chars", type=int, default=120, help="Skip LLM when text shorter than this (fast-path heuristic)")
    p.add_argument("--max-text", type=int, default=12000, help="If text is longer, sample head/mid/tail to this many chars")
    p.add_argument("--force", action="store_true", help="Ignore cache and regenerate everything")
    p.add_argument("--cache-dir", default=".rag_cache", help="Directory for per-record JSON cache")
    p.add_argument("--dry-run", action="store_true", help="Do the work but do not write outputs")
    p.add_argument("--verbose", action="store_true", help="Verbose logging")
    return p.parse_args()

# -------------------------
# Main
# -------------------------

def run_enrich(inp: Path, out: Path, shadow_out: Path, *,
               summary_lang: str = "auto",
               on_progress: Optional[Callable[[str, float, str], None]] = None,
               cancellation_event: Optional[threading.Event] = None, **opts) -> dict:
    deep_source_paths = set()
    for value in opts.get("deep_source_paths", []) or []:
        try:
            deep_source_paths.add(str(Path(str(value)).expanduser().resolve()))
        except Exception:
            deep_source_paths.add(str(value))

    args = Args(
        inp=str(inp),
        out=str(out),
        shadow_out=str(shadow_out),
        ollama=opts.get("ollama", "http://localhost:11434"),
        model=opts.get("model", "phi4:latest"),
        summary_lang=summary_lang,
        concurrency=opts.get("concurrency", max(2, (os.cpu_count() or 4)//2)),
        keep_alive=opts.get("keep_alive", "15m"),
        timeout=opts.get("timeout", 120),
        min_chars=opts.get("min_chars", 120),
        max_text=opts.get("max_text", 12000),
        force=opts.get("force", False),
        cache_dir=opts.get("cache_dir", ".rag_cache"),
        verbose=opts.get("verbose", False),
        deep_source_paths=frozenset(deep_source_paths),
    )

    src = Path(args.inp).expanduser().resolve()
    if not src.exists():
        return {"status": "error", "message": f"Input not found: {src}"}

    out_path = Path(args.out).expanduser().resolve()
    shadow_path = Path(args.shadow_out).expanduser().resolve()
    out_path.parent.mkdir(parents=True, exist_ok=True)
    shadow_path.parent.mkdir(parents=True, exist_ok=True)

    out_path.write_text("", encoding="utf-8")
    shadow_path.write_text("", encoding="utf-8")

    cache_main = Cache(Path(args.cache_dir), prefix="enrich_")
    cache_tr   = Cache(Path(args.cache_dir), prefix="trans_")

    sem = threading.BoundedSemaphore(max(1, args.concurrency))
    lock_out = threading.Lock()
    lock_sh = threading.Lock()

    if on_progress:
        on_progress("load", 0.05, "Loading records...")
    records = list(iter_jsonl(src))
    total = len(records)
    if total == 0:
        if on_progress:
            on_progress("done", 1.0, "Empty input.")
        return {"status": "warning", "message": "Empty input."}

    stats = {
        "cache_hits": 0,
        "fallbacks": 0,
        "qa_topped_up": 0,
        "translated_fields": 0,
        "processed": 0,
    }

    def _worker(rec: Dict[str, Any]) -> None:
        if cancellation_event and cancellation_event.is_set():
            return # Exit early if cancelled
        try:
            enriched, shadow = enrich_one(
                rec, args=args, cache_main=cache_main, cache_tr=cache_tr,
                sem=sem, stats=stats
            )
            write_line(out_path, enriched, lock_out)
            write_line(shadow_path, shadow, lock_sh)
        except Exception as e:
            passthru = dict(rec)
            passthru["enrichment_meta"] = {
                "model": None, "prompt_version": PROMPT_VERSION, "cached": False,
                "strategy": "error", "ok": False, "error": f"{type(e).__name__}: {e}",
                "quality_flags": ["error"],
            }
            write_line(out_path, passthru, lock_out)
        finally:
            stats["processed"] += 1
            if on_progress:
                pct = 0.1 + 0.8 * (stats["processed"] / total)
                on_progress("enrich", pct, f"Processed {stats['processed']}/{total}")

    if on_progress:
        on_progress("enrich", 0.1, f"Enriching {total} records...")
    with cf.ThreadPoolExecutor(max_workers=max(1, args.concurrency)) as ex:
        futs = [ex.submit(_worker, r) for r in records]
        for fut in cf.as_completed(futs):
            if cancellation_event and cancellation_event.is_set():
                for f in futs:
                    f.cancel() # Attempt to cancel remaining futures
                if on_progress:
                    on_progress("done", 1.0, "Enrichment cancelled.")
                return {"status": "cancelled", "message": "Enrichment cancelled."}
            _ = fut.result()

    if on_progress:
        on_progress("done", 1.0, "Enrichment complete.")
    
    return {"status": "ok", "stats": stats}

def main():
    a = parse_args()
    run_enrich(
        Path(a.inp), Path(a.out), Path(a.shadow_out),
        summary_lang=a.summary_lang,
        on_progress=lambda p, pct, d: print(f"[{p}] {pct*100:.1f}%: {d}"),
        **vars(a)
    )

if __name__ == "__main__":
    main()