Coverage for agentos/core/task_router.py: 0%
158 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-05 23:00 +0800
1"""
2AgentOS Task Router — Intelligent Multi-Model Task Routing
3━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
5Production-grade task routing engine that decides which language model
6to use per-request based on:
7 - Task complexity classification (simple/medium/complex)
8 - Latency budget constraints
9 - Cost optimization (cheapest model that meets quality bar)
10 - Fallback chains (primary → fallback → degraded)
11 - Model capability registry (tool-calling, vision, reasoning, etc.)
13Architecture:
14 TaskClassifier → classify task complexity from prompt analysis
15 ModelRegistry → register model capabilities and routing weights
16 RouteDecision → immutable routing decision with reasoning
17 TaskRouter → orchestrates classification + model selection
18"""
20from __future__ import annotations
22import re
23import time
24from dataclasses import dataclass, field
25from enum import Enum
26from typing import Any, Callable, Dict, List, Optional, Set, Tuple
28from agentos.core.cost_tracker import CostTracker
31# ---------------------------------------------------------------------------
32# Task Complexity
33# ---------------------------------------------------------------------------
36class TaskComplexity(str, Enum):
37 """Estimated complexity of a user task."""
38 SIMPLE = "simple" # Single-step, well-defined (translate, summarize)
39 MEDIUM = "medium" # Multi-step reasoning (analysis, code review)
40 COMPLEX = "complex" # Multi-agent orchestration, long-form generation
43class TaskCategory(str, Enum):
44 """Semantic category of the task."""
45 CHAT = "chat"
46 CODE = "code"
47 REASONING = "reasoning"
48 CREATIVE = "creative"
49 ANALYSIS = "analysis"
50 TOOL_USE = "tool_use"
51 VISION = "vision"
52 TRANSLATION = "translation"
55# ---------------------------------------------------------------------------
56# Model Capability
57# ---------------------------------------------------------------------------
60@dataclass
61class ModelSpec:
62 """Specification of a model's capabilities and routing profile."""
63 model_id: str
64 provider: str
65 tier: str # "premium" / "standard" / "budget" / "fallback"
66 max_tokens: int = 128000
67 latency_category: str = "medium" # "fast" / "medium" / "slow"
69 # Capability flags
70 supports_tool_calling: bool = False
71 supports_vision: bool = False
72 supports_reasoning: bool = False
73 supports_code: bool = False
74 supports_function_calling: bool = False
76 # Routing weights (higher = more likely to be selected for that category)
77 category_weights: Dict[TaskCategory, float] = field(default_factory=dict)
79 # Quality bar: tasks of this complexity must use this model or better
80 min_complexity: TaskComplexity = TaskComplexity.SIMPLE
82 # Cost consideration (relative cost multiplier, GPT-4o-mini = 1.0 baseline)
83 cost_multiplier: float = 1.0
85 # Excluded task categories
86 excluded_categories: Set[TaskCategory] = field(default_factory=set)
89# Default model registry
90DEFAULT_MODEL_SPECS: List[ModelSpec] = [
91 ModelSpec(
92 model_id="gpt-4o",
93 provider="openai",
94 tier="premium",
95 supports_tool_calling=True,
96 supports_vision=True,
97 supports_function_calling=True,
98 min_complexity=TaskComplexity.COMPLEX,
99 cost_multiplier=4.5,
100 latency_category="medium",
101 category_weights={
102 TaskCategory.REASONING: 0.9,
103 TaskCategory.ANALYSIS: 0.85,
104 TaskCategory.CODE: 0.8,
105 TaskCategory.TOOL_USE: 0.95,
106 },
107 ),
108 ModelSpec(
109 model_id="gpt-4o-mini",
110 provider="openai",
111 tier="standard",
112 supports_tool_calling=True,
113 supports_function_calling=True,
114 min_complexity=TaskComplexity.MEDIUM,
115 cost_multiplier=1.0,
116 latency_category="fast",
117 category_weights={
118 TaskCategory.CHAT: 1.0,
119 TaskCategory.TRANSLATION: 0.9,
120 TaskCategory.CODE: 0.7,
121 },
122 ),
123 ModelSpec(
124 model_id="claude-sonnet-4-20250514",
125 provider="anthropic",
126 tier="premium",
127 max_tokens=200000,
128 supports_tool_calling=True,
129 supports_reasoning=True,
130 supports_code=True,
131 min_complexity=TaskComplexity.COMPLEX,
132 cost_multiplier=3.0,
133 latency_category="medium",
134 category_weights={
135 TaskCategory.CODE: 1.0,
136 TaskCategory.REASONING: 0.95,
137 TaskCategory.ANALYSIS: 0.9,
138 TaskCategory.CREATIVE: 0.85,
139 },
140 ),
141 ModelSpec(
142 model_id="claude-haiku-3-5-sonnet-20241022",
143 provider="anthropic",
144 tier="standard",
145 supports_tool_calling=True,
146 latency_category="fast",
147 min_complexity=TaskComplexity.MEDIUM,
148 cost_multiplier=0.8,
149 category_weights={
150 TaskCategory.CHAT: 0.9,
151 TaskCategory.TRANSLATION: 0.8,
152 },
153 ),
154 ModelSpec(
155 model_id="deepseek-v3",
156 provider="deepseek",
157 tier="standard",
158 supports_tool_calling=False,
159 min_complexity=TaskComplexity.COMPLEX,
160 cost_multiplier=0.25,
161 latency_category="fast",
162 category_weights={
163 TaskCategory.CODE: 0.85,
164 TaskCategory.REASONING: 0.8,
165 TaskCategory.CHAT: 0.7,
166 },
167 ),
168 ModelSpec(
169 model_id="gemini-2.5-flash",
170 provider="google",
171 tier="standard",
172 supports_tool_calling=True,
173 supports_vision=True,
174 min_complexity=TaskComplexity.MEDIUM,
175 cost_multiplier=0.15,
176 latency_category="fast",
177 category_weights={
178 TaskCategory.CHAT: 0.85,
179 TaskCategory.VISION: 0.95,
180 TaskCategory.TRANSLATION: 0.7,
181 },
182 ),
183 ModelSpec(
184 model_id="llama-4-maverick",
185 provider="meta",
186 tier="budget",
187 min_complexity=TaskComplexity.SIMPLE,
188 cost_multiplier=0.1,
189 latency_category="fast",
190 category_weights={
191 TaskCategory.CHAT: 0.6,
192 },
193 ),
194]
197# ---------------------------------------------------------------------------
198# Task Classifier
199# ---------------------------------------------------------------------------
202class TaskClassifier:
203 """
204 Classify task complexity based on prompt heuristics.
206 Heuristics (no LLM call needed):
207 - Length: > 500 tokens → MEDIUM+; > 2000 tokens → COMPLEX
208 - Multi-step indicators: "first"/"then"/"next"/"finally" → MEDIUM+
209 - Agent keywords: "plan"/"orchestrate"/"coordinate"/"multi-agent" → COMPLEX
210 - Code: code blocks, "function"/"class"/"import" → MEDIUM+
211 - Reasoning: "why"/"explain"/"analyze"/"compare" → MEDIUM+
212 - Tool use: "search"/"fetch"/"call"/"run"/"execute" → MEDIUM+
213 """
215 # Token estimate: ~1.3 tokens per word (rough heuristic)
216 TOKENS_PER_WORD = 1.3
218 SIMPLE_THRESHOLD_TOKENS = 500
219 COMPLEX_THRESHOLD_TOKENS = 2000
221 MULTI_STEP_PATTERNS = [
222 re.compile(r"\b(first|then|next|finally|after that|step \d)\b", re.IGNORECASE),
223 re.compile(r"\d+\.\s+\w+"), # Numbered list
224 ]
226 AGENT_PATTERNS = [
227 re.compile(r"\b(plan|orchestrate|coordinate|multi.agent|dispatch|delegate)\b", re.IGNORECASE),
228 ]
230 CODE_PATTERNS = [
231 re.compile(r"```"),
232 re.compile(r"\b(def |class |import |function |const |let |var )", re.IGNORECASE),
233 ]
235 REASONING_PATTERNS = [
236 re.compile(r"\b(analy[sz]e|explain|compare|contrast|evaluate|assess|why)\b", re.IGNORECASE),
237 ]
239 TOOL_PATTERNS = [
240 re.compile(r"\b(search|fetch|call|run|execute|invoke|download|upload)\b", re.IGNORECASE),
241 ]
243 def classify(
244 self,
245 prompt: str,
246 available_tools: Optional[List[str]] = None,
247 ) -> Tuple[TaskComplexity, List[TaskCategory]]:
248 """
249 Classify a task from its prompt.
251 Returns (complexity, categories).
252 """
253 words = prompt.split()
254 estimated_tokens = max(1, int(len(words) * self.TOKENS_PER_WORD))
256 # Determine categories
257 categories: Set[TaskCategory] = set()
258 categories.add(TaskCategory.CHAT) # Default
260 if self._match_any(self.CODE_PATTERNS, prompt):
261 categories.add(TaskCategory.CODE)
263 if self._match_any(self.REASONING_PATTERNS, prompt):
264 categories.add(TaskCategory.REASONING)
265 categories.add(TaskCategory.ANALYSIS)
267 if self._match_any(self.TOOL_PATTERNS, prompt) or available_tools:
268 categories.add(TaskCategory.TOOL_USE)
270 # Determine complexity
271 is_multi_step = self._match_any(self.MULTI_STEP_PATTERNS, prompt)
272 is_agent = self._match_any(self.AGENT_PATTERNS, prompt)
274 if estimated_tokens > self.COMPLEX_THRESHOLD_TOKENS or is_agent:
275 complexity = TaskComplexity.COMPLEX
276 elif estimated_tokens > self.SIMPLE_THRESHOLD_TOKENS or is_multi_step or len(categories) > 2:
277 complexity = TaskComplexity.MEDIUM
278 else:
279 complexity = TaskComplexity.SIMPLE
281 return complexity, list(categories)
283 @staticmethod
284 def _match_any(patterns: List[re.Pattern], text: str) -> bool:
285 """Check if any pattern matches the text."""
286 return any(p.search(text) for p in patterns)
289# ---------------------------------------------------------------------------
290# Route Decision
291# ---------------------------------------------------------------------------
294@dataclass(frozen=True)
295class RouteDecision:
296 """Immutable routing decision."""
297 selected_model: str
298 provider: str
299 complexity: TaskComplexity
300 categories: List[TaskCategory]
301 fallback_models: List[str]
302 reasoning: str
303 latency_ms: float = 0.0
304 timestamp: float = field(default_factory=time.time)
306 def to_dict(self) -> Dict[str, Any]:
307 return {
308 "selected_model": self.selected_model,
309 "provider": self.provider,
310 "complexity": self.complexity.value,
311 "categories": [c.value for c in self.categories],
312 "fallback_models": self.fallback_models,
313 "reasoning": self.reasoning,
314 "latency_ms": round(self.latency_ms * 1000, 2),
315 }
318# ---------------------------------------------------------------------------
319# Task Router
320# ---------------------------------------------------------------------------
323class TaskRouter:
324 """
325 Routes tasks to the optimal model based on task analysis.
327 Strategy (in order):
328 1. Classify task complexity and categories
329 2. Filter models by capability requirements (tool-calling, vision, etc.)
330 3. Filter by min_complexity (only use models that meet quality bar)
331 4. Score remaining models by category weight × (1 / cost_multiplier)
332 5. Select top model + 2 fallbacks
333 """
335 def __init__(
336 self,
337 model_specs: Optional[List[ModelSpec]] = None,
338 cost_tracker: Optional[CostTracker] = None,
339 ):
340 self._models = model_specs or DEFAULT_MODEL_SPECS
341 self._classifier = TaskClassifier()
342 self._cost_tracker = cost_tracker
343 self._decision_log: List[RouteDecision] = []
345 @property
346 def models(self) -> List[ModelSpec]:
347 return list(self._models)
349 def register_model(self, spec: ModelSpec) -> None:
350 """Register a new model for routing."""
351 self._models.append(spec)
353 def route(
354 self,
355 prompt: str,
356 *,
357 required_capabilities: Optional[Set[str]] = None,
358 latency_budget_ms: Optional[int] = None,
359 available_tools: Optional[List[str]] = None,
360 ) -> RouteDecision:
361 """
362 Route a task to the best model.
364 Args:
365 prompt: The user's task prompt
366 required_capabilities: Set of required capabilities (e.g., {"tool_calling", "vision"})
367 latency_budget_ms: Maximum acceptable latency in ms
368 available_tools: Tools available for this task
370 Returns:
371 RouteDecision with selected model and fallbacks
372 """
373 start = time.time()
375 complexity, categories = self._classifier.classify(prompt, available_tools)
377 # Filter candidates
378 candidates = self._filter_candidates(complexity, categories, required_capabilities)
380 if not candidates:
381 # Fallback: drop min_complexity constraint
382 candidates = self._filter_candidates(complexity, categories, required_capabilities, relax_complexity=True)
384 if not candidates:
385 # Ultimate fallback: cheapest model
386 cheapest = min(self._models, key=lambda m: m.cost_multiplier)
387 candidates = [cheapest]
389 # Score and rank
390 ranked = self._rank_candidates(candidates, categories, latency_budget_ms)
391 selected = ranked[0]
392 fallbacks = [r.model_id for r in ranked[1:3]] if len(ranked) > 1 else []
394 decision = RouteDecision(
395 selected_model=selected.model_id,
396 provider=selected.provider,
397 complexity=complexity,
398 categories=categories,
399 fallback_models=fallbacks,
400 reasoning=self._build_reasoning(selected, complexity, categories),
401 latency_ms=time.time() - start,
402 )
403 self._decision_log.append(decision)
404 return decision
406 def _filter_candidates(
407 self,
408 complexity: TaskComplexity,
409 categories: List[TaskCategory],
410 required_capabilities: Optional[Set[str]] = None,
411 relax_complexity: bool = False,
412 ) -> List[ModelSpec]:
413 """Filter models by complexity bar, category exclusion, and capabilities."""
414 caps = required_capabilities or set()
415 candidates: List[ModelSpec] = []
417 for model in self._models:
418 # Complexity bar (can relax if no candidates)
419 if not relax_complexity:
420 complexity_levels = {
421 TaskComplexity.SIMPLE: 0,
422 TaskComplexity.MEDIUM: 1,
423 TaskComplexity.COMPLEX: 2,
424 }
425 model_level = complexity_levels[model.min_complexity]
426 task_level = complexity_levels[complexity]
427 if model_level < task_level:
428 continue
430 # Category exclusion
431 if any(cat in model.excluded_categories for cat in categories):
432 continue
434 # Capability requirements
435 if "tool_calling" in caps and not model.supports_tool_calling:
436 continue
437 if "vision" in caps and not model.supports_vision:
438 continue
439 if "reasoning" in caps and not model.supports_reasoning:
440 continue
442 candidates.append(model)
444 return candidates
446 def _rank_candidates(
447 self,
448 candidates: List[ModelSpec],
449 categories: List[TaskCategory],
450 latency_budget_ms: Optional[int] = None,
451 ) -> List[ModelSpec]:
452 """Score and rank candidates by category weight / cost."""
453 LATENCY_SCORES = {"fast": 1.2, "medium": 1.0, "slow": 0.7}
455 def score(model: ModelSpec) -> float:
456 # Category affinity
457 cat_score = sum(
458 model.category_weights.get(cat, 0.0) for cat in categories
459 ) / max(1, len(categories))
461 # Cost efficiency (inverse)
462 cost_score = 1.0 / max(0.01, model.cost_multiplier)
464 # Latency bonus
465 latency_score = LATENCY_SCORES.get(model.latency_category, 1.0)
467 # Final: weighted sum
468 return cat_score * 0.5 + cost_score * 0.3 + latency_score * 0.2
470 ranked = sorted(candidates, key=score, reverse=True)
472 # Filter by latency budget if specified
473 if latency_budget_ms is not None:
474 SLOW_THRESHOLD_MS = 2000
475 if latency_budget_ms < SLOW_THRESHOLD_MS:
476 ranked = [m for m in ranked if m.latency_category != "slow"]
478 return ranked
480 def _build_reasoning(
481 self,
482 selected: ModelSpec,
483 complexity: TaskComplexity,
484 categories: List[TaskCategory],
485 ) -> str:
486 """Build human-readable reasoning for the routing decision."""
487 cat_names = ", ".join(c.value for c in categories)
488 return (
489 f"Task complexity: {complexity.value}. Categories: {cat_names}. "
490 f"Selected {selected.model_id} ({selected.provider}, {selected.tier} tier) "
491 f"for optimal quality/cost balance."
492 )
494 def get_decision_log(self) -> List[RouteDecision]:
495 return list(self._decision_log)
497 def clear_log(self) -> None:
498 self._decision_log.clear()
500 def get_statistics(self) -> Dict[str, Any]:
501 """Return routing statistics."""
502 if not self._decision_log:
503 return {}
505 model_counts: Dict[str, int] = {}
506 complexity_counts: Dict[str, int] = {}
507 total_latency = 0.0
509 for d in self._decision_log:
510 model_counts[d.selected_model] = model_counts.get(d.selected_model, 0) + 1
511 complexity_counts[d.complexity.value] = complexity_counts.get(d.complexity.value, 0) + 1
512 total_latency += d.latency_ms
514 return {
515 "total_requests": len(self._decision_log),
516 "model_distribution": model_counts,
517 "complexity_distribution": complexity_counts,
518 "avg_latency_ms": round(total_latency * 1000 / len(self._decision_log), 2),
519 }