diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index b7c6f805233..809829db5a8 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -14,6 +14,7 @@ import litellm from litellm._logging import verbose_logger from litellm.constants import KUBERNETES_POD_ROUTING_KEY from litellm.litellm_core_utils.llm_response_utils.get_headers import get_provider_request_id +from litellm.types.litellm_params import MCP_HANDLER_KWARG_NAMES from litellm.types.llms.openai import AllMessageValues, OpenAIChatCompletionFinishReason if TYPE_CHECKING: @@ -807,36 +808,9 @@ def filter_exceptions_from_params(data: object, max_depth: int = 20) -> Any: return data -def filter_internal_params(data: dict, additional_internal_params: set | None = None) -> dict: - """ - Filter out LiteLLM internal parameters that shouldn't be sent to provider APIs. - - This removes internal/MCP-related parameters that are used by LiteLLM internally - but should not be included in API requests to providers. - - Args: - data: Dictionary of parameters to filter - additional_internal_params: Optional set of additional internal parameter names to filter - - Returns: - Filtered dictionary with internal parameters removed - """ - if not isinstance(data, dict): - return data - - # Known internal parameters that should never be sent to provider APIs - internal_params: Final = { - "skip_mcp_handler", - "mcp_handler_context", - "_skip_mcp_handler", - } - - # Add any additional internal params if provided - if additional_internal_params: - internal_params.update(additional_internal_params) - - # Filter out internal parameters - return {k: v for k, v in data.items() if k not in internal_params} +def filter_internal_params(data: dict) -> dict: + """Drop the MCP handler keys only, since fallback callers send the rest of the kwargs back into acompletion.""" + return {k: v for k, v in data.items() if k not in MCP_HANDLER_KWARG_NAMES} def redact_nested_match_and_regex_keys( diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index 85557090ae9..88eb8e73525 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -97,7 +97,12 @@ _OPTIONAL_KWARGS_KEYS: Final = OPTIONAL_KWARGS_KEYS _CONTROL_OPTIONS: Final = TypeAdapter(ControlOptions) _CONTROL_OPTION_NAMES: Final = tuple(field.name for field in fields(ControlOptions)) _MAX_SHOWN_INT_BITS: Final = 64 -_EXPECTED: Final = f"expected a positive integer of at most {MAX_CONTROL_INT_DIGITS} digits" +_EXPECTED: Final = MappingProxyType( + { + "stream_chunk_size": f"expected a positive integer of at most {MAX_CONTROL_INT_DIGITS} digits", + "fake_stream": "expected true or false", + } +) class _BoundedRepr(reprlib.Repr): @@ -123,7 +128,7 @@ def parse_control_options(kwargs: Mapping[str, object]) -> ControlOptions | Inva except ValidationError as e: param: Final = str(e.errors(include_url=False)[0]["loc"][0]) return InvalidControlOption( - param=param, message=f"Invalid {param}={_BOUNDED_REPR.repr(given[param])}: {_EXPECTED}" + param=param, message=f"Invalid {param}={_BOUNDED_REPR.repr(given[param])}: {_EXPECTED[param]}" ) diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index 7d02b9ef230..258bb536d9d 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -299,9 +299,10 @@ class BedrockConverseLLM(BaseAWSLLM): ): ## SETUP ## stream: Final = optional_params.pop("stream", None) - stream_chunk_size: Final = stored_control_options(litellm_params).stream_chunk_size if stream is True else None + control_options: Final = stored_control_options(litellm_params) + stream_chunk_size: Final = control_options.stream_chunk_size if stream is True else None unencoded_model_id: Final = optional_params.pop("model_id", None) - fake_stream = optional_params.pop("fake_stream", False) + fake_stream = optional_params.pop("fake_stream", False) or control_options.fake_stream is True json_mode: Final = optional_params.get("json_mode") is True if unencoded_model_id is not None: modelId = self.encode_model_id(model_id=unencoded_model_id) diff --git a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py index 87d4fc0b32e..f8fcf552dbf 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py @@ -30,6 +30,7 @@ from litellm.llms.custom_httpx.http_handler import ( HTTPHandler, get_httpx_client, ) +from litellm.types.litellm_params import MCP_HANDLER_KWARG_NAMES from litellm.types.llms.bedrock import GuardrailConfigBlock from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ModelResponse, Usage @@ -202,7 +203,11 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): custom_prompt_dict=custom_prompt_dict, ) inference_params = copy.deepcopy(optional_params) - inference_params = {k: v for k, v in inference_params.items() if k not in self.aws_authentication_params} + inference_params = { + k: v + for k, v in inference_params.items() + if k not in self.aws_authentication_params and k not in MCP_HANDLER_KWARG_NAMES + } request_data: dict = {} if provider == "cohere": if model.startswith("cohere.command-r"): diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index ac2ac1fe31b..b6385e3f73b 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -51,7 +51,7 @@ from litellm.litellm_core_utils.audio_utils.subtitle_utils import ( ) from litellm.litellm_core_utils.core_helpers import set_provider_response_headers_in_hidden_params from litellm.litellm_core_utils.dot_notation_indexing import delete_nested_value -from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS +from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS, stored_control_options from litellm.litellm_core_utils.llm_request_utils import serialize_multipart_form_fields from litellm.litellm_core_utils.realtime_errors import ( close_after_upstream_handshake_refusal, @@ -734,6 +734,7 @@ class BaseLLMHTTPHandler: fake_stream = ( fake_stream or optional_params.pop("fake_stream", False) + or stored_control_options(litellm_params).fake_stream is True or provider_config.should_fake_stream(model=model, custom_llm_provider=custom_llm_provider, stream=stream) ) diff --git a/litellm/main.py b/litellm/main.py index d647cab24e3..64206560ae8 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -152,6 +152,7 @@ from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ( BaseUrlRoutingHint, CustomPricingLiteLLMParams, + DynamicPromptManagementParamLiteral, ModelResponseStream, ProviderRoutingHints, RawRequestTypedDict, @@ -6693,7 +6694,9 @@ def embedding( "encoding_format", ] default_params: Final = [*openai_params, "aembedding", "extra_headers"] - non_default_params: Final = filter_out_litellm_params(kwargs, excluding=default_params) + non_default_params: Final = filter_out_litellm_params( + kwargs, excluding=(*default_params, *DynamicPromptManagementParamLiteral.list_all_params()) + ) litellm_params_dict: Final = get_litellm_params(**kwargs) diff --git a/litellm/types/litellm_params.py b/litellm/types/litellm_params.py index 5f318cc3c69..7e34264ba88 100644 --- a/litellm/types/litellm_params.py +++ b/litellm/types/litellm_params.py @@ -307,6 +307,7 @@ class ControlOptions: ] | None ) = None + fake_stream: bool | None = None @dataclass(frozen=True, slots=True, kw_only=True) @@ -345,6 +346,13 @@ class CallState: context_compaction_state: "CompactionState | None" = field(default=None, metadata=wire("_context_compaction_state")) +@dataclass(frozen=True, slots=True, kw_only=True) +class McpHandlerState: + recursion_guard: bool | None = field(default=None, metadata=wire("_skip_mcp_handler")) + skip_mcp_handler: bool | None = None + mcp_handler_context: Mapping[str, object] | None = None + + @dataclass(frozen=True, slots=True, kw_only=True) class AgenticLoopState: depth: int | None = field(default=None, metadata=wire("_agentic_loop_depth")) @@ -416,6 +424,7 @@ class EntrypointState: class InternalState: call: CallState agentic_loop: AgenticLoopState + mcp_handler: McpHandlerState router: RouterState proxy: ProxyRequestState entrypoint: EntrypointState @@ -443,3 +452,4 @@ def owned_wire_names(root: type) -> tuple[str, ...]: OWNED_KWARG_NAMES: Final = tuple(chain.from_iterable(owned_wire_names(root) for root in LITELLM_OWNED_ROOTS)) AGENTIC_LOOP_KWARG_NAMES: Final = (*wire_names(AgenticLoopState), *wire_names(AgenticLoopOptions)) BEDROCK_BATCH_KWARG_NAMES: Final = wire_names(BedrockBatchConnection) +MCP_HANDLER_KWARG_NAMES: Final = frozenset(wire_names(McpHandlerState)) diff --git a/tests/integration/providers/test_internal_params_wire.py b/tests/integration/providers/test_internal_params_wire.py index f9f5d1e5478..aab400309e8 100644 --- a/tests/integration/providers/test_internal_params_wire.py +++ b/tests/integration/providers/test_internal_params_wire.py @@ -4,15 +4,29 @@ import json import os import struct import zlib -from collections.abc import Callable, Mapping +from collections.abc import Callable, Iterator, Mapping +from contextlib import suppress +from dataclasses import dataclass from pathlib import Path -from typing import Final +from types import MappingProxyType +from typing import Final, get_args +from urllib.parse import parse_qs +import openai import pytest from integration._support.upstream import INTERNAL_FIELDS from integration._support.wire import Reply, Request, wire_server +from pydantic import TypeAdapter import litellm +from litellm.constants import BEDROCK_EMBEDDING_PROVIDERS_LITERAL, BEDROCK_INVOKE_PROVIDERS_LITERAL +from litellm.exceptions import LiteLLMUnknownProvider +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.bedrock.common_utils import BedrockModelInfo, BedrockRoute +from litellm.types import litellm_params +from litellm.types.litellm_params import OWNED_KWARG_NAMES, wire_names +from litellm.types.utils import LlmProviders from tests._support.stream_chunk_size import keys_at_every_depth TEXT: Final = "wire control" @@ -313,3 +327,469 @@ async def test_internal_params_never_reach_provider_body( assert not INTERNAL_FIELDS.intersection(keys) assert not frozenset(key for key in keys if key.startswith("_litellm_")), keys assert _custom_key(body, provider) == 1 + + +WALK_TEXT: Final = "walk every provider" +UNDECLARED_SENTINEL: Final = "_litellm_undeclared_sentinel" + + +def _ignore_log_details(_details: Mapping[str, object]) -> None: + return None + + +WALK_SEEDS: Final[Mapping[str, object]] = MappingProxyType( + { + "aws_batch_role_arn": "arn:aws:iam::000000000000:role/wire", + "s3_bucket_name": "wire-bucket", + "s3_region_name": "us-east-1", + "s3_endpoint_url": "https://s3.wire.invalid", + "s3_output_bucket_name": "wire-output", + "s3_bucket_owner": "000000000000", + "s3_access_key_id": "wire-access-key", + "s3_secret_access_key": "wire-secret-key", + "s3_encryption_key_id": "wire-kms-key", + "bedrock_tags": [{"key": "wire", "value": "walk"}], + "model_info": {"id": "wire-deployment"}, + "rpm": 7, + "tpm": 7, + "itpm": 7, + "otpm": 7, + "default_api_key_rpm_limit": 7, + "default_api_key_tpm_limit": 7, + "max_parallel_requests": 1, + "weight": 1, + "order": 1, + "tag_regex": ["^wire$"], + "max_file_size_mb": 1.0, + "auto_router_config_path": "wire-router.json", + "auto_router_config": "wire-router", + "auto_router_default_model": "wire-default", + "auto_router_embedding_model": "wire-embedding", + "auto_router_max_input_chars": 64, + "auto_router_routing_compression": "wire", + "auto_router_model_compression": "wire", + "complexity_router_config": {"wire": 1}, + "complexity_router_default_model": "wire-default", + "adaptive_router_config": {"wire": 1}, + "adaptive_router_default_model": "wire-default", + "quality_router_config": {"wire": 1}, + "quality_router_default_model": "wire-default", + "cost_per_query": 0.5, + "max_budget": 100.0, + "budget_duration": "30d", + "id": "wire-id", + "metadata": {"wire": 1}, + "litellm_metadata": {"wire": 1}, + "tags": ["wire"], + "litellm_trace_id": "wire-trace", + "litellm_session_id": "wire-session", + "litellm_request_debug": False, + "logger_fn": _ignore_log_details, + "verbose": False, + "no-log": False, + "log_client_error_tracebacks": False, + "max_agentic_loops": 1, + "merge_reasoning_content_in_choices": False, + "enable_json_schema_validation": False, + "complete_response": False, + "keepalive_seconds": 30.0, + "allow_client_keepalive_override": False, + "stream_chunk_size": 64, + "litellm_call_id": "wire-call", + "completion_call_id": "wire-completion", + "preset_cache_key": "wire-cache-key", + "cache_key": "wire-cache-key", + "fake_stream": False, + "_skip_mcp_handler": True, + "skip_mcp_handler": True, + "mcp_handler_context": {"server": "wire"}, + "fallback_depth": 1, + "max_fallbacks": 1, + UNDECLARED_SENTINEL: "internal", + } +) +OWNED_LEAVES_NOT_SEEDED: Final = MappingProxyType( + { + litellm_params.ProviderConnection: "the route needs its own connection settings to send anything", + litellm_params.AnthropicFederationConnection: "switches the route to workload identity federation", + litellm_params.OpenAIFederationConnection: "switches the route to workload identity federation", + litellm_params.DispatchOptions: "chooses which provider route handles the call", + litellm_params.RoutingOptions: "retries and fallbacks change which requests are sent", + litellm_params.CachingOptions: "turns response caching on or off", + litellm_params.GuardrailOptions: "runs the named guardrails before the call", + litellm_params.PromptOptions: "changes how the prompt is built or loads a managed prompt", + litellm_params.MockOptions: "answers without sending a request", + litellm_params.AgenticLoopState: "drives the agentic loop between follow-up calls", + litellm_params.ProxyRequestState: "set by the proxy for its own request handling", + litellm_params.EntrypointState: "picks the sync or async entry point", + } +) +OWNED_NAMES_NOT_SEEDED: Final = MappingProxyType( + { + "silent_model": "sends a second request to another model", + "base_model": "stands in for the model name when a provider config resolves the model", + "model_alias_map": "renames the model before the route is chosen", + "data_residency": "picks a regional endpoint", + "litellm_logging_obj": "must be the call's own logging object", + "stream_response": "holds a response object, not request input", + "_context_compaction_state": "holds compaction state, not request input", + "_router_weights": "holds router weight state, not request input", + "attempted_targets": "holds the router's fallback bookkeeping, not request input", + } +) +EMBEDDING_WALK_SEEDS: Final[Mapping[str, object]] = MappingProxyType( + { + **WALK_SEEDS, + "cache_control_injection_points": [{"location": "message", "role": "user"}], + "knowledge_bases": ["kb_wire"], + "vector_store_ids": ["vs_wire"], + } +) +WALK_REJECTION: Final = Reply( + status=400, + body=json.dumps({"error": {"message": "rejected by the wire", "type": "invalid_request_error"}}).encode(), +) +SYNTHETIC_KEY: Final = MappingProxyType({"api_key": "synthetic-key"}) +AZURE_CONNECTION: Final = MappingProxyType({**SYNTHETIC_KEY, "api_version": "2025-01-01-preview"}) +AWS_CONNECTION: Final = MappingProxyType( + {"aws_access_key_id": "fake", "aws_secret_access_key": "fake", "aws_region_name": "us-east-1"} +) +AWS_BASE_PARAMS: Final = ("api_base", "aws_bedrock_runtime_endpoint") +ONE_PIXEL_GIF: Final = "data:image/gif;base64,R0lGODlhAQABAAAAACw=" + + +@dataclass(frozen=True, slots=True) +class Route: + model: str + connection: Mapping[str, object] = SYNTHETIC_KEY + base_params: tuple[str, ...] = ("api_base",) + base_suffix: str = "" + request: Mapping[str, object] = MappingProxyType({}) + protocol_fields: frozenset[str] = frozenset() + streams: bool = True + + def parameters(self, wire_url: str) -> dict[str, object]: + return { + "model": self.model, + **self.connection, + **dict.fromkeys(self.base_params, wire_url + self.base_suffix), + **self.request, + } + + +def _bedrock(model: str, **connection: object) -> Route: + return Route(f"bedrock/{model}", MappingProxyType({**AWS_CONNECTION, **connection}), AWS_BASE_PARAMS) + + +IMPORTED_MODEL: Final = "arn:aws:bedrock:us-east-1:000000000000:imported-model/wire" +BEDROCK_INVOKE_MODELS: Final = MappingProxyType( + { + "anthropic": "anthropic.claude-haiku-4-5-20251001-v1:0", + "cohere": "cohere.command-r-plus-v1:0", + "mistral": "mistral.mistral-large-2402-v1:0", + "amazon": "amazon.titan-text-express-v1", + "meta": "meta.llama3-1-70b-instruct-v1:0", + "llama": f"llama/{IMPORTED_MODEL}", + "ai21": "ai21.jamba-1-5-large-v1:0", + "nova": "amazon.nova-pro-v1:0", + "deepseek_r1": f"deepseek_r1/{IMPORTED_MODEL}", + "qwen3": f"qwen3/{IMPORTED_MODEL}", + "qwen2": f"qwen2/{IMPORTED_MODEL}", + "twelvelabs": "twelvelabs.pegasus-1-2-v1:0", + "openai": "openai.gpt-oss-120b-1:0", + "moonshot": "moonshot.kimi-k2-thinking", + } +) +BEDROCK_INVOKE_FAMILIES_WITHOUT_CHAT: Final = MappingProxyType({"stability": "generates images only"}) +BEDROCK_ROUTES: Final = MappingProxyType( + { + "converse": _bedrock("converse/anthropic.claude-haiku-4-5-20251001-v1:0"), + "converse_like": _bedrock("converse_like/anthropic.claude-haiku-4-5-20251001-v1:0"), + "openai": _bedrock("openai/wire-model"), + "chat_completions": _bedrock("chat_completions/openai.gpt-oss-120b-1:0"), + "mantle": _bedrock("mantle/openai.gpt-oss-120b"), + "claude_platform": _bedrock("claude_platform/claude-haiku-4-5", workspace_id="wrk_wire"), + "agent": _bedrock("agent/AGENTID/ALIASID"), + "agentcore": _bedrock("agentcore/arn:aws:bedrock-agentcore:us-east-1:000000000000:runtime/wire"), + **{f"invoke/{family}": _bedrock(f"invoke/{model}") for family, model in BEDROCK_INVOKE_MODELS.items()}, + } +) +BEDROCK_ROUTES_WITHOUT_CHAT: Final = MappingProxyType({"async_invoke": "starts an async job instead of a chat call"}) + +API_BASE_CHAT_PROVIDERS: Final = ( + "ai21", "ai21_chat", "aiml", "aiohttp_openai", "amazon_nova", "anthropic", "anthropic_text", "apertis", + "assemblyai", "azure_ai", "baseten", "cerebras", "charity_engine", "chutes", "cloudflare", "codestral", + "cognition", "cohere", "cohere_chat", "cometapi", "compactifai", "cortecs", "custom", "custom_openai", + "darkbloom", "dashscope", "databricks", "datarobot", "deepinfra", "deepseek", "docker_model_runner", "edenai", + "empower", "featherless_ai", "fireworks_ai", "friendliai", "galadriel", "gemini", "github", "groq", "helicone", + "heroku", "hosted_vllm", "huggingface", "hyperbolic", "inception", "lambda_ai", "langflow", "langgraph", + "lemonade", "libertai", "litellm_proxy", "llamafile", "lm_studio", "maritalk", "meta", "meta_llama", "minimax", + "mistral", "modelscope", "moonshot", "morph", "nadir", "nano-gpt", "nebius", "neosantara", "novita", "nscale", + "nvidia_nim", "ollama", "ollama_chat", "oobabooga", "openai", "openrouter", "ovhcloud", "parasail", "perplexity", + "petals", "pinstripes", "poe", "prism", "publicai", "qwen_ai_platform", "qwencloud", "reka", "sail", "sambanova", + "scaleway", "scx-ai", "snowflake", "synthetic", "tencent", "tensormesh", "text-completion-codestral", + "text-completion-inception", "text-completion-openai", "together_ai", "v0", "vercel_ai_gateway", "volcengine", + "wandb", "xai", "xiaomi_mimo", "xinference", "zai", +) # fmt: skip +CHAT_ROUTES: Final = MappingProxyType( + { + **{provider: Route(f"{provider}/wire-model") for provider in API_BASE_CHAT_PROVIDERS}, + "a2a": Route("a2a/wire-model", protocol_fields=frozenset({"id"})), + "azure": Route("azure/wire-model", AZURE_CONNECTION), + "azure_text": Route("azure_text/wire-model", AZURE_CONNECTION), + "bedrock_mantle": Route("bedrock_mantle/wire-model", AWS_CONNECTION, AWS_BASE_PARAMS), + "clarifai": Route("clarifai/user.app.wire-model"), + "fal_ai": Route( + "fal_ai/wire-model", + request=MappingProxyType( + { + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": WALK_TEXT}, + {"type": "image_url", "image_url": {"url": ONE_PIXEL_GIF}}, + ], + } + ] + } + ), + streams=False, + ), + "gradient_ai": Route("gradient_ai/wire-model"), + "nlp_cloud": Route("nlp_cloud/chatdolphin", base_suffix="/"), + "predibase": Route("predibase/wire-model", MappingProxyType({**SYNTHETIC_KEY, "predibase_tenant_id": "wire"})), + "ragflow": Route("ragflow/chat/wire-chat/wire-model"), + "scaledown": Route("scaledown/summarize"), + "triton": Route("triton/wire-model", base_suffix="/generate"), + **{f"bedrock/{route}": target for route, target in BEDROCK_ROUTES.items()}, + } +) + +VENDOR_HOST: Final = "calls the vendor host even when api_base is set" +VENDOR_AUTH: Final = "needs vendor credentials or a token exchange the wire cannot fake" +DEVICE_LOGIN: Final = "needs an interactive device-code login" +CHAT_PROVIDERS_OFF_THE_WIRE: Final = MappingProxyType( + { + **dict.fromkeys(("bytez", "replicate", "sagemaker", "sagemaker_chat", "sagemaker_nova"), VENDOR_HOST), + **dict.fromkeys( + ("gdc", "gigachat", "microsoft_365_copilot", "oci", "sap", "vertex_ai", "vertex_ai_beta", "watsonx", + "watsonx_text"), + VENDOR_AUTH, + ), + **dict.fromkeys(("chatgpt", "github_copilot"), DEVICE_LOGIN), + "vllm": "runs the local vllm package instead of calling a server", + } +) # fmt: skip + +EMBEDDING_ROUTES: Final = MappingProxyType( + { + **{ + provider: Route(f"{provider}/wire-model") + for provider in ( + "azure_ai", "cohere", "cohere_chat", "cometapi", "dashscope", "databricks", "edenai", "fireworks_ai", + "gemini", "hosted_vllm", "infinity", "jina_ai", "litellm_proxy", "llamafile", "lm_studio", "mistral", + "nebius", "nvidia_nim", "ollama", "oobabooga", "openai", "openai_like", "openrouter", "ovhcloud", + "perplexity", "qwen_ai_platform", "qwencloud", "sambanova", "snowflake", "together_ai", "triton", + "vercel_ai_gateway", "volcengine", "voyage", "wandb", "xinference", + ) + }, + "azure": Route("azure/wire-model", AZURE_CONNECTION), + "huggingface": Route("huggingface/wire-model", request=MappingProxyType({"input_type": "feature-extraction"})), + "bedrock/amazon": _bedrock("amazon.titan-embed-image-v1"), + "bedrock/cohere": _bedrock("cohere.embed-english-v3"), + "bedrock/nova": _bedrock("amazon.nova-2-multimodal-embeddings-v1:0"), + "bedrock/twelvelabs": _bedrock("twelvelabs.marengo-embed-2-7-v1:0"), + } +) # fmt: skip +EMBEDDING_PROVIDERS_OFF_THE_WIRE: Final = MappingProxyType( + { + "sagemaker": VENDOR_HOST, + **dict.fromkeys(("gigachat", "oci", "sap", "vertex_ai", "watsonx"), VENDOR_AUTH), + **dict.fromkeys(("chatgpt", "github_copilot"), DEVICE_LOGIN), + } +) +REFUSALS: Final = (openai.OpenAIError, BaseLLMException, ValueError) + +CHAT_ROUTE_BUGS: Final = MappingProxyType( + { + "gradient_ai": "BUG: rejects the max_retries that completion() derives from num_retries, before sending", + "huggingface": "BUG: sends the max_retries that completion() derives from num_retries in the request body", + "ragflow": "BUG: sends the max_retries that completion() derives from num_retries in the request body", + } +) +JSON_BODY: Final = TypeAdapter(dict[str, object]) + + +def _provider(route_id: str) -> str: + return route_id.split("/", 1)[0] + + +def _chat_cases() -> Iterator[tuple[str, bool]]: + for route_id, route in CHAT_ROUTES.items(): + yield route_id, False + if route.streams: + yield route_id, True + + +def _unwalked(routes: Mapping[str, Route], off_the_wire: Mapping[str, str]) -> tuple[str, ...]: + walked: Final = frozenset(_provider(route_id) for route_id in routes) + return tuple(provider.value for provider in LlmProviders if provider.value not in walked | frozenset(off_the_wire)) + + +def _unrouted_model(provider: str) -> str: + return CHAT_ROUTES[provider].model if provider in CHAT_ROUTES else f"{provider}/wire-model" + + +def _complete(parameters: Mapping[str, object]) -> object: + return litellm.completion(**{"messages": [{"role": "user", "content": WALK_TEXT}], **parameters}) + + +def _embed(parameters: Mapping[str, object]) -> object: + return litellm.embedding(**{"input": [WALK_TEXT], **parameters}) + + +def _unrouted_parameters(provider: str, wire_url: str) -> dict[str, object]: + return {"model": _unrouted_model(provider), **SYNTHETIC_KEY, "api_base": wire_url, "num_retries": 0, "timeout": 5} + + +def _sent_body(request: Request) -> Mapping[str, object]: + if request.headers.get("content-type", "").startswith("application/x-www-form-urlencoded"): + fields: Final = parse_qs(request.body.decode(), keep_blank_values=True, strict_parsing=True) + return {name: values[0] for name, values in fields.items()} + return JSON_BODY.validate_json(request.body) + + +def _walked_bodies( + route: Route, seeds: Mapping[str, object], send: Callable[[dict[str, object]], object] +) -> tuple[Mapping[str, object], ...]: + with wire_server(lambda _request: WALK_REJECTION) as wire: + with suppress(openai.OpenAIError, BaseLLMException): + result: Final = send({**seeds, **route.parameters(wire.url), "num_retries": 0, "timeout": 5}) + if isinstance(result, Iterator): + for _chunk in result: + pass + requests: Final = wire.drain() + assert requests, f"{route.model} sent nothing to the wire" + return tuple(_sent_body(request) for request in requests if request.body) + + +def _assert_only_provider_fields( + bodies: tuple[Mapping[str, object], ...], seeds: Mapping[str, object], route: Route +) -> None: + assert any(WALK_TEXT in json.dumps(body) for body in bodies), bodies + for body in bodies: + assert not (frozenset(seeds) - route.protocol_fields) & keys_at_every_depth(body), body + assert not (frozenset(OWNED_KWARG_NAMES) - route.protocol_fields) & frozenset(body), body + assert not tuple(key for key in body if key.startswith("_")), body + + +def test_the_form_body_parser_keeps_a_blank_leaked_field() -> None: + form: Final = Request( + method="POST", + target="/", + headers={"content-type": "application/x-www-form-urlencoded"}, + body=b"inputs=walk+every+provider&skip_mcp_handler=", + ) + + with pytest.raises(AssertionError): + _assert_only_provider_fields((_sent_body(form),), WALK_SEEDS, Route("petals/wire-model")) + + +def test_every_owned_name_is_seeded_or_has_a_reason_it_is_not() -> None: + seeded: Final = frozenset(WALK_SEEDS) - {UNDECLARED_SENTINEL} + unseeded: Final = frozenset(OWNED_NAMES_NOT_SEEDED).union(*map(wire_names, OWNED_LEAVES_NOT_SEEDED)) + + assert seeded.isdisjoint(unseeded), seeded & unseeded + assert seeded | unseeded == frozenset(OWNED_KWARG_NAMES), frozenset(OWNED_KWARG_NAMES) ^ (seeded | unseeded) + + +@pytest.mark.allow_hosts(["127.0.0.1"]) +@pytest.mark.parametrize(("route_id", "stream"), tuple(_chat_cases())) +def test_every_chat_route_keeps_internal_params_out_of_the_provider_body( + provider_wire_environment: None, route_id: str, stream: bool +) -> None: + if route_id in CHAT_ROUTE_BUGS: + pytest.skip(CHAT_ROUTE_BUGS[route_id]) + + bodies: Final = _walked_bodies( + CHAT_ROUTES[route_id], WALK_SEEDS, lambda parameters: _complete({**parameters, "stream": stream}) + ) + + _assert_only_provider_fields(bodies, WALK_SEEDS, CHAT_ROUTES[route_id]) + + +@pytest.mark.allow_hosts(["127.0.0.1"]) +@pytest.mark.parametrize("route_id", EMBEDDING_ROUTES) +def test_every_embedding_route_keeps_internal_params_out_of_the_provider_body( + provider_wire_environment: None, route_id: str +) -> None: + bodies: Final = _walked_bodies(EMBEDDING_ROUTES[route_id], EMBEDDING_WALK_SEEDS, _embed) + + _assert_only_provider_fields(bodies, EMBEDDING_WALK_SEEDS, EMBEDDING_ROUTES[route_id]) + + +@pytest.mark.allow_hosts(["127.0.0.1"]) +@pytest.mark.parametrize("provider", CHAT_PROVIDERS_OFF_THE_WIRE) +def test_a_chat_provider_kept_off_the_wire_sends_nothing_to_it(provider_wire_environment: None, provider: str) -> None: + with wire_server(lambda _request: WALK_REJECTION) as wire: + with suppress(*REFUSALS): + _complete(_unrouted_parameters(provider, wire.url)) + sent: Final = wire.drain() + + assert not sent, f"{provider} reaches the wire, so walk it in CHAT_ROUTES instead" + + +@pytest.mark.allow_hosts(["127.0.0.1"]) +@pytest.mark.parametrize("provider", _unwalked(CHAT_ROUTES, CHAT_PROVIDERS_OFF_THE_WIRE)) +def test_a_provider_without_a_chat_route_has_no_chat_handler(provider_wire_environment: None, provider: str) -> None: + with wire_server(lambda _request: WALK_REJECTION) as wire: + with pytest.raises(LiteLLMUnknownProvider): + _complete(_unrouted_parameters(provider, wire.url)) + sent: Final = wire.drain() + + assert not sent, sent + + +@pytest.mark.allow_hosts(["127.0.0.1"]) +@pytest.mark.parametrize("provider", EMBEDDING_PROVIDERS_OFF_THE_WIRE) +def test_an_embedding_provider_kept_off_the_wire_sends_nothing_to_it( + provider_wire_environment: None, provider: str +) -> None: + with wire_server(lambda _request: WALK_REJECTION) as wire: + with suppress(*REFUSALS): + _embed(_unrouted_parameters(provider, wire.url)) + sent: Final = wire.drain() + + assert not sent, f"{provider} reaches the wire, so walk it in EMBEDDING_ROUTES instead" + + +@pytest.mark.allow_hosts(["127.0.0.1"]) +@pytest.mark.parametrize("provider", _unwalked(EMBEDDING_ROUTES, EMBEDDING_PROVIDERS_OFF_THE_WIRE)) +def test_a_provider_without_an_embedding_route_has_no_embedding_handler( + provider_wire_environment: None, provider: str +) -> None: + with wire_server(lambda _request: WALK_REJECTION) as wire: + with pytest.raises(LiteLLMUnknownProvider): + _embed(_unrouted_parameters(provider, wire.url)) + sent: Final = wire.drain() + + assert not sent, sent + + +def test_every_bedrock_route_and_invoke_family_is_walked() -> None: + routes: Final = { + route: BedrockModelInfo.get_bedrock_route(target.model) for route, target in BEDROCK_ROUTES.items() + } + families: Final = { + family: BaseAWSLLM.get_bedrock_invoke_provider(model=f"invoke/{model}") + for family, model in BEDROCK_INVOKE_MODELS.items() + } + + assert routes == {route: route.split("/", 1)[0] for route in BEDROCK_ROUTES} + assert families == {family: family for family in BEDROCK_INVOKE_MODELS} + assert {*routes.values(), *BEDROCK_ROUTES_WITHOUT_CHAT} == frozenset(get_args(BedrockRoute)) + assert {*families, *BEDROCK_INVOKE_FAMILIES_WITHOUT_CHAT} == frozenset(get_args(BEDROCK_INVOKE_PROVIDERS_LITERAL)) + assert {route_id.split("/", 1)[1] for route_id in EMBEDDING_ROUTES if route_id.startswith("bedrock/")} == frozenset( + get_args(BEDROCK_EMBEDDING_PROVIDERS_LITERAL) + ) diff --git a/tests/unit/caching/test_caching.py b/tests/unit/caching/test_caching.py index 91083921130..e01c7fd75af 100644 --- a/tests/unit/caching/test_caching.py +++ b/tests/unit/caching/test_caching.py @@ -442,6 +442,8 @@ def test_provider_specific_cache_key_ignores_litellm_owned_kwargs(monkeypatch: p assert cache.get_cache_key(**request, _litellm_control={"stream_chunk_size": 64}) == base_key assert cache.get_cache_key(**request, litellm_trace_id="trace-1") == base_key + assert cache.get_cache_key(**request, _skip_mcp_handler=True) == base_key + assert cache.get_cache_key(**request, fake_stream=True) == base_key assert cache.get_cache_key(**{**request, "top_k": 6}) != base_key diff --git a/tests/unit/litellm_core_utils/test_core_helpers.py b/tests/unit/litellm_core_utils/test_core_helpers.py index 6459aa98e70..bc0a1af15af 100644 --- a/tests/unit/litellm_core_utils/test_core_helpers.py +++ b/tests/unit/litellm_core_utils/test_core_helpers.py @@ -17,9 +17,10 @@ from litellm.litellm_core_utils.core_helpers import ( budget_reservation_from_metadata, drop_params_env_flag, drop_params_flag, + filter_internal_params, get_litellm_metadata_from_kwargs, - get_parent_otel_span_from_kwargs, get_or_create_metadata_bucket, + get_parent_otel_span_from_kwargs, get_provider_response_headers_from_hidden_params, map_finish_reason, normalize_drop_params, @@ -191,6 +192,27 @@ class TestGetOrCreateMetadataBucket: assert bucket is request_data["litellm_metadata"] +def test_filter_internal_params_drops_only_the_mcp_handler_keys() -> None: + rest_of_the_call: Final = { + "api_key": "sk-test", + "guardrails": ["default"], + "metadata": {"team": "search"}, + "stream_chunk_size": 64, + "temperature": 0.2, + } + + filtered: Final = filter_internal_params( + { + **rest_of_the_call, + "_skip_mcp_handler": True, + "skip_mcp_handler": True, + "mcp_handler_context": {"server": "docs"}, + } + ) + + assert filtered == rest_of_the_call + + def test_reconstruct_model_name_prefers_deployment_value(): """Ensure deployment metadata wins when reconstructing the model name.""" @@ -258,9 +280,7 @@ class TestMapFinishReasonAnthropic: ("content_filtered", "content_filter"), ], ) - def test_anthropic_finish_reasons( - self, provider_reason: str, expected: str - ) -> None: + def test_anthropic_finish_reasons(self, provider_reason: str, expected: str) -> None: assert map_finish_reason(provider_reason) == expected def test_refusal(self): @@ -332,9 +352,7 @@ class TestMapFinishReasonZhipu: class TestMapFinishReasonOpenAIPassthrough: - @pytest.mark.parametrize( - "reason", ["stop", "length", "tool_calls", "function_call", "content_filter"] - ) + @pytest.mark.parametrize("reason", ["stop", "length", "tool_calls", "function_call", "content_filter"]) def test_openai_values_pass_through(self, reason): assert map_finish_reason(reason) == reason @@ -345,9 +363,7 @@ class TestMapFinishReasonGenericError: assert map_finish_reason("error") == "stop" def test_lowercase_error_does_not_warn(self, mocker): - warn = mocker.patch( - "litellm.litellm_core_utils.core_helpers.verbose_logger.warning" - ) + warn = mocker.patch("litellm.litellm_core_utils.core_helpers.verbose_logger.warning") assert map_finish_reason("error") == "stop" warn.assert_not_called() @@ -365,8 +381,7 @@ class TestFinishReasonMapOutputsAreValid: """Every value in _FINISH_REASON_MAP must be a valid OpenAI finish reason.""" for provider_reason, openai_reason in _FINISH_REASON_MAP.items(): assert openai_reason in VALID_OPENAI_FINISH_REASONS, ( - f"Mapped value '{openai_reason}' (from '{provider_reason}') " - f"is not a valid OpenAI finish reason" + f"Mapped value '{openai_reason}' (from '{provider_reason}') is not a valid OpenAI finish reason" ) @@ -376,28 +391,18 @@ class TestRedactNestedMatchAndRegexKeys: "assessments": [ { "sensitiveInformationPolicy": { - "piiEntities": [ - {"type": "NAME", "match": "secret-name", "action": "BLOCKED"} - ] - }, - "wordPolicy": { - "customWords": [{"match": "badword", "action": "BLOCKED"}] + "piiEntities": [{"type": "NAME", "match": "secret-name", "action": "BLOCKED"}] }, + "wordPolicy": {"customWords": [{"match": "badword", "action": "BLOCKED"}]}, } ], "regex": "should-redact-key-named-regex", } out = redact_nested_match_and_regex_keys(payload) - assert out["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0][ - "match" - ] == "[REDACTED]" - assert out["assessments"][0]["wordPolicy"]["customWords"][0]["match"] == ( - "[REDACTED]" - ) + assert out["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0]["match"] == "[REDACTED]" + assert out["assessments"][0]["wordPolicy"]["customWords"][0]["match"] == ("[REDACTED]") assert out["regex"] == "[REDACTED]" - assert payload["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][ - 0 - ]["match"] == "secret-name" + assert payload["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0]["match"] == "secret-name" def test_passes_through_none_and_str(self): assert redact_nested_match_and_regex_keys(None) is None @@ -560,13 +565,17 @@ class TestIsExpectedClientError: assert is_expected_client_error(over_budget) is True litellm_limit = RateLimitError( - message="key over rpm", llm_provider="anthropic", model="claude-haiku-4-5", + message="key over rpm", + llm_provider="anthropic", + model="claude-haiku-4-5", category=RateLimitErrorCategory.LITELLM_RATE_LIMIT, ) assert is_expected_client_error(litellm_limit) is True vendor_limit = RateLimitError( - message="rate limited upstream", llm_provider="anthropic", model="claude-haiku-4-5", + message="rate limited upstream", + llm_provider="anthropic", + model="claude-haiku-4-5", category=RateLimitErrorCategory.VENDOR_RATE_LIMIT, ) assert is_expected_client_error(vendor_limit) is False diff --git a/tests/unit/litellm_core_utils/test_get_litellm_params.py b/tests/unit/litellm_core_utils/test_get_litellm_params.py index 89173956dc6..dfb2cceb21d 100644 --- a/tests/unit/litellm_core_utils/test_get_litellm_params.py +++ b/tests/unit/litellm_core_utils/test_get_litellm_params.py @@ -199,6 +199,9 @@ class TestGetLitellmParamsKwargsExtraction: ({"stream_chunk_size": "64"}, ControlOptions(stream_chunk_size=64)), ({"stream_chunk_size": None}, ControlOptions()), ({"temperature": 0.2}, ControlOptions()), + ({"fake_stream": True}, ControlOptions(fake_stream=True)), + ({"fake_stream": "true", "stream_chunk_size": 64}, ControlOptions(stream_chunk_size=64, fake_stream=True)), + ({"fake_stream": False}, ControlOptions(fake_stream=False)), ], ) def test_control_options_are_read_from_the_request_kwargs(kwargs: dict[str, object], expected: ControlOptions) -> None: @@ -235,6 +238,12 @@ def test_control_options_reject_a_stream_chunk_size_that_is_not_a_positive_int(r ) +def test_control_options_reject_a_fake_stream_that_is_not_a_boolean() -> None: + assert parse_control_options({"fake_stream": "maybe"}) == InvalidControlOption( + param="fake_stream", message="Invalid fake_stream='maybe': expected true or false" + ) + + @pytest.mark.parametrize("raw", [10**18 - 1, "9" * 18], ids=["int", "digit_string"]) def test_control_options_accept_the_largest_18_digit_value(raw: object) -> None: assert parse_control_options({"stream_chunk_size": raw}) == ControlOptions(stream_chunk_size=10**18 - 1) diff --git a/tests/unit/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py index 1890c68b30d..8f342801a3c 100644 --- a/tests/unit/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py +++ b/tests/unit/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py @@ -45,6 +45,23 @@ def test_completion_keeps_stream_chunk_size_out_of_invoke_bodies(model: str) -> assert "stream_chunk_size" not in keys_at_every_depth(json.loads(request.content)), request.content +def test_transform_request_drops_internal_params(): + """LiteLLM-internal MCP params (e.g. skip_mcp_handler) are control flags used + inside LiteLLM and are not valid Bedrock inference parameters. Leaking them + into the provider request body makes Bedrock reject the request. The Converse + path already filters them via filter_internal_params; the invoke path must do + the same.""" + request_body = AmazonInvokeConfig().transform_request( + model="mistral.mistral-7b-instruct-v0:2", + messages=[{"role": "user", "content": "hi"}], + optional_params={"skip_mcp_handler": True, "max_tokens": 10}, + litellm_params={}, + headers={}, + ) + + assert "skip_mcp_handler" not in json.dumps(request_body) + + def test_validate_environment_maps_guardrail_config_to_invoke_headers(): """The InvokeModel API takes the guardrail identifier/version/trace as X-Amzn-Bedrock-* request headers, unlike Converse which takes them in the diff --git a/tests/unit/llms/bedrock/chat/test_bedrock_converse_handler.py b/tests/unit/llms/bedrock/chat/test_bedrock_converse_handler.py index 7d1b502df2b..895cd529b30 100644 --- a/tests/unit/llms/bedrock/chat/test_bedrock_converse_handler.py +++ b/tests/unit/llms/bedrock/chat/test_bedrock_converse_handler.py @@ -22,10 +22,12 @@ from botocore.credentials import Credentials from botocore.exceptions import ClientError import litellm +from litellm.constants import CONTROL_OPTIONS_KEY from litellm.integrations.custom_logger import CustomLogger from litellm.llms.bedrock.chat.converse_handler import BedrockConverseLLM from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.rust_bridge import configuration +from litellm.types.litellm_params import ControlOptions from litellm.types.utils import AdapterCompletionStreamWrapper, ModelResponse from tests.unit.llms.bedrock.event_loop_probe import EventLoopProbe from tests.unit.llms.bedrock.slow_upstream import ( @@ -191,6 +193,26 @@ def _sync_client_returning_converse_response(): return client +def test_a_caller_fake_stream_calls_converse_and_replays_the_answer_as_chunks(): + client: Final = MagicMock() + client.post.side_effect = lambda *_args, **_kwargs: httpx.Response( + 200, + json=CONVERSE_RESPONSE, + request=httpx.Request("POST", "https://bedrock-runtime.us-west-2.amazonaws.com"), + ) + client.__class__ = HTTPHandler + + response: Final = _run( + optional_params={"maxTokens": 16, "stream": True}, + litellm_params={CONTROL_OPTIONS_KEY: ControlOptions(fake_stream=True)}, + client=client, + ) + + assert "".join(chunk.choices[0].delta.content or "" for chunk in response) == "hi" + assert client.post.call_args.args[0].endswith("/converse"), client.post.call_args + assert client.post.call_args.kwargs["stream"] is False, client.post.call_args + + def test_the_sync_python_path_logs_pre_call_once(): logging_obj = MagicMock() response = _run( @@ -327,7 +349,6 @@ async def test_live_stream_with_a_text_stream_timeout_reports_the_stream_timeout assert getattr(excinfo.value, "timeout", None) == STREAM_TIMEOUT_SECONDS - _FAKE_STREAMED_REQUESTS: Final = pytest.mark.parametrize( ("model", "request_params"), [ @@ -338,7 +359,12 @@ _FAKE_STREAMED_REQUESTS: Final = pytest.mark.parametrize( ), ( "bedrock/converse/meta.llama3-3-70b-instruct-v1:0", - {"response_format": {"type": "json_schema", "json_schema": {"name": "answer", "schema": {"type": "object"}}}}, + { + "response_format": { + "type": "json_schema", + "json_schema": {"name": "answer", "schema": {"type": "object"}}, + } + }, ), ("bedrock/invoke/ai21.j2-ultra-v1", {}), ], @@ -491,7 +517,9 @@ async def test_fake_stream_through_a_router_adapter_is_bounded_by_the_request_ti assert isinstance(response, AsyncIterator) contents: Final = [ - chunk.choices[0].delta.content or "" async for chunk in response if isinstance(chunk, litellm.ModelResponseStream) + chunk.choices[0].delta.content or "" + async for chunk in response + if isinstance(chunk, litellm.ModelResponseStream) ] assert "".join(contents) == "hi" @@ -517,7 +545,10 @@ def test_sync_live_stream_through_the_router_fails_at_the_stream_timeout() -> No ) with pytest.raises(litellm.Timeout): router.completion( - model="bedrock", messages=[{"role": "user", "content": "hi"}], stream=True, client=slow_upstream_sync_client() + model="bedrock", + messages=[{"role": "user", "content": "hi"}], + stream=True, + client=slow_upstream_sync_client(), ) @@ -528,7 +559,10 @@ async def test_async_live_stream_through_the_router_fails_at_the_stream_timeout( ) with pytest.raises(litellm.Timeout) as excinfo: await router.acompletion( - model="bedrock", messages=[{"role": "user", "content": "hi"}], stream=True, client=slow_upstream_async_client() + model="bedrock", + messages=[{"role": "user", "content": "hi"}], + stream=True, + client=slow_upstream_async_client(), ) assert getattr(excinfo.value, "timeout", None) == STREAM_TIMEOUT_SECONDS @@ -616,7 +650,6 @@ async def test_failed_native_output_retry_reports_the_request_timeout_it_was_bou assert _sent_paths(respx_mock) == ["converse-stream", "converse"] - @pytest.mark.asyncio @pytest.mark.parametrize( "stream_text", @@ -641,12 +674,15 @@ async def test_failed_stream_on_each_router_surface_reports_the_timeout_it_was_b ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) respx_mock.post(url__regex=_BEDROCK_RUNTIME_URL).mock(side_effect=slow_upstream_answer(upstream_answer)) - router: Final = _bedrock_router(model, timeout=UPSTREAM_FIRST_BYTE_SECONDS / 2, stream_timeout=STREAM_TIMEOUT_SECONDS) + router: Final = _bedrock_router( + model, timeout=UPSTREAM_FIRST_BYTE_SECONDS / 2, stream_timeout=STREAM_TIMEOUT_SECONDS + ) with pytest.raises(openai.APIError) as excinfo: await stream_text(router) assert getattr(excinfo.value, "timeout", None) == bounding_timeout + RECORDED_REASONING: Final = "Tokyo proper has about 14 million people." RECORDED_THINKING_TOKENS: Final = 309 JSON_MODE_THINKING_CONVERSE_RESPONSE: Final = { diff --git a/tests/unit/llms/custom_httpx/test_llm_http_handler.py b/tests/unit/llms/custom_httpx/test_llm_http_handler.py index 63cc45b29df..271ad3283c8 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -161,6 +161,37 @@ async def test_get_search_preserves_tinyfish_http_error_formatting(is_async: boo assert error.value.headers["retry-after"] == "7" +def test_a_caller_fake_stream_is_honored_and_kept_off_the_wire(respx_mock: respx.MockRouter) -> None: + route: Final = respx_mock.post("https://api.mistral.ai/v1/chat/completions").mock( + return_value=httpx.Response( + 200, + json={ + "id": "chatcmpl-fake-stream", + "object": "chat.completion", + "created": 1, + "model": "mistral-small-latest", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }, + ) + ) + + chunks: Final = list( + litellm.completion( + model="mistral/mistral-small-latest", + api_key="fake-mistral-key", + messages=[{"role": "user", "content": "hi"}], + stream=True, + fake_stream=True, + ) + ) + + sent: Final = json.loads(route.calls.last.request.content) + assert "fake_stream" not in sent, sent + assert sent.get("stream") is not True, sent + assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == "ok" + + def test_prepare_fake_stream_request(): # Initialize the BaseLLMHTTPHandler handler = BaseLLMHTTPHandler() diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index f5d13aecb6f..3c5d04a5f11 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -683,6 +683,68 @@ def test_embedding_keeps_an_internal_prefixed_kwarg_out_of_the_provider_request( assert sent["model"] == "text-embedding-3-small" +def test_embedding_keeps_deployment_prompt_management_params_out_of_the_bedrock_body( + respx_mock: respx.MockRouter, +) -> None: + mock_route: Final = respx_mock.post( + "https://bedrock-runtime.us-east-1.amazonaws.com/model/amazon.titan-embed-image-v1/invoke" + ).mock(return_value=httpx.Response(status_code=200, json={"embedding": [0.1, 0.2], "inputTextTokenCount": 1})) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "titan-image", + "litellm_params": { + "model": "bedrock/amazon.titan-embed-image-v1", + "api_key": "fake_bedrock_api_key", + "aws_region_name": "us-east-1", + "cache_control_injection_points": [{"location": "message", "role": "user"}], + "knowledge_bases": ["kb-1"], + "vector_store_ids": [], + }, + } + ] + ) + + response: Final = router.embedding(model="titan-image", input=["Help me !"]) + + assert json.loads(mock_route.calls.last.request.content) == {"inputText": "Help me !"} + assert response.data[0]["embedding"] == [0.1, 0.2] + + +def test_embedding_drops_root_prompt_management_params_but_keeps_explicit_extra_body( + respx_mock: respx.MockRouter, +) -> None: + api_base: Final = "http://localhost:12346/v1" + mock_route: Final = respx_mock.post(url__regex=rf"{api_base}/embeddings.*").mock( + return_value=httpx.Response( + status_code=200, + json={ + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2]}], + "model": "text-embedding-3-small", + "usage": {"prompt_tokens": 1, "total_tokens": 1}, + }, + ) + ) + + litellm.embedding( + model="openai/text-embedding-3-small", + input="hi", + api_base=api_base, + api_key="fake_openai_api_key", + cache_control_injection_points=[{"location": "message", "role": "user"}], + knowledge_bases=["kb-1"], + vector_store_ids=["root-store"], + extra_body={"vector_store_ids": ["caller-store"]}, + ) + + sent: Final = json.loads(mock_route.calls.last.request.content) + assert "cache_control_injection_points" not in sent, sent + assert "knowledge_bases" not in sent, sent + assert sent["vector_store_ids"] == ["caller-store"] + assert sent["input"] == "hi" + + def test_custom_provider_with_extra_headers(): with patch.object(litellm.llms.custom_httpx.http_handler.HTTPHandler, "post") as mock_post: diff --git a/tests/unit/types/test_litellm_params.py b/tests/unit/types/test_litellm_params.py index 788d405f99f..cb646f8ddf8 100644 --- a/tests/unit/types/test_litellm_params.py +++ b/tests/unit/types/test_litellm_params.py @@ -220,6 +220,7 @@ OPTION_NAMES: Final = ( "enable_json_schema_validation", "complete_response", "stream_chunk_size", + "fake_stream", "keepalive_seconds", "allow_client_keepalive_override", "mock_response", @@ -253,6 +254,9 @@ INTERNAL_STATE_NAMES: Final = ( "stream_response", "_context_compaction_state", *AGENTIC_LOOP_STATE_NAMES, + "_skip_mcp_handler", + "skip_mcp_handler", + "mcp_handler_context", "_router_weights", "fallback_depth", "max_fallbacks", @@ -552,7 +556,7 @@ LEAF_SAMPLES: Final[Mapping[type, Mapping[str, object]]] = { litellm_params.GuardrailOptions: {"guardrails": ("default",)}, litellm_params.PromptOptions: {"prompt_id": "prompt", "prompt_variables": {"name": "value"}}, litellm_params.ResponseOptions: {"keepalive_seconds": 1.5}, - litellm_params.ControlOptions: {"stream_chunk_size": 64}, + litellm_params.ControlOptions: {"stream_chunk_size": 64, "fake_stream": True}, litellm_params.MockOptions: {"mock_timeout": True}, litellm_params.CallState: { "completion_call_id": "call", @@ -560,6 +564,7 @@ LEAF_SAMPLES: Final[Mapping[type, Mapping[str, object]]] = { "data_residency": "us", }, litellm_params.AgenticLoopState: {"api_surface": "chat_completions", "depth": 1}, + litellm_params.McpHandlerState: {"recursion_guard": True, "mcp_handler_context": {"server": "docs"}}, litellm_params.RouterState: {"fallback_depth": 1}, litellm_params.ProxyRequestState: { "proxy_server_request": {"path": "/chat/completions"}, @@ -588,6 +593,7 @@ LEAF_BAD_SAMPLES: Final[Mapping[type, Mapping[str, object]]] = { litellm_params.MockOptions: {"mock_timeout": "true"}, litellm_params.CallState: {"completion_call_id": 1}, litellm_params.AgenticLoopState: {"depth": "1"}, + litellm_params.McpHandlerState: {"recursion_guard": "true"}, litellm_params.RouterState: {"fallback_depth": "1"}, litellm_params.ProxyRequestState: {"proxy_server_request": "request"}, litellm_params.EntrypointState: {"acompletion": "true"},