Spaces:
Running on Zero
Running on Zero
| """Structured (JSON) calls: strict schema generation, tolerant parsing, one repair retry.""" | |
| from __future__ import annotations | |
| import copy | |
| import json | |
| import re | |
| from typing import Any | |
| from pydantic import BaseModel, ValidationError | |
| from .base import LLMClient, LLMError, LLMRequest, UsageMeter | |
| def strict_schema(model: type[BaseModel]) -> dict[str, Any]: | |
| """Pydantic JSON schema rewritten for strict structured-output modes. | |
| Inlines ``$defs``, marks every property required and forbids extra properties, | |
| which both Anthropic and OpenAI strict modes need. | |
| """ | |
| schema = model.model_json_schema() | |
| defs = schema.pop("$defs", {}) | |
| def resolve(node: Any) -> Any: | |
| if isinstance(node, dict): | |
| if "$ref" in node: | |
| name = node["$ref"].split("/")[-1] | |
| return resolve(copy.deepcopy(defs[name])) | |
| out = {k: resolve(v) for k, v in node.items() if k not in ("title", "default")} | |
| if out.get("type") == "object" and "properties" in out: | |
| out["required"] = list(out["properties"].keys()) | |
| out["additionalProperties"] = False | |
| return out | |
| if isinstance(node, list): | |
| return [resolve(x) for x in node] | |
| return node | |
| return resolve(schema) | |
| _FENCE = re.compile(r"```(?:json)?\s*(.*?)```", re.DOTALL) | |
| def extract_json(text: str) -> Any: | |
| """Parse JSON from a model reply that may wrap it in prose or code fences.""" | |
| text = text.strip() | |
| try: | |
| return json.loads(text) | |
| except json.JSONDecodeError: | |
| pass | |
| m = _FENCE.search(text) | |
| if m: | |
| try: | |
| return json.loads(m.group(1)) | |
| except json.JSONDecodeError: | |
| pass | |
| start = min((i for i in (text.find("{"), text.find("[")) if i >= 0), default=-1) | |
| if start >= 0: | |
| decoder = json.JSONDecoder() | |
| try: | |
| obj, _ = decoder.raw_decode(text[start:]) | |
| return obj | |
| except json.JSONDecodeError: | |
| pass | |
| raise ValueError("no JSON object found in model output") | |
| def call_structured[T: BaseModel]( | |
| client: LLMClient, | |
| meter: UsageMeter | None, | |
| stage: str, | |
| system: str, | |
| user: str, | |
| out_model: type[T], | |
| max_tokens: int | None = None, | |
| ) -> T: | |
| # Stage caps are floors, never below the model's configured max_tokens: on models that | |
| # think before answering, thinking tokens count against this limit and a low cap would | |
| # truncate the JSON. | |
| max_tokens = max(max_tokens or 0, client.spec.max_tokens) | |
| schema = strict_schema(out_model) | |
| hint = "\n\nRespond with a single JSON object that matches this JSON schema and nothing else:\n" + json.dumps(schema, ensure_ascii=False) | |
| sys_prompt = system if client.spec.json_mode == "schema" else system + hint | |
| request = LLMRequest(system=sys_prompt, user=user, json_schema=schema, schema_name=out_model.__name__, max_tokens=max_tokens) | |
| resp = client.complete(request) | |
| if meter is not None: | |
| meter.record(stage, client.spec, resp) | |
| try: | |
| return out_model.model_validate(extract_json(resp.text)) | |
| except (ValueError, ValidationError) as first_err: | |
| repair = LLMRequest( | |
| system=sys_prompt, | |
| user=(f"{user}\n\nYour previous reply could not be parsed ({str(first_err)[:300]}). Reply again with only the JSON object."), | |
| json_schema=schema, | |
| schema_name=out_model.__name__, | |
| max_tokens=max_tokens, | |
| ) | |
| resp2 = client.complete(repair) | |
| if meter is not None: | |
| meter.record(stage + ":repair", client.spec, resp2) | |
| try: | |
| return out_model.model_validate(extract_json(resp2.text)) | |
| except (ValueError, ValidationError) as e: | |
| # Do not keep unusable replies in the cache, or a resumed run would replay the same failure. | |
| invalidate = getattr(client, "invalidate", None) | |
| if invalidate is not None: | |
| for key in (resp.cache_key, resp2.cache_key): | |
| if key: | |
| invalidate(key) | |
| truncated = resp2.stop_reason in ("max_tokens", "length") | |
| why = f" (output truncated at max_tokens={max_tokens}; raise it in configs/models.yaml)" if truncated else "" | |
| raise LLMError(f"{stage}: unparseable structured output after repair{why}: {e}") from e | |