Model Routing for Multi-Agent Systems: Matching Tasks to the Right Model

Not every task needs the most powerful (and expensive) model available. A customer support classifier does not need the same model as a complex reasoning engine. Model routing — the practice of directing different tasks to different models based on requirements — is becoming essential for multi-agent systems that need to balance cost, speed, and quality.

Why This Matters

Multi-agent systems often contain agents with vastly different computational needs:

• Classification agents need fast, cheap models with structured output
• Reasoning agents need large context windows and strong logical capabilities
• Code generation agents need models trained specifically on code
• Summarization agents need models that handle long inputs efficiently

Running all of these on the same expensive model is wasteful. Running all of them on a small, cheap model produces poor results. The solution is intelligent routing that matches each task to the optimal model.

LiteLLM, an open-source AI Gateway, provides exactly this capability with a unified interface to 100+ LLM providers. It gives you a single API to call any model while handling routing, fallbacks, and cost tracking.

Setting Up LiteLLM for Model Routing

Installation and Basic Configuration

# Install LiteLLM
# pip install litellm

from litellm import completion
import os

# Configure API keys for multiple providers
os.environ["OPENAI_API_KEY"] = "your-openai-key"
os.environ["ANTHROPIC_API_KEY"] = "your-anthropic-key"
os.environ["GEMINI_API_KEY"] = "your-gemini-key"

# Call any model through a unified interface
# Fast and cheap for simple tasks
response = completion(
    model="gemini-2.0-flash",
    messages=[{"role": "user", "content": "Classify this text as positive or negative"}]
)

# More capable for complex reasoning
response = completion(
    model="anthropic/claude-sonnet-4-20250514",
    messages=[{"role": "user", "content": "Analyze this complex business problem"}]
)

# Code-specialized for development tasks
response = completion(
    model="openai/gpt-4o",
    messages=[{"role": "user", "content": "Write a Python function to sort a list"}]
)

Building a Model Router

from litellm import completion
from dataclasses import dataclass
from enum import Enum
import time

class TaskType(Enum):
    CLASSIFICATION = "classification"
    REASONING = "reasoning"
    CODE_GENERATION = "code_generation"
    SUMMARIZATION = "summarization"
    CREATIVE_WRITING = "creative_writing"
    TRANSLATION = "translation"

@dataclass
class ModelConfig:
    model_id: str
    max_tokens: int
    cost_per_1k_input: float
    cost_per_1k_output: float
    avg_latency_ms: float

# Define model configurations for each task type
MODEL_ROUTING_TABLE = {
    TaskType.CLASSIFICATION: ModelConfig(
        model_id="gemini-2.0-flash",
        max_tokens=100,
        cost_per_1k_input=0.000075,
        cost_per_1k_output=0.0003,
        avg_latency_ms=200,
    ),
    TaskType.REASONING: ModelConfig(
        model_id="anthropic/claude-sonnet-4-20250514",
        max_tokens=4096,
        cost_per_1k_input=0.003,
        cost_per_1k_output=0.015,
        avg_latency_ms=2000,
    ),
    TaskType.CODE_GENERATION: ModelConfig(
        model_id="openai/gpt-4o",
        max_tokens=4096,
        cost_per_1k_input=0.0025,
        cost_per_1k_output=0.01,
        avg_latency_ms=1500,
    ),
    TaskType.SUMMARIZATION: ModelConfig(
        model_id="gemini-2.0-flash",
        max_tokens=1024,
        cost_per_1k_input=0.000075,
        cost_per_1k_output=0.0003,
        avg_latency_ms=500,
    ),
    TaskType.CREATIVE_WRITING: ModelConfig(
        model_id="anthropic/claude-sonnet-4-20250514",
        max_tokens=2048,
        cost_per_1k_input=0.003,
        cost_per_1k_output=0.015,
        avg_latency_ms=1800,
    ),
    TaskType.TRANSLATION: ModelConfig(
        model_id="gemini-2.0-flash",
        max_tokens=2048,
        cost_per_1k_input=0.000075,
        cost_per_1k_output=0.0003,
        avg_latency_ms=300,
    ),
}

class ModelRouter:
    def __init__(self, routing_table: dict[TaskType, ModelConfig] = None):
        self.routing_table = routing_table or MODEL_ROUTING_TABLE
        self.call_count = 0
        self.total_cost = 0.0
    
    def route(
        self, 
        task_type: TaskType, 
        messages: list[dict],
        **kwargs
    ) -> dict:
        """Route a request to the appropriate model."""
        config = self.routing_table[task_type]
        
        start_time = time.time()
        
        response = completion(
            model=config.model_id,
            messages=messages,
            max_tokens=min(
                kwargs.get("max_tokens", config.max_tokens),
                config.max_tokens
            ),
            temperature=kwargs.get("temperature", 0.7),
        )
        
        latency_ms = (time.time() - start_time) * 1000
        
        # Track usage
        self.call_count += 1
        input_tokens = response.usage.prompt_tokens
        output_tokens = response.usage.completion_tokens
        
        cost = (
            (input_tokens / 1000) * config.cost_per_1k_input +
            (output_tokens / 1000) * config.cost_per_1k_output
        )
        self.total_cost += cost
        
        return {
            "content": response.choices[0].message.content,
            "model_used": config.model_id,
            "task_type": task_type.value,
            "tokens": {"input": input_tokens, "output": output_tokens},
            "cost_usd": cost,
            "latency_ms": latency_ms,
        }
    
    def get_stats(self) -> dict:
        """Get routing statistics."""
        return {
            "total_calls": self.call_count,
            "total_cost_usd": self.total_cost,
            "avg_cost_per_call": self.total_cost / max(self.call_count, 1),
        }

# Usage
router = ModelRouter()

# Simple classification - fast and cheap
result = router.route(
    TaskType.CLASSIFICATION,
    [{"role": "user", "content": "This product is amazing! Best purchase ever."}]
)
print(f"Classification: {result['content']}")
print(f"Cost: ${result['cost_usd']:.6f}")

# Complex reasoning - more expensive but higher quality
result = router.route(
    TaskType.REASONING,
    [{"role": "user", "content": "Analyze the trade-offs between microservices and monoliths for a 50-person engineering team."}]
)
print(f"Analysis: {result['content'][:200]}...")
print(f"Cost: ${result['cost_usd']:.4f}")

print(f"\nTotal stats: {router.get_stats()}")

Advanced Routing: Dynamic Selection

For more sophisticated routing, you can implement dynamic model selection based on input characteristics:

import re

class DynamicModelRouter(ModelRouter):
    def __init__(self):
        super().__init__()
        self.complexity_thresholds = {
            "simple": 100,      # < 100 words
            "medium": 500,      # 100-500 words
            "complex": 2000,    # 500-2000 words
        }
    
    def assess_complexity(self, text: str) -> str:
        """Assess input complexity based on multiple signals."""
        word_count = len(text.split())
        sentence_count = len(re.split(r'[.!?]+', text))
        avg_sentence_length = word_count / max(sentence_count, 1)
        
        # Check for technical/complex indicators
        has_code = bool(re.search(r'```|def |class |import ', text))
        has_numbers = bool(re.search(r'\d+\.?\d*', text))
        has_multiple_questions = text.count('?') > 1
        
        complexity_score = 0
        if word_count > 200: complexity_score += 2
        if avg_sentence_length > 20: complexity_score += 1
        if has_code: complexity_score += 2
        if has_numbers: complexity_score += 1
        if has_multiple_questions: complexity_score += 1
        
        if complexity_score >= 4:
            return "complex"
        elif complexity_score >= 2:
            return "medium"
        return "simple"
    
    def auto_route(self, messages: list[dict], **kwargs) -> dict:
        """Automatically route based on input complexity."""
        user_message = messages[-1]["content"]
        complexity = self.assess_complexity(user_message)
        
        # Map complexity to task types
        complexity_to_task = {
            "simple": TaskType.CLASSIFICATION,
            "medium": TaskType.SUMMARIZATION,
            "complex": TaskType.REASONING,
        }
        
        task_type = complexity_to_task[complexity]
        
        # Override with explicit task type if provided
        if "task_type" in kwargs:
            task_type = kwargs.pop("task_type")
        
        return self.route(task_type, messages, **kwargs)

# Auto-routing based on complexity
dynamic_router = DynamicModelRouter()

# Simple question - will route to fast model
result = dynamic_router.auto_route(
    [{"role": "user", "content": "What is the capital of France?"}]
)
print(f"Model used: {result['model_used']}")  # gemini-2.0-flash

# Complex analysis - will route to capable model
result = dynamic_router.auto_route(
    [{"role": "user", "content": "Compare the architectural patterns for building a real-time event processing system that needs to handle 100k events per second with exactly-once delivery guarantees. Consider trade-offs between Kafka, RabbitMQ, and custom solutions."}]
)
print(f"Model used: {result['model_used']}")  # claude-sonnet-4-20250514

Cost Optimization with Fallbacks

class CostOptimizedRouter(ModelRouter):
    def __init__(self, budget_per_call: float = 0.05):
        super().__init__()
        self.budget_per_call = budget_per_call
    
    def route_with_fallback(
        self, 
        task_type: TaskType, 
        messages: list[dict],
        quality_threshold: float = 0.7,
        **kwargs
    ) -> dict:
        """Route to cheapest model that meets quality requirements."""
        # Try models in order of increasing cost
        model_tiers = [
            ("gemini-2.0-flash", 0.000075),
            ("openai/gpt-4o-mini", 0.00015),
            ("anthropic/claude-sonnet-4-20250514", 0.003),
            ("openai/gpt-4o", 0.005),
        ]
        
        for model_id, cost_per_1k in model_tiers:
            response = completion(
                model=model_id,
                messages=messages,
                max_tokens=kwargs.get("max_tokens", 1024),
            )
            
            content = response.choices[0].message.content
            input_tokens = response.usage.prompt_tokens
            output_tokens = response.usage.completion_tokens
            
            cost = (
                (input_tokens / 1000) * cost_per_1k +
                (output_tokens / 1000) * cost_per_1k * 3
            )
            
            # Check if response meets quality bar
            if self._meets_quality(content, quality_threshold):
                return {
                    "content": content,
                    "model_used": model_id,
                    "cost_usd": cost,
                    "tier": model_id.split("/")[-1],
                }
        
        # Fallback to most capable model
        return self.route(task_type, messages, **kwargs)
    
    def _meets_quality(self, content: str, threshold: float) -> bool:
        """Simple quality heuristic — replace with your own quality check."""
        if len(content) < 10:
            return False
        if content.count("I don't know") > 0:
            return False
        return True

Best Practices

1. Start with a routing table: Map your task types to specific models before implementing dynamic routing. Simple static routing is easier to debug and optimize.

2. Track cost per task type: Monitor which task types consume the most budget. You may find opportunities to use cheaper models for tasks you assumed needed expensive ones.

3. Implement fallbacks: Models fail. Rate limits hit. Always have a fallback model for critical paths.

4. Cache aggressively: Many agent tasks are repetitive. Cache model responses for identical inputs to reduce costs and latency.

5. Profile regularly: Model performance and pricing change. Review your routing decisions monthly to ensure they are still optimal.

Common Pitfalls

• Over-routing: Adding routing logic for every possible edge case creates unmaintainable complexity. Start simple, optimize later.
• Ignoring latency requirements: A cheap model with high latency may be worse than a slightly more expensive model that responds faster.
• Forgetting about context windows: Not all models handle long contexts well. Match model capabilities to your input sizes.
• Not testing routing in production: A/B test routing decisions with real traffic to validate that your cost savings do not come at the expense of quality.

Conclusion

Model routing transforms multi-agent systems from expensive, one-size-fits-all architectures into cost-optimized, task-appropriate systems. LiteLLM makes this practical by providing a unified interface to hundreds of models with built-in fallbacks, cost tracking, and load balancing.

The key insight is that not all tasks deserve the same computational resources. By matching each task to the right model, you can reduce costs by 80% or more while maintaining or even improving output quality.

Next steps:

• Audit your current agent system to identify task types and their model requirements
• Set up LiteLLM as a unified gateway for your agent infrastructure
• Implement a basic routing table mapping task types to models
• Add cost tracking and monitor your savings over the first month