"""structured-output-json-schema.py -- validate model JSON against a schema and retry with the error. What it does: implements a small JSON Schema validator (the subset most extraction tasks need: type, properties, required, additionalProperties, enum, minimum, maximum, minLength, items, minItems), a tolerant JSON extractor that strips code fences and trailing prose, and a retry loop that feeds the validation error back to the model. The `call_model` function here is a deterministic SIMULATOR that produces the failure modes seen in practice (fences, prose, wrong types, missing keys, out-of-range values) with fixed probabilities from a seeded RNG, so the numbers printed are properties of the loop, not of any model. Replace `call_model` with a real API call and the rest is unchanged. Inputs: none (the schema and the simulator are in the file) Run: python code/structured-output-json-schema.py Requires: Python 3.11+ standard library. For production use, the `jsonschema` package (4.26.0 at the time of writing) implements the full specification. """ from __future__ import annotations import json import random import re SCHEMA = { "type": "object", "properties": { "vendor": {"type": "string", "minLength": 1}, "total": {"type": "number", "minimum": 0}, "currency": {"type": "string", "enum": ["USD", "EUR", "GBP"]}, "line_items": { "type": "array", "minItems": 1, "items": { "type": "object", "properties": {"sku": {"type": "string"}, "qty": {"type": "integer", "minimum": 1}}, "required": ["sku", "qty"], "additionalProperties": False, }, }, }, "required": ["vendor", "total", "currency", "line_items"], "additionalProperties": False, } class SchemaError(ValueError): pass def validate(value, schema, path="$") -> None: t = schema.get("type") if t == "object": if not isinstance(value, dict): raise SchemaError(f"{path}: expected object, got {type(value).__name__}") for key in schema.get("required", []): if key not in value: raise SchemaError(f"{path}: missing required key {key!r}") props = schema.get("properties", {}) if schema.get("additionalProperties") is False: extra = set(value) - set(props) if extra: raise SchemaError(f"{path}: unexpected keys {sorted(extra)}") for key, sub in props.items(): if key in value: validate(value[key], sub, f"{path}.{key}") elif t == "array": if not isinstance(value, list): raise SchemaError(f"{path}: expected array, got {type(value).__name__}") if len(value) < schema.get("minItems", 0): raise SchemaError(f"{path}: expected at least {schema['minItems']} items") for i, item in enumerate(value): validate(item, schema["items"], f"{path}[{i}]") elif t == "string": if not isinstance(value, str): raise SchemaError(f"{path}: expected string, got {type(value).__name__}") if len(value) < schema.get("minLength", 0): raise SchemaError(f"{path}: string shorter than {schema['minLength']}") if "enum" in schema and value not in schema["enum"]: raise SchemaError(f"{path}: {value!r} not in {schema['enum']}") elif t == "integer": if not isinstance(value, int) or isinstance(value, bool): raise SchemaError(f"{path}: expected integer, got {type(value).__name__}") _range(value, schema, path) elif t == "number": if not isinstance(value, (int, float)) or isinstance(value, bool): raise SchemaError(f"{path}: expected number, got {type(value).__name__}") _range(value, schema, path) def _range(value, schema, path): if "minimum" in schema and value < schema["minimum"]: raise SchemaError(f"{path}: {value} below minimum {schema['minimum']}") if "maximum" in schema and value > schema["maximum"]: raise SchemaError(f"{path}: {value} above maximum {schema['maximum']}") FENCE = re.compile(r"```(?:json)?\s*(.*?)```", re.S) def extract_json(text: str): """Fenced block first; otherwise the first balanced {...} span.""" m = FENCE.search(text) if m: text = m.group(1) start = text.find("{") if start < 0: raise SchemaError("no JSON object in response") depth = 0 for i in range(start, len(text)): if text[i] == "{": depth += 1 elif text[i] == "}": depth -= 1 if depth == 0: return json.loads(text[start:i + 1]) raise SchemaError("unbalanced braces in response") # ---- simulator (stands in for the API call) ---------------------------------------------- GOOD = {"vendor": "Acme", "total": 41.5, "currency": "USD", "line_items": [{"sku": "A-1", "qty": 2}]} FAILURES = ["fence", "prose", "type", "missing", "range", "extra", "truncated"] def call_model(messages: list[dict], rng: random.Random) -> str: """Deterministic stand-in. First attempt fails 40% of the time; after an error message is appended the failure rate halves, which is roughly what error-feedback buys in practice (the exact rate is a parameter here, not a measurement).""" attempt = sum(1 for m in messages if m["role"] == "assistant") + 1 p_fail = 0.40 / (2 ** (attempt - 1)) if rng.random() > p_fail: return json.dumps(GOOD) mode = rng.choice(FAILURES) bad = json.loads(json.dumps(GOOD)) if mode == "fence": return "```json\n" + json.dumps(GOOD) + "\n```" if mode == "prose": return "Here is the extraction:\n" + json.dumps(GOOD) + "\nLet me know if you need more." if mode == "type": bad["total"] = "41.50" elif mode == "missing": del bad["currency"] elif mode == "range": bad["line_items"][0]["qty"] = 0 elif mode == "extra": bad["notes"] = "n/a" elif mode == "truncated": return json.dumps(GOOD)[:-12] return json.dumps(bad) # ---- the loop ----------------------------------------------------------------------------- def extract_with_retries(prompt: str, schema: dict, rng: random.Random, max_attempts: int = 3): messages = [{"role": "user", "content": prompt}] for attempt in range(1, max_attempts + 1): raw = call_model(messages, rng) try: obj = extract_json(raw) validate(obj, schema) return obj, attempt, None except (SchemaError, json.JSONDecodeError) as e: err = f"{type(e).__name__}: {e}" messages.append({"role": "assistant", "content": raw}) messages.append({"role": "user", "content": f"That response failed validation: {err}. " "Return only a JSON object matching the schema, no prose, no code fence."}) return None, max_attempts, err def main() -> None: rng = random.Random(7) n = 2000 by_attempt = {1: 0, 2: 0, 3: 0} failed = 0 raw_ok = 0 for _ in range(n): # raw first-shot success = extractor + validator on a single attempt obj, attempts, err = extract_with_retries("Extract the invoice.", SCHEMA, rng) if obj is None: failed += 1 else: by_attempt[attempts] += 1 if attempts == 1: raw_ok += 1 print(f"runs: {n}") print(f"valid on attempt 1: {by_attempt[1]} ({by_attempt[1] / n:.1%})") print(f"valid on attempt 2: {by_attempt[2]} ({by_attempt[2] / n:.1%})") print(f"valid on attempt 3: {by_attempt[3]} ({by_attempt[3] / n:.1%})") print(f"gave up after 3: {failed} ({failed / n:.2%})") calls = by_attempt[1] + 2 * by_attempt[2] + 3 * (by_attempt[3] + failed) print(f"mean model calls per item: {calls / n:.3f}") if __name__ == "__main__": main()