DEV Community

Torkian
Torkian

Posted on • Edited on

Make Your Agent Return Data, Not Prose — Structured Outputs with NVIDIA NIM

#nvidia #ai #python #tutorial

For eight parts the agent has ended every turn the same way: it prints a sentence and we read it. That's fine for a demo and a dead end for a product. The moment you want to build on the agent — render it in a UI, score it in a test, return it from an API — prose fights you. You're reduced to scraping text with regexes and hoping the wording doesn't change.

This post fixes that with one capability: the final answer becomes a validated JSON object with a fixed shape. Same agent, same tools, same memory — but now it returns {"status": "...", "answer": "...", "items": [...], ...} instead of a paragraph. That contract is the hinge the rest of this series turns on: the tracing chapter logs it, the evaluation chapter asserts on its fields, the deployment chapter returns it as the HTTP response body.

There's an honest catch, and it's the real lesson of this post. The tidy answer would be "use response_format with a JSON schema and let the API enforce it." On the hosted API Catalog with an open model, that's not dependable — strict schema mode is a server-side feature that open-model endpoints implement inconsistently, and it interacts badly with tool calling. So we don't lean on it. We ask for JSON in the prompt and then do the grown-up thing: parse it, validate it ourselves, and repair it once if it's wrong. No framework.

(If you self-host NIM — Part 4 — NVIDIA does document server-side structured generation: NIM 1.x accepts extra_body={"nvext": {"guided_json": schema}} for constrained decoding, and the newest releases move to OpenAI-compatible response_format JSON mode. Either way, NVIDIA's own docs tell you to validate the response client-side — which is exactly the ladder this post builds. Docs: https://docs.nvidia.com/nim/large-language-models/latest/structured-generation.html)

I'm B Torkian, NVIDIA Developer Champion at USC. Part 9 of the series.

What you're adding 

Workshop 8:  final answer -> a string you print
Workshop 9:  final answer -> JSON text -> parse -> validate -> (repair once) -> a dict you can use
Enter fullscreen mode Exit fullscreen mode

The Workshop 8 behavior is unchanged — tools, multi-turn memory, trim-by-turns, the streaming accumulator. The one behavioral change is at the final-answer boundary: instead of returning message.content, we turn it into a validated dict. (We also do a small bit of housekeeping — Workshop 8 had two near-identical loops, one in chat() and one in stream(); we factor them into a single shared loop so the JSON finalization lives in exactly one place. More on that in Step 4. One honest trade: live token-by-token display is dropped for the final answer — half-formed JSON is useless to show — while the wire-level streaming and tool-call fragment reassembly from Part 8 stay.)

Step 1 — Decide the contract

Before any code, decide what the agent must return. For the campus assistant, six keys cover every case:

{
  "status": "answered",
  "answer": "The USC AI Club meets every Thursday at 5 PM in the engineering building, room 204.",
  "category": "campus_event",
  "items": [
    {"name": "USC AI Club meeting", "day": "Thursday", "time": "5 PM", "location": "engineering building, room 204"}
  ],
  "missing": [],
  "sources": ["The USC AI Club meets every Thursday at 5 PM in the engineering building, room 204."]
}
Enter fullscreen mode Exit fullscreen mode

status and category are enums (a fixed set of allowed values) so downstream code can branch on them. items is the machine-readable payload. missing names anything the user asked for that wasn't found. sources is the grounding — the exact knowledge-base lines used, which the guardrails chapter's spirit lives on in.

Step 2 — Parse, validate, repair (the part that matters)

We treat the model's output as untrusted. Three small functions, plain Python:

STATUSES = {"answered", "not_found", "needs_clarification"}
CATEGORIES = {"campus_event", "campus_hours", "campus_resource", "comparison", "refusal"}
REQUIRED_KEYS = ("status", "answer", "category", "items", "missing", "sources")

def parse_json_object(text: str) -> dict:
    # Models wrap JSON in prose or fenced code blocks. Take the {...} span.
    start, end = text.find("{"), text.rfind("}")
    if start == -1 or end == -1 or end < start:
        raise ValueError("no JSON object found")
    return json.loads(text[start:end + 1])

def validate_answer(data) -> list:
    if not isinstance(data, dict):
        return ["response is not a JSON object"]
    errors = []
    for key in REQUIRED_KEYS:
        if key not in data:
            errors.append(f"missing required key: {key}")
    if data.get("status") not in STATUSES:
        errors.append(f"status must be one of {sorted(STATUSES)}")
    if data.get("category") not in CATEGORIES:
        errors.append(f"category must be one of {sorted(CATEGORIES)}")
    if "answer" in data and not isinstance(data["answer"], str):
        errors.append("answer must be a string")
    for key in ("items", "missing", "sources"):
        if key in data and not isinstance(data[key], list):
            errors.append(f"{key} must be a list")
    if isinstance(data.get("items"), list) and not all(isinstance(it, dict) for it in data["items"]):
        errors.append("each entry in items must be an object")
    return errors
Enter fullscreen mode Exit fullscreen mode

If validation fails, make exactly one repair attempt — hand the broken output back to the model at temperature=0 with the list of problems and ask for a clean object:

def repair_answer_json(raw_text: str, errors: list) -> dict | None:
    # (prompt abridged — the repo version also restates the required keys and enums)
    try:
        fix = client.chat.completions.create(
            model=MODEL, temperature=0, max_tokens=800,
            messages=[
                {"role": "system", "content": "/no_think\n\nYou fix malformed JSON. Return ONLY a valid JSON object."},
                {"role": "user", "content": f"Problems: {errors}. Preserve all facts. Original:\n{raw_text}\n\nReturn corrected JSON only."},
            ],
        )
        data = parse_json_object(fix.choices[0].message.content or "")
    except Exception:          # API error or unparseable — fall back deterministically
        return None
    return data if not validate_answer(data) else None
Enter fullscreen mode Exit fullscreen mode

Notice what goes into that repair prompt: the specific validation failures ("status must be one of …", "missing key: category"), not a vague "return valid JSON." Telling the model exactly which field failed, and why, is most of the reason a single retry usually lands.

And if even the repair fails, return a deterministic error object so callers never crash on a surprise:

def format_error(missing: str = "valid_json") -> dict:
    return {"status": "needs_clarification",
            "answer": "I couldn't produce a valid structured response. Please ask again.",
            "category": "refusal", "items": [], "missing": [missing], "sources": []}
Enter fullscreen mode Exit fullscreen mode

Parse → validate → repair once → deterministic fallback. That four-step ladder is what makes structured output safe in production without a schema-enforcement framework. (The Part 8 step-limit fallback becomes structured too — in the repo, format_error takes an optional cause-specific answer, so even the give-up path honors the contract.)

"Why not Pydantic?" Fair question — in production Python, Pydantic is the standard way to do exactly this: define the contract as a class, call model_validate_json, and get typed objects plus precise error messages to feed the repair prompt. We hand-rolled validate_answer for the same reason this series hand-rolled retrieval and the agent loop: so you can see what the tool automates. It's ~30 lines, and now Pydantic will never be magic to you. Swapping it in is a genuinely good exercise — class Answer(BaseModel) with Literal types for the enums, and the ladder's parse + validate steps collapse into one Answer.model_validate_json(raw) inside a try. The repair-once and deterministic-fallback steps stay: no validator, Pydantic included, can fix output that never arrived.

Step 3 — Tell the model the format

The tool guidance from Workshops 7–8 is unchanged. We append a FINAL ANSWER FORMAT block to the system prompt that pins the contract:

SYSTEM_PROMPT = (
    "/no_think\n\n"
    "...all the Workshop 7-8 tool + memory guidance...\n\n"
    "FINAL ANSWER FORMAT. When you are done using tools, your final reply MUST be a "
    "single JSON object and NOTHING else — no prose, no code fences. Use exactly these "
    'keys: status (answered|not_found|needs_clarification), answer (string), category '
    "(campus_event|campus_hours|campus_resource|comparison|refusal), items (list), "
    "missing (list), sources (list)."
)
Enter fullscreen mode Exit fullscreen mode

Tool-calling turns are unaffected — the model still emits tool_calls while it gathers facts. Only the final reply has to be JSON.

Step 4 — One shared loop, JSON at the finish line

ChatSession keeps Workshop 8's control flow exactly. The housekeeping mentioned above: Workshop 8 had the agent loop written out twice — once in chat(), once in stream(). We pull the single model call into _complete(stream) (which returns the text, the reassembled tool calls, and the assistant message, for both streaming and non-streaming) and the loop itself into one shared _run_turn(user_message, stream). Now chat() and stream() are one-line delegators:

def chat(self, user_message: str) -> dict:
    return self._run_turn(user_message, stream=False)

def stream(self, user_message: str) -> dict:
    return self._run_turn(user_message, stream=True)
Enter fullscreen mode Exit fullscreen mode

_run_turn is the Workshop 8 loop — run tools until none remain. The only new step is the final branch: instead of returning the text, finalize it into a validated dict with _finalize_json.

def _finalize_json(self, raw_text: str) -> dict:
    try:
        data = parse_json_object(raw_text)
        errors = validate_answer(data)
    except (ValueError, json.JSONDecodeError):
        data, errors = None, ["response was not valid JSON"]
    if errors:
        repaired = repair_answer_json(raw_text, errors)
        data = repaired if repaired is not None else format_error()
    return data

# the final branch of _run_turn, once the model stops calling tools:
#     data = self._finalize_json(text)
#     self.messages.append({"role": "assistant", "content": json.dumps(data)})
#     self._trim()
#     return data
Enter fullscreen mode Exit fullscreen mode

Note we store the canonical json.dumps(data) in history, not the model's raw text — so the next turn's memory is clean, validated JSON too. Both chat() and stream() return a dict now.

Step 5 — Run it 

session = ChatSession(verbose=True)
for q in [
    "When does the USC AI Club meet?",                          # answered, campus_event
    "How many days until that?",                                # memory + tool
    "Which is sooner, that meeting or the AI/ML office hours?",  # comparison
    "What is the campus wifi password?",                        # not_found / refusal
]:
    result = session.chat(q)
    print(json.dumps(result, indent=2))
Enter fullscreen mode Exit fullscreen mode

You get back four well-formed objects. The wifi question is the satisfying one — instead of a refusal sentence, you get a typed refusal your code can act on:

{"status": "not_found", "answer": "I don't have that information — check with the USC AI Club.",
 "category": "refusal", "items": [], "missing": ["USC campus wifi password"], "sources": []}
Enter fullscreen mode Exit fullscreen mode

Memory and multi-step reasoning are untouched — "how many days until that?" still resolves "that", and the comparison still calls the tool once per day.

Step 6 — What you actually built

  • Workshop 1 gave it a brain. 2 memory of facts. 3 judgment. 4 portability. 5 hands. 6 a plan. 7 memory of the conversation. 8 a real-time voice.
  • Workshop 9 gave it a contract — output other software can consume.

That last one is the quiet turning point of the series. Everything before made the agent smarter; this makes it integratable. And it sets up the back half of this series:

  • Next — Traces: log each turn's tools, latency, and this final object as JSONL, so you can see what the agent did after the fact.
  • Then — Evals: with traces on disk and a fixed contract, you can finally test the agent like software — replay the questions and assert status, category, and that missing fires on the wifi question.
  • Then — Durable sessions and Deploy: persist histories and serve this exact JSON body behind an HTTP API.

The agent is still the same while loop. We just taught its last word to be data.

Get the code

Repo: github.com/torkian/nvidia-nim-workshop
One-click Colab: Open part9_structured_output.ipynb
Local Python: part9_structured_output.py in the repo (python3 part9_structured_output.py after pip install -r requirements.txt).

MIT licensed. I run this at USC — fork it, swap the knowledge base, the tools, and the contract for your school, your club, your project.

The full series

A consolidated long-form version of the whole series is on Medium for anyone who'd rather read it in one sitting.

Top comments (4)

Collapse
 
dipankar_sarkar profile image
Dipankar Sarkar •

The parse-validate-repair ladder is the right call, and the tool-calling interaction you flag is the real trap. Strict response_format fights tool use because it forces schema-shaped JSON on every turn, so the model can't emit a tool call mid-run. The pattern that sidestepped it for us: keep tools free-form during the reasoning turns, then do one final deterministic format turn with tools disabled and the schema enforced. Two phases, not one contract stretched across the whole run. One thing worth adding to the repair step: feed the validator's actual error back into the repair prompt, not just 'return valid JSON'. The single retry lands far more often when it knows which field failed and why.

Collapse
 
torkian profile image
Torkian •

Thanks Dipankar, really good points. That two phase split (tools free during the reasoning turns, then one format turn with tools off and the schema enforced) is the cleanest way to do it when you are not streaming, and I almost shipped exactly that. Two reasons I stayed with prompt plus validate instead. First, Part 8 of the series streams, and a dedicated final format turn there means a full blocking generation before the first token, which kills time to first token. Second, on the open models this series targets, even that final enforced turn is not fully dependable, so I end up validating myself either way. On flagship hosted models your two phase is probably the better call.

And you nailed the repair step, that is exactly what the code does here. It feeds the specific validate errors into the retry ("status must be one of ...", "missing key: category") rather than a generic "return valid JSON". You actually pushed me to make that point explicit in the post itself, so thank you for that.

Collapse
 
arvavit profile image
Vadym Arnaut •

"strict schema mode is a server-side feature open-model endpoints implement inconsistently" matches what we saw running audit agents against llama-family models. We parse a findings list out of free-form output (severity, file:line, fix) and still validate/repair the shape ourselves before anything downstream trusts it. Trusting schema enforcement instead of your own validation layer is the footgun once you're off flagship hosted models.

Collapse
 
torkian profile image
Torkian •

Exactly, Vadym. "Trust your own validation layer, not the endpoint's schema enforcement" is the whole lesson once you are off the flagship hosted models. Your audit agent setup (parsing severity, file:line, and fix out of free form output, then validating the shape before anything downstream trusts it) is a perfect real world version of it. One NVIDIA specific footnote for anyone self hosting: NIM does document server side constrained decoding (guided_json via nvext) and it is genuinely useful, but even NVIDIA's own docs tell you to validate client side afterward. Same footgun avoidance you are describing. Validate downstream, always.