scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 199

Merge d4b7a6a494 into 6083cf9776

Browse source
This commit is contained in:
fangkang 2026-10-21 09:04:05 +09:00 • committed by GitHub
commit c61ffa7d17
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
16 changed files with 727 additions and 76 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

34
litellm/litellm_core_utils/core_helpers.py
View file

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

9
litellm/litellm_core_utils/get_litellm_params.py
View file

@ -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]}"
)

5
litellm/llms/bedrock/chat/converse_handler.py
View file

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

7
litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py
View file

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

3
litellm/llms/custom_httpx/llm_http_handler.py
View file

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

5
litellm/main.py
View file

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

10
litellm/types/litellm_params.py
View file

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

484
tests/integration/providers/test_internal_params_wire.py
View file

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

2
tests/unit/caching/test_caching.py
View file

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

67
tests/unit/litellm_core_utils/test_core_helpers.py
View file

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

9
tests/unit/litellm_core_utils/test_get_litellm_params.py
View file

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

17
tests/unit/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py
View file

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

50
tests/unit/llms/bedrock/chat/test_bedrock_converse_handler.py
View file

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

31
tests/unit/llms/custom_httpx/test_llm_http_handler.py
View file

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

62
tests/unit/test_main.py
View file

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

8
tests/unit/types/test_litellm_params.py
View file

@ -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"},