feat: thinking level validation and clamping

cubepi/providers/__init__.py

@@ -26,6 +26,12 @@
     faux_thinking,
     faux_tool_call,
 )
+from cubepi.providers.models import (
+    THINKING_LEVELS,
+    clamp_thinking_level,
+    get_supported_thinking_levels,
+    models_are_equal,
+)
 
 
 # Lazy imports for optional providers
@@ -62,8 +68,12 @@ def get_openai_provider():
     "Usage",
     "UserMessage",
     "adjust_max_tokens_for_thinking",
+    "clamp_thinking_level",
     "faux_assistant_message",
     "faux_text",
     "faux_thinking",
     "faux_tool_call",
+    "get_supported_thinking_levels",
+    "models_are_equal",
+    "THINKING_LEVELS",
 ]

cubepi/providers/anthropic.py

@@ -22,6 +22,7 @@
     UserMessage,
     adjust_max_tokens_for_thinking,
 )
+from cubepi.providers.models import clamp_thinking_level
 
 CacheRetention = Literal["short", "long", "none"]
 
@@ -50,6 +51,7 @@ async def stream(
         signal: asyncio.Event | None = None,
     ) -> MessageStream:
         ms = MessageStream()
+        thinking = clamp_thinking_level(model, thinking)
 
         cache_control = self._get_cache_control()
         api_messages = [self._convert_message(m) for m in messages]

cubepi/providers/base.py

@@ -73,6 +73,7 @@ class Model(BaseModel):
     context_window: int = 200_000
     max_tokens: int = 8192
     cost: ModelCost | None = None
+    thinking_level_map: dict[str, str | None] | None = None
 
 
 class TextContent(BaseModel):

cubepi/providers/models.py

@@ -0,0 +1,97 @@
+"""Thinking-level validation and model comparison utilities.
+
+Mirrors pi-agent-core's ``getSupportedThinkingLevels``, ``clampThinkingLevel``,
+and ``modelsAreEqual`` functions.
+"""
+
+from __future__ import annotations
+
+from cubepi.providers.base import Model, ThinkingLevel
+
+# Ordered list of all thinking levels from lowest to highest.
+THINKING_LEVELS: list[ThinkingLevel] = [
+    "off",
+    "minimal",
+    "low",
+    "medium",
+    "high",
+    "xhigh",
+]
+
+
+def get_supported_thinking_levels(model: Model) -> list[ThinkingLevel]:
+    """Return the thinking levels supported by *model*.
+
+    * Non-reasoning models only support ``["off"]``.
+    * For reasoning models, levels are filtered through the model's
+      ``thinking_level_map``.  A level mapped to ``None`` is unsupported.
+      ``"xhigh"`` is only included when it has an explicit (non-None) mapping.
+      All other levels are included by default when the map omits them.
+    """
+    if not model.reasoning:
+        return ["off"]
+
+    tlm = model.thinking_level_map
+
+    def _is_supported(level: ThinkingLevel) -> bool:
+        if tlm is not None:
+            mapped = tlm.get(level)
+            if mapped is None and level in tlm:
+                # Explicitly mapped to None -> unsupported
+                return False
+            # "xhigh" requires an explicit mapping to be available
+            if level == "xhigh":
+                return level in tlm and tlm[level] is not None
+        else:
+            # No map at all: xhigh is excluded by default
+            if level == "xhigh":
+                return False
+        return True
+
+    return [lvl for lvl in THINKING_LEVELS if _is_supported(lvl)]
+
+
+def clamp_thinking_level(model: Model, level: ThinkingLevel) -> ThinkingLevel:
+    """Clamp *level* to the nearest supported level for *model*.
+
+    If *level* is already supported, return it unchanged.  Otherwise search
+    upward first (higher intensity), then downward, through the ordered level
+    list to find the closest available level.
+    """
+    available = get_supported_thinking_levels(model)
+
+    if level in available:
+        return level
+
+    # Unknown level -> fall back to first available
+    if level not in THINKING_LEVELS:
+        return available[0] if available else "off"
+
+    requested_idx = THINKING_LEVELS.index(level)
+
+    # Search downward first (prefer cheaper/lower intensity)
+    for i in range(requested_idx - 1, -1, -1):
+        candidate = THINKING_LEVELS[i]
+        if candidate in available:
+            return candidate
+
+    # Then upward
+    for i in range(requested_idx + 1, len(THINKING_LEVELS)):
+        candidate = THINKING_LEVELS[i]
+        if candidate in available:
+            return candidate
+
+    return available[0] if available else "off"
+
+
+def models_are_equal(a: Model | None, b: Model | None) -> bool:
+    """Return ``True`` if *a* and *b* refer to the same model.
+
+    Comparison is by ``id`` and ``provider``.  Returns ``False`` when either
+    argument is ``None``.
+    """
+    if a is None and b is None:
+        return True
+    if a is None or b is None:
+        return False
+    return a.id == b.id and a.provider == b.provider

cubepi/providers/openai.py

@@ -22,6 +22,7 @@
     Usage,
     UserMessage,
 )
+from cubepi.providers.models import clamp_thinking_level
 
 
 class OpenAIProvider:
@@ -49,6 +50,7 @@ async def stream(
         signal: asyncio.Event | None = None,
     ) -> MessageStream:
         ms = MessageStream()
+        thinking = clamp_thinking_level(model, thinking)
 
         api_messages: list[dict[str, Any]] = []
         if system_prompt:

tests/providers/test_models.py

@@ -0,0 +1,221 @@
+"""Tests for thinking-level validation, clamping, and model comparison."""
+
+from cubepi.providers.base import Model
+from cubepi.providers.models import (
+    THINKING_LEVELS,
+    clamp_thinking_level,
+    get_supported_thinking_levels,
+    models_are_equal,
+)
+
+
+# ---------------------------------------------------------------------------
+# Fixtures / helpers
+# ---------------------------------------------------------------------------
+
+
+def _model(
+    *,
+    reasoning: bool = False,
+    thinking_level_map: dict[str, str | None] | None = None,
+    provider: str = "anthropic",
+    model_id: str = "claude-sonnet-4-20250514",
+) -> Model:
+    return Model(
+        id=model_id,
+        provider=provider,
+        reasoning=reasoning,
+        thinking_level_map=thinking_level_map,
+    )
+
+
+# ---------------------------------------------------------------------------
+# get_supported_thinking_levels
+# ---------------------------------------------------------------------------
+
+
+class TestGetSupportedThinkingLevels:
+    def test_non_reasoning_model_only_off(self):
+        model = _model(reasoning=False)
+        assert get_supported_thinking_levels(model) == ["off"]
+
+    def test_reasoning_model_no_map_defaults(self):
+        """Without a map, all levels except xhigh are available."""
+        model = _model(reasoning=True)
+        assert get_supported_thinking_levels(model) == [
+            "off",
+            "minimal",
+            "low",
+            "medium",
+            "high",
+        ]
+
+    def test_reasoning_model_with_xhigh_in_map(self):
+        model = _model(
+            reasoning=True,
+            thinking_level_map={"xhigh": "max_tokens_16k"},
+        )
+        levels = get_supported_thinking_levels(model)
+        assert "xhigh" in levels
+        # All standard levels should still be present
+        for lvl in ("off", "minimal", "low", "medium", "high"):
+            assert lvl in levels
+
+    def test_xhigh_mapped_to_none_is_excluded(self):
+        model = _model(
+            reasoning=True,
+            thinking_level_map={"xhigh": None},
+        )
+        assert "xhigh" not in get_supported_thinking_levels(model)
+
+    def test_level_mapped_to_none_is_excluded(self):
+        model = _model(
+            reasoning=True,
+            thinking_level_map={"minimal": None, "low": None},
+        )
+        levels = get_supported_thinking_levels(model)
+        assert "minimal" not in levels
+        assert "low" not in levels
+        # Others remain
+        assert "off" in levels
+        assert "medium" in levels
+        assert "high" in levels
+
+    def test_empty_map_same_as_no_map(self):
+        model = _model(reasoning=True, thinking_level_map={})
+        levels = get_supported_thinking_levels(model)
+        assert levels == ["off", "minimal", "low", "medium", "high"]
+
+    def test_all_levels_disabled_except_off(self):
+        model = _model(
+            reasoning=True,
+            thinking_level_map={
+                "minimal": None,
+                "low": None,
+                "medium": None,
+                "high": None,
+            },
+        )
+        assert get_supported_thinking_levels(model) == ["off"]
+
+
+# ---------------------------------------------------------------------------
+# clamp_thinking_level
+# ---------------------------------------------------------------------------
+
+
+class TestClampThinkingLevel:
+    def test_supported_level_passes_through(self):
+        model = _model(reasoning=True)
+        assert clamp_thinking_level(model, "medium") == "medium"
+
+    def test_off_always_passes_for_non_reasoning(self):
+        model = _model(reasoning=False)
+        assert clamp_thinking_level(model, "off") == "off"
+
+    def test_non_reasoning_clamps_to_off(self):
+        model = _model(reasoning=False)
+        assert clamp_thinking_level(model, "high") == "off"
+
+    def test_clamp_xhigh_without_support_goes_down(self):
+        """xhigh not in map -> clamp down to high."""
+        model = _model(reasoning=True)
+        assert clamp_thinking_level(model, "xhigh") == "high"
+
+    def test_clamp_down_to_nearest(self):
+        """When a level is disabled, search downward first (prefer cheaper)."""
+        model = _model(
+            reasoning=True,
+            thinking_level_map={"low": None},
+        )
+        # "low" disabled -> next down is "minimal"
+        assert clamp_thinking_level(model, "low") == "minimal"
+
+    def test_clamp_down_when_nothing_above(self):
+        """When all higher levels are disabled, clamp down."""
+        model = _model(
+            reasoning=True,
+            thinking_level_map={"high": None},
+        )
+        # "high" disabled, xhigh not in map -> search down -> "medium"
+        assert clamp_thinking_level(model, "high") == "medium"
+
+    def test_clamp_unknown_level_returns_first_available(self):
+        model = _model(reasoning=True)
+        # An invalid level should fall back to the first available
+        result = clamp_thinking_level(model, "ultra")  # type: ignore[arg-type]
+        assert result == "off"
+
+    def test_clamp_with_only_off_available(self):
+        model = _model(
+            reasoning=True,
+            thinking_level_map={
+                "minimal": None,
+                "low": None,
+                "medium": None,
+                "high": None,
+            },
+        )
+        assert clamp_thinking_level(model, "medium") == "off"
+
+    def test_clamp_preserves_exact_match_with_map(self):
+        model = _model(
+            reasoning=True,
+            thinking_level_map={"medium": "budget_4096", "xhigh": "budget_max"},
+        )
+        assert clamp_thinking_level(model, "medium") == "medium"
+        assert clamp_thinking_level(model, "xhigh") == "xhigh"
+
+
+# ---------------------------------------------------------------------------
+# models_are_equal
+# ---------------------------------------------------------------------------
+
+
+class TestModelsAreEqual:
+    def test_equal_models(self):
+        a = _model(model_id="claude-sonnet-4-20250514", provider="anthropic")
+        b = _model(model_id="claude-sonnet-4-20250514", provider="anthropic")
+        assert models_are_equal(a, b) is True
+
+    def test_different_id(self):
+        a = _model(model_id="claude-sonnet-4-20250514", provider="anthropic")
+        b = _model(model_id="claude-opus-4-20250514", provider="anthropic")
+        assert models_are_equal(a, b) is False
+
+    def test_different_provider(self):
+        a = _model(model_id="claude-sonnet-4-20250514", provider="anthropic")
+        b = _model(model_id="claude-sonnet-4-20250514", provider="bedrock")
+        assert models_are_equal(a, b) is False
+
+    def test_none_a(self):
+        b = _model()
+        assert models_are_equal(None, b) is False
+
+    def test_none_b(self):
+        a = _model()
+        assert models_are_equal(a, None) is False
+
+    def test_both_none(self):
+        assert models_are_equal(None, None) is True
+
+
+# ---------------------------------------------------------------------------
+# THINKING_LEVELS ordering
+# ---------------------------------------------------------------------------
+
+
+class TestThinkingLevelsOrdering:
+    def test_correct_order(self):
+        assert THINKING_LEVELS == [
+            "off",
+            "minimal",
+            "low",
+            "medium",
+            "high",
+            "xhigh",
+        ]
+
+    def test_all_levels_present(self):
+        expected = {"off", "minimal", "low", "medium", "high", "xhigh"}
+        assert set(THINKING_LEVELS) == expected