rex/model_router.py
Rama master · réplica de lectura del repositorio original.
"""Model router — classify user messages into task categories for model selection.
Classification is heuristic-only (keyword/pattern matching) with no LLM call,
giving deterministic, zero-latency routing.
Ollama availability is checked at init time (if needed) and refreshed every
60 seconds via a daemon background thread. OpenAI and local-Transformers
models are assumed always available — no network call is made for them.
"""
from __future__ import annotations
import json
import logging
import re
import threading
import time
import urllib.error
import urllib.request
from collections.abc import Callable, Sequence
from dataclasses import dataclass
from enum import StrEnum, auto
from rex.provider_reliability import ProviderFailureKind, ProviderReliability
logger = logging.getLogger(__name__)
class TaskCategory(StrEnum):
coding = auto()
reasoning = auto()
search = auto()
vision = auto()
fast = auto()
default = auto()
@dataclass(frozen=True, slots=True)
class ModelRouteDecision:
"""Privacy-safe explanation of one fast/deep model-routing decision."""
category: TaskCategory
complexity: str
confidence: float
evidence: tuple[str, ...]
tier: str
model: str
escalation_count: int = 0
fallback_reason: str | None = None
def to_metadata(self) -> dict[str, object]:
return {
"category": str(self.category),
"complexity": self.complexity,
"confidence": round(self.confidence, 3),
"evidence": list(self.evidence),
"route": self.tier,
"model": self.model,
"escalation_count": self.escalation_count,
"fallback_reason": self.fallback_reason,
}
@dataclass(frozen=True, slots=True)
class ProviderRouteCandidate:
"""One configured provider/model candidate in deterministic preference order."""
provider: str
model: str
@dataclass(frozen=True, slots=True)
class ProviderRouteSelection:
"""Privacy-safe result of provider health filtering and fallback."""
provider: str
model: str
fallback_reason: str | None = None
evidence: tuple[str, ...] = ()
# ---------------------------------------------------------------------------
# Keyword / pattern tables
# ---------------------------------------------------------------------------
_CODING_PATTERNS: Sequence[str] = [
r"\bcode\b",
r"\bwrite\s+(a\s+)?(function|class|script|program|module|test)\b",
r"\bdebugg?\b",
r"\brefactor\b",
r"\bimplement\b",
r"\bsyntax\s+error\b",
r"\bpython\b",
r"\bjavascript\b",
r"\btypescript\b",
r"\bgolang?\b",
r"\brust\b",
r"\bsql\b",
r"\bshell\s+script\b",
r"\bbash\b",
r"\bregex\b",
r"\bapi\s+endpoint\b",
r"\bunit\s+test\b",
r"\bpytest\b",
r"\bgit\s+(commit|push|pull|merge|rebase|branch)\b",
r"\bdocker(file)?\b",
r"\bkubernetes\b",
]
_REASONING_PATTERNS: Sequence[str] = [
r"\banalyze\b",
r"\bcompare\b",
r"\bpros?\s+and\s+cons?\b",
r"\btrade.?offs?\b",
r"\bstep.by.step\b",
r"\bexplain\s+(how|why|the\s+reason)\b",
r"\bplan\b",
r"\bstrateg(y|ize)\b",
r"\bevaluat(e|ion)\b",
r"\bbreak\s+(down|it\s+down)\b",
r"\bcomplex\b",
r"\bmulti.step\b",
r"\bdiagnos(e|is)\b",
r"\bdecid(e|ing)\b",
r"\bweigh\s+(the\s+)?options?\b",
]
_SEARCH_PATTERNS: Sequence[str] = [
r"\bsearch\s+(for|the\s+web|online|internet)\b",
r"\blook\s+up\b",
r"\bfind\s+(me\s+)?(information|info|news|articles?)\b",
r"\bwhat('s|\s+is)\s+the\s+latest\b",
r"\bcurrent\s+(news|events?|status)\b",
r"\brecent(ly)?\b.*\b(happened|released|announced)\b",
r"\bbrowse\b",
r"\bweb\s+(search|results?)\b",
r"\bgoogle\b",
r"\bwikipedia\b",
]
_AMBIGUOUS_TOOL_PATTERNS: Sequence[str] = [
r"\bwhichever\s+tool\b",
r"\bwhatever\s+tool\b",
r"\buse\s+(?:a|some)\s+tool\b",
r"\bchoose\s+(?:the\s+)?(?:right|best)\s+tool\b",
]
_CLOUD_PROVIDERS = frozenset({"openai", "openrouter", "anthropic"})
_VISION_PATTERNS: Sequence[str] = [
r"\bimage\b",
r"\bphoto\b",
r"\bpicture\b",
r"\bscreenshot\b",
r"\bdescribe\s+(this|what\s+you\s+see|the\s+image)\b",
r"\bwhat('s|\s+is)\s+in\s+(the\s+)?(image|photo|picture)\b",
r"\bvisual\b",
r"\bocr\b",
r"\bread\s+(the\s+)?text\s+(in|from)\b",
r"\bdiagram\b",
r"\bchart\b",
]
_FAST_PATTERNS: Sequence[str] = [
r"^(hi|hello|hey|howdy|yo|sup)\b",
r"^(thanks?|thank\s+you)\b",
r"^(ok(ay)?|got\s+it|sounds?\s+good)\b",
r"\bwhat\s+(time|day|date)\s+is\s+it\b",
r"\bwhat('s|\s+is)\s+(today|the\s+time|the\s+date)\b",
r"\bhow\s+(are\s+you|old|tall|far|long|many|much)\b",
r"^\s*\d[\d\s\+\-\*\/\.\(\)]*\s*[=\?]?\s*$", # arithmetic
r"\byes\b",
r"\bno\b",
r"\bmaybe\b",
r"\brepeat\s+(that|it|yourself)\b",
r"\bwhat\s+did\s+you\s+say\b",
]
# ---------------------------------------------------------------------------
# Ollama model detection
# ---------------------------------------------------------------------------
# Model ID prefixes that identify non-Ollama providers.
_OPENAI_PREFIXES = (
"gpt-",
"text-",
"o1",
"o3",
"babbage",
"davinci",
"curie",
"ada",
"claude-",
"gemini-",
"mistral-", # Mistral cloud API (vs local Ollama "mistral")
)
_ROUTING_FIELDS = ("default", "coding", "reasoning", "search", "vision", "fast")
def _is_ollama_model(model_id: str) -> bool:
"""Return True if *model_id* looks like an Ollama model (not OpenAI / HuggingFace)."""
if not model_id:
return False
lower = model_id.lower()
# HuggingFace / local Transformers paths contain a slash (org/model-name)
if "/" in model_id:
return False
for prefix in _OPENAI_PREFIXES:
if lower.startswith(prefix):
return False
return True
def _matches_any(text: str, patterns: Sequence[str]) -> bool:
lower = text.lower()
return any(re.search(p, lower) for p in patterns)
# ---------------------------------------------------------------------------
# ModelRouter
# ---------------------------------------------------------------------------
class ModelRouter:
"""Classifies user messages and resolves target model with Ollama availability check.
Classification uses keyword heuristics (no LLM call). Priority order:
vision > coding > search > reasoning > fast > default.
Ollama availability is probed at init time only when at least one routing
target is an Ollama model. OpenAI and HuggingFace models skip the check.
The available-model list is refreshed every *refresh_interval* seconds via
a daemon thread.
"""
_DEFAULT_OLLAMA_URL = "http://localhost:11434"
_DEFAULT_REFRESH_INTERVAL = 60
_DEFAULT_COOLDOWN_SECONDS = 3600 # 1 hour
def __init__(
self,
ollama_base_url: str = _DEFAULT_OLLAMA_URL,
routing_config: object = None,
refresh_interval: int = _DEFAULT_REFRESH_INTERVAL,
cooldown_seconds: int = _DEFAULT_COOLDOWN_SECONDS,
notify_callback: Callable[[str], None] | None = None,
provider_reliability: ProviderReliability | None = None,
) -> None:
self._ollama_base_url = ollama_base_url.rstrip("/")
self._refresh_interval = refresh_interval
self._available_ollama_models: set[str] = set()
self._stop_event = threading.Event()
self._refresh_thread: threading.Thread | None = None
# Cloud rate-limit cooldown state
self._cooldown_seconds = cooldown_seconds
self._cloud_limited_until: float = 0.0 # monotonic seconds
self._notify_callback = notify_callback
self._provider_reliability = provider_reliability or ProviderReliability(
cooldown_seconds=cooldown_seconds
)
# Determine whether any routing target could be an Ollama model.
self._needs_ollama = self._any_ollama_target(routing_config)
if self._needs_ollama:
self._fetch_ollama_models()
self._start_refresh_thread()
# ------------------------------------------------------------------
# Ollama availability helpers
# ------------------------------------------------------------------
@staticmethod
def _any_ollama_target(routing_config: object) -> bool:
if routing_config is None:
return False
return any(_is_ollama_model(getattr(routing_config, f, "") or "") for f in _ROUTING_FIELDS)
def _fetch_ollama_models(self) -> None:
"""Fetch available Ollama models from ``/api/tags`` and update the cache."""
url = f"{self._ollama_base_url}/api/tags"
try:
req = urllib.request.Request(url, headers={"Accept": "application/json"})
with urllib.request.urlopen(req, timeout=5) as resp: # noqa: S310
data = json.loads(resp.read())
entries: list[dict[str, str]] = data.get("models", [])
names: set[str] = set()
for entry in entries:
name = entry.get("name", "")
if name:
names.add(name)
# Also store the name without the tag (e.g. "llama3" from "llama3:latest")
names.add(name.split(":")[0])
self._available_ollama_models = names
logger.debug("model_router: Ollama models available: %s", sorted(names))
except Exception as exc:
logger.warning("model_router: failed to fetch Ollama model list from %s: %s", url, exc)
self._available_ollama_models = set()
def _start_refresh_thread(self) -> None:
def _loop() -> None:
while not self._stop_event.wait(self._refresh_interval):
self._fetch_ollama_models()
t = threading.Thread(target=_loop, daemon=True, name="ModelRouter-ollama-refresh")
t.start()
self._refresh_thread = t
def _is_available(self, model_id: str) -> bool:
"""Return True if *model_id* is usable.
Non-Ollama models (OpenAI, HuggingFace) are always considered available.
Ollama models are checked against the cached ``/api/tags`` response.
"""
if not model_id:
return False
if not _is_ollama_model(model_id):
return True # OpenAI / HuggingFace — always available
name_no_tag = model_id.split(":")[0]
return (
model_id in self._available_ollama_models
or name_no_tag in self._available_ollama_models
)
# ------------------------------------------------------------------
# Cloud rate-limit / quota fallback (US-045)
# ------------------------------------------------------------------
def _reliability(self) -> ProviderReliability:
reliability = getattr(self, "_provider_reliability", None)
if reliability is None:
reliability = ProviderReliability(
cooldown_seconds=int(
getattr(self, "_cooldown_seconds", self._DEFAULT_COOLDOWN_SECONDS)
)
)
self._provider_reliability = reliability
return reliability
def cloud_limit_hit(self, status_code: int, *, provider: str = "cloud") -> None:
"""Record legacy cloud quota/rate-limit evidence and provider cooldown."""
if status_code not in (429, 402):
return
self._cloud_limited_until = time.monotonic() + self._cooldown_seconds
self._reliability().record_failure(
provider,
ProviderFailureKind.RATE_LIMIT,
cooldown_seconds=self._cooldown_seconds,
)
msg = (
f"Cloud limit reached (HTTP {status_code}), switching to local model. "
f"Cloud will be retried in {self._cooldown_seconds // 60} min."
)
logger.warning("model_router: %s", msg)
if self._notify_callback is not None:
self._notify_callback(msg)
def _cloud_in_cooldown(self) -> bool:
"""Return True if cloud is currently under the legacy global cooldown."""
return time.monotonic() < self._cloud_limited_until
@property
def provider_reliability(self) -> ProviderReliability:
return self._reliability()
def record_provider_success(
self, provider: str, *, latency_ms: float | int | None = None
) -> None:
self._reliability().record_success(provider, latency_ms=latency_ms)
def record_provider_failure(
self,
provider: str,
kind: ProviderFailureKind | str,
*,
latency_ms: float | int | None = None,
) -> None:
self._reliability().record_failure(provider, kind, latency_ms=latency_ms)
def provider_diagnostics(self) -> list[dict[str, object]]:
return self._reliability().diagnostics()
def select_provider(
self, candidates: Sequence[ProviderRouteCandidate]
) -> ProviderRouteSelection:
"""Choose the first healthy configured provider in stable preference order."""
evidence: list[str] = []
fallback_reason: str | None = None
for index, candidate in enumerate(candidates):
if not candidate.provider or not candidate.model:
continue
status = self._reliability().status(candidate.provider)
if status.available:
if index == 0 and not evidence:
evidence.append("provider_primary_selected")
else:
evidence.append("provider_fallback_selected")
return ProviderRouteSelection(
provider=candidate.provider,
model=candidate.model,
fallback_reason=fallback_reason,
evidence=tuple(evidence),
)
evidence.append(f"provider_{status.provider}_{status.state}")
fallback_reason = (
"provider_cooldown" if status.state == "cooldown" else "provider_unavailable"
)
evidence.append("all_providers_unavailable")
return ProviderRouteSelection(
provider="",
model="",
fallback_reason="all_providers_unavailable",
evidence=tuple(evidence),
)
# ------------------------------------------------------------------
# Public API
# ------------------------------------------------------------------
def classify(self, message: str) -> TaskCategory:
"""Return the TaskCategory that best fits *message*."""
if _matches_any(message, _VISION_PATTERNS):
return TaskCategory.vision
if _matches_any(message, _CODING_PATTERNS):
return TaskCategory.coding
if _matches_any(message, _SEARCH_PATTERNS):
return TaskCategory.search
if _matches_any(message, _REASONING_PATTERNS):
return TaskCategory.reasoning
if _matches_any(message, _FAST_PATTERNS):
return TaskCategory.fast
return TaskCategory.default
def resolve_model(self, category: TaskCategory, routing_config: object) -> str:
"""Resolve the model identifier for *category* with availability fallback.
Returns an empty string when no routing override is configured,
meaning the caller should use the globally-configured LLM model.
If the preferred model for *category* is unavailable (Ollama not
running or model not pulled), falls back to ``routing_config.default``
and logs a warning.
"""
if routing_config is None:
return ""
model: str = getattr(routing_config, str(category), "") or ""
default: str = getattr(routing_config, "default", "") or ""
if model:
if self._is_available(model):
return model
logger.warning(
"model_router: model %r for category %r is not available in Ollama; "
"falling back to default %r",
model,
str(category),
default,
)
# Return default regardless of its availability (caller handles final fallback)
return default
# ------------------------------------------------------------------
# ModelRouter 2.0 decision contract (US-110)
# ------------------------------------------------------------------
_LOW_CONFIDENCE_THRESHOLD = 0.60
def _decision_signals(
self,
message: str,
*,
requires_tools: bool,
tool_choice_confidence: float | None,
) -> tuple[TaskCategory, str, float, list[str]]:
"""Derive bounded, content-free routing signals from one request."""
category = self.classify(message)
evidence: list[str] = []
ambiguous_tool_request = _matches_any(message, _AMBIGUOUS_TOOL_PATTERNS)
effective_requires_tools = requires_tools or ambiguous_tool_request
if ambiguous_tool_request:
evidence.append("ambiguous_tool_request")
complexity, confidence = self._base_complexity_confidence(
category,
message,
requires_tools=effective_requires_tools,
evidence=evidence,
)
if effective_requires_tools:
complexity = "complex"
evidence.append("tools_required")
confidence = self._apply_tool_confidence(
confidence,
ambiguous_tool_request=ambiguous_tool_request,
tool_choice_confidence=tool_choice_confidence,
evidence=evidence,
)
return category, complexity, confidence, evidence
def _base_complexity_confidence(
self,
category: TaskCategory,
message: str,
*,
requires_tools: bool,
evidence: list[str],
) -> tuple[str, float]:
if category is TaskCategory.fast:
evidence.append("fast_pattern")
return "simple", 0.98
if category in {
TaskCategory.coding,
TaskCategory.reasoning,
TaskCategory.search,
TaskCategory.vision,
}:
evidence.append(f"category_{category}")
return "complex", 0.90
evidence.append("no_category_match")
return self.estimate_complexity(message, requires_tools=requires_tools), 0.80
def _apply_tool_confidence(
self,
confidence: float,
*,
ambiguous_tool_request: bool,
tool_choice_confidence: float | None,
evidence: list[str],
) -> float:
effective_confidence = tool_choice_confidence
if ambiguous_tool_request and effective_confidence is None:
effective_confidence = 0.40
if effective_confidence is None:
return confidence
bounded = max(0.0, min(1.0, float(effective_confidence)))
if bounded < self._LOW_CONFIDENCE_THRESHOLD:
evidence.append("tool_choice_low_confidence")
return bounded
@staticmethod
def _configured_models(
routing_config: object,
category: TaskCategory,
current_model: str,
) -> tuple[str, str]:
config = routing_config
default_model = (getattr(config, "default", "") or "") if config is not None else ""
fast_model = (getattr(config, "fast", "") or "") if config is not None else ""
fast_model = fast_model or default_model or current_model
category_model = ""
if config is not None and category not in {TaskCategory.fast, TaskCategory.default}:
category_model = getattr(config, str(category), "") or ""
reasoning_model = (getattr(config, "reasoning", "") or "") if config is not None else ""
return fast_model, category_model or reasoning_model or default_model
@staticmethod
def _routing_policy_conflict(routing_mode: str, provider: str) -> str | None:
normalized_mode = (
routing_mode
if routing_mode in {"local_preferred", "local_only", "cloud_only"}
else "local_preferred"
)
normalized_provider = provider.strip().lower()
provider_is_cloud = normalized_provider in _CLOUD_PROVIDERS
if normalized_mode == "local_only" and provider_is_cloud:
return "local_only_provider_conflict"
if normalized_mode == "cloud_only" and normalized_provider and not provider_is_cloud:
return "cloud_provider_unconfigured"
return None
def _deep_fallback_reason(
self,
*,
deep_provider_available: bool,
deep_model: str,
fast_model: str,
) -> str | None:
if not deep_provider_available:
return "deep_provider_unavailable"
if not deep_model or deep_model == fast_model:
return "deep_model_unconfigured"
if not self._is_available(deep_model):
return "deep_model_unavailable"
return None
@staticmethod
def _decision(
*,
category: TaskCategory,
complexity: str,
confidence: float,
evidence: list[str],
tier: str,
model: str,
escalation_count: int = 0,
fallback_reason: str | None = None,
) -> ModelRouteDecision:
return ModelRouteDecision(
category=category,
complexity=complexity,
confidence=confidence,
evidence=tuple(dict.fromkeys(evidence)),
tier=tier,
model=model,
escalation_count=min(escalation_count, 1),
fallback_reason=fallback_reason,
)
def decide(
self,
message: str,
*,
routing_config: object = None,
current_model: str = "",
requires_tools: bool = False,
tool_choice_confidence: float | None = None,
deep_provider_available: bool = True,
routing_mode: str = "local_preferred",
provider: str = "",
) -> ModelRouteDecision:
"""Choose a fast/deep route and return bounded, privacy-safe evidence."""
category, complexity, confidence, evidence = self._decision_signals(
message,
requires_tools=requires_tools,
tool_choice_confidence=tool_choice_confidence,
)
fast_model, deep_model = self._configured_models(routing_config, category, current_model)
low_confidence = confidence < self._LOW_CONFIDENCE_THRESHOLD
wants_deep = complexity == "complex" or low_confidence
escalation_count = 1 if low_confidence and wants_deep else 0
policy_conflict = self._routing_policy_conflict(routing_mode, provider)
if policy_conflict is not None:
evidence.append(policy_conflict)
return self._decision(
category=category,
complexity=complexity,
confidence=confidence,
evidence=evidence,
tier="fast",
model=current_model or fast_model,
fallback_reason=policy_conflict,
)
if not wants_deep:
return self._decision(
category=category,
complexity=complexity,
confidence=confidence,
evidence=evidence,
tier="fast",
model=fast_model,
)
provider_fallback_reason: str | None = None
normalized_provider = provider.strip().lower()
if deep_provider_available and normalized_provider:
provider_status = self._reliability().status(normalized_provider)
if not provider_status.available:
deep_provider_available = False
evidence.append(f"provider_{provider_status.provider}_{provider_status.state}")
provider_fallback_reason = f"deep_provider_{provider_status.state}"
fallback_reason = provider_fallback_reason or self._deep_fallback_reason(
deep_provider_available=deep_provider_available,
deep_model=deep_model,
fast_model=fast_model,
)
if fallback_reason is not None:
evidence.append(fallback_reason)
return self._decision(
category=category,
complexity=complexity,
confidence=confidence,
evidence=evidence,
tier="fast",
model=fast_model or current_model,
escalation_count=escalation_count,
fallback_reason=fallback_reason,
)
evidence.append("deep_route_selected")
return self._decision(
category=category,
complexity=complexity,
confidence=confidence,
evidence=evidence,
tier="deep",
model=deep_model,
escalation_count=escalation_count,
)
# ------------------------------------------------------------------
# Smart local/cloud routing (US-044)
# ------------------------------------------------------------------
_SIMPLE_TOKEN_THRESHOLD = 200
@staticmethod
def estimate_complexity(message: str, *, requires_tools: bool = False) -> str:
"""Return ``"simple"`` or ``"complex"`` for *message*.
A query is **simple** when:
- estimated token count < 200 (approximated as ``len(message.split()) * 1.3``)
- no tools are required
Everything else is **complex**.
"""
estimated_tokens = len(message.split()) * 1.3
if requires_tools or estimated_tokens >= ModelRouter._SIMPLE_TOKEN_THRESHOLD:
return "complex"
return "simple"
def route(
self,
message: str,
*,
local_model: str,
cloud_model: str,
routing_mode: str = "local_preferred",
requires_tools: bool = False,
) -> str:
"""Return the model identifier to use for *message*.
Args:
message: The user message to route.
local_model: Ollama model identifier for local inference.
cloud_model: Cloud provider model identifier (e.g. ``"gpt-4o"``).
routing_mode: One of ``"local_preferred"``, ``"cloud_only"``,
``"local_only"``. Invalid values fall back to
``"local_preferred"`` with a warning.
requires_tools: Whether the query needs tool dispatch.
Returns:
Model identifier string.
"""
valid_modes = {"local_preferred", "cloud_only", "local_only"}
if routing_mode not in valid_modes:
logger.warning(
"model_router: unknown routing_mode %r; falling back to 'local_preferred'",
routing_mode,
)
routing_mode = "local_preferred"
cloud_available = cloud_model and not self._cloud_in_cooldown()
if routing_mode == "cloud_only":
if cloud_available:
return cloud_model
# Cloud is in cooldown — fall back to local with notification
logger.warning(
"model_router: cloud_only mode but cloud is in cooldown; "
"falling back to local model %r",
local_model,
)
return local_model or cloud_model
if routing_mode == "local_only":
if local_model and self._is_available(local_model):
return local_model
logger.warning(
"model_router: local_only mode but local model %r is unavailable; "
"cloud fallback is disabled by policy",
local_model,
)
return local_model
# local_preferred
complexity = self.estimate_complexity(message, requires_tools=requires_tools)
local_available = bool(local_model) and self._is_available(local_model)
if complexity == "simple":
if local_available:
return local_model
logger.warning(
"model_router: local model %r is unavailable; routing simple query to cloud",
local_model,
)
return cloud_model
else:
# complex query → prefer cloud if configured and not in cooldown
if cloud_available:
return cloud_model
if local_available:
return local_model
logger.warning(
"model_router: no cloud model available and local model %r is unavailable",
local_model,
)
return local_model or cloud_model