"""Batch orchestrator for claim extraction. Walks all Type/Document atoms in Atomic, runs extraction on each, and pushes the resulting claims as new Type/Claim atoms. Idempotent across runs via the state file in extractor/_extracted.json. """ from __future__ import annotations import asyncio from collections.abc import AsyncIterator from dataclasses import dataclass, field from typing import Any from urllib.parse import unquote from shared.atomic_api import AtomicClient from shared.config import settings from shared.embedding_client import EmbeddingClient # noqa: F401 (future use) from shared.llm_client import LlmClient, LlmError from shared.logging import get_logger from shared.taxonomy import TagResolver from extractor._state import DocExtractionRecord, ExtractionState, now_iso from extractor.extract import ( PROMPT_VERSION, ExtractionResult, extract_claims_from_atom, ) from extractor.push import push_claim log = get_logger(__name__) @dataclass(slots=True) class BatchStats: docs_seen: int = 0 docs_skipped_already_done: int = 0 docs_processed: int = 0 docs_failed: int = 0 claims_raw: int = 0 claims_valid: int = 0 claims_created: int = 0 claims_duplicate: int = 0 claims_error: int = 0 rejected_reasons: dict[str, int] = field(default_factory=dict) # ====================================================== document selection async def _list_documents_to_process( atomic: AtomicClient, resolver: TagResolver, *, limit: int = 1000 ) -> list[dict[str, Any]]: """Return all atoms tagged Type/Document, with their tags inlined. We page through /api/atoms?tag_id= and pull metadata for each, since the list endpoint already returns tags inline. """ type_doc_id = resolver.require("Type/Document") page_size = 50 out: list[dict[str, Any]] = [] offset = 0 while True: result = await atomic.list_atoms( limit=page_size, offset=offset, tag_id=type_doc_id ) atoms = result.get("atoms") or result.get("data") or (result if isinstance(result, list) else []) if not atoms: break for a in atoms: out.append(a) if len(out) >= limit: return out if len(atoms) < page_size: break offset += page_size return out def _title_from_atom(atom: dict[str, Any]) -> str: """Best-effort title: prefer the Markdown H1 in content, fall back to URL slug.""" content = atom.get("content") or "" # Look for the first '# ...' line at the start for line in content.lstrip().splitlines(): line = line.strip() if line.startswith("# "): return line[2:].strip() if line: break # first non-empty isn't a header → fall through to URL url = atom.get("source_url") or "" if url: last = url.rstrip("/").rsplit("/", 1)[-1] return unquote(last).replace("_", " ") return atom.get("id", "?")[:8] def _language_from_atom(atom: dict[str, Any]) -> str: """Read the Language/ tag if present, default 'EN'.""" for tag in atom.get("tags") or []: name = tag.get("name", "") # The tag list returns just `name`, not the full path. Languages are # short codes (RO/EN/RU/...) so direct match works. if name in {"RO", "EN", "RU", "UA", "FR", "DE", "ES", "IT", "PL"}: return name return "EN" # ============================================================== one document async def process_one( *, llm: LlmClient, atomic: AtomicClient, resolver: TagResolver, state: ExtractionState, atom: dict[str, Any], stats: BatchStats, ) -> None: atom_id = atom["id"] if state.has(atom_id, PROMPT_VERSION): stats.docs_skipped_already_done += 1 return # /api/atoms (list) returns summary objects WITHOUT full content. We have # to fetch the full atom individually to get the body for extraction. full_atom = await atomic.get_atom(atom_id) content = full_atom.get("content") or "" title = _title_from_atom(full_atom) language = _language_from_atom(full_atom) if not content: log.warning("doc_no_content", atom_id=atom_id) stats.docs_failed += 1 return log.info("extracting", atom_id=atom_id[:8], title=title, lang=language, chars=len(content)) try: result: ExtractionResult = await extract_claims_from_atom( llm, title=title, language=language, content=content, ) except LlmError as e: log.error("extraction_failed", atom_id=atom_id, error=str(e), body=(e.body or "")[:300]) stats.docs_failed += 1 return except Exception as e: # noqa: BLE001 log.error("extraction_crashed", atom_id=atom_id, error=f"{type(e).__name__}: {e}") stats.docs_failed += 1 return stats.docs_processed += 1 stats.claims_raw += result.raw_count stats.claims_valid += len(result.valid) for k, v in result.rejected.items(): stats.rejected_reasons[k] = stats.rejected_reasons.get(k, 0) + v pushed = 0 duplicates = 0 errors = 0 for c in result.valid: _, status = await push_claim( atomic, parent_atom=full_atom, claim=c, parent_title=title, resolver=resolver, ) if status == "created": pushed += 1 elif status == "duplicate": duplicates += 1 else: errors += 1 stats.claims_created += pushed stats.claims_duplicate += duplicates stats.claims_error += errors state.upsert( DocExtractionRecord( atom_id=atom_id, source_url=full_atom.get("source_url", ""), extracted_at=now_iso(), prompt_version=PROMPT_VERSION, raw_count=result.raw_count, valid_count=len(result.valid), pushed_count=pushed, rejected=dict(result.rejected), ) ) state.save() # save after each doc so a crash doesn't lose progress # ============================================================ batch entry async def run_batch( *, limit: int | None = None, only_atom_ids: set[str] | None = None, ) -> BatchStats: if not settings.atomic_token: raise RuntimeError("ATOMIC_TOKEN missing") resolver = TagResolver() if "Type/Claim" not in resolver: raise RuntimeError("taxonomy not seeded — run scripts/04_seed_taxonomy.py") state = ExtractionState() stats = BatchStats() async with AtomicClient() as atomic, LlmClient() as llm: docs = await _list_documents_to_process(atomic, resolver, limit=limit or 1000) if only_atom_ids: docs = [d for d in docs if d["id"] in only_atom_ids] stats.docs_seen = len(docs) log.info("batch_start", docs=len(docs), prompt_version=PROMPT_VERSION) for atom in docs: await process_one( llm=llm, atomic=atomic, resolver=resolver, state=state, atom=atom, stats=stats, ) return stats