1"""Job handlers (registered with `@handler`): LLM extraction of stored snapshots, embeddings, reprocessing, quality recompute."""2from __future__ import annotations34import logging5from datetime import UTC6from typing import Any78from aiatlas.db import execute, fetch_one, transaction9from aiatlas.sdk import archive10from aiatlas.sdk.facts import EntityRef, Facts11from aiatlas.sdk.writer import FactWriter12from aiatlas.schemas import schema_for13from aiatlas.services.jobs import handler14from aiatlas.services.llm import LLMUnavailable, gateway1516log = logging.getLogger(__name__)171819@handler("llm_extract")20async def llm_extract(payload: dict[str, Any], job: dict[str, Any]) -> dict[str, Any] | None:21 """Stage 2–4: run the local LLM on a stored snapshot's cleaned text and write facts with extractor='llm'."""22 if not gateway.available:23 raise LLMUnavailable("LLM not configured; leaving snapshot llm_pending")24 snapshot_id = payload["snapshot_id"]25 async with transaction() as conn:26 snap = await fetch_one(conn, """select s.*, d.doc_type, d.entity_id as doc_entity_id, d.connector_name, d.source_id, e.entity_type, e.canonical_name27 from snapshots s join documents d on d.id = s.document_id left join entities e on e.id = d.entity_id where s.id = :id""", id=snapshot_id)28 if not snap or not snap["text_path"]:29 return {"skipped": "no text"}30 text = archive.load_text(snap["text_path"])31 if len(text) < 200:32 async with transaction() as conn:33 await execute(conn, "update snapshots set processing_status = 'llm_done' where id = :id", id=snapshot_id)34 return {"skipped": "too short"}35 task = payload.get("task") or "auto"36 if task == "auto":37 task = _guess_task(snap["doc_type"], snap["entity_type"])38 if task == "classify_then_extract":39 label = await gateway.classify(text=text, labels=["model_release", "pricing", "research_paper", "company_news", "framework_release", "hardware", "other"],40 snapshot_id=snapshot_id)41 task = {"model_release": "release_announcement", "pricing": "pricing", "research_paper": "paper_passport", "company_news": "release_announcement",42 "framework_release": "release_announcement", "hardware": "hardware_spec"}.get(label or "", "")43 if not task:44 async with transaction() as conn:45 await execute(conn, "update snapshots set processing_status = 'llm_done' where id = :id", id=snapshot_id)46 return {"classified": label, "extracted": False}47 schema, stage = schema_for(task)48 res = await gateway.extract(task_type=task, document=text, schema=schema, stage=stage, snapshot_id=snapshot_id, entity_id=snap["doc_entity_id"], job_id=job["id"])49 if not res.ok or not res.data:50 async with transaction() as conn:51 await execute(conn, "update snapshots set processing_status = 'failed' where id = :id", id=snapshot_id)52 raise RuntimeError(f"llm extraction failed: {res.error}")53 facts = facts_from_llm(task, res.data, entity_id=snap["doc_entity_id"], entity_type=snap["entity_type"], entity_name=snap["canonical_name"], url=snap["url"])54 async with transaction() as conn:55 tier = await fetch_one(conn, "select tier from sources where id = :id", id=snap["source_id"]) if snap["source_id"] else None56 # LLM output never outranks a deterministic statement from the same source: one tier lower (disagreement → conflicting + review)57 writer = FactWriter(conn, source_id=snap["source_id"], snapshot_id=snapshot_id, source_url=snap["url"], tier=min(4, (tier["tier"] if tier else 2) + 1),58 connector_name=snap["connector_name"], extractor="llm", extractor_version=res.model, observed_at=snap["observed_at"].astimezone(UTC))59 ws = await writer.write(facts)60 await execute(conn, "update snapshots set processing_status = 'llm_done' where id = :id", id=snapshot_id)61 return {"task": task, "model": res.model, **ws.as_dict()}626364def _guess_task(doc_type: str | None, entity_type: str | None) -> str:65 if doc_type in ("model_card", "model_page") or entity_type == "model":66 return "model_passport"67 if doc_type == "pricing":68 return "pricing"69 if doc_type in ("paper", "pdf"):70 return "paper_passport"71 if doc_type == "leaderboard":72 return "benchmark_results"73 if entity_type in ("company", "organization"):74 return "company_passport"75 if entity_type == "hardware" or doc_type == "hardware":76 return "hardware_spec"77 if doc_type in ("release", "feed_item", "news"):78 return "release_announcement"79 return "classify_then_extract"808182GENERIC_MODEL_NAMES = {"claude", "gpt", "gemini", "llama", "mistral", "qwen", "deepseek", "grok", "gemma", "phi", "command", "codex", "sora", "veo", "imagen",83 "model", "models", "the model", "new model", "ai model", "llm", "openai", "anthropic", "google", "meta", "microsoft", "nvidia"}84STATUS_ALIASES = {"available": "active", "ga": "active", "generally available": "active", "live": "active", "released": "active", "beta": "preview",85 "experimental": "preview", "coming soon": "announced", "sunset": "retired", "discontinued": "retired", "legacy": "deprecated"}868788def plausible_model_name(name: str | None) -> bool:89 """Reject family/vendor names and vague phrases the LLM sometimes returns as a model name."""90 if not name:91 return False92 n = name.strip()93 low = n.lower()94 if low in GENERIC_MODEL_NAMES or len(n) < 3 or len(n) > 80:95 return False96 return any(ch.isdigit() for ch in n) or len(n.split()) >= 2 or "-" in n979899def _norm_status(v: Any) -> Any:100 return STATUS_ALIASES.get(v.strip().lower(), v.strip().lower()) if isinstance(v, str) else v101102103def facts_from_llm(task: str, data: dict[str, Any], *, entity_id: str | None, entity_type: str | None, entity_name: str | None, url: str) -> Facts:104 """Map a validated LLM output onto facts. LLM claims default to 'medium' confidence and are written one tier below their source,105 so they never outrank deterministic statements (the writer stores disagreements as conflicting)."""106 facts = Facts()107 conf = "medium"108 if "status" in data:109 data = {**data, "status": _norm_status(data.get("status"))}110111 def model_ref(name: str, developer: str | None = None) -> EntityRef:112 org = facts.entity("company", developer) if developer else None113 return facts.entity("model", name, organization=org)114115 if task == "model_passport":116 name = data.get("name") or entity_name117 if not name or (not (entity_id and entity_type == "model") and not plausible_model_name(name)):118 return facts119 ref = EntityRef(entity_type="model", name=name, id=entity_id) if entity_id and entity_type == "model" else model_ref(name, data.get("developer"))120 if ref not in facts.entities:121 facts.entities.append(ref)122 for prop in ("family", "version", "release_date", "status", "openness", "license", "architecture", "parameter_count", "active_parameter_count", "is_moe",123 "context_length", "max_output_tokens", "knowledge_cutoff", "tool_calling", "structured_output", "reasoning", "vision", "audio",124 "fine_tuning_available", "tokenizer", "hardware_requirements", "safety_notes", "training_data_notes"):125 facts.claim(ref, prop, data.get(prop), confidence=conf)126 mods = sorted(set((data.get("modalities_input") or []) + (data.get("modalities_output") or [])))127 facts.claim(ref, "modalities", mods, confidence=conf)128 facts.claim(ref, "modalities_input", data.get("modalities_input"), confidence=conf)129 facts.claim(ref, "modalities_output", data.get("modalities_output"), confidence=conf)130 facts.claim(ref, "languages", data.get("languages"), confidence=conf)131 for key, prop in (("paper_url", "paper_url"), ("model_card_url", "model_card_url"), ("repository_url", "repository_url"), ("official_page_url", "official_url")):132 facts.claim(ref, prop, data.get(key), confidence=conf)133 if data.get("developer") and not (entity_id and entity_type == "model"):134 pass135 elif data.get("developer"):136 org = facts.entity("company", data["developer"])137 facts.relate(org, "develops", ref, confidence=conf)138 if data.get("base_model"):139 base = facts.entity("model", data["base_model"])140 facts.relate(ref, "derived_from", base, confidence="low")141 if data.get("predecessor"):142 pred = facts.entity("model", data["predecessor"])143 facts.relate(pred, "superseded_by", ref, confidence="low")144 facts.document_entity = ref145 elif task == "pricing":146 provider_name = data.get("provider")147 if not provider_name:148 return facts149 provider = facts.entity("provider", provider_name)150 for line in data.get("prices") or []:151 if not line.get("model"):152 continue153 model = facts.entity("model", line["model"])154 facts.price(model=model, provider=provider, provider_model_id=line.get("provider_model_id"), input_per_mtok=line.get("input_per_mtok"),155 output_per_mtok=line.get("output_per_mtok"), cached_input_per_mtok=line.get("cached_input_per_mtok"),156 cache_write_per_mtok=line.get("cache_write_per_mtok"), batch_input_per_mtok=line.get("batch_input_per_mtok"),157 batch_output_per_mtok=line.get("batch_output_per_mtok"), per_image=line.get("per_image"), currency=data.get("currency") or "USD",158 context_length=line.get("context_length"), max_output_tokens=line.get("max_output_tokens"), meta={"extractor": "llm", "notes": line.get("notes")})159 elif task == "company_passport":160 name = data.get("name") or entity_name161 if not name:162 return facts163 ref = EntityRef(entity_type="company", name=name, id=entity_id) if entity_id and entity_type in ("company", "organization") else facts.entity("company", name)164 if ref not in facts.entities:165 facts.entities.append(ref)166 for prop in ("legal_name", "country", "headquarters", "founded", "founders", "leadership", "website", "employee_count", "funding_total_usd"):167 facts.claim(ref, prop, data.get(prop), confidence=conf)168 if data.get("description"):169 facts.claim(ref, "description", data["description"], confidence=conf)170 for m in data.get("models") or []:171 if plausible_model_name(m):172 facts.relate(ref, "develops", facts.entity("model", m, organization=ref), confidence="low")173 for inv in data.get("investors") or []:174 facts.relate(ref, "funded_by", facts.entity("company", inv), confidence="low")175 if data.get("parent_company"):176 facts.relate(facts.entity("company", data["parent_company"]), "owns", ref, confidence="low")177 facts.document_entity = ref178 elif task == "paper_passport":179 title = data.get("title") or entity_name180 if not title:181 return facts182 ref = EntityRef(entity_type="paper", name=title, id=entity_id) if entity_id and entity_type == "paper" else facts.entity("paper", title)183 if ref not in facts.entities:184 facts.entities.append(ref)185 for prop in ("authors", "affiliations", "date", "field", "summary", "methods", "key_claims", "results", "limitations", "code_url"):186 facts.claim(ref, prop, data.get(prop), confidence=conf)187 for m in data.get("models") or []:188 if plausible_model_name(m):189 facts.relate(facts.entity("model", m), "described_by", ref, confidence="low")190 for d in data.get("datasets") or []:191 facts.relate(ref, "uses_dataset", facts.entity("dataset", d), confidence="low")192 for b in data.get("benchmarks") or []:193 facts.relate(ref, "evaluates_on", facts.entity("benchmark", b), confidence="low")194 facts.document_entity = ref195 elif task == "benchmark_results":196 bname = data.get("benchmark") or entity_name197 if not bname:198 return facts199 bench = facts.entity("benchmark", bname)200 for row in data.get("rows") or []:201 if row.get("model") and row.get("score") is not None:202 facts.result(model=facts.entity("model", row["model"]), benchmark=bench, score=float(row["score"]), metric=row.get("metric") or data.get("metric"),203 higher_is_better=bool(data.get("higher_is_better", True)), config={"config": row.get("config"), "extractor": "llm"}, confidence="low")204 elif task == "hardware_spec":205 name = data.get("name") or entity_name206 if not name:207 return facts208 org = facts.entity("company", data["manufacturer"]) if data.get("manufacturer") else None209 ref = EntityRef(entity_type="hardware", name=name, id=entity_id) if entity_id and entity_type == "hardware" else facts.entity("hardware", name, organization=org)210 if ref not in facts.entities:211 facts.entities.append(ref)212 for prop in ("kind", "architecture", "release_date", "memory_gb", "memory_type", "memory_bandwidth_gbs", "compute_fp16_tflops", "compute_fp8_tflops",213 "compute_int8_tops", "tdp_watts", "form_factor", "price_usd", "interconnect"):214 facts.claim(ref, prop, data.get(prop), confidence=conf)215 if org:216 facts.relate(org, "manufactures", ref, confidence=conf)217 facts.document_entity = ref218 elif task == "release_announcement":219 org = facts.entity("company", data["organization"]) if data.get("organization") else None220 for m in data.get("models") or []:221 if not plausible_model_name(m):222 continue223 ref = facts.entity("model", m, organization=org)224 if org:225 facts.relate(org, "develops", ref, confidence="low")226 if data.get("model_passport"):227 facts.extend(facts_from_llm("model_passport", data["model_passport"], entity_id=None, entity_type=None, entity_name=None, url=url))228 if data.get("pricing"):229 facts.extend(facts_from_llm("pricing", data["pricing"], entity_id=None, entity_type=None, entity_name=None, url=url))230 return facts231232233@handler("embed_entity")234async def embed_entity(payload: dict[str, Any], job: dict[str, Any]) -> dict[str, Any] | None:235 from aiatlas.services.embeddings import embed_entities236237 return await embed_entities(payload.get("entity_ids") or [payload["entity_id"]])238239240@handler("reprocess_snapshot")241async def reprocess_snapshot(payload: dict[str, Any], job: dict[str, Any]) -> dict[str, Any] | None:242 from aiatlas.connectors import get243244 connector = get(payload["connector"])245 ctx = await connector.run(reprocess=True, only_urls=[payload["url"]] if payload.get("url") else None, force=True)246 return {k: v for k, v in ctx.stats.__dict__.items() if k != "meta"}247248249@handler("recompute_quality")250async def recompute_quality(payload: dict[str, Any], job: dict[str, Any]) -> dict[str, Any] | None:251 from aiatlas.services.quality import recompute252253 return await recompute(entity_ids=payload.get("entity_ids"))254255256__all__ = ["llm_extract", "facts_from_llm", "embed_entity", "reprocess_snapshot", "recompute_quality"]257