diff --git a/litellm/batch_completion/main.py b/litellm/batch_completion/main.py index 702dd194fda..7ec71dcd4bb 100644 --- a/litellm/batch_completion/main.py +++ b/litellm/batch_completion/main.py @@ -1,13 +1,29 @@ +from collections.abc import Mapping from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait from typing import Final import litellm from litellm._logging import print_verbose -from litellm.utils import get_optional_params +from litellm.utils import get_model_info, get_optional_params from ..llms.vllm.completion import handler as vllm_handler +def _model_info_for_batch( + *, + model: str, + custom_llm_provider: str | None, + kwargs: Mapping[str, object], +) -> Mapping[str, object] | None: + from_kwargs: Final = kwargs.get("model_info") + if isinstance(from_kwargs, Mapping): + return from_kwargs + try: + return get_model_info(model=model, custom_llm_provider=custom_llm_provider) + except Exception: # noqa: BLE001 # get_model_info raises Exception for unmapped models + return None + + def batch_completion( model: str, # Optional OpenAI params: see https://platform.openai.com/docs/api-reference/chat/create @@ -79,9 +95,13 @@ def batch_completion( frequency_penalty=frequency_penalty, logit_bias=logit_bias, user=user, - # params to identify the model model=model, custom_llm_provider=custom_llm_provider, + model_info=_model_info_for_batch( + model=model, + custom_llm_provider=custom_llm_provider, + kwargs=kwargs, + ), ) results = vllm_handler.batch_completions( model=model, diff --git a/litellm/litellm_core_utils/thinking_param_translation.py b/litellm/litellm_core_utils/thinking_param_translation.py new file mode 100644 index 00000000000..f9916ed9a86 --- /dev/null +++ b/litellm/litellm_core_utils/thinking_param_translation.py @@ -0,0 +1,241 @@ +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final + +_SEND_VIA_EXTRA_BODY: Final = "extra_body" +_SEND_VIA_PROVIDER_MAPPED: Final = "provider_mapped" +_SEND_VIA_VALUES: Final = frozenset((_SEND_VIA_EXTRA_BODY, _SEND_VIA_PROVIDER_MAPPED)) +_THINKING_ENABLED_STRINGS: Final = frozenset(("enabled", "true", "1", "auto")) +_THINKING_TYPE_ENABLED: Final = frozenset(("enabled", "auto", "true")) + +_EFFORT_FALLBACKS: Final[Mapping[str, tuple[str, ...]]] = MappingProxyType( + { + "xhigh": ("max", "high"), + "max": ("xhigh", "high"), + "medium": ("high", "low"), + "minimal": ("low", "none"), + "none": ("low",), + } +) + + +@dataclass(frozen=True, slots=True) +class ThinkingParamsState: + thinking: object | None + reasoning_effort: object | None + extra_body: Mapping[str, object] + + +def _as_str_tuple(value: object) -> tuple[str, ...]: + if not isinstance(value, (list, tuple)): + return () + return tuple(item for item in value if isinstance(item, str)) + + +def _thinking_enabled(thinking: object) -> bool: + if isinstance(thinking, bool): + return thinking + if isinstance(thinking, str): + return thinking.lower() in _THINKING_ENABLED_STRINGS + if isinstance(thinking, Mapping): + typ: Final = thinking.get("type") + if isinstance(typ, str): + return typ.lower() in _THINKING_TYPE_ENABLED + enabled: Final = thinking.get("enabled") + if isinstance(enabled, bool): + return enabled + return False + + +def _thinking_type_value(thinking: object, allowed: Sequence[str]) -> str | None: + if isinstance(thinking, str): + candidate = thinking + elif isinstance(thinking, Mapping): + raw = thinking.get("type") + candidate = raw if isinstance(raw, str) else None + elif isinstance(thinking, bool): + candidate = "enabled" if thinking else "disabled" + else: + candidate = None + if candidate is None: + return None + if not allowed or candidate in allowed: + return candidate + if candidate == "auto" and "enabled" in allowed: + return "enabled" + return None + + +def _thinking_payload(thinking: object, typ: str) -> Mapping[str, object]: + if not isinstance(thinking, Mapping): + return MappingProxyType({"type": typ}) + budget: Final = thinking.get("budget_tokens") + if isinstance(budget, int): + return MappingProxyType({"type": typ, "budget_tokens": budget}) + return MappingProxyType({"type": typ}) + + +def _clamp_effort(value: object, allowed: Sequence[str]) -> str | None: + if not isinstance(value, str): + return None + if not allowed: + return None + if value in allowed: + return value + for fallback in _EFFORT_FALLBACKS.get(value, ()): + if fallback in allowed: + return fallback + return None + + +def _mapping_or_none(value: object) -> Mapping[str, object] | None: + if isinstance(value, Mapping): + return value + return None + + +def _merged_mapping_value(left: Mapping[str, object], right: Mapping[str, object], key: str) -> object: + if key not in right: + return left[key] + if key not in left: + return right[key] + left_map: Final = _mapping_or_none(left[key]) + right_map: Final = _mapping_or_none(right[key]) + if left_map is not None and right_map is not None: + return _deep_merge_pair(left_map, right_map) + return right[key] + + +def _deep_merge_pair(left: Mapping[str, object], right: Mapping[str, object]) -> Mapping[str, object]: + keys: Final = frozenset(left) | frozenset(right) + return MappingProxyType({key: _merged_mapping_value(left, right, key) for key in keys}) + + +def _map_thinking_to_extra_body( + *, + thinking_param: str | None, + thinking: object, + thinking_values: Sequence[str], +) -> Mapping[str, object]: + match thinking_param: + case "thinking.type": + typ: Final = _thinking_type_value(thinking, thinking_values) + if typ is None: + return MappingProxyType({}) + return MappingProxyType({"thinking": _thinking_payload(thinking, typ)}) + case "thinking": + if isinstance(thinking, Mapping): + return MappingProxyType({"thinking": MappingProxyType({k: thinking[k] for k in thinking})}) + typ_only: Final = _thinking_type_value(thinking, thinking_values or ("enabled", "disabled")) + if typ_only is None: + return MappingProxyType({}) + return MappingProxyType({"thinking": MappingProxyType({"type": typ_only})}) + case "enable_thinking": + return MappingProxyType({"enable_thinking": _thinking_enabled(thinking)}) + case "chat_template_kwargs": + return MappingProxyType( + {"chat_template_kwargs": MappingProxyType({"enable_thinking": _thinking_enabled(thinking)})} + ) + case None: + return MappingProxyType({}) + case _: + return MappingProxyType({}) + + +def _map_effort_to_extra_body( + *, + effort: object, + effort_values: Sequence[str], + send_via: object, +) -> Mapping[str, object]: + clamped: Final = _clamp_effort(effort, effort_values) + if clamped is not None: + return MappingProxyType({"reasoning_effort": clamped}) + if not effort_values and send_via == _SEND_VIA_EXTRA_BODY and isinstance(effort, str): + return MappingProxyType({"reasoning_effort": effort}) + return MappingProxyType({}) + + +def translate_thinking_params( + *, + model_info: Mapping[str, object] | None, + state: ThinkingParamsState, +) -> ThinkingParamsState: + if model_info is None: + return state + + send_via: Final = model_info.get("thinking_send_via") + if send_via not in _SEND_VIA_VALUES: + return state + + supports_reasoning: Final = model_info.get("supports_reasoning") is True + thinking_param_raw: Final = model_info.get("thinking_param") + thinking_param: Final = thinking_param_raw if isinstance(thinking_param_raw, str) else None + thinking_values: Final = _as_str_tuple(model_info.get("thinking_values")) + effort_values: Final = _as_str_tuple(model_info.get("reasoning_effort_values")) + + if not supports_reasoning and send_via != _SEND_VIA_PROVIDER_MAPPED: + return state + + thinking: Final = state.thinking + effort: Final = state.reasoning_effort + if thinking is None and effort is None: + return state + + keep_thinking: Final = send_via == _SEND_VIA_PROVIDER_MAPPED + thinking_patch: Final = ( + _map_thinking_to_extra_body( + thinking_param=thinking_param, + thinking=thinking, + thinking_values=thinking_values, + ) + if thinking is not None and send_via == _SEND_VIA_EXTRA_BODY + else MappingProxyType({}) + ) + effort_patch: Final = ( + _map_effort_to_extra_body( + effort=effort, + effort_values=effort_values, + send_via=send_via, + ) + if effort is not None + else MappingProxyType({}) + ) + patch: Final = MappingProxyType({**thinking_patch, **effort_patch}) + if not patch: + return state + + thinking_mapped: Final = any(key in patch for key in ("thinking", "enable_thinking", "chat_template_kwargs")) + next_thinking: Final = thinking if (keep_thinking or not thinking_mapped) else None + next_effort: Final = None if "reasoning_effort" in patch else effort + merged_extra: Final = _deep_merge_pair(state.extra_body, patch) + return ThinkingParamsState( + thinking=next_thinking, + reasoning_effort=next_effort, + extra_body=merged_extra, + ) + + +def apply_thinking_param_translation( + *, + model_info: Mapping[str, object] | None, + thinking: object | None, + reasoning_effort: object | None, + existing_extra_body: Mapping[str, object] | None, +) -> ThinkingParamsState: + base_extra: Final = ( + MappingProxyType({k: existing_extra_body[k] for k in existing_extra_body}) + if isinstance(existing_extra_body, Mapping) + else MappingProxyType({}) + ) + return translate_thinking_params( + model_info=model_info, + state=ThinkingParamsState( + thinking=thinking, + reasoning_effort=reasoning_effort, + extra_body=base_extra, + ), + ) diff --git a/litellm/main.py b/litellm/main.py index 7f4b34d28a0..6378cd6096f 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -5572,6 +5572,7 @@ def completion( "prompt_cache_key": prompt_cache_key, "allowed_openai_params": allowed_openai_params, "base_model": base_model, + "model_info": model_info if isinstance(model_info, dict) else None, } optional_params = get_optional_params(**optional_param_args, **non_default_params) processed_non_default_params: Final = pre_process_non_default_params( diff --git a/litellm/utils.py b/litellm/utils.py index dd35c17809f..99e8c69d86f 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -294,6 +294,9 @@ from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast, runtime_c from typing_extensions import assert_never from litellm import utils as litellm_utils +from litellm.litellm_core_utils.thinking_param_translation import ( + apply_thinking_param_translation, +) # These are lazy loaded via __getattr__ from litellm.llms.base_llm.base_utils import ( @@ -4294,6 +4297,49 @@ def remove_sensitive_keys_from_dict(d: dict) -> dict: return d +def _apply_model_info_thinking_translation( + *, + model_info: Mapping[str, object] | None, + passed_params: dict, + non_default_params: dict, +) -> None: + prior_thinking: Final = non_default_params.get("thinking") + prior_effort: Final = non_default_params.get("reasoning_effort") + existing_extra_raw: Final = passed_params.get("extra_body") + existing_extra: Final = existing_extra_raw if isinstance(existing_extra_raw, Mapping) else None + translated: Final = apply_thinking_param_translation( + model_info=model_info, + thinking=prior_thinking, + reasoning_effort=prior_effort, + existing_extra_body=existing_extra, + ) + prior_extra: Final = dict(existing_extra) if existing_extra is not None else {} # mutable-ok: equality snapshot + if ( + translated.thinking is prior_thinking + and translated.reasoning_effort is prior_effort + and dict(translated.extra_body) == prior_extra # mutable-ok: MappingProxyType equality snapshot + ): + return + + # mutable-ok: get_optional_params already mutates passed_params / non_default_params in place + if translated.thinking is None: + non_default_params.pop("thinking", None) + passed_params["thinking"] = None + else: + non_default_params["thinking"] = translated.thinking + passed_params["thinking"] = translated.thinking + + if translated.reasoning_effort is None: + non_default_params.pop("reasoning_effort", None) + passed_params["reasoning_effort"] = None + else: + non_default_params["reasoning_effort"] = translated.reasoning_effort + passed_params["reasoning_effort"] = translated.reasoning_effort + + if translated.extra_body: + passed_params["extra_body"] = dict(translated.extra_body) # mutable-ok: openai extra_body is a dict + + def pre_process_optional_params(passed_params: dict, non_default_params: dict, custom_llm_provider: str) -> dict: """For .completion(), preprocess optional params""" optional_params: dict = {} @@ -4416,6 +4462,7 @@ def get_optional_params( store: bool | None = None, prompt_cache_key: str | None = None, base_model: str | None = None, + model_info: Mapping[str, object] | None = None, **kwargs, ): drop_params = normalize_drop_params(drop_params) # rebind-ok: config and DB deployments pass "true" as a string @@ -4425,6 +4472,7 @@ def get_optional_params( # non_default_params / _check_valid_arg — it's a routing hint, not an # OpenAI param. passed_params.pop("base_model", None) + model_info_for_translation: Final = passed_params.pop("model_info", None) provider_config: BaseConfig | None = None if custom_llm_provider is not None and custom_llm_provider in [provider.value for provider in LlmProviders]: provider_config = ProviderConfigManager.get_provider_chat_config( @@ -4440,6 +4488,11 @@ def get_optional_params( model=model, provider_config=provider_config, ) + _apply_model_info_thinking_translation( + model_info=model_info_for_translation if isinstance(model_info_for_translation, Mapping) else None, + passed_params=passed_params, + non_default_params=non_default_params, + ) optional_params = pre_process_optional_params( passed_params=passed_params, non_default_params=non_default_params, diff --git a/tests/test_litellm/litellm_core_utils/test_thinking_param_translation.py b/tests/test_litellm/litellm_core_utils/test_thinking_param_translation.py new file mode 100644 index 00000000000..c5c91181278 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_thinking_param_translation.py @@ -0,0 +1,201 @@ +import importlib +from types import MappingProxyType + +from litellm.litellm_core_utils.thinking_param_translation import ( + ThinkingParamsState, + apply_thinking_param_translation, + translate_thinking_params, +) +from litellm.utils import get_optional_params + + +def _extra_body_model_info(**overrides: object) -> dict[str, object]: + base: dict[str, object] = { + "supports_reasoning": True, + "thinking_param": "thinking.type", + "thinking_values": ["enabled", "disabled"], + "reasoning_effort_values": ["low", "high", "max"], + "thinking_send_via": "extra_body", + } + return {**base, **overrides} + + +def test_translate_thinking_type_and_effort_to_extra_body(): + result = apply_thinking_param_translation( + model_info=_extra_body_model_info(), + thinking={"type": "enabled", "budget_tokens": 1024}, + reasoning_effort="high", + existing_extra_body=None, + ) + assert result.thinking is None + assert result.reasoning_effort is None + assert dict(result.extra_body) == { + "thinking": {"type": "enabled", "budget_tokens": 1024}, + "reasoning_effort": "high", + } + + +def test_translate_enable_thinking_bool(): + result = apply_thinking_param_translation( + model_info=_extra_body_model_info( + thinking_param="enable_thinking", + thinking_values=["true", "false"], + ), + thinking={"type": "enabled"}, + reasoning_effort=None, + existing_extra_body=None, + ) + assert result.thinking is None + assert dict(result.extra_body) == {"enable_thinking": True} + + +def test_translate_chat_template_kwargs(): + result = apply_thinking_param_translation( + model_info=_extra_body_model_info( + thinking_param="chat_template_kwargs", + thinking_values=[], + reasoning_effort_values=["low", "medium", "high"], + ), + thinking={"type": "disabled"}, + reasoning_effort="medium", + existing_extra_body=None, + ) + assert dict(result.extra_body) == { + "chat_template_kwargs": {"enable_thinking": False}, + "reasoning_effort": "medium", + } + + +def test_translate_chat_template_kwargs_preserves_existing_nested_keys(): + result = apply_thinking_param_translation( + model_info=_extra_body_model_info( + thinking_param="chat_template_kwargs", + thinking_values=[], + ), + thinking={"type": "enabled"}, + reasoning_effort=None, + existing_extra_body={"chat_template_kwargs": {"reasoning_budget": 512}}, + ) + assert result.extra_body["chat_template_kwargs"]["reasoning_budget"] == 512 + assert result.extra_body["chat_template_kwargs"]["enable_thinking"] is True + + +def test_translate_clamps_effort_aliases(): + result = apply_thinking_param_translation( + model_info=_extra_body_model_info(reasoning_effort_values=["low", "high", "max"]), + thinking=None, + reasoning_effort="xhigh", + existing_extra_body=None, + ) + assert result.reasoning_effort is None + assert result.extra_body["reasoning_effort"] == "max" + + +def test_translate_provider_mapped_keeps_thinking_moves_effort(): + result = translate_thinking_params( + model_info=_extra_body_model_info(thinking_send_via="provider_mapped"), + state=ThinkingParamsState( + thinking={"type": "enabled"}, + reasoning_effort="high", + extra_body=MappingProxyType({}), + ), + ) + assert result.thinking == {"type": "enabled"} + assert result.reasoning_effort is None + assert dict(result.extra_body) == {"reasoning_effort": "high"} + + +def test_translate_noop_without_model_info(): + state = ThinkingParamsState( + thinking={"type": "enabled"}, + reasoning_effort="high", + extra_body=MappingProxyType({}), + ) + assert translate_thinking_params(model_info=None, state=state) is state + + +def test_translate_noop_when_send_via_na(): + result = apply_thinking_param_translation( + model_info=_extra_body_model_info(thinking_send_via="n/a"), + thinking={"type": "enabled"}, + reasoning_effort="high", + existing_extra_body=None, + ) + assert result.thinking == {"type": "enabled"} + assert result.reasoning_effort == "high" + assert dict(result.extra_body) == {} + + +def test_get_optional_params_openai_drop_translates_via_model_info(): + optional_params = get_optional_params( + model="deepseek-v4-flash", + custom_llm_provider="openai", + drop_params=True, + thinking={"type": "enabled"}, + reasoning_effort="high", + model_info=_extra_body_model_info(), + ) + assert optional_params.get("thinking") is None + assert optional_params.get("reasoning_effort") is None + assert optional_params["extra_body"]["thinking"] == {"type": "enabled"} + assert optional_params["extra_body"]["reasoning_effort"] == "high" + + +def test_get_optional_params_openai_drop_without_model_info_drops_params(): + optional_params = get_optional_params( + model="gpt-4o", + custom_llm_provider="openai", + drop_params=True, + thinking={"type": "enabled"}, + reasoning_effort="high", + ) + extra_body = optional_params.get("extra_body") or {} + assert "thinking" not in extra_body + assert "reasoning_effort" not in extra_body + assert optional_params.get("thinking") is None + assert optional_params.get("reasoning_effort") is None + + +def test_get_optional_params_does_not_reintroduce_dropped_thinking(): + optional_params = get_optional_params( + model="deepseek-v4-flash", + custom_llm_provider="openai", + drop_params=True, + thinking={"type": "enabled"}, + additional_drop_params=["thinking"], + model_info=_extra_body_model_info( + thinking_param="chat_template_kwargs", + thinking_values=[], + ), + ) + extra_body = optional_params.get("extra_body") or {} + assert optional_params.get("thinking") is None + assert "enable_thinking" not in extra_body + assert "chat_template_kwargs" not in extra_body + + +def test_batch_completion_vllm_passes_model_info(monkeypatch): + batch_completion_mod = importlib.import_module("litellm.batch_completion.main") + + captured: dict[str, object] = {} + looked_up: dict[str, object] = _extra_body_model_info(thinking_param="enable_thinking") + + def fake_get_optional_params(**kwargs: object) -> dict[str, object]: + captured.update(kwargs) + return {} + + def fake_batch_completions(**kwargs: object) -> list[str]: + return ["ok"] + + def fake_get_model_info(**kwargs: object) -> dict[str, object]: + return looked_up + + monkeypatch.setattr(batch_completion_mod, "get_optional_params", fake_get_optional_params) + monkeypatch.setattr(batch_completion_mod.vllm_handler, "batch_completions", fake_batch_completions) + monkeypatch.setattr(batch_completion_mod, "get_model_info", fake_get_model_info) + + batch_completion_mod.batch_completion( + model="vllm/some-model", + messages=[[{"role": "user", "content": "hi"}]], + ) + assert captured.get("model_info") == looked_up