spb/company-atlas

1"""Strict output schemas for every LLM task (spec §24: structured JSON validated against schemas). The gateway validates model output2against these models and retries once with the validation error; anything that still fails is recorded as a failed job — never stored3as an event. Schema versions are written to `events.schema_version` / `llm_jobs.result`.4"""5from __future__ import annotations67from typing import Any89from pydantic import BaseModel, ConfigDict, Field, field_validator1011from companyatlas.taxonomy import EVENT_SUBTYPES, FORBIDDEN_WORDING1213SCHEMA_VERSIONS = {"ChangeClassification": "classify-v1", "EventSummary": "summary-v1", "LegalDiffSummary": "legal-v1", "IndustryTags": "industry-v1",14                   "AskRoute": "ask-v1", "CompanyProfileText": "profile-v1"}151617def _clean_text(value: str | None, limit: int) -> str | None:18    if value is None:19        return None20    text = " ".join(str(value).split())21    for bad in FORBIDDEN_WORDING:22        if bad in text.lower():23            raise ValueError(f"forbidden wording: {bad!r} — use 'no longer listed' / 'observed' phrasing")24    return text[:limit] or None252627class _Strict(BaseModel):28    model_config = ConfigDict(extra="ignore", str_strip_whitespace=True)293031class ChangeClassification(_Strict):32    event_subtype: str = Field(description="One of the Company Atlas event subtypes, or OTHER")33    importance: float = Field(ge=0, le=1)34    confidence: float = Field(ge=0, le=1)35    title: str = Field(min_length=3, max_length=120)36    summary: str | None = Field(default=None, max_length=400)37    old_value: str | None = Field(default=None, max_length=200)38    new_value: str | None = Field(default=None, max_length=200)39    entities: dict[str, Any] = Field(default_factory=dict)40    tags: list[str] = Field(default_factory=list, max_length=12)41    language: str | None = Field(default=None, max_length=8)42    is_material: bool = True4344    @field_validator("event_subtype", mode="before")45    @classmethod46    def _subtype(cls, v: Any) -> str:47        s = str(v or "OTHER").strip().upper().replace(" ", "_").replace("-", "_")48        return s if s in EVENT_SUBTYPES else "OTHER"4950    @field_validator("title", mode="before")51    @classmethod52    def _title(cls, v: Any) -> str:53        return _clean_text(str(v or ""), 120) or ""5455    @field_validator("summary", "old_value", "new_value", mode="before")56    @classmethod57    def _texts(cls, v: Any) -> str | None:58        return _clean_text(None if v is None else str(v), 400)5960    @field_validator("tags", mode="before")61    @classmethod62    def _tags(cls, v: Any) -> list[str]:63        if not v:64            return []65        out = []66        for t in list(v)[:12]:67            s = str(t).strip().lower()[:40]68            if s and s not in out:69                out.append(s)70        return out717273class EventSummary(_Strict):74    summary: str = Field(min_length=3, max_length=400)75    key_points: list[str] = Field(default_factory=list, max_length=5)76    confidence: float = Field(default=0.7, ge=0, le=1)77    language: str | None = Field(default=None, max_length=8)7879    @field_validator("summary", mode="before")80    @classmethod81    def _summary(cls, v: Any) -> str:82        return _clean_text(str(v or ""), 400) or ""8384    @field_validator("key_points", mode="before")85    @classmethod86    def _points(cls, v: Any) -> list[str]:87        return [p for p in (_clean_text(str(x), 160) for x in (v or [])[:5]) if p]888990class LegalSection(_Strict):91    section: str = Field(max_length=120)92    change: str = Field(max_length=240)939495class LegalDiffSummary(_Strict):96    sections_changed: list[LegalSection] = Field(default_factory=list, max_length=12)97    materiality: str = Field(default="unclear")        # editorial | minor | material | unclear98    summary: str = Field(min_length=3, max_length=400)99    user_impact: str | None = Field(default=None, max_length=240)100    confidence: float = Field(default=0.6, ge=0, le=1)101102    @field_validator("materiality", mode="before")103    @classmethod104    def _mat(cls, v: Any) -> str:105        s = str(v or "unclear").strip().lower()106        return s if s in ("editorial", "minor", "material", "unclear") else "unclear"107108    @field_validator("summary", "user_impact", mode="before")109    @classmethod110    def _texts(cls, v: Any) -> str | None:111        return _clean_text(None if v is None else str(v), 400)112113114class IndustryTags(_Strict):115    industries: list[str] = Field(default_factory=list, max_length=5)     # industry slugs from the registry116    primary: str | None = None117    keywords: list[str] = Field(default_factory=list, max_length=10)118    confidence: float = Field(default=0.6, ge=0, le=1)119120    @field_validator("industries", "keywords", mode="before")121    @classmethod122    def _slugs(cls, v: Any) -> list[str]:123        out: list[str] = []124        for x in (v or [])[:10]:125            s = str(x).strip().lower().replace(" ", "-")[:60]126            if s and s not in out:127                out.append(s)128        return out129130131class CompanyProfileText(_Strict):132    """Grounded 2–3 sentence company description (services/enrichment.py, job kind `company_profile`). Numbers are additionally checked133    against the source text by the caller (`enrichment.numbers_grounded`) — any figure absent from the source rejects the output."""134    description: str = Field(min_length=40, max_length=700)135    language: str | None = Field(default=None, max_length=8)136    confidence: float = Field(default=0.6, ge=0, le=1)137138    @field_validator("description", mode="before")139    @classmethod140    def _description(cls, v: Any) -> str:141        return _clean_text(str(v or ""), 700) or ""142143144class AskRoute(_Strict):145    """LLM refinement of the deterministic question parser (services/llm/ask.py). It may only tighten filters, never invent results."""146    intent: str = Field(default="search")           # search | hiring | pricing | ai | launch | leadership | expansion | compare | trend147    countries: list[str] = Field(default_factory=list, max_length=8)148    industries: list[str] = Field(default_factory=list, max_length=8)149    event_types: list[str] = Field(default_factory=list, max_length=8)150    event_subtypes: list[str] = Field(default_factory=list, max_length=8)151    companies: list[str] = Field(default_factory=list, max_length=8)152    window_days: int | None = Field(default=None, ge=1, le=3650)153    keywords: list[str] = Field(default_factory=list, max_length=8)154    answer_style: str = Field(default="list")       # list | count | compare | timeline155    confidence: float = Field(default=0.6, ge=0, le=1)156157    @field_validator("event_subtypes", mode="before")158    @classmethod159    def _subtypes(cls, v: Any) -> list[str]:160        return [s for s in (str(x).strip().upper() for x in (v or [])[:8]) if s in EVENT_SUBTYPES]161162    @field_validator("countries", mode="before")163    @classmethod164    def _countries(cls, v: Any) -> list[str]:165        return [s for s in (str(x).strip().upper()[:2] for x in (v or [])[:8]) if len(s) == 2 and s.isalpha()]166167168__all__ = ["SCHEMA_VERSIONS", "AskRoute", "ChangeClassification", "CompanyProfileText", "EventSummary", "IndustryTags", "LegalDiffSummary", "LegalSection"]169