"""LLM enrichment worker (spec §24, §2.4): drains `llm_jobs` with SKIP LOCKED, sticky by kind (one model per stream so the on-demand server does not thrash), builds a bounded context from the change diff and structured deltas, calls the gateway, validates the JSON and writes events / summaries with full provenance (`origin`, model, prompt version, schema version, tokens, latency). Job kinds - classify_change (ref = change id, small model) → new event `origin='llm'` when material, else recorded as non-material. - summarize_event (ref = event id, medium model) → `events.summary` (+ payload.llm), `origin='hybrid'`; LEGAL events use the legal-diff prompt. - classify_industry (ref = company id, small model) → `companies.source_meta.llm_industries` suggestion only (never overwrites registry data). Graceful when `settings.llm_configured` is False: jobs stay pending, the worker logs once and returns. """ from __future__ import annotations import logging from datetime import UTC, datetime from typing import Any from pydantic import BaseModel from companyatlas.config import settings from companyatlas.db import execute, fetch_all, fetch_one, fetch_val, jsonb, transaction from companyatlas.ids import new_id, stable_hash from companyatlas.services.clustering import attach_to_cluster from companyatlas.services.llm.gateway import LLMError, LLMNotConfigured, LLMValidationError, get_provider from companyatlas.services.llm.prompts import load_prompt from companyatlas.services.llm.schemas import SCHEMA_VERSIONS, ChangeClassification, EventSummary, IndustryTags, LegalDiffSummary from companyatlas.services.periodic import periodic from companyatlas.taxonomy import EVENT_SUBTYPES, EventType, confidence_label log = logging.getLogger(__name__) PROVIDER_NAME = "openai-compatible" KIND_ORDER = ("classify_change", "summarize_event", "classify_industry") MAX_BLOCKS = 12 MATERIAL_MIN_IMPORTANCE = 0.25 _sticky: dict[str, str | None] = {"kind": None} _warned = {"unconfigured": False} # ---------------------------------------------------------------------------------------------------------------- context def _cut(text: str | None, limit: int) -> str | None: if not text: return None b = text.encode("utf-8") return text if len(b) <= limit else b[:limit].decode("utf-8", "ignore") + "…" def build_change_context(change: dict[str, Any], company: dict[str, Any], sensor: dict[str, Any]) -> dict[str, Any]: """Bounded, structured context: company meta, surface, ≤ 12 blocks (≤ 3 kB each side), structured deltas (trimmed lists).""" limit = settings.llm_context_block_bytes diff = change.get("diff") or {} blocks: list[dict[str, Any]] = [] for bucket in ("added", "removed", "modified"): for d in (diff.get(bucket) or [])[:MAX_BLOCKS]: if len(blocks) >= MAX_BLOCKS: break blocks.append({"op": bucket, "kind": d.get("kind"), "path": d.get("path"), "before": _cut(d.get("before"), limit), "after": _cut(d.get("after"), limit)}) delta = change.get("structured_delta") or {} trimmed: dict[str, Any] = {} for key, value in delta.items(): if isinstance(value, dict): trimmed[key] = {k: (v[:10] if isinstance(v, list) else v) for k, v in value.items()} else: trimmed[key] = value return { "company": {"name": company.get("display_name"), "domain": company.get("canonical_domain"), "country": company.get("country"), "industries": list(company.get("industries") or [])[:5], "description": _cut(company.get("description"), 400)}, "surface": change.get("surface") or sensor.get("surface"), "source_url": sensor.get("url"), "detected_at": str(change.get("detected_at")), "significance": change.get("significance"), "change_kind": change.get("kind"), "counts": diff.get("counts") or {"added": change.get("blocks_added"), "removed": change.get("blocks_removed"), "modified": change.get("blocks_modified")}, "blocks": blocks, "structured_delta": trimmed, } async def _load_change_bundle(conn, change_id: str) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]] | None: # type: ignore[no-untyped-def] change = await fetch_one(conn, "select * from changes where id = :id", id=change_id) if change is None: return None sensor = await fetch_one(conn, "select id, url, surface, connector_id from sensors where id = :id", id=change["sensor_id"]) or {} company = await fetch_one(conn, "select id, slug, display_name, canonical_domain, country, industries, description from companies where id = :id", id=change["company_id"]) or {} return change, company, sensor # ---------------------------------------------------------------------------------------------------------------- claiming async def _budget_left(conn) -> int: # type: ignore[no-untyped-def] used = await fetch_val(conn, "select count(*) from llm_jobs where finished_at >= date_trunc('day', now() at time zone 'utc') and status in ('done', 'failed')") return max(0, settings.llm_daily_budget - int(used or 0)) async def _pick_kind(conn) -> str | None: # type: ignore[no-untyped-def] rows = await fetch_all(conn, "select kind, count(*) as n, min(created_at) as oldest from llm_jobs where status = 'pending' group by kind") if not rows: return None pending = {r["kind"]: r for r in rows} if _sticky["kind"] in pending: return _sticky["kind"] ordered = sorted(rows, key=lambda r: (KIND_ORDER.index(r["kind"]) if r["kind"] in KIND_ORDER else 99, r["oldest"])) _sticky["kind"] = ordered[0]["kind"] return _sticky["kind"] async def claim_jobs(conn, kind: str, limit: int) -> list[dict[str, Any]]: # type: ignore[no-untyped-def] return await fetch_all(conn, """ update llm_jobs set status = 'running', started_at = now(), attempts = attempts + 1 where id in (select id from llm_jobs where status = 'pending' and kind = :kind order by created_at limit :limit for update skip locked) returning *""", kind=kind, limit=limit) async def _finish(conn, job: dict[str, Any], *, status: str, model: str | None, prompt_version: str | None, result: dict[str, Any] | None, # type: ignore[no-untyped-def] error: str | None, req: int = 0, resp: int = 0, latency: int = 0) -> None: await execute(conn, """ update llm_jobs set status = :status, model = :model, prompt_version = :pv, result = cast(:result as jsonb), error = :error, request_tokens = coalesce(request_tokens, 0) + :req, response_tokens = coalesce(response_tokens, 0) + :resp, latency_ms = :latency, finished_at = case when :status in ('done', 'failed') then now() else finished_at end where id = :id""", status=status, model=model, pv=prompt_version, result=jsonb(result) if result is not None else None, error=error, req=req, resp=resp, latency=latency, id=job["id"]) if model and (req or resp): await execute(conn, """ insert into cost_ledger (day, dimension, key, units, cost_estimate) values (current_date, 'llm', :key, :units, 0) on conflict (day, dimension, key) do update set units = cost_ledger.units + excluded.units""", key=model, units=float(req + resp)) # ---------------------------------------------------------------------------------------------------------------- handlers async def _handle_classify_change(conn, job: dict[str, Any]) -> tuple[dict[str, Any], str, int, int, int]: # type: ignore[no-untyped-def] bundle = await _load_change_bundle(conn, job["ref_id"]) if bundle is None: raise LLMError("change not found") change, company, sensor = bundle prompt = load_prompt("change-classifier") context = build_change_context(change, company, sensor) res = await get_provider().complete_json("small", prompt.system, jsonb(context), ChangeClassification, max_tokens=700) c: ChangeClassification = res.data result: dict[str, Any] = {"schema_version": SCHEMA_VERSIONS["ChangeClassification"], "classification": c.model_dump(), "repaired": res.repaired} material = c.is_material and c.event_subtype != "OTHER" and c.importance >= MATERIAL_MIN_IMPORTANCE if material: event_id = await _insert_llm_event(conn, change, company, sensor, c, model=res.model, prompt_version=prompt.ref) result["event_id"] = event_id await execute(conn, "update changes set status = 'enriched' where id = :id", id=change["id"]) return result, res.model, res.request_tokens, res.response_tokens, res.latency_ms async def _insert_llm_event(conn, change: dict[str, Any], company: dict[str, Any], sensor: dict[str, Any], c: ChangeClassification, *, # type: ignore[no-untyped-def] model: str, prompt_version: str) -> str | None: event_type = str(EVENT_SUBTYPES.get(c.event_subtype, (EventType.OTHER, 0.3))[0]) # A deterministic event of the same subtype already exists for this change → enrich it instead of adding a near-duplicate. existing = await fetch_one(conn, """select id, summary from events where change_id = :c and event_subtype = :st and origin in ('deterministic', 'hybrid') and status in ('active', 'review') order by created_at limit 1""", c=change["id"], st=c.event_subtype) if existing is not None: await execute(conn, """update events set summary = coalesce(summary, :summary), origin = 'hybrid', model_provider = :provider, model_name = :model, prompt_version = :pv, payload = payload || cast(:extra as jsonb) where id = :id""", summary=c.summary, provider=PROVIDER_NAME, model=model, pv=prompt_version, id=existing["id"], extra=jsonb({"llm_classification": {"title": c.title, "importance": c.importance, "confidence": c.confidence, "tags": c.tags}})) return None dedupe = stable_hash(company["id"], "llm", change["id"], c.event_subtype, length=40) event_id = new_id("event") detected_at = change.get("detected_at") or datetime.now(UTC) payload = {"llm": {"importance": c.importance, "confidence": c.confidence, "is_material": c.is_material}, "significance": change.get("significance"), "change_kind": change.get("kind"), "entity_key": c.title} row = await fetch_one(conn, """ insert into events (id, company_id, sensor_id, change_id, surface, event_type, event_subtype, importance, confidence, confidence_label, title, summary, old_value, new_value, payload, entities, tags, detected_at, source_url, snapshot_before, snapshot_after, language, origin, model_provider, model_name, prompt_version, schema_version, status, dedupe_key) values (:id, :company_id, :sensor_id, :change_id, :surface, :event_type, :subtype, :importance, :confidence, :label, :title, :summary, :old_value, :new_value, cast(:payload as jsonb), cast(:entities as jsonb), cast(:tags as text[]), :detected_at, :source_url, :snap_before, :snap_after, :language, 'llm', :provider, :model, :pv, :sv, :status, :dedupe) on conflict (dedupe_key) do nothing returning id""", id=event_id, company_id=company["id"], sensor_id=change.get("sensor_id"), change_id=change["id"], surface=change.get("surface"), event_type=event_type, subtype=c.event_subtype, importance=round(min(1.0, c.importance), 4), confidence=round(min(0.9, c.confidence), 4), label=confidence_label(min(0.9, c.confidence)), title=c.title, summary=c.summary, old_value=c.old_value, new_value=c.new_value, payload=jsonb(payload), entities=jsonb(c.entities), tags=list(dict.fromkeys(c.tags + ["llm"])), detected_at=detected_at, source_url=sensor.get("url"), snap_before=change.get("snapshot_before"), snap_after=change.get("snapshot_after"), language=c.language, provider=PROVIDER_NAME, model=model, pv=prompt_version, sv=SCHEMA_VERSIONS["ChangeClassification"], status="review" if c.confidence < 0.5 else "active", dedupe=dedupe) if row is None: return None if sensor.get("url"): await execute(conn, """insert into event_sources (event_id, sensor_id, source_url, snapshot_id, surface, detected_at, kind) values (:e, :s, :u, :snap, :surface, :at, 'primary') on conflict do nothing""", e=event_id, s=change.get("sensor_id"), u=sensor["url"], snap=change.get("snapshot_after"), surface=change.get("surface"), at=detected_at) await attach_to_cluster(conn, {"id": event_id, "company_id": company["id"], "sensor_id": change.get("sensor_id"), "surface": change.get("surface"), "event_type": event_type, "event_subtype": c.event_subtype, "title": c.title, "confidence": min(0.9, c.confidence), "detected_at": detected_at, "source_url": sensor.get("url"), "snapshot_after": change.get("snapshot_after")}, entity_key=c.title) if c.confidence < 0.5: await execute(conn, "insert into review_queue (id, kind, ref_id, company_id, payload) values (:id, 'low_confidence', :ref, :c, cast(:p as jsonb))", id=new_id("review"), ref=event_id, c=company["id"], p=jsonb({"origin": "llm", "title": c.title, "confidence": c.confidence})) await execute(conn, "update sensors set event_count = event_count + 1 where id = :s", s=change.get("sensor_id")) await execute(conn, "update companies set last_event_at = greatest(coalesce(last_event_at, cast(:at as timestamptz)), cast(:at as timestamptz)) where id = :c", at=detected_at, c=company["id"]) return event_id async def _handle_summarize_event(conn, job: dict[str, Any]) -> tuple[dict[str, Any], str, int, int, int]: # type: ignore[no-untyped-def] event = await fetch_one(conn, "select * from events where id = :id", id=job["ref_id"]) if event is None: raise LLMError("event not found") bundle = await _load_change_bundle(conn, event["change_id"]) if event.get("change_id") else None if bundle is None: raise LLMError("event has no change context") change, company, sensor = bundle context = build_change_context(change, company, sensor) context["event"] = {"type": event["event_type"], "subtype": event["event_subtype"], "title": event["title"], "old_value": event.get("old_value"), "new_value": event.get("new_value"), "entities": event.get("entities") or {}} legal = event["event_type"] == EventType.LEGAL prompt = load_prompt("legal-diff" if legal else "event-summarizer") schema: type[BaseModel] = LegalDiffSummary if legal else EventSummary res = await get_provider().complete_json("medium", prompt.system, jsonb(context), schema, max_tokens=700) data = res.data payload = dict(event.get("payload") or {}) if event.get("summary"): payload.setdefault("summary_prev", event["summary"]) llm_block: dict[str, Any] = {"model": res.model, "prompt_version": prompt.ref, "confidence": getattr(data, "confidence", None), "repaired": res.repaired} if legal: d: LegalDiffSummary = data # type: ignore[assignment] llm_block.update({"materiality": d.materiality, "sections_changed": [s.model_dump() for s in d.sections_changed], "user_impact": d.user_impact}) schema_version = SCHEMA_VERSIONS["LegalDiffSummary"] summary = d.summary tags = list(event.get("tags") or []) + [f"materiality:{d.materiality}"] else: e: EventSummary = data # type: ignore[assignment] llm_block.update({"key_points": e.key_points}) schema_version = SCHEMA_VERSIONS["EventSummary"] summary = e.summary tags = list(event.get("tags") or []) payload["llm"] = llm_block await execute(conn, """ update events set summary = :summary, origin = case when origin = 'deterministic' then 'hybrid' else origin end, model_provider = :provider, model_name = :model, prompt_version = :pv, schema_version = :sv, payload = cast(:payload as jsonb), tags = cast(:tags as text[]), language = coalesce(language, :lang) where id = :id""", summary=summary, provider=PROVIDER_NAME, model=res.model, pv=prompt.ref, sv=schema_version, payload=jsonb(payload), tags=list(dict.fromkeys(tags)), lang=getattr(data, "language", None), id=event["id"]) await execute(conn, "update changes set status = 'enriched' where id = :id and status = 'processed'", id=change["id"]) return {"schema_version": schema_version, "summary": data.model_dump(), "repaired": res.repaired}, res.model, res.request_tokens, res.response_tokens, res.latency_ms async def _handle_classify_industry(conn, job: dict[str, Any]) -> tuple[dict[str, Any], str, int, int, int]: # type: ignore[no-untyped-def] company = await fetch_one(conn, "select * from companies where id = :id", id=job["ref_id"]) if company is None: raise LLMError("company not found") allowed = await fetch_all(conn, "select slug, name from industries order by sort_order, slug") products = await fetch_all(conn, "select name from products where company_id = :c and status = 'listed' order by first_seen_at desc limit 10", c=company["id"]) jobs = await fetch_all(conn, "select title from jobs where company_id = :c and status = 'open' order by first_seen_at desc limit 10", c=company["id"]) home = await fetch_one(conn, """select s.title, s.extracted->'meta' as meta from snapshots s join sensors se on se.id = s.sensor_id where se.company_id = :c and se.surface = 'homepage' order by s.fetched_at desc limit 1""", c=company["id"]) context = {"company": {"name": company["display_name"], "domain": company["canonical_domain"], "description": _cut(company.get("description"), 600)}, "homepage": {"title": (home or {}).get("title"), "meta": (home or {}).get("meta")}, "products": [p["name"] for p in products], "job_titles": [j["title"] for j in jobs], "allowed_industries": [{"slug": r["slug"], "name": r["name"]} for r in allowed]} prompt = load_prompt("industry-tagger") res = await get_provider().complete_json("small", prompt.system, jsonb(context), IndustryTags, max_tokens=400) allowed_slugs = {r["slug"] for r in allowed} tags: IndustryTags = res.data industries = [s for s in tags.industries if s in allowed_slugs] meta = dict(company.get("source_meta") or {}) meta["llm_industries"] = {"industries": industries, "primary": tags.primary if tags.primary in allowed_slugs else None, "keywords": tags.keywords, "confidence": tags.confidence, "model": res.model, "prompt_version": prompt.ref, "at": datetime.now(UTC).isoformat()} await execute(conn, "update companies set source_meta = cast(:m as jsonb), updated_at = now() where id = :id", m=jsonb(meta), id=company["id"]) return {"schema_version": SCHEMA_VERSIONS["IndustryTags"], "industries": meta["llm_industries"]}, res.model, res.request_tokens, res.response_tokens, res.latency_ms HANDLERS = {"classify_change": _handle_classify_change, "summarize_event": _handle_summarize_event, "classify_industry": _handle_classify_industry} # ---------------------------------------------------------------------------------------------------------------- worker async def run_llm_jobs(limit: int = 10) -> dict[str, int]: """Process up to `limit` jobs of one kind. Returns counters. Safe to call concurrently (SKIP LOCKED).""" stats = {"claimed": 0, "done": 0, "failed": 0, "retried": 0, "skipped_budget": 0} if not settings.llm_configured: if not _warned["unconfigured"]: log.info("llm enrichment disabled (CA_LLM_BASE_URL unset or CA_LLM_ENABLED=0); jobs stay pending") _warned["unconfigured"] = True return stats async with transaction() as conn: left = await _budget_left(conn) if left <= 0: stats["skipped_budget"] = 1 return stats kind = await _pick_kind(conn) if kind is None: return stats jobs = await claim_jobs(conn, kind, min(limit, left)) stats["claimed"] = len(jobs) new_events: list[str] = [] for job in jobs: handler = HANDLERS.get(job["kind"]) async with transaction() as conn: if handler is None: await _finish(conn, job, status="failed", model=None, prompt_version=None, result=None, error=f"unknown kind {job['kind']}") stats["failed"] += 1 continue try: result, model, req, resp, latency = await handler(conn, job) except LLMNotConfigured: await _finish(conn, job, status="pending", model=None, prompt_version=None, result=None, error="not configured") return stats except LLMValidationError as exc: await _finish(conn, job, status="failed", model=None, prompt_version=None, result=None, error=str(exc)[:500]) stats["failed"] += 1 continue except LLMError as exc: retry = exc.retryable and job["attempts"] < settings.llm_job_max_attempts await _finish(conn, job, status="pending" if retry else "failed", model=None, prompt_version=None, result=None, error=str(exc)[:500]) stats["retried" if retry else "failed"] += 1 if exc.retryable and not retry: log.warning("llm job exhausted", extra={"job": job["id"], "error": str(exc)[:200]}) if exc.retryable: break # server is unhealthy: stop the batch, the periodic task returns continue except Exception as exc: log.exception("llm job crashed", extra={"job": job["id"]}) await _finish(conn, job, status="failed", model=None, prompt_version=None, result=None, error=f"{exc.__class__.__name__}: {exc}"[:500]) stats["failed"] += 1 continue pv = None if job["kind"] == "classify_change": pv = load_prompt("change-classifier").ref elif job["kind"] == "summarize_event": pv = None await _finish(conn, job, status="done", model=model, prompt_version=pv, result=result, error=None, req=req, resp=resp, latency=latency) stats["done"] += 1 if result.get("event_id"): new_events.append(result["event_id"]) if new_events: try: from companyatlas.services.alerts import evaluate_alerts await evaluate_alerts(new_events) except Exception: log.exception("alert evaluation failed after llm events") return stats async def enqueue_llm_job(kind: str, ref_id: str, company_id: str | None) -> bool: async with transaction() as conn: exists = await fetch_val(conn, "select 1 from llm_jobs where kind = :k and ref_id = :r and status in ('pending', 'running')", k=kind, r=ref_id) if exists: return False await execute(conn, "insert into llm_jobs (id, kind, ref_id, company_id, status) values (:id, :k, :r, :c, 'pending')", id=new_id("llm_job"), k=kind, r=ref_id, c=company_id) return True @periodic("llm-enrich", every_s=15, initial_delay_s=20) async def llm_enrich_task() -> None: stats = await run_llm_jobs(limit=10) if stats["claimed"]: log.info("llm-enrich", extra=stats) __all__ = ["HANDLERS", "build_change_context", "claim_jobs", "enqueue_llm_job", "run_llm_jobs"]