Extraction schemas: tolerant coercion of LLM output (null lists, string booleans/ints)
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
1 changed file +73 −11
modified
src/aiatlas/schemas/extraction.py
+73 −11
@@ -1,12 +1,74 @@ | |||||
| 1 | from __future__ import annotations | 1 | from __future__ import annotations | ||
| 2 | 2 | ||||
| 3 | −from pydantic import BaseModel, Field | 3 | +import types | ||
| 4 | +import typing | ||||
| 5 | +from typing import Any | ||||
| 6 | + | ||||
| 7 | +from pydantic import BaseModel, Field, model_validator | ||||
| 4 | 8 | ||||
| 5 | OPENNESS = ("open-weights", "open-source", "proprietary", "restricted", "unknown") | 9 | OPENNESS = ("open-weights", "open-source", "proprietary", "restricted", "unknown") | ||
| 6 | MODALITIES = ("text", "image", "audio", "video", "code", "embedding", "3d", "multimodal") | 10 | MODALITIES = ("text", "image", "audio", "video", "code", "embedding", "3d", "multimodal") | ||
| 7 | 11 | ||||
| 8 | 12 | ||||
| 9 | −class DocumentClassification(BaseModel): | 13 | +def _is_list_type(tp: Any) -> bool: | ||
| 14 | + origin = typing.get_origin(tp) | ||||
| 15 | + if origin is list: | ||||
| 16 | + return True | ||||
| 17 | + if origin in (typing.Union, types.UnionType): | ||||
| 18 | + return any(_is_list_type(a) for a in typing.get_args(tp)) | ||||
| 19 | + return False | ||||
| 20 | + | ||||
| 21 | + | ||||
| 22 | +def _is_type(tp: Any, target: type) -> bool: | ||||
| 23 | + if tp is target: | ||||
| 24 | + return True | ||||
| 25 | + if typing.get_origin(tp) in (typing.Union, types.UnionType): | ||||
| 26 | + return any(a is target for a in typing.get_args(tp)) | ||||
| 27 | + return False | ||||
| 28 | + | ||||
| 29 | + | ||||
| 30 | +class Tolerant(BaseModel): | ||||
| 31 | + """LLM outputs are messy: null for lists, "yes"/"no" for booleans, "70B" for integers. Coerce instead of failing (and escalating).""" | ||||
| 32 | + | ||||
| 33 | + @model_validator(mode="before") | ||||
| 34 | + @classmethod | ||||
| 35 | + def _coerce(cls, data: Any) -> Any: | ||||
| 36 | + if not isinstance(data, dict): | ||||
| 37 | + return data | ||||
| 38 | + out = dict(data) | ||||
| 39 | + for name, field in cls.model_fields.items(): | ||||
| 40 | + if name not in out: | ||||
| 41 | + continue | ||||
| 42 | + v = out[name] | ||||
| 43 | + tp = field.annotation | ||||
| 44 | + if _is_list_type(tp): | ||||
| 45 | + if v is None or v == "" or v == "null": | ||||
| 46 | + out[name] = [] | ||||
| 47 | + elif isinstance(v, (str, int, float)): | ||||
| 48 | + out[name] = [x.strip() for x in str(v).split(",") if x.strip()] if isinstance(v, str) else [v] | ||||
| 49 | + elif isinstance(v, dict): | ||||
| 50 | + out[name] = [str(x) for x in v.values() if x not in (None, "")] | ||||
| 51 | + elif _is_type(tp, bool) and isinstance(v, str): | ||||
| 52 | + low = v.strip().lower() | ||||
| 53 | + out[name] = True if low in ("yes", "true", "supported", "y") else False if low in ("no", "false", "unsupported", "n", "none") else None | ||||
| 54 | + elif _is_type(tp, int) and isinstance(v, str): | ||||
| 55 | + from aiatlas.sdk.extract.numbers import parse_context_length, parse_param_count | ||||
| 56 | + | ||||
| 57 | + parsed = parse_param_count(v) if "param" in name else parse_context_length(v) | ||||
| 58 | + out[name] = parsed if parsed is not None else None | ||||
| 59 | + elif _is_type(tp, float) and isinstance(v, str): | ||||
| 60 | + import re | ||||
| 61 | + | ||||
| 62 | + m = re.search(r"-?\d+(?:\.\d+)?", v.replace(",", "")) | ||||
| 63 | + out[name] = float(m.group(0)) if m else None | ||||
| 64 | + elif _is_type(tp, str) and isinstance(v, (int, float)) and not isinstance(v, bool): | ||||
| 65 | + out[name] = str(v) | ||||
| 66 | + elif v == "null": | ||||
| 67 | + out[name] = None | ||||
| 68 | + return out | ||||
| 69 | + | ||||
| 70 | + | ||||
| 71 | +class DocumentClassification(Tolerant): | ||||
| 10 | doc_class: str = Field(description="one of: model_release, model_update, pricing, research_paper, company_news, product_launch, " | 72 | doc_class: str = Field(description="one of: model_release, model_update, pricing, research_paper, company_news, product_launch, " | ||
| 11 | "framework_release, dataset_release, benchmark, hardware, regulation, incident, funding, acquisition, other") | 73 | "framework_release, dataset_release, benchmark, hardware, regulation, incident, funding, acquisition, other") | ||
| 12 | is_primary_source: bool | None = None | 74 | is_primary_source: bool | None = None | ||
@@ -16,7 +78,7 @@ class DocumentClassification(BaseModel): | |||||
| 16 | relevance: float = Field(default=0.5, ge=0, le=1, description="relevance to the AI ecosystem") | 78 | relevance: float = Field(default=0.5, ge=0, le=1, description="relevance to the AI ecosystem") | ||
| 17 | 79 | ||||
| 18 | 80 | ||||
| 19 | −class ModelPassport(BaseModel): | 81 | +class ModelPassport(Tolerant): | ||
| 20 | name: str | None = None | 82 | name: str | None = None | ||
| 21 | developer: str | None = Field(default=None, description="organization that developed the model") | 83 | developer: str | None = Field(default=None, description="organization that developed the model") | ||
| 22 | family: str | None = None | 84 | family: str | None = None | ||
@@ -55,7 +117,7 @@ class ModelPassport(BaseModel): | |||||
| 55 | evidence: list[str] = Field(default_factory=list, description="short quotes supporting the most important fields") | 117 | evidence: list[str] = Field(default_factory=list, description="short quotes supporting the most important fields") | ||
| 56 | 118 | ||||
| 57 | 119 | ||||
| 58 | −class PriceLine(BaseModel): | 120 | +class PriceLine(Tolerant): | ||
| 59 | model: str | 121 | model: str | ||
| 60 | provider_model_id: str | None = None | 122 | provider_model_id: str | None = None | ||
| 61 | input_per_mtok: float | None = Field(default=None, description="USD per 1M input tokens") | 123 | input_per_mtok: float | None = Field(default=None, description="USD per 1M input tokens") | ||
@@ -70,14 +132,14 @@ class PriceLine(BaseModel): | |||||
| 70 | notes: str | None = None | 132 | notes: str | None = None | ||
| 71 | 133 | ||||
| 72 | 134 | ||||
| 73 | −class PricingExtraction(BaseModel): | 135 | +class PricingExtraction(Tolerant): | ||
| 74 | provider: str | None = None | 136 | provider: str | None = None | ||
| 75 | currency: str = "USD" | 137 | currency: str = "USD" | ||
| 76 | effective_date: str | None = None | 138 | effective_date: str | None = None | ||
| 77 | prices: list[PriceLine] = Field(default_factory=list) | 139 | prices: list[PriceLine] = Field(default_factory=list) | ||
| 78 | 140 | ||||
| 79 | 141 | ||||
| 80 | −class CompanyPassport(BaseModel): | 142 | +class CompanyPassport(Tolerant): | ||
| 81 | name: str | None = None | 143 | name: str | None = None | ||
| 82 | legal_name: str | None = None | 144 | legal_name: str | None = None | ||
| 83 | country: str | None = Field(default=None, description="ISO 3166-1 alpha-2 if determinable") | 145 | country: str | None = Field(default=None, description="ISO 3166-1 alpha-2 if determinable") | ||
@@ -96,7 +158,7 @@ class CompanyPassport(BaseModel): | |||||
| 96 | funding_total_usd: float | None = None | 158 | funding_total_usd: float | None = None | ||
| 97 | 159 | ||||
| 98 | 160 | ||||
| 99 | −class PaperPassport(BaseModel): | 161 | +class PaperPassport(Tolerant): | ||
| 100 | title: str | None = None | 162 | title: str | None = None | ||
| 101 | authors: list[str] = Field(default_factory=list) | 163 | authors: list[str] = Field(default_factory=list) | ||
| 102 | affiliations: list[str] = Field(default_factory=list) | 164 | affiliations: list[str] = Field(default_factory=list) | ||
@@ -113,14 +175,14 @@ class PaperPassport(BaseModel): | |||||
| 113 | code_url: str | None = None | 175 | code_url: str | None = None | ||
| 114 | 176 | ||||
| 115 | 177 | ||||
| 116 | −class BenchmarkRow(BaseModel): | 178 | +class BenchmarkRow(Tolerant): | ||
| 117 | model: str | 179 | model: str | ||
| 118 | score: float | 180 | score: float | ||
| 119 | metric: str | None = None | 181 | metric: str | None = None | ||
| 120 | config: str | None = None | 182 | config: str | None = None | ||
| 121 | 183 | ||||
| 122 | 184 | ||||
| 123 | −class BenchmarkResultExtraction(BaseModel): | 185 | +class BenchmarkResultExtraction(Tolerant): | ||
| 124 | benchmark: str | None = None | 186 | benchmark: str | None = None | ||
| 125 | metric: str | None = None | 187 | metric: str | None = None | ||
| 126 | higher_is_better: bool | None = True | 188 | higher_is_better: bool | None = True | ||
@@ -128,7 +190,7 @@ class BenchmarkResultExtraction(BaseModel): | |||||
| 128 | rows: list[BenchmarkRow] = Field(default_factory=list) | 190 | rows: list[BenchmarkRow] = Field(default_factory=list) | ||
| 129 | 191 | ||||
| 130 | 192 | ||||
| 131 | −class HardwareSpec(BaseModel): | 193 | +class HardwareSpec(Tolerant): | ||
| 132 | name: str | None = None | 194 | name: str | None = None | ||
| 133 | manufacturer: str | None = None | 195 | manufacturer: str | None = None | ||
| 134 | kind: str | None = Field(default=None, description="gpu | cpu | npu | tpu | asic | soc | accelerator | system") | 196 | kind: str | None = Field(default=None, description="gpu | cpu | npu | tpu | asic | soc | accelerator | system") | ||
@@ -146,7 +208,7 @@ class HardwareSpec(BaseModel): | |||||
| 146 | interconnect: str | None = None | 208 | interconnect: str | None = None | ||
| 147 | 209 | ||||
| 148 | 210 | ||||
| 149 | −class ReleaseAnnouncement(BaseModel): | 211 | +class ReleaseAnnouncement(Tolerant): | ||
| 150 | kind: str | None = Field(default=None, description="new_model | model_update | pricing_change | deprecation | new_product | framework_release | research | other") | 212 | kind: str | None = Field(default=None, description="new_model | model_update | pricing_change | deprecation | new_product | framework_release | research | other") | ||
| 151 | title: str | None = None | 213 | title: str | None = None | ||
| 152 | date: str | None = None | 214 | date: str | None = None | ||
| 153 | 215 | ||||