HTML 77.2%
TypeScript 10.5%
Python 9.6%
JavaScript 2.5%
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.schemas import schema_for10from aiatlas.sdk import archive11from aiatlas.sdk.facts import EntityRef, Facts12from aiatlas.sdk.writer import FactWriter13from 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 run_id=snap["run_id"] or job["id"])60 ws = await writer.write(facts)61 await execute(conn, "update snapshots set processing_status = 'llm_done' where id = :id", id=snapshot_id)62 return {"task": task, "model": res.model, **ws.as_dict()}636465def _guess_task(doc_type: str | None, entity_type: str | None) -> str:66 if doc_type in ("model_card", "model_page") or entity_type == "model":67 return "model_passport"68 if doc_type == "pricing":69 return "pricing"70 if doc_type in ("paper", "pdf"):71 return "paper_passport"72 if doc_type == "leaderboard":73 return "benchmark_results"74 if entity_type in ("company", "organization"):75 return "company_passport"76 if entity_type == "hardware" or doc_type == "hardware":77 return "hardware_spec"78 if doc_type in ("release", "feed_item", "news"):79 return "release_announcement"80 return "classify_then_extract"818283GENERIC_MODEL_NAMES = {"claude", "gpt", "gemini", "llama", "mistral", "qwen", "deepseek", "grok", "gemma", "phi", "command", "codex", "sora", "veo", "imagen",84 "model", "models", "the model", "new model", "ai model", "llm", "openai", "anthropic", "google", "meta", "microsoft", "nvidia"}85STATUS_ALIASES = {"available": "active", "ga": "active", "generally available": "active", "live": "active", "released": "active", "beta": "preview",86 "experimental": "preview", "coming soon": "announced", "sunset": "retired", "discontinued": "retired", "legacy": "deprecated"}878889def plausible_model_name(name: str | None) -> bool:90 """Reject family/vendor names and vague phrases the LLM sometimes returns as a model name."""91 if not name:92 return False93 n = name.strip()94 low = n.lower()95 if low in GENERIC_MODEL_NAMES or len(n) < 3 or len(n) > 80:96 return False97 return any(ch.isdigit() for ch in n) or len(n.split()) >= 2 or "-" in n9899100def _norm_status(v: Any) -> Any:101 return STATUS_ALIASES.get(v.strip().lower(), v.strip().lower()) if isinstance(v, str) else v102103104def facts_from_llm(task: str, data: dict[str, Any], *, entity_id: str | None, entity_type: str | None, entity_name: str | None, url: str) -> Facts:105 """Map a validated LLM output onto facts. LLM claims default to 'medium' confidence and are written one tier below their source,106 so they never outrank deterministic statements (the writer stores disagreements as conflicting)."""107 facts = Facts()108 conf = "medium"109 if "status" in data:110 data = {**data, "status": _norm_status(data.get("status"))}111112 def model_ref(name: str, developer: str | None = None) -> EntityRef:113 org = facts.entity("company", developer) if developer else None114 return facts.entity("model", name, organization=org)115116 if task == "model_passport":117 name = data.get("name") or entity_name118 if not name or (not (entity_id and entity_type == "model") and not plausible_model_name(name)):119 return facts120 ref = EntityRef(entity_type="model", name=name, id=entity_id) if entity_id and entity_type == "model" else model_ref(name, data.get("developer"))121 if ref not in facts.entities:122 facts.entities.append(ref)123 for prop in ("family", "version", "release_date", "status", "openness", "license", "architecture", "parameter_count", "active_parameter_count", "is_moe",124 "context_length", "max_output_tokens", "knowledge_cutoff", "tool_calling", "structured_output", "reasoning", "vision", "audio",125 "fine_tuning_available", "tokenizer", "hardware_requirements", "safety_notes", "training_data_notes"):126 facts.claim(ref, prop, data.get(prop), confidence=conf)127 mods = sorted(set((data.get("modalities_input") or []) + (data.get("modalities_output") or [])))128 facts.claim(ref, "modalities", mods, confidence=conf)129 facts.claim(ref, "modalities_input", data.get("modalities_input"), confidence=conf)130 facts.claim(ref, "modalities_output", data.get("modalities_output"), confidence=conf)131 facts.claim(ref, "languages", data.get("languages"), confidence=conf)132 for key, prop in (("paper_url", "paper_url"), ("model_card_url", "model_card_url"), ("repository_url", "repository_url"), ("official_page_url", "official_url")):133 facts.claim(ref, prop, data.get(key), confidence=conf)134 if data.get("developer") and not (entity_id and entity_type == "model"):135 pass136 elif data.get("developer"):137 org = facts.entity("company", data["developer"])138 facts.relate(org, "develops", ref, confidence=conf)139 if data.get("base_model"):140 base = facts.entity("model", data["base_model"])141 facts.relate(ref, "derived_from", base, confidence="low")142 if data.get("predecessor"):143 pred = facts.entity("model", data["predecessor"])144 facts.relate(pred, "superseded_by", ref, confidence="low")145 facts.document_entity = ref146 elif task == "pricing":147 provider_name = data.get("provider")148 if not provider_name:149 return facts150 provider = facts.entity("provider", provider_name)151 for line in data.get("prices") or []:152 if not line.get("model"):153 continue154 model = facts.entity("model", line["model"])155 facts.price(model=model, provider=provider, provider_model_id=line.get("provider_model_id"), input_per_mtok=line.get("input_per_mtok"),156 output_per_mtok=line.get("output_per_mtok"), cached_input_per_mtok=line.get("cached_input_per_mtok"),157 cache_write_per_mtok=line.get("cache_write_per_mtok"), batch_input_per_mtok=line.get("batch_input_per_mtok"),158 batch_output_per_mtok=line.get("batch_output_per_mtok"), per_image=line.get("per_image"), currency=data.get("currency") or "USD",159 context_length=line.get("context_length"), max_output_tokens=line.get("max_output_tokens"), meta={"extractor": "llm", "notes": line.get("notes")})160 elif task == "company_passport":161 name = data.get("name") or entity_name162 if not name:163 return facts164 ref = EntityRef(entity_type="company", name=name, id=entity_id) if entity_id and entity_type in ("company", "organization") else facts.entity("company", name)165 if ref not in facts.entities:166 facts.entities.append(ref)167 for prop in ("legal_name", "country", "headquarters", "founded", "founders", "leadership", "website", "employee_count", "funding_total_usd"):168 facts.claim(ref, prop, data.get(prop), confidence=conf)169 if data.get("description"):170 facts.claim(ref, "description", data["description"], confidence=conf)171 for m in data.get("models") or []:172 if plausible_model_name(m):173 facts.relate(ref, "develops", facts.entity("model", m, organization=ref), confidence="low")174 for inv in data.get("investors") or []:175 facts.relate(ref, "funded_by", facts.entity("company", inv), confidence="low")176 if data.get("parent_company"):177 facts.relate(facts.entity("company", data["parent_company"]), "owns", ref, confidence="low")178 facts.document_entity = ref179 elif task == "paper_passport":180 title = data.get("title") or entity_name181 if not title:182 return facts183 ref = EntityRef(entity_type="paper", name=title, id=entity_id) if entity_id and entity_type == "paper" else facts.entity("paper", title)184 if ref not in facts.entities:185 facts.entities.append(ref)186 for prop in ("authors", "affiliations", "date", "field", "summary", "methods", "key_claims", "results", "limitations", "code_url"):187 facts.claim(ref, prop, data.get(prop), confidence=conf)188 for m in data.get("models") or []:189 if plausible_model_name(m):190 facts.relate(facts.entity("model", m), "described_by", ref, confidence="low")191 for d in data.get("datasets") or []:192 facts.relate(ref, "uses_dataset", facts.entity("dataset", d), confidence="low")193 for b in data.get("benchmarks") or []:194 facts.relate(ref, "evaluates_on", facts.entity("benchmark", b), confidence="low")195 facts.document_entity = ref196 elif task == "benchmark_results":197 bname = data.get("benchmark") or entity_name198 if not bname:199 return facts200 bench = facts.entity("benchmark", bname)201 for row in data.get("rows") or []:202 if row.get("model") and row.get("score") is not None:203 facts.result(model=facts.entity("model", row["model"]), benchmark=bench, score=float(row["score"]), metric=row.get("metric") or data.get("metric"),204 higher_is_better=bool(data.get("higher_is_better", True)), config={"config": row.get("config"), "extractor": "llm"}, confidence="low")205 elif task == "hardware_spec":206 name = data.get("name") or entity_name207 if not name:208 return facts209 org = facts.entity("company", data["manufacturer"]) if data.get("manufacturer") else None210 ref = EntityRef(entity_type="hardware", name=name, id=entity_id) if entity_id and entity_type == "hardware" else facts.entity("hardware", name, organization=org)211 if ref not in facts.entities:212 facts.entities.append(ref)213 for prop in ("kind", "architecture", "release_date", "memory_gb", "memory_type", "memory_bandwidth_gbs", "compute_fp16_tflops", "compute_fp8_tflops",214 "compute_int8_tops", "tdp_watts", "form_factor", "price_usd", "interconnect"):215 facts.claim(ref, prop, data.get(prop), confidence=conf)216 if org:217 facts.relate(org, "manufactures", ref, confidence=conf)218 facts.document_entity = ref219 elif task == "release_announcement":220 org = facts.entity("company", data["organization"]) if data.get("organization") else None221 for m in data.get("models") or []:222 if not plausible_model_name(m):223 continue224 ref = facts.entity("model", m, organization=org)225 if org:226 facts.relate(org, "develops", ref, confidence="low")227 if data.get("model_passport"):228 facts.extend(facts_from_llm("model_passport", data["model_passport"], entity_id=None, entity_type=None, entity_name=None, url=url))229 if data.get("pricing"):230 facts.extend(facts_from_llm("pricing", data["pricing"], entity_id=None, entity_type=None, entity_name=None, url=url))231 return facts232233234@handler("embed_entity")235async def embed_entity(payload: dict[str, Any], job: dict[str, Any]) -> dict[str, Any] | None:236 from aiatlas.services.embeddings import embed_entities237238 return await embed_entities(payload.get("entity_ids") or [payload["entity_id"]])239240241@handler("reprocess_snapshot")242async def reprocess_snapshot(payload: dict[str, Any], job: dict[str, Any]) -> dict[str, Any] | None:243 from aiatlas.connectors import get244245 connector = get(payload["connector"])246 ctx = await connector.run(reprocess=True, only_urls=[payload["url"]] if payload.get("url") else None, force=True)247 return {k: v for k, v in ctx.stats.__dict__.items() if k != "meta"}248249250@handler("recompute_quality")251async def recompute_quality(payload: dict[str, Any], job: dict[str, Any]) -> dict[str, Any] | None:252 from aiatlas.services.quality import recompute253254 return await recompute(entity_ids=payload.get("entity_ids"))255256257__all__ = ["embed_entity", "facts_from_llm", "llm_extract", "recompute_quality", "reprocess_snapshot"]258