diff --git a/.circleci/scripts/run_integration.sh b/.circleci/scripts/run_integration.sh index 1380b07de53..fb91bcce77f 100644 --- a/.circleci/scripts/run_integration.sh +++ b/.circleci/scripts/run_integration.sh @@ -128,6 +128,8 @@ fi sudo iptables -L integration_only -n -v -x > "$results/egress-guard.txt" awk '$3 == "REJECT" && $1 > 0 { rejected=1 } END { exit !rejected }' "$results/egress-guard.txt" +export UV_OFFLINE=true + if [ "$mode" = parity ]; then export INTEGRATION_WORKERS=1 else diff --git a/litellm-rust/crates/model-catalog/src/model_info.rs b/litellm-rust/crates/model-catalog/src/model_info.rs index 8f7549c1c6d..141702bf857 100644 --- a/litellm-rust/crates/model-catalog/src/model_info.rs +++ b/litellm-rust/crates/model-catalog/src/model_info.rs @@ -186,6 +186,9 @@ pub struct ModelInfoFields { pub supports_image_size: Option, pub supported_audio_formats: Option, pub token_count_framing: Option, + pub token_count_tool_framing: Option, + pub token_count_responses_framing: Option, + pub token_count_responses_tool_framing: Option, pub token_count_model: Option, pub token_count_tokenizer: Option, pub vertex_ai_audio_api: Option, diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 91cf9bb2cfe..3b7ea66ad6b 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -45,13 +45,27 @@ from .redis_semantic_cache import RedisSemanticCache from .s3_cache import S3Cache RESPONSE_CACHE_TARGET: Final = "llm_response" -_ALWAYS_KEYED_PROVIDER_PARAMS: Final = frozenset({"continuation_token"}) +_ALWAYS_KEYED_PROVIDER_PARAMS: Final = frozenset({"continuation_token", "return_token_ids"}) _DECODED_LITERAL: Final = TypeAdapter(object) _REDIS_STARTUP_NODES_ADAPTER: Final = TypeAdapter( list[dict[str, object]], config=ConfigDict(strict=True, hide_input_in_errors=True) ) +def _keyed_provider_params_in_extra_body(kwargs: Mapping[str, object]) -> str: + raw_extra_body: Final = kwargs.get("extra_body") + if not isinstance(raw_extra_body, dict): + return "" + extra_body: Final = cast( # cast-ok: isinstance leaves the key and value types unknown + Mapping[str, object], raw_extra_body + ) + return "".join( + f"extra_body.{param}: {extra_body[param]}" + for param in sorted(_ALWAYS_KEYED_PROVIDER_PARAMS) + if extra_body.get(param) + ) + + @contextmanager def response_cache_phase(operation: Literal["get", "set"]) -> Generator[None]: """The ``cache.get llm_response`` / ``cache.set llm_response`` span a response-cache read or write runs @@ -429,7 +443,10 @@ class Cache: "" if embeds_thought_signature_in_tool_call_id(kwargs) else "thought_signature_in_tool_call_id: False" ) hashed_cache_key = Cache._get_hashed_cache_key( - cache_key + reasoning_summary_material + tool_call_id_shape_material + cache_key + + _keyed_provider_params_in_extra_body(kwargs) + + reasoning_summary_material + + tool_call_id_shape_material ) hashed_cache_key = self._add_namespace_to_cache_key(hashed_cache_key, **kwargs) verbose_logger.debug( diff --git a/litellm/constants.py b/litellm/constants.py index cd7c0176428..14a2a9f9eca 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -699,6 +699,9 @@ STREAM_SSE_DONE_STRING: Final[str] = "[DONE]" STREAM_SSE_DATA_PREFIX: Final[str] = "data: " STREAM_SSE_KEEPALIVE_PING_CHUNK: Final[str] = 'event: ping\ndata: {"type": "ping"}\n\n' STREAM_SSE_KEEPALIVE_PING_BYTES: Final[bytes] = STREAM_SSE_KEEPALIVE_PING_CHUNK.encode("utf-8") +STREAM_FINISH_CHOICE_FIELDS: Final = ("stop_reason",) +STREAM_CHUNK_CHOICE_FIELDS: Final = ("token_ids",) +VLLM_CHOICE_FIELDS: Final = STREAM_FINISH_CHOICE_FIELDS + STREAM_CHUNK_CHOICE_FIELDS ### SPEND TRACKING ### DEFAULT_REPLICATE_GPU_PRICE_PER_SECOND: Final = get_env_float( "DEFAULT_REPLICATE_GPU_PRICE_PER_SECOND", 0.001400 diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 42fe0837c29..1c9a443b546 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -3036,7 +3036,7 @@ class PrometheusLogger(CustomLogger): async def async_post_call_failure_hook( self, - request_data: dict, + request_data: dict[str, object], original_exception: Exception, user_api_key_dict: UserAPIKeyAuth, traceback_str: str | None = None, @@ -3059,10 +3059,14 @@ class PrometheusLogger(CustomLogger): try: _tags: Final = StandardLoggingPayloadSetup.get_request_tags( litellm_params=request_data, - proxy_server_request=request_data.get("proxy_server_request", {}), + proxy_server_request=_logged_fields(request_data.get("proxy_server_request")), + ) + _metadata: Final = _logged_fields(request_data.get("metadata")) + model_id: Final = ( + _logged_fields(_metadata.get("model_info")).get("id") + or _logged_fields(request_data.get("model_info")).get("id") + or _logged_fields(request_data.get("standard_logging_object")).get("model_id") ) - _metadata: Final = request_data.get("metadata", {}) or {} - model_id: Final = _metadata.get("model_info", {}).get("id") or request_data.get("model_info", {}).get("id") rate_limit_category, rate_limit_type = self._extract_rate_limit_labels(original_exception) api_provider: Final = self._extract_api_provider_from_request_data( request_data diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 639003c9a6a..5f909f639c4 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -918,6 +918,7 @@ class Logging(LiteLLMLoggingBaseClass): # post_call guardrails have run; the @client decorator then stores the # enqueue closure here instead of firing it immediately. self.defer_async_logging: bool = False + self.drops_stream_token_ids: bool = False self.enqueue_deferred_logging: Callable[[], None] | None = None self._async_success_scheduled: bool = False self._async_success_enqueued: bool = False diff --git a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py index 58fe5b4297f..6bb036e9080 100644 --- a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py +++ b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py @@ -3,7 +3,8 @@ import json import re import time import traceback -from collections.abc import Generator, Iterator, Mapping, Sequence +from collections.abc import Generator, Iterable, Iterator, Mapping, Sequence +from types import MappingProxyType from typing import Final, Literal, cast from pydantic import ConfigDict, TypeAdapter, ValidationError, with_config @@ -11,7 +12,12 @@ from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger -from litellm.constants import RESPONSE_FORMAT_TOOL_NAME +from litellm.constants import ( + RESPONSE_FORMAT_TOOL_NAME, + STREAM_CHUNK_CHOICE_FIELDS, + STREAM_FINISH_CHOICE_FIELDS, + VLLM_CHOICE_FIELDS, +) from litellm.litellm_core_utils.hidden_params import get_hidden_params from litellm.litellm_core_utils.prompt_templates.common_utils import ( extract_reasoning_content, @@ -170,6 +176,28 @@ def _split_assembled_content_for_replay(content: str | None) -> list[str]: return _REPLAY_CONTENT_SLICE_RE.findall(content) +def _drop_choice_fields(choice: StreamingChoices, fields: Iterable[str]) -> None: + for field in fields: + if hasattr(choice, field): + delattr(choice, field) + + +def _vllm_choice_fields(choice: Mapping[str, object]) -> Mapping[str, object]: + cached_fields: Final = _choice_provider_specific_fields(choice) + return MappingProxyType( + { + field: choice[field] if field in choice else cached_fields[field] + for field in VLLM_CHOICE_FIELDS + if field in choice or field in cached_fields + } + ) + + +def _choice_provider_specific_fields(choice: Mapping[str, object]) -> Mapping[str, object]: + fields: Final = choice.get("provider_specific_fields") + return _FIELD_MAP.validate_python(fields) if isinstance(fields, dict) else MappingProxyType({}) + + _REPLAY_REASONING_FIELDS: Final = frozenset({"reasoning_content", "thinking_blocks", "reasoning_items"}) _REPLAY_LEAD_FIELDS: Final = _REPLAY_REASONING_FIELDS | {"role", "provider_specific_fields"} _REPLAY_TOOL_FIELDS: Final = frozenset({"tool_calls", "function_call"}) @@ -220,8 +248,9 @@ def _replay_frame( choice.finish_reason = finish_reason if not keeps_choice_metadata: choice.logprobs = None - if hasattr(choice, "enhancements"): - del choice.enhancements + _drop_choice_fields(choice, ("enhancements", *STREAM_CHUNK_CHOICE_FIELDS)) + if finish_reason is None: + _drop_choice_fields(choice, STREAM_FINISH_CHOICE_FIELDS) return frame @@ -325,7 +354,9 @@ async def convert_to_streaming_response_async( logprobs = choice.get("logprobs", None) - choice = StreamingChoices(finish_reason=finish_reason, index=idx, delta=delta, logprobs=logprobs) + choice = StreamingChoices( + finish_reason=finish_reason, index=idx, delta=delta, logprobs=logprobs, **_vllm_choice_fields(choice) + ) choice_list.append(choice) model_response_object.choices = choice_list @@ -388,6 +419,7 @@ def convert_to_streaming_response( delta=delta, logprobs=logprobs, enhancements=enhancements, + **_vllm_choice_fields(choice), ) choice_list.append(choice) @@ -872,6 +904,7 @@ def convert_to_model_response_object( logprobs=logprobs, enhancements=enhancements, provider_specific_fields=provider_specific_fields, + **_vllm_choice_fields(choice), ) if provider_sent_logprobs and logprobs is None: choice.logprobs = None diff --git a/litellm/litellm_core_utils/prompt_templates/chat_template_tools.py b/litellm/litellm_core_utils/prompt_templates/chat_template_tools.py index 60134456577..49ffae03c40 100644 --- a/litellm/litellm_core_utils/prompt_templates/chat_template_tools.py +++ b/litellm/litellm_core_utils/prompt_templates/chat_template_tools.py @@ -10,12 +10,18 @@ class ChatTemplateToolSection: preamble_tokens: int implicit_system_turn_tokens: int render_tool_schemas: Callable[[Sequence[Mapping[str, object]]], str] + counts_tool_choice: bool = False + renders_strict_only_when_true: bool = False -def _deepseek_v4_tool_schemas(functions: Sequence[Mapping[str, object]]) -> str: +def _json_line_tool_schemas(functions: Sequence[Mapping[str, object]]) -> str: return "\n".join(json.dumps(function, ensure_ascii=False) for function in functions) +def _compact_json_line_tool_schemas(functions: Sequence[Mapping[str, object]]) -> str: + return "\n".join(json.dumps(function, ensure_ascii=False, separators=(",", ":")) for function in functions) + + def _qwen3_5_tool_schemas(functions: Sequence[Mapping[str, object]]) -> str: return "".join( "\n" + json.dumps({"type": "function", "function": function}, ensure_ascii=False) for function in functions @@ -23,13 +29,29 @@ def _qwen3_5_tool_schemas(functions: Sequence[Mapping[str, object]]) -> str: _TOOL_SECTIONS: Final = ( + ( + re.compile(r"deepseek-v4(?:[-._]|p)1(?!\d)", re.IGNORECASE), + ChatTemplateToolSection(211, 0, _json_line_tool_schemas), + ), ( re.compile( r"deepseek-v4(?![.\d]|p\d)|^(?:deepseek/)?deepseek-(?:chat|reasoner|flash)(?::free)?$", re.IGNORECASE ), - ChatTemplateToolSection(222, 0, _deepseek_v4_tool_schemas), + ChatTemplateToolSection(222, 0, _json_line_tool_schemas), ), (re.compile(r"qwen3[._]?5(?!\d)", re.IGNORECASE), ChatTemplateToolSection(196, 5, _qwen3_5_tool_schemas)), + ( + re.compile(r"grok-4[-.]1(?![.\d])", re.IGNORECASE), + ChatTemplateToolSection( + 0, 0, _compact_json_line_tool_schemas, counts_tool_choice=True, renders_strict_only_when_true=True + ), + ), + ( + re.compile(r"grok-4[-.]3(?![.\d])", re.IGNORECASE), + ChatTemplateToolSection( + 0, 0, _json_line_tool_schemas, counts_tool_choice=True, renders_strict_only_when_true=True + ), + ), ) diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 98348e434b2..ab1d2dea74e 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -34,6 +34,7 @@ from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import split_leading_system_run, string_field, ) +from litellm.litellm_core_utils.tool_fields import tools_without_hosted_mcp from litellm.litellm_core_utils.url_utils import async_safe_get, safe_get from litellm.llms.custom_httpx.http_handler import HTTPHandler, get_async_httpx_client, header_value from litellm.types.files import get_file_extension_from_mime_type @@ -5916,8 +5917,13 @@ def function_call_prompt( llm_provider: str, drop_params: bool, ) -> list[dict[str, object]]: + prompt_functions: Final = tools_without_hosted_mcp( + functions, model=model, route=llm_provider, drop_params=drop_params + ) + if not prompt_functions: + return messages header: Final = """To call a function, reply with JSON ONLY in this format {"name": "function_name", "arguments":{"argument_name": "argument_value"}}. Once a function result answers the request, reply to the user in plain text instead of calling a function again. The following functions are available to you:""" - function_prompt: Final = header + "".join(f"\n{function}\n" for function in functions) + function_prompt: Final = header + "".join(f"\n{function}\n" for function in prompt_functions) if not any(_is_system_message(message) for message in messages): return [*messages, {"role": "system", "content": function_prompt}] return [ diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index 981859c620b..2f8ecdcc9b9 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -196,6 +196,9 @@ def _redact_function_call(function_call) -> None: def _redact_choice_content(choice): """Helper to redact content in a choice (message or delta).""" + if isinstance(choice, (litellm.Choices, litellm.utils.StreamingChoices, litellm.utils.TextChoices)): + drop_token_ids(choice) + _redact_provider_specific_fields(choice) if isinstance(choice, litellm.Choices): if choice.message.content is not None: choice.message.content = REDACTED_BY_LITELLM @@ -216,6 +219,9 @@ def _redact_choice_content(choice): _redact_provider_specific_fields(choice.delta) _redact_tool_calls(getattr(choice.delta, "tool_calls", None)) _redact_function_call(getattr(choice.delta, "function_call", None)) + elif isinstance(choice, litellm.utils.TextChoices): + _redact_field_when_present(choice, "text") + _redact_field_when_present(choice, "reasoning_content") def _as_mapping(value: object) -> Mapping[str, object] | None: @@ -227,14 +233,27 @@ def _redact_provider_specific_fields(message: object) -> None: setattr(message, "provider_specific_fields", _redacted_provider_specific_fields(fields)) # noqa: B010 # typed as object +def _drop_provider_specific_token_ids(holder: object) -> None: + fields: Final = _as_mapping(getattr(holder, "provider_specific_fields", None)) + if fields is not None and fields.get("token_ids") is not None: + setattr(holder, "provider_specific_fields", {**fields, "token_ids": None}) # noqa: B010 # typed as object + + +def drop_token_ids(choice: object) -> None: + if getattr(choice, "token_ids", None) is not None: + setattr(choice, "token_ids", None) # noqa: B010 # pydantic extra, not a declared field + for holder in (choice, getattr(choice, "delta", None), getattr(choice, "message", None)): + _drop_provider_specific_token_ids(holder) + + def _redact_provider_specific_fields_dict(choice: Mapping[str, object]) -> None: - for message in (choice.get("message"), choice.get("delta")): + for holder in (choice, choice.get("message"), choice.get("delta")): if ( - (readable := _as_mapping(message)) is not None + (readable := _as_mapping(holder)) is not None and (fields := _as_mapping(readable.get("provider_specific_fields"))) is not None - and isinstance(message, dict) + and isinstance(holder, dict) ): - message["provider_specific_fields"] = _redacted_provider_specific_fields(fields) + holder["provider_specific_fields"] = _redacted_provider_specific_fields(fields) def _redacted_provider_specific_fields(fields: Mapping[str, object]) -> Mapping[str, object]: @@ -251,6 +270,7 @@ def _redacted_provider_specific_fields(fields: Mapping[str, object]) -> Mapping[ "compaction_blocks", "compaction_start", "compaction_delta", + "token_ids", ) if key in fields }, @@ -422,6 +442,8 @@ def _redact_model_response_dict_choices(choices, redacted_str: str): for choice in choices: if isinstance(choice, dict): _redact_provider_specific_fields_dict(choice) + if choice.get("token_ids") is not None: + choice["token_ids"] = None if "message" in choice and isinstance(choice["message"], dict): if choice["message"].get("content") is not None: choice["message"]["content"] = redacted_str diff --git a/litellm/litellm_core_utils/stream_usage_chunks.py b/litellm/litellm_core_utils/stream_usage_chunks.py index eda57c7ea12..d073ddc2724 100644 --- a/litellm/litellm_core_utils/stream_usage_chunks.py +++ b/litellm/litellm_core_utils/stream_usage_chunks.py @@ -1,6 +1,7 @@ from collections.abc import Iterable from typing import Final +from litellm.constants import VLLM_CHOICE_FIELDS from litellm.types.utils import ModelResponseStream, StreamingChoices @@ -9,12 +10,19 @@ def is_empty_streaming_choice(choice: StreamingChoices) -> bool: return False if getattr(choice, "logprobs", None) is not None: return False + if not _all_values_none(_vllm_choice_values(choice)): + return False delta: Final = choice.delta if hasattr(choice, "delta") else None if delta is None: return True return _all_values_none(delta.model_dump().values()) +def _vllm_choice_values(choice: StreamingChoices) -> Iterable[object]: + choice_extras: Final = choice.model_extra or {} + return (choice_extras.get(field) for field in VLLM_CHOICE_FIELDS) + + def _all_values_none(values: Iterable[object]) -> bool: return all(value is None for value in values) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 7191b2ce8e8..c87357122ce 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -12,6 +12,7 @@ from typing_extensions import ReadOnly, Required import litellm from litellm._logging import verbose_logger +from litellm.constants import VLLM_CHOICE_FIELDS from litellm.litellm_core_utils.core_helpers import ( PROVIDERS_BILLED_FROM_THEIR_OWN_USAGE_FIELDS, UPSTREAM_USAGE_COST, @@ -31,6 +32,7 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_content_from_model_response, # pyright: ignore[reportUnknownVariableType] # response param accepts an untyped dict ) +from litellm.litellm_core_utils.redact_messages import drop_token_ids from litellm.litellm_core_utils.token_counter import usage_tokenizer_model from litellm.types.llms.openai import ( AllMessageValues, @@ -434,6 +436,8 @@ class ChunkProcessor: chunks, unreported="tool_calls" if _carries_tool_calls(chunks) else "stop" ), message={"role": ChunkProcessor._get_role_from_chunks(chunks), "content": ""}, + provider_specific_fields=None, + **_streamed_vllm_choice_fields(chain.from_iterable(chunk.get("choices") or () for chunk in chunks)), ) @staticmethod @@ -504,6 +508,7 @@ class ChunkProcessor: "index": 0, "message": {"role": role, "content": ""}, "finish_reason": finish_reason, + **_streamed_vllm_choice_fields(_first_choices(chunks)), } ] ), @@ -1371,6 +1376,7 @@ def _rebuilt_text_choice(index: int, choices: Sequence[TextChoices]) -> TextChoi index=index, logprobs=_merged_logprobs(choices), finish_reason=next((choice.finish_reason for choice in reversed(choices) if choice.finish_reason), None), + **_streamed_vllm_choice_fields(choices), ) @@ -1504,6 +1510,31 @@ def _field_of(value: object, name: str) -> object: return getattr(value, name, None) +def _first_choices(chunks: Iterable[object]) -> tuple[object, ...]: + return tuple(next(iter(object_items(_field_of(chunk, "choices"))), None) for chunk in chunks) + + +def _choice_extras(choice: object) -> Mapping[str, object]: + if isinstance(choice, BaseModel): + return choice.model_extra or {} + return object_mapping(choice) or {} + + +def _streamed_vllm_choice_fields(choices: Iterable[object]) -> Mapping[str, object]: + extras: Final = tuple(map(_choice_extras, choices)) + token_id_runs: Final = tuple(object_items(fields.get("token_ids")) for fields in extras) + stop_reasons: Final = tuple(reason for fields in extras if (reason := fields.get("stop_reason")) is not None) + has_token_ids: Final = any(fields.get("token_ids") is not None for fields in extras) + streamed_values: Final[Mapping[str, object]] = MappingProxyType( + { + "token_ids": list(chain.from_iterable(token_id_runs)) if has_token_ids else None, + "stop_reason": stop_reasons[-1] if stop_reasons else None, + } + ) + present_fields: Final = frozenset(chain.from_iterable(fields.keys() for fields in extras)) + return MappingProxyType({field: streamed_values[field] for field in VLLM_CHOICE_FIELDS if field in present_fields}) + + def _delta_reasoning_items(delta: object) -> Sequence[object]: items: Final = _field_of(delta, "reasoning_items") return cast(Sequence[object], items) if isinstance(items, list) else () # cast-ok: items forwarded verbatim @@ -1776,6 +1807,19 @@ def stream_chunk_builder( end_time: object = None, logging_obj: Optional["Logging"] = None, count_prompt_tokens: Callable[[], int] | None = None, +) -> ModelResponse | TextCompletionResponse | None: + response: Final = _assembled_stream_response(chunks, messages, logging_obj, count_prompt_tokens) + if response is not None and getattr(logging_obj, "drops_stream_token_ids", False) is True: + for choice in response.choices: + drop_token_ids(choice) + return response + + +def _assembled_stream_response( + chunks: Sequence[object], + messages: Sequence[AllMessageValues | Message] | None, + logging_obj: Optional["Logging"], + count_prompt_tokens: Callable[[], int] | None, ) -> ModelResponse | TextCompletionResponse | None: try: if chunks is None: @@ -1788,7 +1832,7 @@ def stream_chunk_builder( if not chunks: return None - processor: Final = ChunkProcessor(chunks, messages) + processor: Final = ChunkProcessor(list(chunks), messages) chunks = processor.chunks ### BASE-CASE ### diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 9bb5b761fd5..a9ef06c3b4f 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -23,6 +23,7 @@ from typing_extensions import NotRequired, TypedDict import litellm from litellm import verbose_logger from litellm._uuid import uuid +from litellm.constants import STREAM_CHUNK_CHOICE_FIELDS, STREAM_FINISH_CHOICE_FIELDS from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.core_helpers import ( UPSTREAM_USAGE_COST, @@ -37,10 +38,11 @@ from litellm.litellm_core_utils.hidden_params import ( set_hidden_param, stored_attribute, ) +from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import object_items from litellm.litellm_core_utils.model_response_utils import ( is_model_response_stream_empty, ) -from litellm.litellm_core_utils.redact_messages import LiteLLMLoggingObject +from litellm.litellm_core_utils.redact_messages import LiteLLMLoggingObject, drop_token_ids from litellm.litellm_core_utils.thread_pool_executor import executor from litellm.llms.ollama.common_utils import OllamaResponseIdNormalizer from litellm.types.llms.base import LiteLLMBaseModel @@ -137,6 +139,70 @@ def _prepend_upstream_choice_extras( choice.__pydantic_fields_set__.update(upstream_extras) +_NO_UPSTREAM_CHOICE_FIELDS: Final[Mapping[str, object]] = MappingProxyType({}) +_NO_CHOICE_FIELDS_BY_INDEX: Final[Mapping[int, Mapping[str, object]]] = MappingProxyType({}) + + +def _upstream_first_choice(upstream_chunk: object) -> BaseModel | None: + match getattr(upstream_chunk, "choices", None): + case [BaseModel() as first_choice, *_]: + return first_choice + case _: + return None + + +def _upstream_choice_fields(upstream_chunk: object, fields: Sequence[str]) -> Mapping[str, object]: + first_choice: Final = _upstream_first_choice(upstream_chunk) + if first_choice is None: + return _NO_UPSTREAM_CHOICE_FIELDS + upstream_extras: Final[Mapping[str, object]] = first_choice.model_extra or {} + return MappingProxyType({field: upstream_extras[field] for field in fields if field in upstream_extras}) + + +def _upstream_choice_index(upstream_chunk: object) -> int: + upstream_index: Final[object] = getattr(_upstream_first_choice(upstream_chunk), "index", 0) + return upstream_index if isinstance(upstream_index, int) else 0 + + +def _response_choice_index(response_obj: Mapping[str, object]) -> int: + choice_index: Final = response_obj.get("choice_index") + if isinstance(choice_index, int): + return choice_index + return _upstream_choice_index(response_obj.get("original_chunk")) + + +def _non_null_upstream_choice_fields(upstream_chunk: object, fields: Sequence[str]) -> Mapping[str, object]: + first_choice: Final = _upstream_first_choice(upstream_chunk) + upstream_extras: Final = first_choice.model_extra if first_choice is not None else None + if not upstream_extras: + return _NO_UPSTREAM_CHOICE_FIELDS + return MappingProxyType( + {field: upstream_extras[field] for field in fields if upstream_extras.get(field) is not None} + ) + + +def _apply_choice_fields(choice: StreamingChoices, fields: Mapping[str, object]) -> None: + for field, value in fields.items(): + setattr(choice, field, value) + + +def _choice_fields_chunk( + model_response: ModelResponseStream, upstream_index: int, fields: Mapping[str, object] +) -> ModelResponseStream: + model_response.choices[0].delta = Delta(content="") # rebind-ok: chunks are built in place here + model_response.choices[0].index = upstream_index # rebind-ok: chunks are built in place here + _apply_choice_fields(model_response.choices[0], fields) + return model_response + + +def _without_token_ids_under_post_call_rules(chunk: ModelResponseStream | None) -> ModelResponseStream | None: + if chunk is None or not litellm.post_call_rules: + return chunk + for choice in chunk.choices: + drop_token_ids(choice) + return chunk + + def _holds_only_scalars(fields: Iterable[tuple[str, object]]) -> bool: return all(value is None or isinstance(value, (str, int, float, bool)) for _, value in fields) @@ -378,6 +444,8 @@ class CustomStreamWrapper: self._provider_response_model: str | None = None self._upstream_provider: str | None = None self.received_finish_reason: str | None = None + self.received_finish_choice_fields: Mapping[int, Mapping[str, object]] = _NO_CHOICE_FIELDS_BY_INDEX + self.carried_chunk_choice_fields: Mapping[int, Mapping[str, object]] = _NO_CHOICE_FIELDS_BY_INDEX self.intermittent_finish_reason: str | None = None # finish reasons that show up mid-stream self.received_any_finish_reason = False # a later chunk can reset intermittent_finish_reason self.special_tokens = [ @@ -823,6 +891,7 @@ class CustomStreamWrapper: "logprobs": logprobs, "original_chunk": str_line, "usage": usage, + "chunk_choice_fields": _non_null_upstream_choice_fields(str_line, STREAM_CHUNK_CHOICE_FIELDS), } except Exception as e: raise e @@ -847,7 +916,7 @@ class CustomStreamWrapper: except Exception as e: raise e - def handle_openai_text_completion_chunk(self, chunk): + def handle_openai_text_completion_chunk(self, chunk: object): try: text = "" is_finished = False @@ -860,11 +929,15 @@ class CustomStreamWrapper: is_finished = True finish_reason = choices[0].finish_reason usage = getattr(chunk, "usage", None) + choice_index: Final = _upstream_choice_index(chunk) return { "text": text, "is_finished": is_finished, "finish_reason": finish_reason, "usage": usage, + "chunk_choice_fields": _non_null_upstream_choice_fields(chunk, STREAM_CHUNK_CHOICE_FIELDS), + "choice_index": choice_index, + "sibling_finish_reason": finish_reason if choice_index != 0 else None, } except Exception as e: @@ -1183,6 +1256,13 @@ class CustomStreamWrapper: ) is_chunk_non_empty: Final = self.is_chunk_non_empty(completion_obj, model_response, response_obj) + upstream_chunk_choice_fields: Final[Mapping[str, object]] = ( + _NO_UPSTREAM_CHOICE_FIELDS + if litellm.post_call_rules + else response_obj.get("chunk_choice_fields", _NO_UPSTREAM_CHOICE_FIELDS) + ) + chunk_choice_fields: Final = self._take_carried_chunk_choice_fields(response_obj, upstream_chunk_choice_fields) + sibling_finish_reason: Final[str | None] = response_obj.get("sibling_finish_reason") if is_chunk_non_empty: # cannot set content of an OpenAI Object to be an empty string self.raise_on_model_repetition() @@ -1270,16 +1350,40 @@ class CustomStreamWrapper: if response_obj.get("provider_specific_fields") is not None: completion_obj["provider_specific_fields"] = response_obj["provider_specific_fields"] model_response.choices[0].delta = Delta(**completion_obj) - _index: Final[int | None] = completion_obj.get("index") + _index: Final[int | None] = completion_obj.get("index", response_obj.get("choice_index")) if _index is not None: model_response.choices[0].index = _index + _apply_choice_fields(model_response.choices[0], chunk_choice_fields) + if sibling_finish_reason is not None: + self._finish_sibling_choice(model_response.choices[0], sibling_finish_reason) + elif response_obj.get("is_finished"): + self._apply_finish_choice_fields(model_response.choices[0]) self._optional_combine_thinking_block_in_choices(model_response=model_response) return model_response else: return + elif sibling_finish_reason is not None: + sibling_chunk: Final = _choice_fields_chunk( + model_response, _response_choice_index(response_obj), chunk_choice_fields + ) + self._finish_sibling_choice(sibling_chunk.choices[0], sibling_finish_reason) + return sibling_chunk elif self.received_finish_reason is not None: + chunk_upstream_index: Final = _response_choice_index(response_obj) + if chunk_choice_fields and self._streams_past_the_pending_finish( + model_response, response_obj, chunk_upstream_index + ): + self.carried_chunk_choice_fields = MappingProxyType( + {**self.carried_chunk_choice_fields, chunk_upstream_index: chunk_choice_fields} + ) + else: + another_choice_fields_chunk: Final = self._another_choice_fields_chunk( + model_response, response_obj, chunk_choice_fields + ) + if another_choice_fields_chunk is not None: + return another_choice_fields_chunk if self.sent_last_chunk is True: # Bedrock returns the guardrail trace in the last chunk - we want to return this here if self.custom_llm_provider == "bedrock" and "trace" in model_response: @@ -1317,12 +1421,20 @@ class CustomStreamWrapper: model_response.choices[0].finish_reason = self._normalize_tool_call_finish_reason( map_finish_reason(finish_reason=self.received_finish_reason) ) # ensure consistent output to openai + self._apply_finish_choice_fields(model_response.choices[0]) + if ( + response_obj.get("is_finished") + and _response_choice_index(response_obj) == model_response.choices[0].index + ): + _apply_choice_fields(model_response.choices[0], chunk_choice_fields) self.sent_last_chunk = True return model_response elif self._has_special_delta_content(model_response): return self._handle_special_delta_content(model_response) + elif chunk_choice_fields: + return _choice_fields_chunk(model_response, _response_choice_index(response_obj), chunk_choice_fields) else: if hasattr(model_response, "usage"): self.chunks.append(model_response) @@ -1545,6 +1657,8 @@ class CustomStreamWrapper: completion_obj["content"] = response_obj["text"] print_verbose(f"completion obj content: {completion_obj['content']}") if response_obj["is_finished"]: + self._record_finish_choice_fields(chunk) + if response_obj["is_finished"] and response_obj["sibling_finish_reason"] is None: self.received_finish_reason = response_obj["finish_reason"] if response_obj["usage"] is not None: _text_completion_usage: Final[Usage] = response_obj["usage"] @@ -1608,6 +1722,7 @@ class CustomStreamWrapper: self.system_fingerprint = cached_chunk.system_fingerprint if response_obj["is_finished"]: self.received_finish_reason = response_obj["finish_reason"] + self._record_finish_choice_fields(cached_chunk) else: # openai / azure chat model self.uses_openai_chat_stream = True if isinstance(chunk, ChatCompletionChunk): @@ -1637,6 +1752,7 @@ class CustomStreamWrapper: f"{self.custom_llm_provider} raised a streaming error - finish_reason: error, no content string given. Received Chunk={response_obj}" ) self.received_finish_reason = response_obj["finish_reason"] + self._record_finish_choice_fields(chunk) if response_obj.get("original_chunk", None) is not None: if hasattr(response_obj["original_chunk"], "id"): model_response = self.set_model_id(response_obj["original_chunk"].id, model_response) @@ -1704,7 +1820,7 @@ class CustomStreamWrapper: completion_obj=completion_obj, ) if isinstance(dispatch_result, _ProviderChunkEarlyReturn): - return dispatch_result.value + return _without_token_ids_under_post_call_rules(dispatch_result.value) response_obj = dispatch_result.response_obj model_response.model = self._served_openrouter_model() or self.model @@ -1794,7 +1910,7 @@ class CustomStreamWrapper: response_obj=response_obj, plain_delta_created=plain_delta_created, ) - return result + return _without_token_ids_under_post_call_rules(result) except StopIteration: raise StopIteration @@ -1987,6 +2103,75 @@ class CustomStreamWrapper: return "tool_calls" return finish_reason + def _record_finish_choice_fields(self, upstream_chunk: object) -> None: + self.received_finish_choice_fields = MappingProxyType( + { + **self.received_finish_choice_fields, + _upstream_choice_index(upstream_chunk): _upstream_choice_fields( + upstream_chunk, STREAM_FINISH_CHOICE_FIELDS + ), + } + ) + + def _apply_finish_choice_fields(self, choice: StreamingChoices) -> None: + _apply_choice_fields(choice, self.received_finish_choice_fields.get(choice.index, _NO_UPSTREAM_CHOICE_FIELDS)) + + def _finish_sibling_choice(self, choice: StreamingChoices, finish_reason: str) -> None: + choice.finish_reason = map_finish_reason(finish_reason) # rebind-ok: chunks are built in place here + self._apply_finish_choice_fields(choice) + + def _streams_past_the_pending_finish( + self, model_response: ModelResponseStream, response_obj: Mapping[str, object], upstream_index: int + ) -> bool: + finishing_index: Final = model_response.choices[0].index + return ( + not self.sent_last_chunk + and not response_obj.get("is_finished") + and upstream_index != finishing_index + and finishing_index in self.received_finish_choice_fields + ) + + def _take_carried_chunk_choice_fields( + self, response_obj: Mapping[str, object], chunk_choice_fields: Mapping[str, object] + ) -> Mapping[str, object]: + if not self.carried_chunk_choice_fields: + return chunk_choice_fields + upstream_index: Final = _response_choice_index(response_obj) + carried: Final = self.carried_chunk_choice_fields.get(upstream_index) + if carried is None: + return chunk_choice_fields + self.carried_chunk_choice_fields = MappingProxyType( + {index: fields for index, fields in self.carried_chunk_choice_fields.items() if index != upstream_index} + ) + return MappingProxyType( + { + field: [*object_items(carried.get(field)), *object_items(chunk_choice_fields.get(field))] + for field in STREAM_CHUNK_CHOICE_FIELDS + if field in carried or field in chunk_choice_fields + } + ) + + def _another_choice_fields_chunk( + self, + model_response: ModelResponseStream, + response_obj: Mapping[str, object], + chunk_choice_fields: Mapping[str, object], + ) -> ModelResponseStream | None: + upstream_index: Final = _response_choice_index(response_obj) + is_finished: Final = bool(response_obj.get("is_finished")) + if is_finished and upstream_index == model_response.choices[0].index and not self.sent_last_chunk: + return None + finish_fields: Final = ( + self.received_finish_choice_fields.get(upstream_index, _NO_UPSTREAM_CHOICE_FIELDS) + if is_finished + else _NO_UPSTREAM_CHOICE_FIELDS + ) + if not chunk_choice_fields and all(value is None for value in finish_fields.values()): + return None + return _choice_fields_chunk( + model_response, upstream_index, MappingProxyType({**finish_fields, **chunk_choice_fields}) + ) + def finish_reason_handler(self): model_response: Final = self.model_response_creator( hidden_params=None if self.received_model_id is None else {"received_model_id": self.received_model_id} @@ -2005,6 +2190,7 @@ class CustomStreamWrapper: model_response.choices[0].finish_reason = self._normalize_tool_call_finish_reason( _finish_reason if _finish_reason is not None else "stop" ) + self._apply_finish_choice_fields(model_response.choices[0]) close_tag: Final = self._close_open_think_block() if close_tag is not None: model_response.choices[0].delta.content = close_tag diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 323af72940a..df8038e756c 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -990,7 +990,11 @@ def _count_extra( tool_section (Optional[ChatTemplateToolSection]): How the model's chat template renders tools, if known. """ reply_primer_tokens: Final = 3 - tool_choice_tokens: Final = 0 if tool_section is not None else _count_tool_choice(count_function, tool_choice) + tool_choice_tokens: Final = ( + _count_tool_choice(count_function, tool_choice) + if tool_section is None or tool_section.counts_tool_choice + else 0 + ) return ( reply_primer_tokens + _count_tool_definitions(count_function, tools, includes_system_message, tool_section) @@ -1029,14 +1033,29 @@ def _count_tool_choice( return 0 +def _named_function_definitions(tools: Sequence[object]) -> Iterable[Mapping[str, object]]: + return (function for function in _function_definitions(_tools_sent_as_functions(tools)) if function.get("name")) + + +def sends_function_tools(tools: Sequence[object] | None) -> bool: + return tools is not None and any(_named_function_definitions(tools)) + + +def _without_false_strict(function: Mapping[str, object]) -> Mapping[str, object]: + return {key: value for key, value in function.items() if key != "strict" or value} + + def _count_chat_template_tools( count_function: TokenCounterFunction, tools: Sequence[object], includes_system_message: bool, tool_section: ChatTemplateToolSection, ) -> int: - functions: Final = tuple( - function for function in _function_definitions(_tools_sent_as_functions(tools)) if function.get("name") + named_functions: Final = tuple(_named_function_definitions(tools)) + functions: Final = ( + tuple(map(_without_false_strict, named_functions)) + if tool_section.renders_strict_only_when_true + else named_functions ) if not functions: return 0 diff --git a/litellm/litellm_core_utils/tool_fields.py b/litellm/litellm_core_utils/tool_fields.py index f90c8f5d174..5ee96ea9133 100644 --- a/litellm/litellm_core_utils/tool_fields.py +++ b/litellm/litellm_core_utils/tool_fields.py @@ -1,13 +1,16 @@ -from collections.abc import Callable, Mapping +from collections.abc import Callable, Mapping, Sequence +from copy import copy from dataclasses import dataclass from functools import reduce -from typing import Annotated, Final, TypeAlias +from typing import Annotated, Final, TypeAlias, TypeVar from pydantic import BeforeValidator, ConfigDict, Json, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict import litellm -from litellm.litellm_core_utils.core_helpers import normalize_drop_params +from litellm._logging import verbose_logger +from litellm.exceptions import UnsupportedParamsError +from litellm.litellm_core_utils.core_helpers import drops_unsupported_params, normalize_drop_params ROUTES_WITH_TOOLS: Final = frozenset( { @@ -162,3 +165,145 @@ def resolve_deployment_tool_fields( if "proxy_server_request" not in kwargs or route not in ROUTES_WITH_TOOLS: return kwargs return resolve_tool_fields(kwargs, lambda: (drop_policy(kwargs),)) + + +_ToolT: Final = TypeVar("_ToolT") + +HOSTED_MCP_DROP_WARNING: Final = ( + "Dropping the hosted MCP tool(s) %s for model=%s on %s (drop_params=True) because this route cannot run a " + "hosted MCP tool." +) + + +def _hosted_mcp_server_label(tool: object) -> str | None: + match tool: + case {"type": "mcp", "server_label": str() as server_label}: + return server_label + case {"type": "mcp"}: + return "unlabeled" + case _: + return None + + +def _drop_params_hint(dropped: str) -> str: + return ( + f"To drop {dropped} instead, set `litellm.drop_params=True` or for proxy:\n\n" + "`litellm_settings:\n drop_params: true`\n" + ) + + +def _one_log_line(text: str) -> str: + return text.replace("\r", "").replace("\n", "") + + +def _unrunnable_mcp_error(message: str, model: str, route: str) -> UnsupportedParamsError: + return UnsupportedParamsError(message=message, model=model, llm_provider=route, status_code=400) + + +def tools_without_hosted_mcp(tools: Sequence[_ToolT], model: str, route: str, drop_params: bool) -> Sequence[_ToolT]: + server_labels: Final = tuple(label for label in map(_hosted_mcp_server_label, tools) if label is not None) + if not server_labels: + return tools + named_servers: Final = ", ".join(f"`{label}`" for label in server_labels) + if not drops_unsupported_params(drop_params): + raise _unrunnable_mcp_error( + f'{model} on {route} cannot run the hosted MCP tool {named_servers} (`type: "mcp"`). To have LiteLLM ' + "run the MCP server for this model, register it on the LiteLLM MCP gateway and send " + '`server_url: "litellm_proxy/mcp/"`. ' + _drop_params_hint("hosted MCP tools"), + model=model, + route=route, + ) + verbose_logger.warning( + HOSTED_MCP_DROP_WARNING, _one_log_line(named_servers), _one_log_line(model), _one_log_line(route) + ) + return [tool for tool in tools if _hosted_mcp_server_label(tool) is None] + + +MCP_SERVERS_DROP_WARNING: Final = ( + "Dropping mcp_servers %s and their mcp_toolset tools for model=%s on %s (drop_params=True) because LiteLLM " + "cannot pass mcp_servers to this route." +) + +ROUTES_WITHOUT_MCP_SERVERS: Final = frozenset( + {"bedrock", "databricks", "deepseek", "gemini", "gigachat", "ollama", "ollama_chat", "vertex_ai", "vertex_ai_beta"} +) + + +def _mcp_server_name(server: object) -> str: + match server: + case {"name": str() as name}: + return name + case _: + return "unnamed" + + +def _mcp_server_names(mcp_servers: object) -> tuple[str, ...]: + if not mcp_servers: + return () + if not isinstance(mcp_servers, list): + return ("unnamed",) + return tuple(map(_mcp_server_name, _TOOLS.validate_python(mcp_servers))) + + +def _is_mcp_toolset(tool: object) -> bool: + match tool: + case {"type": "mcp_toolset"}: + return True + case _: + return False + + +def drops_mcp_servers(request_kwargs: Mapping[str, object]) -> bool: + additional_drop_params: Final = request_kwargs.get("additional_drop_params") + return drops_unsupported_params(request_kwargs.get("drop_params")) or ( + isinstance(additional_drop_params, list) and "mcp_servers" in additional_drop_params + ) + + +def tools_without_mcp_servers( + mcp_servers: object, tools: Sequence[_ToolT], model: str, route: str, drop_params: bool +) -> Sequence[_ToolT]: + server_names: Final = _mcp_server_names(mcp_servers) + if not server_names: + return tools + named_servers: Final = ", ".join(f"`{name}`" for name in server_names) + if not drops_unsupported_params(drop_params): + raise _unrunnable_mcp_error( + f"LiteLLM cannot pass `mcp_servers` to {model} on {route}, so the MCP server {named_servers} would not " + 'run. To have LiteLLM run it for this model, register it on the LiteLLM MCP gateway and send a `type: "mcp"` ' + 'tool with `server_url: "litellm_proxy/mcp/"`. ' + _drop_params_hint("`mcp_servers`"), + model=model, + route=route, + ) + verbose_logger.warning( + MCP_SERVERS_DROP_WARNING, _one_log_line(named_servers), _one_log_line(model), _one_log_line(route) + ) + return [tool for tool in tools if not _is_mcp_toolset(tool)] + + +_TOOL_SELECTION_PARAMS: Final = frozenset({"tools", "tool_choice", "parallel_tool_calls"}) +_ParamsT: Final = TypeVar("_ParamsT", bound=dict[str, object]) + + +def tool_params_without_hosted_mcp( + params: _ParamsT, model: str, route: str, drop_params: bool, runs_hosted_mcp: Callable[[], bool] | None = None +) -> _ParamsT: + tools: Final = params.get("tools") + if not isinstance(tools, list) or not tools: + return params + sent_tools: Final = _TOOLS.validate_python(tools) + if all(_hosted_mcp_server_label(tool) is None for tool in sent_tools) or ( + runs_hosted_mcp is not None and runs_hosted_mcp() + ): + return params + kept_tools: Final = tools_without_hosted_mcp(tools=sent_tools, model=model, route=route, drop_params=drop_params) + if len(kept_tools) == len(sent_tools): + return params + kept_params: Final = copy(params) + if kept_tools: + kept_params["tools"] = list(kept_tools) + return kept_params + dropped_keys: Final = ("tools",) if params.get("functions") else _TOOL_SELECTION_PARAMS + for key in dropped_keys: + _ = kept_params.pop(key, None) + return kept_params diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index f198d76f72a..01a4c508cf5 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -13,6 +13,7 @@ from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict import litellm +from litellm.anthropic_beta_headers_manager import get_unsupported_headers from litellm.constants import ( ANTHROPIC_BILLING_METADATA_PREFIX, ANTHROPIC_MIN_THINKING_BUDGET_TOKENS, @@ -56,6 +57,7 @@ from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import split_leading_system_run, ) from litellm.litellm_core_utils.reasoning_effort_utils import OPAQUE_VALUE, warn_non_string_reasoning_effort +from litellm.litellm_core_utils.tool_fields import tool_params_without_hosted_mcp from litellm.llms.base_llm.base_utils import ( ResponseFormatSchema, as_json_object, @@ -1315,6 +1317,11 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): *(tool for tool in tools[first_search_tool_index + 1 :] if not _is_openai_tool_search_tool(tool)), ] + def runs_hosted_mcp(self) -> bool: + return ANTHROPIC_BETA_HEADER_VALUES.MCP_CLIENT_2025_04_04.value not in get_unsupported_headers( + self._resolved_provider + ) + def _tools_param_with_openai_tool_search_mapped( self, non_default_params: Mapping[str, object], model: str, drop_params: bool ) -> Mapping[str, object]: @@ -1903,7 +1910,14 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): model: str, drop_params: bool, ) -> dict: - is_thinking_enabled: Final = self.is_thinking_enabled(non_default_params=non_default_params) + request_params: Final = tool_params_without_hosted_mcp( + non_default_params, + model=model, + route=self._resolved_provider, + drop_params=drop_params, + runs_hosted_mcp=self.runs_hosted_mcp, + ) + is_thinking_enabled: Final = self.is_thinking_enabled(non_default_params=request_params) # NB: ``map_openai_params`` deliberately does NOT sanitize tool names # here. Names are the *original* OpenAI names at this stage, and must @@ -1916,7 +1930,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): # via ``data = {**optional_params}`` and would 400 with # ``Extra inputs are not permitted``). - for param, value in non_default_params.items(): + for param, value in request_params.items(): if param == "max_tokens" or param == "max_completion_tokens": optional_params["max_tokens"] = value if isinstance(value, int) else max(1, int(round(value))) elif param == "tools": @@ -1930,8 +1944,8 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): optional_params["mcp_servers"] = mcp_servers elif param == "tool_choice" or param == "parallel_tool_calls": _tool_choice: AnthropicMessagesToolChoice | None = self._map_tool_choice( - tool_choice=non_default_params.get("tool_choice"), - parallel_tool_use=non_default_params.get("parallel_tool_calls"), + tool_choice=request_params.get("tool_choice"), + parallel_tool_use=request_params.get("parallel_tool_calls"), ) if _tool_choice is not None: @@ -1957,7 +1971,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): if self.uses_native_structured_output_for_schema( model, self._extract_json_schema_from_response_format(value), - REQUEST_PARAMS.validate_python(non_default_params), + REQUEST_PARAMS.validate_python(request_params), ): _output_format = self.map_response_format_to_anthropic_output_format(value) if _output_format is not None: @@ -1993,7 +2007,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): # `thinking={type: enabled, budget_tokens}` interface a # pre-4.6 model actually supports instead of forwarding a # shape the model will reject. - max_tokens = non_default_params.get("max_completion_tokens") or non_default_params.get("max_tokens") + max_tokens = request_params.get("max_completion_tokens") or request_params.get("max_tokens") legacy_thinking = AnthropicConfig._map_reasoning_effort( reasoning_effort="medium", model=model, @@ -2087,7 +2101,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): ## handle thinking tokens self.update_optional_params_with_thinking_tokens( - non_default_params=non_default_params, optional_params=optional_params + non_default_params=request_params, optional_params=optional_params ) return optional_params diff --git a/litellm/llms/anthropic/messages/adapters/handler.py b/litellm/llms/anthropic/messages/adapters/handler.py index d13fd4f47ff..31a4e9d9d73 100644 --- a/litellm/llms/anthropic/messages/adapters/handler.py +++ b/litellm/llms/anthropic/messages/adapters/handler.py @@ -15,6 +15,11 @@ from litellm._logging import verbose_logger from litellm.litellm_core_utils.asyncify import run_async_function from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject +from litellm.litellm_core_utils.tool_fields import ( + ROUTES_WITHOUT_MCP_SERVERS, + drops_mcp_servers, + tools_without_mcp_servers, +) from litellm.llms.anthropic.messages.adapters.transformation import ( AnthropicAdapter, ) @@ -610,6 +615,23 @@ class LiteLLMMessagesToCompletionTransformationHandler: Logging as LiteLLMLoggingObject, ) + provider_kwargs: Final = extra_kwargs or {} + custom_llm_provider: Final = provider_kwargs.get("custom_llm_provider") + target_provider: Final = custom_llm_provider if isinstance(custom_llm_provider, str) else None + drop_params: Final = drops_unsupported_params(provider_kwargs.get("drop_params")) + strips_mcp_servers: Final = target_provider in ROUTES_WITHOUT_MCP_SERVERS + request_tools: Final = ( + tools_without_mcp_servers( + provider_kwargs.get("mcp_servers"), + tools or (), + model=model, + route=target_provider, + drop_params=drops_mcp_servers(provider_kwargs), + ) + if strips_mcp_servers + else tools + ) + request_data: Final[dict[str, object]] = { "model": model, "messages": messages, @@ -626,10 +648,10 @@ class LiteLLMMessagesToCompletionTransformationHandler: request_data["temperature"] = temperature if thinking: request_data["thinking"] = thinking - if tool_choice: + if tool_choice and (request_tools or not tools): request_data["tool_choice"] = tool_choice - if tools: - request_data["tools"] = tools + if request_tools: + request_data["tools"] = request_tools if top_k is not None: request_data["top_k"] = top_k if top_p is not None: @@ -647,16 +669,15 @@ class LiteLLMMessagesToCompletionTransformationHandler: if "output_config" in extra_kwargs: request_data["output_config"] = extra_kwargs["output_config"] - custom_llm_provider: Final = extra_kwargs.get("custom_llm_provider") api_base: Final = extra_kwargs.get("api_base") ( openai_request, tool_name_mapping, ) = ANTHROPIC_ADAPTER.translate_completion_input_params_with_tool_mapping( request_data, - custom_llm_provider=custom_llm_provider if isinstance(custom_llm_provider, str) else None, + custom_llm_provider=target_provider, api_base=api_base if isinstance(api_base, str) else None, - drop_params=drops_unsupported_params(extra_kwargs.get("drop_params")), + drop_params=drop_params, ) if openai_request is None: @@ -684,7 +705,11 @@ class LiteLLMMessagesToCompletionTransformationHandler: # Maintainability: when adding a new Anthropic-only request param to # ``AnthropicMessagesRequestOptionalParams``, also extend # ``ANTHROPIC_ONLY_REQUEST_KEYS`` here so it doesn't silently leak. - excluded_keys: Final = ANTHROPIC_ONLY_REQUEST_KEYS | CONVERTED_STREAM_KEYS | {"anthropic_messages"} + excluded_keys: Final = ( + ANTHROPIC_ONLY_REQUEST_KEYS + | CONVERTED_STREAM_KEYS + | ({"anthropic_messages", "mcp_servers"} if strips_mcp_servers else {"anthropic_messages"}) + ) # NOTE: extra_kwargs was already coerced from None to {} at the top of # this method (line ~220). It is guaranteed to be a dict here. for key, value in extra_kwargs.items(): diff --git a/litellm/llms/anthropic/messages/adapters/streaming_iterator.py b/litellm/llms/anthropic/messages/adapters/streaming_iterator.py index 61d23c2de0d..c41deb00f1d 100644 --- a/litellm/llms/anthropic/messages/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/messages/adapters/streaming_iterator.py @@ -493,10 +493,23 @@ class _CombinedChunkSplitter: for position, block in enumerate(blocks) ) + @staticmethod + def _without_sibling_finish(chunk: "ModelResponseStream") -> "ModelResponseStream": + choices: Final = _optional_attr_sequence(chunk, "choices") + if ( + len(choices) != 1 + or _optional_attr(choices[0], "index") in (None, 0) + or _optional_attr(choices[0], "finish_reason") is None + ): + return chunk + unfinished: Final = copy.deepcopy(chunk) + unfinished.choices[0].finish_reason = None + return unfinished + def _expand(self, chunk: "ModelResponseStream") -> "tuple[ModelResponseStream, ...]": split_chunks: Final = tuple( split_chunk - for combined_chunk in self._split(chunk) + for combined_chunk in self._split(self._without_sibling_finish(chunk)) for split_chunk in self._split_by_payload_kind(combined_chunk) # comprehension-ok: flattens nested splits ) block_chunks: Final = tuple(itertools.chain.from_iterable(map(self._split_thinking_blocks, split_chunks))) diff --git a/litellm/llms/anthropic/messages/responses_adapters/handler.py b/litellm/llms/anthropic/messages/responses_adapters/handler.py index 08b00dd5fcf..8b638f2c03a 100644 --- a/litellm/llms/anthropic/messages/responses_adapters/handler.py +++ b/litellm/llms/anthropic/messages/responses_adapters/handler.py @@ -24,6 +24,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( prompt_cache_breakpoint_with_readable_mode, with_prompt_cache_breakpoint, ) +from litellm.litellm_core_utils.tool_fields import drops_mcp_servers, tools_without_mcp_servers from litellm.responses.utils import auto_reasoning_summary from litellm.types.integrations.custom_logger import CONVERTED_STREAM_KEYS from litellm.types.llms.anthropic import ( @@ -272,6 +273,14 @@ def _build_responses_kwargs( """ forwarded_kwargs: Final = _forwarded_kwargs(extra_kwargs) reject_malformed_tool_result_text(messages, model, forwarded_kwargs) + target_provider: Final = forwarded_kwargs.get("custom_llm_provider") + request_tools: Final = tools_without_mcp_servers( + forwarded_kwargs.get("mcp_servers"), + tools or (), + model=model, + route=target_provider if isinstance(target_provider, str) else "this route", + drop_params=drops_mcp_servers(forwarded_kwargs), + ) block_rewrite: Final = _prompt_cache_block_rewrite(model, forwarded_kwargs) request_messages: Final[AnthropicRequestMessages] = ( [ @@ -304,10 +313,10 @@ def _build_responses_kwargs( request_data["temperature"] = temperature if thinking: request_data["thinking"] = thinking - if tool_choice: + if tool_choice and (request_tools or not tools): request_data["tool_choice"] = tool_choice - if tools: - request_data["tools"] = tools + if request_tools: + request_data["tools"] = list(request_tools) if top_p is not None: request_data["top_p"] = top_p if output_format: @@ -347,7 +356,7 @@ def _build_responses_kwargs( responses_kwargs["stream"] = True # Forward litellm-specific kwargs (api_key, api_base, logging obj, etc.) - excluded: Final = CONVERTED_STREAM_KEYS | {"anthropic_messages"} + excluded: Final = CONVERTED_STREAM_KEYS | {"anthropic_messages", "mcp_servers"} for key, value in forwarded_kwargs.items(): if key == "litellm_logging_obj" and value is not None: from litellm.litellm_core_utils.litellm_logging import ( diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 21af97e523e..9a345ccdc9d 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -52,6 +52,7 @@ from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import parts_of, ) from litellm.litellm_core_utils.reasoning_effort_utils import OPAQUE_VALUE, warn_dropped_reasoning +from litellm.litellm_core_utils.tool_fields import tool_params_without_hosted_mcp from litellm.llms.anthropic.chat.transformation import ( DROP_UNSUPPORTED_ADAPTIVE_THINKING_WARNING, DROP_UNSUPPORTED_OUTPUT_CONFIG_WARNING, @@ -1068,20 +1069,23 @@ class AmazonConverseConfig(BaseConfig): drop_params: bool, model_info: Mapping[str, object] | None = None, ) -> dict: - is_thinking_enabled: Final = self.is_thinking_enabled(non_default_params) + request_params: Final = tool_params_without_hosted_mcp( + non_default_params, model=model, route="bedrock", drop_params=drop_params + ) + is_thinking_enabled: Final = self.is_thinking_enabled(request_params) base_model: Final = BedrockModelInfo.get_base_model(model) drop_thinking_param: Final = self._is_deepseek_model(model=model, base_model=base_model) drop_reasoning_effort_param: Final = self._model_rejects_reasoning_effort_param( model=model, base_model=base_model ) - for param, value in non_default_params.items(): + for param, value in request_params.items(): if param == "response_format" and isinstance(value, dict): optional_params = self._translate_response_format_param( value=value, model=model, optional_params=optional_params, - non_default_params=non_default_params, + non_default_params=request_params, is_thinking_enabled=is_thinking_enabled, model_info=model_info, ) @@ -1135,7 +1139,7 @@ class AmazonConverseConfig(BaseConfig): ), ), model=model, - non_default_params=non_default_params, + non_default_params=request_params, optional_params=optional_params, ) if param == "tool_choice": @@ -1160,7 +1164,7 @@ class AmazonConverseConfig(BaseConfig): and value.get("type") == "adaptive" and not AnthropicConfig.is_adaptive_thinking_model(model, "bedrock") ): - max_tokens = non_default_params.get("max_completion_tokens") or non_default_params.get("max_tokens") + max_tokens = request_params.get("max_completion_tokens") or request_params.get("max_tokens") legacy_thinking = AnthropicConfig.map_reasoning_effort( reasoning_effort="medium", model=model, @@ -1230,7 +1234,7 @@ class AmazonConverseConfig(BaseConfig): # Nova 2 handles token budgeting differently through reasoningConfig if "gpt-oss" not in model and not self._is_nova_2_model(model): self.update_optional_params_with_thinking_tokens( - non_default_params=non_default_params, optional_params=optional_params + non_default_params=request_params, optional_params=optional_params ) final_is_thinking_enabled: Final = self.is_thinking_enabled(optional_params) diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude_transformation.py index e32df0bbdbb..e150791b09e 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude_transformation.py @@ -1,11 +1,13 @@ -from collections.abc import Mapping +from collections.abc import Callable, Mapping from typing import TYPE_CHECKING, Any, Final import httpx from pydantic import TypeAdapter +import litellm from litellm.anthropic_beta_gated_fields import strip_beta_gated_request_body from litellm.anthropic_beta_headers_manager import filter_and_transform_beta_headers +from litellm.constants import DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import object_mapping from litellm.litellm_core_utils.prompt_templates.factory import ( convert_to_anthropic_image_obj, @@ -25,7 +27,9 @@ from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation ) from litellm.llms.bedrock.common_utils import ( BEDROCK_INVOKE_TRACE_FIELD, + BedrockModelInfo, apply_bedrock_invoke_structured_output, + bedrock_capability_model, bedrock_chat_uses_native_structured_output, bedrock_cleaned_output_schema, bedrock_supports_tool_search, @@ -35,6 +39,7 @@ from litellm.llms.bedrock.common_utils import ( normalize_custom_field_on_tools, normalize_tool_input_schema_types_for_bedrock_invoke, requests_opt_in_native_output, + strip_bedrock_routing_prefix, strip_unsupported_bedrock_invoke_output_config_keys, tools_without_eager_input_streaming, ) @@ -59,6 +64,19 @@ _NON_DEFAULT_PARAMS: Final = TypeAdapter(Mapping[str, object]) _LITELLM_PARAMS: Final = TypeAdapter(Mapping[str, object]) +def _anthropic_config_max_tokens() -> int | None: + return litellm.AnthropicConfig.max_tokens + + +def _mapped_max_tokens(model: str) -> int | None: + row: Final = object_mapping(litellm.model_cost.get(model)) or {} # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # cost map rows are untyped JSON + match row.get("max_output_tokens"): + case int() as max_tokens: + return max_tokens + case _: + return None + + class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): """ Reference: @@ -79,9 +97,18 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): anthropic_version: str = "bedrock-2023-05-31" - def __init__(self, model_info: Mapping[str, object] | None = None) -> None: + def __init__( + self, + model_info: Mapping[str, object] | None = None, + read_configured_max_tokens: Callable[[], int | None] = _anthropic_config_max_tokens, + ) -> None: super().__init__() self._deployment_model_info: Final = model_info + self._read_configured_max_tokens: Final = read_configured_max_tokens + + @classmethod + def get_config(cls, *, model: str | None = None) -> dict[str, object]: # mutable-ok: overrides BaseConfig's dict + return {key: value for key, value in super().get_config(model=model).items() if key != "max_tokens"} @property def custom_llm_provider(self) -> str | None: @@ -241,6 +268,36 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): headers=headers, ) + def _default_max_tokens_models(self, model: str, litellm_params: Mapping[str, object]) -> tuple[str, ...]: + return ( + BedrockModelInfo.get_base_model( + bedrock_capability_model(strip_bedrock_routing_prefix(model), litellm_params.get("base_model")) + ), + ) + + def _default_max_tokens(self, model: str, litellm_params: Mapping[str, object]) -> int: + configured_max_tokens: Final = self._read_configured_max_tokens() + match deployment_model_info(litellm_params).get("max_output_tokens"): + case int() as deployment_max_tokens: + return ( + deployment_max_tokens + if configured_max_tokens is None + else min(deployment_max_tokens, configured_max_tokens) + ) + case _ if configured_max_tokens is not None: + return configured_max_tokens + case _: + return next( + ( + max_tokens + for max_tokens in map( + _mapped_max_tokens, self._default_max_tokens_models(model, litellm_params) + ) + if max_tokens is not None + ), + DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS, + ) + def _build_bedrock_anthropic_request_base( self, model: str, @@ -257,13 +314,21 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): model=model, output_config=filtered_params["output_config"], ) - filtered_params = self._normalize_bedrock_tool_search_tools(filtered_params) + normalized_params: Final = self._normalize_bedrock_tool_search_tools(filtered_params) + default_max_tokens: Final = ( + None + if "max_tokens" in normalized_params + else self._default_max_tokens(model, _LITELLM_PARAMS.validate_python(litellm_params)) + ) + request_params: Final = ( + normalized_params if default_max_tokens is None else {**normalized_params, "max_tokens": default_max_tokens} + ) anthropic_request: Final = AnthropicConfig.transform_request( self, model=model, messages=messages, - optional_params=filtered_params, + optional_params=request_params, litellm_params=litellm_params, headers=headers, ) @@ -377,7 +442,7 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): "data": image_chunk["data"], } - def _normalize_bedrock_tool_search_tools(self, optional_params: dict) -> dict: + def _normalize_bedrock_tool_search_tools(self, optional_params: dict) -> dict[str, object]: """ Convert tool search entries to the format supported by the Bedrock Invoke API. """ diff --git a/litellm/llms/bedrock/chat/mantle/transformation.py b/litellm/llms/bedrock/chat/mantle/transformation.py index b5a352bab48..cbde6f918bc 100644 --- a/litellm/llms/bedrock/chat/mantle/transformation.py +++ b/litellm/llms/bedrock/chat/mantle/transformation.py @@ -54,6 +54,9 @@ class AmazonMantleConfig(AmazonAnthropicClaudeConfig): region=region, ) + def _default_max_tokens_models(self, model: str, litellm_params: Mapping[str, object]) -> tuple[str, ...]: + return (f"bedrock_mantle/{model}", *super()._default_max_tokens_models(model, litellm_params)) + def validate_environment( self, headers: dict, diff --git a/litellm/llms/bedrock/claude_platform/transformation.py b/litellm/llms/bedrock/claude_platform/transformation.py index a5ecc37d394..6efc4b506e8 100644 --- a/litellm/llms/bedrock/claude_platform/transformation.py +++ b/litellm/llms/bedrock/claude_platform/transformation.py @@ -30,6 +30,9 @@ class BedrockClaudePlatformConfig(BedrockClaudePlatformMixin, AnthropicConfig): def should_strip_billing_metadata(self) -> bool: return True + def runs_hosted_mcp(self) -> bool: + return True + def validate_environment( self, headers: dict, diff --git a/litellm/llms/cohere/chat/v2_transformation.py b/litellm/llms/cohere/chat/v2_transformation.py index cf033e10b41..141a174e558 100644 --- a/litellm/llms/cohere/chat/v2_transformation.py +++ b/litellm/llms/cohere/chat/v2_transformation.py @@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx import litellm +from litellm.litellm_core_utils.core_helpers import map_finish_reason from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig from litellm.types.llms.cohere import CohereV2ChatResponse @@ -248,6 +249,9 @@ class CohereV2ChatConfig(OpenAIGPTConfig): prompt_tokens: Final = token_usage.get("input_tokens", 0) completion_tokens: Final = token_usage.get("output_tokens", 0) + model_response.choices[0].finish_reason = map_finish_reason( # rebind-ok: ModelResponse is filled in place + cohere_v2_chat_response.get("finish_reason", "COMPLETE") + ) model_response.created = int(time.time()) model_response.model = model usage: Final = Usage( diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index 60081c38073..b5393fa1603 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -4,6 +4,7 @@ Translates from OpenAI's `/v1/chat/completions` to Databricks' `/chat/completion import os from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence +from itertools import chain, takewhile from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, cast, overload @@ -12,6 +13,7 @@ from pydantic import BaseModel, ConfigDict, TypeAdapter from litellm.anthropic_beta_gated_fields import log_beta_gated_strips, strip_beta_gated_thinking from litellm.constants import RESPONSE_FORMAT_TOOL_NAME +from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( handle_invalid_parallel_tool_calls, should_convert_tool_call_to_json_mode, @@ -23,6 +25,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( strip_name_from_message, strip_prefill_prefix_field, ) +from litellm.litellm_core_utils.tool_fields import tool_params_without_hosted_mcp from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.base_utils import ( @@ -65,7 +68,7 @@ from ...anthropic.chat.transformation import ( tools_without_unsupported_tool_search, ) from ...openai_like.chat.transformation import OpenAILikeChatConfig -from ..common_utils import DatabricksBase, DatabricksException +from ..common_utils import DATABRICKS_HOSTED_MODEL_PREFIXES, DatabricksBase, DatabricksException _RESPONSE_FORMAT_ADAPTER: ( # mutable-ok: mirrors the dict return contract of get_json_schema_from_pydantic_object Final[TypeAdapter[dict[str, object] | None]] @@ -81,15 +84,139 @@ def _is_bare_assistant_message(message_dict: Mapping[str, object]) -> bool: ) +_CONTENT_BLOCKS: Final = TypeAdapter(tuple[object, ...]) + + +def _content_blocks(message: object) -> tuple[object, ...]: + match message.model_dump(exclude_none=True) if isinstance(message, BaseModel) else message: + case {"content": list() as content}: + return _CONTENT_BLOCKS.validate_python(content) + case _: + return () + + +def _plain_text(block: object) -> str | None: + match block: + case {"type": "text", "text": str() as text} if block == {"type": "text", "text": text}: + return text + case _: + return None + + +def _plain_texts(message: Mapping[str, object]) -> tuple[str, ...]: + blocks: Final = _content_blocks(message) + texts: Final = tuple(text for block in blocks if (text := _plain_text(block)) is not None) + return texts if len(texts) == len(blocks) else () + + def _collapse_single_text_block(message: Mapping[str, object]) -> Mapping[str, object]: - content: Final = message.get("content") - match content: - case [{"type": "text", "text": str() as text}] if content == [{"type": "text", "text": text}]: + match _plain_texts(message): + case (text,): return {**message, "content": text} case _: return message +def _without_text_cache_control(block: object) -> object: + match block: + case {"type": "text", "text": str() as text, "cache_control": object() as cache_control} if block == { + "type": "text", + "text": text, + "cache_control": cache_control, + }: + return {"type": "text", "text": text} + case _: + return block + + +def _without_cache_control(message: AllMessageValues) -> AllMessageValues: + uncached: Final = {key: value for key, value in message.items() if key != "cache_control"} + match content := uncached.get("content"): + case list(): + blocks: Final = [_without_text_cache_control(block) for block in _CONTENT_BLOCKS.validate_python(content)] + return cast(AllMessageValues, {**uncached, "content": blocks}) # cast-ok: spread keeps the shape + case _: + return cast(AllMessageValues, uncached) # cast-ok: only cache_control was left out + + +def _has_string_content(message: Mapping[str, object]) -> bool: + return isinstance(message.get("content"), str) + + +def _is_system(message: Mapping[str, object]) -> bool: + return message.get("role") == "system" + + +def _json_mode_message(message: AllMessageValues, drop_params: bool) -> AllMessageValues: + candidate: Final = _without_cache_control(message) if drop_params else message + texts: Final = _plain_texts(candidate) + if "cache_control" in candidate or not (texts or _has_string_content(candidate)): + return message + if sum(1 for text in texts if text.strip()) < 2: + return candidate + return cast(AllMessageValues, {**candidate, "content": "".join(texts)}) # cast-ok: spread keeps the shape + + +def _is_dropped_on_send(message: Mapping[str, object]) -> bool: + sanitized: Final = {**message} + _sanitize_empty_content(sanitized) + return _is_bare_assistant_message(sanitized) + + +def _is_system_or_dropped(message: Mapping[str, object]) -> bool: + return _is_system(message) or _is_dropped_on_send(message) + + +def _json_checked_start(messages: Sequence[AllMessageValues]) -> int: + trailing: Final = tuple(takewhile(_is_system_or_dropped, reversed(messages))) + last_sent_offset: Final = 0 if any(_is_system(message) for message in trailing) else 1 + return max(len(messages) - len(trailing) - last_sent_offset, 0) + + +def _json_object_ready_messages( + model: str, + messages: list[AllMessageValues], # mutable-ok: inherited transform_request contract + response_format: object, + drop_params: bool, +) -> list[AllMessageValues]: # mutable-ok: inherited transform_request contract + match response_format: + case {"type": "json_object"} if "claude" not in model and model.startswith(DATABRICKS_HOSTED_MODEL_PREFIXES): + checked_start: Final = _json_checked_start(messages) + return [ + *messages[:checked_start], + *(_json_mode_message(message, drop_params) for message in messages[checked_start:]), + ] + case _: + return messages + + +def _is_empty_text_block(block: object) -> bool: + match block: + case {"type": "text", "text": str() as text}: + return not text.strip() + case {"type": "text", "text": object() as text}: + return not text + case {"type": "text"}: + return True + case _: + return False + + +def _non_string_text_type(block: object) -> str | None: + match block: + case {"type": "text", "text": str()}: + return None + case {"type": "text", "text": object() as text} if text: + return type(text).__name__ + case _: + return None + + +def _first_non_string_text_type(messages: Sequence[AllMessageValues]) -> str | None: + blocks: Final = chain.from_iterable(map(_content_blocks, messages)) + return next((type_name for block in blocks if (type_name := _non_string_text_type(block))), None) + + def _sanitize_empty_content(message_dict: dict[str, object]) -> None: """ Remove or filter content so empty text blocks are not sent. @@ -108,9 +235,7 @@ def _sanitize_empty_content(message_dict: dict[str, object]) -> None: message_dict.pop("content") return filtered: Final = [ - block - for block in content - if not (isinstance(block, dict) and block.get("type") == "text" and not (block.get("text") or "").strip()) + block for block in _CONTENT_BLOCKS.validate_python(content) if not _is_empty_text_block(block) ] if not filtered: message_dict.pop("content") @@ -343,7 +468,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): use_ai_gateway: Final = self.uses_ai_gateway(model) return join_api_base(self._get_api_base(api_base, use_ai_gateway=use_ai_gateway), "chat/completions") - def get_supported_openai_params(self, model: str | None = None) -> list: + def get_supported_openai_params(self, model: str | None = None) -> list[str]: return [ "stream", "stop", @@ -449,8 +574,13 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): drop_params: bool, replace_max_completion_tokens_with_max_tokens: bool = True, ) -> dict: - is_thinking_enabled: Final = self.is_thinking_enabled(non_default_params) - mapped_params: Final = super().map_openai_params(non_default_params, optional_params, model, drop_params) + request_params: Final = ( + tool_params_without_hosted_mcp(non_default_params, model=model, route="databricks", drop_params=drop_params) + if "claude" in model + else non_default_params + ) + is_thinking_enabled: Final = self.is_thinking_enabled(request_params) + mapped_params: Final = super().map_openai_params(request_params, optional_params, model, drop_params) if "claude" in model: AnthropicConfig.translate_legacy_thinking_for_adaptive_model( model=model, optional_params=mapped_params, custom_llm_provider="databricks", drop_params=drop_params @@ -459,16 +589,16 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): mapped_params["tools"] = self._map_openai_to_dbrx_tool( model=model, tools=mapped_params["tools"], drop_params=drop_params ) - if "max_completion_tokens" in non_default_params and replace_max_completion_tokens_with_max_tokens: - mapped_params["max_tokens"] = non_default_params[ + if "max_completion_tokens" in request_params and replace_max_completion_tokens_with_max_tokens: + mapped_params["max_tokens"] = request_params[ "max_completion_tokens" ] # most openai-compatible providers support 'max_tokens' not 'max_completion_tokens' mapped_params.pop("max_completion_tokens", None) - if "response_format" in non_default_params and "claude" in model: + if "response_format" in request_params and "claude" in model: _tool: Final = self.map_response_format_to_databricks_tool( model, - non_default_params["response_format"], + request_params["response_format"], mapped_params, is_thinking_enabled, ) @@ -486,8 +616,8 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): "response_format", None ) # unsupported for claude models - if json_schema -> convert to tool call - if "reasoning_effort" in non_default_params and self._uses_anthropic_thinking_param(model): - reasoning_effort_value: Final = non_default_params.get("reasoning_effort") + if "reasoning_effort" in request_params and self._uses_anthropic_thinking_param(model): + reasoning_effort_value: Final = request_params.get("reasoning_effort") mapped_thinking: Final = AnthropicConfig._map_reasoning_effort( reasoning_effort=reasoning_effort_value, model=model, @@ -513,7 +643,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): optional_params.pop("reasoning_effort", None) ## handle thinking tokens self.update_optional_params_with_thinking_tokens( - non_default_params=non_default_params, optional_params=mapped_params + non_default_params=request_params, optional_params=mapped_params ) gated_thinking: Final = strip_beta_gated_thinking( mapped_params.get("thinking"), # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # mapped_params is the base config's bare dict @@ -535,6 +665,35 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): return False + def transform_request( + self, + model: str, + messages: list[AllMessageValues], # mutable-ok: inherited contract + optional_params: dict[str, object], # mutable-ok: inherited contract + litellm_params: dict[str, object], # mutable-ok: inherited contract + headers: dict[str, object], # mutable-ok: inherited contract + ) -> dict[str, object]: # mutable-ok: inherited contract + match _first_non_string_text_type(messages): + case str() as type_name: + raise DatabricksException( + status_code=400, + message=f"A text content block needs a string 'text', got {type_name}", + ) + case None: + pass + return super().transform_request( + model=model, + messages=_json_object_ready_messages( + model, + messages, + optional_params.get("response_format"), + drop_params=drops_unsupported_params(litellm_params.get("drop_params")), + ), + optional_params=optional_params, + litellm_params=litellm_params, + headers=headers, + ) + @overload def _transform_messages( self, messages: list[AllMessageValues], model: str, is_async: Literal[True] diff --git a/litellm/llms/databricks/common_utils.py b/litellm/llms/databricks/common_utils.py index b1562bef6b8..e8313d9a003 100644 --- a/litellm/llms/databricks/common_utils.py +++ b/litellm/llms/databricks/common_utils.py @@ -56,6 +56,7 @@ DATABRICKS_WORKSPACE_ROUTE_PREFIXES: Final = ( DATABRICKS_AI_GATEWAY_MESSAGES_PATH, ) DATABRICKS_NATIVE_MESSAGES_MODEL_PREFIXES: Final = ("databricks-claude-", "system.ai.claude-") +DATABRICKS_HOSTED_MODEL_PREFIXES: Final = ("databricks-", "system.ai.") DATABRICKS_ENDPOINT_PATHS: Final = MappingProxyType( { "chat_completions": "/chat/completions", diff --git a/litellm/llms/deepseek/chat/transformation.py b/litellm/llms/deepseek/chat/transformation.py index 5b7c2a2c410..9c222f03c8c 100644 --- a/litellm/llms/deepseek/chat/transformation.py +++ b/litellm/llms/deepseek/chat/transformation.py @@ -12,6 +12,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( extract_search_results_text, ) from litellm.litellm_core_utils.reasoning_effort_utils import OPAQUE_VALUE, warn_dropped_reasoning +from litellm.litellm_core_utils.tool_fields import tool_params_without_hosted_mcp from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues @@ -33,17 +34,15 @@ def _warn_on_dropped_budget(thinking: object, model: str) -> None: class DeepSeekChatConfig(OpenAIGPTConfig): - def get_supported_openai_params(self, model: str) -> list: + def get_supported_openai_params(self, model: str) -> list[str]: """ DeepSeek reasoner models support thinking parameter. """ - params: Final = super().get_supported_openai_params(model) - params.extend(["thinking", "reasoning_effort"]) - return params + return [*super().get_supported_openai_params(model), "thinking", "reasoning_effort"] def map_openai_params( self, - non_default_params: dict, + non_default_params: dict[str, object], optional_params: dict, model: str, drop_params: bool, @@ -59,7 +58,12 @@ class DeepSeekChatConfig(OpenAIGPTConfig): Reference: https://api-docs.deepseek.com/guides/thinking_mode """ # Let parent handle standard params first - optional_params = super().map_openai_params(non_default_params, optional_params, model, drop_params) + optional_params = super().map_openai_params( + tool_params_without_hosted_mcp(non_default_params, model=model, route="deepseek", drop_params=drop_params), + optional_params, + model, + drop_params, + ) # Pop thinking/reasoning_effort from optional_params first (parent may have added them) # Then re-add only if valid for DeepSeek @@ -265,6 +269,7 @@ class DeepSeekChatConfig(OpenAIGPTConfig): tools: Final = optional_params.get("tools") if not isinstance(tools, list) or not tools: return optional_params + sent_tools: Final[tuple[object, ...]] = tuple(tools) def _is_function_tool(tool: object) -> bool: return isinstance(tool, dict) and tool.get("type") == "function" @@ -289,14 +294,14 @@ class DeepSeekChatConfig(OpenAIGPTConfig): name: Final = function.get("name") return isinstance(name, str) and name in function_tool_names - function_tools: Final = [tool for tool in tools if _is_function_tool(tool)] - if len(function_tools) == len(tools): + function_tools: Final = [tool for tool in sent_tools if _is_function_tool(tool)] + if len(function_tools) == len(sent_tools): return optional_params dropped_types: Final = sorted( { str(tool.get("type")) if isinstance(tool, dict) else type(tool).__name__ - for tool in tools + for tool in sent_tools if not _is_function_tool(tool) } ) diff --git a/litellm/llms/gemini/chat/transformation.py b/litellm/llms/gemini/chat/transformation.py index 936919af83d..1539f13455f 100644 --- a/litellm/llms/gemini/chat/transformation.py +++ b/litellm/llms/gemini/chat/transformation.py @@ -1,5 +1,5 @@ from collections.abc import Mapping -from typing import Final, cast +from typing import ClassVar, Final, cast import litellm from litellm.litellm_core_utils.prompt_templates.factory import ( @@ -49,6 +49,8 @@ class GoogleAIStudioGeminiConfig(VertexGeminiConfig): Note: Please make sure to modify the default parameters as required for your use case. """ + _route: ClassVar[str] = "gemini" + temperature: float | None = None max_output_tokens: int | None = None top_p: float | None = None diff --git a/litellm/llms/gigachat/chat/transformation.py b/litellm/llms/gigachat/chat/transformation.py index 81a186eadf4..7bc783b960f 100644 --- a/litellm/llms/gigachat/chat/transformation.py +++ b/litellm/llms/gigachat/chat/transformation.py @@ -17,6 +17,7 @@ import httpx from pydantic import TypeAdapter from litellm._logging import verbose_logger +from litellm.litellm_core_utils.tool_fields import tool_params_without_hosted_mcp from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.llms.gigachat.utils import convert_usage, get_api_base @@ -169,7 +170,10 @@ class GigaChatConfig(BaseConfig): drop_params: bool, ) -> dict: # mutable-ok: base class contract returns dict """Map OpenAI parameters to GigaChat parameters.""" - for param, value in non_default_params.items(): + request_params: Final = tool_params_without_hosted_mcp( + dict(non_default_params), model=model, route="gigachat", drop_params=drop_params + ) + for param, value in request_params.items(): if param == "stream": optional_params["stream"] = value elif param == "temperature": diff --git a/litellm/llms/ollama/chat/transformation.py b/litellm/llms/ollama/chat/transformation.py index 168b562b358..59020b058a2 100644 --- a/litellm/llms/ollama/chat/transformation.py +++ b/litellm/llms/ollama/chat/transformation.py @@ -4,7 +4,7 @@ from collections.abc import AsyncIterator, Iterator, Mapping from typing import TYPE_CHECKING, Any, Final, cast from httpx._models import Headers, Response -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter import litellm from litellm._uuid import uuid @@ -15,6 +15,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( extract_reasoning_content, ) from litellm.litellm_core_utils.reasoning_effort_utils import OPAQUE_VALUE +from litellm.litellm_core_utils.tool_fields import tool_params_without_hosted_mcp from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException @@ -40,6 +41,8 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_REQUEST_PARAMS: Final = TypeAdapter(dict[str, object]) + class OllamaChatConfig(BaseConfig): """ @@ -152,7 +155,13 @@ class OllamaChatConfig(BaseConfig): model: str, drop_params: bool, ) -> dict: - for param, value in non_default_params.items(): + request_params: Final = tool_params_without_hosted_mcp( + _REQUEST_PARAMS.validate_python(non_default_params), + model=model, + route="ollama_chat", + drop_params=drop_params, + ) + for param, value in request_params.items(): if param == "max_tokens" or param == "max_completion_tokens": optional_params["num_predict"] = value if param == "stream": diff --git a/litellm/llms/opencode/chat/transformation.py b/litellm/llms/opencode/chat/transformation.py index dcb30b5835c..7e16f93d7cb 100644 --- a/litellm/llms/opencode/chat/transformation.py +++ b/litellm/llms/opencode/chat/transformation.py @@ -1,5 +1,5 @@ from collections.abc import AsyncIterator, Iterator, Mapping, Sequence -from typing import TYPE_CHECKING, Final +from typing import TYPE_CHECKING, ClassVar, Final import httpx from pydantic import TypeAdapter @@ -159,6 +159,7 @@ class OpenCodeGoMessagesChatConfig(_GoReportedCostMixin, AnthropicConfig): class OpenCodeGoGeminiChatConfig(_GoReportedCostMixin, GoogleAIStudioGeminiConfig): go_usage_key: str = "usageMetadata" + _route: ClassVar[str] = "opencode_go" @property def supports_stream_param_in_request_body(self) -> bool: diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index d0710339464..9f5bec37ac0 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -9,7 +9,7 @@ from collections.abc import Awaitable, Callable, Mapping, Sequence from copy import deepcopy from functools import lru_cache, partial from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, Union, cast, get_args +from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, Protocol, Union, cast, get_args import httpx from pydantic import JsonValue, TypeAdapter @@ -34,6 +34,7 @@ from litellm.litellm_core_utils.json_fragment_accumulator import JSONFragmentAcc from litellm.litellm_core_utils.model_param_helper import embeds_thought_signature_in_tool_call_id from litellm.litellm_core_utils.prompt_templates.factory import encode_tool_call_id_with_signature from litellm.litellm_core_utils.reasoning_effort_utils import warn_dropped_reasoning +from litellm.litellm_core_utils.tool_fields import tool_params_without_hosted_mcp from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -614,6 +615,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): Note: Please make sure to modify the default parameters as required for your use case. """ + _route: ClassVar[str] = "vertex_ai" + temperature: float | None = None max_output_tokens: int | None = None top_p: float | None = None @@ -1667,9 +1670,12 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): drop_params: bool, allow_integer_enums: bool = True, ) -> dict: - self._apply_include_server_side_tool_invocations(non_default_params, optional_params) + request_params: Final = tool_params_without_hosted_mcp( + non_default_params, model=model, route=self._route, drop_params=drop_params + ) + self._apply_include_server_side_tool_invocations(request_params, optional_params) supports_sampling_params: Final = self.supports_sampling_params(model) - for param, value in non_default_params.items(): + for param, value in request_params.items(): if param in ("temperature", "top_p", "top_k") and not supports_sampling_params: continue if param == "temperature": @@ -1733,7 +1739,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): if _tool_choice_value is not None: optional_params["tool_choice"] = _tool_choice_value elif param == "parallel_tool_calls": - tools_list = non_default_params.get("tools", non_default_params.get("functions")) + tools_list = request_params.get("tools", request_params.get("functions")) num_tools = len(tools_list) if isinstance(tools_list, list) else 0 # Gemini does not support parallel_tool_calls=False with multiple # tools. Drop the param instead of failing — Responses API clients diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 02bf7dedcef..e6f65d6b416 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -13478,6 +13478,10 @@ "input_cost_per_token": 2e-07, "output_cost_per_token": 5e-07, "litellm_provider": "azure_ai", + "token_count_framing": 149, + "token_count_tool_framing": 302, + "token_count_responses_framing": 145, + "token_count_responses_tool_framing": 100, "max_input_tokens": 128000, "max_output_tokens": 128000, "mode": "chat", @@ -54457,6 +54461,8 @@ "input_cost_per_token_above_200k_tokens_batches": 2e-06, "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", + "token_count_framing": 185, + "token_count_tool_framing": 214, "max_input_tokens": 1000000, "max_output_tokens": 1000000, "mode": "chat", @@ -62640,7 +62646,9 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "token_count_framing": 23, + "token_count_tool_framing": 12 }, "deepseek/deepseek-v4-flash-vision-exp": { "cache_creation_input_token_cost": 0.0, @@ -62746,7 +62754,9 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": false + "supports_vision": false, + "token_count_framing": 76, + "token_count_tool_framing": 75 }, "tencent/deepseek-v4-pro": { "cache_creation_input_token_cost": 0.0, @@ -64032,6 +64042,7 @@ "input_cost_per_token": 3e-07, "input_cost_per_token_priority": 3.75e-07, "litellm_provider": "fireworks_ai", + "token_count_framing": 23, "max_input_tokens": 1048576, "max_output_tokens": 393216, "mode": "chat", @@ -64133,6 +64144,7 @@ "input_cost_per_token": 3e-07, "input_cost_per_token_priority": 3.75e-07, "litellm_provider": "fireworks_ai", + "token_count_framing": 23, "max_input_tokens": 1048576, "max_output_tokens": 393216, "mode": "chat", diff --git a/litellm/proxy/client/cli/commands/agents.py b/litellm/proxy/client/cli/commands/agents.py index 12d97ef6cce..b70432d1e72 100644 --- a/litellm/proxy/client/cli/commands/agents.py +++ b/litellm/proxy/client/cli/commands/agents.py @@ -770,9 +770,10 @@ def codex_model_catalog( A model Codex does not know gets the fallback entry, with the same base instructions Codex itself uses so the agent never runs without a system prompt, and offers the reasoning efforts `groups` gives its model group - that the installed Codex, at `client_version`, can parse. Apart from aliases, the catalog is - the one `/v1/models?client_version=...` serves, built from the installed Codex's own entries - instead of the proxy's vendored ones. + that the installed Codex, at `client_version`, can parse. Apart from aliases, and the window the + deployments' `model_info.max_input_tokens` gives a model Codex knows, the catalog is the one + `/v1/models?client_version=...` serves, built from the installed Codex's own entries instead of + the proxy's vendored ones. """ rows: Final = tuple(_codex_row(m, upstreams.get(m.id), groups) for m in models) body: Final = codex_models_response_json( diff --git a/litellm/proxy/common_utils/codex_model_catalog.py b/litellm/proxy/common_utils/codex_model_catalog.py index 63fb27d62cf..7b48aa07365 100644 --- a/litellm/proxy/common_utils/codex_model_catalog.py +++ b/litellm/proxy/common_utils/codex_model_catalog.py @@ -7,8 +7,10 @@ never the OpenAI list shape. This module builds that body from the rows ``/v1/mo lists: a model Codex knows keeps Codex's own entry for it, any other model gets the fallback entry Codex itself uses for an unknown slug, and ``model_info.service_tiers`` becomes the entry's ``service_tiers``, which Codex turns into slash commands (``/ultrafast``) that send ``service_tier`` -upstream. A fallback entry offers the model group's ``supported_reasoning_efforts`` as its reasoning -levels, trimmed to the ones the requesting Codex release can parse. +upstream. A ``model_info.max_input_tokens`` the operator sets replaces a known model's +``context_window``, up to the most Codex lets its own ``model_context_window`` raise that model to. +A fallback entry offers the model group's ``supported_reasoning_efforts`` as its reasoning levels, +trimmed to the ones the requesting Codex release can parse. The stock entries are Codex 0.159.3's bundled catalog, vendored unchanged from https://github.com/openai/codex/blob/rust-v0.159.3/codex-rs/models-manager/models.json (Apache-2.0) @@ -158,6 +160,7 @@ class CodexStockModel(LiteLLMBaseModel): supports_search_tool: bool = False supported_reasoning_levels: tuple[CodexReasoningLevel, ...] = () context_window: int | None = None + max_context_window: int | None = None class CodexStockCatalog(LiteLLMBaseModel): @@ -173,8 +176,9 @@ class CodexCatalogRow: """What the catalog needs to know about one listed model: its public id, the listing's mode and input limit, the upstream model behind it, the reasoning effort its deployments pin, and the `model_info` values the operator set; `service_tiers` holds one raw `model_info.service_tiers` - value per deployment behind the id, None where a deployment sets none. `supported_reasoning_efforts` - is the model group's, None when unknown.""" + value per deployment behind the id, None where a deployment sets none, and + `configured_max_input_tokens` one `model_info.max_input_tokens` per deployment the same way. + `supported_reasoning_efforts` is the model group's, None when unknown.""" id: str mode: str | None = None @@ -185,6 +189,7 @@ class CodexCatalogRow: service_tiers: tuple[object, ...] = () supported_reasoning_efforts: tuple[str, ...] | None = None search_tool_support: tuple[bool | None, ...] = () + configured_max_input_tokens: tuple[int | None, ...] = () @cache @@ -333,6 +338,25 @@ def _alias_context_window(row: CodexCatalogRow, stock_model: CodexStockModel) -> return MappingProxyType({"context_window": min(stock_model.context_window, row.max_input_tokens)}) +def _configured_window(row: CodexCatalogRow, stock_model: CodexStockModel) -> Mapping[str, int]: + """The `context_window` the operator's `model_info.max_input_tokens` gives a known model's entry, + or nothing, which keeps `stock_model`'s own. + + A request to the row can route to any of its deployments, so the window is the smallest one of them + takes: its `max_input_tokens` where the operator set a positive one, Codex's own window where not. A + limit set on only some deployments can therefore lower the window but never raise it. The window never + goes past `stock_model`'s `max_context_window`, the most Codex lets its own `model_context_window` + setting raise that model to.""" + per_deployment: Final = row.configured_max_input_tokens + declared: Final = tuple(limit for limit in per_deployment if limit is not None and limit > 0) + unset_window: Final = () if len(declared) == len(per_deployment) else (stock_model.context_window,) + if not declared or None in unset_window: + return MappingProxyType({}) + smallest: Final = min(limit for limit in (*declared, *unset_window) if limit is not None) + ceiling: Final = stock_model.max_context_window + return MappingProxyType({"context_window": smallest if ceiling is None else min(smallest, ceiling)}) + + def _alias_reasoning_level(row: CodexCatalogRow, stock_model: CodexStockModel) -> str | None: efforts: Final = frozenset(level.effort for level in stock_model.supported_reasoning_levels) return row.reasoning_effort if row.reasoning_effort in efforts else None @@ -398,6 +422,7 @@ def _entry( "upgrade": upgrade, **({} if search_tool is None else {"supports_search_tool": search_tool}), **(_deployment_defaults(row, stock_model) if is_alias and alias_defaults == "deployment" else {}), + **_configured_window(row, stock_model), } if configured is None: return stock_model.model_copy(update=listing) @@ -473,10 +498,12 @@ def codex_models_response_json( byte limit silently and keeps its bundled one, so the body holds only the entries that fit, the ones offering a service tier taken first and the rest in listing order, an entry too large for the bytes left passed over, every kept entry at its listing position; `byte_limit=None` - keeps every entry. A row listed under another name than the Codex model its `upstream_model` - names (an alias) takes that model's entry; `alias_defaults="deployment"` drops the defaults - Codex picks for that model in favor of what the alias's deployments pin, while `"codex"` keeps - them, as Codex would use them for the model's own name. `client_version` is the requesting Codex's version, which decides the reasoning + keeps every entry. A row whose deployments set `model_info.max_input_tokens` gets the smallest + window any of them takes as its known model's `context_window`, capped at that model's + `max_context_window`. A row listed under another name than the Codex model its `upstream_model` names (an alias) takes that model's + entry; `alias_defaults="deployment"` drops the defaults Codex picks for that model in favor of + what the alias's deployments pin, while `"codex"` keeps them, as Codex would use them for the + model's own name. `client_version` is the requesting Codex's version, which decides the reasoning levels a fallback entry may offer; None counts as a version too old for custom levels. """ known: Final = bundled_codex_models() if stock is None else stock @@ -519,6 +546,9 @@ def _catalog_row( search_tool_support=( registry.resolve(lookup_id, kind="search_tool_support", team_id=team_id) if registry is not None else () ), + configured_max_input_tokens=( + registry.group(lookup_id, field="max_input_tokens", team_id=team_id) if registry is not None else () + ), ) @@ -532,9 +562,9 @@ def codex_catalog_rows( """`rows` joined with the router's configured metadata, looked up by each entry's internal id so team-scoped rows resolve the way the Anthropic listing's display names do, a `model_group_alias` reading its target's deployments under its own name; `team_id` is the requesting key's team, so the - upstream model that picks Codex's stock entry and the tiers are read only off the deployments its - requests can route to, and the reasoning efforts are the ones `/model_group/info` reports for the - entry. Display names are read under `label_entries`' ids when given.""" + upstream model that picks Codex's stock entry, the tiers, and the configured input limits are read + only off the deployments its requests can route to, and the reasoning efforts are the ones + `/model_group/info` reports for the entry. Display names are read under `label_entries`' ids when given.""" lookup_ids: Final = MappingProxyType(dict(entries)) label_ids: Final = lookup_ids if label_entries is None else MappingProxyType(dict(label_entries)) return tuple( diff --git a/litellm/proxy/common_utils/token_count_settings.py b/litellm/proxy/common_utils/token_count_settings.py index dd04e531f2d..4e3059cfe26 100644 --- a/litellm/proxy/common_utils/token_count_settings.py +++ b/litellm/proxy/common_utils/token_count_settings.py @@ -56,6 +56,22 @@ class TokenCountSettings: framing: int | None count_model: TokenCountModel | None tokenizer: TokenCountTokenizer | None = None + tool_framing: int | None = None + responses_framing: int | None = None + responses_tool_framing: int | None = None + + @property + def frames_counts(self) -> bool: + return any( + framing is not None + for framing in (self.framing, self.tool_framing, self.responses_framing, self.responses_tool_framing) + ) + + def hidden_prompt_tokens(self, carries_tools: bool, responses_api: bool) -> int: + chat: Final = self.tool_framing if carries_tools else None + responses: Final = self.responses_tool_framing if carries_tools else self.responses_framing + narrower: Final = responses if responses_api else chat + return (narrower if narrower is not None else self.framing) or 0 @dataclass(frozen=True, slots=True) @@ -105,6 +121,15 @@ def _framing(value: object) -> int | None: return None +def _first_framing(source: Mapping[str, object], keys: tuple[str, ...]) -> int | None: + return next((framing for framing in map(_framing, map(source.get, keys)) if framing is not None), None) + + +def _configured_framing(configured: Mapping[str, object], default: Mapping[str, object], *keys: str) -> int | None: + operator_sets_one: Final = _first_framing(configured, (*keys, "token_count_framing")) is not None + return _first_framing(configured if operator_sets_one else default, keys) + + def _non_blank(value: object) -> str | None: return value if isinstance(value, str) and value.strip() else None @@ -175,13 +200,21 @@ def _tokenizer(value: object) -> TokenCountTokenizer | None: def resolve_token_count_settings( configured: Mapping[str, object], default: Mapping[str, object], deployment_region: Callable[[], str] ) -> TokenCountSettings: - framing: Final = _framing(configured.get("token_count_framing")) return TokenCountSettings( - framing=framing if framing is not None else _framing(default.get("token_count_framing")), + framing=_configured_framing(configured, default, "token_count_framing"), count_model=_resolved_count_model( configured.get("token_count_model"), default.get("token_count_model"), deployment_region ), tokenizer=_tokenizer(default.get("token_count_tokenizer")), + tool_framing=_configured_framing(configured, default, "token_count_tool_framing"), + responses_framing=_configured_framing(configured, default, "token_count_responses_framing"), + responses_tool_framing=_configured_framing( + configured, + default, + "token_count_responses_tool_framing", + "token_count_responses_framing", + "token_count_tool_framing", + ), ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/mcp_security/mcp_security_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/mcp_security/mcp_security_guardrail.py index 9c0a574ae3b..b4652e010ff 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/mcp_security/mcp_security_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/mcp_security/mcp_security_guardrail.py @@ -107,9 +107,20 @@ class MCPSecurityGuardrail(CustomGuardrail): from litellm.proxy.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) + from litellm.proxy.proxy_server import prisma_client async with global_mcp_server_manager.catalog.operation(): - registry: Final = global_mcp_server_manager.get_registry() - registered_names: Final = set(registry.keys()) - - return frozenset(requested_servers - registered_names) + not_servers: Final = frozenset( + name + for name in requested_servers + if global_mcp_server_manager.get_mcp_server_by_identifier(name) is None + ) + if prisma_client is None: + return not_servers + return frozenset( + [ + name + for name in not_servers + if await global_mcp_server_manager.get_toolset_by_name_cached(prisma_client, name) is None + ] + ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c17cb5403c6..c7f357a6848 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -90,6 +90,7 @@ from litellm.constants import ( LITELLM_UI_SESSION_DURATION, OPENAI_LIVE_SESSION_START_TIMEOUT_SECONDS, RUNTIME_UPDATABLE_ROUTER_SETTINGS, + VLLM_CHOICE_FIELDS, ) from litellm.litellm_core_utils.access_windows import access_windows_config_error from litellm.litellm_core_utils.add_retry_fallback_headers import ( @@ -107,6 +108,7 @@ from litellm.litellm_core_utils.litellm_logging import ( initialize_sentry, ) from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER, LoggingWorker +from litellm.litellm_core_utils.prompt_templates.chat_template_tools import chat_template_tool_section from litellm.litellm_core_utils.routing_common_utils import resolve_model_group_alias from litellm.litellm_core_utils.routing_groups import parse_routing_groups from litellm.litellm_core_utils.safe_json_dumps import safe_dumps @@ -115,6 +117,7 @@ from litellm.litellm_core_utils.token_counter import ( UnsupportedContentTypeError, deployment_openai_model, offload_token_count, + sends_function_tools, ) from litellm.llms.bedrock.count_tokens.kimi_k3 import ( HUGGING_FACE_SOURCE, @@ -1007,6 +1010,7 @@ from litellm.proxy.video_endpoints.endpoints import router as video_router from litellm.repositories.base_repository import SupportsModelDump from litellm.repositories.credentials_repository import CredentialsRepository from litellm.repositories.prisma_protocols import TableActions +from litellm.responses.main import provider_stores_the_conversation from litellm.responses.streaming_iterator import drain_logging_tasks from litellm.router import ( AssistantsTypedDict, @@ -10223,6 +10227,11 @@ def _restamp_nested_response_model_in_sse_line(line: bytes, target_model: str) - return _SSE_DATA_FIELD + b" " + serialized_event.encode() + line_ending +def _vllm_choice_fields(choice: BaseModel) -> Mapping[str, object]: + choice_extras: Final = choice.model_extra or {} + return MappingProxyType({field: choice_extras[field] for field in VLLM_CHOICE_FIELDS if field in choice_extras}) + + def _json_encodes(value: object) -> bool: try: to_jsonable_python(value) @@ -10278,11 +10287,52 @@ def _serialize_chat_chunk(chunk: ModelResponseStream) -> bytes: def _serialize_streaming_chunk(chunk: BaseModel) -> bytes: + choices: Final = _streaming_chunk_choices(chunk) + null_choice_fields: Final = tuple(_null_vllm_choice_fields(choice) for choice in choices) + if any(null_choice_fields): + serialized_chunk: Final = _serialize_keeping_null_choice_fields(chunk, choices, null_choice_fields) + if serialized_chunk is not None: + return serialized_chunk if isinstance(chunk, ModelResponseStream): return _serialize_chat_chunk(chunk) return chunk.model_dump_json(exclude_none=True, exclude_unset=True).encode() +def _serialize_keeping_null_choice_fields( + chunk: BaseModel, choices: Sequence[BaseModel], null_choice_fields: Sequence[Mapping[str, None]] +) -> bytes | None: + try: + return json.dumps( + { + **_dump_set_fields(chunk, exclude=frozenset({"choices"})), + "choices": [ + {**_dump_set_fields(choice), **nulls} + for choice, nulls in zip(choices, null_choice_fields, strict=True) + ], + }, + separators=(",", ":"), + ensure_ascii=False, + ).encode() + except ValueError: + return None + + +def _streaming_chunk_choices(chunk: BaseModel) -> Sequence[BaseModel]: + match chunk: + case ModelResponseStream() | TextCompletionResponse(): + return chunk.choices + case _: + return () + + +def _null_vllm_choice_fields(choice: BaseModel) -> Mapping[str, None]: + return MappingProxyType({field: None for field, value in _vllm_choice_fields(choice).items() if value is None}) + + +def _dump_set_fields(model: BaseModel, exclude: AbstractSet[str] = frozenset()) -> Mapping[str, object]: + return model.model_dump(mode="json", exclude_none=True, exclude_unset=True, exclude=set(exclude)) + + async def _apply_streaming_chunk_hooks( *, chunk: Any, @@ -15053,7 +15103,7 @@ async def _provider_token_count( framing_can_stand_in: Final = (count_model is None or not count_model.configured) and ( failure.cannot_count_model or (pinned_count_model is not None and failure.denied_in_region) ) - if not (framing_can_stand_in and settings.framing is not None and local_estimate_allowed): + if not (framing_can_stand_in and settings.frames_counts and local_estimate_allowed): raise failure.to_proxy_exception() if pinned_count_model is not None: verbose_proxy_logger.warning( @@ -15253,6 +15303,7 @@ async def _tokenizer_or_provider_error( async def _local_token_estimate( model: str, + template_model: str, prompt: str | None, messages: Sequence[AllMessageValues] | None, tools: list[ChatCompletionToolParam] | None, @@ -15265,7 +15316,7 @@ async def _local_token_estimate( else await _tokenizer_or_provider_error(model, custom_tokenizer, provider_failure) ) count: Final = offload_token_count(litellm.token_counter)( - model=model, + model=template_model, text=prompt, messages=messages, tools=tools, @@ -15331,10 +15382,15 @@ def _deployment_aws_region(deployment: _TokenCountDeployment) -> str: ) +def _deployment_base_model(deployment: _TokenCountDeployment) -> str | None: + return _non_empty_str(deployment.get("model_info", {}).get("base_model")) or _non_empty_str( + deployment.get("litellm_params", {}).get("base_model") + ) + + def _token_count_settings_model(deployment: _TokenCountDeployment, request_model: str) -> tuple[str, str | None]: - model_info: Final = deployment.get("model_info", {}) params: Final = deployment.get("litellm_params", {}) - base_model: Final = _non_empty_str(model_info.get("base_model")) or _non_empty_str(params.get("base_model")) + base_model: Final = _deployment_base_model(deployment) if base_model is not None: return base_model, None model: Final = _non_empty_str(params.get("model")) @@ -15486,12 +15542,15 @@ def _responses_routing_kwargs(responses_request: Mapping[str, object] | None) -> return {"previous_response_id": previous_response_id} if isinstance(previous_response_id, str) else {} +def _inference_params( + litellm_params: Mapping[str, object], responses_request: Mapping[str, object] +) -> Mapping[str, object]: + return {**litellm_params, **_router_default_litellm_params(llm_router), **responses_request} + + def _count_drop_policy(deployment: Mapping[str, object] | None, responses_request: Mapping[str, object]) -> DropPolicy: litellm_params: Final = as_str_mapping(None if deployment is None else deployment.get("litellm_params")) or {} - router_defaults: Final[Mapping[str, object]] = ( - {} if llm_router is None else _router_default_litellm_params(llm_router) - ) - return drop_policy({**litellm_params, **router_defaults, **responses_request}) + return drop_policy(_inference_params(litellm_params, responses_request)) def _native_responses_count( @@ -15514,8 +15573,54 @@ def _native_responses_count( ) -def _deployment_api_base(deployment: _TokenCountDeployment) -> str | None: - return _non_empty_str(deployment.get("litellm_params", {}).get("api_base")) +def _deployment_api_base( + deployment: _TokenCountDeployment, responses_request: Mapping[str, object] | None = None +) -> str | None: + return _non_empty_str( + _inference_params(deployment.get("litellm_params", {}), responses_request or {}).get("api_base") + ) + + +def _deployment_provider( + deployment: _TokenCountDeployment, request_model: str, responses_request: Mapping[str, object] | None +) -> str | None: + params: Final = deployment.get("litellm_params", {}) + try: + return litellm.get_llm_provider( + model=_non_empty_str(params.get("model")) or request_model, + custom_llm_provider=_non_empty_str(params.get("custom_llm_provider")), + api_base=_deployment_api_base(deployment, responses_request), + )[1] + except litellm.BadRequestError: + return None + + +def _serves_native_responses( + deployment: _TokenCountDeployment, request_model: str, responses_request: Mapping[str, object] +) -> bool: + params: Final = deployment.get("litellm_params", {}) + inference: Final = _inference_params(params, responses_request) + try: + return provider_stores_the_conversation( + _non_empty_str(params.get("model")) or request_model, + _non_empty_str(params.get("custom_llm_provider")), + { + "use_chat_completions_api": bool(inference.get("use_chat_completions_api")), + "model_info": dict(deployment.get("model_info", {})), + "api_base": _non_empty_str(inference.get("api_base")), + }, + ) + except litellm.BadRequestError: + return False + + +_AZURE_PROVIDERS: Final = frozenset((litellm.LlmProviders.AZURE.value, litellm.LlmProviders.AZURE_AI.value)) + + +def _with_strict_by_default(tools: Sequence[Mapping[str, object]]) -> tuple[Mapping[str, object], ...]: + return tuple( + {**tool, "strict": True} if tool.get("type") == "function" and "strict" not in tool else tool for tool in tools + ) def _system_message(system: object) -> ChatCompletionSystemMessage | None: @@ -15720,7 +15825,7 @@ async def _count_request_tokens( else await chat_count_request( model_to_use, custom_llm_provider, - _deployment_api_base(deployment_view), + _deployment_api_base(deployment_view, routed_request), routed_request, _count_drop_policy(routed_deployment, routed_request), ) @@ -15818,10 +15923,31 @@ async def _count_request_tokens( counted_messages: Final = ( typed_messages if typed_messages is None or system_message is None else (system_message, *typed_messages) ) - counted_tools: Final = cast( # cast-ok: raw OpenAI or Anthropic tool dicts, both of which token_counter formats + native_responses: Final = responses_request is not None and _serves_native_responses( + deployment_view, request.model, responses_request + ) + sent_tools: Final = cast( # cast-ok: raw OpenAI or Anthropic tool dicts, both of which token_counter formats list[ChatCompletionToolParam] | None, tools if counted_messages is not None else None ) - framing: Final = settings.framing if counted_messages is not None and settings.framing is not None else 0 + openai_model: Final = deployment_openai_model(deployment_view) + template_model: Final = openai_model or _deployment_base_model(deployment_view) or model_to_use + tool_section: Final = chat_template_tool_section(template_model) + counted_tools: Final = ( + cast(list[ChatCompletionToolParam], list(_with_strict_by_default(sent_tools))) # cast-ok: strict added + if sent_tools is not None + and native_responses + and _deployment_provider(deployment_view, request.model, responses_request) in _AZURE_PROVIDERS + and tool_section is not None + and tool_section.renders_strict_only_when_true + else sent_tools + ) + framing: Final = ( + 0 + if counted_messages is None + else settings.hidden_prompt_tokens( + carries_tools=sends_function_tools(counted_tools), responses_api=native_responses + ) + ) kimi_k3_messages: Final = counted_messages if settings.tokenizer == "kimi_k3" else None exact: Final = ( await _kimi_k3_count( @@ -15845,7 +15971,8 @@ async def _count_request_tokens( count_source="local_estimate", ) estimate: Final = await _local_token_estimate( - model=deployment_openai_model(deployment_view) or model_to_use, + model=openai_model or model_to_use, + template_model=template_model, prompt=prompt if counted_messages is None else None, messages=counted_messages, tools=counted_tools, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index a1b36e5b7a1..0668c882a67 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -68,6 +68,7 @@ from litellm.constants import ( SPEND_LOG_QUEUE_MAX_BYTES, SPEND_LOG_WRITE_BATCH_MAX_BYTES, SPEND_LOG_WRITE_BATCH_MAX_ROWS, + VLLM_CHOICE_FIELDS, ) from litellm.cost_calculator import BaseTokenUsageProcessor from litellm.litellm_core_utils.bug_report import ( @@ -170,6 +171,7 @@ from litellm.litellm_core_utils.litellm_logging import ( deployment_pricing_model_info, ) from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token +from litellm.litellm_core_utils.redact_messages import drop_token_ids from litellm.litellm_core_utils.routing_common_utils import resolve_model_group_alias from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads @@ -179,6 +181,7 @@ from litellm.litellm_core_utils.served_output_texts import ( ) from litellm.litellm_core_utils.token_counter import offload_token_count, usage_tokenizer_model from litellm.llms import load_guardrail_translation_mappings +from litellm.llms.base_llm.guardrail_translation.utils import stream_item_field, stream_item_items from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import ( AlertType, @@ -707,6 +710,56 @@ def _pipeline_step_streams_live(step: "PipelineStep") -> bool: return not unified_guardrail.resolve_streaming_flag(callback, "streaming_buffer_until_moderated", True) +def _drop_stream_chunk_token_ids(chunk: object) -> None: + for choice in stream_item_items(chunk, "choices"): + drop_token_ids(choice) + + +def _drop_built_stream_token_ids(request_data: Mapping[str, object]) -> None: + logging_obj: Final = request_data.get("litellm_logging_obj") + if isinstance(logging_obj, Logging): + logging_obj.drops_stream_token_ids = True + + +def _drop_token_ids_ahead_of_chunk_check( + callback: CustomLogger, chunk: object, request_data: Mapping[str, object] +) -> None: + if _overrides_hook(callback, "async_post_call_streaming_hook"): + _drop_stream_chunk_token_ids(chunk) + _drop_built_stream_token_ids(request_data) + + +_PLAIN_STREAM_FIELDS: Final = TypeAdapter(dict[str, object]) +_EMPTY_STREAM_VALUES: Final[tuple[object, ...]] = (None, "", [], {}) + + +def _holds_stream_payload(value: object) -> bool: + if isinstance(value, dict): + nested: Final = _PLAIN_STREAM_FIELDS.validate_python(value).values() + return any(item not in _EMPTY_STREAM_VALUES for item in nested) + return value not in _EMPTY_STREAM_VALUES + + +def _choice_carries_only_token_ids(choice: object) -> bool: + if stream_item_field(choice, "token_ids") is None or not isinstance(choice, BaseModel): + return False + fields: Final = _PLAIN_STREAM_FIELDS.validate_python(choice.model_dump()) + return not any(_holds_stream_payload(value) for key, value in fields.items() if key not in ("index", "token_ids")) + + +def _is_token_id_only_chunk(chunk: object) -> bool: + choices: Final = stream_item_items(chunk, "choices") + return bool(choices) and all(_choice_carries_only_token_ids(choice) for choice in choices) + + +async def _without_token_id_only_chunks(response: AsyncIterable[object]) -> AsyncGenerator[object, None]: + async for chunk in response: + if _is_token_id_only_chunk(chunk): + _drop_stream_chunk_token_ids(chunk) + continue + yield chunk + + def _pipelines_stream_live(pipelines: Sequence[tuple[str, "GuardrailPipeline"]]) -> bool: return all(_pipeline_step_streams_live(step) for step in _pipeline_steps(pipelines)) @@ -4710,6 +4763,7 @@ class ProxyLogging: else: _callback = callback if _callback is not None and isinstance(_callback, CustomLogger): + _drop_token_ids_ahead_of_chunk_check(_callback, response, data) if str_so_far is not None: complete_response = str_so_far + response_str else: @@ -4731,7 +4785,7 @@ class ProxyLogging: self, response, user_api_key_dict: UserAPIKeyAuth, - request_data: dict, + request_data: dict[str, object], ): """ Allow user to modify outgoing streaming data -> Given a whole response iterator. @@ -4765,20 +4819,26 @@ class ProxyLogging: # Merge model-level guardrails before checking which guardrails to run request_data = check_and_merge_model_level_guardrails(data=request_data, llm_router=llm_router) - current_response = response stream_needs_translation: Final = ProxyLogging._stream_requires_guardrail_translation(user_api_key_dict) pipeline_gated_names: Final = _pipeline_step_guardrail_names(post_call_pipelines) + running_overrides: Final = tuple( + (callback, kind) + for callback, kind in caps.iterator_overrides + if not isinstance(callback, CustomGuardrail) + or ( + callback.guardrail_name not in pipeline_gated_names + and callback.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_call) is True + ) + ) + drops_token_ids: Final = bool(post_call_pipelines) or any( + isinstance(callback, CustomGuardrail) for callback, _ in running_overrides + ) + if drops_token_ids: + _drop_built_stream_token_ids(request_data) + current_response = _without_token_id_only_chunks(response) if drops_token_ids else response guarded_layers: Final[list[AsyncGenerator[object, None]]] = [] # mutable-ok: closed on disconnect - for resolved_callback, kind in caps.iterator_overrides: - if isinstance(resolved_callback, CustomGuardrail): - if resolved_callback.guardrail_name in pipeline_gated_names: - continue - if ( - resolved_callback.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_call) - is not True - ): - continue + for resolved_callback, kind in running_overrides: effective_kind = ( "apply_guardrail" if ( @@ -4836,6 +4896,8 @@ class ProxyLogging: try: async for chunk in current_response: served_chunks.append(chunk) + if drops_token_ids: + _drop_stream_chunk_token_ids(served_chunks[-1]) yield chunk except (GeneratorExit, asyncio.CancelledError): await ProxyLogging._close_guarded_layers(guarded_layers) @@ -10203,7 +10265,7 @@ _PRESERVED_NONE_FIELDS: Final[list[tuple[str, str]]] = [ ("message", "refusal"), ("delta", "content"), # null in streaming chunks ] -_PRESERVED_NONE_CHOICE_FIELDS: Final = ("logprobs",) +_PRESERVED_NONE_CHOICE_FIELDS: Final = ("logprobs", *VLLM_CHOICE_FIELDS) _PRESERVED_NONE_RESPONSE_FIELDS: Final = ("system_fingerprint",) diff --git a/litellm/router_api.py b/litellm/router_api.py index c013a99573f..3007b5345c7 100644 --- a/litellm/router_api.py +++ b/litellm/router_api.py @@ -155,6 +155,7 @@ GroupField: TypeAlias = Literal[ "display_name", "mode", "service_tiers", + "max_input_tokens", "token_limits", "listing_info", ] @@ -855,6 +856,16 @@ class RouterRegistry: team_id: str | None = None, ) -> tuple[object, ...]: ... + @overload + def group( + self, + name: str, + *, + field: Literal["max_input_tokens"], + include_hidden: bool = False, + team_id: str | None = None, + ) -> tuple[int | None, ...]: ... + @overload def group( self, @@ -882,9 +893,7 @@ class RouterRegistry: field: GroupField = "info", include_hidden: bool = False, team_id: str | None = None, - ) -> ( - "ModelGroupInfo | DeploymentModelListingInfo | tuple[int | None, int | None] | tuple[object, ...] | str | None" - ): + ) -> "ModelGroupInfo | DeploymentModelListingInfo | tuple[int | None, ...] | tuple[object, ...] | str | None": """ What the router knows about a model group. ``field``: @@ -894,6 +903,7 @@ class RouterRegistry: (were ``get_configured_display_name``, ``get_configured_mode``, ``get_configured_token_limits``, ``get_model_listing_info``) - ``service_tiers``: (was ``get_configured_service_tiers``) + - ``max_input_tokens``: each routable deployment's own ``model_info.max_input_tokens`` """ r: Final = self._router if field == "info": @@ -906,6 +916,8 @@ class RouterRegistry: return r._core_get_configured_mode(name) if field == "service_tiers": return r._core_get_configured_service_tiers(name, team_id=team_id) + if field == "max_input_tokens": + return r._core_get_configured_max_input_tokens(name, team_id=team_id) if field == "token_limits": return r._core_get_configured_token_limits(name) if field == "listing_info": diff --git a/litellm/router_python_core.py b/litellm/router_python_core.py index c63e538c998..98c9e641278 100644 --- a/litellm/router_python_core.py +++ b/litellm/router_python_core.py @@ -141,6 +141,7 @@ from litellm.litellm_core_utils.get_llm_provider_logic import ( is_registered_custom_provider, resolve_model_and_provider_for_metadata, ) +from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.litellm_core_utils.get_retry_from_policy import ( get_num_retries_from_retry_policy as _get_num_retries_from_retry_policy, ) @@ -475,6 +476,7 @@ from litellm.types.router import ( from litellm.types.services import ServiceTypes from litellm.types.utils import ( AUTOROUTER_CLASSIFIER_CALL_ORIGIN, + COST_MAP_LOOKUP_KEY, PROMPT_QUOTING_ROUTING_DECISION_FIELDS, CustomPricingLiteLLMParams, GenericBudgetConfigType, @@ -486,6 +488,7 @@ from litellm.types.utils import ( StandardLoggingRoutingDecision, Usage, all_litellm_params, + echoed_cost_map_fields, is_server_derived_pricing_key, shared_backend_model_info, ) @@ -772,6 +775,17 @@ def _configured_model_info(deployment: Mapping[str, object]) -> Mapping[str, obj return _MODEL_INFO_ADAPTER.validate_python(deployment.get("model_info") or {}) +def _operator_max_input_tokens(deployment: Mapping[str, object]) -> int | None: + model_info: Final = _configured_model_info(deployment) + catalog_key: Final = model_info.get(COST_MAP_LOOKUP_KEY) + catalog_entry: Final = ( + GetModelCostMap.loaded_model_cost_map().get(catalog_key) if isinstance(catalog_key, str) else None + ) + if catalog_entry is not None and "max_input_tokens" in echoed_cost_map_fields(model_info, catalog_entry): + return None + return coerce_token_limit(model_info.get("max_input_tokens")) + + def _search_tool_support(deployment: DeploymentTypedDict) -> bool | None: from litellm.responses.litellm_completion_transformation.tool_search import chat_config_builds_anthropic_request @@ -12693,6 +12707,24 @@ class PythonCore(DeprecatedRouterMembers): for deployment in self._routable_deployments(model_name, team_id) ) + def _core_get_configured_max_input_tokens( + self, model_name: str, team_id: str | None = None + ) -> tuple[int | None, ...]: + """ + Return the max_input_tokens each deployment a request from team_id to + model_name can route to (see _routable_deployments) sets in its own + model_info, in model_list order, None for a deployment that sets none + or a malformed one. Unlike get_model_listing_info, neither the cost map + nor a discovered limit is read, so only an operator's own value counts, + and a value a saved /model/info response copied from the cost map is + skipped too. + + Returns an empty tuple for wildcard-expanded or unknown names. + """ + return tuple( + _operator_max_input_tokens(deployment) for deployment in self._routable_deployments(model_name, team_id) + ) + def _core_get_search_tool_support(self, model_name: str, team_id: str | None = None) -> tuple[bool | None, ...]: """ Whether each deployment a request from team_id to model_name can route diff --git a/litellm/text_completions/streaming_iterator.py b/litellm/text_completions/streaming_iterator.py index 61e47e9d7af..6bba56a5afe 100644 --- a/litellm/text_completions/streaming_iterator.py +++ b/litellm/text_completions/streaming_iterator.py @@ -1,5 +1,6 @@ from typing import Final +from litellm.constants import VLLM_CHOICE_FIELDS from litellm.litellm_core_utils.stream_usage_chunks import is_injected_stream_usage_artifact from litellm.types.utils import Choices, ModelResponse, ModelResponseStream, TextChoices, TextCompletionResponse @@ -37,16 +38,19 @@ class TextCompletionStreamWrapper: response["object"] = "text_completion" response["created"] = chunk.get("created", None) response["model"] = chunk.get("model", None) - text_choices: Final = TextChoices() if isinstance(chunk, Choices): # chunk should always be of type StreamingChoices raise Exception if chunk["choices"]: - delta: Final = chunk["choices"][0]["delta"] + chat_choice: Final = chunk["choices"][0] + text_choices: Final = TextChoices( + **{field: chat_choice[field] for field in VLLM_CHOICE_FIELDS if field in chat_choice} + ) + delta: Final = chat_choice["delta"] text_choices["text"] = delta["content"] text_choices["reasoning_content"] = delta.get("reasoning_content") - text_choices["index"] = chunk["choices"][0]["index"] - text_choices["finish_reason"] = chunk["choices"][0]["finish_reason"] + text_choices["index"] = chat_choice["index"] + text_choices["finish_reason"] = chat_choice["finish_reason"] response["choices"] = [text_choices] else: response["choices"] = [] diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index dc35af81091..5737d717661 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -849,6 +849,7 @@ class ANTHROPIC_BETA_HEADER_VALUES(str, Enum): ADVISOR_TOOL_2026_03_01 = "advisor-tool-2026-03-01" PER_TURN_CONTROL_2026_07_01 = "per-turn-control-2026-07-01" DANGEROUS_TOOL_USE_2026_09_03 = "dangerous-tool-use-2026-09-03" + MCP_CLIENT_2025_04_04 = "mcp-client-2025-04-04" # Tool search beta header constant (for Anthropic direct API and Microsoft Foundry) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index da1d99957a7..bf69445fdd1 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -246,6 +246,9 @@ class ProviderSpecificModelInfo(TypedDict, total=False): bedrock_converse_supports_strict_tools: bool | None supports_regex_lookaround: ReadOnly[bool | None] token_count_framing: ReadOnly[int | None] + token_count_tool_framing: ReadOnly[int | None] + token_count_responses_framing: ReadOnly[int | None] + token_count_responses_tool_framing: ReadOnly[int | None] token_count_model: ReadOnly[TokenCountModelInfo | None] token_count_tokenizer: ReadOnly[TokenCountTokenizer | None] @@ -4236,7 +4239,15 @@ DEPLOYMENT_SCOPED_CAPABILITY_FIELDS: Final[frozenset[str]] = frozenset( } ) -DEPLOYMENT_SCOPED_TOKEN_COUNT_FIELDS: Final[frozenset[str]] = frozenset({"token_count_framing", "token_count_model"}) +DEPLOYMENT_SCOPED_TOKEN_COUNT_FIELDS: Final[frozenset[str]] = frozenset( + { + "token_count_framing", + "token_count_tool_framing", + "token_count_responses_framing", + "token_count_responses_tool_framing", + "token_count_model", + } +) SHARED_BACKEND_MODEL_INFO_FIELDS: Final[frozenset[str]] = ( frozenset(ModelInfoBase.__required_keys__ | ModelInfoBase.__optional_keys__) diff --git a/litellm/utils.py b/litellm/utils.py index b0826f84006..f03b1eb7ad9 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -98,7 +98,11 @@ from litellm.litellm_core_utils.model_retirement import is_hidden_as_retired from litellm.litellm_core_utils.offline_mode import is_offline_mode from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload from litellm.litellm_core_utils.tokenizer import strip_special_tokens -from litellm.litellm_core_utils.tool_fields import MalformedField, resolve_deployment_tool_fields +from litellm.litellm_core_utils.tool_fields import ( + MalformedField, + resolve_deployment_tool_fields, + tools_without_hosted_mcp, +) from litellm.rust_bridge import tokenizer as tokenizer_dispatch from litellm.rust_bridge.catalog import decision from litellm.rust_bridge.configuration import Decision @@ -5186,7 +5190,12 @@ def remove_sensitive_keys_from_dict(d: dict) -> dict: return d -def pre_process_optional_params(passed_params: dict, non_default_params: dict, custom_llm_provider: str) -> dict: +_TOOL_LIST: Final = TypeAdapter(list[object]) + + +def pre_process_optional_params( + passed_params: dict, non_default_params: dict, custom_llm_provider: str, model: str, drop_params: bool | None +) -> dict: """For .completion(), preprocess optional params""" optional_params: dict = {} @@ -5243,14 +5252,20 @@ def pre_process_optional_params(passed_params: dict, non_default_params: dict, c and custom_llm_provider not in litellm.openai_compatible_providers ): if custom_llm_provider == "ollama": + if "tools" in non_default_params: + sent_tools: Final = _TOOL_LIST.validate_python(non_default_params.pop("tools")) + non_default_params.pop("tool_choice", None) # causes ollama requests to hang + prompt_tools: Final = tools_without_hosted_mcp( + sent_tools, model=model, route=custom_llm_provider, drop_params=bool(drop_params) + ) + if sent_tools and not prompt_tools: + return optional_params + optional_params["functions_unsupported_model"] = prompt_tools + elif "functions" in non_default_params: + optional_params["functions_unsupported_model"] = non_default_params.pop("functions") # ollama actually supports json output optional_params["format"] = "json" litellm.add_function_to_prompt = True # so that main.py adds the function call to the prompt - if "tools" in non_default_params: - optional_params["functions_unsupported_model"] = non_default_params.pop("tools") - non_default_params.pop("tool_choice", None) # causes ollama requests to hang - elif "functions" in non_default_params: - optional_params["functions_unsupported_model"] = non_default_params.pop("functions") elif litellm.add_function_to_prompt: # if user opts to add it to prompt instead optional_params["functions_unsupported_model"] = non_default_params.pop( "tools", non_default_params.pop("functions", None) @@ -5385,6 +5400,8 @@ def get_optional_params( passed_params=passed_params, non_default_params=non_default_params, custom_llm_provider=custom_llm_provider, + model=model, + drop_params=drop_params, ) def _check_valid_arg(supported_params: list[str]): @@ -7398,6 +7415,9 @@ def get_model_info_helper( supports_image_size=_model_info.get("supports_image_size", None), supported_audio_formats=_model_info.get("supported_audio_formats", None), token_count_framing=_model_info.get("token_count_framing", None), + token_count_tool_framing=_model_info.get("token_count_tool_framing", None), + token_count_responses_framing=_model_info.get("token_count_responses_framing", None), + token_count_responses_tool_framing=_model_info.get("token_count_responses_tool_framing", None), token_count_model=_model_info.get("token_count_model", None), token_count_tokenizer=_model_info.get("token_count_tokenizer", None), vertex_ai_audio_api=_model_info.get("vertex_ai_audio_api", None), diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 02bf7dedcef..e6f65d6b416 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -13478,6 +13478,10 @@ "input_cost_per_token": 2e-07, "output_cost_per_token": 5e-07, "litellm_provider": "azure_ai", + "token_count_framing": 149, + "token_count_tool_framing": 302, + "token_count_responses_framing": 145, + "token_count_responses_tool_framing": 100, "max_input_tokens": 128000, "max_output_tokens": 128000, "mode": "chat", @@ -54457,6 +54461,8 @@ "input_cost_per_token_above_200k_tokens_batches": 2e-06, "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", + "token_count_framing": 185, + "token_count_tool_framing": 214, "max_input_tokens": 1000000, "max_output_tokens": 1000000, "mode": "chat", @@ -62640,7 +62646,9 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "token_count_framing": 23, + "token_count_tool_framing": 12 }, "deepseek/deepseek-v4-flash-vision-exp": { "cache_creation_input_token_cost": 0.0, @@ -62746,7 +62754,9 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": false + "supports_vision": false, + "token_count_framing": 76, + "token_count_tool_framing": 75 }, "tencent/deepseek-v4-pro": { "cache_creation_input_token_cost": 0.0, @@ -64032,6 +64042,7 @@ "input_cost_per_token": 3e-07, "input_cost_per_token_priority": 3.75e-07, "litellm_provider": "fireworks_ai", + "token_count_framing": 23, "max_input_tokens": 1048576, "max_output_tokens": 393216, "mode": "chat", @@ -64133,6 +64144,7 @@ "input_cost_per_token": 3e-07, "input_cost_per_token_priority": 3.75e-07, "litellm_provider": "fireworks_ai", + "token_count_framing": 23, "max_input_tokens": 1048576, "max_output_tokens": 393216, "mode": "chat", diff --git a/tests/integration/caching/test_response_cache_empty_output.py b/tests/integration/caching/test_response_cache_empty_output.py index ce5a4d64ca7..7a20e712d4e 100644 --- a/tests/integration/caching/test_response_cache_empty_output.py +++ b/tests/integration/caching/test_response_cache_empty_output.py @@ -23,7 +23,7 @@ from typing import Final, TypeAlias, TypeVar import httpx import pytest from anthropic import Anthropic, AsyncAnthropic -from openai import AsyncOpenAI, OpenAI +from openai import AsyncOpenAI, InternalServerError, OpenAI from pydantic import JsonValue from tests.integration._support.client import JSON_OBJECT, Gateway, eventually, object_value, string_value @@ -420,14 +420,25 @@ def _streamed_chat(gateway: Gateway, model: str, text: str) -> _Answer: return ids, streamed, tuple(chunk.object for chunk in chunks) -def test_streamed_chat_without_choices_is_never_assembled_or_cached(gateway: Gateway) -> None: +def test_streamed_chat_without_choices_fails_and_is_never_cached(gateway: Gateway) -> None: upstream: Final = Upstream(gateway.upstream_url) + store: Final = redis_store() text: Final = prompt() with gateway.scenario() as scenario: handle: Final = scripted(scenario, CHAT_STREAM_WITHOUT_CHOICES) model: Final = scenario.model(api_base=handle.api_base()) - empty: Final = _never_replayed(partial(_streamed_chat, gateway, model, text), upstream, handle.scenario_id) - assert all(answer.text == "" for answer in empty), empty + stream: Final = partial(_streamed_chat, gateway, model, text) + for attempt in range(1, EMPTY_STREAMS + 1): + with pytest.raises(InternalServerError, match="Stream ended before a finish_reason was received"): + stream() + _assert_prompt_reached_upstream(upstream, handle.scenario_id, text, attempt) + rescript(handle, CHAT_STREAM) + real_ids, real_text, _ = stream() + assert real_text == "streamed response" and len(real_ids) == 1, (real_ids, real_text) + _assert_prompt_reached_upstream(upstream, handle.scenario_id, text, EMPTY_STREAMS + 1) + await_entry(store, next(iter(real_ids))) + replay: Final = _replayed(stream, upstream, handle.scenario_id, real_ids) + assert replay.text == real_text, replay def _streamed_text(gateway: Gateway, model: str, text: str) -> _Answer: diff --git a/tests/integration/compatibility/test_a2a_wire_versions.py b/tests/integration/compatibility/test_a2a_wire_versions.py index 4471322d1e9..b3ca0c99392 100644 --- a/tests/integration/compatibility/test_a2a_wire_versions.py +++ b/tests/integration/compatibility/test_a2a_wire_versions.py @@ -1023,9 +1023,18 @@ def test_a2a_task_methods_pass_through_verbatim_and_push_callbacks_are_validated {"taskId": "t1", "url": callback}, ): rejected = call("tasks/pushNotificationConfig/set", params) - assert (rejected.status_code, rejected.json()) == (400, {"detail": detail}), ( - f"{params}: {rejected.text}" - ) + assert (rejected.status_code, rejected.json()) == ( + 400, + { + "jsonrpc": "2.0", + "id": marker + "tasks/pushNotificationConfig/set", + "error": { + "code": -32000, + "message": detail, + "data": {"type": "invalid_request_error", "param": None, "code": "400"}, + }, + }, + ), f"{params}: {rejected.text}" assert tuple(item for item in wire.drain() if item.method == "POST") == () public: Final = {"taskId": "t1", "pushNotificationConfig": {"url": "https://1.1.1.1/callback"}} @@ -1118,7 +1127,15 @@ def test_a2a_v1_task_methods_and_extended_card_lower_to_0_3_upstream_and_answer_ rejected = call("CreateTaskPushNotificationConfig", params) assert (rejected.status_code, rejected.json()) == ( 400, - {"detail": "Push notification URL must use HTTPS"}, + { + "jsonrpc": "2.0", + "id": marker + "CreateTaskPushNotificationConfig", + "error": { + "code": -32000, + "message": "Push notification URL must use HTTPS", + "data": {"type": "invalid_request_error", "param": None, "code": "400"}, + }, + }, ), f"{params}: {rejected.text}" assert forwarded() == () diff --git a/tests/integration/compatibility/test_codex_model_catalog.py b/tests/integration/compatibility/test_codex_model_catalog.py index a3ba032e612..fbf05b0f526 100644 --- a/tests/integration/compatibility/test_codex_model_catalog.py +++ b/tests/integration/compatibility/test_codex_model_catalog.py @@ -351,6 +351,41 @@ def test_stock_model_keeps_codex_entry(gateway: Gateway) -> None: assert catalog[name] == expected, catalog[name] +def test_stock_model_alias_takes_configured_max_input_tokens(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + stock: Final = _bundled_stock()["gpt-6-sol"] + stock_window: Final = stock["context_window"] + ceiling: Final = stock["max_context_window"] + assert isinstance(stock_window, int) and isinstance(ceiling, int) and stock_window < ceiling, stock + raised: Final = _model_name("codex-raised-window") + lowered: Final = _model_name("codex-lowered-window") + half_set: Final = _model_name("codex-half-set-window") + half_lowered: Final = _model_name("codex-half-lowered-window") + unset: Final = _model_name("codex-unset-window") + names: Final = (raised, lowered, half_set, half_lowered, unset) + _new_model(scenario, raised, model="openai/gpt-6-sol", model_info={"max_input_tokens": ceiling + 50000}) + _new_model(scenario, lowered, model="openai/gpt-6-sol", model_info={"max_input_tokens": stock_window // 2}) + _new_model(scenario, half_set, model="openai/gpt-6-sol", model_info={"max_input_tokens": ceiling}) + _new_model(scenario, half_set, model="openai/gpt-6-sol") + _new_model(scenario, half_lowered, model="openai/gpt-6-sol", model_info={"max_input_tokens": stock_window // 3}) + _new_model(scenario, half_lowered, model="openai/gpt-6-sol") + _new_model(scenario, unset, model="openai/gpt-6-sol") + expected: Final = { + raised: ceiling, + lowered: stock_window // 2, + half_set: stock_window, + half_lowered: stock_window // 3, + unset: stock_window, + } + windows: Final = _settled( + lambda: { + slug: entry["context_window"] for slug, entry in _by_slug(_catalog(gateway)).items() if slug in names + }, + lambda value: value == expected, + ) + assert windows == expected, windows + + def test_stock_model_configured_tiers_override(gateway: Gateway) -> None: with gateway.scenario() as scenario: stock: Final = _bundled_stock()["gpt-6-sol"] diff --git a/tests/integration/cost_calculation/cases/azure.py b/tests/integration/cost_calculation/cases/azure.py index d9f17c12bda..6495199b8bb 100644 --- a/tests/integration/cost_calculation/cases/azure.py +++ b/tests/integration/cost_calculation/cases/azure.py @@ -1588,6 +1588,7 @@ AZURE_PINNED_GPT_5_4_MINI_STREAM: Final = CostTrackingTestCase( content_type="text/event-stream", frames=( 'data: {"id":"chatcmpl-$REQUEST_ID","object":"chat.completion.chunk","created":1,"model":"$MODEL","choices":[{"index":0,"delta":{"role":"assistant","content":"ok"},"finish_reason":null}]}\n\n', + 'data: {"id":"chatcmpl-$REQUEST_ID","object":"chat.completion.chunk","created":1,"model":"$MODEL","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}\n\n', 'data: {"id":"chatcmpl-$REQUEST_ID","object":"chat.completion.chunk","created":1,"model":"$MODEL","choices":[],"usage":{"prompt_tokens":1840,"completion_tokens":412,"total_tokens":2252}}\n\n', "data: [DONE]\n\n", ), diff --git a/tests/integration/cost_calculation/cases/openai.py b/tests/integration/cost_calculation/cases/openai.py index bc7fb3486be..270bb6e105d 100644 --- a/tests/integration/cost_calculation/cases/openai.py +++ b/tests/integration/cost_calculation/cases/openai.py @@ -7952,6 +7952,7 @@ GPT_5_6_STREAM_USAGE_IN_EMPTY_CHOICES_CHUNK: Final = CostTrackingTestCase( frames=( 'data: {"id":"chatcmpl-$REQUEST_ID","object":"chat.completion.chunk","created":1789788263,"model":"gpt-5.6","choices":[{"index":0,"delta":{"role":"assistant"},"finish_reason":null}],"usage":null}', 'data: {"id":"chatcmpl-$REQUEST_ID","object":"chat.completion.chunk","created":1789788263,"model":"gpt-5.6","choices":[{"index":0,"delta":{"content":"scripted answer"},"finish_reason":null}],"usage":null}', + 'data: {"id":"chatcmpl-$REQUEST_ID","object":"chat.completion.chunk","created":1789788263,"model":"gpt-5.6","choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":null}', 'data: {"id":"chatcmpl-$REQUEST_ID","object":"chat.completion.chunk","created":1789788263,"model":"gpt-5.6","choices":[],"usage":{"prompt_tokens":1840,"completion_tokens":412,"total_tokens":2252}}', "data: [DONE]", ), diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_pdf_document_count_tokens_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_pdf_document_count_tokens_wire.py index 8eeaab66553..210b754f626 100644 --- a/tests/integration/messages_endpoint/providers/anthropic/test_pdf_document_count_tokens_wire.py +++ b/tests/integration/messages_endpoint/providers/anthropic/test_pdf_document_count_tokens_wire.py @@ -199,14 +199,14 @@ def test_utils_token_counter_adds_title_and_context_once_per_document( @pytest.mark.parametrize("label", list(_UNREADABLE)) -def test_utils_token_counter_prices_unreadable_pdf_bytes_like_one_image( +def test_utils_token_counter_uses_the_same_bounded_fallback_for_unreadable_pdf_bytes( gateway: Gateway, provider: SharedProvider, label: str ) -> None: data: Final = _UNREADABLE[label] as_pdf: Final = _local(gateway, provider, [document(base64_source(data))]) - as_png: Final = _local(gateway, provider, [document(base64_source(data, media_type="image/png"))]) + fallback: Final = _local(gateway, provider, [document(base64_source(encoded(b"unreadable-document")))]) without: Final = _local(gateway, provider, []) - assert as_pdf == as_png, (as_pdf, as_png) + assert as_pdf == fallback, (as_pdf, fallback) assert 0 <= as_pdf - without < rendered_tokens((LETTER,)), (as_pdf, without) diff --git a/tests/integration/observability/test_cache_hit_guardrail_metrics.py b/tests/integration/observability/test_cache_hit_guardrail_metrics.py index c52649c46df..13cd2133d4f 100644 --- a/tests/integration/observability/test_cache_hit_guardrail_metrics.py +++ b/tests/integration/observability/test_cache_hit_guardrail_metrics.py @@ -253,6 +253,16 @@ def _count(samples: tuple, name: str, model_id: str) -> float: ) +def _deployment_states(rig: Rig) -> tuple[float, ...]: + return tuple( + sample.value + for sample in _metric_samples(rig.candidate, rig.model_name) + if sample.name == DEPLOYMENT_STATE + and sample.labels.get("model_id") == rig.deployment_id + and sample.labels.get("api_base") == rig.provider.url + "/v1" + ) + + def _populated_failures(samples: tuple, rig: Rig, api_provider: str) -> float: return float( sum( @@ -558,7 +568,8 @@ def test_cache_hit_post_call_guardrail_outage_keeps_deployment_labels( ) assert reject.status_code >= 400, reject.text assert rig.provider.received.qsize() == 1, rig.provider.drain() - _expect_metrics(rig, 0, 0, pf_populated=1, pf_blank=0) + _expect_metrics(rig, 1, 0) + _assert_spend(_call_id(reject), rig) def test_two_identical_cache_hit_rejects_increment_populated_series(gateway: Gateway, tmp_path: Path) -> None: @@ -635,6 +646,7 @@ def test_cache_hit_reject_deployment_metric_set_diff(gateway: Gateway, tmp_path: text: Final = "cache hit control e4 " + marker warm: Final = rig.candidate.request("POST", "/v1/chat/completions", _chat_body(rig.model_name, text, None)) assert warm.status_code == 200, warm.text + eventually(lambda: _deployment_states(rig), lambda states: states == (0.0,)) reject: Final = rig.candidate.request( "POST", "/v1/chat/completions", _chat_body(rig.model_name, text, rig.guardrail_name) ) @@ -642,11 +654,4 @@ def test_cache_hit_reject_deployment_metric_set_diff(gateway: Gateway, tmp_path: samples: Final = _expect_metrics(rig, 1, 0) blank: Final = tuple(sample for sample in samples if sample.labels.get("model_id") == "") assert blank == (), blank - states: Final = tuple( - sample.value - for sample in samples - if sample.name == DEPLOYMENT_STATE - and sample.labels.get("model_id") == rig.deployment_id - and sample.labels.get("api_base") == "" - ) - assert states == (1.0,), states + eventually(lambda: _deployment_states(rig), lambda states: states == (1.0,)) diff --git a/tests/integration/observability/test_prometheus_provider_auth_failure.py b/tests/integration/observability/test_prometheus_provider_auth_failure.py index 13d264ca6b0..de1f19bc523 100644 --- a/tests/integration/observability/test_prometheus_provider_auth_failure.py +++ b/tests/integration/observability/test_prometheus_provider_auth_failure.py @@ -995,7 +995,7 @@ def _gated(gateway: Gateway, cell: str, key: str | None) -> httpx.Response: def recorded_before(cell: str) -> list[tuple[str, str, str, float]]: - if cell in ("invalid-key", "expired-key"): + if cell in ("invalid-key", "missing-key", "expired-key", "blocked-key"): return [] return sorted([(FAILURES, "", "", 1.0), (TOTALS, "", "", 1.0), (LLM_FAILED, identity(cell), "", 1.0)]) @@ -1020,7 +1020,9 @@ def test_a_request_the_proxy_rejects_itself_never_marks_the_deployment_down_and_ assert response.status_code == (403 if cell == "no-model-access" else 401), response.text assert rig.upstream.drain() == () samples: Final = until(rig.gateway, PROXY_FAILED, requested_model=name(cell)) + assert values(samples, PROXY_FAILED, requested_model=name(cell)) == [1.0], recorded(samples) assert deployment_rows(samples, cell) == recorded_before(cell), recorded(samples, requested_model=name(cell)) + assert values(samples, STATE, model_id=identity(cell)) == [], recorded(samples) def test_a_key_rate_limit_rejected_at_the_proxy_is_still_counted_as_before(rig: Rig) -> None: diff --git a/tests/integration/pricing/test_price_data_reload_prices.py b/tests/integration/pricing/test_price_data_reload_prices.py index 0aa3a3bcba2..a91f5cba1ad 100644 --- a/tests/integration/pricing/test_price_data_reload_prices.py +++ b/tests/integration/pricing/test_price_data_reload_prices.py @@ -284,7 +284,7 @@ def _charged_as_listed(proxy: Gateway, deployment: _Deployment, scale: float) -> assert _listed_rates(proxy, deployment.name) == _expected_listing(rates), deployment.name request_id, charged = _billed(proxy.client, proxy.key, deployment) expected: Final = _expected_cost(rates, deployment.usage) - assert charged == (None if rates is None else pytest.approx(expected, rel=1e-6)), deployment.name + assert charged == pytest.approx(expected, rel=1e-6), deployment.name return _Charge(request_id=request_id, expected=expected) diff --git a/tests/integration/providers/test_bedrock_mantle_claude_chat_wire.py b/tests/integration/providers/test_bedrock_mantle_claude_chat_wire.py index 9c73fcc3b6e..b8bd734ece5 100644 --- a/tests/integration/providers/test_bedrock_mantle_claude_chat_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_claude_chat_wire.py @@ -70,10 +70,18 @@ def _chat_messages(marker: str) -> list[JsonValue]: return [{"role": "user", "content": _prompt(marker)}] -def _outbound(backend: str, messages: list[JsonValue], **fields: JsonValue) -> dict[str, JsonValue]: +def _default_max_tokens(gateway: Gateway, backend: str) -> int: + match gateway.get("/public/litellm_model_cost_map").get(f"bedrock_mantle/{backend}"): + case {"max_output_tokens": int() as mapped}: + return mapped + case _: + return DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS + + +def _outbound(gateway: Gateway, backend: str, messages: list[JsonValue], **fields: JsonValue) -> dict[str, JsonValue]: return { "messages": messages, - "max_tokens": DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS, + "max_tokens": _default_max_tokens(gateway, backend), "anthropic_version": _BRIDGE_VERSION, "model": backend, **fields, @@ -290,7 +298,7 @@ def test_chat_completion_on_a_mantle_claude_id_posts_native_messages_and_spends_ gateway: Gateway, backend: str ) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(backend, [_user_turn(marker)]) + expected: Final = _outbound(gateway, backend, [_user_turn(marker)]) with wire_server(_bearer_peer(expected, _text_reply(marker, backend))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{backend}", api_base=wire.url, api_key=_API_KEY) response: Final = gateway.request( @@ -314,7 +322,7 @@ def test_chat_stream_on_a_mantle_claude_id_relays_openai_chunks_from_the_native_ gateway: Gateway, stream_options: dict[str, JsonValue], usage_chunks: list[tuple[JsonValue, ...]] ) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)], stream=True) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)], stream=True) with wire_server(_bearer_peer(expected, _text_stream(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{_HAIKU}", api_base=wire.url, api_key=_API_KEY) response: Final = gateway.request( @@ -355,7 +363,7 @@ def test_chat_request_fields_are_translated_to_the_native_messages_shape( outbound_fields: dict[str, JsonValue], ) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)], **outbound_fields) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)], **outbound_fields) with wire_server(_bearer_peer(expected, _text_reply(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{_HAIKU}", api_base=wire.url, api_key=_API_KEY) response: Final = gateway.request( @@ -377,7 +385,7 @@ def test_chat_tools_reach_mantle_as_input_schema_and_the_tool_use_block_returns_ marker: Final = uuid4().hex backend: Final = _mantle_claude_where(gateway, "supports_forced_tool_use", disabled=False) expected: Final = _outbound( - backend, [_user_turn(marker)], tools=[_WEATHER_TOOL_OUTBOUND], tool_choice=outbound_choice + gateway, backend, [_user_turn(marker)], tools=[_WEATHER_TOOL_OUTBOUND], tool_choice=outbound_choice ) reply: Final = _tool_reply(marker, backend, "get_weather") with wire_server(_bearer_peer(expected, reply)) as wire, gateway.scenario() as scenario: @@ -407,7 +415,7 @@ def test_chat_tools_reach_mantle_as_input_schema_and_the_tool_use_block_returns_ def test_chat_stream_relays_a_native_tool_use_block_as_tool_call_chunks(gateway: Gateway) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)], tools=[_WEATHER_TOOL_OUTBOUND], stream=True) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)], tools=[_WEATHER_TOOL_OUTBOUND], stream=True) with wire_server(_bearer_peer(expected, _tool_stream(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{_HAIKU}", api_base=wire.url, api_key=_API_KEY) response: Final = gateway.request( @@ -439,7 +447,7 @@ def test_chat_tool_result_turn_reaches_mantle_as_tool_use_and_tool_result_blocks marker: Final = uuid4().hex tool_call: Final = f"toolu_{marker}" expected: Final = _outbound( - _HAIKU, + gateway, _HAIKU, [ _user_turn(marker), { @@ -484,7 +492,7 @@ def test_chat_json_schema_response_format_is_sent_as_a_forced_tool_and_returned_ ) -> None: marker: Final = uuid4().hex expected: Final = _outbound( - _HAIKU, + gateway, _HAIKU, [_user_turn(marker)], tools=[{"name": "json_tool_call", "input_schema": _CITY_SCHEMA}], tool_choice={"name": "json_tool_call", "type": "tool"}, @@ -537,7 +545,7 @@ def _sigv4_peer(expected: Mapping[str, JsonValue], reply: Reply) -> Callable[[Re def test_chat_completion_signs_the_native_messages_request_with_the_deployment_aws_keys(gateway: Gateway) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)]) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)]) with wire_server(_sigv4_peer(expected, _text_reply(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model( model=f"bedrock_mantle/{_HAIKU}", @@ -556,7 +564,7 @@ def test_chat_completion_signs_the_native_messages_request_with_the_deployment_a def test_region_prefix_in_a_mantle_claude_id_is_routing_only_and_keeps_the_mantle_price(gateway: Gateway) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)]) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)]) with wire_server(_bearer_peer(expected, _text_reply(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/us-east-2/{_HAIKU}", api_base=wire.url, api_key=_API_KEY) response: Final = gateway.request( @@ -577,7 +585,7 @@ def test_region_prefix_in_a_mantle_claude_id_is_routing_only_and_keeps_the_mantl def test_openai_sdk_chat_completion_on_a_mantle_claude_id(gateway: Gateway) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)]) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)]) with wire_server(_bearer_peer(expected, _text_reply(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{_HAIKU}", api_base=wire.url, api_key=_API_KEY) with OpenAI( @@ -606,7 +614,7 @@ def test_openai_sdk_chat_completion_on_a_mantle_claude_id(gateway: Gateway) -> N def test_openai_sdk_chat_stream_on_a_mantle_claude_id(gateway: Gateway) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)], stream=True) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)], stream=True) with wire_server(_bearer_peer(expected, _text_stream(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{_HAIKU}", api_base=wire.url, api_key=_API_KEY) with OpenAI( @@ -645,7 +653,7 @@ def test_openai_sdk_chat_stream_on_a_mantle_claude_id(gateway: Gateway) -> None: async def test_async_openai_sdk_chat_completion_on_a_mantle_claude_id(gateway: Gateway) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)]) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)]) with wire_server(_bearer_peer(expected, _text_reply(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{_HAIKU}", api_base=wire.url, api_key=_API_KEY) async with AsyncOpenAI( @@ -674,7 +682,7 @@ async def test_async_openai_sdk_chat_completion_on_a_mantle_claude_id(gateway: G async def test_async_openai_sdk_chat_stream_on_a_mantle_claude_id(gateway: Gateway) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)], stream=True) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)], stream=True) with wire_server(_bearer_peer(expected, _text_stream(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{_HAIKU}", api_base=wire.url, api_key=_API_KEY) async with AsyncOpenAI( @@ -714,7 +722,7 @@ def test_identical_chat_request_is_served_from_the_response_cache_and_logged_as_ gateway: Gateway, ) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)]) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)]) with wire_server(_bearer_peer(expected, _text_reply(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{_HAIKU}", api_base=wire.url, api_key=_API_KEY) request: Final[dict[str, JsonValue]] = {"model": model, "messages": _chat_messages(marker)} @@ -802,7 +810,7 @@ def test_open_weight_mantle_id_keeps_the_openai_compatible_chat_route_and_its_sa ) def test_api_base_with_a_route_suffix_still_posts_to_the_native_messages_path(gateway: Gateway, suffix: str) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)]) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)]) with wire_server(_bearer_peer(expected, _text_reply(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{_HAIKU}", api_base=wire.url + suffix, api_key=_API_KEY) response: Final = gateway.request( @@ -814,7 +822,7 @@ def test_api_base_with_a_route_suffix_still_posts_to_the_native_messages_path(ga def test_responses_api_on_a_mantle_claude_id_bridges_to_native_messages(gateway: Gateway) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)]) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)]) with wire_server(_bearer_peer(expected, _text_reply(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{_HAIKU}", api_base=wire.url, api_key=_API_KEY) response: Final = gateway.request("POST", "/v1/responses", {"model": model, "input": _prompt(marker)}) @@ -839,7 +847,7 @@ def test_responses_api_on_a_mantle_claude_id_bridges_to_native_messages(gateway: def test_responses_api_stream_on_a_mantle_claude_id_bridges_to_the_native_stream(gateway: Gateway) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)], stream=True) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)], stream=True) with wire_server(_bearer_peer(expected, _text_stream(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{_HAIKU}", api_base=wire.url, api_key=_API_KEY) response: Final = gateway.request( @@ -1001,7 +1009,7 @@ def test_unsupported_claude_param_is_dropped_under_drop_params( ) -> None: marker: Final = uuid4().hex backend: Final = _backend_without(gateway, flag) - expected: Final = _outbound(backend, [_user_turn(marker)], **outbound_fields) + expected: Final = _outbound(gateway, backend, [_user_turn(marker)], **outbound_fields) with wire_server(_bearer_peer(expected, _text_reply(marker, backend))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{backend}", api_base=wire.url, api_key=_API_KEY) response: Final = gateway.request( @@ -1119,7 +1127,7 @@ def _wildcard_prefix(gateway: Gateway, scenario: Scenario, api_base: str) -> str def test_uppercase_claude_id_through_a_mantle_wildcard_is_sent_to_native_messages(gateway: Gateway) -> None: marker: Final = uuid4().hex backend: Final = _HAIKU.upper() - expected: Final = _outbound(backend, [_user_turn(marker)]) + expected: Final = _outbound(gateway, backend, [_user_turn(marker)]) with wire_server(_bearer_peer(expected, _text_reply(marker, backend))) as wire, gateway.scenario() as scenario: model: Final = f"{_wildcard_prefix(gateway, scenario, wire.url)}/{backend}" response: Final = gateway.request( @@ -1134,7 +1142,7 @@ def test_five_kilobyte_claude_id_through_a_mantle_wildcard_reaches_mantle_and_it ) -> None: marker: Final = uuid4().hex backend: Final = f"{_HAIKU}-{'x' * 5000}" - expected: Final = _outbound(backend, [_user_turn(marker)]) + expected: Final = _outbound(gateway, backend, [_user_turn(marker)]) unknown: Final = _error_reply(404, "not_found_error", "model: synthetic unknown model") with wire_server(_bearer_peer(expected, unknown)) as wire, gateway.scenario() as scenario: model: Final = f"{_wildcard_prefix(gateway, scenario, wire.url)}/{backend}" diff --git a/tests/integration/providers/test_chat_template_tool_count_wire.py b/tests/integration/providers/test_chat_template_tool_count_wire.py index a8b97fa732e..552df3a87bc 100644 --- a/tests/integration/providers/test_chat_template_tool_count_wire.py +++ b/tests/integration/providers/test_chat_template_tool_count_wire.py @@ -39,7 +39,7 @@ _API_KEY: Final = "scripted-provider-key" _DEEPSEEK_V4: Final = "azure_ai/DeepSeek-V4-Flash-0731" _DEEPSEEK_R1: Final = "azure_ai/DeepSeek-R1-0528" _QWEN3_5: Final = "databricks/databricks-qwen35-122b-a10b" -_GROK: Final = "azure_ai/grok-4-1" +_GROK_3: Final = "azure_ai/grok-3" _DEEPSEEK_V4_TOOL_PREAMBLE: Final = 222 _QWEN3_5_TOOL_PREAMBLE: Final = 196 _QWEN3_5_OPENED_SYSTEM_TURN: Final = 5 @@ -250,7 +250,7 @@ def test_messages_count_tokens_skips_web_search_and_counts_unmapped_tool_keys_fo def test_messages_count_tokens_keeps_openai_tool_framing_for_other_models(gateway: Gateway) -> None: with wire_server(_completion) as wire, gateway.scenario() as scenario: - model: Final = _deployment(scenario, "count", _GROK, wire.url) + model: Final = _deployment(scenario, "count", _GROK_3, wire.url) without_tools: Final = _count(gateway, _MESSAGES_COUNT, _MESSAGES_COUNT.body(model, ())) plain: Final = _count(gateway, _MESSAGES_COUNT, _MESSAGES_COUNT.body(model, (_ANTHROPIC_TOOL,))) deferred: Final = _count( @@ -428,7 +428,7 @@ def test_project_itpm_reserves_deepseek_v4_tools_the_way_its_chat_template_rende def test_project_itpm_keeps_openai_tool_framing_for_other_models(gateway: Gateway) -> None: chat: Final = _ENDPOINTS["chat-completions"] with wire_server(_completion) as wire, gateway.scenario() as scenario: - group: Final = _deployment(scenario, "grok-4-1", _GROK, wire.url) + group: Final = _deployment(scenario, "grok-3", _GROK_3, wire.url) response: Final = gateway.request( "POST", chat.path, @@ -504,8 +504,8 @@ def _agent_body(group: str, marker: str, file_body: str, trigger: int) -> dict[s @pytest.mark.parametrize( ("model", "files_kept"), - ((_DEEPSEEK_V4, (False, False, True)), (_GROK, (True, True, True))), - ids=("deepseek-v4", "grok-4-1"), + ((_DEEPSEEK_V4, (False, False, True)), (_GROK_3, (True, True, True))), + ids=("deepseek-v4", "grok-3"), ) def test_messages_context_management_trigger_counts_tools_the_way_the_models_chat_template_renders_them( gateway: Gateway, model: str, files_kept: tuple[bool, ...] diff --git a/tests/integration/providers/test_databricks_chat_wire.py b/tests/integration/providers/test_databricks_chat_wire.py index b93cebdbc4c..3d9d749eb04 100644 --- a/tests/integration/providers/test_databricks_chat_wire.py +++ b/tests/integration/providers/test_databricks_chat_wire.py @@ -353,9 +353,172 @@ def _shape_reply(identity: str, stream: bool) -> Reply: "/chat/completions", id="A5-responses-list-input", ), + pytest.param( + "json-object-two-blocks", + "/v1/chat/completions", + "databricks-meta-llama-3-3-70b-instruct", + { + "messages": [ + {"role": "user", "content": [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}]} + ], + "response_format": {"type": "json_object"}, + }, + [{"role": "user", "content": "Reply in JSON"}], + False, + "/chat/completions", + id="A6-chat-json-object-two-blocks-joined", + ), + pytest.param( + "json-object-two-blocks-stream", + "/v1/chat/completions", + "databricks-qwen35-122b-a10b", + { + "messages": [ + {"role": "user", "content": [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}]} + ], + "response_format": {"type": "json_object"}, + "stream": True, + }, + [{"role": "user", "content": "Reply in JSON"}], + False, + "/chat/completions", + id="A7-chat-stream-json-object-two-blocks-joined", + ), + pytest.param( + "json-object-history", + "/v1/chat/completions", + "databricks-meta-llama-3-3-70b-instruct", + { + "messages": [ + {"role": "user", "content": [{"type": "text", "text": "Code: A"}, {"type": "text", "text": "BC"}]}, + {"role": "assistant", "content": "Got it."}, + {"role": "user", "content": [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}]}, + ], + "response_format": {"type": "json_object"}, + }, + [ + {"role": "user", "content": [{"type": "text", "text": "Code: A"}, {"type": "text", "text": "BC"}]}, + {"role": "assistant", "content": "Got it."}, + {"role": "user", "content": "Reply in JSON"}, + ], + False, + "/chat/completions", + id="A8-chat-json-object-earlier-turns-kept", + ), + pytest.param( + "json-object-trailing-system", + "/v1/chat/completions", + "databricks-meta-llama-3-3-70b-instruct", + { + "messages": [ + {"role": "user", "content": "Say hi"}, + {"role": "system", "content": [{"type": "text", "text": "Alpha"}, {"type": "text", "text": " rule."}]}, + {"role": "system", "content": [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}]}, + ], + "response_format": {"type": "json_object"}, + }, + [{"role": "user", "content": "Say hi"}, {"role": "system", "content": "Alpha rule.\n\nReply in JSON"}], + False, + "/chat/completions", + id="A8b-chat-json-object-trailing-system-keeps-break", + ), + pytest.param( + "json-object-drop-params-cache-control", + "/v1/chat/completions", + "databricks-meta-llama-3-3-70b-instruct", + { + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Reply"}, + {"type": "text", "text": " in JSON", "cache_control": {"type": "ephemeral"}}, + ], + } + ], + "response_format": {"type": "json_object"}, + "drop_params": True, + }, + [{"role": "user", "content": "Reply in JSON"}], + False, + "/chat/completions", + id="A8c-chat-json-object-drop-params-drops-cache-control", + ), + pytest.param( + "json-object-trailing-thinking-turn", + "/v1/chat/completions", + "databricks-meta-llama-3-3-70b-instruct", + { + "messages": [ + {"role": "user", "content": [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}]}, + { + "role": "assistant", + "content": "", + "thinking_blocks": [{"type": "thinking", "thinking": "Plan.", "signature": "sig"}], + }, + ], + "response_format": {"type": "json_object"}, + }, + [{"role": "user", "content": "Reply in JSON"}], + False, + "/chat/completions", + id="A8d-chat-json-object-trailing-thinking-turn-dropped", + ), + pytest.param( + "json-object-external-endpoint", + "/v1/chat/completions", + "openai-gpt-external", + { + "messages": [ + {"role": "user", "content": [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}]} + ], + "response_format": {"type": "json_object"}, + }, + [{"role": "user", "content": [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}]}], + False, + "/chat/completions", + id="A8e-chat-json-object-external-endpoint-kept-as-sent", + ), + pytest.param( + "json-object-responses-two-parts", + "/v1/responses", + "databricks-qwen35-122b-a10b", + { + "input": [ + { + "role": "user", + "content": [{"type": "input_text", "text": "Reply"}, {"type": "input_text", "text": " in JSON"}], + } + ], + "text": {"format": {"type": "json_object"}}, + }, + [{"role": "user", "content": "Reply in JSON"}], + False, + "/chat/completions", + id="A9-responses-json-object-two-parts-joined", + ), + pytest.param( + "json-object-responses-trailing-reasoning", + "/v1/responses", + "databricks-qwen35-122b-a10b", + { + "input": [ + { + "role": "user", + "content": [{"type": "input_text", "text": "Reply"}, {"type": "input_text", "text": " in JSON"}], + }, + {"type": "reasoning", "id": "rs_1", "summary": [{"type": "summary_text", "text": "Plan."}]}, + ], + "text": {"format": {"type": "json_object"}}, + }, + [{"role": "user", "content": "Reply in JSON"}], + False, + "/chat/completions", + id="A9b-responses-json-object-trailing-reasoning-dropped", + ), ], ) -def test_databricks_non_claude_single_text_block_reaches_upstream_as_string( +def test_databricks_non_claude_plain_text_content_reaches_upstream_as_string( gateway: Gateway, case: str, endpoint: str, @@ -365,7 +528,7 @@ def test_databricks_non_claude_single_text_block_reaches_upstream_as_string( responses_mode: bool, upstream_path: str, ) -> None: - identity: Final = f"databricks-single-text-{case}-{uuid.uuid4().hex}" + identity: Final = f"databricks-plain-text-{case}-{uuid.uuid4().hex}" streaming: Final = request_body.get("stream") is True def respond(request: Request) -> Reply: @@ -1453,7 +1616,7 @@ def test_databricks_non_string_text_block_is_rejected_before_reaching_upstream(g ) received: Final = wire.drain() assert response.status_code == 400, response.text - assert "'int' object has no attribute 'strip'" in response.text + assert "A text content block needs a string 'text', got int" in response.text assert received == () diff --git a/tests/integration/providers/test_decisions_wire.py b/tests/integration/providers/test_decisions_wire.py index 5af1532e9ed..4e091460667 100644 --- a/tests/integration/providers/test_decisions_wire.py +++ b/tests/integration/providers/test_decisions_wire.py @@ -262,11 +262,28 @@ def _upstream_calls(gateway: Gateway, handle: ScenarioHandle) -> list[dict[str, return _calls_to(_observed_requests(gateway), handle) -def _sdk_inference_call(gateway: Gateway, handle: ScenarioHandle, provider: _Provider) -> dict[str, JsonValue]: +def _sdk_inference_call( + gateway: Gateway, + handle: ScenarioHandle, + provider: _Provider, + *, + metadata_authorization: str | None = None, +) -> dict[str, JsonValue]: observed: Final = _upstream_calls(gateway, handle) inference_path: Final = f"/{handle.scenario_id}/v1/systemone" (call,) = tuple(request for request in observed if request["path"] == inference_path) - metadata_calls: Final = tuple(request for request in observed if request["path"] != inference_path) + environment_metadata: Final = { + "path": f"/{handle.scenario_id}/api/show", + "authorization": metadata_authorization, + "body": {"name": provider.body_model}, + "method": "POST", + "api_key": "", + } + metadata_calls: Final = tuple( + request + for request in observed + if request["path"] != inference_path and (metadata_authorization is None or request != environment_metadata) + ) assert metadata_calls in ( (), ( @@ -556,7 +573,9 @@ def test_ollama_decisions_sdk_respects_environment_and_explicit_auth( api_key=api_key, ) assert response.model_dump(mode="json") == _answer_body(provider) - call: Final = _sdk_inference_call(gateway, handle, provider) + call: Final = _sdk_inference_call( + gateway, handle, provider, metadata_authorization=authorization if environment_base else None + ) assert (call["path"], call["authorization"], call["body"]) == ( f"/{handle.scenario_id}/v1/systemone", authorization, diff --git a/tests/integration/providers/test_device_code_login_guard_boot.py b/tests/integration/providers/test_device_code_login_guard_boot.py index 8aba6841cf9..91f502179f3 100644 --- a/tests/integration/providers/test_device_code_login_guard_boot.py +++ b/tests/integration/providers/test_device_code_login_guard_boot.py @@ -28,7 +28,6 @@ from pydantic import JsonValue, TypeAdapter pytestmark: Final = [pytest.mark.timeout(900), pytest.mark.shares_worker_state] _STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") -_DROPPED: Final = re.compile(r"original model: (\S+), ignoring and continuing") _FIXED: Final[Mapping[str, str]] = MappingProxyType( {"chatgpt-fixed": "chatgpt/gpt-5.6-terra", "copilot-fixed": "github_copilot/gpt-5.4"} ) @@ -41,8 +40,6 @@ _CONTROL_CALLS_BEFORE_THE_RELOAD: Final = 1 _SHAPES: Final[Mapping[str, bool]] = MappingProxyType( {"docs-shapes": False, "docs-shapes-beside-a-universal-wildcard": True} ) -_NOT_FOUND: Final = "Invalid model name passed in model=" -_NO_HEALTHY: Final = "There are no healthy deployments for this model" _CONTROL: Final = "device-login-boot-control" _EXTRA: Final[Mapping[str, JsonValue]] = MappingProxyType({"num_retries": 0, "cache": {"no-cache": True}}) _CALL_ID: Final = "x-litellm-call-id" @@ -174,15 +171,11 @@ def _lists_login_models(listed: tuple[str, ...]) -> bool: ) -def _every_worker_lists( - owned: OwnedProxy, workers: tuple[int, ...], *, login_models: bool -) -> Mapping[int, tuple[str, ...]]: +def _every_worker_lists(owned: OwnedProxy, workers: tuple[int, ...]) -> Mapping[int, tuple[str, ...]]: def settled(by_worker: Mapping[int, tuple[str, ...]]) -> bool: if set(by_worker) != set(workers): return False - if login_models: - return all(_lists_login_models(listed) for listed in by_worker.values()) - return all(_login_names_listed(listed) == frozenset() for listed in by_worker.values()) + return all(_lists_login_models(listed) for listed in by_worker.values()) return eventually(lambda: _listing_by_worker(owned, workers), settled, seconds=90) @@ -212,16 +205,8 @@ def _reload(owned: OwnedProxy) -> None: assert response.status_code == 200, response.text -def _refusal_before_the_reload(name: str, universal: bool) -> str: - if name in _FIXED: - return _NO_HEALTHY - if not universal: - return _NOT_FOUND - return dl.CHATGPT_REFUSAL if name.startswith("chatgpt/") else dl.COPILOT_REFUSAL - - @pytest.mark.parametrize("universal", _SHAPES.values(), ids=_SHAPES.keys()) -def test_booting_without_tokens_drops_the_login_deployments_and_a_reload_after_mounting_brings_them_back( +def test_booting_without_tokens_lists_login_deployments_and_serves_only_after_tokens_are_mounted( gateway: Gateway, tmp_path: Path, universal: bool ) -> None: dirs: Final = _token_dirs(tmp_path) @@ -240,20 +225,18 @@ def test_booting_without_tokens_drops_the_login_deployments_and_a_reload_after_m workers=2, ) as owned: workers: Final = _workers(owned) - _every_worker_lists(owned, workers, login_models=False) + _every_worker_lists(owned, workers) boot_log: Final = owned.log.read_text() assert peers.auth_connections() == () assert dl.CHATGPT_USER_CODE not in boot_log and dl.COPILOT_USER_CODE not in boot_log, boot_log[-3000:] - dropped_models: Final = _DROPPED.findall(boot_log) - expected_drops: Final = frozenset((*_FIXED.values(), *_WILDCARDS)) - assert frozenset(dropped_models) == expected_drops, boot_log[-3000:] - assert all(dropped_models.count(model) >= len(workers) for model in expected_drops), boot_log[-3000:] with owned.gateway.scenario() as scenario: key: Final = scenario.key() for name in _LOGIN_NAMES: refused: Final = _chat(_base_url(owned), key, name, uuid.uuid4().hex) assert refused.status == 400, refused.text - assert _refusal_before_the_reload(name, universal) in refused.message(), refused.text + refusal: Final = dl.CHATGPT_REFUSAL if name.startswith("chatgpt") else dl.COPILOT_REFUSAL + assert refusal in refused.message(), refused.text + assert peers.api.drain() == () control: Final = _chat(_base_url(owned), key, _CONTROL, uuid.uuid4().hex) assert control.status == 200, control.text assert peers.auth_connections() == () @@ -261,7 +244,7 @@ def test_booting_without_tokens_drops_the_login_deployments_and_a_reload_after_m dirs.store_valid_tokens(peers.api.url) _reload(owned) - _every_worker_lists(owned, workers, login_models=True) + _every_worker_lists(owned, workers) for served_logins, name in enumerate(_LOGIN_NAMES, start=1): served: Final = _chat(_base_url(owned), key, name, uuid.uuid4().hex) assert served.status == 200, served.text diff --git a/tests/integration/providers/test_device_code_login_guard_wire.py b/tests/integration/providers/test_device_code_login_guard_wire.py index b4044674ed5..39e36fd95fa 100644 --- a/tests/integration/providers/test_device_code_login_guard_wire.py +++ b/tests/integration/providers/test_device_code_login_guard_wire.py @@ -746,10 +746,7 @@ def test_a_recent_device_code_request_does_not_hold_the_worker_for_the_rest_of_t assert scene.rig.peers.auth_connections() == () -_NOT_LIVE: Final = re.compile(r"saved to the database, but the model id\(s\) \['([0-9a-f-]+)'\] are not live") - - -def test_model_created_without_a_token_is_saved_not_live_and_serves_once_the_token_is_mounted(scene: _Scene) -> None: +def test_model_created_without_a_token_refuses_requests_and_serves_once_the_token_is_mounted(scene: _Scene) -> None: gateway: Final = scene.rig.owned.gateway name: Final = f"device-login-db-{uuid.uuid4().hex}" identity: Final = str(uuid.uuid4()) @@ -763,14 +760,18 @@ def test_model_created_without_a_token_is_saved_not_live_and_serves_once_the_tok }, ) try: - assert created.status_code == 500, created.text - not_live: Final = _NOT_LIVE.search(created.text) - assert not_live is not None and not_live.group(1) == identity, created.text + assert created.status_code == 200, created.text + assert created.json()["model_info"]["id"] == identity, created.text assert scene.rig.peers.auth_connections() == () - scene.store_valid_token("chatgpt") call: Final = _Call( scene.rig.base_url(), scene.key, name, "responses", False, "httpx", uuid.uuid4().hex, SimpleQueue() ) + refused: Final = _serve(call) + assert refused.status == 400, refused.text + assert dl.CHATGPT_REFUSAL in _error_message(refused.text), refused.text + assert scene.rig.peers.auth_connections() == () + assert scene.rig.peers.api.drain() == () + scene.store_valid_token("chatgpt") served: Final = eventually(lambda: _serve(call), lambda answer: answer.status == 200, seconds=90) assert set(rv.MARKER.findall(served.text)) == {call.marker}, served.text assert scene.rig.peers.auth_connections() == () diff --git a/tests/integration/sdk/test_anthropic_json_tool_cut_stream_sync_sdk.py b/tests/integration/sdk/test_anthropic_json_tool_cut_stream_sync_sdk.py index da091b1c12e..54a6df02edf 100644 --- a/tests/integration/sdk/test_anthropic_json_tool_cut_stream_sync_sdk.py +++ b/tests/integration/sdk/test_anthropic_json_tool_cut_stream_sync_sdk.py @@ -10,6 +10,7 @@ from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue import litellm +from litellm.exceptions import MidStreamFallbackError from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper _JSON_TOOL: Final = "json_tool_call" @@ -122,12 +123,10 @@ def _credentials(model: str) -> dict[str, str]: @pytest.mark.parametrize("model", [_ANTHROPIC_MODEL, _BEDROCK_INVOKE_MODEL], ids=["anthropic", "bedrock_invoke"]) @pytest.mark.parametrize( ("shape", "cut"), - [("plain", False), ("plain", True), ("wrapped", True)], - ids=["plain_complete", "plain_cut", "wrapped_cut"], + [("plain", False), ("wrapped", False), ("plain", True), ("wrapped", True)], + ids=["plain_complete", "wrapped_complete", "plain_cut", "wrapped_cut"], ) -def test_sync_stream_cut_before_the_json_tool_block_closes_still_yields_the_whole_answer( - model: str, shape: str, cut: bool -) -> None: +def test_sync_json_tool_stream_requires_a_provider_finish_reason(model: str, shape: str, cut: bool) -> None: marker_hex: Final = uuid.uuid4().hex with wire_server(_peer(marker_hex, shape, cut)) as wire: response: Final = litellm.completion( @@ -141,6 +140,11 @@ def test_sync_stream_cut_before_the_json_tool_block_closes_still_yields_the_whol **_credentials(model), ) assert isinstance(response, CustomStreamWrapper), type(response) + if cut: + with pytest.raises(MidStreamFallbackError, match="Stream ended before a finish_reason was received"): + tuple(response) + assert len(wire.drain()) == 1 + return text: Final = "".join(str(chunk.choices[0].delta.content or "") for chunk in tuple(response)) assert len(wire.drain()) == 1 assert json.loads(text) == _payload(marker_hex), text diff --git a/tests/integration/security/test_files_batches_access.py b/tests/integration/security/test_files_batches_access.py index 119055c64c7..93ccc742d58 100644 --- a/tests/integration/security/test_files_batches_access.py +++ b/tests/integration/security/test_files_batches_access.py @@ -25,7 +25,7 @@ from cryptography.hazmat.primitives import serialization from cryptography.hazmat.primitives.asymmetric import rsa from integration._support.client import Gateway, Scenario, eventually, object_value, string_value from integration._support.database import read_rows -from integration._support.process import owned_proxy +from integration._support.process import OwnedState, owned_proxy from integration._support.tls import server_context, write_self_signed_cert from integration._support.wire import Reply, Request, Wire, wire_server from openai import OpenAI, PermissionDeniedError @@ -250,15 +250,20 @@ def test_stranger_cannot_read_download_or_delete_another_teams_managed_file(gate assert _drained_other_than_model_list_probes(wire) == () -def test_stranger_cannot_retrieve_or_cancel_another_teams_managed_batch(gateway: Gateway) -> None: +@pytest.mark.timeout(180) +def test_stranger_cannot_retrieve_or_cancel_another_teams_managed_batch(gateway: Gateway, tmp_path: Path) -> None: key: Final = "provider-key-" + uuid.uuid4().hex[:8] - with wire_server(_openai_backend()) as wire, gateway.scenario() as scenario: + with ( + owned_proxy(gateway, tmp_path, {"PROXY_BATCH_POLLING_ENABLED": "false"}, state=OwnedState.PRIVATE) as candidate, + wire_server(_openai_backend()) as wire, + candidate.scenario() as scenario, + ): model: Final = scenario.model(model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key=key) - _wait_until_every_worker_serves(gateway, model) + _wait_until_every_worker_serves(candidate, model) owner: Final = _member(scenario, model) stranger: Final = _member(scenario, model) - uploaded: Final = gateway.request_multipart( + uploaded: Final = candidate.request_multipart( "/v1/files", {"purpose": "batch", "target_model_names": model}, {"file": (INPUT_FILENAME, _jsonl(model), "application/jsonl")}, @@ -266,7 +271,7 @@ def test_stranger_cannot_retrieve_or_cancel_another_teams_managed_batch(gateway: ) assert uploaded.status_code == 200, uploaded.text managed_input: Final = string_value(_json(uploaded)["id"]) - created: Final = gateway.request( + created: Final = candidate.request( "POST", "/v1/batches", {"input_file_id": managed_input, "endpoint": "/v1/chat/completions", "completion_window": "24h"}, @@ -277,13 +282,13 @@ def test_stranger_cannot_retrieve_or_cancel_another_teams_managed_batch(gateway: wire.drain() expected_message: Final = f"User {stranger.user} does not have access to the object {unified_batch}" - stranger_get: Final = gateway.request("GET", f"/v1/batches/{unified_batch}", key=stranger.key) + stranger_get: Final = candidate.request("GET", f"/v1/batches/{unified_batch}", key=stranger.key) assert _error_message(stranger_get, 403) == expected_message, stranger_get.text - stranger_cancel: Final = gateway.request("POST", f"/v1/batches/{unified_batch}/cancel", {}, key=stranger.key) + stranger_cancel: Final = candidate.request("POST", f"/v1/batches/{unified_batch}/cancel", {}, key=stranger.key) assert _error_message(stranger_cancel, 403) == expected_message, stranger_cancel.text assert _drained_other_than_model_list_probes(wire) == (), "stranger batch calls must not reach the provider" - owner_cancel: Final = gateway.request("POST", f"/v1/batches/{unified_batch}/cancel", {}, key=owner.key) + owner_cancel: Final = candidate.request("POST", f"/v1/batches/{unified_batch}/cancel", {}, key=owner.key) assert owner_cancel.status_code == 200, owner_cancel.text assert string_value(_json(owner_cancel)["id"]) == unified_batch, owner_cancel.text (cancel_request,) = _drained_other_than_model_list_probes(wire) diff --git a/tests/unit/caching/test_caching.py b/tests/unit/caching/test_caching.py index 0f6ef2304fb..0911d5ef888 100644 --- a/tests/unit/caching/test_caching.py +++ b/tests/unit/caching/test_caching.py @@ -472,6 +472,30 @@ def test_cache_key_tells_apart_requests_that_differ_only_in_continuation_token() assert cache.get_cache_key(**request, continuation_token="") == first_segment_key +def test_cache_key_tells_apart_requests_that_ask_for_token_ids() -> None: + cache: Final = Cache(type=LiteLLMCacheType.LOCAL) + request: Final = {"model": "hosted_vllm/qwen3", "messages": [{"role": "user", "content": "count"}]} + + plain_key: Final = cache.get_cache_key(**request) + + assert cache.get_cache_key(**request, return_token_ids=True) != plain_key + assert cache.get_cache_key(**request, return_token_ids=False) == plain_key + assert cache.get_cache_key(**request, return_token_ids=None) == plain_key + + +def test_cache_key_tells_apart_requests_that_ask_for_token_ids_through_extra_body() -> None: + cache: Final = Cache(type=LiteLLMCacheType.LOCAL) + request: Final = {"model": "hosted_vllm/qwen3", "messages": [{"role": "user", "content": "count"}]} + + plain_key: Final = cache.get_cache_key(**request) + other_extra_body_key: Final = cache.get_cache_key(**request, extra_body={"top_k": 5}) + + assert cache.get_cache_key(**request, extra_body={"return_token_ids": True}) != plain_key + assert cache.get_cache_key(**request, extra_body={"top_k": 5, "return_token_ids": True}) != other_extra_body_key + assert cache.get_cache_key(**request, extra_body={"return_token_ids": False}) == plain_key + assert cache.get_cache_key(**request, extra_body={"continuation_token": "abc"}) != plain_key + + def test_cache_key_tells_apart_requests_that_differ_only_in_tool_call_id_shape() -> None: cache: Final = Cache(type=LiteLLMCacheType.LOCAL) request: Final = { diff --git a/tests/unit/integrations/test_prometheus_labels.py b/tests/unit/integrations/test_prometheus_labels.py index 878f19c66f8..04eb7c3a5a4 100644 --- a/tests/unit/integrations/test_prometheus_labels.py +++ b/tests/unit/integrations/test_prometheus_labels.py @@ -5,6 +5,7 @@ Unit tests for prometheus metric labels configuration from typing import Final, Literal import pytest +from pydantic import JsonValue import litellm from litellm.types.integrations.prometheus import ( @@ -802,6 +803,47 @@ async def test_failure_hook_emits_api_provider_value_on_failed_requests_metric() _clear_prometheus_registry() +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("metadata_id", "request_id", "logged_id", "expected"), + ( + (None, None, "logged-deployment", "logged-deployment"), + ("", None, "logged-deployment", "logged-deployment"), + ("metadata-deployment", "request-deployment", "logged-deployment", "metadata-deployment"), + (None, "request-deployment", "logged-deployment", "request-deployment"), + (None, None, None, "None"), + ), +) +async def test_failure_counters_keep_logged_deployment_attribution( + metadata_id: str | None, request_id: str | None, logged_id: str | None, expected: str +) -> None: + from fastapi import HTTPException + + from litellm.integrations.prometheus import PrometheusLogger + from litellm.proxy._types import UserAPIKeyAuth + + request_data: Final[dict[str, JsonValue]] = { + "model": "metric-test-alias", + "metadata": {"model_info": {"id": metadata_id}}, + "model_info": {"id": request_id}, + "standard_logging_object": {"model_id": logged_id, "custom_llm_provider": "openai"}, + } + _clear_prometheus_registry() + try: + await PrometheusLogger().async_post_call_failure_hook( + request_data=request_data, + original_exception=HTTPException(status_code=400, detail="guardrail rejected"), + user_api_key_dict=UserAPIKeyAuth(token="metric-test-token"), + ) + for metric_name in ("litellm_proxy_failed_requests_metric_total", "litellm_proxy_total_requests_metric_total"): + samples: Final = _collected_samples(metric_name) + assert [(sample.labels["model_id"], sample.value) for sample in samples] == [ + (expected, 1.0) + ] + finally: + _clear_prometheus_registry() + + @pytest.mark.asyncio async def test_failure_hook_labels_a_cli_session_with_the_per_user_alias_not_the_login_token(): from litellm.integrations.prometheus import PrometheusLogger diff --git a/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py index d0f5bfc6b9f..3e197740e98 100644 --- a/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py +++ b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py @@ -3089,6 +3089,73 @@ def test_convert_to_model_response_object_function_output(): assert result._hidden_params is not None +@pytest.mark.parametrize( + "vllm_choice_fields", + [{"stop_reason": ", 5", "token_ids": [16, 11, 220]}, {"stop_reason": 151645}, {"stop_reason": None, "token_ids": None}], + ids=["stop-string-with-token-ids", "stop-token", "natural-stop"], +) +def test_vllm_choice_fields_stay_on_the_choice_through_a_cache_round_trip( + vllm_choice_fields: dict[str, object], +) -> None: + upstream: Final = { + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 1, + "model": "Qwen/Qwen3-0.6B", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "1, 2, 3, 4"}, + "finish_reason": "stop", + **vllm_choice_fields, + } + ], + } + response: Final = convert_to_model_response_object(response_object=upstream, model_response_object=ModelResponse()) + assert isinstance(response, ModelResponse) + cache_hit: Final = convert_to_model_response_object( + response_object=json.loads(response.model_dump_json()), model_response_object=ModelResponse() + ) + assert isinstance(cache_hit, ModelResponse) + + dumped_choices: Final = [json.loads(rebuilt.model_dump_json())["choices"][0] for rebuilt in (response, cache_hit)] + assert [{field: choice.get(field, "missing") for field in vllm_choice_fields} for choice in dumped_choices] == [ + vllm_choice_fields, + vllm_choice_fields, + ] + + +@pytest.mark.parametrize( + "vllm_choice_fields", + [{"stop_reason": ", 5", "token_ids": [16, 11, 220]}, {"stop_reason": None, "token_ids": None}], + ids=["stop-string-with-token-ids", "natural-stop"], +) +def test_vllm_choice_fields_cached_only_under_provider_specific_fields_come_back_on_the_choice( + vllm_choice_fields: dict[str, object], +) -> None: + cached_before_upgrade: Final = { + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 1, + "model": "Qwen/Qwen3-0.6B", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "1, 2, 3, 4"}, + "finish_reason": "stop", + "provider_specific_fields": vllm_choice_fields, + } + ], + } + cache_hit: Final = convert_to_model_response_object( + response_object=cached_before_upgrade, model_response_object=ModelResponse() + ) + assert isinstance(cache_hit, ModelResponse) + + dumped_choice: Final = json.loads(cache_hit.model_dump_json())["choices"][0] + assert {field: dumped_choice.get(field, "missing") for field in vllm_choice_fields} == vllm_choice_fields + + @pytest.mark.parametrize( "reported_cost", [ diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_chat_template_tools.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_chat_template_tools.py index 51249e01d4d..a8162860684 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_chat_template_tools.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_chat_template_tools.py @@ -9,7 +9,10 @@ _BARE_FUNCTION: Final = {"name": "g", "parameters": {}} # Tool schemas as rendered by render_tools in encoding/encoding_dsv4.py (deepseek-ai/DeepSeek-V4-Flash) and by -# chat_template.jinja (Qwen/Qwen3.5-122B-A10B), both read on Hugging Face on 2026-10-17 +# chat_template.jinja (Qwen/Qwen3.5-122B-A10B), both read on Hugging Face on 2026-10-17, and by chat_template.jinja +# (deepseek-ai/DeepSeek-V4.1-Flash) read on 2026-10-29. Grok's template is closed, so its two renderings are the ones +# whose token counts left a constant preamble across seven tool shapes billed by Azure AI Foundry +# grok-4-1-fast-reasoning and xAI grok-4.3 on 2026-10-29 @pytest.mark.parametrize( ("model", "rendered"), [ @@ -17,6 +20,18 @@ _BARE_FUNCTION: Final = {"name": "g", "parameters": {}} "DeepSeek-V4-Flash-0731", '{"name": "f", "description": "é", "parameters": {}}\n{"name": "g", "parameters": {}}', ), + ( + "DeepSeek-V4.1-Flash", + '{"name": "f", "description": "é", "parameters": {}}\n{"name": "g", "parameters": {}}', + ), + ( + "grok-4-1-fast-reasoning", + '{"name":"f","description":"é","parameters":{}}\n{"name":"g","parameters":{}}', + ), + ( + "grok-4.3", + '{"name": "f", "description": "é", "parameters": {}}\n{"name": "g", "parameters": {}}', + ), ( "databricks-qwen35-122b-a10b", '\n{"type": "function", "function": {"name": "f", "description": "é", "parameters": {}}}' diff --git a/tests/unit/litellm_core_utils/test_redact_messages.py b/tests/unit/litellm_core_utils/test_redact_messages.py index 9e385c588ae..0aee61d8365 100644 --- a/tests/unit/litellm_core_utils/test_redact_messages.py +++ b/tests/unit/litellm_core_utils/test_redact_messages.py @@ -561,6 +561,83 @@ class TestPerformRedaction: {"text": "redacted-by-litellm", "finish_reason": "length", "index": 1}, ] + def test_redacts_token_ids_that_decode_to_the_completion(self): + details = { + "stream": True, + "complete_streaming_response": litellm.ModelResponse( + choices=[ + litellm.Choices(message=litellm.Message(content="secret", role="assistant"), token_ids=[20, 11]) + ] + ), + "standard_logging_object": { + "response": { + "choices": [ + { + "token_ids": [20, 11], + "message": {"content": "secret", "provider_specific_fields": {"token_ids": [20, 11]}}, + }, + {"text": " Paris.", "token_ids": [20, 11], "index": 1}, + ] + } + }, + } + + perform_redaction(details, None) + + logged_choices = details["standard_logging_object"]["response"]["choices"] + assert [choice["token_ids"] for choice in logged_choices] == [None, None] + assert logged_choices[0]["message"]["provider_specific_fields"]["token_ids"] is None + assert details["complete_streaming_response"].choices[0].get("token_ids") is None + + def test_redacts_text_completion_stream_text_and_the_token_ids_that_decode_to_it(self): + details = { + "stream": True, + "async_complete_streaming_response": litellm.TextCompletionResponse( + choices=[ + litellm.utils.TextChoices( + index=0, text=" Paris.", reasoning_content="thinking", token_ids=[20, 11], finish_reason="stop" + ) + ] + ), + } + + perform_redaction(details, None) + + logged_choice = details["async_complete_streaming_response"].choices[0] + assert (logged_choice.text, logged_choice.reasoning_content) == ("redacted-by-litellm", "redacted-by-litellm") + assert logged_choice.get("token_ids") is None + assert logged_choice.finish_reason == "stop" + + def test_redacts_token_ids_a_cache_hit_keeps_on_the_choice_provider_specific_fields(self): + result = litellm.ModelResponse( + choices=[ + litellm.Choices( + message=litellm.Message(content="secret", role="assistant"), + provider_specific_fields={"token_ids": [20, 11], "stop_reason": "END"}, + ) + ] + ) + details = { + "standard_logging_object": { + "response": { + "choices": [ + { + "message": {"content": "secret"}, + "provider_specific_fields": {"token_ids": [20, 11], "stop_reason": "END"}, + } + ] + } + } + } + + redacted = perform_redaction(details, result) + + logged_fields = details["standard_logging_object"]["response"]["choices"][0]["provider_specific_fields"] + assert logged_fields == {"token_ids": None, "stop_reason": "END"} + assert redacted.choices[0].provider_specific_fields["token_ids"] is None + assert redacted.choices[0].provider_specific_fields["stop_reason"] == "END" + assert result.choices[0].provider_specific_fields["token_ids"] == [20, 11] + def test_redacts_object_choices_inside_model_response_dict(self): result = { "choices": [ diff --git a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py index bbf1897b3a7..527c83d5760 100644 --- a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -2757,6 +2757,95 @@ def test_stream_chunk_builder_carries_the_bedrock_guardrail_trace_and_records_it assert entry["guardrail_cost"] == pytest.approx(2 * 0.00015) +def _vllm_chat_chunk(index: int, content: str, finish_reason: str | None = None, **fields: object) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-1", + created=1, + model="qwen3", + choices=[StreamingChoices(index=index, delta=Delta(content=content), finish_reason=finish_reason, **fields)], + ) + + +@pytest.mark.parametrize("finish_content", ["", "4"], ids=["finish-only", "content-then-finish"]) +def test_stream_chunk_builder_keeps_the_vllm_stop_reason_and_every_token_id_of_each_choice( + finish_content: str, +) -> None: + built: Final = stream_chunk_builder( + chunks=[ + _vllm_chat_chunk(0, "1,", stop_reason=None, token_ids=[16, 11]), + _vllm_chat_chunk(1, "7,", stop_reason=None, token_ids=[22, 11]), + _vllm_chat_chunk(0, "2,3", stop_reason=None, token_ids=[17, 11, 18]), + _vllm_chat_chunk(0, finish_content, "stop", stop_reason=",5", token_ids=[19]), + ] + ) + + assert built is not None + assert [ + {field: choice[field] for field in ("index", "stop_reason", "token_ids")} + for choice in built.model_dump()["choices"] + ] == [ + {"index": 0, "stop_reason": ",5", "token_ids": [16, 11, 17, 11, 18, 19]}, + {"index": 1, "stop_reason": None, "token_ids": [22, 11]}, + ] + + +def test_stream_chunk_builder_keeps_the_vllm_fields_of_a_sole_choice_streamed_on_a_nonzero_index() -> None: + built: Final = stream_chunk_builder( + chunks=[ + _vllm_chat_chunk(1, "7,", stop_reason=None, token_ids=[22, 11]), + _vllm_chat_chunk(1, "8", "stop", stop_reason=",9", token_ids=[23]), + ] + ) + + assert built is not None + assert {field: built.model_dump()["choices"][0].get(field) for field in ("stop_reason", "token_ids")} == { + "stop_reason": ",9", + "token_ids": [22, 11, 23], + } + + +def test_stream_chunk_builder_keeps_null_vllm_fields_and_adds_none_for_other_providers() -> None: + vllm_built: Final = stream_chunk_builder( + chunks=[ + _vllm_chat_chunk(0, "hi", stop_reason=None, token_ids=None), + _vllm_chat_chunk(0, "", "stop", stop_reason=None, token_ids=None), + ] + ) + other_built: Final = stream_chunk_builder( + chunks=[_vllm_chat_chunk(0, "hi"), _vllm_chat_chunk(0, "", "stop")] + ) + + assert vllm_built is not None and other_built is not None + assert { + field: vllm_built.model_dump()["choices"][0][field] for field in ("stop_reason", "token_ids") + } == {"stop_reason": None, "token_ids": None} + assert {"stop_reason", "token_ids"}.isdisjoint(other_built.model_dump()["choices"][0]) + + +def test_stream_chunk_builder_text_completion_keeps_each_choices_vllm_stop_reason_and_token_ids() -> None: + def text_chunk(index: int, text: str, **fields: object) -> TextCompletionResponse: + return TextCompletionResponse( + id="cmpl-1", + object="text_completion", + created=1, + model="qwen3", + choices=[TextChoices(index=index, text=text, **fields)], + ) + + built: Final = stream_chunk_builder_text_completion( + chunks=[ + text_chunk(0, "1,", token_ids=[16, 11], stop_reason=None), + text_chunk(1, "7,", token_ids=[22, 11], stop_reason=None), + text_chunk(0, "2", token_ids=[17], stop_reason=", 3"), + text_chunk(1, "8", token_ids=[23], stop_reason=151645), + ] + ) + + assert [ + (choice["index"], choice["stop_reason"], choice["token_ids"]) for choice in built.model_dump()["choices"] + ] == [(0, ", 3", [16, 11, 17]), (1, 151645, [22, 11, 23])] + + @pytest.mark.parametrize( ("base_model", "counted_as"), [("azure/gpt-6.1-sol", "azure/gpt-6.1-sol"), ("claude-sonnet-4-5", "my-gpt-6")] ) diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index da88f1358ad..cf8de9c40f5 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -10,6 +10,7 @@ import threading import time from collections.abc import Mapping, Sequence from dataclasses import dataclass +from itertools import chain from pathlib import Path from types import CodeType, FrameType, MappingProxyType, TracebackType from typing import AsyncIterator, Callable, Final, Iterator, Literal, NoReturn, Optional, cast @@ -27,6 +28,10 @@ from litellm.exceptions import MidStreamFallbackError from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.get_litellm_params import get_litellm_params from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( + convert_to_streaming_response, + convert_to_streaming_response_async, +) from litellm.litellm_core_utils.streaming_handler import ( CustomStreamWrapper, _ProviderChunkEarlyReturn, @@ -34,6 +39,7 @@ from litellm.litellm_core_utils.streaming_handler import ( ) from litellm.llms.anthropic.chat.handler import ModelResponseIterator as AnthropicModelResponseIterator from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.llms.hosted_vllm.chat.transformation import HostedVLLMChatConfig from litellm.llms.openai.chat.gpt_transformation import OpenAIChatCompletionStreamingHandler from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( ModelResponseIterator as GeminiModelResponseIterator, @@ -6971,6 +6977,327 @@ def test_chunk_creator_keeps_upstream_index_on_single_choice_chunk(logging_obj: assert result.choices[0].delta.content == "hello" +_TOKEN_IDS: Final = TypeAdapter(list[int] | None) +_CHOICE_INDEX: Final = TypeAdapter(int) +_DELTA: Final = TypeAdapter(Mapping[str, object]) + + +def _joined_token_ids(choices: Sequence[Mapping[str, object]]) -> tuple[int, ...]: + return tuple(chain.from_iterable(_TOKEN_IDS.validate_python(choice.get("token_ids")) or () for choice in choices)) + + +def _hosted_vllm_stream( + logging_obj: Logging, upstream_choices: Sequence[Mapping[str, object]] +) -> CustomStreamWrapper: + sse_lines: Final = ( + *( + "data: " + json.dumps({"id": "chatcmpl-1", "object": "chat.completion.chunk", "model": "m", "choices": [c]}) + for c in upstream_choices + ), + "data: [DONE]", + ) + completion_stream: Final = HostedVLLMChatConfig().get_model_response_iterator( + streaming_response=iter(sse_lines), sync_stream=True, json_mode=False + ) + return CustomStreamWrapper( + completion_stream=completion_stream, model="m", custom_llm_provider="hosted_vllm", logging_obj=logging_obj + ) + + +def _hosted_vllm_stream_choices( + logging_obj: Logging, upstream_choices: Sequence[Mapping[str, object]] +) -> tuple[dict[str, object], ...]: + return tuple(chunk.choices[0].model_dump() for chunk in _hosted_vllm_stream(logging_obj, upstream_choices)) + + +@pytest.mark.parametrize("stop_reason", [",5", 151645]) +@pytest.mark.parametrize("final_content", [",3", ""], ids=["content-then-finish", "finish-only"]) +def test_hosted_vllm_stream_keeps_stop_reason_and_every_token_id( + logging_obj: Logging, final_content: str, stop_reason: str | int +) -> None: + emitted: Final = _hosted_vllm_stream_choices( + logging_obj, + ( + {"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}, + {"index": 0, "delta": {"content": ""}, "finish_reason": None, "token_ids": [16]}, + {"index": 0, "delta": {"content": "1,2"}, "finish_reason": None, "token_ids": [11, 17]}, + { + "index": 0, + "delta": {"content": final_content}, + "finish_reason": "stop", + "stop_reason": stop_reason, + "token_ids": [20], + }, + ), + ) + + assert [choice.get("stop_reason") for choice in emitted if choice["finish_reason"] is not None] == [stop_reason] + assert [ + (_DELTA.validate_python(choice["delta"]).get("content"), choice["stop_reason"]) + for choice in emitted + if choice["finish_reason"] is None and choice.get("stop_reason") is not None + ] == ([(final_content, stop_reason)] if final_content else []) + assert _joined_token_ids(emitted) == (16, 11, 17, 20) + + +def test_hosted_vllm_stream_sends_empty_content_on_token_id_only_chunks(logging_obj: Logging) -> None: + emitted: Final = _hosted_vllm_stream_choices( + logging_obj, + ( + {"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}, + {"index": 0, "delta": {"content": ""}, "finish_reason": None, "token_ids": [11]}, + {"index": 0, "delta": {"content": ""}, "finish_reason": None, "token_ids": [220]}, + {"index": 0, "delta": {"content": "1"}, "finish_reason": None, "token_ids": [16]}, + ), + ) + + assert [ + (_DELTA.validate_python(choice["delta"]).get("content"), choice.get("token_ids")) + for choice in emitted + if choice.get("token_ids") + ] == [("", [11]), ("", [220]), ("1", [16])] + + +def test_hosted_vllm_stream_keeps_held_back_token_ids_on_their_own_choice(logging_obj: Logging) -> None: + emitted: Final = _hosted_vllm_stream_choices( + logging_obj, + ( + {"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}, + {"index": 1, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}, + {"index": 1, "delta": {"content": ""}, "finish_reason": None, "token_ids": [16]}, + {"index": 0, "delta": {"content": "1"}, "finish_reason": None, "token_ids": [17]}, + {"index": 1, "delta": {"content": "2"}, "finish_reason": None, "token_ids": [18]}, + ), + ) + + assert [(choice["index"], choice.get("token_ids")) for choice in emitted if choice.get("token_ids")] == [ + (1, [16]), + (0, [17]), + (1, [18]), + ] + + +def _rejects_a_two(text: str) -> bool: + return "2" not in text + + +def _choices_sent_before_rule_failure(stream: CustomStreamWrapper) -> Iterator[dict[str, object]]: + try: + yield from (chunk.choices[0].model_dump() for chunk in stream) + except MidStreamFallbackError: + return + + +def test_hosted_vllm_stream_sends_no_token_ids_ahead_of_a_failing_post_call_rule( + logging_obj: Logging, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "post_call_rules", [_rejects_a_two]) + stream: Final = _hosted_vllm_stream( + logging_obj, + ( + {"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}, + {"index": 0, "delta": {"content": ""}, "finish_reason": None, "token_ids": [16]}, + {"index": 0, "delta": {"content": "1"}, "finish_reason": None, "token_ids": [11, 17]}, + {"index": 0, "delta": {"content": ",2"}, "finish_reason": "stop", "stop_reason": None, "token_ids": [20]}, + ), + ) + + assert [ + (_DELTA.validate_python(choice["delta"]).get("content"), choice.get("token_ids")) + for choice in _choices_sent_before_rule_failure(stream) + ] == [("", None), ("1", None)] + + +def _held_back_stop_string_stream(token_ids: Mapping[str, None]) -> tuple[Mapping[str, object], ...]: + return ( + {"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}, + {"index": 0, "delta": {"content": ""}, "finish_reason": None, **token_ids}, + {"index": 0, "delta": {"content": "1,2"}, "finish_reason": None, **token_ids}, + {"index": 0, "delta": {"content": ""}, "finish_reason": "stop", "stop_reason": ",3", **token_ids}, + ) + + +def test_hosted_vllm_stream_with_null_token_ids_emits_the_same_chunks_as_without_them(logging_obj: Logging) -> None: + with_null_ids: Final = _hosted_vllm_stream_choices(logging_obj, _held_back_stop_string_stream({"token_ids": None})) + without_ids: Final = _hosted_vllm_stream_choices(logging_obj, _held_back_stop_string_stream({})) + + assert [(c["delta"], c["finish_reason"]) for c in with_null_ids] == [ + (c["delta"], c["finish_reason"]) for c in without_ids + ] + + +_TWO_CHOICE_STOP_REASONS: Final[Mapping[int, str | int]] = MappingProxyType({0: ", 3", 1: 151645}) +_TWO_CHOICE_TOKEN_IDS: Final[Mapping[int, tuple[int, int]]] = MappingProxyType({0: (16, 17), 1: (26, 27)}) + + +def _two_choice_stream(first_finished: int, finish_contents: tuple[str, str]) -> tuple[Mapping[str, object], ...]: + return ( + *({"index": index, "delta": {"role": "assistant", "content": ""}, "finish_reason": None} for index in (0, 1)), + *( + {"index": index, "delta": {"content": "1"}, "finish_reason": None, "token_ids": [ids[0]]} + for index, ids in _TWO_CHOICE_TOKEN_IDS.items() + ), + *( + { + "index": index, + "delta": {"content": content}, + "finish_reason": "stop", + "stop_reason": _TWO_CHOICE_STOP_REASONS[index], + "token_ids": [_TWO_CHOICE_TOKEN_IDS[index][1]], + } + for index, content in zip((first_finished, 1 - first_finished), finish_contents) + ), + ) + + +def _carries_another_choices_fields(choice: Mapping[str, object]) -> bool: + own_index: Final = _CHOICE_INDEX.validate_python(choice["index"]) + token_ids: Final = _TOKEN_IDS.validate_python(choice.get("token_ids")) or () + return not set(token_ids) <= set(_TWO_CHOICE_TOKEN_IDS[own_index]) or choice.get("stop_reason") not in ( + None, + _TWO_CHOICE_STOP_REASONS[own_index], + ) + + +@pytest.mark.parametrize("first_finished", [0, 1]) +@pytest.mark.parametrize( + "finish_contents", + [("", ""), ("2", ""), ("2", "2")], + ids=["both-finish-only", "content-then-finish-only", "both-content-then-finish"], +) +def test_hosted_vllm_stream_with_two_choices_keeps_ids_and_stop_reason_on_their_own_choice( + logging_obj: Logging, first_finished: int, finish_contents: tuple[str, str] +) -> None: + emitted: Final = _hosted_vllm_stream_choices(logging_obj, _two_choice_stream(first_finished, finish_contents)) + + assert [ + (choice["index"], choice.get("token_ids"), choice.get("stop_reason")) + for choice in emitted + if _carries_another_choices_fields(choice) + ] == [] + + +@pytest.mark.parametrize("first_finished", [0, 1]) +@pytest.mark.parametrize( + "finish_contents", [("2", ""), ("2", "2")], ids=["content-then-finish-only", "both-content-then-finish"] +) +def test_hosted_vllm_stream_with_two_choices_keeps_each_stop_reason_on_its_final_content_chunk( + logging_obj: Logging, first_finished: int, finish_contents: tuple[str, str] +) -> None: + emitted: Final = _hosted_vllm_stream_choices(logging_obj, _two_choice_stream(first_finished, finish_contents)) + content_chunk_stop_reasons: Final = frozenset( + (choice["index"], choice["stop_reason"]) + for choice in emitted + if _DELTA.validate_python(choice["delta"]).get("content") and choice.get("stop_reason") is not None + ) + finished_with_content: Final = frozenset( + (index, _TWO_CHOICE_STOP_REASONS[index]) + for index, content in zip((first_finished, 1 - first_finished), finish_contents) + if content + ) + + assert content_chunk_stop_reasons == finished_with_content + + +@pytest.mark.parametrize("first_finished", [0, 1]) +@pytest.mark.parametrize( + "finish_contents", + [("", ""), ("2", ""), ("", "2"), ("2", "2")], + ids=["both-finish-only", "content-then-finish-only", "finish-only-then-content", "both-content-then-finish"], +) +def test_hosted_vllm_stream_with_two_choices_delivers_every_token_id_and_stop_reason_of_each_choice( + logging_obj: Logging, first_finished: int, finish_contents: tuple[str, str] +) -> None: + emitted: Final = _hosted_vllm_stream_choices(logging_obj, _two_choice_stream(first_finished, finish_contents)) + + assert { + index: _joined_token_ids([choice for choice in emitted if choice["index"] == index]) + for index in _TWO_CHOICE_TOKEN_IDS + } == _TWO_CHOICE_TOKEN_IDS + assert frozenset( + (choice["index"], choice["stop_reason"]) for choice in emitted if choice.get("stop_reason") is not None + ) == frozenset(_TWO_CHOICE_STOP_REASONS.items()) + + +def test_hosted_vllm_stream_with_two_choices_finishes_the_first_choice_before_the_other_choice_streams_on( + logging_obj: Logging, +) -> None: + emitted: Final = _hosted_vllm_stream_choices( + logging_obj, + ( + *({"index": index, "delta": {"role": "assistant", "content": ""}, "finish_reason": None} for index in (0, 1)), + {"index": 0, "delta": {"content": "1"}, "finish_reason": "stop", "stop_reason": ",2", "token_ids": [16]}, + {"index": 1, "delta": {"content": ""}, "finish_reason": None, "token_ids": [26]}, + {"index": 1, "delta": {"content": ""}, "finish_reason": None, "token_ids": [11]}, + {"index": 1, "delta": {"content": "1,"}, "finish_reason": None, "token_ids": [27]}, + {"index": 1, "delta": {"content": "3"}, "finish_reason": "stop", "stop_reason": 151645, "token_ids": [28]}, + ), + ) + first_choice_finish: Final = next( + position for position, choice in enumerate(emitted) if choice["index"] == 0 and choice["finish_reason"] + ) + + assert emitted[first_choice_finish]["stop_reason"] == ",2" + assert _joined_token_ids([choice for choice in emitted[first_choice_finish:] if choice["index"] == 1]) == ( + 26, + 11, + 27, + 28, + ) + + +_CACHED_VLLM_TOKEN_IDS: Final = (16, 11, 220, 17, 11, 220, 18, 11, 220, 19) + + +def _cached_vllm_completion() -> dict[str, object]: + return { + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 1, + "model": "Qwen/Qwen3-0.6B", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "1, 2, 3, 4"}, + "finish_reason": "stop", + "stop_reason": ", 5", + "token_ids": list(_CACHED_VLLM_TOKEN_IDS), + } + ], + } + + +async def _replayed_cache_hit_choices(logging_obj: Logging, replay: str) -> tuple[dict[str, object], ...]: + if replay == "async": + async_wrapper: Final = CustomStreamWrapper( + completion_stream=convert_to_streaming_response_async(_cached_vllm_completion()), + model="m", + custom_llm_provider="cached_response", + logging_obj=logging_obj, + ) + return tuple([chunk.choices[0].model_dump() async for chunk in async_wrapper if chunk.choices]) + sync_wrapper: Final = CustomStreamWrapper( + completion_stream=convert_to_streaming_response(_cached_vllm_completion()), + model="m", + custom_llm_provider="cached_response", + logging_obj=logging_obj, + ) + return tuple(chunk.choices[0].model_dump() for chunk in sync_wrapper if chunk.choices) + + +@pytest.mark.parametrize("replay", ["sync", "async"]) +@pytest.mark.asyncio +async def test_a_streamed_cache_hit_keeps_the_vllm_stop_reason_and_every_token_id( + logging_obj: Logging, replay: str +) -> None: + emitted: Final = await _replayed_cache_hit_choices(logging_obj, replay) + + assert [ + (choice["finish_reason"], choice["stop_reason"]) for choice in emitted if choice.get("stop_reason") is not None + ] == [("stop", ", 5")] + assert _joined_token_ids(emitted) == _CACHED_VLLM_TOKEN_IDS + + def _later_chunk_outputs( custom_llm_provider: str, later_chunks: Sequence[ChatCompletionChunk], logging_obj: Logging ) -> tuple[ModelResponseStream, ...]: diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index 697ac76b4e5..8ef118154d1 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -638,8 +638,14 @@ def test_token_counter_tools_match_billed_prompt_for_json_tool_chat_templates( ("deepseek/deepseek-chat:free", "DeepSeek-V4-Flash-0731"), ("deepseek-reasoner", "DeepSeek-V4-Flash-0731"), ("deepseek/deepseek-flash", "DeepSeek-V4-Flash-0731"), - ("DeepSeek-V4.1-Flash", "deepseek-v3.2"), - ("accounts/fireworks/models/deepseek-v4p1-flash", "deepseek-v3.2"), + ("DeepSeek-V4.1-Flash", "accounts/fireworks/models/deepseek-v4p1-flash"), + ("databricks-deepseek-v4-1-flash", "accounts/fireworks/models/deepseek-v4p1-flash"), + ("deepseek-ai/deepseek-v4.1-flash-maas", "accounts/fireworks/models/deepseek-v4p1-flash"), + ("deepseek-v4p1-flash-1017", "accounts/fireworks/models/deepseek-v4p1-flash"), + ("grok-4.1-fast-non-reasoning", "grok-4-1-fast-reasoning"), + ("xai.grok-4.1-fast", "grok-4-1-fast-reasoning"), + ("grok-4.3-latest", "grok-4.3"), + ("xai.grok-4-3", "grok-4.3"), ("Qwen/Qwen3.5-122B-A10B", "databricks-qwen35-122b-a10b"), ("qwen3-235b-a22b", "mistral-large-latest"), ], @@ -704,6 +710,48 @@ def test_token_counter_tool_choice_adds_nothing_for_deepseek_v4(tool_choice: obj ) == token_counter_new(model="DeepSeek-V4-Flash-0731", messages=list(_USER_MESSAGES), tools=tools) +# A request with a one-function tool billed 283 prompt tokens on xAI grok-4.3 without tool_choice and 297 with a +# named one, and 357 and 358 on Azure AI Foundry grok-4-1-fast-reasoning, on 2026-10-29 +@pytest.mark.parametrize("model", ["grok-4.3", "grok-4-1-fast-reasoning"]) +def test_token_counter_counts_a_named_tool_choice_for_grok(model: str): + tools: Final = [_openai_tool(_WEATHER_FUNCTION)] + named: Final = {"type": "function", "function": {"name": "get_weather"}} + + assert token_counter_new( + model=model, messages=list(_USER_MESSAGES), tools=tools, tool_choice=named + ) > token_counter_new(model=model, messages=list(_USER_MESSAGES), tools=tools) + + +def _counted_with_strict(model: str, **strict: bool) -> int: + return token_counter_new( + model=model, messages=list(_USER_MESSAGES), tools=[_openai_tool({**_WEATHER_FUNCTION, **strict})] + ) + + +# A one-function tool billed the same prompt tokens with strict unset and false on 2026-10-29: 357 through Azure AI +# Foundry grok-4-1-fast-reasoning chat and 283 through xAI grok-4.3 +@pytest.mark.parametrize("model", ["grok-4-1-fast-reasoning", "grok-4.3"]) +def test_token_counter_renders_a_false_strict_like_an_unset_one_for_grok(model: str): + assert ( + _counted_with_strict(model, strict=False) + == _counted_with_strict(model) + < _counted_with_strict(model, strict=True) + ) + + +# A one-function tool billed 5 prompt tokens more with strict false than unset, the same as with strict true, on +# 2026-10-30: 310 against 305 through Baseten DeepSeek-V4.1-Flash and 309 against 304 through Together +@pytest.mark.parametrize( + "model", ["accounts/fireworks/models/deepseek-v4p1-flash", "deepseek-v4-flash", "qwen3.5-plus"] +) +def test_token_counter_renders_a_false_strict_like_a_true_one_for_other_chat_templates(model: str): + assert ( + _counted_with_strict(model, strict=False) + == _counted_with_strict(model, strict=True) + > _counted_with_strict(model) + ) + + @pytest.mark.parametrize( ("tool_choice", "expected_tool_choice"), [ @@ -759,9 +807,17 @@ def test_token_counter_counts_anthropic_tools_as_sent_to_chat_template_families( # Qwen3.5's chat_template.jinja opens a system turn for the tool schemas when the request has none, 5 tokens that a # request with a system message already pays, while DeepSeek V4's encoding_dsv4.py appends them to the first turn -# either way (both read on Hugging Face on 2026-10-17) +# either way (both read on Hugging Face on 2026-10-17). A system message cost the same with and without a tool on +# Fireworks deepseek-v4p1-flash, Azure AI Foundry grok-4-1-fast-reasoning and xAI grok-4.3 on 2026-10-29 @pytest.mark.parametrize( - ("model", "implicit_system_turn_tokens"), [("DeepSeek-V4-Flash-0731", 0), ("databricks-qwen35-122b-a10b", 5)] + ("model", "implicit_system_turn_tokens"), + [ + ("DeepSeek-V4-Flash-0731", 0), + ("databricks-qwen35-122b-a10b", 5), + ("accounts/fireworks/models/deepseek-v4p1-flash", 0), + ("grok-4-1-fast-reasoning", 0), + ("grok-4.3", 0), + ], ) def test_token_counter_counts_the_tools_system_turn_only_without_a_system_message( model: str, implicit_system_turn_tokens: int diff --git a/tests/unit/litellm_core_utils/test_tool_fields.py b/tests/unit/litellm_core_utils/test_tool_fields.py index 23d7de98fbc..a99278d18dd 100644 --- a/tests/unit/litellm_core_utils/test_tool_fields.py +++ b/tests/unit/litellm_core_utils/test_tool_fields.py @@ -1,4 +1,5 @@ import datetime +import logging from collections.abc import Callable, Mapping from typing import Final @@ -6,7 +7,13 @@ import pytest import litellm from litellm import Router +from litellm._logging import verbose_logger from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.litellm_core_utils.tool_fields import ( + tool_params_without_hosted_mcp, + tools_without_hosted_mcp, + tools_without_mcp_servers, +) from litellm.utils import async_pre_call_deployment_hook _GOOD_TOOL: Final = {"type": "function", "function": {"name": "get_weather", "parameters": {"type": "object"}}} @@ -185,3 +192,92 @@ async def test_tools_core_accepts_reach_the_deployment(tools: object) -> None: sent: Final = await async_pre_call_deployment_hook(dict(_proxy_request("keeps", tools=tools)), "acompletion") assert sent["tools"] is tools + + +def test_a_hosted_mcp_error_names_every_server_and_marks_an_unlabeled_one(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(litellm.UnsupportedParamsError) as raised: + tools_without_hosted_mcp( + tools=[ + {"type": "mcp", "server_label": "deepwiki", "server_url": "https://mcp.deepwiki.com/mcp"}, + {"type": "mcp", "server_url": "https://mcp.example.com/mcp"}, + ], + model="gemini-3.8-flash", + route="gemini", + drop_params=False, + ) + assert "the hosted MCP tool `deepwiki`, `unlabeled` " in str(raised.value) + + +_HOSTED_MCP_TOOL: Final = {"type": "mcp", "server_label": "deepwiki", "server_url": "https://mcp.deepwiki.com/mcp"} + + +def test_dropping_a_hosted_mcp_tool_keeps_the_other_tool_params_and_leaves_the_caller_dict_alone() -> None: + tool_choice: Final = {"type": "function", "function": {"name": "get_weather"}} + sent: Final = { + "tools": [_HOSTED_MCP_TOOL, _GOOD_TOOL], + "tool_choice": tool_choice, + "parallel_tool_calls": False, + "temperature": 0.2, + } + params: Final = tool_params_without_hosted_mcp(sent, model="deepseek-v4-flash", route="deepseek", drop_params=True) + + assert params == { + "tools": [_GOOD_TOOL], + "tool_choice": tool_choice, + "parallel_tool_calls": False, + "temperature": 0.2, + } + assert sent["tools"] == [_HOSTED_MCP_TOOL, _GOOD_TOOL] + + +def test_dropping_every_hosted_mcp_tool_clears_the_tool_selection_params() -> None: + params: Final = tool_params_without_hosted_mcp( + {"tools": [_HOSTED_MCP_TOOL], "tool_choice": "required", "parallel_tool_calls": True, "temperature": 0.2}, + model="deepseek-v4-flash", + route="deepseek", + drop_params=True, + ) + + assert params == {"temperature": 0.2} + + +class _Records(logging.Handler): + def __init__(self) -> None: + super().__init__(level=logging.WARNING) + self.messages: Final[list[str]] = [] + + def emit(self, record: logging.LogRecord) -> None: + self.messages.append(record.getMessage()) + + +@pytest.mark.parametrize( + "drop", + [ + lambda model: tools_without_hosted_mcp( + [{"type": "mcp", "server_label": "wiki\r\nFORGED", "server_url": "https://mcp.example.com/mcp"}], + model=model, + route="deepseek", + drop_params=True, + ), + lambda model: tools_without_mcp_servers( + [{"type": "url", "name": "wiki\r\nFORGED", "url": "https://mcp.example.com/mcp"}], + tools=[], + model=model, + route="bedrock", + drop_params=True, + ), + ], + ids=["hosted_mcp_tool", "mcp_servers"], +) +def test_a_drop_warning_stays_on_one_log_line(drop: Callable[[str], object]) -> None: + records: Final = _Records() + verbose_logger.addHandler(records) + try: + drop("deepseek-v4-flash\nFORGED level=CRITICAL") + finally: + verbose_logger.removeHandler(records) + + assert len(records.messages) == 1 + assert "\n" not in records.messages[0] and "\r" not in records.messages[0] + assert "FORGED level=CRITICAL" in records.messages[0] diff --git a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py index 63a304ab77d..395255f0608 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -10014,6 +10014,85 @@ def test_transform_request_tool_config_point_without_tools_drops_the_key(): assert litellm_params["metadata"] == {} +_HOSTED_MCP_TOOL: Final = { + "type": "mcp", + "server_label": "deepwiki", + "server_url": "https://mcp.deepwiki.com/mcp", + "require_approval": "never", +} +_WEATHER_FUNCTION_TOOL: Final = { + "type": "function", + "function": {"name": "get_weather", "parameters": {"type": "object", "properties": {}}}, +} +_ROUTES_WITHOUT_THE_MCP_CONNECTOR: Final = ( + pytest.param(AmazonAnthropicClaudeConfig, "bedrock", id="bedrock_invoke"), + pytest.param(VertexAIAnthropicConfig, "vertex_ai", id="vertex_ai"), +) + + +@pytest.mark.parametrize("config_class", [AnthropicConfig, AzureAnthropicConfig, BedrockClaudePlatformConfig]) +def test_routes_with_the_mcp_connector_send_a_hosted_mcp_tool_as_an_mcp_server(monkeypatch, config_class): + monkeypatch.setattr(litellm, "drop_params", False) + mapped: Final = config_class().map_openai_params( + non_default_params={"tools": [_WEATHER_FUNCTION_TOOL, _HOSTED_MCP_TOOL]}, + optional_params={}, + model="claude-haiku-5-5", + drop_params=False, + ) + assert [server["url"] for server in mapped["mcp_servers"]] == [_HOSTED_MCP_TOOL["server_url"]] + assert [tool["name"] for tool in mapped["tools"]] == ["get_weather"] + + +@pytest.mark.parametrize(("config_class", "route"), _ROUTES_WITHOUT_THE_MCP_CONNECTOR) +def test_a_hosted_mcp_tool_answers_400_on_claude_routes_without_the_mcp_connector(monkeypatch, config_class, route): + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(litellm.utils.UnsupportedParamsError) as raised: + config_class().map_openai_params( + non_default_params={"tools": [_WEATHER_FUNCTION_TOOL, _HOSTED_MCP_TOOL]}, + optional_params={}, + model="claude-haiku-5-5", + drop_params=False, + ) + assert raised.value.status_code == 400 + assert f"claude-haiku-5-5 on {route} " in str(raised.value) + assert "`deepwiki`" in str(raised.value) + + +@pytest.mark.parametrize(("config_class", "route"), _ROUTES_WITHOUT_THE_MCP_CONNECTOR) +def test_claude_routes_without_the_mcp_connector_drop_a_hosted_mcp_tool_under_drop_params( + monkeypatch, caplog, config_class, route +): + monkeypatch.setattr(litellm, "drop_params", False) + with caplog.at_level("WARNING"): + mapped: Final = config_class().map_openai_params( + non_default_params={"tools": [_HOSTED_MCP_TOOL, _WEATHER_FUNCTION_TOOL]}, + optional_params={}, + model="claude-haiku-5-5", + drop_params=True, + ) + assert [tool["name"] for tool in mapped["tools"]] == ["get_weather"] + assert "mcp_servers" not in mapped + assert any( + "`deepwiki`" in record.getMessage() and f" on {route} (" in record.getMessage() for record in caplog.records + ) + + +@pytest.mark.parametrize(("config_class", "route"), _ROUTES_WITHOUT_THE_MCP_CONNECTOR) +def test_claude_routes_without_the_mcp_connector_send_no_tools_when_drop_params_drops_every_one( + monkeypatch, config_class, route +): + monkeypatch.setattr(litellm, "drop_params", False) + mapped: Final = config_class().map_openai_params( + non_default_params={"tools": [_HOSTED_MCP_TOOL], "tool_choice": "required"}, + optional_params={}, + model="claude-haiku-5-5", + drop_params=True, + ) + assert "tools" not in mapped + assert "mcp_servers" not in mapped + assert "tool_choice" not in mapped + + _KEPT_SEARCH_CALL: Final = { "role": "assistant", "content": None, diff --git a/tests/unit/llms/anthropic/messages/adapters/test_handler_output_config_passthrough.py b/tests/unit/llms/anthropic/messages/adapters/test_handler_output_config_passthrough.py index c34cb686928..12769e78329 100644 --- a/tests/unit/llms/anthropic/messages/adapters/test_handler_output_config_passthrough.py +++ b/tests/unit/llms/anthropic/messages/adapters/test_handler_output_config_passthrough.py @@ -27,11 +27,14 @@ Tests cover (consolidating PRs #23706 and #22727): import os import sys +from collections.abc import Mapping from typing import Final from unittest.mock import MagicMock, patch import pytest +import litellm + # Anchor sys.path to this file's location — not the working-directory-relative # pattern Greptile flagged on PR #23706. Resolves correctly regardless of # where pytest is invoked from. @@ -67,8 +70,8 @@ def _call_prepare(extra_kwargs, model="gpt-4o", output_format=None, **overrides) system=None, temperature=None, thinking=None, - tool_choice=None, - tools=None, + tool_choice=overrides.get("tool_choice"), + tools=overrides.get("tools"), top_k=None, top_p=None, output_format=output_format, @@ -263,3 +266,112 @@ async def test_messages_layer_converted_stream_flag_still_gets_a_message_back(co ) assert response["content"][0]["text"] == "hi" + + +_DEEPWIKI_SERVER: Final = {"type": "url", "url": "https://mcp.deepwiki.com/mcp", "name": "deepwiki"} + + +@pytest.mark.parametrize( + ("mcp_servers", "named"), + [ + ([_DEEPWIKI_SERVER], "`deepwiki`"), + ([_DEEPWIKI_SERVER, {"type": "url", "url": "https://mcp.example.com/mcp"}], "`deepwiki`, `unnamed`"), + ({"name": "deepwiki"}, "`unnamed`"), + ], + ids=["named", "one-unnamed", "not-a-list"], +) +def test_mcp_servers_a_chat_route_cannot_run_answer_400_naming_each_server( + monkeypatch: pytest.MonkeyPatch, mcp_servers: object, named: str +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(litellm.UnsupportedParamsError) as raised: + _call_prepare( + extra_kwargs={"mcp_servers": mcp_servers, "custom_llm_provider": "gemini"}, model="gemini-3.8-flash" + ) + + assert raised.value.status_code == 400 + assert f"to gemini-3.8-flash on gemini, so the MCP server {named} would not run" in str(raised.value) + + +@pytest.mark.parametrize( + ("route", "forwards"), + [("litellm_proxy", True), ("hosted_vllm", True), ("bedrock", False), ("vertex_ai", False)], +) +def test_mcp_servers_reach_routes_that_carry_them_and_answer_400_elsewhere( + monkeypatch: pytest.MonkeyPatch, route: str, forwards: bool +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + extra_kwargs: Final = {"mcp_servers": [_DEEPWIKI_SERVER], "custom_llm_provider": route} + if not forwards: + with pytest.raises(litellm.UnsupportedParamsError, match=f"on {route}, so the MCP server `deepwiki`"): + _call_prepare(extra_kwargs=extra_kwargs, model="some-model") + return + + completion_kwargs, _ = _call_prepare(extra_kwargs=extra_kwargs, model="some-model") + + assert completion_kwargs["mcp_servers"] == [_DEEPWIKI_SERVER] + + +_DEEPWIKI_TOOLSET: Final = {"type": "mcp_toolset", "mcp_server_name": "deepwiki"} +_WEATHER_TOOL: Final = { + "name": "get_weather", + "description": "Get the weather for a city", + "input_schema": {"type": "object", "properties": {"city": {"type": "string"}}}, +} + + +@pytest.mark.parametrize( + ("tools", "kept_tool_names"), + [([_DEEPWIKI_TOOLSET, _WEATHER_TOOL], ["get_weather"]), ([_DEEPWIKI_TOOLSET], None)], + ids=["with-function", "toolset-only"], +) +def test_drop_params_drops_the_mcp_toolset_tools_of_dropped_mcp_servers( + monkeypatch: pytest.MonkeyPatch, tools: list[dict[str, object]], kept_tool_names: list[str] | None +) -> None: + monkeypatch.setattr(litellm, "drop_params", True) + completion_kwargs, _ = _call_prepare( + extra_kwargs={"mcp_servers": [_DEEPWIKI_SERVER], "custom_llm_provider": "gemini"}, + model="gemini-3.8-flash", + tools=tools, + tool_choice={"type": "auto"}, + ) + + sent_tools: Final = completion_kwargs.get("tools") + sent_names: Final = [tool["function"]["name"] for tool in sent_tools] if sent_tools else None + assert sent_names == kept_tool_names + assert ("tool_choice" in completion_kwargs) == bool(kept_tool_names) + + +_DROP_FORMS: Final[Mapping[str, Mapping[str, object]]] = { + "request": {"drop_params": True}, + "global": {}, + "additional_drop_params": {"additional_drop_params": ["mcp_servers"]}, +} + + +@pytest.mark.parametrize("drop_params", list(_DROP_FORMS)) +def test_drop_params_drops_mcp_servers_instead_of_forwarding_them( + monkeypatch: pytest.MonkeyPatch, drop_params: str +) -> None: + monkeypatch.setattr(litellm, "drop_params", drop_params == "global") + completion_kwargs, _ = _call_prepare( + extra_kwargs={ + "mcp_servers": [_DEEPWIKI_SERVER], + "custom_llm_provider": "gemini", + **_DROP_FORMS[drop_params], + }, + model="gemini-3.8-flash", + ) + + assert "mcp_servers" not in completion_kwargs + assert completion_kwargs["messages"] == MESSAGES + + +@pytest.mark.parametrize("mcp_servers", [None, []], ids=["absent", "empty"]) +def test_no_mcp_servers_needs_no_drop_params(monkeypatch: pytest.MonkeyPatch, mcp_servers: object) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + completion_kwargs, _ = _call_prepare( + extra_kwargs={"mcp_servers": mcp_servers, "custom_llm_provider": "gemini"}, model="gemini-3.8-flash" + ) + + assert "mcp_servers" not in completion_kwargs diff --git a/tests/unit/llms/anthropic/messages/adapters/test_streaming_iterator_combined_chunk.py b/tests/unit/llms/anthropic/messages/adapters/test_streaming_iterator_combined_chunk.py index d4911427e73..a57011b6410 100644 --- a/tests/unit/llms/anthropic/messages/adapters/test_streaming_iterator_combined_chunk.py +++ b/tests/unit/llms/anthropic/messages/adapters/test_streaming_iterator_combined_chunk.py @@ -204,6 +204,20 @@ def test_splitter_splits_combined_chunk_into_content_then_finish(): assert finish_chunk.choices[0].delta.content is None +def test_a_finish_on_another_choice_index_does_not_end_the_message(): + chunks: Final = [ + ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content="Hello"))]), + ModelResponseStream(choices=[StreamingChoices(index=1, delta=Delta(content="Hi"))]), + ModelResponseStream(choices=[StreamingChoices(index=1, delta=Delta(content=""), finish_reason="length")]), + ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content=""), finish_reason="stop")]), + ] + + events: Final = list(AnthropicStreamWrapper(completion_stream=iter(chunks), model="qwen3")) + + stop_reasons: Final = [event["delta"]["stop_reason"] for event in events if event.get("type") == "message_delta"] + assert stop_reasons == ["end_turn"] + + def test_is_combined_false_when_choices_empty(): """A metadata-only chunk with no choices is never treated as combined.""" assert _CombinedChunkSplitter._is_combined(SimpleNamespace(choices=[])) is False diff --git a/tests/unit/llms/anthropic/messages/responses_adapters/test_responses_adapters_handler.py b/tests/unit/llms/anthropic/messages/responses_adapters/test_responses_adapters_handler.py index 382917fc756..c06c9d721ac 100644 --- a/tests/unit/llms/anthropic/messages/responses_adapters/test_responses_adapters_handler.py +++ b/tests/unit/llms/anthropic/messages/responses_adapters/test_responses_adapters_handler.py @@ -764,6 +764,57 @@ async def test_recomputed_cost_of_a_served_bridge_response_prices_the_reasoning_ assert logging_obj.cost_breakdown["reasoning_cost"] == pytest.approx(expected_reasoning) +_DEEPWIKI_SERVER: Final = {"type": "url", "url": "https://mcp.deepwiki.com/mcp", "name": "deepwiki"} + + +def test_mcp_servers_the_responses_bridge_cannot_run_answer_400(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(litellm.UnsupportedParamsError) as rejected: + _build_responses_kwargs( + model="gpt-5.6", + extra_kwargs={**OPENAI_API, "mcp_servers": [_DEEPWIKI_SERVER]}, + max_tokens=1024, + messages=MESSAGES, + ) + + assert rejected.value.status_code == 400 + assert "to gpt-5.6 on openai, so the MCP server `deepwiki` would not run" in rejected.value.message + + +@pytest.mark.parametrize( + "drop_form", + [{"drop_params": True}, {"additional_drop_params": ["mcp_servers"]}], + ids=["drop_params", "additional_drop_params"], +) +def test_drop_params_keeps_mcp_servers_off_the_responses_request( + monkeypatch: pytest.MonkeyPatch, drop_form: Mapping[str, object] +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + responses_kwargs: Final = _build_responses_kwargs( + model="gpt-5.6", + extra_kwargs={**OPENAI_API, "mcp_servers": [_DEEPWIKI_SERVER], **drop_form}, + max_tokens=1024, + messages=MESSAGES, + ) + + assert "mcp_servers" not in responses_kwargs + assert responses_kwargs["input"] + + +def test_drop_params_drops_the_mcp_toolset_tools_of_dropped_mcp_servers(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + weather_tool: Final = {"name": "get_weather", "input_schema": {"type": "object", "properties": {}}} + responses_kwargs: Final = _build_responses_kwargs( + model="gpt-5.6", + extra_kwargs={**OPENAI_API, "mcp_servers": [_DEEPWIKI_SERVER], "drop_params": True}, + max_tokens=1024, + messages=MESSAGES, + tools=[{"type": "mcp_toolset", "mcp_server_name": "deepwiki"}, weather_tool], + ) + + assert [tool["name"] for tool in responses_kwargs["tools"]] == ["get_weather"] + + TEXT_RESPONSE: Final = { "id": "resp_cleared", "object": "response", diff --git a/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude_transformation.py index d40abb4860b..4850a402103 100644 --- a/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude_transformation.py +++ b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude_transformation.py @@ -14,6 +14,7 @@ from pydantic import BaseModel, ConfigDict # Ensure the project root is on the import path so `litellm` can be imported when # tests are executed from any working directory. import litellm +from litellm.constants import DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude_transformation import ( AmazonAnthropicClaudeConfig, ) @@ -2010,6 +2011,123 @@ def test_invoke_uses_opt_in_native_output_is_false_when_native_output_did_not_co ) +BEDROCK_HAIKU_45_BASE_ID: Final = "anthropic.claude-haiku-4-5-20251001-v1:0" + + +@pytest.mark.parametrize( + "model", + [ + f"invoke/us.{BEDROCK_HAIKU_45_BASE_ID}", + f"invoke/{BEDROCK_HAIKU_45_BASE_ID}", + ], +) +def test_bedrock_invoke_claude_default_max_tokens_comes_from_base_model(local_model_cost_map, model: str): + mapped_cap: Final = litellm.get_max_tokens(BEDROCK_HAIKU_45_BASE_ID) + assert mapped_cap != DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS, "the model's mapped cap must differ from the fallback" + + body: Final = AmazonAnthropicClaudeConfig().transform_request( + model=model, + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert body["max_tokens"] == mapped_cap, f"default max_tokens should be the base model's cap, got {body['max_tokens']}" + + +def test_bedrock_invoke_claude_explicit_max_tokens_is_kept(local_model_cost_map): + body: Final = AmazonAnthropicClaudeConfig().transform_request( + model=f"invoke/us.{BEDROCK_HAIKU_45_BASE_ID}", + messages=[{"role": "user", "content": "hi"}], + optional_params={"max_tokens": 123}, + litellm_params={}, + headers={}, + ) + + assert body["max_tokens"] == 123, f"explicit max_tokens must not be overridden, got {body['max_tokens']}" + + +def test_bedrock_invoke_claude_default_max_tokens_takes_the_deployments_own_cap(local_model_cost_map): + deployment_cap: Final = 8192 + assert litellm.get_max_tokens(BEDROCK_HAIKU_45_BASE_ID) != deployment_cap + + body: Final = AmazonAnthropicClaudeConfig().transform_request( + model=f"invoke/us.{BEDROCK_HAIKU_45_BASE_ID}", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={"model_info": {"max_output_tokens": deployment_cap}}, + headers={}, + ) + + assert body["max_tokens"] == deployment_cap, f"the deployment's own cap should win, got {body['max_tokens']}" + + +def test_bedrock_invoke_claude_default_max_tokens_ignores_a_sibling_deployments_cap(local_model_cost_map): + regional_id: Final = f"us.{BEDROCK_HAIKU_45_BASE_ID}" + mapped_cap: Final = litellm.get_max_tokens(BEDROCK_HAIKU_45_BASE_ID) + litellm.register_model(model_cost={regional_id: {"max_output_tokens": 50}}, persist_across_reloads=False) + + body: Final = AmazonAnthropicClaudeConfig().transform_request( + model=f"invoke/{regional_id}", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert body["max_tokens"] == mapped_cap, f"a sibling's cap on the shared id leaked, got {body['max_tokens']}" + + +def test_bedrock_invoke_claude_default_max_tokens_yields_to_anthropic_config(local_model_cost_map): + configured_max_tokens: Final = 1000 + assert litellm.get_max_tokens(BEDROCK_HAIKU_45_BASE_ID) != configured_max_tokens + + body: Final = AmazonAnthropicClaudeConfig( + read_configured_max_tokens=lambda: configured_max_tokens + ).transform_request( + model=f"invoke/us.{BEDROCK_HAIKU_45_BASE_ID}", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert body["max_tokens"] == configured_max_tokens, f"AnthropicConfig(max_tokens) was ignored, got {body['max_tokens']}" + + +@pytest.mark.parametrize(("configured_max_tokens", "deployment_cap"), [(100000, 50), (1000, 8192)]) +def test_bedrock_invoke_claude_default_max_tokens_is_the_lower_of_anthropic_config_and_the_deployments_cap( + local_model_cost_map, configured_max_tokens: int, deployment_cap: int +): + body: Final = AmazonAnthropicClaudeConfig( + read_configured_max_tokens=lambda: configured_max_tokens + ).transform_request( + model=f"invoke/us.{BEDROCK_HAIKU_45_BASE_ID}", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={"model_info": {"max_output_tokens": deployment_cap}}, + headers={}, + ) + + assert body["max_tokens"] == min(configured_max_tokens, deployment_cap), f"got {body['max_tokens']}" + + +def test_bedrock_invoke_claude_default_max_tokens_reads_an_inference_profile_arns_base_model(local_model_cost_map): + mapped_cap: Final = litellm.get_max_tokens(BEDROCK_HAIKU_45_BASE_ID) + assert mapped_cap != DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS, "the model's mapped cap must differ from the fallback" + + body: Final = AmazonAnthropicClaudeConfig().transform_request( + model="invoke/anthropic/arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc123", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={"base_model": BEDROCK_HAIKU_45_BASE_ID}, + headers={}, + ) + + assert body["max_tokens"] == mapped_cap, f"an ARN should take its base_model's cap, got {body['max_tokens']}" + + @pytest.mark.usefixtures("local_model_cost_map", "local_beta_headers_config") @pytest.mark.parametrize("drop_params", [True, False]) def test_bedrock_chat_invoke_adds_the_binding_beta_so_block_binding_is_kept( diff --git a/tests/unit/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py b/tests/unit/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py index 154be9f0370..56d6b039788 100644 --- a/tests/unit/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py +++ b/tests/unit/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py @@ -2,6 +2,7 @@ import base64 import json import uuid from types import SimpleNamespace +from typing import Final import httpx import pytest @@ -243,3 +244,21 @@ async def test_mantle_claude_async_completion_streams_a_wrapped_json_schema_answ chunks = tuple([chunk async for chunk in stream]) assert _streamed_answer(chunks) == BARE_STREAMED_ANSWER + + +def test_mantle_claude_default_max_tokens_comes_from_the_mantle_row(local_model_cost_map): + from litellm.constants import DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS + from litellm.llms.bedrock.chat.mantle.transformation import AmazonMantleConfig + + mapped_cap: Final = litellm.get_max_tokens("bedrock_mantle/anthropic.claude-haiku-4-5") + assert mapped_cap != DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS, "the model's mapped cap must differ from the fallback" + + body: Final = AmazonMantleConfig().transform_request( + model="mantle/anthropic.claude-haiku-4-5", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert body["max_tokens"] == mapped_cap, f"Mantle should send its mapped cap, got {body['max_tokens']}" diff --git a/tests/unit/llms/bedrock/chat/test_converse_transformation.py b/tests/unit/llms/bedrock/chat/test_converse_transformation.py index 9ab7060d766..44d3427be91 100644 --- a/tests/unit/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/unit/llms/bedrock/chat/test_converse_transformation.py @@ -11097,6 +11097,70 @@ def test_a_deployment_with_a_batch_control_plane_endpoint_keeps_it_out_of_chat_r assert "aws_bedrock_control_plane_endpoint" not in route.calls[0].request.content.decode() +_HOSTED_MCP_TOOL: Final = { + "type": "mcp", + "server_label": "deepwiki", + "server_url": "https://mcp.deepwiki.com/mcp", + "require_approval": "never", +} +_WEATHER_FUNCTION_TOOL: Final = { + "type": "function", + "function": {"name": "get_weather", "parameters": {"type": "object", "properties": {}}}, +} + + +def _converse_request_for_tools( + model: str, non_default_params: Mapping[str, object], drop_params: bool +) -> Mapping[str, object]: + config: Final = AmazonConverseConfig() + return config._transform_request_helper( + model=model, + system_content_blocks=[], + optional_params=config.map_openai_params( + non_default_params=dict(non_default_params), optional_params={}, model=model, drop_params=drop_params + ), + messages=[{"role": "user", "content": "hi"}], + ) + + +@pytest.mark.parametrize( + "model", ["us.openai.gpt-oss-120b-1:0", "us.anthropic.claude-haiku-5-5", "us.amazon.nova-2-lite-v1:0"] +) +def test_a_hosted_mcp_tool_answers_400_on_converse_without_drop_params(monkeypatch, model): + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(litellm.utils.UnsupportedParamsError) as raised: + _converse_request_for_tools(model, {"tools": [_WEATHER_FUNCTION_TOOL, _HOSTED_MCP_TOOL]}, drop_params=False) + assert raised.value.status_code == 400 + assert model in str(raised.value) + assert "`deepwiki`" in str(raised.value) + assert 'server_url: "litellm_proxy/mcp/"' in str(raised.value) + assert "drop_params: true" in str(raised.value) + + +def test_converse_drops_a_hosted_mcp_tool_under_drop_params_and_keeps_the_function_tool(monkeypatch, caplog): + monkeypatch.setattr(litellm, "drop_params", False) + model: Final = "us.openai.gpt-oss-120b-1:0" + with caplog.at_level("WARNING"): + request: Final = _converse_request_for_tools( + model, {"tools": [_HOSTED_MCP_TOOL, _WEATHER_FUNCTION_TOOL]}, drop_params=True + ) + assert [tool["toolSpec"]["name"] for tool in request["toolConfig"]["tools"]] == ["get_weather"] + assert any( + "`deepwiki`" in record.getMessage() and f"model={model} on bedrock (" in record.getMessage() + for record in caplog.records + ) + + +def test_converse_sends_no_tools_when_drop_params_drops_every_hosted_mcp_tool(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + request: Final = _converse_request_for_tools( + "us.openai.gpt-oss-120b-1:0", {"tools": [_HOSTED_MCP_TOOL], "tool_choice": "auto"}, drop_params=True + ) + assert "toolConfig" not in request + assert "toolChoice" not in json.dumps(request) + assert "tool_choice" not in json.dumps(request) + + def _warnings_naming(caplog: pytest.LogCaptureFixture, param: str) -> tuple[str, ...]: return tuple( r.getMessage() for r in caplog.records if r.levelno == logging.WARNING and f"Dropping {param}=" in r.getMessage() diff --git a/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_transformation.py b/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_transformation.py index f252d9886d3..6005acffc0a 100644 --- a/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_transformation.py +++ b/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_transformation.py @@ -1096,6 +1096,25 @@ class TestBedrockMantleClaudeChatRoute: ) assert response._hidden_params["response_cost"] != pytest.approx(_row_cost("anthropic.claude-opus-5-5", 10, 5)) + def test_claude_completion_without_max_tokens_sends_the_mantle_rows_output_limit(self, monkeypatch, local_cost_map): + from litellm.constants import DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + mapped_cap: Final = litellm.model_cost["bedrock_mantle/anthropic.claude-haiku-4-5"]["max_output_tokens"] + assert mapped_cap != DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS, "the model's mapped cap must differ from the fallback" + monkeypatch.setenv("BEDROCK_MANTLE_API_KEY", "mantle-key") + monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) + handler: Final = Mock(side_effect=_anthropic_message) + + litellm.completion( + model="bedrock_mantle/anthropic.claude-haiku-4-5", + messages=[{"role": "user", "content": "hello"}], + aws_region_name="us-east-1", + client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(handler))), + ) + + assert json.loads(handler.call_args.args[0].content)["max_tokens"] == mapped_cap + def test_claude_streaming_completion_uses_native_messages_endpoint(self, monkeypatch, local_cost_map): from litellm.llms.custom_httpx.http_handler import HTTPHandler diff --git a/tests/unit/llms/cohere/chat/test_cohere_transformation.py b/tests/unit/llms/cohere/chat/test_cohere_transformation.py index d32b3bcbeca..fcb3d3f1cc0 100644 --- a/tests/unit/llms/cohere/chat/test_cohere_transformation.py +++ b/tests/unit/llms/cohere/chat/test_cohere_transformation.py @@ -576,6 +576,54 @@ def test_cohere_v2_tool_calling(respx_mock: respx.MockRouter) -> None: assert [(call.id, call.function.name, call.function.arguments) for call in tool_calls] == [ ("get_weather_k9q2xw1n", "get_weather", '{"location":"New York, NY"}') ] + assert response.choices[0].finish_reason == "tool_calls" + + +@pytest.mark.usefixtures("_cohere_httpx_transport") +@pytest.mark.respx(assert_all_called=True) +@pytest.mark.parametrize( + ("cohere_finish_reason", "expected_finish_reason"), + [("COMPLETE", "stop"), ("MAX_TOKENS", "length"), ("STOP_SEQUENCE", "stop")], +) +def test_cohere_v2_maps_finish_reason( + respx_mock: respx.MockRouter, cohere_finish_reason: str, expected_finish_reason: str +) -> None: + respx_mock.post(COHERE_V2_CHAT_URL).mock( + return_value=httpx.Response( + 200, + json=_cohere_v2_response( + {"content": [{"type": "text", "text": "Hello"}]}, + finish_reason=cohere_finish_reason, + ), + ) + ) + + response: Final = litellm.completion( + model="cohere_chat/v2/command-a-03-2025", + messages=[{"role": "user", "content": "Say hello"}], + api_key="cohere-test-key", + ) + + assert response.choices[0].finish_reason == expected_finish_reason + + +@pytest.mark.usefixtures("_cohere_httpx_transport") +@pytest.mark.respx(assert_all_called=True) +def test_cohere_v2_missing_finish_reason_defaults_to_stop(respx_mock: respx.MockRouter) -> None: + body_without_finish_reason: Final = { + key: value + for key, value in _cohere_v2_response({"content": [{"type": "text", "text": "Hello"}]}).items() + if key != "finish_reason" + } + respx_mock.post(COHERE_V2_CHAT_URL).mock(return_value=httpx.Response(200, json=body_without_finish_reason)) + + response: Final = litellm.completion( + model="cohere_chat/v2/command-a-03-2025", + messages=[{"role": "user", "content": "Say hello"}], + api_key="cohere-test-key", + ) + + assert (response.choices[0].message.content, response.choices[0].finish_reason) == ("Hello", "stop") @pytest.mark.usefixtures("_cohere_httpx_transport") diff --git a/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py index 7a6f9c7dc18..d2b78235c29 100644 --- a/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py +++ b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py @@ -1,14 +1,14 @@ import copy import json from collections.abc import AsyncIterator, Iterator -from typing import Final +from typing import Final, cast from unittest.mock import MagicMock, patch import httpx import pytest import respx from openai.lib.streaming.chat import ChatCompletionStreamState -from openai.types.chat import ChatCompletionChunk +from openai.types.chat import ChatCompletionChunk, ChatCompletionMessage from pydantic import TypeAdapter from fastapi.testclient import TestClient from pydantic import BaseModel, ConfigDict @@ -21,6 +21,7 @@ from litellm.constants import ( DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET, ) from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.llms.databricks.common_utils import DatabricksException from litellm.llms.databricks.chat.transformation import ( DatabricksChatResponseIterator, DatabricksConfig, @@ -539,13 +540,6 @@ def test_transform_request_preserves_unity_model_service_name(): [{"type": "text", "text": "Reply in JSON", "cache_control": {"type": "ephemeral"}}], id="cache-control", ), - pytest.param( - "databricks-meta-llama-3-3-70b-instruct", - [{"role": "user", "content": [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}]}], - {"response_format": {"type": "json_object"}}, - [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}], - id="multiple-text-blocks", - ), pytest.param( "databricks-meta-llama-3-3-70b-instruct", [ @@ -603,6 +597,283 @@ def test_transform_request_collapses_single_plain_text_block_for_non_claude_mode assert outgoing_messages == [{**messages[0], "content": expected_content}] +_TWO_TEXT_BLOCKS: Final = [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}] +_JSON_OBJECT: Final = {"response_format": {"type": "json_object"}} +_IMAGE_BLOCK: Final = {"type": "image_url", "image_url": {"url": "https://example.com/image.png"}} +_CACHED_BLOCK: Final = {"type": "text", "text": " in JSON", "cache_control": {"type": "ephemeral"}} + + +def _outgoing_messages( + model: str, messages: list[AllMessageValues], optional_params: dict[str, object], drop_params: bool = False +) -> list[dict[str, object]]: + request: Final = TypeAdapter(dict[str, object]).validate_python( + DatabricksConfig().transform_request( # pyright: ignore[reportUnknownMemberType] # inherited hook is untyped + model=model, + messages=messages, + optional_params=optional_params, + litellm_params={"drop_params": True} if drop_params else {}, + headers={}, + ) + ) + return TypeAdapter(list[dict[str, object]]).validate_python(request["messages"]) + + +@pytest.mark.parametrize( + "model,messages,optional_params,expected_content", + [ + pytest.param( + "databricks-meta-llama-3-3-70b-instruct", + [{"role": "user", "content": _TWO_TEXT_BLOCKS}], + _JSON_OBJECT, + "Reply in JSON", + id="json-object-two-blocks-joined", + ), + pytest.param( + "databricks-qwen35-122b-a10b", + [{"role": "user", "content": [{"type": "text", "text": "A"}, {"type": "text", "text": "B json"}] * 2}], + _JSON_OBJECT, + "AB jsonAB json", + id="json-object-four-blocks-joined-as-sent", + ), + pytest.param( + "databricks-meta-llama-3-3-70b-instruct", + [{"role": "user", "content": _TWO_TEXT_BLOCKS}], + {}, + _TWO_TEXT_BLOCKS, + id="no-response-format-stays-list", + ), + pytest.param( + "databricks-meta-llama-3-3-70b-instruct", + [{"role": "user", "content": _TWO_TEXT_BLOCKS}], + {"response_format": {"type": "json_schema", "json_schema": {"name": "r", "schema": {"type": "object"}}}}, + _TWO_TEXT_BLOCKS, + id="json-schema-stays-list", + ), + pytest.param( + "databricks-meta-llama-3-3-70b-instruct", + [{"role": "user", "content": [{"type": "text", "text": "Reply"}, _CACHED_BLOCK]}], + _JSON_OBJECT, + [{"type": "text", "text": "Reply"}, _CACHED_BLOCK], + id="json-object-cache-control-stays-list", + ), + pytest.param( + "databricks-meta-llama-3-3-70b-instruct", + [{"role": "user", "content": "Reply in JSON", "cache_control": {"type": "ephemeral"}}], + _JSON_OBJECT, + [{"type": "text", "text": "Reply in JSON", "cache_control": {"type": "ephemeral"}}], + id="json-object-message-cache-control-stays-block", + ), + pytest.param( + "databricks-meta-llama-3-3-70b-instruct", + [{"role": "user", "content": [{"type": "text", "text": "Reply in JSON"}, _IMAGE_BLOCK]}], + _JSON_OBJECT, + [{"type": "text", "text": "Reply in JSON"}, _IMAGE_BLOCK], + id="json-object-text-and-image-stays-list", + ), + pytest.param( + "databricks-claude-sonnet-5", + [{"role": "user", "content": _TWO_TEXT_BLOCKS}], + _JSON_OBJECT, + _TWO_TEXT_BLOCKS, + id="json-object-claude-stays-list", + ), + pytest.param( + "system.ai.deepseek-v4-flash-0731", + [{"role": "user", "content": _TWO_TEXT_BLOCKS}], + _JSON_OBJECT, + "Reply in JSON", + id="json-object-ai-gateway-model-joined", + ), + pytest.param( + "openai-gpt-external", + [{"role": "user", "content": _TWO_TEXT_BLOCKS}], + _JSON_OBJECT, + _TWO_TEXT_BLOCKS, + id="json-object-external-endpoint-stays-list", + ), + ], +) +def test_transform_request_joins_plain_text_blocks_only_for_hosted_non_claude_json_object( + model: str, + messages: list[AllMessageValues], + optional_params: dict[str, object], + expected_content: object, +) -> None: + assert _outgoing_messages(model, messages, optional_params) == [{"role": "user", "content": expected_content}] + + +@pytest.mark.parametrize("response_format", [None, {"type": "json_object"}], ids=["no-response-format", "json-object"]) +def test_transform_request_rejects_a_text_block_whose_text_is_a_truthy_non_string( + response_format: dict[str, str] | None, +) -> None: + with pytest.raises(DatabricksException) as error: + _outgoing_messages( + "databricks-meta-llama-3-3-70b-instruct", + [{"role": "user", "content": [{"type": "text", "text": "Reply in JSON"}, {"type": "text", "text": 7}]}], + {} if response_format is None else {"response_format": response_format}, + ) + + assert error.value.status_code == 400 + + +def test_transform_request_accepts_an_openai_message_object_in_history() -> None: + earlier_answer: Final = ChatCompletionMessage(role="assistant", content="Earlier answer") + + outgoing: Final = _outgoing_messages( + "databricks-meta-llama-3-3-70b-instruct", + [ + {"role": "user", "content": "Earlier question"}, + cast(AllMessageValues, earlier_answer), # cast-ok: SDK callers pass response messages back as history + {"role": "user", "content": "Next question"}, + ], + {}, + ) + + assert [message.get("content") for message in outgoing] == ["Earlier question", "Earlier answer", "Next question"] + + +def test_transform_request_json_object_keeps_earlier_turns_as_sent() -> None: + earlier_turn: Final[AllMessageValues] = { + "role": "user", + "content": [{"type": "text", "text": "Code: A"}, {"type": "text", "text": "BC"}], + } + reply: Final[AllMessageValues] = {"role": "assistant", "content": "Got it."} + + assert _outgoing_messages( + "databricks-meta-llama-3-3-70b-instruct", + [earlier_turn, reply, {"role": "user", "content": _TWO_TEXT_BLOCKS}], + _JSON_OBJECT, + ) == [earlier_turn, reply, {"role": "user", "content": "Reply in JSON"}] + + +def test_transform_request_json_object_joins_a_trailing_system_run_and_keeps_its_break() -> None: + messages: Final[list[AllMessageValues]] = [ + {"role": "user", "content": "Say hi"}, + {"role": "system", "content": [{"type": "text", "text": "Alpha"}, {"type": "text", "text": " rule."}]}, + {"role": "system", "content": _TWO_TEXT_BLOCKS}, + ] + + assert _outgoing_messages("databricks-meta-llama-3-3-70b-instruct", messages, _JSON_OBJECT) == [ + {"role": "user", "content": "Say hi"}, + {"role": "system", "content": "Alpha rule.\n\nReply in JSON"}, + ] + + +_CACHED_TEXT_AND_IMAGE: Final = [{**_CACHED_BLOCK, "text": "Reply in JSON"}, _IMAGE_BLOCK] + + +@pytest.mark.parametrize( + "message,expected", + [ + pytest.param( + {"role": "user", "content": [{"type": "text", "text": "Reply"}, _CACHED_BLOCK]}, + {"role": "user", "content": "Reply in JSON"}, + id="block-cache-control-dropped-and-joined", + ), + pytest.param( + {"role": "user", "content": "Reply in JSON", "cache_control": {"type": "ephemeral"}}, + {"role": "user", "content": "Reply in JSON"}, + id="message-cache-control-dropped", + ), + pytest.param( + {"role": "user", "content": _CACHED_TEXT_AND_IMAGE}, + {"role": "user", "content": _CACHED_TEXT_AND_IMAGE}, + id="image-keeps-cache-control", + ), + ], +) +def test_transform_request_json_object_drops_last_message_cache_control_under_drop_params( + message: AllMessageValues, expected: dict[str, object] +) -> None: + assert _outgoing_messages("databricks-meta-llama-3-3-70b-instruct", [message], _JSON_OBJECT, drop_params=True) == [ + expected + ] + + +_THINKING_BLOCK: Final = {"type": "thinking", "thinking": "Plan the reply.", "signature": "sig"} + + +@pytest.mark.parametrize( + "dropped_turn", + [ + pytest.param({"role": "assistant", "content": "", "thinking_blocks": [_THINKING_BLOCK]}, id="thinking-only"), + pytest.param({"role": "assistant", "content": None, "reasoning_content": "Plan."}, id="reasoning-only"), + pytest.param({"role": "assistant", "content": [{"type": "text", "text": " "}]}, id="blank-text-part"), + ], +) +def test_transform_request_json_object_joins_the_message_before_a_dropped_assistant_turn( + dropped_turn: AllMessageValues, +) -> None: + assert _outgoing_messages( + "databricks-meta-llama-3-3-70b-instruct", + [{"role": "user", "content": _TWO_TEXT_BLOCKS}, dropped_turn], + _JSON_OBJECT, + ) == [{"role": "user", "content": "Reply in JSON"}] + + +def test_transform_request_json_object_joins_a_system_run_split_by_a_dropped_assistant_turn() -> None: + earlier_turn: Final[AllMessageValues] = { + "role": "user", + "content": [{"type": "text", "text": "Say"}, {"type": "text", "text": " hi"}], + } + messages: Final[list[AllMessageValues]] = [ + earlier_turn, + {"role": "system", "content": [{"type": "text", "text": "Alpha"}, {"type": "text", "text": " rule."}]}, + {"role": "assistant", "content": "", "thinking_blocks": [_THINKING_BLOCK]}, + {"role": "system", "content": _TWO_TEXT_BLOCKS}, + ] + + assert _outgoing_messages("databricks-meta-llama-3-3-70b-instruct", messages, _JSON_OBJECT) == [ + earlier_turn, + {"role": "system", "content": "Alpha rule.\n\nReply in JSON"}, + ] + + +@pytest.mark.parametrize( + "content,expected_content", + [ + pytest.param( + [{"type": "text", "text": "Reply in JSON"}, {"type": "text", "text": "\n\n"}], + "Reply in JSON", + id="one-non-blank-part-goes-out-as-before", + ), + pytest.param( + [{"type": "text", "text": "Reply"}, {"type": "text", "text": " "}, {"type": "text", "text": "in JSON"}], + "Reply in JSON", + id="blank-separator-between-parts-kept", + ), + ], +) +def test_transform_request_json_object_joins_only_two_or_more_non_blank_parts( + content: list[object], expected_content: str +) -> None: + assert _outgoing_messages( + "databricks-meta-llama-3-3-70b-instruct", [{"role": "user", "content": content}], _JSON_OBJECT + ) == [{"role": "user", "content": expected_content}] + + +def test_transform_request_json_object_keeps_message_cache_control_parts_as_sent_without_drop_params() -> None: + message: Final[AllMessageValues] = { + "role": "user", + "content": _TWO_TEXT_BLOCKS, + "cache_control": {"type": "ephemeral"}, + } + + assert _outgoing_messages("databricks-meta-llama-3-3-70b-instruct", [message], {}) == _outgoing_messages( + "databricks-meta-llama-3-3-70b-instruct", [message], _JSON_OBJECT + ) + + +def test_transform_request_json_object_reads_the_global_drop_params(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "drop_params", True) + + assert _outgoing_messages( + "databricks-meta-llama-3-3-70b-instruct", + [{"role": "user", "content": [{"type": "text", "text": "Reply"}, _CACHED_BLOCK]}], + _JSON_OBJECT, + ) == [{"role": "user", "content": "Reply in JSON"}] + + def test_transform_request_strips_thinking_blocks_and_reasoning_content(): """Regression for LIT-6762: replaying an assistant turn that litellm decorated with `thinking_blocks` / `reasoning_content` made Databricks 400 with @@ -2775,6 +3046,74 @@ def test_claude_json_schema_stream_split_across_json_tools_keeps_the_non_stream_ assert [choice.finish_reason for choice in choices if choice.finish_reason] == [answer.finish_reason] +_HOSTED_MCP_TOOL: Final = { + "type": "mcp", + "server_label": "deepwiki", + "server_url": "https://mcp.deepwiki.com/mcp", + "require_approval": "never", +} +_WEATHER_FUNCTION_TOOL: Final = { + "type": "function", + "function": {"name": "get_weather", "parameters": {"type": "object", "properties": {}}}, +} + + +def test_databricks_claude_answers_400_for_a_hosted_mcp_tool_without_drop_params(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(litellm.utils.UnsupportedParamsError) as raised: + DatabricksConfig().map_openai_params( + non_default_params={"tools": [_WEATHER_FUNCTION_TOOL, _HOSTED_MCP_TOOL]}, + optional_params={}, + model="databricks-claude-opus-5-5", + drop_params=False, + ) + assert raised.value.status_code == 400 + assert "databricks-claude-opus-5-5 on databricks " in str(raised.value) + assert "`deepwiki`" in str(raised.value) + + +def test_databricks_claude_drops_a_hosted_mcp_tool_under_drop_params_and_keeps_the_function_tool(monkeypatch, caplog): + monkeypatch.setattr(litellm, "drop_params", False) + with caplog.at_level("WARNING"): + mapped: Final = DatabricksConfig().map_openai_params( + non_default_params={"tools": [_HOSTED_MCP_TOOL, _WEATHER_FUNCTION_TOOL]}, + optional_params={}, + model="databricks-claude-opus-5-5", + drop_params=True, + ) + assert [tool["function"]["name"] for tool in mapped["tools"]] == ["get_weather"] + assert any( + "`deepwiki`" in record.getMessage() and " on databricks (" in record.getMessage() for record in caplog.records + ) + + +def test_databricks_claude_sends_no_tools_when_drop_params_drops_every_hosted_mcp_tool(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + mapped: Final = DatabricksConfig().map_openai_params( + non_default_params={"tools": [_HOSTED_MCP_TOOL], "tool_choice": "required"}, + optional_params={}, + model="databricks-claude-opus-5-5", + drop_params=True, + ) + assert "tools" not in mapped + assert "tool_choice" not in mapped + + +def test_databricks_claude_still_sends_max_completion_tokens_as_max_tokens_when_dropping_a_hosted_mcp_tool( + monkeypatch, +): + monkeypatch.setattr(litellm, "drop_params", False) + mapped: Final = DatabricksConfig().map_openai_params( + non_default_params={"tools": [_HOSTED_MCP_TOOL, _WEATHER_FUNCTION_TOOL], "max_completion_tokens": 64}, + optional_params={}, + model="databricks-claude-opus-5-5", + drop_params=True, + ) + assert mapped["max_tokens"] == 64 + assert "max_completion_tokens" not in mapped + assert [tool["function"]["name"] for tool in mapped["tools"]] == ["get_weather"] + + @pytest.mark.parametrize( ("model", "base_model"), [("databricks-claude-haiku-5-5", None), ("support-bot-endpoint", "databricks/databricks-claude-haiku-5-5")], diff --git a/tests/unit/llms/deepseek/chat/test_deepseek_chat_transformation.py b/tests/unit/llms/deepseek/chat/test_deepseek_chat_transformation.py index 56af47aeebe..160ed3950b3 100644 --- a/tests/unit/llms/deepseek/chat/test_deepseek_chat_transformation.py +++ b/tests/unit/llms/deepseek/chat/test_deepseek_chat_transformation.py @@ -955,6 +955,51 @@ def test_get_complete_url_ignores_a_trailing_slash_on_api_base(api_base: str): ) +_HOSTED_MCP_TOOL: Final = { + "type": "mcp", + "server_label": "deepwiki", + "server_url": "https://mcp.deepwiki.com/mcp", + "require_approval": "never", +} + + +def test_a_hosted_mcp_tool_answers_400_on_deepseek_without_drop_params(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(litellm.utils.UnsupportedParamsError) as raised: + DeepSeekChatConfig().map_openai_params( + {"tools": [_function_tool("shell"), _HOSTED_MCP_TOOL]}, {}, "deepseek-v4-flash", False + ) + assert raised.value.status_code == 400 + assert "deepseek-v4-flash on deepseek " in str(raised.value) + assert "`deepwiki`" in str(raised.value) + + +def test_deepseek_drops_a_hosted_mcp_tool_under_drop_params_and_keeps_the_function_tool(monkeypatch, caplog): + monkeypatch.setattr(litellm, "drop_params", False) + with caplog.at_level("WARNING"): + result: Final = DeepSeekChatConfig().map_openai_params( + {"tools": [_HOSTED_MCP_TOOL, _function_tool("shell")], "tool_choice": "auto"}, {}, "deepseek-v4-flash", True + ) + assert result["tools"] == [_function_tool("shell")] + assert result["tool_choice"] == "auto" + assert any( + "`deepwiki`" in record.getMessage() and " on deepseek (" in record.getMessage() for record in caplog.records + ) + + +def test_deepseek_sends_no_tools_when_drop_params_drops_every_hosted_mcp_tool(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + result: Final = DeepSeekChatConfig().map_openai_params( + {"tools": [_HOSTED_MCP_TOOL], "tool_choice": "required", "parallel_tool_calls": True}, + {}, + "deepseek-v4-flash", + True, + ) + assert "tools" not in result + assert "tool_choice" not in result + assert "parallel_tool_calls" not in result + + def _budget_warnings(caplog: pytest.LogCaptureFixture) -> tuple[str, ...]: return tuple( r.getMessage() diff --git a/tests/unit/llms/gigachat/chat/test_transformation.py b/tests/unit/llms/gigachat/chat/test_transformation.py index 9116a4d1979..4d8a916ea87 100644 --- a/tests/unit/llms/gigachat/chat/test_transformation.py +++ b/tests/unit/llms/gigachat/chat/test_transformation.py @@ -7,10 +7,12 @@ Run with: pytest tests/unit/llms/gigachat/chat/test_transformation.py -v import asyncio import importlib +from typing import Final import pytest import litellm +from litellm.llms.gigachat.chat.transformation import GigaChatConfig from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -575,3 +577,56 @@ _SCALAR_DEFAULTS = { "api_key": getattr(litellm, "api_key", None), "cohere_key": getattr(litellm, "cohere_key", None), } + + +_HOSTED_MCP_TOOL: Final = { + "type": "mcp", + "server_label": "deepwiki", + "server_url": "https://mcp.deepwiki.com/mcp", + "require_approval": "never", +} +_WEATHER_FUNCTION_TOOL: Final = { + "type": "function", + "function": {"name": "get_weather", "parameters": {"type": "object", "properties": {}}}, +} + + +def test_a_hosted_mcp_tool_answers_400_on_gigachat_without_drop_params(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(litellm.utils.UnsupportedParamsError) as raised: + GigaChatConfig().map_openai_params( + non_default_params={"tools": [_WEATHER_FUNCTION_TOOL, _HOSTED_MCP_TOOL]}, + optional_params={}, + model="GigaChat-2-Max", + drop_params=False, + ) + assert raised.value.status_code == 400 + assert "GigaChat-2-Max on gigachat " in str(raised.value) + assert "`deepwiki`" in str(raised.value) + + +def test_gigachat_drops_a_hosted_mcp_tool_under_drop_params_and_keeps_the_function_tool(monkeypatch, caplog): + monkeypatch.setattr(litellm, "drop_params", False) + with caplog.at_level("WARNING"): + mapped: Final = GigaChatConfig().map_openai_params( + non_default_params={"tools": [_HOSTED_MCP_TOOL, _WEATHER_FUNCTION_TOOL]}, + optional_params={}, + model="GigaChat-2-Max", + drop_params=True, + ) + assert [function["name"] for function in mapped["functions"]] == ["get_weather"] + assert any( + "`deepwiki`" in record.getMessage() and " on gigachat (" in record.getMessage() for record in caplog.records + ) + + +def test_gigachat_sends_no_functions_when_drop_params_drops_every_hosted_mcp_tool(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + mapped: Final = GigaChatConfig().map_openai_params( + non_default_params={"tools": [_HOSTED_MCP_TOOL], "tool_choice": "auto"}, + optional_params={}, + model="GigaChat-2-Max", + drop_params=True, + ) + assert "functions" not in mapped + assert "function_call" not in mapped diff --git a/tests/unit/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py b/tests/unit/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py index fc5ae3631ff..94f1873faba 100644 --- a/tests/unit/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py +++ b/tests/unit/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py @@ -3,14 +3,15 @@ import json from typing import Final from unittest.mock import MagicMock, patch +import httpx import pytest - from litellm.constants import ( DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET, DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET, ) from litellm.llms.hosted_vllm.chat.transformation import HostedVLLMChatConfig +from litellm.types.utils import ModelResponse def test_hosted_vllm_chat_transformation_file_url(): @@ -401,6 +402,64 @@ def test_hosted_vllm_custom_tools_use_top_level_input_schema(): assert tools[0]["function"]["parameters"] == input_schema +def _transform_vllm_chat_response(choice_extras: dict[str, object]) -> ModelResponse: + payload: Final = { + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 1, + "model": "glm", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "1, 2, 3, 4"}, + "logprobs": None, + "finish_reason": "stop", + **choice_extras, + } + ], + "usage": {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8}, + } + raw_response: Final = httpx.Response( + 200, json=payload, request=httpx.Request("POST", "http://vllm.local/v1/chat/completions") + ) + return HostedVLLMChatConfig().transform_response( + model="hosted_vllm/glm", + raw_response=raw_response, + model_response=ModelResponse(), + logging_obj=MagicMock(), + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=None, + ) + + +@pytest.mark.parametrize( + "choice_extras", + [ + {"stop_reason": ", 5", "token_ids": [16, 11, 220]}, + {"stop_reason": 128009}, + {"stop_reason": None}, + ], +) +def test_hosted_vllm_keeps_vllm_choice_fields_on_the_choice(choice_extras: dict[str, object]): + response: Final = _transform_vllm_chat_response(choice_extras) + + dumped_choice: Final[dict[str, dict[str, object]]] = json.loads(response.model_dump_json())["choices"][0] + for field, value in choice_extras.items(): + assert dumped_choice[field] == value + assert dumped_choice["provider_specific_fields"][field] == value + + +def test_hosted_vllm_does_not_invent_vllm_choice_fields(): + response: Final = _transform_vllm_chat_response({}) + + dumped_choice: Final[dict[str, object]] = json.loads(response.model_dump_json())["choices"][0] + assert "stop_reason" not in dumped_choice + assert "token_ids" not in dumped_choice + + def _strict_weather_tool() -> dict[str, object]: return { "type": "function", diff --git a/tests/unit/llms/ollama/test_ollama_chat_transformation.py b/tests/unit/llms/ollama/test_ollama_chat_transformation.py index 6a80f29ac43..98ca5aee407 100644 --- a/tests/unit/llms/ollama/test_ollama_chat_transformation.py +++ b/tests/unit/llms/ollama/test_ollama_chat_transformation.py @@ -6,6 +6,7 @@ import sys from typing import Final, cast import pytest +import respx from pydantic import BaseModel sys.path.insert( @@ -1211,6 +1212,72 @@ def test_get_complete_url_ignores_a_trailing_slash_on_api_base(api_base: str): ) +_DEEPWIKI_MCP_TOOL: Final = { + "type": "mcp", + "server_label": "deepwiki", + "server_url": "https://mcp.deepwiki.com/mcp", + "require_approval": "never", +} +_WEATHER_TOOL: Final = { + "type": "function", + "function": {"name": "get_weather", "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}}, +} + + +def _ollama_chat_route(respx_mock: respx.MockRouter) -> respx.Route: + return respx_mock.post("http://ollama.test:11434/api/chat").respond( + json={ + "model": "qwen3:0.6b", + "message": {"role": "assistant", "content": "Sunny"}, + "done": True, + "prompt_eval_count": 5, + "eval_count": 3, + } + ) + + +def test_ollama_chat_answers_400_for_a_hosted_mcp_tool_it_cannot_run( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + route: Final = _ollama_chat_route(respx_mock) + + with pytest.raises(litellm.UnsupportedParamsError) as raised: + litellm.completion( + model="ollama_chat/qwen3:0.6b", + messages=[{"role": "user", "content": "weather in Paris?"}], + tools=[_DEEPWIKI_MCP_TOOL, _WEATHER_TOOL], + api_base="http://ollama.test:11434", + ) + + assert raised.value.status_code == 400 + assert "qwen3:0.6b on ollama_chat cannot run the hosted MCP tool `deepwiki`" in str(raised.value) + assert not route.called + + +@pytest.mark.parametrize( + ("tools", "sent_tools"), + [([_DEEPWIKI_MCP_TOOL, _WEATHER_TOOL], [_WEATHER_TOOL]), ([_DEEPWIKI_MCP_TOOL], None)], + ids=["beside-a-function", "alone"], +) +def test_ollama_chat_drop_params_sends_only_the_tools_ollama_can_call( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, tools: list[object], sent_tools: object +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + route: Final = _ollama_chat_route(respx_mock) + + response: Final = litellm.completion( + model="ollama_chat/qwen3:0.6b", + messages=[{"role": "user", "content": "weather in Paris?"}], + tools=tools, + api_base="http://ollama.test:11434", + drop_params=True, + ) + + assert json.loads(route.calls.last.request.content).get("tools") == sent_tools + assert response.choices[0].message.content == "Sunny" + + def _warnings_naming(caplog: pytest.LogCaptureFixture, param: str) -> tuple[str, ...]: return tuple( r.getMessage() for r in caplog.records if r.levelno == logging.WARNING and f"Dropping {param}=" in r.getMessage() diff --git a/tests/unit/llms/ollama/test_ollama_completion_transformation.py b/tests/unit/llms/ollama/test_ollama_completion_transformation.py index edea39bb10c..26addf1b24c 100644 --- a/tests/unit/llms/ollama/test_ollama_completion_transformation.py +++ b/tests/unit/llms/ollama/test_ollama_completion_transformation.py @@ -993,6 +993,77 @@ def test_get_complete_url_ignores_a_trailing_slash_on_api_base(api_base: str): == "http://gateway.example.com/ollama/api/generate" ) + +_DEEPWIKI_MCP_TOOL: Final = { + "type": "mcp", + "server_label": "deepwiki", + "server_url": "https://mcp.deepwiki.com/mcp", + "require_approval": "never", +} +_WEATHER_TOOL: Final = { + "type": "function", + "function": { + "name": "get_weather", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}, + }, +} + + +def _ollama_generate_route(respx_mock: respx.MockRouter) -> respx.Route: + return respx_mock.post("http://ollama.test:11434/api/generate").respond( + json={"model": "llama3.2", "response": "Sunny", "done": True, "prompt_eval_count": 5, "eval_count": 3} + ) + + +def test_ollama_answers_400_for_a_hosted_mcp_tool_instead_of_pasting_it_into_the_prompt( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + monkeypatch.setattr(litellm, "add_function_to_prompt", False) + route: Final = _ollama_generate_route(respx_mock) + + with pytest.raises(litellm.UnsupportedParamsError) as raised: + litellm.completion( + model="ollama/llama3.2", + messages=[{"role": "user", "content": "weather in Paris?"}], + tools=[_DEEPWIKI_MCP_TOOL, _WEATHER_TOOL], + api_base="http://ollama.test:11434", + ) + + assert raised.value.status_code == 400 + assert "cannot run the hosted MCP tool `deepwiki`" in str(raised.value) + assert not route.called + + +@pytest.mark.parametrize( + ("tools", "prompt_names_weather"), + [([_DEEPWIKI_MCP_TOOL, _WEATHER_TOOL], True), ([_DEEPWIKI_MCP_TOOL], False)], + ids=["beside-a-function", "alone"], +) +def test_ollama_drop_params_keeps_a_hosted_mcp_tool_out_of_the_prompt( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, tools: list[object], prompt_names_weather: bool +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + monkeypatch.setattr(litellm, "add_function_to_prompt", False) + route: Final = _ollama_generate_route(respx_mock) + + litellm.completion( + model="ollama/llama3.2", + messages=[{"role": "user", "content": "weather in Paris?"}], + tools=tools, + api_base="http://ollama.test:11434", + drop_params=True, + ) + + body: Final = json.loads(route.calls.last.request.content) + prompt: Final = body["prompt"] + assert (body.get("format") == "json") is prompt_names_weather + assert "deepwiki" not in prompt + assert ("get_weather" in prompt) is prompt_names_weather + assert prompt.startswith("### User:\nweather in Paris?") + assert prompt_names_weather or prompt == "### User:\nweather in Paris?\n\n" + + def _reasoning_effort_warnings(caplog: pytest.LogCaptureFixture) -> tuple[str, ...]: return tuple( r.getMessage() diff --git a/tests/unit/llms/opencode/chat/test_transformation.py b/tests/unit/llms/opencode/chat/test_transformation.py index 0a5e49ca1b2..91361a3d81d 100644 --- a/tests/unit/llms/opencode/chat/test_transformation.py +++ b/tests/unit/llms/opencode/chat/test_transformation.py @@ -354,6 +354,17 @@ def test_gemini_custom_base_normalizes_version_and_models(go_catalog: None, suff assert response.choices[0].message.content == "hello" +def test_gemini_hosted_mcp_error_names_the_opencode_go_route(go_catalog: None) -> None: + with pytest.raises(litellm.UnsupportedParamsError, match="fixture-gemini on opencode_go cannot run"): + litellm.completion( + model="opencode_go/fixture-gemini", + api_key="fake-go-key", + api_base="https://gateway.example/google", + messages=[{"role": "user", "content": "hello"}], + tools=[{"type": "mcp", "server_label": "deepwiki", "server_url": "https://mcp.deepwiki.com/mcp"}], + ) + + def test_gemini_sync_stream_preserves_tool_calls_and_usage(go_catalog: None) -> None: chunk: Final = { "candidates": [ diff --git a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index ef71f09fc7c..2eecb83e2df 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -10861,6 +10861,90 @@ async def test_gemini_sends_a_boolean_for_a_gateway_default_written_as_text_or_n assert sent["includeThoughts"] is expected +_HOSTED_MCP_TOOL: Final = { + "type": "mcp", + "server_label": "deepwiki", + "server_url": "https://mcp.deepwiki.com/mcp", + "require_approval": "never", +} +_WEATHER_FUNCTION_TOOL: Final = { + "type": "function", + "function": {"name": "get_weather", "parameters": {"type": "object", "properties": {}}}, +} +_GEMINI_ROUTES: Final = ( + pytest.param(VertexGeminiConfig, "vertex_ai", id="vertex_ai"), + pytest.param(GoogleAIStudioGeminiConfig, "gemini", id="gemini"), +) + + +@pytest.mark.parametrize(("config_class", "route"), _GEMINI_ROUTES) +def test_a_hosted_mcp_tool_answers_400_on_gemini_without_drop_params(monkeypatch, config_class, route): + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(litellm.utils.UnsupportedParamsError) as raised: + config_class().map_openai_params( + non_default_params={"tools": [_WEATHER_FUNCTION_TOOL, _HOSTED_MCP_TOOL]}, + optional_params={}, + model="gemini-3.8-flash", + drop_params=False, + ) + assert raised.value.status_code == 400 + assert f"gemini-3.8-flash on {route} " in str(raised.value) + assert "`deepwiki`" in str(raised.value) + assert "litellm_proxy/mcp/" in str(raised.value) + + +@pytest.mark.parametrize(("config_class", "route"), _GEMINI_ROUTES) +def test_gemini_drops_a_hosted_mcp_tool_under_drop_params_and_keeps_the_function_tool( + monkeypatch, caplog, config_class, route +): + monkeypatch.setattr(litellm, "drop_params", False) + with caplog.at_level("WARNING"): + mapped: Final = config_class().map_openai_params( + non_default_params={"tools": [_HOSTED_MCP_TOOL, _WEATHER_FUNCTION_TOOL]}, + optional_params={}, + model="gemini-3.8-flash", + drop_params=True, + ) + assert "deepwiki" not in json.dumps(mapped) + assert "get_weather" in json.dumps(mapped["tools"]) + assert any( + "`deepwiki`" in record.getMessage() and f" on {route} (" in record.getMessage() for record in caplog.records + ) + + +@pytest.mark.parametrize(("config_class", "route"), _GEMINI_ROUTES) +def test_gemini_sends_no_tools_when_drop_params_drops_every_hosted_mcp_tool(monkeypatch, config_class, route): + monkeypatch.setattr(litellm, "drop_params", False) + mapped: Final = config_class().map_openai_params( + non_default_params={"tools": [_HOSTED_MCP_TOOL], "tool_choice": "required"}, + optional_params={}, + model="gemini-3.8-flash", + drop_params=True, + ) + assert "tools" not in mapped + assert "tool_choice" not in mapped + + +@pytest.mark.parametrize(("config_class", "route"), _GEMINI_ROUTES) +def test_gemini_keeps_tool_choice_none_for_legacy_functions_after_dropping_every_hosted_mcp_tool( + monkeypatch, config_class, route +): + monkeypatch.setattr(litellm, "drop_params", False) + mapped: Final = config_class().map_openai_params( + non_default_params={ + "tools": [_HOSTED_MCP_TOOL], + "functions": [_WEATHER_FUNCTION_TOOL["function"]], + "tool_choice": "none", + }, + optional_params={}, + model="gemini-3.8-flash", + drop_params=True, + ) + assert "deepwiki" not in json.dumps(mapped) + assert "get_weather" in json.dumps(mapped["tools"]) + assert mapped["tool_choice"] == {"functionCallingConfig": {"mode": "NONE"}} + + class _BodyInChunks(httpx.SyncByteStream, httpx.AsyncByteStream): def __init__(self, chunks: tuple[bytes, ...]) -> None: self._chunks: Final = chunks diff --git a/tests/unit/proxy/common_utils/test_codex_model_catalog.py b/tests/unit/proxy/common_utils/test_codex_model_catalog.py index 509303d9635..e3d941a6aa4 100644 --- a/tests/unit/proxy/common_utils/test_codex_model_catalog.py +++ b/tests/unit/proxy/common_utils/test_codex_model_catalog.py @@ -1,9 +1,12 @@ import json +from collections.abc import Mapping from types import MappingProxyType +from typing import Final, Literal import pytest import litellm +from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.proxy.common_utils.codex_model_catalog import ( CODEX_CATALOG_BYTE_LIMIT, CodexCatalogRow, @@ -233,6 +236,140 @@ def test_alias_with_deployment_defaults_caps_codexs_window_at_its_listed_max_inp assert json.loads(body.json)["models"][0]["context_window"] == expected +_WIDE_STOCK = MappingProxyType( + { + "gpt-6-sol": _stock("gpt-6-sol").model_copy(update={"context_window": 272000, "max_context_window": 872000}), + "gpt-5.5": _stock("gpt-5.5").model_copy(update={"context_window": 272000}), + } +) + + +@pytest.mark.parametrize( + ("configured", "expected"), + [ + ((500000,), 500000), + ((100000,), 100000), + ((1050000,), 872000), + ((600000, 400000), 400000), + ((600000, None), 272000), + ((100000, None), 100000), + ((None,), 272000), + ((), 272000), + ((0,), 272000), + ((-5, 400000), 272000), + ((0, 100000), 100000), + ], +) +@pytest.mark.parametrize("alias_defaults", ["codex", "deployment"]) +def test_a_known_model_and_its_alias_take_the_smallest_window_any_deployment_takes( + configured: tuple[int | None, ...], expected: int, alias_defaults: Literal["codex", "deployment"] +) -> None: + rows: Final = ( + _row("gpt-6-sol", max_input_tokens=1050000, configured_max_input_tokens=configured), + _row( + "sol-1m", + upstream_model="openai/gpt-6-sol", + max_input_tokens=1050000, + configured_max_input_tokens=configured, + ), + ) + + body: Final = codex_models_response_json( + rows, stock=_WIDE_STOCK, instructions=_PROMPT, alias_defaults=alias_defaults + ) + + assert {entry["slug"]: entry["context_window"] for entry in json.loads(body.json)["models"]} == { + "gpt-6-sol": expected, + "sol-1m": expected, + } + + +def test_a_configured_window_goes_past_codexs_own_when_the_stock_entry_sets_no_ceiling() -> None: + (entry,) = json.loads( + codex_models_response_json( + (_row("gpt-5.5", configured_max_input_tokens=(1050000,)),), stock=_WIDE_STOCK, instructions=_PROMPT + ).json + )["models"] + + assert entry["context_window"] == 1050000 + + +def test_a_deployment_without_a_limit_keeps_a_windowless_stock_entry_as_codex_ships_it() -> None: + windowless: Final = MappingProxyType({"gpt-5.5": _stock("gpt-5.5").model_copy(update={"context_window": None})}) + + (entry,) = json.loads( + codex_models_response_json( + (_row("gpt-5.5", configured_max_input_tokens=(1050000, None)),), stock=windowless, instructions=_PROMPT + ).json + )["models"] + + assert entry["context_window"] is None + + +def test_model_list_body_sizes_a_known_models_window_off_the_deployments_the_key_team_routes_to() -> None: + router: Final = litellm.Router( + model_list=[ + { + "model_name": "gpt-6-astra", + "litellm_params": {"model": "openai/gpt-6-astra"}, + "model_info": {"max_input_tokens": 922000}, + }, + { + "model_name": "gpt-6-astra", + "litellm_params": {"model": "openai/gpt-6-astra", "api_base": "https://team-2.example"}, + "model_info": {"team_id": "team-2", "max_input_tokens": 400000}, + }, + { + "model_name": "gpt-6-astra", + "litellm_params": {"model": "openai/gpt-6-astra", "api_base": "https://team-3.example"}, + "model_info": {"team_id": "team-3"}, + }, + ], + model_group_alias={"astra-1m": "gpt-6-astra"}, + ) + listing: Final = tuple( + {"id": name, "object": "model", "created": 0, "owned_by": "openai", "mode": "chat"} + for name in ("gpt-6-astra", "astra-1m") + ) + entries: Final = (("gpt-6-astra", "gpt-6-astra"), ("astra-1m", "astra-1m")) + stock: Final = bundled_codex_models()["gpt-6-astra"] + assert stock.context_window is not None and stock.max_context_window is not None + assert stock.context_window < 400000 < stock.max_context_window < 922000 + + def windows(team_id: str) -> dict[str, int | None]: + models: Final = json.loads(codex_model_list_body(listing, entries, router, team_id))["models"] + return {entry["slug"]: entry["context_window"] for entry in models} + + assert windows("team-1") == {"gpt-6-astra": stock.max_context_window, "astra-1m": stock.max_context_window} + assert windows("team-2") == {"gpt-6-astra": 400000, "astra-1m": 400000} + assert windows("team-3") == {"gpt-6-astra": stock.context_window, "astra-1m": stock.context_window} + + +def test_model_list_body_skips_a_limit_a_saved_model_info_response_copied_from_the_cost_map() -> None: + catalog_limit: Final = GetModelCostMap.loaded_model_cost_map()["gpt-6-astra"]["max_input_tokens"] + stock: Final = bundled_codex_models()["gpt-6-astra"] + assert isinstance(catalog_limit, int) and stock.max_context_window is not None + assert stock.context_window is not None and stock.context_window < 600000 < stock.max_context_window + + def window(model_info: Mapping[str, object]) -> int | None: + router: Final = litellm.Router( + model_list=[ + { + "model_name": "gpt-6-astra", + "litellm_params": {"model": "openai/gpt-6-astra"}, + "model_info": dict(model_info), + } + ] + ) + listing: Final = ({"id": "gpt-6-astra", "object": "model", "created": 0, "owned_by": "openai", "mode": "chat"},) + (entry,) = json.loads(codex_model_list_body(listing, (("gpt-6-astra", "gpt-6-astra"),), router))["models"] + return entry["context_window"] + + assert window({"key": "gpt-6-astra", "max_input_tokens": catalog_limit}) == stock.context_window + assert window({"key": "gpt-6-astra", "max_input_tokens": 600000}) == 600000 + assert window({"max_input_tokens": catalog_limit}) == min(catalog_limit, stock.max_context_window) + + def test_alias_with_deployment_defaults_drops_codexs_tier_and_upgrade_while_the_native_name_keeps_them(): rows = (_row("gpt-6-sol"), _row("team-sol", upstream_model="gpt-6-sol"), _row("team-55", upstream_model="gpt-5.5")) codex = {entry["slug"]: entry for entry in _models(*rows)} @@ -599,9 +736,14 @@ def test_catalog_rows_read_the_router_by_each_entry_lookup_id(): "model_name_team-1_c0ffee" ).supported_reasoning_efforts, search_tool_support=(None, None), + configured_max_input_tokens=(None, None), ), CodexCatalogRow( - id="plain", upstream_model="openai/some-unmapped-model", service_tiers=(None,), search_tool_support=(None,) + id="plain", + upstream_model="openai/some-unmapped-model", + service_tiers=(None,), + search_tool_support=(None,), + configured_max_input_tokens=(None,), ), ) @@ -631,6 +773,7 @@ def test_catalog_rows_read_an_alias_off_its_target_under_the_alias_name(): service_tiers=(["ultrafast"],), supported_reasoning_efforts=router.get_model_group_info("gpt-6-astra").supported_reasoning_efforts, search_tool_support=(None,), + configured_max_input_tokens=(None,), ) (entry,) = json.loads(codex_model_list_body(listing, (("gpt-6", "gpt-6"),), router))["models"] assert (entry["slug"], entry["display_name"], [tier["id"] for tier in entry["service_tiers"]]) == ( diff --git a/tests/unit/proxy/common_utils/test_token_count_settings.py b/tests/unit/proxy/common_utils/test_token_count_settings.py index 43edc7e449f..68910f52b02 100644 --- a/tests/unit/proxy/common_utils/test_token_count_settings.py +++ b/tests/unit/proxy/common_utils/test_token_count_settings.py @@ -131,6 +131,107 @@ def test_shipped_australian_count_region_stays_with_an_australian_deployment() - assert settings.count_model == TokenCountModel(model="au.anthropic.claude-opus-5-5", region="ap-southeast-4") +_ABSENT: Final = object() + + +@pytest.mark.parametrize( + ("key", "field"), + [ + ("token_count_tool_framing", "tool_framing"), + ("token_count_responses_framing", "responses_framing"), + ("token_count_responses_tool_framing", "responses_tool_framing"), + ], +) +@pytest.mark.parametrize( + ("configured", "default", "expected"), + [ + (_ABSENT, 151, 151), + (40, 151, 40), + (0, 151, 0), + (40.0, _ABSENT, 40), + (-1, 151, 151), + (True, 151, 151), + ("40", _ABSENT, None), + ], +) +def test_each_request_kind_framing_reads_its_value_like_framing( + key: str, field: str, configured: object, default: object, expected: int | None +) -> None: + settings: Final = resolve_token_count_settings( + configured={} if configured is _ABSENT else {key: configured}, + default={} if default is _ABSENT else {key: default}, + deployment_region=_in_us_west_2, + ) + assert getattr(settings, field) == expected + + +def _hidden_prompts(settings: TokenCountSettings) -> tuple[int, int, int, int]: + return ( + settings.hidden_prompt_tokens(carries_tools=False, responses_api=False), + settings.hidden_prompt_tokens(carries_tools=True, responses_api=False), + settings.hidden_prompt_tokens(carries_tools=False, responses_api=True), + settings.hidden_prompt_tokens(carries_tools=True, responses_api=True), + ) + + +@pytest.mark.parametrize( + ("default", "expected"), + [ + ({}, (0, 0, 0, 0)), + ({"token_count_framing": 23}, (23, 23, 23, 23)), + ({"token_count_framing": 185, "token_count_tool_framing": 214}, (185, 214, 185, 214)), + ( + { + "token_count_framing": 149, + "token_count_tool_framing": 300, + "token_count_responses_framing": 145, + "token_count_responses_tool_framing": 104, + }, + (149, 300, 145, 104), + ), + ({"token_count_framing": 149, "token_count_tool_framing": 0}, (149, 0, 149, 0)), + ({"token_count_framing": 149, "token_count_responses_tool_framing": 0}, (149, 149, 149, 0)), + ({"token_count_framing": 149, "token_count_responses_framing": 145}, (149, 149, 145, 145)), + ], +) +def test_hidden_prompt_is_the_most_specific_framing_set_for_the_request_api_and_tools( + default: dict[str, object], expected: tuple[int, int, int, int] +) -> None: + settings: Final = resolve_token_count_settings(configured={}, default=default, deployment_region=_in_us_west_2) + assert _hidden_prompts(settings) == expected + + +_PER_API_COST_MAP: Final = { + "token_count_framing": 149, + "token_count_tool_framing": 300, + "token_count_responses_framing": 145, + "token_count_responses_tool_framing": 104, +} + + +@pytest.mark.parametrize( + ("configured", "expected"), + [ + ({}, (149, 300, 145, 104)), + ({"token_count_framing": 500}, (500, 500, 500, 500)), + ({"token_count_framing": 0}, (0, 0, 0, 0)), + ({"token_count_tool_framing": 250}, (149, 250, 145, 250)), + ({"token_count_framing": 500, "token_count_tool_framing": 250}, (500, 250, 500, 250)), + ({"token_count_responses_framing": 120}, (149, 300, 120, 120)), + ({"token_count_tool_framing": 250, "token_count_responses_framing": 120}, (149, 250, 120, 120)), + ({"token_count_responses_tool_framing": 90}, (149, 300, 145, 90)), + ({"token_count_framing": -1}, (149, 300, 145, 104)), + ], +) +def test_an_operator_framing_applies_to_every_request_kind_the_cost_map_splits_more_finely( + configured: dict[str, object], expected: tuple[int, int, int, int] +) -> None: + settings: Final = resolve_token_count_settings( + configured=configured, default=_PER_API_COST_MAP, deployment_region=_in_us_west_2 + ) + assert _hidden_prompts(settings) == expected + + def test_no_setting_anywhere_means_the_model_counts_itself_without_framing() -> None: assert resolve_token_count_settings( configured={}, default={}, deployment_region=_in_us_west_2 diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_mcp_security.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_mcp_security.py index 0e6670c26e7..43238e37f63 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_mcp_security.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_mcp_security.py @@ -6,19 +6,24 @@ and allows requests with only registered servers. Covers both /chat/completions and /responses API paths (same pre_call_hook logic, different call_type). """ -from typing import Literal -from unittest.mock import MagicMock, patch +from typing import Final, Literal +from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import HTTPException import litellm +from litellm.caching.caching import DualCache from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.mcp_security import initialize_guardrail from litellm.proxy.guardrails.guardrail_hooks.mcp_security.mcp_security_guardrail import ( MCPSecurityGuardrail, ) +from litellm.proxy.mcp_server.mcp_server_manager import MCPServerManager from litellm.types.guardrails import Guardrail, GuardrailEventHooks, LitellmParams +from litellm.types.mcp import MCPTransport +from litellm.types.mcp_server.mcp_server_manager import MCPServer +from litellm.types.mcp_server.mcp_toolset import MCPToolset @pytest.fixture @@ -60,93 +65,67 @@ class TestExtractMCPServerNames: assert names == set() +def _manager_with_saved_server() -> MCPServerManager: + manager: Final = MCPServerManager() + server: Final = MCPServer( + server_id="3f6b2c1e-saved-zapier", + name="zapier", + server_name="zapier", + alias="zap", + transport=MCPTransport.http, + url="https://upstream.example/mcp", + ) + manager.registry[server.server_id] = server + return manager + + class TestMCPSecurityGuardrailPreCall: @pytest.mark.asyncio - @patch( - "litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager" - ) - async def test_blocks_unregistered_server_chat_completions( - self, mock_manager, guardrail - ): - """Simulates /chat/completions with an unregistered MCP server.""" - mock_manager.get_registry.return_value = {"zapier": MagicMock()} - + @pytest.mark.parametrize("call_type", ["acompletion", "aresponses"]) + async def test_blocks_unregistered_server(self, guardrail, call_type): data = { "tools": [ {"type": "mcp", "server_url": "litellm_proxy/mcp/zapier"}, {"type": "mcp", "server_url": "litellm_proxy/mcp/evil_server"}, ], "model": "gpt-4", - "messages": [{"role": "user", "content": "hi"}], "guardrails": ["test-mcp-security"], } - with pytest.raises(HTTPException) as exc_info: + with ( + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager", _manager_with_saved_server()), + pytest.raises(HTTPException) as exc_info, + ): await guardrail.async_pre_call_hook( user_api_key_dict=UserAPIKeyAuth(), cache=MagicMock(), data=data, - call_type="acompletion", + call_type=call_type, ) assert exc_info.value.status_code == 400 - assert "evil_server" in str(exc_info.value.detail) + assert exc_info.value.detail["unregistered_servers"] == ["evil_server"] @pytest.mark.asyncio - @patch( - "litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager" - ) - async def test_blocks_unregistered_server_responses_api( - self, mock_manager, guardrail - ): - """Simulates /responses with an unregistered MCP server.""" - mock_manager.get_registry.return_value = {"github": MagicMock()} - + @pytest.mark.parametrize("call_type", ["acompletion", "aresponses"]) + @pytest.mark.parametrize("reference", ["zapier", "zap", "3f6b2c1e-saved-zapier"]) + async def test_allows_saved_server_by_any_name_the_gateway_routes(self, guardrail, call_type, reference): data = { - "tools": [ - {"type": "mcp", "server_url": "litellm_proxy/mcp/unknown_server"}, - ], - "model": "gpt-4o", - "input": "What can you do?", - "guardrails": ["test-mcp-security"], - } - - with pytest.raises(HTTPException) as exc_info: - await guardrail.async_pre_call_hook( - user_api_key_dict=UserAPIKeyAuth(), - cache=MagicMock(), - data=data, - call_type="aresponses", - ) - assert exc_info.value.status_code == 400 - assert "unknown_server" in str(exc_info.value.detail) - - @pytest.mark.asyncio - @patch( - "litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager" - ) - async def test_allows_registered_servers(self, mock_manager, guardrail): - """All MCP servers are registered - request passes through.""" - mock_manager.get_registry.return_value = { - "zapier": MagicMock(), - "github": MagicMock(), - } - - data = { - "tools": [ - {"type": "mcp", "server_url": "litellm_proxy/mcp/zapier"}, - {"type": "mcp", "server_url": "litellm_proxy/mcp/github"}, - ], + "tools": [{"type": "mcp", "server_url": f"litellm_proxy/mcp/{reference}"}], "model": "gpt-4", - "messages": [{"role": "user", "content": "hi"}], "guardrails": ["test-mcp-security"], } - result = await guardrail.async_pre_call_hook( - user_api_key_dict=UserAPIKeyAuth(), - cache=MagicMock(), - data=data, - call_type="acompletion", - ) + with ( + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager", _manager_with_saved_server()), + ): + result = await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=MagicMock(), + data=data, + call_type=call_type, + ) assert result == data @pytest.mark.asyncio @@ -222,11 +201,13 @@ async def test_guardrail_observes_saved_server_creation_and_deletion_on_another_ prisma.writer_db = prisma.db prisma.db.litellm_mcpservertable.find_many = AsyncMock(side_effect=([row], [])) prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) + prisma.db.litellm_mcptoolsettable.find_first = AsyncMock(return_value=None) data = {"tools": [{"type": "mcp", "server_url": "litellm_proxy/mcp/peer-server"}], "guardrails": ["test-mcp-security"]} with ( patch("litellm.proxy.proxy_server.prisma_client", prisma), patch("litellm.proxy.proxy_server.general_settings", {}), + patch("litellm.proxy.proxy_server.user_api_key_cache", DualCache()), patch("litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager", manager), ): result = await guardrail.async_pre_call_hook(UserAPIKeyAuth(), MagicMock(), data, call_type) @@ -235,3 +216,34 @@ async def test_guardrail_observes_saved_server_creation_and_deletion_on_another_ await guardrail.async_pre_call_hook(UserAPIKeyAuth(), MagicMock(), data, call_type) assert exc.value.status_code == 400 assert exc.value.detail["unregistered_servers"] == ["peer-server"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call_type", ["acompletion", "aresponses"]) +async def test_allows_saved_toolset_by_name_and_blocks_unknown_name(guardrail, call_type): + saved_toolset: Final = MCPToolset(toolset_id="7c1d-research", toolset_name="research") + + async def find_toolset(*, where: dict[str, str]) -> MCPToolset | None: + return saved_toolset if where["toolset_name"] == saved_toolset.toolset_name else None + + prisma: Final = MagicMock() + prisma.writer_db = prisma.db + prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[]) + prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) + prisma.db.litellm_mcptoolsettable.find_first = AsyncMock(side_effect=find_toolset) + data: Final = { + "tools": [ + {"type": "mcp", "server_url": "litellm_proxy/mcp/research"}, + {"type": "mcp", "server_url": "litellm_proxy/mcp/evil_server"}, + ], + "guardrails": ["test-mcp-security"], + } + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), + patch("litellm.proxy.proxy_server.general_settings", {}), + patch("litellm.proxy.proxy_server.user_api_key_cache", DualCache()), + patch("litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager", MCPServerManager()), + pytest.raises(HTTPException) as exc, + ): + await guardrail.async_pre_call_hook(UserAPIKeyAuth(), MagicMock(), data, call_type) + assert exc.value.detail["unregistered_servers"] == ["evil_server"] diff --git a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py index ab6cfcfebfe..44c43713438 100644 --- a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py @@ -4896,6 +4896,9 @@ class TestModelInfoCostMapEchoFilter: [ ("bedrock/us.anthropic.claude-sonnet-5-5", "token_count_model", {}), ("bedrock/us.moonshotai.kimi-k3", "token_count_framing", 0), + ("azure_ai/grok-4-1-fast-reasoning", "token_count_tool_framing", 0), + ("azure_ai/grok-4-1-fast-reasoning", "token_count_responses_framing", 0), + ("azure_ai/grok-4-1-fast-reasoning", "token_count_responses_tool_framing", 0), ], ) def test_token_count_setting_is_stored_only_when_the_operator_changed_it( diff --git a/tests/unit/proxy/proxy_server/test_routes_models.py b/tests/unit/proxy/proxy_server/test_routes_models.py index 874f471f3e7..ff74cd3fa1d 100644 --- a/tests/unit/proxy/proxy_server/test_routes_models.py +++ b/tests/unit/proxy/proxy_server/test_routes_models.py @@ -58,6 +58,7 @@ def patched_models(monkeypatch): router.get_routable_upstream_model = MagicMock(return_value="gpt-4") router.get_configured_display_name = MagicMock(return_value=None) router.get_configured_service_tiers = MagicMock(return_value=()) + router.get_configured_max_input_tokens = MagicMock(return_value=()) router.get_model_group_info = MagicMock(return_value=None) router.get_search_tool_support = MagicMock(return_value=()) route_entry_points_to_stubs(router) diff --git a/tests/unit/proxy/proxy_server/test_streaming_helpers.py b/tests/unit/proxy/proxy_server/test_streaming_helpers.py index 6e081d2c73c..fdf58ec2fd7 100644 --- a/tests/unit/proxy/proxy_server/test_streaming_helpers.py +++ b/tests/unit/proxy/proxy_server/test_streaming_helpers.py @@ -71,6 +71,8 @@ from litellm.types.utils import ( Function, ModelResponseStream, StreamingChoices, + TextChoices, + TextCompletionResponse, Usage, ) from tests._support.clock_loops import AutojumpClockLoop @@ -909,6 +911,102 @@ def test_serialize_streaming_chunk_simple_payload(): } +@pytest.mark.parametrize( + ("choice_extras", "serialized_extras"), + [ + ({"stop_reason": ",5", "token_ids": [20]}, {"stop_reason": ",5", "token_ids": [20]}), + ({"finish_reason": "stop", "stop_reason": None}, {"finish_reason": "stop", "stop_reason": None}), + ], + ids=["vllm-fields", "natural-stop"], +) +def test_serialize_streaming_chunk_keeps_the_vllm_choice_fields( + choice_extras: dict[str, object], serialized_extras: dict[str, object] +): + chunk: Final = ModelResponseStream( + id="chatcmpl-test", + choices=[StreamingChoices(index=0, delta=Delta(content="1,2,3,4", role="assistant"), **choice_extras)], + created=0, + model="gpt-4", + object="chat.completion.chunk", + ) + + assert json.loads(_serialize_streaming_chunk(chunk))["choices"] == [ + {"index": 0, "delta": {"role": "assistant", "content": "1,2,3,4"}, **serialized_extras} + ] + + +def _text_completion_chunk(choice_extras: Mapping[str, object]) -> TextCompletionResponse: + return TextCompletionResponse( + id="cmpl-1", + choices=[TextChoices(finish_reason="length", index=0, text="", **choice_extras)], + created=1, + model="Qwen/Qwen3-0.6B", + ) + + +def _chat_chunk_with_fingerprint(choice_extras: Mapping[str, object]) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-1", + choices=[StreamingChoices(finish_reason="length", index=0, delta=Delta(content=None), **choice_extras)], + created=1, + model="Qwen/Qwen3-0.6B", + object="chat.completion.chunk", + system_fingerprint="fp_1", + ) + + +_VLLM_SHAPED_CHUNKS: Final = pytest.mark.parametrize( + "build_chunk", [_text_completion_chunk, _chat_chunk_with_fingerprint], ids=["text-completion", "chat-fingerprint"] +) + + +@_VLLM_SHAPED_CHUNKS +@pytest.mark.parametrize( + "vllm_choice_fields", + [{"stop_reason": None}, {"stop_reason": None, "token_ids": [21]}, {"stop_reason": None, "token_ids": None}], + ids=["null-stop-reason", "null-stop-reason-with-token-ids", "null-stop-reason-and-token-ids"], +) +def test_serialize_streaming_chunk_keeps_null_vllm_choice_fields( + build_chunk: Callable[[Mapping[str, object]], TextCompletionResponse | ModelResponseStream], + vllm_choice_fields: dict[str, object], +) -> None: + chunk: Final = build_chunk(vllm_choice_fields) + plain_chunk: Final = json.loads(build_chunk({}).model_dump_json(exclude_none=True, exclude_unset=True)) + + assert json.loads(_serialize_streaming_chunk(chunk)) == { + **plain_chunk, + "choices": [{**plain_chunk["choices"][0], **vllm_choice_fields}], + } + + +@_VLLM_SHAPED_CHUNKS +@pytest.mark.parametrize("choice_extras", [{}, {"stop_reason": ",5", "token_ids": [20]}], ids=["no-vllm", "set-vllm"]) +def test_serialize_streaming_chunk_without_null_vllm_fields_matches_model_dump_json( + build_chunk: Callable[[Mapping[str, object]], TextCompletionResponse | ModelResponseStream], + choice_extras: dict[str, object], +) -> None: + chunk: Final = build_chunk(choice_extras) + + assert _serialize_streaming_chunk(chunk) == chunk.model_dump_json(exclude_none=True, exclude_unset=True).encode() + + +def test_serialize_streaming_chunk_keeps_streaming_a_null_stop_reason_chunk_with_a_value_json_cannot_encode() -> None: + chunk: Final = ModelResponseStream( + id="chatcmpl-vllm", + choices=[StreamingChoices(index=0, delta=Delta(content="a"), stop_reason=None, token_ids=[21])], + created=0, + model="Qwen/Qwen3-0.6B", + object="chat.completion.chunk", + ) + setattr(chunk, "trace_chunk", hashlib.sha256(b"trace").digest()) + + payload: Final = json.loads(_serialize_streaming_chunk(chunk)) + + assert payload["choices"][0]["delta"] == {"content": "a"} + assert payload["choices"][0]["token_ids"] == [21] + assert "trace_chunk" not in payload + + def test_serialize_streaming_chunk_invalid_input_raises_attribute_error(): """The helper is typed as ``BaseModel`` — handing it a plain dict trips the attribute-access path (no ``model_dump_json``).""" diff --git a/tests/unit/proxy/router_entry_point_stubs.py b/tests/unit/proxy/router_entry_point_stubs.py index 6912049d225..51b5bf13f16 100644 --- a/tests/unit/proxy/router_entry_point_stubs.py +++ b/tests/unit/proxy/router_entry_point_stubs.py @@ -123,6 +123,8 @@ def route_entry_points_to_stubs(router: _RouterDouble) -> _RouterDouble: return router.get_configured_display_name(name) if field == "service_tiers": return router.get_configured_service_tiers(name, **_given(team_id=team_id)) + if field == "max_input_tokens": + return router.get_configured_max_input_tokens(name, **_given(team_id=team_id)) if field == "listing_info": return router.get_model_listing_info(name) if field == "mode": diff --git a/tests/unit/proxy/test_model_dump_with_preserved_fields.py b/tests/unit/proxy/test_model_dump_with_preserved_fields.py index 31531dccd16..92c7e58126b 100644 --- a/tests/unit/proxy/test_model_dump_with_preserved_fields.py +++ b/tests/unit/proxy/test_model_dump_with_preserved_fields.py @@ -545,3 +545,31 @@ def test_provider_refusal_string_is_preserved_on_message() -> None: "object": "chat.completion", } assert model_response.choices[0].message.refusal == "I can't help with that." + + +def test_vllm_null_stop_reason_is_kept_only_on_the_choice_that_carries_it() -> None: + model_response: Final = convert_to_model_response_object( + response_object={ + "id": "chatcmpl-vllm-natural-stop", + "created": 4, + "model": "Qwen/Qwen3-0.6B", + "object": "chat.completion", + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": {"role": "assistant", "content": "1, 2"}, + "stop_reason": None, + }, + {"finish_reason": "stop", "index": 1, "message": {"role": "assistant", "content": "1, 2"}}, + ], + }, + model_response_object=ModelResponse(), + ) + + dumped_choices: Final = model_dump_with_preserved_fields(model_response, exclude_unset=True)["choices"] + + assert [(choice["index"], "stop_reason" in choice, choice.get("stop_reason")) for choice in dumped_choices] == [ + (0, True, None), + (1, False, None), + ] diff --git a/tests/unit/proxy/test_proxy_logging_hook_detection.py b/tests/unit/proxy/test_proxy_logging_hook_detection.py index 182ef833069..afcb89010fe 100644 --- a/tests/unit/proxy/test_proxy_logging_hook_detection.py +++ b/tests/unit/proxy/test_proxy_logging_hook_detection.py @@ -347,6 +347,241 @@ async def test_post_call_stream_guardrail_keeps_own_iterator_on_chat_completions assert delivered_text != "" +async def _served_token_ids(monkeypatch, guardrail_default_on: bool) -> tuple[object, ...]: + from litellm.caching.caching import DualCache + from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ( + ContentFilterGuardrail, + ) + from litellm.types.guardrails import BlockedWord, ContentFilterAction + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + guardrail = ContentFilterGuardrail( + guardrail_name="output-filter", + blocked_words=[BlockedWord(keyword="zebra", action=ContentFilterAction("MASK"))], + event_hook="post_call", + default_on=guardrail_default_on, + ) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + + def chunk(content: str, token_ids: list[int], finish_reason: str | None = None) -> ModelResponseStream: + streamed = ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta(content=content), finish_reason=finish_reason)] + ) + setattr(streamed.choices[0], "token_ids", token_ids) + return streamed + + async def fake_stream(): + yield chunk("the zebra runs", [1, 2, 3]) + yield chunk("", [4], finish_reason="stop") + + served = [ + served_chunk + async for served_chunk in ProxyLogging(user_api_key_cache=DualCache()).async_post_call_streaming_iterator_hook( + response=fake_stream(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876", request_route="/chat/completions"), + request_data={"model": "gpt-4o-mini", "metadata": {}}, + ) + ] + return tuple(getattr(served_chunk.choices[0], "token_ids", None) for served_chunk in served) + + +@pytest.mark.asyncio +async def test_post_call_stream_guardrail_drops_token_ids_that_decode_to_the_unmasked_text(monkeypatch): + served_token_ids = await _served_token_ids(monkeypatch, guardrail_default_on=True) + + assert served_token_ids + assert all(token_ids is None for token_ids in served_token_ids) + + +@pytest.mark.asyncio +async def test_post_call_stream_keeps_token_ids_when_no_guardrail_runs_on_the_request(monkeypatch): + served_token_ids = await _served_token_ids(monkeypatch, guardrail_default_on=False) + + assert served_token_ids == ([1, 2, 3], [4]) + + +class _ChecksEachChunkText(CustomGuardrail): + async def async_post_call_streaming_hook(self, user_api_key_dict, response): + return None + + +async def _per_chunk_hook_token_ids(monkeypatch, guardrail_default_on: bool) -> tuple[object, bool]: + import time + + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + monkeypatch.setattr( + litellm, + "callbacks", + [_ChecksEachChunkText(guardrail_name="chunk-check", event_hook="post_call", default_on=guardrail_default_on)], + ) + logging_obj = Logging( + model="m", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="acompletion", + start_time=time.time(), + litellm_call_id="per-chunk-call", + function_id="per-chunk-call", + ) + held_back_text_chunk = ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content=""))]) + setattr(held_back_text_chunk.choices[0], "token_ids", [220, 19, 11]) + + served = await ProxyLogging(user_api_key_cache=DualCache()).async_post_call_streaming_hook( + data={"model": "m", "metadata": {}, "litellm_logging_obj": logging_obj}, + response=held_back_text_chunk, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876", request_route="/chat/completions"), + ) + return getattr(served.choices[0], "token_ids", None), logging_obj.drops_stream_token_ids + + +@pytest.mark.asyncio +async def test_per_chunk_stream_check_drops_token_ids_served_ahead_of_the_text_it_checks(monkeypatch): + assert await _per_chunk_hook_token_ids(monkeypatch, guardrail_default_on=True) == (None, True) + + +@pytest.mark.asyncio +async def test_per_chunk_stream_keeps_token_ids_when_its_check_does_not_run_on_the_request(monkeypatch): + assert await _per_chunk_hook_token_ids(monkeypatch, guardrail_default_on=False) == ([220, 19, 11], False) + + +class _ScansWholeStreamFirst(CustomGuardrail): + async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data): + for chunk in [chunk async for chunk in response]: + yield chunk + + +async def _logged_stream_token_ids(monkeypatch, guardrail_default_on: bool) -> object: + import asyncio + import json + import time + + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + from litellm.llms.hosted_vllm.chat.transformation import HostedVLLMChatConfig + + guardrail = _ScansWholeStreamFirst( + guardrail_name="whole-stream-scan", event_hook="post_call", default_on=guardrail_default_on + ) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + logging_obj = Logging( + model="m", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="acompletion", + start_time=time.time(), + litellm_call_id="token-ids-call", + function_id="token-ids-call", + ) + logged = asyncio.get_running_loop().create_future() + + async def capture_logged_response(complete_streaming_response, cache_hit): + logged.set_result(complete_streaming_response) + + logging_obj._on_deferred_stream_complete = capture_logged_response + + async def sse_lines(): + for choice in ( + {"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}, + {"index": 0, "delta": {"content": ""}, "finish_reason": None, "token_ids": [16]}, + {"index": 0, "delta": {"content": "1,2"}, "finish_reason": None, "token_ids": [11, 17]}, + {"index": 0, "delta": {"content": ""}, "finish_reason": "stop", "token_ids": [20]}, + ): + yield "data: " + json.dumps( + {"id": "chatcmpl-1", "object": "chat.completion.chunk", "model": "m", "choices": [choice]} + ) + yield "data: [DONE]" + + stream = CustomStreamWrapper( + completion_stream=HostedVLLMChatConfig().get_model_response_iterator( + streaming_response=sse_lines(), sync_stream=False, json_mode=False + ), + model="m", + custom_llm_provider="hosted_vllm", + logging_obj=logging_obj, + ) + async for _ in ProxyLogging(user_api_key_cache=DualCache()).async_post_call_streaming_iterator_hook( + response=stream, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876", request_route="/chat/completions"), + request_data={"model": "m", "metadata": {}, "litellm_logging_obj": logging_obj}, + ): + pass + return (await asyncio.wait_for(logged, timeout=5)).choices[0].token_ids + + +@pytest.mark.asyncio +async def test_post_call_stream_guardrail_drops_token_ids_from_the_response_it_logs_and_caches(monkeypatch): + assert await _logged_stream_token_ids(monkeypatch, guardrail_default_on=True) is None + + +@pytest.mark.asyncio +async def test_post_call_stream_logs_every_token_id_when_no_guardrail_runs_on_the_request(monkeypatch): + assert await _logged_stream_token_ids(monkeypatch, guardrail_default_on=False) == [16, 11, 17, 20] + + +def _completion_chunk(text: str, token_ids: list[int], finish_reason: str | None = None) -> object: + from litellm.types.utils import TextChoices, TextCompletionResponse + + return TextCompletionResponse( + object="text_completion", + choices=[TextChoices(index=0, text=text, finish_reason=finish_reason, token_ids=token_ids)], + ) + + +def _guarded_completions_stream(monkeypatch, action: str, held_back: tuple[object, ...]): + from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ( + ContentFilterGuardrail, + ) + from litellm.types.guardrails import BlockedWord, ContentFilterAction + + monkeypatch.setattr( + litellm, + "callbacks", + [ + ContentFilterGuardrail( + guardrail_name="output-filter", + blocked_words=[BlockedWord(keyword="7", action=ContentFilterAction(action))], + event_hook="post_call", + default_on=True, + ) + ], + ) + + async def upstream(): + for chunk in held_back: + yield chunk + yield _completion_chunk(", 6, 7", [21, 11, 22], finish_reason="stop") + + return ProxyLogging(user_api_key_cache=DualCache()).async_post_call_streaming_iterator_hook( + response=upstream(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876", request_route="/v1/completions"), + request_data={"model": "qwen3", "metadata": {}}, + ) + + +@pytest.mark.asyncio +async def test_post_call_stream_guardrail_blocks_completions_before_serving_token_id_only_chunks(monkeypatch): + stream = _guarded_completions_stream( + monkeypatch, "BLOCK", (_completion_chunk("", [11]), _completion_chunk("", [220])) + ) + + with pytest.raises(HTTPException): + await stream.__anext__() + + +@pytest.mark.asyncio +async def test_post_call_stream_guardrail_serves_no_chunk_left_empty_by_the_token_id_drop(monkeypatch): + held_back = (_completion_chunk("", [11]), _completion_chunk("", [220])) + + served = [chunk async for chunk in _guarded_completions_stream(monkeypatch, "MASK", held_back)] + + assert [(chunk.choices[0].text, chunk.choices[0].token_ids) for chunk in served] == [ + (", 6, [KEYWORD_REDACTED]", None) + ] + assert [chunk.choices[0].token_ids for chunk in held_back] == [None, None] + + @pytest.mark.asyncio async def test_post_call_stream_records_masked_text_for_deferred_logging(monkeypatch): from litellm.caching.caching import DualCache diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 3647146eb4e..76d94628d81 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -13900,7 +13900,7 @@ def _stream_usage_generator_chunks(): def test_is_injected_stream_usage_artifact(): from litellm.litellm_core_utils.stream_usage_chunks import is_injected_stream_usage_artifact - from litellm.types.utils import ModelResponseStream, Usage + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices, Usage content_chunk, finish_chunk, empty_choices_usage_chunk = _stream_usage_test_chunks() assert is_injected_stream_usage_artifact(empty_choices_usage_chunk) is True @@ -13923,6 +13923,21 @@ def test_is_injected_stream_usage_artifact(): assert is_injected_stream_usage_artifact({"usage": {"prompt_tokens": 1}}) is False + held_back_token_ids_chunk: Final = ModelResponseStream( + model="qwen3", choices=[StreamingChoices(delta=Delta(content=None), token_ids=[16])] + ) + assert is_injected_stream_usage_artifact(held_back_token_ids_chunk) is False + + null_token_ids_chunk: Final = ModelResponseStream( + model="qwen3", choices=[StreamingChoices(delta=Delta(content=None), token_ids=None)] + ) + assert is_injected_stream_usage_artifact(null_token_ids_chunk) is True + + unrelated_extra_chunk: Final = ModelResponseStream( + model="qwen3", choices=[StreamingChoices(delta=Delta(content=None), text="")] + ) + assert is_injected_stream_usage_artifact(unrelated_extra_chunk) is True + async def _collect_async_data_generator_frames(request_data: dict) -> list: from litellm.proxy.proxy_server import async_data_generator diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 608efe20c50..2f85d1c7516 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -1925,6 +1925,23 @@ async def test_uncountable_model_without_framing_raises(): assert raised.value.code == "404" +@pytest.mark.parametrize( + "settings", + [ + TokenCountSettings(framing=None, count_model=None, tool_framing=300), + TokenCountSettings(framing=None, count_model=None, responses_framing=145), + TokenCountSettings(framing=None, count_model=None, responses_tool_framing=100), + ], +) +@pytest.mark.asyncio +async def test_uncountable_model_with_only_a_request_kind_framing_falls_through_to_the_local_estimate( + settings: TokenCountSettings, +): + counter: Final = _ScriptedCounter({"us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", status_code=404)}) + + assert isinstance(await _ladder(counter, settings), ProviderCountFailure) + + @pytest.mark.parametrize( ("count_model", "counted_as"), [ @@ -2770,6 +2787,480 @@ async def test_framing_comes_from_the_row_the_deployment_provider_or_base_model_ assert aliased.total_tokens == litellm.token_counter(model="gpt-5.4-mini", messages=messages) + framing +@pytest.mark.asyncio +async def test_a_tool_request_counts_its_framing_in_place_of_the_plain_one_per_api(monkeypatch: pytest.MonkeyPatch): + def deployment(name: str, model_info: dict[str, object]) -> dict[str, object]: + return { + "model_name": name, + "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": f"sk-unused-{name}"}, + "model_info": model_info, + } + + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + deployment("plain", {}), + deployment("chat", {"token_count_framing": 7, "token_count_tool_framing": 50}), + deployment( + "per-api", + { + "token_count_framing": 7, + "token_count_tool_framing": 50, + "token_count_responses_framing": 5, + "token_count_responses_tool_framing": 2, + }, + ), + ] + ), + ) + messages: Final = [{"role": "user", "content": "What is the weather in Paris?"}] + tools: Final = [_KIMI_WEATHER] + + async def count(model: str, request_tools: list[dict[str, object]] | None, responses: bool) -> int: + response: Final = await count_request_tokens( + request=TokenCountRequest(model=model, messages=messages, tools=request_tools), + responses_request={"model": model, "input": "What is the weather in Paris?"} if responses else None, + ) + return response.total_tokens + + plain_with_tools: Final = await count("plain", tools, responses=False) + plain_without_tools: Final = await count("plain", None, responses=False) + + assert await count("plain", tools, responses=True) == plain_with_tools + assert await count("chat", tools, responses=False) == plain_with_tools + 50 + assert await count("chat", None, responses=False) == plain_without_tools + 7 + assert await count("chat", [], responses=False) == plain_without_tools + 7 + assert await count("chat", tools, responses=True) == plain_with_tools + 50 + assert await count("chat", None, responses=True) == plain_without_tools + 7 + assert await count("per-api", tools, responses=False) == plain_with_tools + 50 + assert await count("per-api", tools, responses=True) == plain_with_tools + 2 + assert await count("per-api", None, responses=True) == plain_without_tools + 5 + + +_GET_WEATHER: Final = { + "name": "get_weather", + "description": "Get the current weather for a city", + "input_schema": { + "type": "object", + "properties": {"city": {"type": "string", "description": "City name"}}, + "required": ["city"], + }, +} +_THREE_TOOLS: Final = ( + _GET_WEATHER, + { + "name": "get_time", + "description": "Get the local time in a timezone", + "input_schema": { + "type": "object", + "properties": {"timezone": {"type": "string", "description": "IANA timezone"}}, + "required": ["timezone"], + }, + }, + { + "name": "search_docs", + "description": "Search internal documents and return the top matches", + "input_schema": { + "type": "object", + "properties": {"query": {"type": "string"}, "limit": {"type": "integer", "description": "Max results"}}, + "required": ["query"], + }, + }, +) +_FORM_FIELDS: Final = tuple(f"field_{index}" for index in range(10)) +_FILL_FORM: Final = { + "name": "fill_form", + "description": "Fill a form", + "input_schema": { + "type": "object", + "properties": { + field: {"type": "string", "description": f"Value for field number {index}"} + for index, field in enumerate(_FORM_FIELDS) + }, + "required": list(_FORM_FIELDS), + }, +} +_WEATHER_SYSTEM: Final = "You are a helpful assistant that answers weather questions briefly." +_AZURE_GROK_4_1: Final = "azure_ai/grok-4-1-fast-reasoning" +_FIREWORKS_DEEPSEEK_V4_1: Final = "fireworks_ai/accounts/fireworks/models/deepseek-v4p1-flash" +_XAI_GROK_4_3: Final = "xai/grok-4.3" +_NATIVE_RESPONSES_BASE: Final = { + _AZURE_GROK_4_1: "https://unused.services.ai.azure.com/models", + _FIREWORKS_DEEPSEEK_V4_1: "https://unused.invalid", + _XAI_GROK_4_3: "https://unused.invalid", +} + + +def _responses_function(tool: Mapping[str, object], **strict: bool | None) -> dict[str, object]: + return { + "type": "function", + "name": tool["name"], + "description": tool["description"], + "parameters": tool["input_schema"], + **strict, + } + + +def _get_weather_copies(count: int) -> tuple[dict[str, object], ...]: + return tuple({**_GET_WEATHER, "name": f"get_weather_{index}"} for index in range(count)) + + +# usage.input_tokens plus cache_read_input_tokens that POST /v1/messages billed for these requests with the user turn +# "What is the weather in Paris?" through Azure AI Foundry grok-4-1-fast-reasoning, Fireworks deepseek-v4p1-flash and +# xAI grok-4.3 on 2026-10-29, and usage.prompt_tokens that POST /v1/chat/completions billed through DeepSeek's API on +# 2026-10-30 +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("model", "system", "tools", "billed_prompt_tokens"), + [ + (_AZURE_GROK_4_1, None, (), 163), + (_AZURE_GROK_4_1, None, (_GET_WEATHER,), 357), + (_AZURE_GROK_4_1, None, _THREE_TOOLS, 445), + (_AZURE_GROK_4_1, None, (_FILL_FORM,), 550), + (_AZURE_GROK_4_1, _WEATHER_SYSTEM, (_GET_WEATHER,), 368), + (_FIREWORKS_DEEPSEEK_V4_1, None, (), 37), + (_FIREWORKS_DEEPSEEK_V4_1, None, (_GET_WEATHER,), 304), + (_FIREWORKS_DEEPSEEK_V4_1, None, _THREE_TOOLS, 430), + (_FIREWORKS_DEEPSEEK_V4_1, None, (_FILL_FORM,), 550), + (_FIREWORKS_DEEPSEEK_V4_1, _WEATHER_SYSTEM, (_GET_WEATHER,), 315), + (_XAI_GROK_4_3, None, (), 199), + (_XAI_GROK_4_3, None, (_GET_WEATHER,), 283), + (_XAI_GROK_4_3, None, _THREE_TOOLS, 402), + (_XAI_GROK_4_3, None, (_FILL_FORM,), 530), + (_XAI_GROK_4_3, _WEATHER_SYSTEM, (_GET_WEATHER,), 297), + ("deepseek/deepseek-v4-flash", None, (), 37), + ("deepseek/deepseek-v4-flash", None, (_GET_WEATHER,), 304), + ("deepseek/deepseek-v4-flash", None, _THREE_TOOLS, 430), + ("deepseek/deepseek-v4-flash", _WEATHER_SYSTEM, (_GET_WEATHER,), 315), + ("deepseek/deepseek-v4-pro", None, (), 90), + ("deepseek/deepseek-v4-pro", None, (_GET_WEATHER,), 367), + ("deepseek/deepseek-v4-pro", None, _THREE_TOOLS, 493), + ("deepseek/deepseek-v4-pro", _WEATHER_SYSTEM, (_GET_WEATHER,), 378), + ], +) +async def test_local_estimate_matches_the_billed_prompt_with_the_shipped_framing( + monkeypatch: pytest.MonkeyPatch, + model: str, + system: str | None, + tools: tuple[dict[str, object], ...], + billed_prompt_tokens: int, +): + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "billed", + "litellm_params": {"model": model, "api_key": "unused", "api_base": "https://unused.invalid"}, + } + ] + ), + ) + + counted: Final = await token_counter( + request=TokenCountRequest( + model="billed", + messages=[{"role": "user", "content": "What is the weather in Paris?"}], + tools=list(tools) or None, + system=system, + ) + ) + + assert counted.count_source == "local_estimate" + assert abs(counted.total_tokens - billed_prompt_tokens) <= billed_prompt_tokens * 0.05 + + +# usage.input_tokens that POST /v1/responses billed for these requests with the input "What is the weather in Paris?" +# through Azure AI Foundry grok-4-1-fast-reasoning, Fireworks deepseek-v4p1-flash and xAI grok-4.3 on 2026-10-29 +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("model", "system", "tools", "billed_prompt_tokens"), + [ + (_AZURE_GROK_4_1, None, (), 159), + (_AZURE_GROK_4_1, None, (_GET_WEATHER,), 159), + (_AZURE_GROK_4_1, None, _THREE_TOOLS, 257), + (_AZURE_GROK_4_1, None, (_FILL_FORM,), 352), + (_AZURE_GROK_4_1, _WEATHER_SYSTEM, (_GET_WEATHER,), 170), + (_FIREWORKS_DEEPSEEK_V4_1, None, (_GET_WEATHER,), 304), + (_XAI_GROK_4_3, None, (_GET_WEATHER,), 283), + ], +) +async def test_responses_estimate_matches_the_billed_responses_prompt_with_the_shipped_framing( + monkeypatch: pytest.MonkeyPatch, + model: str, + system: str | None, + tools: tuple[dict[str, object], ...], + billed_prompt_tokens: int, +): + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "billed", + "litellm_params": {"model": model, "api_key": "unused", "api_base": _NATIVE_RESPONSES_BASE[model]}, + } + ] + ), + ) + + counted: Final = await count_request_tokens( + request=TokenCountRequest( + model="billed", + messages=[{"role": "user", "content": "What is the weather in Paris?"}], + tools=[_responses_function(tool) for tool in tools] or None, + system=system, + ), + responses_request={"model": "billed", "input": "What is the weather in Paris?"}, + ) + + assert counted.count_source == "local_estimate" + assert abs(counted.total_tokens - billed_prompt_tokens) <= billed_prompt_tokens * 0.05 + + +def _single_deployment_router(litellm_params: dict[str, object], model_info: dict[str, object]) -> Router: + return Router( + model_list=[ + { + "model_name": "counted", + "litellm_params": {"api_key": "unused", **litellm_params}, + "model_info": model_info, + } + ] + ) + + +async def _count_with_router( + monkeypatch: pytest.MonkeyPatch, router: Router, tools: tuple[dict[str, object], ...], responses: bool +) -> int: + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) + counted: Final = await count_request_tokens( + request=TokenCountRequest( + model="counted", + messages=[{"role": "user", "content": "What is the weather in Paris?"}], + tools=list(tools) or None, + ), + responses_request={"model": "counted", "input": "What is the weather in Paris?"} if responses else None, + ) + return counted.total_tokens + + +@pytest.mark.asyncio +@pytest.mark.parametrize("tools", [(), (_GET_WEATHER,), _THREE_TOOLS, (_FILL_FORM,)]) +@pytest.mark.parametrize("responses", [False, True]) +async def test_a_custom_deployment_name_counts_like_the_base_model_it_names( + monkeypatch: pytest.MonkeyPatch, tools: tuple[dict[str, object], ...], responses: bool +): + api_base: Final = _NATIVE_RESPONSES_BASE[_AZURE_GROK_4_1] + custom_name: Final = _single_deployment_router( + {"model": "azure_ai/prod-grok", "api_base": api_base}, {"base_model": _AZURE_GROK_4_1} + ) + row_name: Final = _single_deployment_router({"model": _AZURE_GROK_4_1, "api_base": api_base}, {}) + + assert await _count_with_router(monkeypatch, custom_name, tools, responses) == await _count_with_router( + monkeypatch, row_name, tools, responses + ) + + +# usage.input_tokens that POST /v1/responses billed through Azure AI Foundry grok-4-1-fast-reasoning on 2026-10-29 for +# the input "What is the weather in Paris?" with these function tools, strict left unset or sent as false or null, +# the same on the .services.ai.azure.com and .openai.azure.com hosts +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("tools", "strict", "billed_prompt_tokens"), + [ + ((_GET_WEATHER,), {}, 159), + ((_GET_WEATHER,), {"strict": False}, 155), + ((_GET_WEATHER,), {"strict": None}, 155), + (_THREE_TOOLS, {}, 257), + (_get_weather_copies(5), {}, 345), + (_get_weather_copies(10), {}, 575), + (_get_weather_copies(10), {"strict": False}, 535), + (_get_weather_copies(10), {"strict": None}, 535), + ], +) +@pytest.mark.parametrize( + ("litellm_params", "model_info"), + [ + ({"model": _AZURE_GROK_4_1, "api_base": _NATIVE_RESPONSES_BASE[_AZURE_GROK_4_1]}, {}), + ( + {"model": "azure/grok-4-1-fast-reasoning", "api_base": "https://unused.openai.azure.com/openai/v1"}, + {"base_model": _AZURE_GROK_4_1}, + ), + ], +) +async def test_azure_responses_estimate_stays_at_or_just_over_the_billed_prompt_as_function_tools_are_added( + monkeypatch: pytest.MonkeyPatch, + tools: tuple[dict[str, object], ...], + strict: dict[str, bool | None], + billed_prompt_tokens: int, + litellm_params: dict[str, object], + model_info: dict[str, object], +): + router: Final = _single_deployment_router(litellm_params, model_info) + + counted: Final = await _count_with_router( + monkeypatch, router, tuple(_responses_function(tool, **strict) for tool in tools), responses=True + ) + + assert billed_prompt_tokens <= counted <= billed_prompt_tokens * 1.02 + + +# POST /v1/responses billed the same usage.input_tokens with strict unset and false on 2026-10-29, 283 for one tool +# through xAI grok-4.3 +@pytest.mark.asyncio +async def test_a_responses_host_that_bills_an_unset_strict_like_false_counts_them_alike( + monkeypatch: pytest.MonkeyPatch, +): + router: Final = _single_deployment_router( + {"model": _XAI_GROK_4_3, "api_base": _NATIVE_RESPONSES_BASE[_XAI_GROK_4_3]}, {} + ) + + unset: Final = await _count_with_router(monkeypatch, router, (_responses_function(_GET_WEATHER),), responses=True) + + assert unset == await _count_with_router( + monkeypatch, router, (_responses_function(_GET_WEATHER, strict=False),), responses=True + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("built_in_tool", "responses"), + [({"type": "web_search"}, True), ({"type": "web_search_20250305", "name": "web_search"}, False)], +) +async def test_a_request_whose_only_tools_are_built_in_counts_like_one_without_tools( + monkeypatch: pytest.MonkeyPatch, built_in_tool: dict[str, object], responses: bool +): + router: Final = _single_deployment_router( + {"model": _AZURE_GROK_4_1, "api_base": _NATIVE_RESPONSES_BASE[_AZURE_GROK_4_1]}, {} + ) + + assert await _count_with_router(monkeypatch, router, (built_in_tool,), responses) == await _count_with_router( + monkeypatch, router, (), responses + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("tools", [(), (_GET_WEATHER,), _THREE_TOOLS]) +async def test_a_responses_count_the_proxy_bridges_to_chat_completions_frames_like_the_chat_count( + monkeypatch: pytest.MonkeyPatch, tools: tuple[dict[str, object], ...] +): + bridged: Final = _single_deployment_router( + {"model": _AZURE_GROK_4_1, "api_base": "https://unused.cognitiveservices.azure.com/models"}, {} + ) + native: Final = _single_deployment_router( + {"model": _AZURE_GROK_4_1, "api_base": _NATIVE_RESPONSES_BASE[_AZURE_GROK_4_1]}, {} + ) + chat_count: Final = await _count_with_router(monkeypatch, bridged, tools, responses=False) + + assert await _count_with_router(monkeypatch, bridged, tools, responses=True) == chat_count + assert await _count_with_router(monkeypatch, native, tools, responses=True) != chat_count + + +@pytest.mark.asyncio +@pytest.mark.parametrize("from_router_defaults", [False, True]) +async def test_a_responses_count_that_asks_for_the_chat_completions_bridge_frames_like_the_chat_count( + monkeypatch: pytest.MonkeyPatch, from_router_defaults: bool +): + router: Final = Router( + model_list=[ + { + "model_name": "counted", + "litellm_params": { + "api_key": "unused", + "model": _AZURE_GROK_4_1, + "api_base": _NATIVE_RESPONSES_BASE[_AZURE_GROK_4_1], + }, + "model_info": {}, + } + ], + default_litellm_params={"use_chat_completions_api": True} if from_router_defaults else None, + ) + chat_count: Final = await _count_with_router(monkeypatch, router, (_GET_WEATHER,), responses=False) + + counted: Final = await count_request_tokens( + request=TokenCountRequest( + model="counted", + messages=[{"role": "user", "content": "What is the weather in Paris?"}], + tools=[_GET_WEATHER], + ), + responses_request={ + "model": "counted", + "input": "What is the weather in Paris?", + **({} if from_router_defaults else {"use_chat_completions_api": True}), + }, + ) + + assert counted.total_tokens == chat_count + + +@pytest.mark.asyncio +async def test_a_responses_count_frames_for_the_api_base_router_defaults_send_inference_to( + monkeypatch: pytest.MonkeyPatch, +): + router: Final = Router( + model_list=[ + { + "model_name": "counted", + "litellm_params": { + "api_key": "unused", + "model": _AZURE_GROK_4_1, + "api_base": _NATIVE_RESPONSES_BASE[_AZURE_GROK_4_1], + }, + "model_info": {}, + } + ], + default_litellm_params={"api_base": "https://unused.cognitiveservices.azure.com/models"}, + ) + chat_count: Final = await _count_with_router(monkeypatch, router, (_GET_WEATHER,), responses=False) + + assert await _count_with_router(monkeypatch, router, (_GET_WEATHER,), responses=True) == chat_count + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_chat_completions_api", ["maybe", {"a": 1}]) +async def test_a_responses_count_with_a_malformed_bridge_flag_reads_it_the_way_the_responses_route_does( + monkeypatch: pytest.MonkeyPatch, use_chat_completions_api: object +): + router: Final = _single_deployment_router( + {"model": _AZURE_GROK_4_1, "api_base": _NATIVE_RESPONSES_BASE[_AZURE_GROK_4_1]}, {} + ) + chat_count: Final = await _count_with_router(monkeypatch, router, (_GET_WEATHER,), responses=False) + + counted: Final = await count_request_tokens( + request=TokenCountRequest( + model="counted", + messages=[{"role": "user", "content": "What is the weather in Paris?"}], + tools=[_GET_WEATHER], + ), + responses_request={ + "model": "counted", + "input": "What is the weather in Paris?", + "use_chat_completions_api": use_chat_completions_api, + }, + ) + + assert counted.total_tokens == chat_count + + +@pytest.mark.asyncio +async def test_an_azure_responses_count_leaves_strict_unset_for_templates_not_measured_there( + monkeypatch: pytest.MonkeyPatch, +): + router: Final = _single_deployment_router( + {"model": "azure_ai/deepseek-v4.1-flash", "api_base": "https://unused.services.ai.azure.com/models"}, {} + ) + + async def counted(**strict: bool) -> int: + return await _count_with_router(monkeypatch, router, (_responses_function(_GET_WEATHER, **strict),), True) + + assert await counted() < await counted(strict=True) + + _KIMI_WEATHER: Final = { "type": "function", "function": { diff --git a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py index b4cd606374a..d5619515491 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -7192,6 +7192,29 @@ async def test_bridge_partial_response_on_disconnect_bills_only_the_input_since_ ) == await _input_tokens_billed_on_disconnect(since_compaction) +def test_a_hosted_mcp_tool_sent_to_responses_on_bedrock_converse_answers_400(monkeypatch) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(litellm.UnsupportedParamsError) as raised: + litellm.responses( + model="bedrock/converse/us.openai.gpt-oss-120b-1:0", + input="What does the deepwiki MCP server say about litellm?", + tools=[ + { + "type": "mcp", + "server_label": "deepwiki", + "server_url": "https://mcp.deepwiki.com/mcp", + "require_approval": "never", + } + ], + aws_access_key_id="AKIAHOSTEDMCP", + aws_secret_access_key="hosted-mcp-secret", + aws_region_name="us-east-1", + ) + assert raised.value.status_code == 400 + assert "`deepwiki`" in str(raised.value) + assert "litellm_proxy/mcp/" in str(raised.value) + + _ANTHROPIC_MODEL: Final = "anthropic/claude-sonnet-5-5" _WEATHER_ASK: Final = {"role": "user", "content": "Weather in Paris?"} diff --git a/tests/unit/test_router_model_cost_isolation.py b/tests/unit/test_router_model_cost_isolation.py index 91f0b2339c7..447a4dfb81d 100644 --- a/tests/unit/test_router_model_cost_isolation.py +++ b/tests/unit/test_router_model_cost_isolation.py @@ -672,7 +672,7 @@ def test_pdf_opt_out_on_one_deployment_does_not_block_its_sibling(respx_mock: re def test_token_count_settings_stay_on_the_deployment_that_set_them() -> None: - """A deployment's ``token_count_model`` and ``token_count_framing`` must not land on + """A deployment's ``token_count_model`` and framing settings must not land on the shared ``{provider}/{model}`` key, or a sibling deployment of that model would count with them.""" from litellm.proxy.common_utils.token_count_settings import cost_map_entry @@ -681,6 +681,9 @@ def test_token_count_settings_stay_on_the_deployment_that_set_them() -> None: configured: Final = { "token_count_model": {"model": "anthropic.claude-opus-5-5", "region": "us-west-2"}, "token_count_framing": 7, + "token_count_tool_framing": 3, + "token_count_responses_framing": 5, + "token_count_responses_tool_framing": 2, } builtin: Final = {field: cost_map_entry(backend_model).get(field) for field in configured} assert all(builtin[field] != value for field, value in configured.items()), builtin diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 5cf872f4ea7..61825e463f8 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -1178,6 +1178,9 @@ _COST_MAP_ROW_SCHEMA: Final = { "supports_vision": {"type": "boolean"}, "supports_web_search": {"type": "boolean"}, "token_count_framing": {"type": "integer"}, + "token_count_tool_framing": {"type": "integer"}, + "token_count_responses_framing": {"type": "integer"}, + "token_count_responses_tool_framing": {"type": "integer"}, "token_count_model": { "type": "object", "properties": {"model": {"type": "string"}, "region": {"type": "string"}}, diff --git a/tests/unit/text_completions/test_streaming_iterator.py b/tests/unit/text_completions/test_streaming_iterator.py index e5459a5a0cd..571451164d9 100644 --- a/tests/unit/text_completions/test_streaming_iterator.py +++ b/tests/unit/text_completions/test_streaming_iterator.py @@ -1,11 +1,15 @@ -from collections.abc import AsyncIterator, Iterator +from collections.abc import AsyncIterator, Iterator, Sequence +from itertools import chain from typing import Final import pytest +from openai.types.completion import Completion from litellm.exceptions import APIConnectionError +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.text_completions.streaming_iterator import TextCompletionStreamWrapper -from litellm.types.utils import Choices, Delta, ModelResponseStream, StreamingChoices, Usage +from litellm.types.utils import Choices, Delta, ModelResponseStream, StreamingChoices, TextCompletionResponse, Usage def test_reasoning_content_preserved_in_text_completion_wrapper(): @@ -85,6 +89,146 @@ async def test_text_completion_stream_usage_is_opt_in( ] +def _completion_chunk(text: str, finish_reason: str | None, **choice_extras: object) -> Completion: + return Completion.construct( + id="cmpl-1", + created=1, + model="m", + object="text_completion", + choices=[{"index": 0, "text": text, "finish_reason": finish_reason, "logprobs": None, **choice_extras}], + ) + + +def _text_completion_stream(upstream: Sequence[Completion]) -> list[TextCompletionResponse]: + logging_obj: Final = Logging( + model="m", + messages=[{"role": "user", "content": "1, 2, 3, 4"}], + stream=True, + call_type="text_completion", + start_time=0.0, + litellm_call_id="call-1", + function_id="function-1", + ) + chat_stream: Final = CustomStreamWrapper( + completion_stream=iter(upstream), + model="m", + custom_llm_provider="text-completion-openai", + logging_obj=logging_obj, + ) + return list(TextCompletionStreamWrapper(completion_stream=chat_stream, model="m")) + + +@pytest.mark.parametrize("final_text", [", 5, 6", ""], ids=["text-then-finish", "finish-only"]) +def test_text_completion_stream_keeps_stop_reason_on_the_finish_chunk(final_text: str) -> None: + responses: Final = _text_completion_stream( + (_completion_chunk(", 5", None), _completion_chunk(final_text, "stop", stop_reason=", 7")) + ) + + assert responses[-1].choices[0].model_dump(exclude_none=True) == { + "finish_reason": "stop", + "index": 0, + "stop_reason": ", 7", + } + + +def test_text_completion_stream_puts_stop_reason_on_the_final_text_chunk() -> None: + responses: Final = _text_completion_stream( + (_completion_chunk(", 5", None), _completion_chunk(", 6", "stop", stop_reason=", 7")) + ) + + assert [(response.choices[0].get("text"), response.choices[0].get("stop_reason")) for response in responses] == [ + (", 5", None), + (", 6", ", 7"), + (None, ", 7"), + ] + + +@pytest.mark.parametrize("final_text", [", 6", ""], ids=["text-then-finish", "finish-only"]) +def test_text_completion_stream_keeps_held_back_token_ids_and_the_stop_reason(final_text: str) -> None: + responses: Final = _text_completion_stream( + ( + _completion_chunk("", None, token_ids=[11]), + _completion_chunk("", None, token_ids=[220]), + _completion_chunk(", 5", None, token_ids=[20]), + _completion_chunk(final_text, "stop", stop_reason=", 7", token_ids=[11, 220, 21]), + ) + ) + choices: Final = [response.choices[0] for response in responses] + + assert list(chain.from_iterable(choice.get("token_ids") or () for choice in choices)) == [11, 220, 20, 11, 220, 21] + assert [choice.get("stop_reason") for choice in choices if choice.finish_reason is not None] == [", 7"] + + +def test_text_completion_stream_sends_empty_text_on_token_id_only_chunks() -> None: + responses: Final = _text_completion_stream( + ( + _completion_chunk("", None, token_ids=[11]), + _completion_chunk("", None, token_ids=[220]), + _completion_chunk(", 5", None, token_ids=[20]), + ) + ) + + choices: Final = [response.choices[0] for response in responses] + + assert [(choice.text, choice.get("token_ids")) for choice in choices if choice.get("token_ids")] == [ + ("", [11]), + ("", [220]), + (", 5", [20]), + ] + + +def test_text_completion_stream_keeps_each_choices_text_and_token_ids_on_its_own_index() -> None: + responses: Final = _text_completion_stream( + ( + _completion_chunk(" 5", None, index=0, token_ids=[20]), + _completion_chunk(" 7", None, index=1, token_ids=[22]), + _completion_chunk(", 6", None, index=0, token_ids=[11, 21]), + _completion_chunk("", None, index=1, token_ids=[11]), + _completion_chunk("", "stop", index=0, stop_reason=", 9", token_ids=[11, 220]), + _completion_chunk(", 8", "length", index=1, token_ids=[220, 23]), + ) + ) + choices: Final = tuple(chain.from_iterable(response.choices for response in responses)) + + assert { + index: ( + "".join(choice.text or "" for choice in choices if choice.index == index), + list(chain.from_iterable(choice.get("token_ids") or () for choice in choices if choice.index == index)), + ) + for index in (0, 1) + } == {0: (" 5, 6", [20, 11, 21, 11, 220]), 1: (" 7, 8", [22, 11, 220, 23])} + + +@pytest.mark.parametrize( + "upstream", + [ + ( + _completion_chunk(" a", None, index=0), + _completion_chunk(" b", None, index=1), + _completion_chunk("", "length", index=2), + _completion_chunk(" red", "stop", index=1, stop_reason=","), + _completion_chunk(" s", "stop", index=0, stop_reason=","), + ), + ( + _completion_chunk(" a", None, index=0), + _completion_chunk(" b", None, index=1), + _completion_chunk(" s", "stop", index=0, stop_reason=","), + _completion_chunk("", "length", index=2), + _completion_chunk(" red", "stop", index=1, stop_reason=","), + ), + ], + ids=["other-choices-finish-first", "choice-0-finishes-first"], +) +def test_text_completion_stream_finishes_each_choice_on_its_own_index(upstream: Sequence[Completion]) -> None: + finished: Final = [ + (choice.index, choice.finish_reason, choice.get("stop_reason")) + for choice in chain.from_iterable(response.choices for response in _text_completion_stream(upstream)) + if choice.finish_reason is not None + ] + + assert sorted(finished) == [(0, "stop", ","), (1, "stop", ","), (2, "length", None)] + + def test_text_completion_conversion_rejects_non_streaming_choices() -> None: chunk: Final = Choices() wrapper: Final = TextCompletionStreamWrapper(completion_stream=None, model="test-model") diff --git a/ui/litellm-dashboard/src/components/llm_calls/responses_api.test.tsx b/ui/litellm-dashboard/src/components/llm_calls/responses_api.test.tsx index f9c5e3b0181..9285a49c93f 100644 --- a/ui/litellm-dashboard/src/components/llm_calls/responses_api.test.tsx +++ b/ui/litellm-dashboard/src/components/llm_calls/responses_api.test.tsx @@ -207,6 +207,46 @@ describe("responses_api", () => { expect(onTimingData).not.toHaveBeenCalled(); }); + it.each([ + [["__all__"], [], [{ server_label: "litellm", server_url: "litellm_proxy/mcp" }]], + [ + ["server-1"], + [{ server_id: "server-1", server_name: "my server" }], + [{ server_label: "my server", server_url: "litellm_proxy/mcp/my server" }], + ], + ])( + "sends MCP selection %j as gateway references the proxy runs for any model", + async (selected, servers, expected) => { + await makeOpenAIResponsesRequest( + messages, + mockUpdateTextUI, + "gemini-pro", + "test-token", + undefined, // tags + undefined, // signal + undefined, // onReasoningContent + undefined, // onTimingData + undefined, // onUsageData + undefined, // traceId + undefined, // vector_store_ids + undefined, // guardrails + undefined, // policies + selected, + undefined, // previousResponseId + undefined, // onResponseId + undefined, // onMCPEvent + undefined, // codeInterpreterEnabled + undefined, // onCodeInterpreterResult + undefined, // customBaseUrl + servers as never, + ); + + expect(mockResponsesCreate.mock.calls[0][0].tools).toEqual( + expected.map((tool) => ({ type: "mcp", require_approval: "never", ...tool })), + ); + }, + ); + it("should report total latency in both streaming and non-streaming modes", async () => { const onTotalLatency = vi.fn(); const callWithStreaming = (streamingEnabled: boolean) => @@ -475,14 +515,14 @@ describe("responses_api", () => { { type: "mcp", server_label: "Alpha", - server_url: "https://backend.example.com/mcp/Alpha", + server_url: "litellm_proxy/mcp/Alpha", require_approval: "never", allowed_tools: ["toolA"], }, { type: "mcp", server_label: "Beta", - server_url: "https://backend.example.com/mcp/Beta", + server_url: "litellm_proxy/mcp/Beta", require_approval: "never", allowed_tools: ["toolB", "toolC"], }, diff --git a/ui/litellm-dashboard/src/components/llm_calls/responses_api.tsx b/ui/litellm-dashboard/src/components/llm_calls/responses_api.tsx index f25ac939331..ed1aeb25074 100644 --- a/ui/litellm-dashboard/src/components/llm_calls/responses_api.tsx +++ b/ui/litellm-dashboard/src/components/llm_calls/responses_api.tsx @@ -7,11 +7,12 @@ import type { } from "openai/resources/responses/responses"; import { MessageType } from "../chat_ui/types"; import { TokenUsage } from "../chat_ui/ResponseMetrics"; -import { getGatewayBaseUrl, getProxyBaseUrl } from "@/components/networking"; +import { getGatewayBaseUrl } from "@/components/networking"; import { toast } from "@/lib/toast"; import { extractPromptCacheTokens } from "@/utils/promptCacheUsage"; import { parseUsageCost } from "./usage_cost"; import { buildPlaygroundHeaders, type CustomHeaders } from "./request_headers"; +import { buildMcpToolBlocks, type BuildMcpToolBlocksArgs } from "./mcp_tool_blocks"; import type { MCPEvent } from "@/components/mcp_tools/types"; import { MCPServer, MCPToolset } from "@/components/mcp_tools/types"; import { @@ -132,7 +133,6 @@ export async function makeOpenAIResponsesRequest( } const proxyBaseUrl = customBaseUrl || getGatewayBaseUrl(); - const mcpBaseUrl = customBaseUrl || getProxyBaseUrl(); const headers = buildPlaygroundHeaders(tags, customHeaders); const client = new openai.OpenAI({ @@ -178,53 +178,13 @@ export async function makeOpenAIResponsesRequest( }); // Build tools array - const tools: any[] = []; - - // Add MCP servers if selected - if (selectedMCPServers && selectedMCPServers.length > 0) { - if (selectedMCPServers.includes("__all__")) { - // All MCP Servers selected - const allServersTool = { - type: "mcp", - server_label: "litellm", - server_url: `${mcpBaseUrl}/mcp`, - require_approval: "never", - }; - tools.push(allServersTool); - } else { - // Individual servers/toolsets selected - create one entry per item - selectedMCPServers.forEach((serverId) => { - if (serverId.startsWith("toolset:")) { - // Toolset: same /{name}/mcp pattern as individual servers - const toolsetId = serverId.slice("toolset:".length); - const toolset = mcpToolsets?.find((t) => t.toolset_id === toolsetId); - const toolsetName = toolset?.toolset_name || toolsetId; - const toolsetTool = { - type: "mcp", - server_label: toolsetName, - server_url: `${mcpBaseUrl}/mcp/${encodeURIComponent(toolsetName)}`, - require_approval: "never", - }; - tools.push(toolsetTool); - } else { - const server = mcpServers?.find((s) => s.server_id === serverId); - // Use server_name for both routing and labelling. server_name is the - // unique registered identifier; aliases can collide across servers. - const routeName = server?.server_name || serverId; - const allowedTools = mcpServerToolRestrictions?.[serverId] || []; - - const serverTool = { - type: "mcp", - server_label: routeName, // unique per request — collisions cause silent tool-routing failures - server_url: `${mcpBaseUrl}/mcp/${encodeURIComponent(routeName)}`, - require_approval: "never", - ...(allowedTools.length > 0 ? { allowed_tools: allowedTools } : {}), - }; - tools.push(serverTool); - } - }); - } - } + const mcpSelection: BuildMcpToolBlocksArgs = { + selectedMCPServers, + mcpServers, + mcpToolsets, + mcpServerToolRestrictions, + }; + const tools: any[] = buildMcpToolBlocks(mcpSelection); // Add code_interpreter tool if enabled (OpenAI auto-creates container) if (codeInterpreterEnabled) {