Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit c05160e

Browse files
thygrrrthygrrr
committed
feat: add MiMo V2 Pro reasoning support via OpenRouter
Adds reasoning_details round-trip for OpenRouter's MiMo V2 Pro in the same style as the existing reasoning_content path for DeepSeek-R1/Kimi. Request-side reasoning is passed via extra_body so litellm.drop_params can't silently strip it, and _apply_model_overrides now consults the active gateway spec so registry overrides fire for gateway-routed models (previously a latent dead-code bug for any override attached to OpenRouter or AiHubMix).
1 parent 833341a commit c05160e

8 files changed

Lines changed: 291 additions & 5 deletions

File tree

‎pocketfox/agent/context.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -766,6 +766,7 @@ def add_assistant_message(
766766
content: str | None,
767767
tool_calls: list[dict[str, Any]] | None = None,
768768
reasoning_content: str | None = None,
769+
reasoning_details: list[dict[str, Any]] | None = None,
769770
) -> list[dict[str, Any]]:
770771
"""
771772
Add an assistant message to the message list.
@@ -775,6 +776,8 @@ def add_assistant_message(
775776
content: Message content.
776777
tool_calls: Optional tool calls.
777778
reasoning_content: Thinking output (Kimi, DeepSeek-R1, etc.).
779+
reasoning_details: OpenRouter reasoning segments (MiMo etc.) — must
780+
round-trip verbatim in subsequent turns.
778781
779782
Returns:
780783
Updated message list.
@@ -788,5 +791,8 @@ def add_assistant_message(
788791
if reasoning_content:
789792
msg["reasoning_content"] = reasoning_content
790793

794+
if reasoning_details:
795+
msg["reasoning_details"] = reasoning_details
796+
791797
messages.append(msg)
792798
return messages

‎pocketfox/agent/loop.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -586,6 +586,7 @@ async def _run_llm_loop(
586586
response.content,
587587
tool_call_dicts,
588588
reasoning_content=response.reasoning_content,
589+
reasoning_details=response.reasoning_details,
589590
)
590591

591592
# Collect any image blocks returned by multimodal tools (e.g.

‎pocketfox/cli/tty.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -211,6 +211,7 @@ async def process(self, content: str, max_iterations: int = 50) -> str:
211211
response.content,
212212
tool_call_dicts,
213213
reasoning_content=response.reasoning_content,
214+
reasoning_details=response.reasoning_details,
214215
)
215216

216217
# Execute tools. Collect any image blocks returned by

‎pocketfox/providers/base.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ class LLMResponse:
2323
finish_reason: str = "stop"
2424
usage: dict[str, int] = field(default_factory=dict)
2525
reasoning_content: str | None = None # Kimi, DeepSeek-R1 etc.
26+
reasoning_details: list[dict[str, Any]] | None = None # OpenRouter MiMo etc. — round-trip verbatim
2627

2728
@property
2829
def has_tool_calls(self) -> bool:

‎pocketfox/providers/litellm_provider.py‎

Lines changed: 36 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,15 @@
1111
from pocketfox.providers.registry import find_by_model, find_gateway
1212

1313

14+
def _merge_overrides(kwargs: dict[str, Any], overrides: dict[str, Any]) -> None:
15+
"""Merge registry overrides into kwargs, deep-merging extra_body."""
16+
for key, val in overrides.items():
17+
if key == "extra_body" and isinstance(kwargs.get("extra_body"), dict) and isinstance(val, dict):
18+
kwargs["extra_body"] = {**kwargs["extra_body"], **val}
19+
else:
20+
kwargs[key] = val
21+
22+
1423
class LiteLLMProvider(LLMProvider):
1524
"""
1625
LLM provider using LiteLLM for multi-provider support.
@@ -90,13 +99,19 @@ def _resolve_model(self, model: str) -> str:
9099
return model
91100

92101
def _apply_model_overrides(self, model: str, kwargs: dict[str, Any]) -> None:
93-
"""Apply model-specific parameter overrides from the registry."""
102+
"""Apply model-specific parameter overrides from the registry.
103+
104+
Consults both the matched standard provider (if any) and the active
105+
gateway spec (if any) — gateways route arbitrary models, so their
106+
per-model overrides must also fire (e.g. OpenRouter → xiaomi/mimo).
107+
"""
94108
model_lower = model.lower()
95-
spec = find_by_model(model)
96-
if spec:
109+
for spec in (find_by_model(model), self._gateway):
110+
if not spec:
111+
continue
97112
for pattern, overrides in spec.model_overrides:
98113
if pattern in model_lower:
99-
kwargs.update(overrides)
114+
_merge_overrides(kwargs, overrides)
100115
return
101116

102117
async def chat(
@@ -192,13 +207,30 @@ def _parse_response(self, response: Any) -> LLMResponse:
192207
usage["cache_read_input_tokens"] = response.usage.cache_read_input_tokens
193208

194209
reasoning_content = getattr(message, "reasoning_content", None)
210+
reasoning_details = getattr(message, "reasoning_details", None)
211+
212+
# LiteLLM 1.83.1 may not know about OpenRouter's reasoning_details field
213+
# and stash it in pydantic's model_extra instead of exposing it directly.
214+
if reasoning_details is None:
215+
extra = getattr(message, "model_extra", None) or {}
216+
reasoning_details = extra.get("reasoning_details")
217+
218+
# Synthesize a human-readable reasoning_content from details when the
219+
# server populates details-only (MiMo via OpenRouter). Keeps downstream
220+
# display/logging code working without caring about the shape.
221+
if not reasoning_content and reasoning_details:
222+
synthesized = "\n".join(
223+
d.get("text", "") for d in reasoning_details if isinstance(d, dict)
224+
).strip()
225+
reasoning_content = synthesized or None
195226

196227
return LLMResponse(
197228
content=message.content,
198229
tool_calls=tool_calls,
199230
finish_reason=choice.finish_reason or "stop",
200231
usage=usage,
201232
reasoning_content=reasoning_content,
233+
reasoning_details=reasoning_details,
202234
)
203235

204236
def get_default_model(self) -> str:

‎pocketfox/providers/registry.py‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -78,7 +78,11 @@ def label(self) -> str:
7878
detect_by_base_keyword="openrouter",
7979
default_api_base="https://openrouter.ai/api/v1",
8080
strip_model_prefix=False,
81-
model_overrides=(),
81+
# MiMo V2 Pro: pass reasoning via extra_body so LiteLLM's drop_params
82+
# can't silently strip it before the request leaves the library.
83+
model_overrides=(
84+
("xiaomi/mimo", {"extra_body": {"reasoning": {"enabled": True}}}),
85+
),
8286
),
8387
# AiHubMix: global gateway, OpenAI-compatible interface.
8488
# strip_model_prefix=True: it doesn't understand "anthropic/claude-3",

‎scripts/test_mimo.py‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
"""Manual smoke test for MiMo V2 Pro via OpenRouter.
2+
3+
Requires OPENROUTER_API_KEY in env. Not part of the pytest suite — it hits the
4+
real API and will spend tokens.
5+
6+
Usage:
7+
python scripts/test_mimo.py
8+
9+
Verifies end-to-end:
10+
1. `extra_body={"reasoning": {"enabled": True}}` reaches OpenRouter.
11+
2. `reasoning_details` is parsed off the response.
12+
3. A follow-up turn with the full `reasoning_details` round-tripped in
13+
conversation history is accepted by OpenRouter without complaint.
14+
"""
15+
16+
import asyncio
17+
import os
18+
import sys
19+
20+
from pocketfox.agent.context import ContextBuilder # noqa: F401 (unused direct, but mirrors prod code path)
21+
from pocketfox.providers.litellm_provider import LiteLLMProvider
22+
23+
24+
def _require_key() -> str:
25+
key = os.environ.get("OPENROUTER_API_KEY")
26+
if not key:
27+
print("ERROR: set OPENROUTER_API_KEY before running this script", file=sys.stderr)
28+
sys.exit(1)
29+
return key
30+
31+
32+
async def main() -> None:
33+
provider = LiteLLMProvider(
34+
api_key=_require_key(),
35+
default_model="openrouter/xiaomi/mimo-v2-pro",
36+
provider_name="openrouter",
37+
)
38+
39+
# --- Turn 1 -------------------------------------------------------------
40+
messages = [{"role": "user", "content": "Think step by step: what's 17 * 23?"}]
41+
print("[turn 1] calling MiMo V2 Pro...")
42+
response1 = await provider.chat(messages=messages)
43+
44+
print(f"[turn 1] content: {response1.content!r}")
45+
print(f"[turn 1] reasoning_content (first 500 chars): {(response1.reasoning_content or '')[:500]!r}")
46+
if response1.reasoning_details:
47+
print(f"[turn 1] reasoning_details[0]: {response1.reasoning_details[0]!r}")
48+
print(f"[turn 1] reasoning_details count: {len(response1.reasoning_details)}")
49+
else:
50+
print("[turn 1] reasoning_details: None <-- likely a problem")
51+
52+
if not response1.reasoning_details and not response1.reasoning_content:
53+
print(
54+
"[turn 1] WARNING: no reasoning returned. The extra_body path may not "
55+
"have reached OpenRouter. Try setting `litellm.set_verbose = True` "
56+
"in this script and re-running to see the outgoing HTTP body.",
57+
file=sys.stderr,
58+
)
59+
60+
# --- Turn 2 -------------------------------------------------------------
61+
assistant_msg: dict = {"role": "assistant", "content": response1.content or ""}
62+
if response1.reasoning_content:
63+
assistant_msg["reasoning_content"] = response1.reasoning_content
64+
if response1.reasoning_details:
65+
assistant_msg["reasoning_details"] = response1.reasoning_details
66+
67+
messages.append(assistant_msg)
68+
messages.append({"role": "user", "content": "Now double it."})
69+
70+
print("\n[turn 2] calling MiMo V2 Pro with round-tripped reasoning_details...")
71+
response2 = await provider.chat(messages=messages)
72+
73+
print(f"[turn 2] content: {response2.content!r}")
74+
print(f"[turn 2] finish_reason: {response2.finish_reason}")
75+
if response2.finish_reason == "error":
76+
print("[turn 2] FAILED — OpenRouter rejected the history", file=sys.stderr)
77+
sys.exit(2)
78+
79+
print("\nOK")
80+
81+
82+
if __name__ == "__main__":
83+
asyncio.run(main())

‎tests/test_litellm_reasoning.py‎

Lines changed: 158 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,158 @@
1+
"""Tests for reasoning_content / reasoning_details parsing and gateway override routing."""
2+
3+
from types import SimpleNamespace
4+
from typing import Any
5+
6+
from pocketfox.providers.litellm_provider import LiteLLMProvider, _merge_overrides
7+
8+
9+
def _fake_response(
10+
*,
11+
content: str = "answer",
12+
reasoning_content: Any = None,
13+
reasoning_details: Any = None,
14+
model_extra: dict[str, Any] | None = None,
15+
) -> SimpleNamespace:
16+
"""Build an object shaped enough like a LiteLLM ModelResponse for _parse_response."""
17+
message = SimpleNamespace(
18+
content=content,
19+
tool_calls=None,
20+
reasoning_content=reasoning_content,
21+
reasoning_details=reasoning_details,
22+
model_extra=model_extra,
23+
)
24+
choice = SimpleNamespace(message=message, finish_reason="stop")
25+
return SimpleNamespace(choices=[choice], usage=None)
26+
27+
28+
def _provider() -> LiteLLMProvider:
29+
# provider_name="openrouter" routes through the OpenRouter gateway spec
30+
# without requiring a real api_key / network access.
31+
return LiteLLMProvider(
32+
api_key="sk-or-test",
33+
default_model="openrouter/xiaomi/mimo-v2-pro",
34+
provider_name="openrouter",
35+
)
36+
37+
38+
# ---------------------------------------------------------------------------
39+
# _parse_response — reasoning extraction
40+
# ---------------------------------------------------------------------------
41+
42+
43+
def test_parse_reasoning_details_populated_synthesizes_content():
44+
provider = _provider()
45+
details = [
46+
{"type": "reasoning.text", "text": "first step"},
47+
{"type": "reasoning.text", "text": "second step"},
48+
]
49+
response = _fake_response(
50+
content="final",
51+
reasoning_content=None,
52+
reasoning_details=details,
53+
)
54+
55+
parsed = provider._parse_response(response)
56+
57+
assert parsed.content == "final"
58+
assert parsed.reasoning_details == details
59+
assert parsed.reasoning_content == "first step\nsecond step"
60+
61+
62+
def test_parse_reasoning_content_only_backwards_compatible():
63+
provider = _provider()
64+
response = _fake_response(
65+
content="final",
66+
reasoning_content="thinking...",
67+
reasoning_details=None,
68+
)
69+
70+
parsed = provider._parse_response(response)
71+
72+
assert parsed.reasoning_content == "thinking..."
73+
assert parsed.reasoning_details is None
74+
75+
76+
def test_parse_reasoning_details_via_model_extra_fallback():
77+
provider = _provider()
78+
details = [{"type": "reasoning.text", "text": "hidden in extras"}]
79+
response = _fake_response(
80+
content="final",
81+
reasoning_content=None,
82+
reasoning_details=None,
83+
model_extra={"reasoning_details": details},
84+
)
85+
86+
parsed = provider._parse_response(response)
87+
88+
assert parsed.reasoning_details == details
89+
assert parsed.reasoning_content == "hidden in extras"
90+
91+
92+
def test_parse_no_reasoning_at_all():
93+
provider = _provider()
94+
response = _fake_response(content="final")
95+
96+
parsed = provider._parse_response(response)
97+
98+
assert parsed.content == "final"
99+
assert parsed.reasoning_content is None
100+
assert parsed.reasoning_details is None
101+
102+
103+
# ---------------------------------------------------------------------------
104+
# _apply_model_overrides — gateway routing
105+
# ---------------------------------------------------------------------------
106+
107+
108+
def test_gateway_override_fires_for_openrouter_mimo():
109+
provider = _provider()
110+
kwargs: dict[str, Any] = {"model": "openrouter/xiaomi/mimo-v2-pro", "messages": []}
111+
112+
provider._apply_model_overrides("openrouter/xiaomi/mimo-v2-pro", kwargs)
113+
114+
assert kwargs["extra_body"] == {"reasoning": {"enabled": True}}
115+
116+
117+
def test_gateway_override_does_not_fire_for_unrelated_openrouter_model():
118+
provider = _provider()
119+
kwargs: dict[str, Any] = {"model": "openrouter/anthropic/claude-3", "messages": []}
120+
121+
provider._apply_model_overrides("openrouter/anthropic/claude-3", kwargs)
122+
123+
assert "extra_body" not in kwargs
124+
125+
126+
# ---------------------------------------------------------------------------
127+
# _merge_overrides — extra_body deep merge
128+
# ---------------------------------------------------------------------------
129+
130+
131+
def test_merge_overrides_deep_merges_extra_body():
132+
kwargs: dict[str, Any] = {"extra_body": {"existing": "value"}}
133+
overrides = {"extra_body": {"reasoning": {"enabled": True}}}
134+
135+
_merge_overrides(kwargs, overrides)
136+
137+
assert kwargs["extra_body"] == {
138+
"existing": "value",
139+
"reasoning": {"enabled": True},
140+
}
141+
142+
143+
def test_merge_overrides_sets_extra_body_when_absent():
144+
kwargs: dict[str, Any] = {}
145+
overrides = {"extra_body": {"reasoning": {"enabled": True}}}
146+
147+
_merge_overrides(kwargs, overrides)
148+
149+
assert kwargs["extra_body"] == {"reasoning": {"enabled": True}}
150+
151+
152+
def test_merge_overrides_non_extra_body_keys_overwrite():
153+
kwargs: dict[str, Any] = {"temperature": 0.7}
154+
overrides = {"temperature": 1.0}
155+
156+
_merge_overrides(kwargs, overrides)
157+
158+
assert kwargs["temperature"] == 1.0

0 commit comments

Comments
 (0)