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 writer = FactWriter(conn, source_id=snap["source_id"], snapshot_id=snapshot_id, source_url=snap["url"], tier=(tier["tier"] if tier else 2),57 connector_name=snap["connector_name"], extractor="llm", extractor_version=res.model, observed_at=snap["observed_at"].astimezone(UTC))58 ws = await writer.write(facts)59 await execute(conn, "update snapshots set processing_status = 'llm_done' where id = :id", id=snapshot_id)60 return {"task": task, "model": res.model, **ws.as_dict()}616263def _guess_task(doc_type: str | None, entity_type: str | None) -> str:64 if doc_type in ("model_card", "model_page") or entity_type == "model":65 return "model_passport"66 if doc_type == "pricing":67 return "pricing"68 if doc_type in ("paper", "pdf"):69 return "paper_passport"70 if doc_type == "leaderboard":71 return "benchmark_results"72 if entity_type in ("company", "organization"):73 return "company_passport"74 if entity_type == "hardware" or doc_type == "hardware":75 return "hardware_spec"76 if doc_type in ("release", "feed_item", "news"):77 return "release_announcement"78 return "classify_then_extract"798081def facts_from_llm(task: str, data: dict[str, Any], *, entity_id: str | None, entity_type: str | None, entity_name: str | None, url: str) -> Facts:82 """Map a validated LLM output onto facts. LLM claims default to 'medium' confidence and never outrank tier-1 deterministic ones83 (the writer stores disagreements as conflicting)."""84 facts = Facts()85 conf = "medium"8687 def model_ref(name: str, developer: str | None = None) -> EntityRef:88 org = facts.entity("company", developer) if developer else None89 return facts.entity("model", name, organization=org)9091 if task == "model_passport":92 name = data.get("name") or entity_name93 if not name:94 return facts95 ref = EntityRef(entity_type="model", name=name, id=entity_id) if entity_id and entity_type == "model" else model_ref(name, data.get("developer"))96 if ref not in facts.entities:97 facts.entities.append(ref)98 for prop in ("family", "version", "release_date", "status", "openness", "license", "architecture", "parameter_count", "active_parameter_count", "is_moe",99 "context_length", "max_output_tokens", "knowledge_cutoff", "tool_calling", "structured_output", "reasoning", "vision", "audio",100 "fine_tuning_available", "tokenizer", "hardware_requirements", "safety_notes", "training_data_notes"):101 facts.claim(ref, prop, data.get(prop), confidence=conf)102 mods = sorted(set((data.get("modalities_input") or []) + (data.get("modalities_output") or [])))103 facts.claim(ref, "modalities", mods, confidence=conf)104 facts.claim(ref, "modalities_input", data.get("modalities_input"), confidence=conf)105 facts.claim(ref, "modalities_output", data.get("modalities_output"), confidence=conf)106 facts.claim(ref, "languages", data.get("languages"), confidence=conf)107 for key, prop in (("paper_url", "paper_url"), ("model_card_url", "model_card_url"), ("repository_url", "repository_url"), ("official_page_url", "official_url")):108 facts.claim(ref, prop, data.get(key), confidence=conf)109 if data.get("developer") and not (entity_id and entity_type == "model"):110 pass111 elif data.get("developer"):112 org = facts.entity("company", data["developer"])113 facts.relate(org, "develops", ref, confidence=conf)114 if data.get("base_model"):115 base = facts.entity("model", data["base_model"])116 facts.relate(ref, "derived_from", base, confidence="low")117 if data.get("predecessor"):118 pred = facts.entity("model", data["predecessor"])119 facts.relate(pred, "superseded_by", ref, confidence="low")120 facts.document_entity = ref121 elif task == "pricing":122 provider_name = data.get("provider")123 if not provider_name:124 return facts125 provider = facts.entity("provider", provider_name)126 for line in data.get("prices") or []:127 if not line.get("model"):128 continue129 model = facts.entity("model", line["model"])130 facts.price(model=model, provider=provider, provider_model_id=line.get("provider_model_id"), input_per_mtok=line.get("input_per_mtok"),131 output_per_mtok=line.get("output_per_mtok"), cached_input_per_mtok=line.get("cached_input_per_mtok"),132 cache_write_per_mtok=line.get("cache_write_per_mtok"), batch_input_per_mtok=line.get("batch_input_per_mtok"),133 batch_output_per_mtok=line.get("batch_output_per_mtok"), per_image=line.get("per_image"), currency=data.get("currency") or "USD",134 context_length=line.get("context_length"), max_output_tokens=line.get("max_output_tokens"), meta={"extractor": "llm", "notes": line.get("notes")})135 elif task == "company_passport":136 name = data.get("name") or entity_name137 if not name:138 return facts139 ref = EntityRef(entity_type="company", name=name, id=entity_id) if entity_id and entity_type in ("company", "organization") else facts.entity("company", name)140 if ref not in facts.entities:141 facts.entities.append(ref)142 for prop in ("legal_name", "country", "headquarters", "founded", "founders", "leadership", "website", "employee_count", "funding_total_usd"):143 facts.claim(ref, prop, data.get(prop), confidence=conf)144 if data.get("description"):145 facts.claim(ref, "description", data["description"], confidence=conf)146 for m in data.get("models") or []:147 facts.relate(ref, "develops", facts.entity("model", m, organization=ref), confidence="low")148 for inv in data.get("investors") or []:149 facts.relate(ref, "funded_by", facts.entity("company", inv), confidence="low")150 if data.get("parent_company"):151 facts.relate(facts.entity("company", data["parent_company"]), "owns", ref, confidence="low")152 facts.document_entity = ref153 elif task == "paper_passport":154 title = data.get("title") or entity_name155 if not title:156 return facts157 ref = EntityRef(entity_type="paper", name=title, id=entity_id) if entity_id and entity_type == "paper" else facts.entity("paper", title)158 if ref not in facts.entities:159 facts.entities.append(ref)160 for prop in ("authors", "affiliations", "date", "field", "summary", "methods", "key_claims", "results", "limitations", "code_url"):161 facts.claim(ref, prop, data.get(prop), confidence=conf)162 for m in data.get("models") or []:163 facts.relate(facts.entity("model", m), "described_by", ref, confidence="low")164 for d in data.get("datasets") or []:165 facts.relate(ref, "uses_dataset", facts.entity("dataset", d), confidence="low")166 for b in data.get("benchmarks") or []:167 facts.relate(ref, "evaluates_on", facts.entity("benchmark", b), confidence="low")168 facts.document_entity = ref169 elif task == "benchmark_results":170 bname = data.get("benchmark") or entity_name171 if not bname:172 return facts173 bench = facts.entity("benchmark", bname)174 for row in data.get("rows") or []:175 if row.get("model") and row.get("score") is not None:176 facts.result(model=facts.entity("model", row["model"]), benchmark=bench, score=float(row["score"]), metric=row.get("metric") or data.get("metric"),177 higher_is_better=bool(data.get("higher_is_better", True)), config={"config": row.get("config"), "extractor": "llm"}, confidence="low")178 elif task == "hardware_spec":179 name = data.get("name") or entity_name180 if not name:181 return facts182 org = facts.entity("company", data["manufacturer"]) if data.get("manufacturer") else None183 ref = EntityRef(entity_type="hardware", name=name, id=entity_id) if entity_id and entity_type == "hardware" else facts.entity("hardware", name, organization=org)184 if ref not in facts.entities:185 facts.entities.append(ref)186 for prop in ("kind", "architecture", "release_date", "memory_gb", "memory_type", "memory_bandwidth_gbs", "compute_fp16_tflops", "compute_fp8_tflops",187 "compute_int8_tops", "tdp_watts", "form_factor", "price_usd", "interconnect"):188 facts.claim(ref, prop, data.get(prop), confidence=conf)189 if org:190 facts.relate(org, "manufactures", ref, confidence=conf)191 facts.document_entity = ref192 elif task == "release_announcement":193 org = facts.entity("company", data["organization"]) if data.get("organization") else None194 for m in data.get("models") or []:195 ref = facts.entity("model", m, organization=org)196 if org:197 facts.relate(org, "develops", ref, confidence="low")198 if data.get("model_passport"):199 facts.extend(facts_from_llm("model_passport", data["model_passport"], entity_id=None, entity_type=None, entity_name=None, url=url))200 if data.get("pricing"):201 facts.extend(facts_from_llm("pricing", data["pricing"], entity_id=None, entity_type=None, entity_name=None, url=url))202 return facts203204205@handler("embed_entity")206async def embed_entity(payload: dict[str, Any], job: dict[str, Any]) -> dict[str, Any] | None:207 from aiatlas.services.embeddings import embed_entities208209 return await embed_entities(payload.get("entity_ids") or [payload["entity_id"]])210211212@handler("reprocess_snapshot")213async def reprocess_snapshot(payload: dict[str, Any], job: dict[str, Any]) -> dict[str, Any] | None:214 from aiatlas.connectors import get215216 connector = get(payload["connector"])217 ctx = await connector.run(reprocess=True, only_urls=[payload["url"]] if payload.get("url") else None, force=True)218 return {k: v for k, v in ctx.stats.__dict__.items() if k != "meta"}219220221@handler("recompute_quality")222async def recompute_quality(payload: dict[str, Any], job: dict[str, Any]) -> dict[str, Any] | None:223 from aiatlas.services.quality import recompute224225 return await recompute(entity_ids=payload.get("entity_ids"))226227228__all__ = ["llm_extract", "facts_from_llm", "embed_entity", "reprocess_snapshot", "recompute_quality"]229