spb/uqo-research-db

1# ============================================================2# SPID — extraction engine3#4# Iterates enriched sections, runs the chosen extractor5# (rules or gpt-4o-mini), filters by confidence, and loads6# mentions into DuckDB.  LLM calls are threaded (IO-bound).7# ============================================================8from __future__ import annotations910from concurrent.futures import ThreadPoolExecutor, as_completed1112from spid.config import get_logger, get_settings13from spid.extract.base import get_extractor14from spid.store import Store1516log = get_logger("pipeline")171819def _ctx_from_rows(rows: list[dict]) -> dict:20    ctx = {}21    for r in rows:22        ctx.setdefault(r["accession_number"], {23            "ticker": r.get("ticker"),24            "company_name": r.get("company_name") or r.get("filing_company"),25            "filing_company": r.get("filing_company"),26            "form_type": r.get("form_type"),27            "filing_date": r.get("filing_date"),28            "fiscal_year": r.get("fiscal_year"),29            "sector": r.get("sector"),30        })31    return ctx323334def run_extraction(store: Store, *, backend: str | None = None,35                   ciks: list[str] | None = None, forms: list[str] | None = None,36                   limit_sections: int | None = None, min_words: int = 30,37                   workers: int = 8, min_confidence: float | None = None,38                   flush_every: int = 500) -> dict:39    s = get_settings()40    backend = backend or s.backend41    min_conf = s.min_confidence if min_confidence is None else min_confidence42    extractor = get_extractor(backend)4344    rows = store.sections_for(ciks=ciks, forms=forms, limit=limit_sections, min_words=min_words)45    ctx = _ctx_from_rows(rows)46    log.info("Extracting from %d sections | backend=%s | workers=%d", len(rows), backend, workers)4748    kept = 049    buf = []50    done_rows: list[tuple] = []5152    def work(r: dict):53        return extractor.extract(text=r["text"], accession_number=r["accession_number"],54                                 cik=r["cik"], section_id=r["section_id"],55                                 section_name=r["section_name"])5657    def flush():58        nonlocal kept, buf, done_rows59        if buf:60            kept += store.upsert_mentions(buf, ctx)61            buf = []62        if done_rows:63            store.mark_extracted(done_rows)64            done_rows = []6566    done = 067    if backend == "llm" and workers > 1:68        with ThreadPoolExecutor(max_workers=workers) as ex:69            futs = {ex.submit(work, r): r for r in rows}70            for fut in as_completed(futs):71                done += 172                r = futs[fut]73                try:74                    ms = [m for m in fut.result() if m.confidence >= min_conf]75                    done_rows.append((r["section_id"], r["accession_number"], backend, len(ms)))76                except Exception as e:  # noqa: BLE00177                    log.warning("section failed: %s", e); ms = []78                buf.extend(ms)79                if len(buf) >= flush_every:80                    flush()81                if done % 200 == 0:82                    log.info("  ... %d/%d sections, %d mentions kept", done, len(rows), kept + len(buf))83    else:84        for r in rows:85            done += 186            ms = [m for m in work(r) if m.confidence >= min_conf]87            done_rows.append((r["section_id"], r["accession_number"], backend, len(ms)))88            buf.extend(ms)89            if len(buf) >= flush_every:90                flush()91            if done % 2000 == 0:92                log.info("  ... %d/%d sections, %d mentions kept", done, len(rows), kept + len(buf))93    flush()94    log.info("Extraction complete: %d mentions from %d sections.", kept, len(rows))95    return {"sections": len(rows), "mentions": kept, "backend": backend}96