Spaces:
Running
Running
| """Cost-aware model router. | |
| Sends each task to the cheapest *capable* model, with graceful fallback down the | |
| tiers (frontier → local → mock) so a missing key or a dead Ollama never breaks a | |
| run. Consults the semantic cache before spending tokens, and records every call to | |
| the metrics store. This is the component that operationalizes the cost story. | |
| """ | |
| from __future__ import annotations | |
| from typing import Optional | |
| from .caching import SemanticCache | |
| from .config import Settings | |
| from .metrics import MetricsStore | |
| from .providers import LLMProvider, LLMRequest, LLMResponse, ProviderRegistry | |
| from .providers.base import Usage | |
| # Tasks that are safe + worthwhile to semantic-cache (deterministic-ish, repeatable). | |
| _CACHEABLE_TASKS = {"classify", "normalize", "extract"} | |
| class ModelRouter: | |
| def __init__( | |
| self, | |
| registry: ProviderRegistry, | |
| settings: Settings, | |
| metrics: MetricsStore, | |
| semantic_cache: Optional[SemanticCache] = None, | |
| ) -> None: | |
| self.registry = registry | |
| self.settings = settings | |
| self.metrics = metrics | |
| self.cache = semantic_cache or SemanticCache( | |
| threshold=settings.semantic_cache_threshold, | |
| enabled=settings.enable_semantic_cache, | |
| ) | |
| # --- public API ----------------------------------------------------------- | |
| def run(self, req: LLMRequest, run_id: str = "adhoc") -> LLMResponse: | |
| # 1) Semantic cache short-circuit. | |
| if req.task in _CACHEABLE_TASKS and self.settings.enable_semantic_cache: | |
| cache_key = f"{req.task}:{req.cacheable_prefix()}:{req.user_content}" | |
| cached = self.cache.get(cache_key) | |
| if cached is not None: | |
| resp = LLMResponse.build( | |
| text=cached, model="semantic-cache", provider="cache", | |
| usage=Usage(), latency_ms=0.2, cache_hit=True, | |
| routing_reason="semantic cache hit (LLM call skipped)", | |
| ) | |
| self.metrics.record_call(run_id, resp, req.task) | |
| return resp | |
| # 2) Route to candidates and execute with fallback. | |
| candidates = self._candidates(req) | |
| resp: Optional[LLMResponse] = None | |
| for provider, model, reason in candidates: | |
| resp = provider.complete(req, model) | |
| resp.routing_reason = reason + ( | |
| f" → {resp.routing_reason}" if resp.routing_reason else "" | |
| ) | |
| if not resp.error and resp.text: | |
| break # success | |
| if resp is None: # should never happen (mock is always last) | |
| resp = self.registry.mock.complete(req, "mock") | |
| # 3) Record + populate cache. | |
| self.metrics.record_call(run_id, resp, req.task) | |
| if ( | |
| req.task in _CACHEABLE_TASKS | |
| and self.settings.enable_semantic_cache | |
| and resp.text | |
| and not resp.error | |
| ): | |
| self.cache.put(f"{req.task}:{req.cacheable_prefix()}:{req.user_content}", resp.text) | |
| return resp | |
| # --- routing policy ------------------------------------------------------- | |
| def _candidates(self, req: LLMRequest) -> list[tuple[LLMProvider, str, str]]: | |
| """Ordered (provider, model, reason) list. First success wins; the list | |
| always ends with mock so a run can never hard-fail.""" | |
| policy = self.settings.routing_policy | |
| if policy == "offline": | |
| return [(self.registry.mock, "mock", "policy=offline")] | |
| cheap = self._cheap_chain() | |
| smart = self._smart_chain() | |
| if policy == "cheap": | |
| chain = cheap | |
| elif policy == "smart": | |
| chain = smart | |
| else: # auto | |
| chain = self._auto_chain(req, cheap, smart) | |
| # Always guarantee a terminal mock fallback. | |
| if not chain or chain[-1][0] is not self.registry.mock: | |
| chain = chain + [(self.registry.mock, "mock", "fallback=mock")] | |
| return chain | |
| def _auto_chain(self, req, cheap, smart): | |
| task = req.task | |
| if task in {"classify", "normalize", "validate", "format"}: | |
| return cheap # easy → cheapest | |
| if task == "agent": | |
| return smart + cheap # reasoning-heavy → frontier first | |
| if task == "extract": | |
| complexity = req.context.get("complexity", "simple") | |
| if complexity == "hard": | |
| return smart + cheap | |
| return cheap + smart # simple extract → cheap first, frontier as backup | |
| return cheap | |
| def _cheap_chain(self) -> list[tuple[LLMProvider, str, str]]: | |
| chain: list[tuple[LLMProvider, str, str]] = [] | |
| if self.registry.local and self.registry.local.available(): | |
| chain.append((self.registry.local, self.settings.local_model, | |
| "cheap: local Gemma (≈$0)")) | |
| if self.registry.anthropic and self.registry.anthropic.available(): | |
| chain.append((self.registry.anthropic, self.settings.anthropic_model_cheap, | |
| "cheap: Claude Haiku")) | |
| if self.registry.gemini and self.registry.gemini.available(): | |
| chain.append((self.registry.gemini, self.settings.gemini_model, | |
| "cheap: Gemini Flash")) | |
| if getattr(self.registry, "minicpm", None) and self.registry.minicpm.available(): | |
| chain.append((self.registry.minicpm, self.settings.minicpm_model, | |
| "cheap: MiniCPM-V (small model)")) | |
| return chain | |
| def _smart_chain(self) -> list[tuple[LLMProvider, str, str]]: | |
| chain: list[tuple[LLMProvider, str, str]] = [] | |
| if self.registry.anthropic and self.registry.anthropic.available(): | |
| chain.append((self.registry.anthropic, self.settings.anthropic_model_smart, | |
| "smart: Claude Sonnet")) | |
| if self.registry.gemini and self.registry.gemini.available(): | |
| chain.append((self.registry.gemini, self.settings.gemini_model, | |
| "smart: Gemini Flash")) | |
| if getattr(self.registry, "minicpm", None) and self.registry.minicpm.available(): | |
| chain.append((self.registry.minicpm, self.settings.minicpm_model, | |
| "smart: MiniCPM-V (small model)")) | |
| if self.registry.local and self.registry.local.available(): | |
| chain.append((self.registry.local, self.settings.local_model, | |
| "smart: local Gemma (no frontier available)")) | |
| return chain | |