From 5646dbdd5bcc131deb23bd0208defca50e3d7161 Mon Sep 17 00:00:00 2001 From: Justin Schille <217401759+justinschille@users.noreply.github.com> Date: Fri, 10 Jul 2026 12:53:16 -0600 Subject: [PATCH] fix(moa): project slot reasoning through provider profiles --- agent/auxiliary_client.py | 121 ++++++++++++++++++++++++--- tests/agent/test_auxiliary_client.py | 82 ++++++++++++++++++ 2 files changed, 193 insertions(+), 10 deletions(-) diff --git a/agent/auxiliary_client.py b/agent/auxiliary_client.py index 0a02073f2..92e65812f 100644 --- a/agent/auxiliary_client.py +++ b/agent/auxiliary_client.py @@ -3443,6 +3443,7 @@ def _retry_same_provider_sync( tools: Optional[list], effective_timeout: float, effective_extra_body: dict, + reasoning_config: Optional[dict], ) -> Any: if task == "vision": _, retry_client, retry_model = resolve_vision_provider_client( @@ -3476,6 +3477,7 @@ def _retry_same_provider_sync( tools=tools, timeout=effective_timeout, extra_body=effective_extra_body, + reasoning_config=reasoning_config, base_url=retry_base or resolved_base_url, ) if _is_anthropic_compat_endpoint(resolved_provider, retry_base): @@ -3500,6 +3502,7 @@ async def _retry_same_provider_async( tools: Optional[list], effective_timeout: float, effective_extra_body: dict, + reasoning_config: Optional[dict], ) -> Any: if task == "vision": _, retry_client, retry_model = resolve_vision_provider_client( @@ -3533,6 +3536,7 @@ async def _retry_same_provider_async( tools=tools, timeout=effective_timeout, extra_body=effective_extra_body, + reasoning_config=reasoning_config, base_url=retry_base or resolved_base_url, ) if _is_anthropic_compat_endpoint(resolved_provider, retry_base): @@ -3652,6 +3656,7 @@ def _call_fallback_candidate_sync( tools: Optional[list], effective_timeout: float, effective_extra_body: dict, + reasoning_config: Optional[dict], ) -> Optional[Any]: """Call one fallback candidate with stale-credential recovery. @@ -3673,7 +3678,8 @@ def _call_fallback_candidate_sync( fb_label, fb_model, messages, temperature=temperature, max_tokens=max_tokens, tools=tools, timeout=effective_timeout, - extra_body=effective_extra_body, base_url=fb_base) + extra_body=effective_extra_body, reasoning_config=reasoning_config, + base_url=fb_base) try: return _validate_llm_response( fb_client.chat.completions.create(**fb_kwargs), task) @@ -3689,6 +3695,7 @@ def _call_fallback_candidate_sync( temperature=temperature, max_tokens=max_tokens, tools=tools, timeout=effective_timeout, extra_body=effective_extra_body, + reasoning_config=reasoning_config, base_url=str(getattr(retry_client, "base_url", "") or fb_base)) try: return _validate_llm_response( @@ -3721,6 +3728,7 @@ async def _call_fallback_candidate_async( tools: Optional[list], effective_timeout: float, effective_extra_body: dict, + reasoning_config: Optional[dict], ) -> Optional[Any]: """Async mirror of :func:`_call_fallback_candidate_sync`.""" fb_base = str(getattr(fb_client, "base_url", "") or "") @@ -3728,7 +3736,8 @@ async def _call_fallback_candidate_async( fb_label, fb_model, messages, temperature=temperature, max_tokens=max_tokens, tools=tools, timeout=effective_timeout, - extra_body=effective_extra_body, base_url=fb_base) + extra_body=effective_extra_body, reasoning_config=reasoning_config, + base_url=fb_base) try: return _validate_llm_response( await fb_client.chat.completions.create(**fb_kwargs), task) @@ -3745,6 +3754,7 @@ async def _call_fallback_candidate_async( temperature=temperature, max_tokens=max_tokens, tools=tools, timeout=effective_timeout, extra_body=effective_extra_body, + reasoning_config=reasoning_config, base_url=str(getattr(retry_client, "base_url", "") or fb_base)) try: return _validate_llm_response( @@ -6290,6 +6300,32 @@ def _convert_openai_images_to_anthropic(messages: list) -> list: return converted +_PROFILE_REASONING_KEYS = { + "reasoning", + "reasoning_effort", + "thinking", + "thinking_config", + "thinkingconfig", + "thinking_budget", + "thinkingbudget", + "enable_thinking", + "think", + "verbosity", +} + + +def _contains_profile_reasoning_fields(value: Any) -> bool: + """Return whether a profile payload contains a reasoning wire control.""" + if not isinstance(value, dict): + return False + for key, nested in value.items(): + normalized = str(key).strip().lower() + if normalized in _PROFILE_REASONING_KEYS: + return True + if _contains_profile_reasoning_fields(nested): + return True + return False + def _build_call_kwargs( provider: str, @@ -6384,16 +6420,71 @@ def _build_call_kwargs( _deduped.append(_t) kwargs["tools"] = _deduped - # Provider-specific extra_body + # Build provider-aware reasoning kwargs through the same profile hooks used + # by the standard chat-completions transport. Some providers require + # top-level controls (Kimi/custom ``reasoning_effort``), others use nested + # body fields (Gemini ``thinking_config``), and OpenRouter/Nous use + # ``extra_body.reasoning``. Profiles are the source of truth for those wire + # shapes. Providers without a reasoning-aware profile retain the generic + # ``extra_body.reasoning`` fallback used by Codex-compatible adapters. + effective_base = base_url or ( + _current_custom_base_url() if provider == "custom" else "" + ) + profile_body: Dict[str, Any] = {} + profile_reasoning_extra: Dict[str, Any] = {} + profile_top_level: Dict[str, Any] = {} + profile_handles_reasoning = False + try: + from providers import get_provider_profile + from providers.base import ProviderProfile + + profile = get_provider_profile(str(provider or "").strip().lower()) + if profile is not None: + profile_body = profile.build_extra_body( + model=model, + base_url=effective_base, + reasoning_config=reasoning_config, + ) or {} + profile_reasoning_extra, profile_top_level = ( + profile.build_api_kwargs_extras( + reasoning_config=reasoning_config, + supports_reasoning=reasoning_config is not None, + model=model, + base_url=effective_base, + ) + ) + profile_reasoning_extra = profile_reasoning_extra or {} + profile_top_level = profile_top_level or {} + profile_handles_reasoning = ( + type(profile).build_api_kwargs_extras + is not ProviderProfile.build_api_kwargs_extras + or _contains_profile_reasoning_fields(profile_body) + or _contains_profile_reasoning_fields(profile_reasoning_extra) + or _contains_profile_reasoning_fields(profile_top_level) + ) + except Exception as exc: + logger.debug( + "_build_call_kwargs: provider profile projection failed for %s: %s", + provider, + exc, + ) + + kwargs.update(profile_top_level) merged_extra = dict(extra_body or {}) - if reasoning_config and isinstance(reasoning_config, dict): + merged_extra.update(profile_body) + merged_extra.update(profile_reasoning_extra) + if ( + reasoning_config + and isinstance(reasoning_config, dict) + and not profile_handles_reasoning + ): if reasoning_config.get("enabled") is False: merged_extra["reasoning"] = {"enabled": False} else: effort = reasoning_config.get("effort") or "medium" merged_extra["reasoning"] = {"enabled": True, "effort": effort} - if provider == "nous": - merged_extra.setdefault("tags", []).extend(_nous_portal_tags()) + if provider == "nous" and "tags" not in merged_extra: + merged_extra["tags"] = _nous_portal_tags() if merged_extra: kwargs["extra_body"] = merged_extra @@ -6904,6 +6995,7 @@ def call_llm( tools=tools, effective_timeout=effective_timeout, effective_extra_body=effective_extra_body, + reasoning_config=reasoning_config, ) # ── Same-provider credential-pool recovery ───────────────────── @@ -6946,6 +7038,7 @@ def call_llm( tools=tools, effective_timeout=effective_timeout, effective_extra_body=effective_extra_body, + reasoning_config=reasoning_config, ) except Exception as retry2_err: # The rotated key also hit a quota/auth wall. Mark it @@ -7067,7 +7160,8 @@ def call_llm( task=task, messages=messages, temperature=temperature, max_tokens=max_tokens, tools=tools, effective_timeout=effective_timeout, - effective_extra_body=effective_extra_body) + effective_extra_body=effective_extra_body, + reasoning_config=reasoning_config) if fb_resp is not None: return fb_resp # The candidate had a stale/unrefreshable credential and was @@ -7081,7 +7175,8 @@ def call_llm( task=task, messages=messages, temperature=temperature, max_tokens=max_tokens, tools=tools, effective_timeout=effective_timeout, - effective_extra_body=effective_extra_body) + effective_extra_body=effective_extra_body, + reasoning_config=reasoning_config) if fb_resp is not None: return fb_resp # All fallback layers exhausted — emit a single user-visible @@ -7176,6 +7271,7 @@ async def async_call_llm( tools: list = None, timeout: float = None, extra_body: dict = None, + reasoning_config: Optional[dict] = None, ) -> Any: """Centralized asynchronous LLM call. @@ -7255,6 +7351,7 @@ async def async_call_llm( resolved_provider, final_model, messages, temperature=temperature, max_tokens=max_tokens, tools=tools, timeout=effective_timeout, extra_body=effective_extra_body, + reasoning_config=reasoning_config, base_url=_client_base or resolved_base_url) # Convert image blocks for Anthropic-compatible endpoints (e.g. MiniMax) @@ -7451,6 +7548,7 @@ async def async_call_llm( tools=tools, effective_timeout=effective_timeout, effective_extra_body=effective_extra_body, + reasoning_config=reasoning_config, ) # ── Same-provider credential-pool recovery (mirrors sync) ───── @@ -7488,6 +7586,7 @@ async def async_call_llm( tools=tools, effective_timeout=effective_timeout, effective_extra_body=effective_extra_body, + reasoning_config=reasoning_config, ) except Exception as retry2_err: if (_is_payment_error(retry2_err) or _is_auth_error(retry2_err) @@ -7576,7 +7675,8 @@ async def async_call_llm( task=task, messages=messages, temperature=temperature, max_tokens=max_tokens, tools=tools, effective_timeout=effective_timeout, - effective_extra_body=effective_extra_body) + effective_extra_body=effective_extra_body, + reasoning_config=reasoning_config) if fb_resp is not None: return fb_resp # Stale/unrefreshable candidate credential — quarantined; walk @@ -7592,7 +7692,8 @@ async def async_call_llm( task=task, messages=messages, temperature=temperature, max_tokens=max_tokens, tools=tools, effective_timeout=effective_timeout, - effective_extra_body=effective_extra_body) + effective_extra_body=effective_extra_body, + reasoning_config=reasoning_config) if fb_resp is not None: return fb_resp # All fallback layers exhausted — warn before re-raising. (#26882) diff --git a/tests/agent/test_auxiliary_client.py b/tests/agent/test_auxiliary_client.py index 11cad1f1b..216a79fb3 100644 --- a/tests/agent/test_auxiliary_client.py +++ b/tests/agent/test_auxiliary_client.py @@ -3966,6 +3966,88 @@ class TestAnthropicAuxiliaryReasoningTranslation: assert "_reasoning_config" not in openai_wire_kwargs +class TestAuxiliaryProviderProfileReasoning: + """Auxiliary calls must reuse provider-profile reasoning wire shapes.""" + + def test_kimi_reasoning_uses_top_level_effort(self): + kwargs = _build_call_kwargs( + "kimi-coding", + "kimi-k2-turbo-preview", + [{"role": "user", "content": "hi"}], + reasoning_config={"enabled": True, "effort": "medium"}, + base_url="https://api.moonshot.ai/v1", + ) + + assert kwargs["reasoning_effort"] == "medium" + assert "reasoning" not in kwargs.get("extra_body", {}) + assert "thinking" not in kwargs.get("extra_body", {}) + + def test_gemini_reasoning_uses_thinking_config(self): + kwargs = _build_call_kwargs( + "gemini", + "gemini-3.5-flash", + [{"role": "user", "content": "hi"}], + reasoning_config={"enabled": True, "effort": "high"}, + base_url="https://generativelanguage.googleapis.com/v1beta", + ) + + assert kwargs["extra_body"]["thinking_config"] == { + "includeThoughts": True, + "thinkingLevel": "high", + } + assert "reasoning" not in kwargs["extra_body"] + + def test_custom_openai_compatible_reasoning_uses_top_level_effort(self): + kwargs = _build_call_kwargs( + "custom", + "glm-5.2", + [{"role": "user", "content": "hi"}], + reasoning_config={"enabled": True, "effort": "max"}, + base_url="https://example.test/v1", + ) + + assert kwargs["reasoning_effort"] == "max" + assert "reasoning" not in kwargs.get("extra_body", {}) + + @pytest.mark.asyncio + async def test_async_call_llm_preserves_profile_reasoning_kwargs(self): + response = SimpleNamespace( + choices=[SimpleNamespace(message=SimpleNamespace(content="ok"))] + ) + create = AsyncMock(return_value=response) + client = SimpleNamespace( + base_url="https://api.moonshot.ai/v1", + chat=SimpleNamespace( + completions=SimpleNamespace(create=create), + ), + ) + + with patch( + "agent.auxiliary_client._resolve_task_provider_model", + return_value=( + "kimi-coding", + "kimi-k2-turbo-preview", + "https://api.moonshot.ai/v1", + "test-key", + None, + ), + ), patch( + "agent.auxiliary_client._get_cached_client", + return_value=(client, "kimi-k2-turbo-preview"), + ): + result = await async_call_llm( + provider="kimi-coding", + model="kimi-k2-turbo-preview", + messages=[{"role": "user", "content": "hi"}], + reasoning_config={"enabled": True, "effort": "high"}, + ) + + assert result is response + final_kwargs = create.call_args.kwargs + assert final_kwargs["reasoning_effort"] == "high" + assert "reasoning" not in final_kwargs.get("extra_body", {}) + + class TestCodexAdapterReasoningTranslation: """Verify _CodexCompletionsAdapter translates extra_body.reasoning into the Responses API's top-level reasoning + include fields, matching