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 like llama-3.3-70b, 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

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."]
}

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

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": "You 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

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": []}

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 = (
    "...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)."
)

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)

_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

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))

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": []}

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.