How to Implement Structured AI Agent Output with Gemini and Pydantic
How to Implement Structured AI Agent Output with Gemini and Pydantic
Structured AI agent output is achieved by using constrained decoding and formal schemas like Pydantic to force LLMs to return valid, type-safe JSON. This approach eliminates non-deterministic text formatting errors and ensures seamless integration between probabilistic models and deterministic software systems. By defining a response schema, developers can guarantee that AI responses adhere to specific data structures required by backend APIs.
It was 3:14 AM on a Tuesday when my PagerDuty went off. The error log was a mess of 500 Internal Server Errors originating from my FastAPI backend. The culprit? A background worker responsible for processing customer support tickets using an LLM-based agent. The agent had decided, for the first time in three weeks, to wrap its JSON response in a triple-backtick Markdown block and prefix it with the phrase "Sure, I can help with that! Here is the data:".
My Pydantic validator, expecting a clean JSON string to hydrate a TicketAnalysis model, choked. The entire processing pipeline stalled. By the time I manually cleared the queue and patched the parser, we had dropped 1,200 events. This wasn't a failure of the model's logic; the AI had actually correctly identified the customer's sentiment and intent. It was a failure of communication between a non-deterministic probabilistic engine and a deterministic software system. I realized then that relying on raw LLM responses—even with "Please return only JSON" in the prompt—was a recipe for technical debt and operational instability. Implementing structured AI agent output is the only way to ensure production-grade reliability.
In this post, I want to break down exactly why raw text responses are a liability for production AI agents and how I rebuilt my stack using Gemini 1.5 Pro’s structured output capabilities and Pydantic to ensure type safety. I’ll share the exact patterns that reduced my parsing error rate from 4.2% to 0% over the last 60 days.
Why String-Based Prompt Engineering Fails in Production
Relying on raw text prompts for JSON formatting leads to non-deterministic errors that break production parsers and increase system latency. When I first started building AI agents, I followed the common path: prompt engineering. I would spend hours tweaking a system instruction, adding phrases like "Output must be valid JSON," "Do not include any preamble," or "Strictly adhere to the following schema." For a while, this worked. But as I scaled, the edge cases surfaced. LLMs are trained to be helpful conversationalists, and occasionally, that "helpfulness" manifests as unwanted text.
I tried the regex approach next. I wrote complex patterns to extract anything between { and }. This was a nightmare to maintain. What happens when the LLM includes a nested JSON object? What happens when it hallucinates a trailing comma that makes the JSON invalid? I was writing more code to clean the AI's output than I was writing to implement actual business logic. This increased the complexity of my Go-based gateway, which I previously wrote about in my guide to building a high performance LLM API gateway. The overhead of sanitizing strings was adding 50-100ms of latency per request, which is unacceptable in a high-throughput system.
The fundamental problem is that "JSON-like text" is not "JSON data." In a production environment, your code needs to know—with 100% certainty—that the status field is an enum of ['open', 'closed', 'pending'] and not a creative variation like 'currently_active'. Raw responses don't give you that guarantee, making structured AI agent output a technical necessity.
How to Implement Constrained Decoding with Gemini 1.5 Pro
Constrained decoding uses a formal schema to restrict the model's sampling logic, ensuring every generated token adheres to the defined structure during the inference phase. The breakthrough came when I migrated my agents to use Gemini 1.5 Pro’s response_mime_type and response_schema parameters. Instead of begging the model to format its output correctly, I now define a formal schema that the model's decoding process must follow. The model doesn't just "try" to output JSON; the underlying sampling logic is restricted to only generate tokens that satisfy the schema.
Here is the core pattern I now use for every agent I build. We define our data structure using Pydantic, which serves as our source of truth for both the Python application and the LLM schema.
from pydantic import BaseModel, Field
from typing import List, Optional
import enum
class Sentiment(str, enum.Enum):
POSITIVE = "positive"
NEUTRAL = "neutral"
NEGATIVE = "negative"
class TicketCategory(str, enum.Enum):
BILLING = "billing"
TECHNICAL = "technical"
FEATURES = "feature_request"
OTHER = "other"
class TicketAnalysis(BaseModel):
summary: str = Field(description="A concise summary of the user's issue")
sentiment: Sentiment
category: TicketCategory
confidence_score: float = Field(ge=0, le=1)
suggested_tags: List[str]
requires_escalation: bool
By using Pydantic, I get local validation for free. But the real magic happens when I pass this schema directly to the Gemini API. This eliminates the need for "Please return JSON" prompts entirely, which, as I noted in my article on reducing LLM API costs, actually saves a significant amount of input tokens over thousands of requests.
Connecting FastAPI to Gemini for Type-Safe AI Responses
Integrating Pydantic models with the Gemini SDK allows for automatic JSON schema generation and immediate response validation within your backend services. When integrating this into a FastAPI service, I use the google-generativeai Python SDK. The key is to convert the Pydantic model into a JSON schema format that the Gemini API understands. Since Gemini supports a subset of the JSON Schema spec, I wrote a helper function to bridge the two.
import google.generativeai as genai
import os
# Initialize the model with the schema
model = genai.GenerativeModel('gemini-1.5-pro')
def analyze_ticket(ticket_text: str) -> TicketAnalysis:
prompt = f"Analyze the following support ticket: {ticket_text}"
# We define the response schema here
response = model.generate_content(
prompt,
generation_config=genai.GenerationConfig(
response_mime_type="application/json",
response_schema=TicketAnalysis
)
)
# Because the model is constrained, response.text is guaranteed
# to be valid JSON matching our TicketAnalysis model.
return TicketAnalysis.model_validate_json(response.text)
This approach changed my life. I no longer have to check if response.text starts with a backtick. I don't have to worry about the model adding conversational fluff. The response.text is strictly the JSON object. If the model fails to generate a valid response (which is rare with constrained decoding), the SDK throws a clear error that I can catch and retry, rather than a silent parsing failure deep in my business logic.
How to Enforce Complex Enums and Constraints in AI Schemas
Enforcing enums at the sampling layer makes it mathematically impossible for the LLM to hallucinate invalid field values during generation. In the TicketAnalysis model above, I used a Python Enum. When this is passed to Gemini via the response_schema, the model is physically unable to output a string that isn't in that enum list. This is a massive win. I used to see models return "billing_issue" when the schema expected "billing". With constrained decoding, that hallucination is mathematically impossible during the token selection phase.
According to the official Gemini API documentation, providing a clear description in the Pydantic Field is crucial. The model uses these descriptions as semantic hints to understand what data should go into each field. I've found that spending time on these descriptions is far more effective than writing long system prompts for structured AI agent output.
Performance Benchmarks: Comparing Structured vs. Unstructured AI Outputs
Using schema-constrained outputs increases parsing success rates to 100% while reducing output token usage by approximately 23%. I ran a benchmark over 1,000 simulated support tickets to compare the reliability of "Prompt-based JSON" versus "Schema-constrained JSON." The results were stark.
| Metric | Prompt-based (Raw) | Schema-constrained |
|---|---|---|
| Parsing Success Rate | 95.8% | 100% |
| Avg. Output Tokens | 245 | 188 |
| Latency (p95) | 1.8s | 1.6s |
| Schema Hallucinations | 12 instances | 0 instances |
The reduction in output tokens is particularly interesting. Because the model doesn't output Markdown markers, preambles, or explanations, the payload is leaner. In a high-volume system, a 23% reduction in output tokens translates directly to lower costs and faster time-to-first-token. The slight improvement in latency is likely due to the model not having to "think" about how to format the response—the format is enforced by the sampler.
Managing Nested Data Structures and AI Tool Calling
Including a reasoning field within a structured schema preserves the model's chain-of-thought capabilities without sacrificing data integrity or parseability. Real-world agents are rarely as simple as a single flat JSON object. Often, my agents need to perform "Chain of Thought" reasoning before arriving at a conclusion. Initially, I worried that forcing structured output would kill the model's ability to reason. If the model can only output the final JSON, where does the internal monologue go?
The solution is to include a reasoning or thought_process field in the schema itself. This forces the model to "show its work" inside the JSON structure.
class AdvancedAnalysis(BaseModel):
internal_monologue: str = Field(description="Step-by-step reasoning about the ticket")
root_cause_identified: str
resolution_steps: List[str]
final_sentiment: Sentiment
This keeps the reasoning process structured and parseable. I can store the internal_monologue in my logs for debugging while only presenting the resolution_steps to the end-user. This pattern has been instrumental in debugging why the agent might have miscategorized a complex technical issue. Furthermore, when using Gemini's Function Calling, the same principles apply. I've stopped using "Agent Frameworks" that try to parse these calls manually, relying instead on structured AI agent output to drive my tool execution loop.
Key Takeaways for Building Reliable AI Agents
Shifting from prompting to schema-driven development is the most effective way to ensure AI reliability and reduce operational overhead. The transition from raw text to structured data was a turning point in my AI engineering practice. Here are the core lessons I've integrated into my workflow:
- Stop Prompting for Format: Prompting is for intent and logic. Schemas are for structure. Never use a prompt to do a schema's job.
- Pydantic is Mandatory: Don't pass raw dictionaries to your LLM integration. Use Pydantic models to define your contract. It provides validation on both ends: it documents the schema for the LLM and validates the response for your code.
- Constrained Decoding > Regex: If your model provider supports a
response_schemaor "JSON Mode," use it. It is significantly more reliable than trying to extract JSON from a string using regular expressions. - Descriptions Matter: The
descriptionfield in your schema is the new "system prompt." Be explicit about what each field represents. - Handle Validation Errors Gracefully: Even with structured output, the model might occasionally hit a safety filter or a max token limit. Always wrap your validation logic in a try-except block and implement a retry strategy.
Additional Resources for AI Engineering
- How to Reduce LLM API Costs Across Multiple Model Providers - Learn how structured output and efficient prompting can lower your monthly API bill.
- Building a High Performance LLM API Gateway with Go and Cloud Run - My architectural guide on handling high-throughput LLM traffic with a custom gateway.
Looking forward, I’m keeping a close eye on the development of more complex constrained decoding techniques, such as context-free grammar (CFG) enforcement. While JSON is a great start, the ability to enforce custom DSLs (Domain Specific Languages) directly at the sampling layer will open up even more possibilities for reliable AI agents. For now, moving away from raw LLM responses to a schema-first approach for structured AI agent output has been the single most impactful change I’ve made to my production stability. If you're still parsing strings with regex, it's time to refactor.
Comments
Post a Comment