Speeding Up the Agent › Choosing the Right Model for the Task

from openai import OpenAI
from anthropic import Anthropic

client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

def classify_query_complexity(query):
    """Determine if a query needs a powerful model or a fast one."""
    # Simple heuristic: check for complexity indicators
    complex_indicators = [
        "explain", "analyze", "compare", "why", "how does",
        "step by step", "reasoning", "pros and cons"
    ]
    
    query_lower = query.lower()
    is_complex = any(indicator in query_lower for indicator in complex_indicators)
    
    return "complex" if is_complex else "simple"

def route_to_model(query):
    """Route query to appropriate model based on complexity."""
    complexity = classify_query_complexity(query)
    
    if complexity == "simple":
        # Use GPT-5 for fast, straightforward responses
        model = "gpt-5"
        max_tokens = 150
    else:
        # Use Claude Sonnet 4.5 for complex reasoning
        anthropic_client = Anthropic(api_key=os.getenv("ANTHROPIC_API_KEY"))
        response = anthropic_client.messages.create(
            model="claude-sonnet-4-5",
            max_tokens=1024,
            messages=[{"role": "user", "content": query}]
        )
        return response.content[0].text
    
    # Handle simple queries with GPT-5
    response = client.chat.completions.create(
        model=model,
        max_completion_tokens=max_tokens,
        messages=[{"role": "user", "content": query}]
    )
    
    return response.choices[0].message.content

## Test both paths
print("Simple query:", route_to_model("What's the capital of France?"))
print("Complex query:", route_to_model("Explain why France chose Paris as its capital"))