Spaces:
Running on Zero
Running on Zero
File size: 4,398 Bytes
465b477 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 | """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
|