AI Engineering · 60 min

Validate LLM Output With a Schema

Build a robust function that forces an LLM to return structured JSON, validates it against a schema, and retries on failure — so you never blindly trust raw model output in production.

Get the starter project

Fork or clone it (Ruby/Rails, Python/FastAPI or TypeScript) and make the failing tests pass.

Problem

LLMs are probabilistic text generators, not typed APIs. Ask a model for JSON and, most of the time, it obeys — but "most of the time" is exactly what breaks in production. The same prompt can return: - Prose wrapped around the JSON: `Sure! Here is the object you asked for: {...}` - Markdown fences: ```json ... ``` - A missing required field, or a field with the wrong type (`"3"` instead of `3`) - An enum value you never defined (`"medium-hard"`) - Truncated or invalid JSON when the response hits a token limit The naive approach — `JSON.parse(response)` and trust the result — is a landmine. It works in the demo and blows up on the 200th real request, often deep inside code that assumed the shape was guaranteed. In this lab you will build an **anti-fragile** boundary around the model: define a **JSON Schema** for the exact output you expect, instruct the model to return only that JSON, extract and parse it defensively, **validate** it against the schema, and **retry with feedback** when validation fails (telling the model exactly what was wrong). After a bounded number of attempts, you return a clear failure — never a half-populated object. The caller decides what to do next.

Objectives

Define a precise JSON Schema (fields, types, required, enums) for a real use case
Write a prompt that instructs the model to return only valid JSON matching the schema
Extract and parse JSON defensively, tolerating ```json fences and surrounding prose
Validate parsed output against the schema and build a clear, machine-readable error
Implement a retry loop that feeds the validation error back to the model, with a max attempt cap and backoff
Return a typed failure (Result-style) after exhausting retries, instead of a partial object
Inject the LLM client so the function is deterministic and testable with a stub

Prerequisites

Access to any LLM API (OpenAI, Anthropic, a local model, etc.) — or a fake client for the tests
A JSON Schema validation library for your language (e.g. ajv in JS, jsonschema in Python, json_schemer in Ruby, everit/networknt in Java)
The programming language of your choice (the lab is stack-agnostic; examples are in pseudocode + JSON)
A test runner you already use, to write the stubbed-LLM tests in the final step

Treat every LLM call as an untrusted network boundary. The model is a helpful but
unreliable third party: it will usually give you what you asked for, and your job
is to make the unusual cases safe and self-correcting instead of catastrophic.

The pattern you will build has five moving parts:

A schema — the single source of truth for the shape you accept.
A prompt — that asks for exactly that shape, and nothing else.
A defensive parse — that survives fences and stray prose.
A validate-and-retry loop — that turns a validation error into a corrective
follow-up prompt, bounded by a max attempt count.
A final failure — an honest error when the model can't comply, so the caller
stays in control.

Throughout the lab we use one concrete use case: extracting structured metadata
{ title, tags[], difficulty } from a free-form block of text. Everything
generalizes to any structured-output task (entity extraction, classification,
function arguments, form filling).

Steps

Define the output schema and the use case
Start from the shape you want, not from the model. Write a JSON Schema that
names every field, its type, which fields are required, and any enum
constraints. This schema is the contract: the prompt describes it, the
validator enforces it, and the tests assert against it.

Our use case: given an arbitrary paragraph, extract a title, a list of
tags, and a difficulty.
```
{
  "$schema": "https://json-schema.org/draft/2020-12/schema",
  "title": "ExtractedMetadata",
  "type": "object",
  "properties": {
    "title":      { "type": "string", "minLength": 1, "maxLength": 120 },
    "tags":       {
      "type": "array",
      "items": { "type": "string", "minLength": 1 },
      "minItems": 1,
      "maxItems": 8
    },
    "difficulty": { "type": "string", "enum": ["beginner", "intermediate", "advanced"] }
  },
  "required": ["title", "tags", "difficulty"],
  "additionalProperties": false
}
```
Note the deliberate strictness:
- required lists every field, so a missing field is a hard error.
- enum on difficulty rejects creative values like "medium".
- additionalProperties: false rejects any extra keys the model invents.
- Bounds (minItems, maxLength) catch empty or runaway outputs.
Keep this schema in a constant your code can reference — you will pass it to
both the validator and (in the next step) the prompt.
Write a prompt that demands strict JSON
The prompt is your first line of defense. Make the contract explicit: show the
model the schema, forbid prose and fences, and give a concrete example of a
valid response. If your provider supports structured output (response_format
json_schema), tool/function calling, or a JSON mode, prefer those — they
make well-formed JSON far more likely. But never rely on them alone: still
validate, because "valid JSON" is not the same as "matches your schema".

A portable prompt template:
```
You are a metadata extraction function. Extract structured data from the INPUT.

Return ONLY a single JSON object that conforms to this JSON Schema.
Do not include markdown code fences, comments, or any prose before or after.

JSON Schema:
<SCHEMA_HERE>

Example of a valid response:
{"title":"Intro to Vectors","tags":["math","linear-algebra"],"difficulty":"beginner"}

INPUT:
<TEXT_HERE>
```
Design principles:
- Frame the model as a function, not a chatbot — this suppresses chit-chat.
- Inline the schema so the field names, types, and enums are unambiguous.
- Show one valid example — models imitate format strongly.
- Explicitly forbid fences and prose — but still handle them in step 3,
  because the model will occasionally ignore you.
If you use structured output / tool calling, pass the same schema to the API's
native parameter and keep it for validation. Defense in depth.
Extract and parse the response defensively
Even a perfect prompt gets imperfect responses. Before validating, you must
reliably pull a JSON object out of whatever the model returned. Handle, in
order: a clean object, an object wrapped in ```json fences, and an object
buried in surrounding prose.

A robust extraction routine:
```
function extract_json(raw_text):
    text = raw_text.strip()

    # 1) Strip a fenced code block if present: ```json ... ``` or ``` ... ```
    fence = match text against /```(?:json)?\s*([\s\S]*?)```/i
    if fence found:
        text = fence.group(1).strip()

    # 2) Try a direct parse first (the happy path)
    try:
        return Ok(parse_json(text))
    except ParseError:
        pass

    # 3) Fall back to the first balanced {...} object in the text
    candidate = slice from first '{' to its matching '}'
    if candidate exists:
        try:
            return Ok(parse_json(candidate))
        except ParseError:
            pass

    return Err("no parseable JSON object found")
```
Notes:
- Prefer a balanced-brace scan over a greedy {.*} regex, so nested objects
  don't get truncated. Track depth: +1 on {, -1 on }, stop at zero.
- This function only guarantees parseable JSON, not valid JSON per your
  schema. That is the next step's job — keep the two concerns separate.
- Return a Result (Ok/Err), not an exception, so the retry loop in step 4
  can treat a parse failure the same way it treats a validation failure: as a
  reason to retry with feedback.

Validate against the schema and retry with feedback

Now wire the pieces into a single loop. Each iteration: call the model, extract
JSON, validate against the schema. On any failure — parse or validation —
build a precise error message and feed it back into the next prompt, so the
model corrects itself. Cap the attempts (e.g. 3) and back off between them.

function get_validated_output(input_text, llm, schema, max_attempts = 3):
    prompt = build_prompt(schema, input_text)   # from step 2
    last_error = null

    for attempt in 1..max_attempts:
        raw = llm.complete(prompt)

        parsed = extract_json(raw)               # from step 3
        if parsed is Err:
            last_error = "Response was not parseable JSON: " + parsed.error
        else:
            errors = validate(parsed.value, schema)   # JSON Schema lib
            if errors is empty:
                return Ok(parsed.value)               # success
            last_error = format_validation_errors(errors)

        # Build a corrective follow-up prompt with the exact failure
        prompt = build_prompt(schema, input_text)
                 + "\n\nYour previous output was rejected.\n"
                 + "Error: " + last_error + "\n"
                 + "Return corrected JSON that satisfies the schema. JSON only."

        sleep(backoff(attempt))   # e.g. 0.5s, 1s, 2s — exponential

    return Err({ reason: "validation_failed",
                 attempts: max_attempts,
                 detail: last_error })

What makes this anti-fragile:

The error is specific. format_validation_errors should say
"tags: expected array, got string" or "difficulty: 'medium' is not one of [beginner, intermediate, advanced]" — not a generic "invalid". Specific
feedback dramatically raises the odds the retry succeeds.
Parse and validation failures share one path. Both set last_error and
trigger the same corrective retry.
Attempts are bounded. No infinite loop, no runaway API bill.
Backoff spaces out retries (also helps with rate limits / transient errors).

A typical run: attempt 1 returns difficulty: "medium" → rejected → the follow-up
names the enum → attempt 2 returns "intermediate" → validates → Ok.

Handle final failure and test with a stubbed LLM
The loop must end honestly. When every attempt fails, return the typed Err
from step 4 — never a partial or fabricated object. The caller inspects the
Result and decides: surface an error to the user, fall back to a default, queue
for human review, etc. Your function's contract is simple and total: it
always returns either validated output or a clear failure.

Because you injected the llm client, you can test all three paths
deterministically with a stub — no network, no flakiness, no cost. A stub is
just an object whose complete() returns pre-scripted responses in order.
```
test "succeeds on the first valid response":
    llm = StubLLM(responses: [
        '{"title":"Vectors 101","tags":["math"],"difficulty":"beginner"}'
    ])
    result = get_validated_output("...", llm, SCHEMA)
    assert result is Ok
    assert result.value.difficulty == "beginner"
    assert llm.call_count == 1

test "retries after an invalid response, then succeeds":
    llm = StubLLM(responses: [
        '```json\n{"title":"Vectors 101","tags":[],"difficulty":"medium"}\n```', # invalid: empty tags + bad enum
        '{"title":"Vectors 101","tags":["math"],"difficulty":"intermediate"}'    # valid
    ])
    result = get_validated_output("...", llm, SCHEMA, max_attempts: 3)
    assert result is Ok
    assert result.value.difficulty == "intermediate"
    assert llm.call_count == 2
    # optional: assert the 2nd prompt contained the validation error text

test "returns a failure after exhausting all attempts":
    llm = StubLLM(responses: [
        'I cannot help with that.',   # not JSON
        '{"title":"X"}',              # missing required fields
        '{"title":"X","tags":"nope","difficulty":"medium"}'  # wrong types + bad enum
    ])
    result = get_validated_output("...", llm, SCHEMA, max_attempts: 3)
    assert result is Err
    assert result.error.reason == "validation_failed"
    assert result.error.attempts == 3
    assert llm.call_count == 3
```
Tip: stub sleep/backoff too (inject a no-op clock), so the tests run
instantly. The first response in the retry test deliberately includes ```json
fences to prove your extractor from step 3 is exercised end-to-end.

Submission criteria

You are done when:
1. You have a single function that, given input text and an LLM client, always
  returns either validated output (matching the schema) or a typed failure —
  never a partial object and never a raw, unvalidated parse.
2. The function defines and enforces a JSON Schema, prompts for strict JSON,
  extracts defensively (tolerating fences and prose), validates, and retries
  with the specific validation error, bounded by a max attempt count with backoff.
3. You have automated tests with a stubbed LLM covering all three paths:
  success, retry-then-success, and final failure after exhausting retries,
  asserting the call count for each.

Validate LLM Output With a Schema

Problem

Objectives

Prerequisites

Steps

Submission criteria