To handle non-deterministic JSON output from LLM tools in production, implement a multi-layered strategy focusing on robust parsing, strict schema validation, and intelligent self-correction/retry mechanisms. This approach ensures data integrity and system resilience against malformed or unexpected LLM responses.
Here's a breakdown of the key steps:
1. Prompt Engineering & Model Configuration:
Clear Instructions: Explicitly instruct the LLM to output valid JSON conforming to a specific schema. Provide the schema directly in the prompt.
Few-Shot Examples: Include one or two perfect JSON examples to guide the model.
Model-Specific Features: Utilize API features like OpenAI's response_format={"type": "json_object"} which significantly improves JSON adherence for supported models (e.g., GPT-4o).
2. Robust Parsing:
Employ parsers tolerant of common LLM errors (e.g., trailing commas, comments, unquoted keys). Libraries like json5 are more forgiving than standard json.
```python
import json5
from pydantic import BaseModel, ValidationError
def parse_robustly(llm_output: str):
try:
return json5.loads(llm_output)
except ValueError as e:
raise ValueError(f"Failed to parse malformed JSON: {e}")
```
3. Schema Validation:
Immediately validate the parsed dictionary against a predefined schema using tools like Pydantic. This ensures the structure and data types are correct.
```python
class AgentAction(BaseModel):
tool_name: str
tool_args: dict
action_id: str
def validate_schema(data: dict) -> AgentAction:
try:
return AgentAction(data)
except ValidationError as e:
raise ValueError(f"Schema validation failed: {e}")
```
4. Self-Correction & Retry Mechanisms:
Retry with Backoff: If parsing or validation fails, implement a retry loop with exponential backoff (e.g., 3 attempts).
* LLM Self-Correction: For persistent errors, feed the original prompt, the LLM's faulty output, and the specific error message back to the LLM, asking it to correct its response.
```python
import time
def execute_llm_call_with_retries(prompt_func, max_retries=3):
for attempt in range(max_retries):
llm_output = prompt_func() # Assume this calls the LLM
try:
parsed_data = parse_robustly(llm_output)
validated_data = validate_schema(parsed_data)
return validated_data
except ValueError as e:
if attempt < max_retries - 1:
print(f"Attempt {attempt+1} failed: {e}. Retrying...")
time.sleep(2 attempt) # Exponential backoff
# Optionally, for LLM self-correction:
# prompt_func.add_correction_context(llm_output, str(e))
else:
raise
raise RuntimeError("LLM output could not be parsed or validated after multiple retries.")
```
Practical Gotcha: Be wary of LLMs hallucinating new fields or altering expected field names even when given a schema. Always log the raw LLM output and the validation errors to identify recurring patterns that might necessitate prompt adjustments or model fine-tuning.