spb/company-atlas
Public
Python 66.3%
TypeScript 22.7%
JavaScript 8.6%
HTML 1.4%
CSS 0.7%
1"""LLM enrichment worker (spec §24, §2.4): drains `llm_jobs` with SKIP LOCKED, sticky by kind (one model per stream so the on-demand2server does not thrash), builds a bounded context from the change diff and structured deltas, calls the gateway, validates the JSON3and writes events / summaries with full provenance (`origin`, model, prompt version, schema version, tokens, latency).45Job kinds6- classify_change (ref = change id, small model) → new event `origin='llm'` when material, else recorded as non-material.7- summarize_event (ref = event id, medium model) → `events.summary` (+ payload.llm), `origin='hybrid'`; LEGAL events use the legal-diff prompt.8- classify_industry (ref = company id, small model) → `companies.source_meta.llm_industries` suggestion only (never overwrites registry data).910Graceful when `settings.llm_configured` is False: jobs stay pending, the worker logs once and returns.11"""12from __future__ import annotations1314import logging15from datetime import UTC, datetime16from typing import Any1718from pydantic import BaseModel1920from companyatlas.config import settings21from companyatlas.db import execute, fetch_all, fetch_one, fetch_val, jsonb, transaction22from companyatlas.ids import new_id, stable_hash23from companyatlas.services.clustering import attach_to_cluster24from companyatlas.services.llm.gateway import LLMError, LLMNotConfigured, LLMValidationError, get_provider25from companyatlas.services.llm.prompts import load_prompt26from companyatlas.services.llm.schemas import SCHEMA_VERSIONS, ChangeClassification, EventSummary, IndustryTags, LegalDiffSummary27from companyatlas.services.periodic import periodic28from companyatlas.taxonomy import EVENT_SUBTYPES, EventType, confidence_label2930log = logging.getLogger(__name__)3132PROVIDER_NAME = "openai-compatible"33KIND_ORDER = ("classify_change", "summarize_event", "classify_industry")34MAX_BLOCKS = 1235MATERIAL_MIN_IMPORTANCE = 0.2536_sticky: dict[str, str | None] = {"kind": None}37_warned = {"unconfigured": False}383940# ---------------------------------------------------------------------------------------------------------------- context414243def _cut(text: str | None, limit: int) -> str | None:44 if not text:45 return None46 b = text.encode("utf-8")47 return text if len(b) <= limit else b[:limit].decode("utf-8", "ignore") + "…"484950def build_change_context(change: dict[str, Any], company: dict[str, Any], sensor: dict[str, Any]) -> dict[str, Any]:51 """Bounded, structured context: company meta, surface, ≤ 12 blocks (≤ 3 kB each side), structured deltas (trimmed lists)."""52 limit = settings.llm_context_block_bytes53 diff = change.get("diff") or {}54 blocks: list[dict[str, Any]] = []55 for bucket in ("added", "removed", "modified"):56 for d in (diff.get(bucket) or [])[:MAX_BLOCKS]:57 if len(blocks) >= MAX_BLOCKS:58 break59 blocks.append({"op": bucket, "kind": d.get("kind"), "path": d.get("path"), "before": _cut(d.get("before"), limit), "after": _cut(d.get("after"), limit)})60 delta = change.get("structured_delta") or {}61 trimmed: dict[str, Any] = {}62 for key, value in delta.items():63 if isinstance(value, dict):64 trimmed[key] = {k: (v[:10] if isinstance(v, list) else v) for k, v in value.items()}65 else:66 trimmed[key] = value67 return {68 "company": {"name": company.get("display_name"), "domain": company.get("canonical_domain"), "country": company.get("country"),69 "industries": list(company.get("industries") or [])[:5], "description": _cut(company.get("description"), 400)},70 "surface": change.get("surface") or sensor.get("surface"), "source_url": sensor.get("url"), "detected_at": str(change.get("detected_at")),71 "significance": change.get("significance"), "change_kind": change.get("kind"),72 "counts": diff.get("counts") or {"added": change.get("blocks_added"), "removed": change.get("blocks_removed"), "modified": change.get("blocks_modified")},73 "blocks": blocks, "structured_delta": trimmed,74 }757677async def _load_change_bundle(conn, change_id: str) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]] | None: # type: ignore[no-untyped-def]78 change = await fetch_one(conn, "select * from changes where id = :id", id=change_id)79 if change is None:80 return None81 sensor = await fetch_one(conn, "select id, url, surface, connector_id from sensors where id = :id", id=change["sensor_id"]) or {}82 company = await fetch_one(conn, "select id, slug, display_name, canonical_domain, country, industries, description from companies where id = :id",83 id=change["company_id"]) or {}84 return change, company, sensor858687# ---------------------------------------------------------------------------------------------------------------- claiming888990async def _budget_left(conn) -> int: # type: ignore[no-untyped-def]91 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')")92 return max(0, settings.llm_daily_budget - int(used or 0))939495async def _pick_kind(conn) -> str | None: # type: ignore[no-untyped-def]96 rows = await fetch_all(conn, "select kind, count(*) as n, min(created_at) as oldest from llm_jobs where status = 'pending' group by kind")97 if not rows:98 return None99 pending = {r["kind"]: r for r in rows}100 if _sticky["kind"] in pending:101 return _sticky["kind"]102 ordered = sorted(rows, key=lambda r: (KIND_ORDER.index(r["kind"]) if r["kind"] in KIND_ORDER else 99, r["oldest"]))103 _sticky["kind"] = ordered[0]["kind"]104 return _sticky["kind"]105106107async def claim_jobs(conn, kind: str, limit: int) -> list[dict[str, Any]]: # type: ignore[no-untyped-def]108 return await fetch_all(conn, """109 update llm_jobs set status = 'running', started_at = now(), attempts = attempts + 1110 where id in (select id from llm_jobs where status = 'pending' and kind = :kind order by created_at limit :limit for update skip locked)111 returning *""", kind=kind, limit=limit)112113114async 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]115 error: str | None, req: int = 0, resp: int = 0, latency: int = 0) -> None:116 await execute(conn, """117 update llm_jobs set status = :status, model = :model, prompt_version = :pv, result = cast(:result as jsonb), error = :error,118 request_tokens = coalesce(request_tokens, 0) + :req, response_tokens = coalesce(response_tokens, 0) + :resp, latency_ms = :latency,119 finished_at = case when :status in ('done', 'failed') then now() else finished_at end120 where id = :id""", status=status, model=model, pv=prompt_version, result=jsonb(result) if result is not None else None, error=error,121 req=req, resp=resp, latency=latency, id=job["id"])122 if model and (req or resp):123 await execute(conn, """124 insert into cost_ledger (day, dimension, key, units, cost_estimate) values (current_date, 'llm', :key, :units, 0)125 on conflict (day, dimension, key) do update set units = cost_ledger.units + excluded.units""", key=model, units=float(req + resp))126127128# ---------------------------------------------------------------------------------------------------------------- handlers129130131async def _handle_classify_change(conn, job: dict[str, Any]) -> tuple[dict[str, Any], str, int, int, int]: # type: ignore[no-untyped-def]132 bundle = await _load_change_bundle(conn, job["ref_id"])133 if bundle is None:134 raise LLMError("change not found")135 change, company, sensor = bundle136 prompt = load_prompt("change-classifier")137 context = build_change_context(change, company, sensor)138 res = await get_provider().complete_json("small", prompt.system, jsonb(context), ChangeClassification, max_tokens=700)139 c: ChangeClassification = res.data140 result: dict[str, Any] = {"schema_version": SCHEMA_VERSIONS["ChangeClassification"], "classification": c.model_dump(), "repaired": res.repaired}141 material = c.is_material and c.event_subtype != "OTHER" and c.importance >= MATERIAL_MIN_IMPORTANCE142 if material:143 event_id = await _insert_llm_event(conn, change, company, sensor, c, model=res.model, prompt_version=prompt.ref)144 result["event_id"] = event_id145 await execute(conn, "update changes set status = 'enriched' where id = :id", id=change["id"])146 return result, res.model, res.request_tokens, res.response_tokens, res.latency_ms147148149async def _insert_llm_event(conn, change: dict[str, Any], company: dict[str, Any], sensor: dict[str, Any], c: ChangeClassification, *, # type: ignore[no-untyped-def]150 model: str, prompt_version: str) -> str | None:151 event_type = str(EVENT_SUBTYPES.get(c.event_subtype, (EventType.OTHER, 0.3))[0])152 # A deterministic event of the same subtype already exists for this change → enrich it instead of adding a near-duplicate.153 existing = await fetch_one(conn, """select id, summary from events where change_id = :c and event_subtype = :st and origin in ('deterministic', 'hybrid')154 and status in ('active', 'review') order by created_at limit 1""", c=change["id"], st=c.event_subtype)155 if existing is not None:156 await execute(conn, """update events set summary = coalesce(summary, :summary), origin = 'hybrid', model_provider = :provider, model_name = :model,157 prompt_version = :pv, payload = payload || cast(:extra as jsonb) where id = :id""",158 summary=c.summary, provider=PROVIDER_NAME, model=model, pv=prompt_version, id=existing["id"],159 extra=jsonb({"llm_classification": {"title": c.title, "importance": c.importance, "confidence": c.confidence, "tags": c.tags}}))160 return None161 dedupe = stable_hash(company["id"], "llm", change["id"], c.event_subtype, length=40)162 event_id = new_id("event")163 detected_at = change.get("detected_at") or datetime.now(UTC)164 payload = {"llm": {"importance": c.importance, "confidence": c.confidence, "is_material": c.is_material}, "significance": change.get("significance"),165 "change_kind": change.get("kind"), "entity_key": c.title}166 row = await fetch_one(conn, """167 insert into events (id, company_id, sensor_id, change_id, surface, event_type, event_subtype, importance, confidence, confidence_label, title, summary,168 old_value, new_value, payload, entities, tags, detected_at, source_url, snapshot_before, snapshot_after, language, origin,169 model_provider, model_name, prompt_version, schema_version, status, dedupe_key)170 values (:id, :company_id, :sensor_id, :change_id, :surface, :event_type, :subtype, :importance, :confidence, :label, :title, :summary, :old_value,171 :new_value, cast(:payload as jsonb), cast(:entities as jsonb), cast(:tags as text[]), :detected_at, :source_url, :snap_before, :snap_after,172 :language, 'llm', :provider, :model, :pv, :sv, :status, :dedupe)173 on conflict (dedupe_key) do nothing returning id""",174 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,175 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)),176 title=c.title, summary=c.summary, old_value=c.old_value, new_value=c.new_value, payload=jsonb(payload), entities=jsonb(c.entities),177 tags=list(dict.fromkeys(c.tags + ["llm"])), detected_at=detected_at, source_url=sensor.get("url"), snap_before=change.get("snapshot_before"),178 snap_after=change.get("snapshot_after"), language=c.language, provider=PROVIDER_NAME, model=model, pv=prompt_version,179 sv=SCHEMA_VERSIONS["ChangeClassification"], status="review" if c.confidence < 0.5 else "active", dedupe=dedupe)180 if row is None:181 return None182 if sensor.get("url"):183 await execute(conn, """insert into event_sources (event_id, sensor_id, source_url, snapshot_id, surface, detected_at, kind)184 values (:e, :s, :u, :snap, :surface, :at, 'primary') on conflict do nothing""",185 e=event_id, s=change.get("sensor_id"), u=sensor["url"], snap=change.get("snapshot_after"), surface=change.get("surface"), at=detected_at)186 await attach_to_cluster(conn, {"id": event_id, "company_id": company["id"], "sensor_id": change.get("sensor_id"), "surface": change.get("surface"),187 "event_type": event_type, "event_subtype": c.event_subtype, "title": c.title, "confidence": min(0.9, c.confidence),188 "detected_at": detected_at, "source_url": sensor.get("url"), "snapshot_after": change.get("snapshot_after")},189 entity_key=c.title)190 if c.confidence < 0.5:191 await execute(conn, "insert into review_queue (id, kind, ref_id, company_id, payload) values (:id, 'low_confidence', :ref, :c, cast(:p as jsonb))",192 id=new_id("review"), ref=event_id, c=company["id"], p=jsonb({"origin": "llm", "title": c.title, "confidence": c.confidence}))193 await execute(conn, "update sensors set event_count = event_count + 1 where id = :s", s=change.get("sensor_id"))194 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",195 at=detected_at, c=company["id"])196 return event_id197198199async def _handle_summarize_event(conn, job: dict[str, Any]) -> tuple[dict[str, Any], str, int, int, int]: # type: ignore[no-untyped-def]200 event = await fetch_one(conn, "select * from events where id = :id", id=job["ref_id"])201 if event is None:202 raise LLMError("event not found")203 bundle = await _load_change_bundle(conn, event["change_id"]) if event.get("change_id") else None204 if bundle is None:205 raise LLMError("event has no change context")206 change, company, sensor = bundle207 context = build_change_context(change, company, sensor)208 context["event"] = {"type": event["event_type"], "subtype": event["event_subtype"], "title": event["title"], "old_value": event.get("old_value"),209 "new_value": event.get("new_value"), "entities": event.get("entities") or {}}210 legal = event["event_type"] == EventType.LEGAL211 prompt = load_prompt("legal-diff" if legal else "event-summarizer")212 schema: type[BaseModel] = LegalDiffSummary if legal else EventSummary213 res = await get_provider().complete_json("medium", prompt.system, jsonb(context), schema, max_tokens=700)214 data = res.data215 payload = dict(event.get("payload") or {})216 if event.get("summary"):217 payload.setdefault("summary_prev", event["summary"])218 llm_block: dict[str, Any] = {"model": res.model, "prompt_version": prompt.ref, "confidence": getattr(data, "confidence", None), "repaired": res.repaired}219 if legal:220 d: LegalDiffSummary = data # type: ignore[assignment]221 llm_block.update({"materiality": d.materiality, "sections_changed": [s.model_dump() for s in d.sections_changed], "user_impact": d.user_impact})222 schema_version = SCHEMA_VERSIONS["LegalDiffSummary"]223 summary = d.summary224 tags = list(event.get("tags") or []) + [f"materiality:{d.materiality}"]225 else:226 e: EventSummary = data # type: ignore[assignment]227 llm_block.update({"key_points": e.key_points})228 schema_version = SCHEMA_VERSIONS["EventSummary"]229 summary = e.summary230 tags = list(event.get("tags") or [])231 payload["llm"] = llm_block232 await execute(conn, """233 update events set summary = :summary, origin = case when origin = 'deterministic' then 'hybrid' else origin end, model_provider = :provider,234 model_name = :model, prompt_version = :pv, schema_version = :sv, payload = cast(:payload as jsonb), tags = cast(:tags as text[]),235 language = coalesce(language, :lang)236 where id = :id""", summary=summary, provider=PROVIDER_NAME, model=res.model, pv=prompt.ref, sv=schema_version, payload=jsonb(payload),237 tags=list(dict.fromkeys(tags)), lang=getattr(data, "language", None), id=event["id"])238 await execute(conn, "update changes set status = 'enriched' where id = :id and status = 'processed'", id=change["id"])239 return {"schema_version": schema_version, "summary": data.model_dump(), "repaired": res.repaired}, res.model, res.request_tokens, res.response_tokens, res.latency_ms240241242async def _handle_classify_industry(conn, job: dict[str, Any]) -> tuple[dict[str, Any], str, int, int, int]: # type: ignore[no-untyped-def]243 company = await fetch_one(conn, "select * from companies where id = :id", id=job["ref_id"])244 if company is None:245 raise LLMError("company not found")246 allowed = await fetch_all(conn, "select slug, name from industries order by sort_order, slug")247 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"])248 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"])249 home = await fetch_one(conn, """select s.title, s.extracted->'meta' as meta from snapshots s join sensors se on se.id = s.sensor_id250 where se.company_id = :c and se.surface = 'homepage' order by s.fetched_at desc limit 1""", c=company["id"])251 context = {"company": {"name": company["display_name"], "domain": company["canonical_domain"], "description": _cut(company.get("description"), 600)},252 "homepage": {"title": (home or {}).get("title"), "meta": (home or {}).get("meta")}, "products": [p["name"] for p in products],253 "job_titles": [j["title"] for j in jobs], "allowed_industries": [{"slug": r["slug"], "name": r["name"]} for r in allowed]}254 prompt = load_prompt("industry-tagger")255 res = await get_provider().complete_json("small", prompt.system, jsonb(context), IndustryTags, max_tokens=400)256 allowed_slugs = {r["slug"] for r in allowed}257 tags: IndustryTags = res.data258 industries = [s for s in tags.industries if s in allowed_slugs]259 meta = dict(company.get("source_meta") or {})260 meta["llm_industries"] = {"industries": industries, "primary": tags.primary if tags.primary in allowed_slugs else None, "keywords": tags.keywords,261 "confidence": tags.confidence, "model": res.model, "prompt_version": prompt.ref, "at": datetime.now(UTC).isoformat()}262 await execute(conn, "update companies set source_meta = cast(:m as jsonb), updated_at = now() where id = :id", m=jsonb(meta), id=company["id"])263 return {"schema_version": SCHEMA_VERSIONS["IndustryTags"], "industries": meta["llm_industries"]}, res.model, res.request_tokens, res.response_tokens, res.latency_ms264265266HANDLERS = {"classify_change": _handle_classify_change, "summarize_event": _handle_summarize_event, "classify_industry": _handle_classify_industry}267268269# ---------------------------------------------------------------------------------------------------------------- worker270271272async def run_llm_jobs(limit: int = 10) -> dict[str, int]:273 """Process up to `limit` jobs of one kind. Returns counters. Safe to call concurrently (SKIP LOCKED)."""274 stats = {"claimed": 0, "done": 0, "failed": 0, "retried": 0, "skipped_budget": 0}275 if not settings.llm_configured:276 if not _warned["unconfigured"]:277 log.info("llm enrichment disabled (CA_LLM_BASE_URL unset or CA_LLM_ENABLED=0); jobs stay pending")278 _warned["unconfigured"] = True279 return stats280 async with transaction() as conn:281 left = await _budget_left(conn)282 if left <= 0:283 stats["skipped_budget"] = 1284 return stats285 kind = await _pick_kind(conn)286 if kind is None:287 return stats288 jobs = await claim_jobs(conn, kind, min(limit, left))289 stats["claimed"] = len(jobs)290 new_events: list[str] = []291 for job in jobs:292 handler = HANDLERS.get(job["kind"])293 async with transaction() as conn:294 if handler is None:295 await _finish(conn, job, status="failed", model=None, prompt_version=None, result=None, error=f"unknown kind {job['kind']}")296 stats["failed"] += 1297 continue298 try:299 result, model, req, resp, latency = await handler(conn, job)300 except LLMNotConfigured:301 await _finish(conn, job, status="pending", model=None, prompt_version=None, result=None, error="not configured")302 return stats303 except LLMValidationError as exc:304 await _finish(conn, job, status="failed", model=None, prompt_version=None, result=None, error=str(exc)[:500])305 stats["failed"] += 1306 continue307 except LLMError as exc:308 retry = exc.retryable and job["attempts"] < settings.llm_job_max_attempts309 await _finish(conn, job, status="pending" if retry else "failed", model=None, prompt_version=None, result=None, error=str(exc)[:500])310 stats["retried" if retry else "failed"] += 1311 if exc.retryable and not retry:312 log.warning("llm job exhausted", extra={"job": job["id"], "error": str(exc)[:200]})313 if exc.retryable:314 break # server is unhealthy: stop the batch, the periodic task returns315 continue316 except Exception as exc:317 log.exception("llm job crashed", extra={"job": job["id"]})318 await _finish(conn, job, status="failed", model=None, prompt_version=None, result=None, error=f"{exc.__class__.__name__}: {exc}"[:500])319 stats["failed"] += 1320 continue321 pv = None322 if job["kind"] == "classify_change":323 pv = load_prompt("change-classifier").ref324 elif job["kind"] == "summarize_event":325 pv = None326 await _finish(conn, job, status="done", model=model, prompt_version=pv, result=result, error=None, req=req, resp=resp, latency=latency)327 stats["done"] += 1328 if result.get("event_id"):329 new_events.append(result["event_id"])330 if new_events:331 try:332 from companyatlas.services.alerts import evaluate_alerts333334 await evaluate_alerts(new_events)335 except Exception:336 log.exception("alert evaluation failed after llm events")337 return stats338339340async def enqueue_llm_job(kind: str, ref_id: str, company_id: str | None) -> bool:341 async with transaction() as conn:342 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)343 if exists:344 return False345 await execute(conn, "insert into llm_jobs (id, kind, ref_id, company_id, status) values (:id, :k, :r, :c, 'pending')",346 id=new_id("llm_job"), k=kind, r=ref_id, c=company_id)347 return True348349350@periodic("llm-enrich", every_s=15, initial_delay_s=20)351async def llm_enrich_task() -> None:352 stats = await run_llm_jobs(limit=10)353 if stats["claimed"]:354 log.info("llm-enrich", extra=stats)355356357__all__ = ["HANDLERS", "build_change_context", "claim_jobs", "enqueue_llm_job", "run_llm_jobs"]358