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 | 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 | 9 | OPENNESS = ("open-weights", "open-source", "proprietary", "restricted", "unknown") |
| 6 | 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 | 72 | doc_class: str = Field(description="one of: model_release, model_update, pricing, research_paper, company_news, product_launch, " |
| 11 | 73 | "framework_release, dataset_release, benchmark, hardware, regulation, incident, funding, acquisition, other") |
| 12 | 74 | is_primary_source: bool | None = None |
@@ -16,7 +78,7 @@ class DocumentClassification(BaseModel): | ||
| 16 | 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 | 82 | name: str | None = None |
| 21 | 83 | developer: str | None = Field(default=None, description="organization that developed the model") |
| 22 | 84 | family: str | None = None |
@@ -55,7 +117,7 @@ class ModelPassport(BaseModel): | ||
| 55 | 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 | 121 | model: str |
| 60 | 122 | provider_model_id: str | None = None |
| 61 | 123 | input_per_mtok: float | None = Field(default=None, description="USD per 1M input tokens") |
@@ -70,14 +132,14 @@ class PriceLine(BaseModel): | ||
| 70 | 132 | notes: str | None = None |
| 71 | 133 | |
| 72 | 134 | |
| 73 | −class PricingExtraction(BaseModel): | |
| 135 | +class PricingExtraction(Tolerant): | |
| 74 | 136 | provider: str | None = None |
| 75 | 137 | currency: str = "USD" |
| 76 | 138 | effective_date: str | None = None |
| 77 | 139 | prices: list[PriceLine] = Field(default_factory=list) |
| 78 | 140 | |
| 79 | 141 | |
| 80 | −class CompanyPassport(BaseModel): | |
| 142 | +class CompanyPassport(Tolerant): | |
| 81 | 143 | name: str | None = None |
| 82 | 144 | legal_name: str | None = None |
| 83 | 145 | country: str | None = Field(default=None, description="ISO 3166-1 alpha-2 if determinable") |
@@ -96,7 +158,7 @@ class CompanyPassport(BaseModel): | ||
| 96 | 158 | funding_total_usd: float | None = None |
| 97 | 159 | |
| 98 | 160 | |
| 99 | −class PaperPassport(BaseModel): | |
| 161 | +class PaperPassport(Tolerant): | |
| 100 | 162 | title: str | None = None |
| 101 | 163 | authors: list[str] = Field(default_factory=list) |
| 102 | 164 | affiliations: list[str] = Field(default_factory=list) |
@@ -113,14 +175,14 @@ class PaperPassport(BaseModel): | ||
| 113 | 175 | code_url: str | None = None |
| 114 | 176 | |
| 115 | 177 | |
| 116 | −class BenchmarkRow(BaseModel): | |
| 178 | +class BenchmarkRow(Tolerant): | |
| 117 | 179 | model: str |
| 118 | 180 | score: float |
| 119 | 181 | metric: str | None = None |
| 120 | 182 | config: str | None = None |
| 121 | 183 | |
| 122 | 184 | |
| 123 | −class BenchmarkResultExtraction(BaseModel): | |
| 185 | +class BenchmarkResultExtraction(Tolerant): | |
| 124 | 186 | benchmark: str | None = None |
| 125 | 187 | metric: str | None = None |
| 126 | 188 | higher_is_better: bool | None = True |
@@ -128,7 +190,7 @@ class BenchmarkResultExtraction(BaseModel): | ||
| 128 | 190 | rows: list[BenchmarkRow] = Field(default_factory=list) |
| 129 | 191 | |
| 130 | 192 | |
| 131 | −class HardwareSpec(BaseModel): | |
| 193 | +class HardwareSpec(Tolerant): | |
| 132 | 194 | name: str | None = None |
| 133 | 195 | manufacturer: str | None = None |
| 134 | 196 | kind: str | None = Field(default=None, description="gpu | cpu | npu | tpu | asic | soc | accelerator | system") |
@@ -146,7 +208,7 @@ class HardwareSpec(BaseModel): | ||
| 146 | 208 | interconnect: str | None = None |
| 147 | 209 | |
| 148 | 210 | |
| 149 | −class ReleaseAnnouncement(BaseModel): | |
| 211 | +class ReleaseAnnouncement(Tolerant): | |
| 150 | 212 | kind: str | None = Field(default=None, description="new_model | model_update | pricing_change | deprecation | new_product | framework_release | research | other") |
| 151 | 213 | title: str | None = None |
| 152 | 214 | date: str | None = None |
| 153 | 215 | |