diff --git a/python/omp-rpc/src/omp_rpc/__init__.py b/python/omp-rpc/src/omp_rpc/__init__.py index 3576f18bc..f596d2161 100644 --- a/python/omp-rpc/src/omp_rpc/__init__.py +++ b/python/omp-rpc/src/omp_rpc/__init__.py @@ -50,6 +50,7 @@ from .protocol import ( CompactionSummaryMessage, CustomMessage, DeveloperMessage, + Effort, ExtensionError, ExtensionUiRequest, FileMentionMessage, @@ -114,6 +115,7 @@ __all__ = [ "CompactionSummaryMessage", "CustomMessage", "DeveloperMessage", + "Effort", "ExtensionError", "ExtensionErrorListener", "ExtensionUiRequest", diff --git a/python/omp-rpc/src/omp_rpc/protocol.py b/python/omp-rpc/src/omp_rpc/protocol.py index f29182e8c..a761e4d34 100644 --- a/python/omp-rpc/src/omp_rpc/protocol.py +++ b/python/omp-rpc/src/omp_rpc/protocol.py @@ -11,6 +11,7 @@ JsonValue: TypeAlias = JsonPrimitive | list["JsonValue"] | dict[str, "JsonValue" JsonObject: TypeAlias = dict[str, JsonValue] Attribution: TypeAlias = Literal["user", "agent"] +Effort: TypeAlias = Literal["minimal", "low", "medium", "high", "xhigh"] ThinkingLevel: TypeAlias = Literal["off", "minimal", "low", "medium", "high", "xhigh"] StreamingBehavior: TypeAlias = Literal["steer", "followUp"] SteeringMode: TypeAlias = Literal["all", "one-at-a-time"] @@ -48,9 +49,10 @@ INTERACTIVE_EXTENSION_UI_METHODS: Final[frozenset[InteractiveExtensionUiMethod]] VALUE_EXTENSION_UI_METHODS: Final[frozenset[ValueExtensionUiMethod]] = frozenset( {"select", "input", "editor"} ) -_THINKING_LEVEL_VALUES: Final[frozenset[str]] = frozenset( - {"off", "minimal", "low", "medium", "high", "xhigh"} +_EFFORT_VALUES: Final[frozenset[str]] = frozenset( + {"minimal", "low", "medium", "high", "xhigh"} ) +_THINKING_LEVEL_VALUES: Final[frozenset[str]] = _EFFORT_VALUES | frozenset({"off"}) _STEERING_MODE_VALUES: Final[frozenset[str]] = frozenset({"all", "one-at-a-time"}) _INTERRUPT_MODE_VALUES: Final[frozenset[str]] = frozenset({"immediate", "wait"}) _STOP_REASON_VALUES: Final[frozenset[str]] = frozenset( @@ -696,9 +698,14 @@ class ModelCost: @dataclass(slots=True, frozen=True) class ThinkingConfig: - min_level: ThinkingLevel - max_level: ThinkingLevel mode: str + efforts: tuple[Effort, ...] + default_level: Effort | None = None + effort_map: dict[str, str] | None = None + supports_display: bool | None = None + effort_routing: dict[str, str] | None = None + suppress_when_off: bool | None = None + requires_effort: bool | None = None @dataclass(slots=True, frozen=True) @@ -1128,6 +1135,45 @@ def assistant_text_with_thinking(message: AgentMessage) -> str | None: return assistant_text(message, include_thinking=True) +def _parse_thinking_config(payload: object) -> ThinkingConfig | None: + if not isinstance(payload, dict): + return None + raw_efforts = payload.get("efforts") + if not isinstance(raw_efforts, list): + raise ValueError("model.thinking.efforts must be a list") + efforts: tuple[Effort, ...] = tuple( + cast(Effort, _require_literal(item, _EFFORT_VALUES, field="model.thinking.efforts[]")) + for item in raw_efforts + ) + return ThinkingConfig( + mode=_require_str(cast(JsonObject, payload), "mode"), + efforts=efforts, + default_level=cast( + Effort | None, + _optional_literal( + payload.get("defaultLevel"), + _EFFORT_VALUES, + field="model.thinking.defaultLevel", + ), + ), + effort_map=cast( + dict[str, str] | None, + _optional_json_object( + payload.get("effortMap"), field="model.thinking.effortMap" + ), + ), + supports_display=_optional_bool(cast(JsonObject, payload), "supportsDisplay"), + effort_routing=cast( + dict[str, str] | None, + _optional_json_object( + payload.get("effortRouting"), field="model.thinking.effortRouting" + ), + ), + suppress_when_off=_optional_bool(cast(JsonObject, payload), "suppressWhenOff"), + requires_effort=_optional_bool(cast(JsonObject, payload), "requiresEffort"), + ) + + def parse_model_info(payload: JsonObject | None) -> ModelInfo | None: if payload is None: return None @@ -1168,29 +1214,7 @@ def parse_model_info(payload: JsonObject | None) -> ModelInfo | None: else None ), priority=int(payload["priority"]) if "priority" in payload else None, - thinking=( - ThinkingConfig( - min_level=cast( - ThinkingLevel, - _require_literal( - thinking_payload.get("minLevel"), - _THINKING_LEVEL_VALUES, - field="model.thinking.minLevel", - ), - ), - max_level=cast( - ThinkingLevel, - _require_literal( - thinking_payload.get("maxLevel"), - _THINKING_LEVEL_VALUES, - field="model.thinking.maxLevel", - ), - ), - mode=_require_str(cast(JsonObject, thinking_payload), "mode"), - ) - if isinstance(thinking_payload, dict) - else None - ), + thinking=_parse_thinking_config(thinking_payload), compat=_optional_json_object(compat_payload, field="model.compat"), ) diff --git a/python/omp-rpc/tests/test_protocol.py b/python/omp-rpc/tests/test_protocol.py index ff641f19b..5d7bb158f 100644 --- a/python/omp-rpc/tests/test_protocol.py +++ b/python/omp-rpc/tests/test_protocol.py @@ -35,9 +35,11 @@ class ProtocolParsingTests(unittest.TestCase): "contextWindow": 200000, "maxTokens": 8192, "thinking": { - "minLevel": "minimal", - "maxLevel": "high", "mode": "effort", + "efforts": ["minimal", "low", "medium", "high"], + "defaultLevel": "medium", + "effortMap": {"high": "xhigh"}, + "supportsDisplay": True, }, }, "thinkingLevel": "medium", @@ -85,6 +87,14 @@ class ProtocolParsingTests(unittest.TestCase): # Legacy bare-string systemPrompt is accepted and wrapped to a tuple. self.assertEqual(state.system_prompt, ("You are useful.",)) self.assertEqual(state.dump_tools[0].name, "read") + assert state.model is not None and state.model.thinking is not None + self.assertEqual( + state.model.thinking.efforts, ("minimal", "low", "medium", "high") + ) + self.assertEqual(state.model.thinking.mode, "effort") + self.assertEqual(state.model.thinking.default_level, "medium") + self.assertEqual(state.model.thinking.effort_map, {"high": "xhigh"}) + self.assertTrue(state.model.thinking.supports_display) def test_parse_agent_end_notification(self) -> None: notification = parse_notification( @@ -184,6 +194,26 @@ class ProtocolParsingTests(unittest.TestCase): } ) + def test_parse_model_info_rejects_unknown_effort(self) -> None: + with self.assertRaises(ValueError): + parse_session_state( + { + "sessionId": "session-123", + "steeringMode": "one-at-a-time", + "followUpMode": "one-at-a-time", + "interruptMode": "immediate", + "model": { + "id": "m", + "name": "M", + "api": "anthropic-messages", + "provider": "anthropic", + "baseUrl": "https://api.anthropic.com", + "reasoning": True, + "thinking": {"mode": "effort", "efforts": ["extreme"]}, + }, + } + ) + def test_parse_session_state_accepts_system_prompt_array(self) -> None: state = parse_session_state( {