Getting LLMs to Return Valid JSON: A Production Guide

# Real examples from our logs

# Truncation (hit max_tokens)
{"name": "John", "email": "john@exam

# Schema violation (array instead of object)
["John", "john@example.com", 30]

# Type mismatch (string instead of number)
{"name": "John", "age": "thirty"}

# Missing required field
{"name": "John"}  # Missing email

# Extra field
{"name": "John", "email": "j@e.com", "ssn": "123-45-6789"}  # Security issue!

LLM Response → Parse JSON → Validate Schema → Type Coercion → 
    ↓ (fail)                   ↓ (fail)         ↓ (fail)
    └─→ Repair & Retry ──────────────────────────┘
                                                  ↓
                                            Clean Output

import json
import re
from typing import Any, Optional, Tuple

class RobustJSONParser:
    def parse(self, text: str) -> Tuple[Optional[dict], Optional[str]]:
        """Parse JSON with multiple fallback strategies."""
        
        # Strategy 1: Direct parse
        try:
            return json.loads(text), None
        except json.JSONDecodeError:
            pass
        
        # Strategy 2: Extract JSON from markdown code blocks
        json_match = re.search(r'

', text)
        if json_match:
            try:
                return json.loads(json_match.group(1)), None
            except json.JSONDecodeError:
                pass
        
        # Strategy 3: Find JSON object in text
        # Handle cases where LLM adds explanation before/after
        brace_start = text.find('{')
        brace_end = text.rfind('}')
        if brace_start != -1 and brace_end != -1:
            potential_json = text[brace_start:brace_end + 1]
            try:
                return json.loads(potential_json), None
            except json.JSONDecodeError:
                pass
        
        # Strategy 4: Repair common issues
        repaired = self._repair_json(text)
        if repaired:
            try:
                return json.loads(repaired), None
            except json.JSONDecodeError as e:
                return None, f"Parse failed after repair: {e}"
        
        return None, "Could not parse JSON from response"
    
    def _repair_json(self, text: str) -> Optional[str]:
        """Attempt to repair malformed JSON."""
        
        # Extract just the JSON part
        start = text.find('{')
        if start == -1:
            return None
        
        text = text[start:]
        
        # Fix common issues
        repairs = [
            # Trailing commas
            (r',\s*}', '}'),
            (r',\s*]', ']'),
            
            # Single quotes to double quotes
            (r"'([^']*)':", r'"\1":'),
            (r":\s*'([^']*)'", r': "\1"'),
            
            # Unquoted keys
            (r'(\{|,)\s*(\w+)\s*:', r'\1"\2":'),
            
            # Missing closing braces (truncation)
            # Count open/close braces and add missing ones
        ]
        
        for pattern, replacement in repairs:
            text = re.sub(pattern, replacement, text)
        
        # Handle truncation: count braces and brackets
        open_braces = text.count('{') - text.count('}')
        open_brackets = text.count('[') - text.count(']')
        
        # Add missing closing characters
        text = text.rstrip()
        if text.endswith(','):
            text = text[:-1]
        
        text += ']' * open_brackets + '}' * open_braces
        
        return text

from pydantic import BaseModel, Field, validator
from typing import List, Optional, Literal
from enum import Enum

class Priority(str, Enum):
    LOW = "low"
    MEDIUM = "medium"
    HIGH = "high"
    CRITICAL = "critical"

class ExtractedEntity(BaseModel):
    text: str = Field(..., min_length=1, max_length=500)
    entity_type: Literal["person", "organization", "location", "date", "money"]
    confidence: float = Field(..., ge=0.0, le=1.0)
    
    @validator('confidence', pre=True)
    def coerce_confidence(cls, v):
        # Handle string percentages like "85%"
        if isinstance(v, str):
            v = v.strip('%')
            return float(v) / 100 if float(v) > 1 else float(v)
        return v

class ExtractionResult(BaseModel):
    entities: List[ExtractedEntity]
    summary: Optional[str] = Field(None, max_length=1000)
    priority: Priority = Priority.MEDIUM
    
    @validator('priority', pre=True)
    def normalize_priority(cls, v):
        # Handle various capitalizations and synonyms
        if isinstance(v, str):
            v = v.lower().strip()
            synonyms = {
                "urgent": "critical",
                "important": "high",
                "normal": "medium",
                "minor": "low"
            }
            return synonyms.get(v, v)
        return v

class SchemaValidator:
    def __init__(self, model: type[BaseModel]):
        self.model = model
    
    def validate(self, data: dict) -> Tuple[Optional[BaseModel], List[str]]:
        """Validate data against schema, return validated model or errors."""
        try:
            validated = self.model.model_validate(data)
            return validated, []
        except ValidationError as e:
            errors = [f"{err['loc']}: {err['msg']}" for err in e.errors()]
            return None, errors

class LLMOutputRepairer:
    def __init__(self, llm_client, max_repair_attempts: int = 2):
        self.llm = llm_client
        self.max_attempts = max_repair_attempts
    
    async def repair(self, original_response: str, schema: dict,
                     validation_errors: List[str]) -> Optional[dict]:
        """Use LLM to repair invalid output."""
        
        repair_prompt = f'''
Your previous response was invalid. Please fix it.

EXPECTED SCHEMA:
{json.dumps(schema, indent=2)}

YOUR PREVIOUS RESPONSE:
{original_response}

VALIDATION ERRORS:
{chr(10).join(f"- {e}" for e in validation_errors)}

INSTRUCTIONS:
1. Fix ALL validation errors
2. Return ONLY the corrected JSON, no explanation
3. Ensure all required fields are present
4. Ensure all types match the schema

CORRECTED JSON:
'''
        
        for attempt in range(self.max_attempts):
            response = await self.llm.generate(
                messages=[{"role": "user", "content": repair_prompt}],
                response_format={"type": "json_object"},
                max_tokens=2000,
                temperature=0  # Deterministic for repairs
            )
            
            parsed, parse_error = self.parser.parse(response.content)
            if parsed:
                validated, val_errors = self.validator.validate(parsed)
                if validated:
                    return validated.model_dump()
                
                # Update prompt with new errors for next attempt
                validation_errors = val_errors
        
        return None

class StructuredOutputPipeline:
    def __init__(self, llm_client, schema_model: type[BaseModel]):
        self.llm = llm_client
        self.schema_model = schema_model
        self.parser = RobustJSONParser()
        self.validator = SchemaValidator(schema_model)
        self.repairer = LLMOutputRepairer(llm_client)
        
        # Convert Pydantic model to JSON schema for prompts
        self.json_schema = schema_model.model_json_schema()
    
    async def generate(self, prompt: str, context: str = "") -> dict:
        """Generate structured output with validation and repair."""
        
        # Build prompt with schema
        full_prompt = f'''
{prompt}

{context}

RESPOND WITH JSON MATCHING THIS EXACT SCHEMA:
{json.dumps(self.json_schema, indent=2)}

JSON RESPONSE:
'''
        
        # Initial generation
        response = await self.llm.generate(
            messages=[{"role": "user", "content": full_prompt}],
            response_format={"type": "json_object"},
            max_tokens=2000
        )
        
        # Parse
        parsed, parse_error = self.parser.parse(response.content)
        if not parsed:
            # Try repair
            parsed = await self.repairer.repair(
                response.content,
                self.json_schema,
                [parse_error]
            )
            if not parsed:
                raise StructuredOutputError(f"Failed to parse: {parse_error}")
        
        # Validate
        validated, val_errors = self.validator.validate(parsed)
        if not validated:
            # Try repair
            repaired = await self.repairer.repair(
                json.dumps(parsed),
                self.json_schema,
                val_errors
            )
            if repaired:
                return repaired
            raise StructuredOutputError(f"Validation failed: {val_errors}")
        
        return validated.model_dump()

class TruncationHandler:
    def __init__(self, llm_client, token_counter):
        self.llm = llm_client
        self.token_counter = token_counter
    
    def estimate_output_tokens(self, schema: dict, list_lengths: dict = None) -> int:
        """Estimate tokens needed for output based on schema."""
        
        base_tokens = 50  # JSON structure overhead
        
        for field_name, field_info in schema.get("properties", {}).items():
            field_type = field_info.get("type")
            
            if field_type == "string":
                max_length = field_info.get("maxLength", 100)
                base_tokens += max_length // 4  # ~4 chars per token
            
            elif field_type == "array":
                items_schema = field_info.get("items", {})
                expected_length = list_lengths.get(field_name, 10) if list_lengths else 10
                item_tokens = self.estimate_output_tokens({"properties": items_schema})
                base_tokens += item_tokens * expected_length
            
            elif field_type == "object":
                base_tokens += self.estimate_output_tokens(field_info)
            
            else:
                base_tokens += 10  # numbers, booleans, etc.
        
        return int(base_tokens * 1.2)  # 20% safety margin
    
    async def generate_with_continuation(self, prompt: str, schema: dict,
                                          max_total_tokens: int = 8000) -> str:
        """Generate with automatic continuation if truncated."""
        
        estimated_needed = self.estimate_output_tokens(schema)
        max_tokens = min(estimated_needed, max_total_tokens)
        
        response = await self.llm.generate(
            messages=[{"role": "user", "content": prompt}],
            response_format={"type": "json_object"},
            max_tokens=max_tokens
        )
        
        # Check if truncated
        if response.finish_reason == "length":
            # Continue generation
            continuation_prompt = f'''
Continue this JSON from where it was cut off:

{response.content}

Continue EXACTLY from the last character. Do not restart.
'''
            
            continuation = await self.llm.generate(
                messages=[{"role": "user", "content": continuation_prompt}],
                response_format={"type": "json_object"},
                max_tokens=max_tokens
            )
            
            return response.content + continuation.content
        
        return response.content