233 lines
7.1 KiB
Python
233 lines
7.1 KiB
Python
"""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=<Type/Document> 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/<X> 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
|