fix(moa): project slot reasoning through provider profiles

This commit is contained in:
Justin Schille 2026-07-10 12:53:16 -06:00 committed by Teknium
parent 3dca75b45c
commit 5646dbdd5b
2 changed files with 193 additions and 10 deletions

View file

@ -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)

View file

@ -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