Coverage for agentos/core/task

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 StrEnum

26from typing import Any

28from agentos.core.cost_tracker import CostTracker

30# ---------------------------------------------------------------------------

31# Task Complexity

32# ---------------------------------------------------------------------------

35class TaskComplexity(StrEnum):

36 """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(StrEnum):

44 """Semantic category of the task."""

46 CHAT = "chat"

47 CODE = "code"

48 REASONING = "reasoning"

49 CREATIVE = "creative"

50 ANALYSIS = "analysis"

51 TOOL_USE = "tool_use"

52 VISION = "vision"

53 TRANSLATION = "translation"

56# ---------------------------------------------------------------------------

57# Model Capability

58# ---------------------------------------------------------------------------

61@dataclass

62class ModelSpec:

63 """Specification of a model's capabilities and routing profile."""

65 model_id: str

66 provider: str

67 tier: str # "premium" / "standard" / "budget" / "fallback"

68 max_tokens: int = 128000

69 latency_category: str = "medium" # "fast" / "medium" / "slow"

71 # Capability flags

72 supports_tool_calling: bool = False

73 supports_vision: bool = False

74 supports_reasoning: bool = False

75 supports_code: bool = False

76 supports_function_calling: bool = False

78 # Routing weights (higher = more likely to be selected for that category)

79 category_weights: dict[TaskCategory, float] = field(default_factory=dict)

81 # Quality bar: tasks of this complexity must use this model or better

82 min_complexity: TaskComplexity = TaskComplexity.SIMPLE

84 # Cost consideration (relative cost multiplier, GPT-4o-mini = 1.0 baseline)

85 cost_multiplier: float = 1.0

87 # Excluded task categories

88 excluded_categories: set[TaskCategory] = field(default_factory=set)

91# Default model registry

92DEFAULT_MODEL_SPECS: list[ModelSpec] = [

93 ModelSpec(

94 model_id="gpt-4o",

95 provider="openai",

96 tier="premium",

97 supports_tool_calling=True,

98 supports_vision=True,

99 supports_function_calling=True,

100 min_complexity=TaskComplexity.COMPLEX,

101 cost_multiplier=4.5,

102 latency_category="medium",

103 category_weights={

104 TaskCategory.REASONING: 0.9,

105 TaskCategory.ANALYSIS: 0.85,

106 TaskCategory.CODE: 0.8,

107 TaskCategory.TOOL_USE: 0.95,

108 },

109 ),

110 ModelSpec(

111 model_id="gpt-4o-mini",

112 provider="openai",

113 tier="standard",

114 supports_tool_calling=True,

115 supports_function_calling=True,

116 min_complexity=TaskComplexity.MEDIUM,

117 cost_multiplier=1.0,

118 latency_category="fast",

119 category_weights={

120 TaskCategory.CHAT: 1.0,

121 TaskCategory.TRANSLATION: 0.9,

122 TaskCategory.CODE: 0.7,

123 },

124 ),

125 ModelSpec(

126 model_id="claude-sonnet-4-20250514",

127 provider="anthropic",

128 tier="premium",

129 max_tokens=200000,

130 supports_tool_calling=True,

131 supports_reasoning=True,

132 supports_code=True,

133 min_complexity=TaskComplexity.COMPLEX,

134 cost_multiplier=3.0,

135 latency_category="medium",

136 category_weights={

137 TaskCategory.CODE: 1.0,

138 TaskCategory.REASONING: 0.95,

139 TaskCategory.ANALYSIS: 0.9,

140 TaskCategory.CREATIVE: 0.85,

141 },

142 ),

143 ModelSpec(

144 model_id="claude-haiku-3-5-sonnet-20241022",

145 provider="anthropic",

146 tier="standard",

147 supports_tool_calling=True,

148 latency_category="fast",

149 min_complexity=TaskComplexity.MEDIUM,

150 cost_multiplier=0.8,

151 category_weights={

152 TaskCategory.CHAT: 0.9,

153 TaskCategory.TRANSLATION: 0.8,

154 },

155 ),

156 ModelSpec(

157 model_id="deepseek-v3",

158 provider="deepseek",

159 tier="standard",

160 supports_tool_calling=False,

161 min_complexity=TaskComplexity.COMPLEX,

162 cost_multiplier=0.25,

163 latency_category="fast",

164 category_weights={

165 TaskCategory.CODE: 0.85,

166 TaskCategory.REASONING: 0.8,

167 TaskCategory.CHAT: 0.7,

168 },

169 ),

170 ModelSpec(

171 model_id="gemini-2.5-flash",

172 provider="google",

173 tier="standard",

174 supports_tool_calling=True,

175 supports_vision=True,

176 min_complexity=TaskComplexity.MEDIUM,

177 cost_multiplier=0.15,

178 latency_category="fast",

179 category_weights={

180 TaskCategory.CHAT: 0.85,

181 TaskCategory.VISION: 0.95,

182 TaskCategory.TRANSLATION: 0.7,

183 },

184 ),

185 ModelSpec(

186 model_id="llama-4-maverick",

187 provider="meta",

188 tier="budget",

189 min_complexity=TaskComplexity.SIMPLE,

190 cost_multiplier=0.1,

191 latency_category="fast",

192 category_weights={

193 TaskCategory.CHAT: 0.6,

194 },

195 ),

196]

197

198

199# ---------------------------------------------------------------------------

200# Task Classifier

201# ---------------------------------------------------------------------------

202

203

204class TaskClassifier:

205 """

206 Classify task complexity based on prompt heuristics.

207

208 Heuristics (no LLM call needed):

209 - Length: > 500 tokens → MEDIUM+; > 2000 tokens → COMPLEX

210 - Multi-step indicators: "first"/"then"/"next"/"finally" → MEDIUM+

211 - Agent keywords: "plan"/"orchestrate"/"coordinate"/"multi-agent" → COMPLEX

212 - Code: code blocks, "function"/"class"/"import" → MEDIUM+

213 - Reasoning: "why"/"explain"/"analyze"/"compare" → MEDIUM+

214 - Tool use: "search"/"fetch"/"call"/"run"/"execute" → MEDIUM+

215 """

216

217 # Token estimate: ~1.3 tokens per word (rough heuristic)

218 TOKENS_PER_WORD = 1.3

219

220 SIMPLE_THRESHOLD_TOKENS = 500

221 COMPLEX_THRESHOLD_TOKENS = 2000

222

223 MULTI_STEP_PATTERNS = [

225 re.compile(r"\d+\.\s+\w+"), # Numbered list

226 ]

227

228 AGENT_PATTERNS = [

229 re.compile(

231 ),

232 ]

233

234 CODE_PATTERNS = [

235 re.compile(r"```"),

237 ]

238

239 REASONING_PATTERNS = [

241 ]

242

243 TOOL_PATTERNS = [

245 ]

246

247 def classify(

248 self,

249 prompt: str,

250 available_tools: list[str] | None = None,

251 ) -> tuple[TaskComplexity, list[TaskCategory]]:

252 """

253 Classify a task from its prompt.

254

255 Returns (complexity, categories).

256 """

257 words = prompt.split()

258 estimated_tokens = max(1, int(len(words) * self.TOKENS_PER_WORD))

259

260 # Determine categories

261 categories: set[TaskCategory] = set()

262 categories.add(TaskCategory.CHAT) # Default

263

264 if self._match_any(self.CODE_PATTERNS, prompt):

265 categories.add(TaskCategory.CODE)

266

267 if self._match_any(self.REASONING_PATTERNS, prompt):

268 categories.add(TaskCategory.REASONING)

269 categories.add(TaskCategory.ANALYSIS)

270

271 if self._match_any(self.TOOL_PATTERNS, prompt) or available_tools:

272 categories.add(TaskCategory.TOOL_USE)

273

274 # Determine complexity

275 is_multi_step = self._match_any(self.MULTI_STEP_PATTERNS, prompt)

276 is_agent = self._match_any(self.AGENT_PATTERNS, prompt)

277

278 if estimated_tokens > self.COMPLEX_THRESHOLD_TOKENS or is_agent:

279 complexity = TaskComplexity.COMPLEX

280 elif (

281 estimated_tokens > self.SIMPLE_THRESHOLD_TOKENS or is_multi_step or len(categories) > 2

282 ):

283 complexity = TaskComplexity.MEDIUM

284 else:

285 complexity = TaskComplexity.SIMPLE

286

287 return complexity, list(categories)

288

289 @staticmethod

290 def _match_any(patterns: list[re.Pattern], text: str) -> bool:

291 """Check if any pattern matches the text."""

292 return any(p.search(text) for p in patterns)

293

294

295# ---------------------------------------------------------------------------

296# Route Decision

297# ---------------------------------------------------------------------------

298

299

300@dataclass(frozen=True)

301class RouteDecision:

302 """Immutable routing decision."""

303

304 selected_model: str

305 provider: str

306 complexity: TaskComplexity

307 categories: list[TaskCategory]

308 fallback_models: list[str]

309 reasoning: str

310 latency_ms: float = 0.0

311 timestamp: float = field(default_factory=time.time)

312

313 def to_dict(self) -> dict[str, Any]:

314 return {

315 "selected_model": self.selected_model,

316 "provider": self.provider,

317 "complexity": self.complexity.value,

318 "categories": [c.value for c in self.categories],

319 "fallback_models": self.fallback_models,

320 "reasoning": self.reasoning,

321 "latency_ms": round(self.latency_ms * 1000, 2),

322 }

323

324

325# ---------------------------------------------------------------------------

326# Task Router

327# ---------------------------------------------------------------------------

328

329

330class TaskRouter:

331 """

332 Routes tasks to the optimal model based on task analysis.

333

334 Strategy (in order):

335 1. Classify task complexity and categories

336 2. Filter models by capability requirements (tool-calling, vision, etc.)

337 3. Filter by min_complexity (only use models that meet quality bar)

338 4. Score remaining models by category weight × (1 / cost_multiplier)

339 5. Select top model + 2 fallbacks

340 """

341

342 def __init__(

343 self,

344 model_specs: list[ModelSpec] | None = None,

345 cost_tracker: CostTracker | None = None,

346 ):

347 self._models = model_specs or DEFAULT_MODEL_SPECS

348 self._classifier = TaskClassifier()

349 self._cost_tracker = cost_tracker

350 self._decision_log: list[RouteDecision] = []

351

352 @property

353 def models(self) -> list[ModelSpec]:

354 return list(self._models)

355

356 def register_model(self, spec: ModelSpec) -> None:

357 """Register a new model for routing."""

358 self._models.append(spec)

359

360 def route(

361 self,

362 prompt: str,

363 *,

364 required_capabilities: set[str] | None = None,

365 latency_budget_ms: int | None = None,

366 available_tools: list[str] | None = None,

367 ) -> RouteDecision:

368 """

369 Route a task to the best model.

370

371 Args:

372 prompt: The user's task prompt

373 required_capabilities: Set of required capabilities (e.g., {"tool_calling", "vision"})

374 latency_budget_ms: Maximum acceptable latency in ms

375 available_tools: Tools available for this task

376

377 Returns:

378 RouteDecision with selected model and fallbacks

379 """

380 start = time.time()

381

382 complexity, categories = self._classifier.classify(prompt, available_tools)

383

384 # Filter candidates

385 candidates = self._filter_candidates(complexity, categories, required_capabilities)

386

387 if not candidates:

388 # Fallback: drop min_complexity constraint

389 candidates = self._filter_candidates(

390 complexity, categories, required_capabilities, relax_complexity=True

391 )

392

393 if not candidates:

394 # Ultimate fallback: cheapest model

395 cheapest = min(self._models, key=lambda m: m.cost_multiplier)

396 candidates = [cheapest]

397

398 # Score and rank

399 ranked = self._rank_candidates(candidates, categories, latency_budget_ms)

400 selected = ranked[0]

401 fallbacks = [r.model_id for r in ranked[1:3]] if len(ranked) > 1 else []

402

403 decision = RouteDecision(

404 selected_model=selected.model_id,

405 provider=selected.provider,

406 complexity=complexity,

407 categories=categories,

408 fallback_models=fallbacks,

409 reasoning=self._build_reasoning(selected, complexity, categories),

410 latency_ms=time.time() - start,

411 )

412 self._decision_log.append(decision)

413 return decision

414

415 def _filter_candidates(

416 self,

417 complexity: TaskComplexity,

418 categories: list[TaskCategory],

419 required_capabilities: set[str] | None = None,

420 relax_complexity: bool = False,

421 ) -> list[ModelSpec]:

422 """Filter models by complexity bar, category exclusion, and capabilities."""

423 caps = required_capabilities or set()

424 candidates: list[ModelSpec] = []

425

426 for model in self._models:

427 # Complexity bar (can relax if no candidates)

428 if not relax_complexity:

429 complexity_levels = {

430 TaskComplexity.SIMPLE: 0,

431 TaskComplexity.MEDIUM: 1,

432 TaskComplexity.COMPLEX: 2,

433 }

434 model_level = complexity_levels[model.min_complexity]

435 task_level = complexity_levels[complexity]

436 if model_level < task_level:

437 continue

438

439 # Category exclusion

440 if any(cat in model.excluded_categories for cat in categories):

441 continue

442

443 # Capability requirements

444 if "tool_calling" in caps and not model.supports_tool_calling:

445 continue

446 if "vision" in caps and not model.supports_vision:

447 continue

448 if "reasoning" in caps and not model.supports_reasoning:

449 continue

450

451 candidates.append(model)

452

453 return candidates

454

455 def _rank_candidates(

456 self,

457 candidates: list[ModelSpec],

458 categories: list[TaskCategory],

459 latency_budget_ms: int | None = None,

460 ) -> list[ModelSpec]:

461 """Score and rank candidates by category weight / cost."""

462 LATENCY_SCORES = {"fast": 1.2, "medium": 1.0, "slow": 0.7}

463

464 def score(model: ModelSpec) -> float:

465 # Category affinity

466 cat_score = sum(model.category_weights.get(cat, 0.0) for cat in categories) / max(

467 1, len(categories)

468 )

469

470 # Cost efficiency (inverse)

471 cost_score = 1.0 / max(0.01, model.cost_multiplier)

472

473 # Latency bonus

474 latency_score = LATENCY_SCORES.get(model.latency_category, 1.0)

475

476 # Final: weighted sum

477 return cat_score * 0.5 + cost_score * 0.3 + latency_score * 0.2

478

479 ranked = sorted(candidates, key=score, reverse=True)

480

481 # Filter by latency budget if specified

482 if latency_budget_ms is not None:

483 SLOW_THRESHOLD_MS = 2000

484 if latency_budget_ms < SLOW_THRESHOLD_MS:

485 ranked = [m for m in ranked if m.latency_category != "slow"]

486

487 return ranked

488

489 def _build_reasoning(

490 self,

491 selected: ModelSpec,

492 complexity: TaskComplexity,

493 categories: list[TaskCategory],

494 ) -> str:

495 """Build human-readable reasoning for the routing decision."""

496 cat_names = ", ".join(c.value for c in categories)

497 return (

498 f"Task complexity: {complexity.value}. Categories: {cat_names}. "

499 f"Selected {selected.model_id} ({selected.provider}, {selected.tier} tier) "

500 f"for optimal quality/cost balance."

501 )

502

503 def get_decision_log(self) -> list[RouteDecision]:

504 return list(self._decision_log)

505

506 def clear_log(self) -> None:

507 self._decision_log.clear()

508

509 def get_statistics(self) -> dict[str, Any]:

510 """Return routing statistics."""

511 if not self._decision_log:

512 return {}

513

514 model_counts: dict[str, int] = {}

515 complexity_counts: dict[str, int] = {}

516 total_latency = 0.0

517

518 for d in self._decision_log:

519 model_counts[d.selected_model] = model_counts.get(d.selected_model, 0) + 1

520 complexity_counts[d.complexity.value] = complexity_counts.get(d.complexity.value, 0) + 1

521 total_latency += d.latency_ms

522

523 return {

524 "total_requests": len(self._decision_log),

525 "model_distribution": model_counts,

526 "complexity_distribution": complexity_counts,

527 "avg_latency_ms": round(total_latency * 1000 / len(self._decision_log), 2),

528 }

Coverage for agentos/core/task_router.py: 99%

158 statements