Merge pull request #818 from kvandre12-commits/assist/pr-760-refresh

Refresh per-run model routing hook on current main

Author: mpfaffenberger

AGENTS.md

@@ -76,6 +76,7 @@ approval. With `fail_closed=True` its exception is reported as a block instead.
 | `invoke_agent` | Sub-agent invoked | `(*args, **kwargs) -> None` |
 | `agent_exception` | Unhandled agent error | `(exception, *args, **kwargs) -> None` |
 | `agent_run_start` | Before agent task | `(agent_name, model_name, session_id=None) -> None` |
+| `model_select` | Select a model for one run | `(*, agent_name, current_model, prompt, messages, session_id=None) -> str \| None` — first non-empty result wins |
 | `agent_run_end` | After agent run | `(agent_name, model_name, session_id=None, success=True, error=None, response_text=None, metadata=None) -> None` |
 | `load_prompt` | System prompt assembly | `() -> str \| None` |
 | `run_shell_command` | Before shell exec | `(context, command, cwd=None, timeout=60) -> dict \| None` (return `{"blocked": True}` to block, `{"rewrite": "<new cmd>"}` to transparently transform) |

code_puppy/agents/_runtime.py

@@ -693,6 +693,20 @@ async def _run_with_mcp_impl(
         # Hook failures must never block the run.
         pass
 
+    # Let a ``model_select`` hook route THIS turn to a different model (e.g.
+    # small-vs-large by complexity) before the pydantic agent is built. This
+    # resets any prior turn's auto choice, respects an explicit runtime
+    # override, and invalidates the cached agent when the model changes so the
+    # build below picks it up. No-op (and near-zero cost) if no plugin
+    # registered the hook.
+    try:
+        from code_puppy.model_switching import resolve_run_model_selection
+
+        resolve_run_model_selection(agent, prompt, agent._message_history, group_id)
+    except Exception:
+        # Selection must never block a run.
+        pass
+
     if agent._code_generation_agent is None:
         build_pydantic_agent(agent)
     pydantic_agent = agent._code_generation_agent

code_puppy/agents/base_agent.py

@@ -72,6 +72,11 @@ def __init__(self) -> None:
         self._last_model_name: Optional[str] = None
         self._runtime_model_name_override: Optional[str] = None
         self._runtime_system_prompt_additions: List[str] = []
+        # Model chosen by a ``model_select`` hook for the current run. Slots
+        # below an explicit runtime override but above pinned/JSON/global, and
+        # is reset at the start of every run (see resolve_run_model_selection),
+        # so it never leaks across turns.
+        self._auto_model_override: Optional[str] = None
         self._puppy_rules: Optional[str] = None
         self._mcp_servers: List[Any] = []
         self.cur_model: Optional[pydantic_ai.models.Model] = None
@@ -125,6 +130,14 @@ def set_runtime_model_name_override(self, model_name: Optional[str]) -> None:
         """
         self._runtime_model_name_override = model_name
 
+    def get_auto_model_override(self) -> Optional[str]:
+        """Return the model chosen by a ``model_select`` hook for this run."""
+        return self._auto_model_override
+
+    def set_auto_model_override(self, model_name: Optional[str]) -> None:
+        """Set the ``model_select``-chosen model for this run (not persisted)."""
+        self._auto_model_override = model_name
+
     @contextmanager
     def temporary_model_name_override(
         self, model_name: Optional[str]
@@ -154,6 +167,9 @@ def get_model_name(self) -> Optional[str]:
         override = self.get_runtime_model_name_override()
         if override:
             return override
+        auto = self.get_auto_model_override()
+        if auto:
+            return auto
         pinned = get_agent_pinned_model(self.name)
         return pinned if pinned else get_global_model_name()
 

code_puppy/agents/json_agent.py

@@ -220,6 +220,13 @@ def get_model_name(self) -> Optional[str]:
         if override:
             return override
 
+        # A ``model_select`` hook choice outranks the JSON ``model`` field so
+        # per-turn routing works for JSON agents too (see get_model_name in
+        # BaseAgent for the full precedence ladder).
+        auto = self.get_auto_model_override()
+        if auto:
+            return auto
+
         result = self._config.get("model")
         if result is None or (isinstance(result, str) and not result.strip()):
             result = super().get_model_name()

code_puppy/callbacks.py

@@ -46,6 +46,7 @@
     "agent_run_start",
     "agent_run_end",
     "agent_run_result",
+    "model_select",
     "register_mcp_catalog_servers",
     "register_browser_types",
     "register_model_providers",
@@ -132,6 +133,7 @@ def __repr__(self) -> str:
     "agent_run_start": [],
     "agent_run_end": [],
     "agent_run_result": [],
+    "model_select": [],
     "register_mcp_catalog_servers": [],
     "register_browser_types": [],
     "register_model_providers": [],
@@ -1154,6 +1156,50 @@ async def on_agent_run_start(
     )
 
 
+def on_model_select(
+    *,
+    agent_name: str,
+    current_model: str | None,
+    prompt: str,
+    messages: List[Any],
+    session_id: str | None = None,
+) -> str | None:
+    """Ask plugins to choose the model for the current run.
+
+    Fires once per run, before the pydantic agent is (re)built. Lets a plugin
+    route each turn to a different model based on the agent, the effective
+    ("would-be") model, and the message history -- e.g. a small model for
+    trivial turns and a frontier model when it matters, or a cost/latency/
+    failover policy.
+
+    Precedence: an explicit runtime override still wins over this hook; this
+    hook wins over the pinned / JSON / global model. The first callback to
+    return a non-empty string wins; return ``None`` to defer.
+
+    Args:
+        agent_name: Name of the agent about to run.
+        current_model: The model that would be used absent any hook.
+        prompt: The current user prompt after submission hooks have rewritten it.
+        messages: The agent's prior message history for this run.
+        session_id: Optional per-run identifier.
+
+    Returns:
+        A model name to use for this run, or ``None`` to keep ``current_model``.
+    """
+    results = _trigger_callbacks_sync(
+        "model_select",
+        agent_name=agent_name,
+        current_model=current_model,
+        prompt=prompt,
+        messages=messages,
+        session_id=session_id,
+    )
+    for result in results:
+        if isinstance(result, str) and result.strip():
+            return result
+    return None
+
+
 async def on_agent_run_end(
     agent_name: str,
     model_name: str,

code_puppy/model_switching.py

@@ -2,11 +2,65 @@
 
 from __future__ import annotations
 
-from typing import Optional
+from typing import Any, List, Optional
 
 from code_puppy.config import set_model_name
 
 
+def resolve_run_model_selection(
+    agent: Any,
+    prompt: str,
+    messages: List[Any],
+    session_id: Optional[str] = None,
+) -> Optional[str]:
+    """Apply the ``model_select`` hook for one run and return the chosen model.
+
+    Called once at the start of every run, BEFORE the pydantic agent is built.
+
+    Behaviour / precedence:
+    * The per-run auto override is always cleared first, so a choice never
+      leaks into a later turn.
+    * An explicit runtime override (e.g. ``invoke_agent_with_model``) wins and
+      short-circuits the hook entirely.
+    * Otherwise the hook is asked to pick a model. If it returns a name that
+      differs from the effective model, we install it as the run's auto
+      override and invalidate the cached pydantic agent so the build picks up
+      the new model.
+
+    Returns the chosen model name if the hook changed it, else ``None``.
+    Never raises -- a misbehaving selector must not break a run.
+    """
+    try:
+        # Reset any previous turn's auto choice so it can't leak forward.
+        if agent.get_auto_model_override():
+            agent.set_auto_model_override(None)
+            agent._code_generation_agent = None
+
+        # Explicit runtime override always wins -- don't even ask the hook.
+        if agent.get_runtime_model_name_override():
+            return None
+
+        from code_puppy.callbacks import on_model_select
+
+        current = agent.get_model_name()
+        selected = on_model_select(
+            agent_name=getattr(agent, "name", None),
+            current_model=current,
+            prompt=prompt,
+            messages=messages or [],
+            session_id=session_id,
+        )
+        if selected and selected != current:
+            agent.set_auto_model_override(selected)
+            # Force a rebuild so the new model is actually used this turn.
+            agent._code_generation_agent = None
+            return selected
+    except Exception:
+        # A broken selector must never block the agent run.
+        pass
+    return None
+
+
 def _get_effective_agent_model(agent) -> Optional[str]:
     """Safely fetch the effective model name for an agent."""
     try:

tests/test_model_select_hook.py

@@ -0,0 +1,152 @@
+"""Tests for the ``model_select`` hook and per-run model resolution.
+
+Covers ``callbacks.on_model_select`` and
+``model_switching.resolve_run_model_selection`` -- the two new pieces that let
+a plugin route each turn to a different model (small-vs-large auto mode, cost
+caps, failover, ...).
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from code_puppy.callbacks import (
+    clear_callbacks,
+    on_model_select,
+    register_callback,
+)
+from code_puppy.model_switching import resolve_run_model_selection
+
+
+@pytest.fixture(autouse=True)
+def _clean_hook():
+    clear_callbacks("model_select")
+    yield
+    clear_callbacks("model_select")
+
+
+class FakeAgent:
+    """Minimal stand-in exposing the model-override surface resolve() uses."""
+
+    def __init__(self, base="global-model", runtime=None, pinned=None):
+        self.name = "fake"
+        self._base = base
+        self._runtime = runtime
+        self._auto = None
+        self._code_generation_agent = object()  # non-None = "cached build"
+        self._message_history = []
+
+    def get_runtime_model_name_override(self):
+        return self._runtime
+
+    def get_auto_model_override(self):
+        return self._auto
+
+    def set_auto_model_override(self, name):
+        self._auto = name
+
+    def get_model_name(self):
+        # Mirrors BaseAgent precedence: runtime > auto > base.
+        return self._runtime or self._auto or self._base
+
+
+# ---- on_model_select -------------------------------------------------------
+
+
+def test_on_model_select_no_callbacks_returns_none():
+    assert (
+        on_model_select(
+            agent_name="fake",
+            current_model="m",
+            prompt="hello",
+            messages=[],
+            session_id=None,
+        )
+        is None
+    )
+
+
+def test_on_model_select_returns_first_nonempty_result():
+    register_callback("model_select", lambda **k: None)
+    register_callback("model_select", lambda **k: "")
+    register_callback("model_select", lambda **k: "small-model")
+
+    assert (
+        on_model_select(
+            agent_name="fake",
+            current_model="big",
+            prompt="fix the parser",
+            messages=[],
+            session_id="s",
+        )
+        == "small-model"
+    )
+
+
+def test_on_model_select_passes_context_to_callback():
+    seen = {}
+
+    def selector(**kwargs):
+        seen.update(kwargs)
+        return None
+
+    register_callback("model_select", selector)
+    on_model_select(
+        agent_name="orch",
+        current_model="big",
+        prompt="current turn",
+        messages=[1, 2],
+        session_id="x",
+    )
+    assert seen["agent_name"] == "orch"
+    assert seen["current_model"] == "big"
+    assert seen["prompt"] == "current turn"
+    assert seen["messages"] == [1, 2]
+    assert seen["session_id"] == "x"
+
+
+# ---- resolve_run_model_selection ------------------------------------------
+
+
+def test_resolve_applies_hook_choice_and_invalidates_cache():
+    register_callback("model_select", lambda **k: "small-model")
+    agent = FakeAgent(base="big-model")
+    chosen = resolve_run_model_selection(agent, "current prompt", [], "s")
+    assert chosen == "small-model"
+    assert agent.get_auto_model_override() == "small-model"
+    assert agent._code_generation_agent is None  # forced rebuild
+
+
+def test_resolve_noop_when_hook_picks_same_model():
+    register_callback("model_select", lambda **k: "big-model")
+    agent = FakeAgent(base="big-model")
+    assert resolve_run_model_selection(agent, "current prompt", [], "s") is None
+    assert agent.get_auto_model_override() is None
+    assert agent._code_generation_agent is not None  # no rebuild
+
+
+def test_explicit_runtime_override_beats_hook():
+    register_callback("model_select", lambda **k: "small-model")
+    agent = FakeAgent(base="big-model", runtime="user-picked")
+    assert resolve_run_model_selection(agent, "current prompt", [], "s") is None
+    assert agent.get_auto_model_override() is None  # hook never consulted
+
+
+def test_prior_auto_choice_is_reset_each_run():
+    # No hook registered this run; a stale auto choice must be cleared.
+    agent = FakeAgent(base="big-model")
+    agent.set_auto_model_override("stale-small")
+    assert resolve_run_model_selection(agent, "current prompt", [], "s") is None
+    assert agent.get_auto_model_override() is None
+    assert agent._code_generation_agent is None  # invalidated on reset
+
+
+def test_broken_selector_never_raises():
+    def boom(**k):
+        raise RuntimeError("selector exploded")
+
+    register_callback("model_select", boom)
+    agent = FakeAgent(base="big-model")
+    # Must swallow the error and leave the run on its normal model.
+    assert resolve_run_model_selection(agent, "current prompt", [], "s") is None
+    assert agent.get_auto_model_override() is None