gdmcode 0.1.0__py3-none-any.whl

src/models/router.py

@@ -0,0 +1,418 @@
+"""ModelRouter — smart model tier routing based on task complexity.
+
+Routes tasks to the appropriate model tier (Scout, Coder, Thinker, Reasoner)
+using keyword analysis, context size, file count, and failure history.
+
+Two complementary routing mechanisms:
+  1. route_task(ctx)  — overall task complexity → tier (for planning / first call)
+  2. route_tool(name) — per-tool optimal tier   (for individual tool executions)
+  3. select_tier_for_turn() — per-turn micro-toggle respecting reasoning_mode flag
+"""
+from __future__ import annotations
+
+import logging
+from enum import Enum
+from dataclasses import dataclass, field
+
+from src.models.definitions import ModelTier, Provider
+
+__all__ = ["TaskContext", "ModelRouter", "FailureType", "EscalationContext", "EscalationDecision", "CircuitBreaker"]
+
+log = logging.getLogger(__name__)
+
+# ---------------------------------------------------------------------------
+# Keyword sets for complexity scoring
+# ---------------------------------------------------------------------------
+
+_HEAVY_KEYWORDS: frozenset[str] = frozenset({
+    "architect",
+    "design",
+    "why does",
+    "debug",
+    "analyze",
+    "performance",
+    "security",
+    "refactor",
+    "migrate",
+    "integrate",
+})
+
+_MEDIUM_KEYWORDS: frozenset[str] = frozenset({
+    "implement",
+    "create",
+    "build",
+    "add feature",
+    "fix",
+    "update",
+})
+
+_LIGHT_KEYWORDS: frozenset[str] = frozenset({
+    "rename",
+    "comment",
+    "format",
+    "typo",
+    "add test",
+})
+
+# ---------------------------------------------------------------------------
+# Tool routing sets
+# ---------------------------------------------------------------------------
+
+_SCOUT_TOOLS: frozenset[str] = frozenset({
+    "read_file",
+    "grep",
+    "list_dir",
+    "find_symbol",
+    "git_log",
+    "summarize",
+})
+
+_CODER_TOOLS: frozenset[str] = frozenset({
+    "apply_patch",
+    "write_file",
+    "write_tests",
+    "git_commit",
+})
+
+# Tools that score points toward a higher tier (per-turn classifier)
+_WRITE_TOOLS: frozenset[str] = frozenset({
+    "apply_patch",
+    "write_file",
+    "write_tests",
+})
+
+# ---------------------------------------------------------------------------
+# Scoring constants
+# ---------------------------------------------------------------------------
+
+_HEAVY_SCORE: int = 6
+_MEDIUM_SCORE: int = 4
+_LIGHT_SCORE: int = 1
+
+_LARGE_CONTEXT_THRESHOLD: int = 60_000
+_LARGE_CONTEXT_SCORE: int = 2
+
+_MANY_FILES_THRESHOLD: int = 5
+_MANY_FILES_SCORE: int = 1
+
+_FAILURE_SCORE: int = 2
+_MAX_COMPLEXITY_SCORE: int = 10
+
+_REASONER_THRESHOLD: int = 8
+_THINKER_THRESHOLD: int = 5
+_CODER_THRESHOLD: int = 2
+
+# Per-turn scoring (models-002)
+_TOOL_ONLY_PENALTY: int = -3        # pure tool turn, no text generation
+_WRITE_TOOL_SCORE: int = 4          # apply_patch / write_file / write_tests
+_TEST_FAILURE_SCORE: int = 5        # prompt contains failure output
+_MANY_PATHS_THRESHOLD: int = 3      # N file paths in prompt triggers bonus
+_MANY_PATHS_SCORE: int = 3          # bonus for 3+ file paths
+_FAILURE_ADDITIVE_SCORE: int = 3    # per failed_attempt (additive)
+
+# Injection cap: tool-result-only spike → cap at THINKER
+_INJECTION_CAP_TIER: str = ModelTier.THINKER
+
+# Cost-budget guard threshold
+_FORCED_REASONER_WARNING_THRESHOLD: int = 5
+
+# Escalation ladder — ordered lowest → highest
+_TIER_ORDER: list[str] = [
+    ModelTier.SCOUT,
+    ModelTier.CODER,
+    ModelTier.THINKER,
+    ModelTier.REASONER,
+]
+
+
+
+# ---------------------------------------------------------------------------
+# Failure-type enumeration and escalation dataclasses
+# ---------------------------------------------------------------------------
+
+class FailureType(str, Enum):
+    API_ERROR = "api_error"
+    TOOL_ERROR = "tool_error"
+    TEST_FAILURE = "test_failure"
+    QUALITY_FAILURE = "quality_failure"
+    REFUSAL = "refusal"
+    CONTEXT_OVERFLOW = "context_overflow"
+
+
+@dataclass
+class EscalationContext:
+    current_tier: str
+    failure_type: "FailureType"
+    failure_detail: str = ""
+    attempt_number: int = 1
+    cost_spent_usd: float = 0.0
+    transcript_summary: str = ""
+
+
+@dataclass
+class EscalationDecision:
+    next_tier: "str | None"
+    retry_prompt_addition: str
+    carry_transcript: bool
+
+
+@dataclass
+class CircuitBreaker:
+    max_cost_usd: float = 10.0
+    max_escalations: int = 3
+    _total_cost_usd: float = field(default=0.0, init=False, repr=False)
+    _escalation_count: int = field(default=0, init=False, repr=False)
+
+    def should_halt(self) -> bool:
+        return (self._total_cost_usd >= self.max_cost_usd
+                or self._escalation_count >= self.max_escalations)
+
+    def record_escalation(self, cost_usd: float = 0.0) -> None:
+        self._escalation_count += 1
+        self._total_cost_usd += cost_usd
+
+    def halt_reason(self) -> str:
+        if self._total_cost_usd >= self.max_cost_usd:
+            return f"Cost limit {self._total_cost_usd:.3f}/{self.max_cost_usd:.3f} USD reached"
+        return f"escalations {self._escalation_count}/{self.max_escalations} reached"
+
+    @property
+    def total_cost_usd(self) -> float:
+        return self._total_cost_usd
+
+    @property
+    def escalation_count(self) -> int:
+        return self._escalation_count
+
+@dataclass
+class TaskContext:
+    """Snapshot of agent state used for routing decisions."""
+
+    prompt: str
+    token_count: int = 0
+    files_touched: int = 0
+    failed_attempts: int = 0
+    provider: str = field(default=Provider.GROK)
+
+
+class ModelRouter:
+    """Routes tasks to the appropriate model tier based on complexity.
+
+    Two routing levels:
+    1. route_task(ctx)   — overall task complexity → model tier string
+    2. route_tool(name)  — per-tool optimal tier (Scout for reads, etc.)
+    3. select_tier_for_turn() — per-turn micro-toggle with mode override
+
+    Escalation is implicit in the complexity score: each failed_attempts
+    adds _FAILURE_SCORE points, pushing the task to a higher tier.
+    The explicit escalate() method is available for callers that need to
+    manually bump one tier (e.g. after a tool execution error).
+    """
+
+    def route_task(self, ctx: TaskContext) -> str:
+        """Return ModelTier constant for this task. Never raises."""
+        score = self._complexity_score(ctx)
+        log.debug(
+            "Task complexity score=%d for prompt=%r",
+            score,
+            ctx.prompt[:80],
+        )
+        if score >= _REASONER_THRESHOLD:
+            return ModelTier.REASONER
+        if score >= _THINKER_THRESHOLD:
+            return ModelTier.THINKER
+        if score >= _CODER_THRESHOLD:
+            return ModelTier.CODER
+        return ModelTier.SCOUT
+
+    def route_tool(self, tool_name: str) -> str:
+        """Return optimal ModelTier for a specific tool operation."""
+        if tool_name in _SCOUT_TOOLS:
+            return ModelTier.SCOUT
+        if tool_name in _CODER_TOOLS:
+            return ModelTier.CODER
+        return ModelTier.THINKER
+
+    def escalate(self, current_tier: str) -> str:
+        """Return next tier up from current_tier (caps at REASONER)."""
+        try:
+            idx = _TIER_ORDER.index(current_tier)
+        except ValueError:
+            log.warning(
+                "Unknown tier %r passed to escalate — defaulting to REASONER",
+                current_tier,
+            )
+            return ModelTier.REASONER
+        return _TIER_ORDER[min(idx + 1, len(_TIER_ORDER) - 1)]
+
+    def select_tier_for_turn(
+        self,
+        ctx: TaskContext,
+        *,
+        reasoning_mode: str = "auto",
+        tool_names: list[str] | None = None,
+        is_tool_result_only: bool = False,
+    ) -> str:
+        """Return ModelTier constant for a single agent turn.
+
+        Args:
+            ctx:                  Current task context (prompt, tokens, failures).
+            reasoning_mode:       "on" | "off" | "auto".  "on"/"off" override classifier.
+            tool_names:           Tools about to be called (or just called).  Used in
+                                  the per-turn scoring table.
+            is_tool_result_only:  True when the prompt text is only a tool result (no
+                                  user input this turn).  Enables injection-cap logic.
+
+        Returns:
+            ModelTier constant string.  Never raises.
+        """
+        mode = reasoning_mode.lower() if reasoning_mode else "auto"
+
+        if mode == "off":
+            log.info("gdm.router mode=off → SCOUT (forced)")
+            return ModelTier.SCOUT
+
+        if mode == "on":
+            log.info("gdm.router mode=on → REASONER (forced)")
+            return ModelTier.REASONER
+
+        # ── auto mode: deterministic scoring ──────────────────────────────
+        score, signals = self._per_turn_score(ctx, tool_names or [])
+
+        # Determine tier from score
+        if score >= _REASONER_THRESHOLD:
+            tier = ModelTier.REASONER
+        elif score >= _THINKER_THRESHOLD:
+            tier = ModelTier.THINKER
+        elif score >= _CODER_THRESHOLD:
+            tier = ModelTier.CODER
+        else:
+            tier = ModelTier.SCOUT
+
+        # Prompt-injection cap: if this turn is tool-result-only and the user's
+        # original message scored < 2, cap tier at THINKER.
+        user_score = self._keyword_score(ctx.prompt)
+        if is_tool_result_only and user_score < _CODER_THRESHOLD and tier == ModelTier.REASONER:
+            log.info(
+                "gdm.router injection-cap applied: tool-result score=%d → capping at THINKER",
+                score,
+            )
+            tier = _INJECTION_CAP_TIER
+
+        log.info(
+            "gdm.router turn=auto score=%d tier=%s signals=%s",
+            score,
+            tier,
+            signals,
+        )
+        return tier
+
+    # ------------------------------------------------------------------
+    # Private helpers
+    # ------------------------------------------------------------------
+
+
+    def escalate_with_context(self, ctx: "EscalationContext") -> "EscalationDecision":
+        ft = ctx.failure_type
+        log.warning(
+            "gdm.router escalating failure=%s tier=%s attempt=%d cost=%.3f",
+            ft.value, ctx.current_tier, ctx.attempt_number, ctx.cost_spent_usd,
+        )
+        if ft == FailureType.REFUSAL:
+            return EscalationDecision(next_tier=None, retry_prompt_addition="", carry_transcript=False)
+        if ft == FailureType.API_ERROR:
+            return EscalationDecision(
+                next_tier=ctx.current_tier,
+                retry_prompt_addition=f"Retry: {ctx.failure_detail}",
+                carry_transcript=False,
+            )
+        try:
+            idx = _TIER_ORDER.index(ctx.current_tier)
+        except ValueError:
+            idx = len(_TIER_ORDER) - 1
+        next_tier = _TIER_ORDER[min(idx + 1, len(_TIER_ORDER) - 1)]
+        carry = ft == FailureType.CONTEXT_OVERFLOW
+        if ft == FailureType.TOOL_ERROR:
+            prompt = f"Tool error ({ctx.failure_detail}). Retry with a different approach."
+        elif ft == FailureType.TEST_FAILURE:
+            prompt = f"Fix the failures: {ctx.failure_detail}"
+        elif ft == FailureType.QUALITY_FAILURE:
+            prompt = f"Quality issue: {ctx.failure_detail}. Improve the solution."
+        elif ft == FailureType.CONTEXT_OVERFLOW:
+            prompt = f"Summarize the context and continue. {ctx.failure_detail}".strip()
+        else:
+            prompt = f"Escalating due to {ft.value}: {ctx.failure_detail}"
+        return EscalationDecision(next_tier=next_tier, retry_prompt_addition=prompt, carry_transcript=carry)
+
+    def _per_turn_score(self, ctx: TaskContext, tool_names: list[str]) -> tuple[int, list[str]]:
+        """Compute per-turn score and return (score, signals) tuple."""
+        score = 0
+        signals: list[str] = []
+
+        # Tool-only turn penalty (only when purely read/scout tools, no write tools)
+        has_write_tools = any(t in _WRITE_TOOLS for t in tool_names)
+        if tool_names and not ctx.prompt.strip() and not has_write_tools:
+            score += _TOOL_ONLY_PENALTY
+            signals.append(f"tool_only{_TOOL_ONLY_PENALTY}")
+
+        # Write tools
+        if has_write_tools:
+            score += _WRITE_TOOL_SCORE
+            signals.append(f"write_tool+{_WRITE_TOOL_SCORE}")
+
+        # Test failure output
+        prompt_lower = ctx.prompt.lower()
+        if "error" in prompt_lower and ("traceback" in prompt_lower or "assert" in prompt_lower or "failed" in prompt_lower):
+            score += _TEST_FAILURE_SCORE
+            signals.append(f"test_failure+{_TEST_FAILURE_SCORE}")
+
+        # Many file paths (count occurrences of "/" or "\" path separators)
+        path_count = prompt_lower.count(".py") + prompt_lower.count(".ts") + prompt_lower.count(".js")
+        if path_count >= _MANY_PATHS_THRESHOLD:
+            score += _MANY_PATHS_SCORE
+            signals.append(f"many_files+{_MANY_PATHS_SCORE}")
+
+        # Large context
+        if ctx.token_count > _LARGE_CONTEXT_THRESHOLD:
+            score += _LARGE_CONTEXT_SCORE
+            signals.append(f"large_ctx+{_LARGE_CONTEXT_SCORE}")
+
+        # Failed attempts (additive)
+        if ctx.failed_attempts > 0:
+            fa_score = ctx.failed_attempts * _FAILURE_ADDITIVE_SCORE
+            score += fa_score
+            signals.append(f"failures+{fa_score}")
+
+        # Keywords
+        kw_score = self._keyword_score(ctx.prompt)
+        if kw_score > 0:
+            label = "heavy" if kw_score >= _HEAVY_SCORE else ("medium" if kw_score >= _MEDIUM_SCORE else "light")
+            score += kw_score
+            signals.append(f"{label}_kw+{kw_score}")
+
+        return score, signals
+
+    def _complexity_score(self, ctx: TaskContext) -> int:
+        """Score 0-10 based on keywords, context size, file count, and failure count."""
+        score = self._keyword_score(ctx.prompt)
+        if ctx.token_count > _LARGE_CONTEXT_THRESHOLD:
+            score += _LARGE_CONTEXT_SCORE
+        if ctx.files_touched > _MANY_FILES_THRESHOLD:
+            score += _MANY_FILES_SCORE
+        score += ctx.failed_attempts * _FAILURE_SCORE
+        return min(score, _MAX_COMPLEXITY_SCORE)
+
+    def _keyword_score(self, prompt: str) -> int:
+        """Return keyword-based sub-score for the prompt text.
+
+        Checks tier groups in descending order; returns the score for the
+        highest-severity keyword found (not additive across tiers).
+        """
+        prompt_lower = prompt.lower()
+        if any(kw in prompt_lower for kw in _HEAVY_KEYWORDS):
+            return _HEAVY_SCORE
+        if any(kw in prompt_lower for kw in _MEDIUM_KEYWORDS):
+            return _MEDIUM_SCORE
+        if any(kw in prompt_lower for kw in _LIGHT_KEYWORDS):
+            return _LIGHT_SCORE
+        return 0