Livrare LOT 1 - Didi

ai_platform/modules/didi_brain/extractor/batch.py

@@ -0,0 +1,233 @@
+"""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