Coverage for agentos/agent/model_router.py: 52%
110 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-09 07:12 +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
17__all__ = [
18 "TaskComplexity",
19 "TaskPriority",
20 "ModelSpec",
21 "RequestSpec",
22 "RouteResult",
23 "ModelRouter",
24]
27# ── Enums ──────────────────────────────────────────────────────────
30class TaskComplexity(Enum):
31 TRIVIAL = 0 # weather, time, simple calc
32 SIMPLE = 1 # basic Q&A, short translation
33 MODERATE = 2 # typical assistant tasks
34 COMPLEX = 3 # code review, analysis, research
35 EXPERT = 4 # deep research, architecture design
38class TaskPriority(Enum):
39 LOW = 0
40 NORMAL = 1
41 HIGH = 2
42 CRITICAL = 3
45# ── 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(
84 "gpt-4o",
85 "openai",
86 2.50,
87 10.00,
88 16384,
89 128000,
90 TaskComplexity.COMPLEX,
91 ["gpt", "vision", "best"],
92 ),
93 ModelSpec(
94 "gpt-4o-mini",
95 "openai",
96 0.15,
97 0.60,
98 16384,
99 128000,
100 TaskComplexity.SIMPLE,
101 ["gpt", "cheap", "fast"],
102 ),
103 # Claude family
104 ModelSpec(
105 "claude-3.5-sonnet",
106 "anthropic",
107 3.00,
108 15.00,
109 8192,
110 200000,
111 TaskComplexity.COMPLEX,
112 ["claude", "code", "best"],
113 ),
114 ModelSpec(
115 "claude-3-haiku",
116 "anthropic",
117 0.25,
118 1.25,
119 4096,
120 200000,
121 TaskComplexity.TRIVIAL,
122 ["claude", "cheap", "fast"],
123 ),
124 # DeepSeek
125 ModelSpec(
126 "deepseek-v3",
127 "deepseek",
128 0.27,
129 1.10,
130 8192,
131 64000,
132 TaskComplexity.MODERATE,
133 ["deepseek", "value"],
134 ),
135 ModelSpec(
136 "deepseek-r1",
137 "deepseek",
138 0.55,
139 2.19,
140 32768,
141 128000,
142 TaskComplexity.EXPERT,
143 ["deepseek", "reasoning", "best"],
144 ),
145 # Gemini
146 ModelSpec(
147 "gemini-2.5-pro",
148 "google",
149 1.25,
150 10.00,
151 8192,
152 1048576,
153 TaskComplexity.EXPERT,
154 ["gemini", "best", "context"],
155 ),
156 ModelSpec(
157 "gemini-2.5-flash",
158 "google",
159 0.15,
160 0.60,
161 8192,
162 1048576,
163 TaskComplexity.SIMPLE,
164 ["gemini", "cheap", "fast"],
165 ),
166 # Ollama / local
167 ModelSpec(
168 "llama3.2-3b", "ollama", 0.0, 0.0, 4096, 128000, TaskComplexity.TRIVIAL, ["local", "free"]
169 ),
170 ModelSpec(
171 "qwen2.5-7b", "ollama", 0.0, 0.0, 8192, 128000, TaskComplexity.SIMPLE, ["local", "free"]
172 ),
173]
176# ── Model Router ──────────────────────────────────────────────────
179class ModelRouter:
180 """Intelligent model router with budget-aware selection.
182 Selects the best model for each task based on complexity, priority,
183 and daily budget. Supports fallback chains when the preferred model
184 would exceed budget limits.
186 Usage:
187 router = ModelRouter.with_defaults(daily_budget_usd=50.0)
188 spec = RequestSpec(
189 estimated_input_tokens=2000,
190 estimated_output_tokens=500,
191 complexity=TaskComplexity.COMPLEX,
192 priority=TaskPriority.NORMAL,
193 )
194 result = router.route(spec)
195 if result.success:
196 print(f"Routed to {result.model.name}, cost ~${result.estimated_cost:.4f}")
197 """
199 def __init__(
200 self,
201 models: list[ModelSpec],
202 daily_budget_usd: float = 50.0,
203 complexity_cost_map: dict[TaskComplexity, float] | None = None,
204 ):
205 self.models = models
206 self._daily_budget = daily_budget_usd
207 self._daily_spent = 0.0
208 self._day_start = time.time()
209 self._request_count = 0
210 self._route_cache: dict[str, str] = {} # session_id → model_name
212 # Complexity → max estimated cost multiplier
213 self._complexity_budget: dict[TaskComplexity, float] = complexity_cost_map or {
214 TaskComplexity.TRIVIAL: 0.002,
215 TaskComplexity.SIMPLE: 0.01,
216 TaskComplexity.MODERATE: 0.05,
217 TaskComplexity.COMPLEX: 0.20,
218 TaskComplexity.EXPERT: 0.50,
219 }
221 # Sort models by cost (cheapest first for fallback)
222 self._by_cost = sorted(models, key=lambda m: m.cost_per_1k_input)
224 # ── Factory ────────────────────────────────────────────────────
226 @classmethod
227 def with_defaults(cls, daily_budget_usd: float = 50.0) -> ModelRouter:
228 """Create a ModelRouter with default model registry."""
229 return cls(models=list(DEFAULT_MODELS), daily_budget_usd=daily_budget_usd)
231 # ── Routing ────────────────────────────────────────────────────
233 def route(self, spec: RequestSpec) -> RouteResult:
234 """Select the best model for the given request spec."""
235 self._request_count += 1
236 self._maybe_reset_daily()
238 # Check session affinity: if we've routed this session before, use same model
239 if spec.session_id and spec.session_id in self._route_cache:
240 model_name = self._route_cache[spec.session_id]
241 model = next((m for m in self.models if m.name == model_name), None)
242 if model:
243 est_cost = self._estimate_cost(model, spec)
244 if est_cost + self._daily_spent <= self._daily_budget:
245 return RouteResult(
246 success=True,
247 model=model,
248 reason=f"Session affinity → {model.name}",
249 estimated_cost=est_cost,
250 )
252 # 1. Find candidate models that can handle this complexity
253 candidates = [m for m in self.models if m.min_complexity.value <= spec.complexity.value]
254 if not candidates:
255 # Use cheapest as fallback
256 candidates = [self._by_cost[0]]
258 # 2. For high/critical priority, prefer top-tier
259 if spec.priority in (TaskPriority.HIGH, TaskPriority.CRITICAL):
260 # Filter to best-in-class models
261 tier = [m for m in candidates if "best" in m.tags]
262 if tier:
263 candidates = tier
265 # 3. Budget-aware selection
266 max_budget = self._complexity_budget.get(spec.complexity, 0.05)
268 # Try each candidate, falling back to cheaper ones if over budget
269 fallback_chain: list[str] = []
270 for model in sorted(candidates, key=lambda m: m.cost_per_1k_input):
271 est_cost = self._estimate_cost(model, spec)
272 if est_cost <= max_budget:
273 if self._daily_spent + est_cost <= self._daily_budget:
274 # Cache session → model
275 if spec.session_id:
276 self._route_cache[spec.session_id] = model.name
277 return RouteResult(
278 success=True,
279 model=model,
280 reason=f"Best match within budget ({spec.complexity.name})",
281 estimated_cost=est_cost,
282 fallback_chain=fallback_chain,
283 )
284 fallback_chain.append(model.name)
286 # 4. No model fits budget — use cheapest as emergency fallback
287 cheapest = self._by_cost[0]
288 est_cost = self._estimate_cost(cheapest, spec)
289 return RouteResult(
290 success=True,
291 model=cheapest,
292 reason=f"Budget constrained — fallen back to {cheapest.name}",
293 estimated_cost=est_cost,
294 fallback_chain=fallback_chain,
295 )
297 def record_request(
298 self,
299 model_name: str,
300 success: bool = True,
301 tokens_used: int = 0,
302 cost_usd: float = 0.0,
303 latency_ms: float = 0.0,
304 ):
305 """Record actual usage for budget tracking.
307 Args:
308 model_name: The model that was actually used.
309 success: Whether the request succeeded.
310 tokens_used: Total tokens consumed (input + output).
311 cost_usd: Actual cost in USD.
312 latency_ms: Request latency in milliseconds.
313 """
314 self._daily_spent += cost_usd
315 self._request_count += 1
317 # ── Budget ─────────────────────────────────────────────────────
319 @property
320 def daily_budget_remaining(self) -> float:
321 self._maybe_reset_daily()
322 return max(0.0, self._daily_budget - self._daily_spent)
324 @property
325 def daily_budget_total(self) -> float:
326 return self._daily_budget
328 @property
329 def daily_spent(self) -> float:
330 self._maybe_reset_daily()
331 return self._daily_spent
333 # ── Summary ────────────────────────────────────────────────────
335 def summary(self) -> dict:
336 """Return router state summary."""
337 return {
338 "models_available": len(self.models),
339 "daily_budget_usd": self._daily_budget,
340 "daily_spent": round(self._daily_spent, 6),
341 "daily_remaining": round(self.daily_budget_remaining, 4),
342 "request_count": self._request_count,
343 "cached_sessions": len(self._route_cache),
344 }
346 # ── Internal ───────────────────────────────────────────────────
348 def _estimate_cost(self, model: ModelSpec, spec: RequestSpec) -> float:
349 """Estimate USD cost for a request."""
350 input_cost = (spec.estimated_input_tokens / 1000) * model.cost_per_1k_input
351 output_cost = (spec.estimated_output_tokens / 1000) * model.cost_per_1k_output
352 return input_cost + output_cost
354 def _maybe_reset_daily(self):
355 """Reset daily budget if a day has passed."""
356 if time.time() - self._day_start > 86400:
357 self._daily_spent = 0.0
358 self._day_start = time.time()
359 self._route_cache.clear()