caudate-cli 0.1.0__py3-none-any.whl

llm/router.py

@@ -0,0 +1,373 @@
+"""Dual-process model router — System 1 (fast) vs System 2 (slow).
+
+Cognos runs a tiered LLM stack inspired by Kahneman's Thinking Fast and Slow:
+
+  - System 1: a small, fast Ollama model. Handles routine tool dispatch,
+    short answers, pattern-matching — the "autopilot" of the agent.
+  - System 2: a large, capable model. Handles planning, reflection,
+    meta-learning, complex synthesis, anything that benefits from depth.
+
+The `Router` scores every LLM call against a policy and picks the right tier.
+The `DualLLMProvider` wraps two `LLMProvider`s and exposes the same interface
+as a single provider, so agentic loops and tools don't know or care about
+routing.
+
+Caller tags (passed via the `caller` kwarg that LLMProvider already accepts)
+let deliberate-reasoning modules (planning, reflection, meta, compaction)
+always take System 2, while the agentic loop itself routes heuristically.
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass, field
+from typing import Any, AsyncIterator
+
+from pydantic import BaseModel
+
+from core.schemas import StreamEvent
+from llm.provider import LLMProvider, LLMResponse
+
+logger = logging.getLogger(__name__)
+
+
+# Weights sum to 1.0; tuning note: raising `keyword` weight makes the router
+# escalate on intent more aggressively; raising `depth` escalates on
+# long conversations (synthesis time).
+_WEIGHTS = {"length": 0.25, "depth": 0.10, "keyword": 0.45, "mood": 0.20}
+
+_COMPLEX_KEYWORDS = (
+    "why", "decide", "strategy", "compare", "evaluate", "analyze",
+    "refactor", "design", "architecture", "debug", "explain",
+    "plan", "reason", "approach", "trade-off", "tradeoff", "rationale",
+    "critique", "assess", "diagnose",
+)
+
+
+@dataclass
+class RoutingPolicy:
+    """Policy knobs for the Router. Defaults favor fast when unsure."""
+
+    complexity_threshold: float = 0.40
+    slow_caller_tags: set[str] = field(default_factory=lambda: {
+        "planning", "reflection", "meta", "compaction",
+    })
+    fast_caller_tags: set[str] = field(default_factory=set)
+    escalate_on_stuck: bool = True
+    length_tokens_saturation: int = 4000
+    depth_messages_saturation: int = 20
+
+
+@dataclass
+class RoutingDecision:
+    tier: str  # "fast" | "slow"
+    provider: LLMProvider
+    score: float
+    reasons: list[str]
+
+
+class Router:
+    """Pick a provider for an LLM call based on context + policy."""
+
+    def __init__(
+        self,
+        fast: LLMProvider,
+        slow: LLMProvider,
+        policy: RoutingPolicy | None = None,
+        mood=None,
+    ):
+        self.fast = fast
+        self.slow = slow
+        self.policy = policy or RoutingPolicy()
+        self.mood = mood  # optional MoodState for escalation
+        self.caudate = None  # optional CaudateObserver for ADVISOR+ override
+        self.stats = {"fast": 0, "slow": 0}
+        self.last_decision: RoutingDecision | None = None
+
+    def set_mood(self, mood) -> None:
+        self.mood = mood
+
+    def set_caudate(self, caudate) -> None:
+        """Wire a CaudateObserver in. Once set, her tier prediction
+        overrides the heuristic at ADVISOR level or higher."""
+        self.caudate = caudate
+
+    def choose(
+        self,
+        messages: list[dict[str, Any]] | None = None,
+        tools: list[dict] | None = None,
+        caller: str | None = None,
+    ) -> RoutingDecision:
+        """Return the provider to use for this call."""
+        messages = messages or []
+        tools = tools or []
+
+        # Caller-tag overrides are absolute.
+        if caller and caller in self.policy.slow_caller_tags:
+            decision = RoutingDecision(
+                tier="slow", provider=self.slow, score=1.0,
+                reasons=[f"caller={caller}"],
+            )
+            return self._record(decision)
+        if caller and caller in self.policy.fast_caller_tags:
+            decision = RoutingDecision(
+                tier="fast", provider=self.fast, score=0.0,
+                reasons=[f"caller={caller}"],
+            )
+            return self._record(decision)
+
+        # Stuck mood is an absolute escalation — when frustration/failure
+        # streak is high, the fast model has already proven insufficient.
+        if self.mood is not None and self.policy.escalate_on_stuck:
+            try:
+                if self.mood.should_defer_to_user():
+                    decision = RoutingDecision(
+                        tier="slow", provider=self.slow, score=1.0,
+                        reasons=["mood=stuck"],
+                    )
+                    return self._record(decision)
+            except Exception:
+                pass
+
+        score, reasons = self._score(messages, tools)
+
+        # Caudate override at ADVISOR+ trust. She's earned the right to
+        # pick the tier; the heuristic becomes a fallback if she's silent.
+        caudate_pred = None
+        try:
+            if self.caudate is not None and self.caudate.can_advise():
+                caudate_pred = self.caudate._last_prediction
+        except Exception:
+            caudate_pred = None
+        if caudate_pred is not None and caudate_pred.tier_confidence >= 0.55:
+            tier = caudate_pred.tier
+            reasons.append(f"caudate={tier}@{caudate_pred.tier_confidence:.2f}")
+        else:
+            tier = "slow" if score >= self.policy.complexity_threshold else "fast"
+
+        decision = RoutingDecision(
+            tier=tier,
+            provider=self.slow if tier == "slow" else self.fast,
+            score=score,
+            reasons=reasons,
+        )
+        return self._record(decision)
+
+    def _record(self, decision: RoutingDecision) -> RoutingDecision:
+        self.last_decision = decision
+        self.stats[decision.tier] = self.stats.get(decision.tier, 0) + 1
+        logger.info(
+            f"Router -> {decision.tier} (score={decision.score:.2f}, "
+            f"reasons={', '.join(decision.reasons)})"
+        )
+        return decision
+
+    def _score(
+        self,
+        messages: list[dict[str, Any]],
+        tools: list[dict],
+    ) -> tuple[float, list[str]]:
+        # Length feature — total chars across all messages, saturating.
+        total_chars = sum(len(_message_text(m)) for m in messages)
+        length = min(1.0, total_chars / (self.policy.length_tokens_saturation * 4))
+
+        # Depth feature — how many messages deep are we.
+        depth = min(1.0, len(messages) / self.policy.depth_messages_saturation)
+
+        # Keyword feature — does the latest user message hint at deep reasoning.
+        latest_user = next(
+            (_message_text(m) for m in reversed(messages)
+             if m.get("role") == "user"),
+            "",
+        ).lower()
+        keyword = 1.0 if any(k in latest_user for k in _COMPLEX_KEYWORDS) else 0.0
+
+        # Mood feature — slow down when the agent is uncertain or stuck.
+        mood_score = 0.0
+        if self.mood is not None and self.policy.escalate_on_stuck:
+            try:
+                if self.mood.should_defer_to_user() or self.mood.should_slow_down():
+                    mood_score = 1.0
+            except Exception:
+                mood_score = 0.0
+
+        score = (
+            _WEIGHTS["length"] * length
+            + _WEIGHTS["depth"] * depth
+            + _WEIGHTS["keyword"] * keyword
+            + _WEIGHTS["mood"] * mood_score
+        )
+        reasons = [
+            f"len={length:.2f}",
+            f"depth={depth:.2f}",
+            f"kw={keyword:.2f}",
+            f"mood={mood_score:.2f}",
+        ]
+        return score, reasons
+
+
+class DualLLMProvider:
+    """Drop-in replacement for LLMProvider that routes per call.
+
+    Exposes the same async surface (chat, complete, stream, structured_output).
+    Each method passes the `caller` tag through to the router.
+    """
+
+    def __init__(
+        self,
+        fast: LLMProvider,
+        slow: LLMProvider,
+        policy: RoutingPolicy | None = None,
+        mood=None,
+    ):
+        self._fast = fast
+        self._slow = slow
+        self.router = Router(fast=fast, slow=slow, policy=policy, mood=mood)
+
+    # ------------------------------------------------------------------
+    # Compatibility surface (mirrors LLMProvider)
+    # ------------------------------------------------------------------
+
+    @property
+    def model(self) -> str:
+        """A synthetic id so sessions/logs show the tiered setup."""
+        return f"dual[fast={self._fast.model},slow={self._slow.model}]"
+
+    @property
+    def temperature(self) -> float:
+        return self._fast.temperature
+
+    @property
+    def max_tokens(self) -> int:
+        return self._slow.max_tokens
+
+    def switch_model(self, model: str) -> None:
+        """Switch the slow (System 2) tier — treat that as the primary target."""
+        self._slow.switch_model(model)
+
+    def set_mood(self, mood) -> None:
+        self.router.set_mood(mood)
+
+    # ------------------------------------------------------------------
+    # Hot-swap support — used by /system1 and /system2 slash commands
+    # ------------------------------------------------------------------
+
+    @property
+    def fast_model(self) -> str:
+        return self._fast.model
+
+    @property
+    def slow_model(self) -> str:
+        return self._slow.model
+
+    @property
+    def last_tier(self) -> str | None:
+        """Which tier was used on the most recent call ('fast' / 'slow' / None)."""
+        return self.router.last_decision.tier if self.router.last_decision else None
+
+    @property
+    def last_provider_model(self) -> str | None:
+        """The actual model id that ran on the most recent call."""
+        d = self.router.last_decision
+        if d is None:
+            return None
+        return self._fast.model if d.tier == "fast" else self._slow.model
+
+    def set_fast(self, model: str) -> None:
+        """Hot-swap the fast tier without rebuilding the agent."""
+        self._fast.switch_model(model)
+        # Router holds a reference to the same _fast object, so it
+        # picks up the new model id automatically on next call.
+
+    def set_slow(self, model: str) -> None:
+        """Hot-swap the slow tier without rebuilding the agent."""
+        self._slow.switch_model(model)
+
+    # ------------------------------------------------------------------
+    # Calls
+    # ------------------------------------------------------------------
+
+    async def complete(
+        self,
+        prompt: str,
+        system: str | None = None,
+        temperature: float | None = None,
+        max_tokens: int | None = None,
+        response_format: dict | None = None,
+        caller: str | None = None,
+    ) -> LLMResponse:
+        msgs: list[dict[str, Any]] = []
+        if system:
+            msgs.append({"role": "system", "content": system})
+        msgs.append({"role": "user", "content": prompt})
+        provider = self.router.choose(msgs, None, caller).provider
+        return await provider.complete(
+            prompt=prompt, system=system,
+            temperature=temperature, max_tokens=max_tokens,
+            response_format=response_format,
+        )
+
+    async def chat(
+        self,
+        messages: list[dict[str, Any]],
+        temperature: float | None = None,
+        max_tokens: int | None = None,
+        response_format: dict | None = None,
+        tools: list[dict] | None = None,
+        tool_choice: str | None = None,
+        caller: str | None = None,
+    ) -> LLMResponse:
+        provider = self.router.choose(messages, tools, caller).provider
+        return await provider.chat(
+            messages=messages,
+            temperature=temperature, max_tokens=max_tokens,
+            response_format=response_format,
+            tools=tools, tool_choice=tool_choice,
+        )
+
+    async def stream(
+        self,
+        messages: list[dict[str, Any]],
+        temperature: float | None = None,
+        max_tokens: int | None = None,
+        tools: list[dict] | None = None,
+        tool_choice: str | None = None,
+        caller: str | None = None,
+    ) -> AsyncIterator[StreamEvent]:
+        provider = self.router.choose(messages, tools, caller).provider
+        async for event in provider.stream(
+            messages=messages,
+            temperature=temperature, max_tokens=max_tokens,
+            tools=tools, tool_choice=tool_choice,
+        ):
+            yield event
+
+    async def structured_output(
+        self,
+        prompt: str,
+        system: str | None = None,
+        schema_hint: str = "",
+        response_model: type[BaseModel] | None = None,
+        caller: str | None = None,
+    ) -> Any:
+        # Structured output almost always implies deliberate reasoning.
+        # Default to "structured" caller so the router's policy can map it.
+        tag = caller or "structured"
+        msgs = [{"role": "user", "content": prompt}]
+        if system:
+            msgs.insert(0, {"role": "system", "content": system})
+        provider = self.router.choose(msgs, None, tag).provider
+        return await provider.structured_output(
+            prompt=prompt, system=system,
+            schema_hint=schema_hint,
+            response_model=response_model,
+        )
+
+
+def _message_text(message: dict[str, Any]) -> str:
+    """Extract a single text chunk from a chat message (handling tool_calls)."""
+    out = str(message.get("content") or "")
+    for tc in message.get("tool_calls", []) or []:
+        fn = tc.get("function", {})
+        out += f" {fn.get('name', '')} {fn.get('arguments', '')}"
+    return out