refactor(rust-bridge): project Messages capabilities without mutable dicts

The capability flags and effort tiers were built as dict comprehensions,
which the type-discipline gate counts as mutable construction, and the
asdict call carried a mutable-ok suppression that suppressed nothing. The
flags are now passed one by one and the effort tiers are a frozen dataclass,
which asdict projects to the same map the native side reads
This commit is contained in:
Yujong Lee 2026-09-24 11:59:59 -07:00
parent 92713aca11
commit 8edd279027
2 changed files with 36 additions and 22 deletions

View file

@ -13,15 +13,17 @@ from litellm.rust_bridge import failures
from litellm.rust_bridge.messages.entrypoints import LiteLLMMessagesRequest
from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse
_EFFORT_TIERS: Final = ("minimal", "low", "medium", "high", "xhigh", "max")
_DROP_PATHS: Final = TypeAdapter(list[object])
_CAPABILITY_FLAGS: Final = (
"supports_reasoning",
"supports_adaptive_thinking",
"thinking_always_on",
"supports_legacy_thinking",
"supports_output_config",
)
@dataclass(frozen=True, slots=True)
class EffortTiers:
minimal: bool
low: bool
medium: bool
high: bool
xhigh: bool
max: bool
@dataclass(frozen=True, slots=True)
@ -33,7 +35,7 @@ class ModelCapabilities:
supports_output_config: bool
supports_sampling_params: bool
supports_speed: bool
effort_tiers: Mapping[str, bool]
effort_tiers: EffortTiers
@dataclass(frozen=True, slots=True)
@ -78,18 +80,29 @@ def model_capabilities(model: str, custom_llm_provider: str | None) -> ModelCapa
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
resolved_model, provider = _resolved_provider(model, custom_llm_provider)
flags: Final = {
flag: AnthropicModelInfo._supports_model_capability(model, flag, provider) # pyright: ignore[reportPrivateUsage] # same probes the Python transform runs; forking them would drift
for flag in _CAPABILITY_FLAGS
}
def supports(flag: str) -> bool:
return AnthropicModelInfo._supports_model_capability(model, flag, provider) # pyright: ignore[reportPrivateUsage] # same probes the Python transform runs; forking them would drift
def tier(level: str) -> bool:
return AnthropicConfig._supports_effort_level(model, level, provider) # pyright: ignore[reportPrivateUsage] # same probe the Python transform runs
return ModelCapabilities(
supports_reasoning=supports("supports_reasoning"),
supports_adaptive_thinking=supports("supports_adaptive_thinking"),
thinking_always_on=supports("thinking_always_on"),
supports_legacy_thinking=supports("supports_legacy_thinking"),
supports_output_config=supports("supports_output_config"),
supports_sampling_params=AnthropicModelInfo._supports_sampling_params(resolved_model), # pyright: ignore[reportPrivateUsage] # same gate the handler applies
supports_speed=AnthropicConfig._model_supports_speed_param(resolved_model, provider), # pyright: ignore[reportPrivateUsage] # same gate the handler applies
effort_tiers={
tier: AnthropicConfig._supports_effort_level(model, tier, provider) # pyright: ignore[reportPrivateUsage] # same probe the Python transform runs
for tier in _EFFORT_TIERS
},
**flags,
effort_tiers=EffortTiers(
minimal=tier("minimal"),
low=tier("low"),
medium=tier("medium"),
high=tier("high"),
xhigh=tier("xhigh"),
max=tier("max"),
),
)
@ -106,7 +119,7 @@ def _additional_drop_params(kwargs: Mapping[str, object]) -> tuple[str, ...]:
def shaping(model: str, custom_llm_provider: str | None, kwargs: Mapping[str, object]) -> dict[str, object]:
return asdict( # mutable-ok: the native side depythonizes a plain dict
return asdict(
MessagesShaping(
capabilities=model_capabilities(model, custom_llm_provider),
drop_params=_drop_params(kwargs),

View file

@ -1,3 +1,4 @@
from dataclasses import astuple
from typing import Final
import pytest
@ -39,8 +40,8 @@ def test_capabilities_come_from_the_model_map_under_the_callers_provider(monkeyp
assert capabilities.supports_output_config
assert not capabilities.supports_legacy_thinking
assert not capabilities.supports_sampling_params
assert capabilities.effort_tiers["xhigh"]
assert not capabilities.effort_tiers["max"]
assert capabilities.effort_tiers.xhigh
assert not capabilities.effort_tiers.max
def test_unmapped_model_keeps_sampling_params_and_no_reasoning_features() -> None:
@ -49,7 +50,7 @@ def test_unmapped_model_keeps_sampling_params_and_no_reasoning_features() -> Non
assert capabilities.supports_sampling_params
assert not capabilities.supports_reasoning
assert not capabilities.supports_adaptive_thinking
assert not any(capabilities.effort_tiers.values())
assert not any(astuple(capabilities.effort_tiers))
@pytest.mark.parametrize(