Coverage for agentos/agent/model_router.py: 52%
111 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-05 20:52 +0800
1"""Model Router — intelligent model selection based on task complexity and budget.
3Automatically selects the best LLM model for each task based on:
4 - Task complexity (TRIVIAL → EXPERT)
5 - Priority (LOW → CRITICAL)
6 - Daily budget constraints
7 - Model cost/performance tiers
8 - Fallback chains when budget is tight
9"""
11from __future__ import annotations
13import time
14from dataclasses import dataclass, field
15from enum import Enum, auto
16from typing import Callable, Optional
19__all__ = [
20 "TaskComplexity",
21 "TaskPriority",
22 "ModelSpec",
23 "RequestSpec",
24 "RouteResult",
25 "ModelRouter",
26]
29# ── Enums ──────────────────────────────────────────────────────────
31class TaskComplexity(Enum):
32 TRIVIAL = 0 # weather, time, simple calc
33 SIMPLE = 1 # basic Q&A, short translation
34 MODERATE = 2 # typical assistant tasks
35 COMPLEX = 3 # code review, analysis, research
36 EXPERT = 4 # deep research, architecture design
39class TaskPriority(Enum):
40 LOW = 0
41 NORMAL = 1
42 HIGH = 2
43 CRITICAL = 3
46# ── Dataclasses ────────────────────────────────────────────────────
48@dataclass
49class ModelSpec:
50 name: str
51 provider: str
52 cost_per_1k_input: float # USD per 1k input tokens
53 cost_per_1k_output: float # USD per 1k output tokens
54 max_tokens: int = 4096
55 context_window: int = 128_000
56 min_complexity: TaskComplexity = TaskComplexity.TRIVIAL
57 tags: list[str] = field(default_factory=list)
60@dataclass
61class RequestSpec:
62 estimated_input_tokens: int
63 estimated_output_tokens: int
64 complexity: TaskComplexity
65 priority: TaskPriority
66 task_id: str = ""
67 session_id: str = ""
70@dataclass
71class RouteResult:
72 success: bool
73 model: ModelSpec
74 reason: str = ""
75 estimated_cost: float = 0.0
76 fallback_chain: list[str] = field(default_factory=list)
79# ── Model Registry ────────────────────────────────────────────────
81DEFAULT_MODELS: list[ModelSpec] = [
82 # GPT family
83 ModelSpec("gpt-4o", "openai", 2.50, 10.00, 16384, 128000,
84 TaskComplexity.COMPLEX, ["gpt", "vision", "best"]),
85 ModelSpec("gpt-4o-mini", "openai", 0.15, 0.60, 16384, 128000,
86 TaskComplexity.SIMPLE, ["gpt", "cheap", "fast"]),
87 # Claude family
88 ModelSpec("claude-3.5-sonnet", "anthropic", 3.00, 15.00, 8192, 200000,
89 TaskComplexity.COMPLEX, ["claude", "code", "best"]),
90 ModelSpec("claude-3-haiku", "anthropic", 0.25, 1.25, 4096, 200000,
91 TaskComplexity.TRIVIAL, ["claude", "cheap", "fast"]),
92 # DeepSeek
93 ModelSpec("deepseek-v3", "deepseek", 0.27, 1.10, 8192, 64000,
94 TaskComplexity.MODERATE, ["deepseek", "value"]),
95 ModelSpec("deepseek-r1", "deepseek", 0.55, 2.19, 32768, 128000,
96 TaskComplexity.EXPERT, ["deepseek", "reasoning", "best"]),
97 # Gemini
98 ModelSpec("gemini-2.5-pro", "google", 1.25, 10.00, 8192, 1048576,
99 TaskComplexity.EXPERT, ["gemini", "best", "context"]),
100 ModelSpec("gemini-2.5-flash", "google", 0.15, 0.60, 8192, 1048576,
101 TaskComplexity.SIMPLE, ["gemini", "cheap", "fast"]),
102 # Ollama / local
103 ModelSpec("llama3.2-3b", "ollama", 0.0, 0.0, 4096, 128000,
104 TaskComplexity.TRIVIAL, ["local", "free"]),
105 ModelSpec("qwen2.5-7b", "ollama", 0.0, 0.0, 8192, 128000,
106 TaskComplexity.SIMPLE, ["local", "free"]),
107]
110# ── Model Router ──────────────────────────────────────────────────
112class ModelRouter:
113 """Intelligent model router with budget-aware selection.
115 Selects the best model for each task based on complexity, priority,
116 and daily budget. Supports fallback chains when the preferred model
117 would exceed budget limits.
119 Usage:
120 router = ModelRouter.with_defaults(daily_budget_usd=50.0)
121 spec = RequestSpec(
122 estimated_input_tokens=2000,
123 estimated_output_tokens=500,
124 complexity=TaskComplexity.COMPLEX,
125 priority=TaskPriority.NORMAL,
126 )
127 result = router.route(spec)
128 if result.success:
129 print(f"Routed to {result.model.name}, cost ~${result.estimated_cost:.4f}")
130 """
132 def __init__(
133 self,
134 models: list[ModelSpec],
135 daily_budget_usd: float = 50.0,
136 complexity_cost_map: Optional[dict[TaskComplexity, float]] = None,
137 ):
138 self.models = models
139 self._daily_budget = daily_budget_usd
140 self._daily_spent = 0.0
141 self._day_start = time.time()
142 self._request_count = 0
143 self._route_cache: dict[str, str] = {} # session_id → model_name
145 # Complexity → max estimated cost multiplier
146 self._complexity_budget: dict[TaskComplexity, float] = complexity_cost_map or {
147 TaskComplexity.TRIVIAL: 0.002,
148 TaskComplexity.SIMPLE: 0.01,
149 TaskComplexity.MODERATE: 0.05,
150 TaskComplexity.COMPLEX: 0.20,
151 TaskComplexity.EXPERT: 0.50,
152 }
154 # Sort models by cost (cheapest first for fallback)
155 self._by_cost = sorted(models, key=lambda m: m.cost_per_1k_input)
157 # ── Factory ────────────────────────────────────────────────────
159 @classmethod
160 def with_defaults(cls, daily_budget_usd: float = 50.0) -> "ModelRouter":
161 """Create a ModelRouter with default model registry."""
162 return cls(models=list(DEFAULT_MODELS), daily_budget_usd=daily_budget_usd)
164 # ── Routing ────────────────────────────────────────────────────
166 def route(self, spec: RequestSpec) -> RouteResult:
167 """Select the best model for the given request spec."""
168 self._request_count += 1
169 self._maybe_reset_daily()
171 # Check session affinity: if we've routed this session before, use same model
172 if spec.session_id and spec.session_id in self._route_cache:
173 model_name = self._route_cache[spec.session_id]
174 model = next((m for m in self.models if m.name == model_name), None)
175 if model:
176 est_cost = self._estimate_cost(model, spec)
177 if est_cost + self._daily_spent <= self._daily_budget:
178 return RouteResult(
179 success=True,
180 model=model,
181 reason=f"Session affinity → {model.name}",
182 estimated_cost=est_cost,
183 )
185 # 1. Find candidate models that can handle this complexity
186 candidates = [
187 m for m in self.models
188 if m.min_complexity.value <= spec.complexity.value
189 ]
190 if not candidates:
191 # Use cheapest as fallback
192 candidates = [self._by_cost[0]]
194 # 2. For high/critical priority, prefer top-tier
195 if spec.priority in (TaskPriority.HIGH, TaskPriority.CRITICAL):
196 # Filter to best-in-class models
197 tier = [m for m in candidates if "best" in m.tags]
198 if tier:
199 candidates = tier
201 # 3. Budget-aware selection
202 max_budget = self._complexity_budget.get(spec.complexity, 0.05)
204 # Try each candidate, falling back to cheaper ones if over budget
205 fallback_chain: list[str] = []
206 for model in sorted(candidates, key=lambda m: m.cost_per_1k_input):
207 est_cost = self._estimate_cost(model, spec)
208 if est_cost <= max_budget:
209 if self._daily_spent + est_cost <= self._daily_budget:
210 # Cache session → model
211 if spec.session_id:
212 self._route_cache[spec.session_id] = model.name
213 return RouteResult(
214 success=True,
215 model=model,
216 reason=f"Best match within budget ({spec.complexity.name})",
217 estimated_cost=est_cost,
218 fallback_chain=fallback_chain,
219 )
220 fallback_chain.append(model.name)
222 # 4. No model fits budget — use cheapest as emergency fallback
223 cheapest = self._by_cost[0]
224 est_cost = self._estimate_cost(cheapest, spec)
225 return RouteResult(
226 success=True,
227 model=cheapest,
228 reason=f"Budget constrained — fallen back to {cheapest.name}",
229 estimated_cost=est_cost,
230 fallback_chain=fallback_chain,
231 )
233 def record_request(
234 self,
235 model_name: str,
236 success: bool = True,
237 tokens_used: int = 0,
238 cost_usd: float = 0.0,
239 latency_ms: float = 0.0,
240 ):
241 """Record actual usage for budget tracking.
243 Args:
244 model_name: The model that was actually used.
245 success: Whether the request succeeded.
246 tokens_used: Total tokens consumed (input + output).
247 cost_usd: Actual cost in USD.
248 latency_ms: Request latency in milliseconds.
249 """
250 self._daily_spent += cost_usd
251 self._request_count += 1
253 # ── Budget ─────────────────────────────────────────────────────
255 @property
256 def daily_budget_remaining(self) -> float:
257 self._maybe_reset_daily()
258 return max(0.0, self._daily_budget - self._daily_spent)
260 @property
261 def daily_budget_total(self) -> float:
262 return self._daily_budget
264 @property
265 def daily_spent(self) -> float:
266 self._maybe_reset_daily()
267 return self._daily_spent
269 # ── Summary ────────────────────────────────────────────────────
271 def summary(self) -> dict:
272 """Return router state summary."""
273 return {
274 "models_available": len(self.models),
275 "daily_budget_usd": self._daily_budget,
276 "daily_spent": round(self._daily_spent, 6),
277 "daily_remaining": round(self.daily_budget_remaining, 4),
278 "request_count": self._request_count,
279 "cached_sessions": len(self._route_cache),
280 }
282 # ── Internal ───────────────────────────────────────────────────
284 def _estimate_cost(self, model: ModelSpec, spec: RequestSpec) -> float:
285 """Estimate USD cost for a request."""
286 input_cost = (spec.estimated_input_tokens / 1000) * model.cost_per_1k_input
287 output_cost = (spec.estimated_output_tokens / 1000) * model.cost_per_1k_output
288 return input_cost + output_cost
290 def _maybe_reset_daily(self):
291 """Reset daily budget if a day has passed."""
292 if time.time() - self._day_start > 86400:
293 self._daily_spent = 0.0
294 self._day_start = time.time()
295 self._route_cache.clear()