scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-12 23:01:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 87

chore(release): patch v1.89.0-rc.2 with batch-file auth, CrowdStrike AIDR, Mantle SigV4, NetApp streaming-cost fix, and team-scoped Datadog toward v1.89.0-rc.3 (#30179)

Browse source 
* fix(proxy): authorize batch files using upload target_model_names (LIT-3593) (#30009)

* fix(proxy): authorize batch files using upload target_model_names (LIT-3593)

After replace_model_in_jsonl, body.model is a stripped provider id. Reverse-mapping it via resolve_model_name_from_model_id is first-match on model_list and caused false 403s when multiple deployments share the same stripped name. Use target_model_names from the unified file id instead.

Co-authored-by: Cursor <cursoragent@cursor.com>

* fix(proxy): restore resolve_model_name_from_model_id for JSONL fallback path (LIT-3593)

Restores the reverse-lookup for the JSONL body.model fallback path so that
legacy/pre-target_model_names managed files still map stripped provider IDs
back to proxy aliases before auth. Also cleans up redundant `or None`.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

* Revert "fix(proxy): restore resolve_model_name_from_model_id for JSONL fallback path (LIT-3593)"

This reverts commit 30d2e96f77.

---------

Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
(cherry picked from commit 2cd7e87485)

* feat(guardrails): capture user and model metadata in CrowdStrike AIDR

(cherry picked from commit 6fc715c5bd)

* fix(guardrails): read CrowdStrike AIDR identity from both metadata bags (#29991)

Capture user_id and extra_info from metadata or litellm_metadata. The single-bag read dropped identity whenever a request carried a present litellm_metadata field (null or a user-supplied dict), since /chat/completions routes the authenticated identity into metadata while the guardrail read litellm_metadata first

(cherry picked from commit 1bbaf1c39d)

* feat(bedrock_mantle): add SigV4/IAM auth to Responses API route (#29788)

Applied as the squash diff of PR #29788 (head 9800b2f17c), which landed
upstream inside the litellm_oss_staging_080626 sync (32c88ca74f, #29932)
and has no standalone commit to cherry-pick. The rc line already carries
the prerequisite #29490 Responses route via the 040626 sync.

* fix: completion_cost AttributeError on streaming Anthropic web_search responses (#26153) (#27346)

Cherry-picked from staging squash 4a3860df1f.

The rc line predates the Usage.__init__ server_tool_use dict->ServerToolUse
coercion that staging carries (it landed via the squashed OSS sync #29932 /
32c88ca74f, not as a standalone commit). The calculate_usage
Usage(**returned_usage.model_dump()) round-trip re-serializes server_tool_use
to a plain dict, so without that coercion the rebuilt usage holds a dict and the
regression test asserting a ServerToolUse type fails. Restored the coercion in
litellm/types/utils.py to satisfy the prerequisite -- it matches #27346's own
first commit (coerce server_tool_use dict to ServerToolUse in Usage.__init__),
which was dropped from the squash only because staging already carried it.

* feat(datadog): add team-scoped Datadog callback support (#29947)

Cherry-picked from the PR head 9c049daa1b (single-commit PR, merged to
litellm_oss_branch). Applied cleanly; no conflicts.

Note: black --check in this worktree flags pre-existing multi-line string
formatting in litellm_core_utils/litellm_logging.py (lines ~1006-1050) that is
already present on the patch/v1.89.0-rc.1 base and is untouched by this pick --
left as-is to avoid reformatting unrelated lines.

---------

Co-authored-by: Sameer Kankute <sameer@berri.ai>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-authored-by: Kenan Yildirim <kenan@kenany.me>
Co-authored-by: yuneng-jiang <yuneng@berri.ai>
Co-authored-by: Kent <kingdooo@gmail.com>
Co-authored-by: ishaan-berri <155045088+ishaan-berri@users.noreply.github.com>
Co-authored-by: aanchal22 <12680748+aanchal22@users.noreply.github.com>
This commit is contained in:
Mateo Wang 2026-06-10 22:00:52 -07:00 • committed by GitHub
parent 2978a92bdb
commit fc344fd2a9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
23 changed files with 1948 additions and 95 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

69
litellm/integrations/datadog/datadog.py
View file

@ -92,12 +92,22 @@ class DataDogLogger(
# Class variables or attributes
def __init__(
self,
dd_api_key: Optional[str] = None,
dd_site: Optional[str] = None,
dd_agent_host: Optional[str] = None,
dd_agent_port: Optional[str] = None,
**kwargs,
):
"""
Initializes the datadog logger, checks if the correct env variables are set
Required environment variables (Direct API):
Args:
dd_api_key: Datadog API key. Falls back to DD_API_KEY env var.
dd_site: Datadog site (e.g. "us5.datadoghq.com"). Falls back to DD_SITE env var.
dd_agent_host: Hostname or IP of DataDog agent. Falls back to LITELLM_DD_AGENT_HOST env var.
dd_agent_port: Port of DataDog agent (default: 10518). Falls back to LITELLM_DD_AGENT_PORT env var.
Required environment variables (Direct API) when kwargs not provided:
`DD_API_KEY` - your datadog api key
`DD_SITE` - your datadog site, example = `"us5.datadoghq.com"`
@ -130,12 +140,19 @@ class DataDogLogger(
)
# Configure DataDog endpoint (Agent or Direct API)
# Use LITELLM_DD_AGENT_HOST to avoid conflicts with ddtrace's DD_AGENT_HOST
dd_agent_host = os.getenv("LITELLM_DD_AGENT_HOST")
if dd_agent_host:
self._configure_dd_agent(dd_agent_host=dd_agent_host)
# Prefer explicit kwargs, then fall back to env vars
resolved_agent_host = dd_agent_host or os.getenv("LITELLM_DD_AGENT_HOST")
if resolved_agent_host:
self._configure_dd_agent(
dd_agent_host=resolved_agent_host,
dd_agent_port=dd_agent_port,
dd_api_key=dd_api_key,
)
else:
self._configure_dd_direct_api()
self._configure_dd_direct_api(
dd_api_key=dd_api_key,
dd_site=dd_site,
)
# Optional override for testing
dd_base_url = get_datadog_base_url_from_env()
@ -172,34 +189,54 @@ class DataDogLogger(
).model_dump()
return dict_datadog_params
def _configure_dd_agent(self, dd_agent_host: str) -> None:
def _configure_dd_agent(
self,
dd_agent_host: str,
dd_agent_port: Optional[str] = None,
dd_api_key: Optional[str] = None,
) -> None:
"""
Configure DataDog Agent for log forwarding
Args:
dd_agent_host: Hostname or IP of DataDog agent
dd_agent_port: Port of DataDog agent. Falls back to LITELLM_DD_AGENT_PORT env var (default: 10518).
dd_api_key: Datadog API key. Falls back to DD_API_KEY env var. Optional when using agent.
"""
dd_agent_port = os.getenv(
resolved_port = dd_agent_port or os.getenv(
"LITELLM_DD_AGENT_PORT", "10518"
) # default port for logs
self.intake_url = f"http://{dd_agent_host}:{dd_agent_port}/api/v2/logs"
self.DD_API_KEY = os.getenv("DD_API_KEY") # Optional when using agent
self.intake_url = f"http://{dd_agent_host}:{resolved_port}/api/v2/logs"
self.DD_API_KEY = dd_api_key or os.getenv(
"DD_API_KEY"
) # Optional when using agent
verbose_logger.debug(f"Datadog: Using DD Agent at {self.intake_url}")
def _configure_dd_direct_api(self) -> None:
def _configure_dd_direct_api(
self,
dd_api_key: Optional[str] = None,
dd_site: Optional[str] = None,
) -> None:
"""
Configure direct DataDog API connection
Args:
dd_api_key: Datadog API key. Falls back to DD_API_KEY env var.
dd_site: Datadog site. Falls back to DD_SITE env var.
Raises:
Exception: If required environment variables are not set
Exception: If required credentials are not provided via args or env vars
"""
if os.getenv("DD_API_KEY", None) is None:
resolved_api_key = dd_api_key or os.getenv("DD_API_KEY")
resolved_site = dd_site or os.getenv("DD_SITE")
if resolved_api_key is None:
raise Exception("DD_API_KEY is not set, set 'DD_API_KEY=<>")
if os.getenv("DD_SITE", None) is None:
if resolved_site is None:
raise Exception("DD_SITE is not set in .env, set 'DD_SITE=<>")
self.DD_API_KEY = os.getenv("DD_API_KEY")
self.intake_url = f"https://http-intake.logs.{os.getenv('DD_SITE')}/api/v2/logs"
self.DD_API_KEY = resolved_api_key
self.intake_url = f"https://http-intake.logs.{resolved_site}/api/v2/logs"
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
"""

117
litellm/integrations/datadog/datadog_team_handler.py Normal file
View file

@ -0,0 +1,117 @@
"""
DataDog Team Handler
Used to get the DataDogLogger for a given request.
Handles Key/Team Based Datadog Logging, following the same pattern as LangFuseHandler.
"""
from typing import TYPE_CHECKING, Any, Dict, Optional, TypedDict
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.litellm_logging import StandardCallbackDynamicParams
from .datadog import DataDogLogger
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import DynamicLoggingCache
else:
DynamicLoggingCache = Any
class DatadogLoggingConfig(TypedDict):
dd_api_key: Optional[str]
dd_site: Optional[str]
dd_agent_host: Optional[str]
dd_agent_port: Optional[str]
class DataDogHandler:
@staticmethod
def get_datadog_logger_for_request(
standard_callback_dynamic_params: StandardCallbackDynamicParams,
in_memory_dynamic_logger_cache: DynamicLoggingCache,
) -> DataDogLogger:
"""
Get a team-scoped DataDogLogger for a given request.
Resolves and caches per-team DataDogLogger instances using DynamicLoggingCache,
keyed by the team's DD credentials. Each unique set of credentials gets its own
logger instance with its own batch/flush loop.
Note: This handler is only called when team-scoped DD credentials are present.
The global (env-var based) DataDogLogger is managed separately by
_init_custom_logger_compatible_class via _in_memory_loggers.
"""
_credentials = DataDogHandler.get_dynamic_datadog_logging_config(
standard_callback_dynamic_params=standard_callback_dynamic_params,
)
credentials_dict = dict(_credentials)
# check if datadog logger is already cached
temp_datadog_logger = in_memory_dynamic_logger_cache.get_cache(
credentials=credentials_dict, service_name="datadog"
)
# if not cached, create a new datadog logger and cache it
if temp_datadog_logger is None:
temp_datadog_logger = (
DataDogHandler._create_datadog_logger_from_credentials(
credentials=credentials_dict,
in_memory_dynamic_logger_cache=in_memory_dynamic_logger_cache,
)
)
return temp_datadog_logger
@staticmethod
def _create_datadog_logger_from_credentials(
credentials: Dict,
in_memory_dynamic_logger_cache: DynamicLoggingCache,
) -> DataDogLogger:
"""
Create a DataDogLogger from the credentials and cache it.
"""
datadog_logger = DataDogLogger(
dd_api_key=credentials.get("dd_api_key"),
dd_site=credentials.get("dd_site"),
dd_agent_host=credentials.get("dd_agent_host"),
dd_agent_port=credentials.get("dd_agent_port"),
)
in_memory_dynamic_logger_cache.set_cache(
credentials=credentials,
service_name="datadog",
logging_obj=datadog_logger,
)
verbose_logger.debug(
"Datadog: Created and cached new DataDogLogger for team-scoped credentials"
)
return datadog_logger
@staticmethod
def get_dynamic_datadog_logging_config(
standard_callback_dynamic_params: StandardCallbackDynamicParams,
) -> DatadogLoggingConfig:
"""
Get the Datadog logging config for a given request from dynamic params.
"""
return DatadogLoggingConfig(
dd_api_key=standard_callback_dynamic_params.get("dd_api_key"),
dd_site=standard_callback_dynamic_params.get("dd_site"),
dd_agent_host=standard_callback_dynamic_params.get("dd_agent_host"),
dd_agent_port=standard_callback_dynamic_params.get("dd_agent_port"),
)
@staticmethod
def _dynamic_datadog_credentials_are_passed(
standard_callback_dynamic_params: StandardCallbackDynamicParams,
) -> bool:
"""
Check if dynamic Datadog credentials are passed in standard_callback_dynamic_params.
"""
if (
standard_callback_dynamic_params.get("dd_api_key") is not None
or standard_callback_dynamic_params.get("dd_site") is not None
or standard_callback_dynamic_params.get("dd_agent_host") is not None
):
return True
return False

8
litellm/litellm_core_utils/initialize_dynamic_callback_params.py
View file

@ -53,11 +53,19 @@ _supported_callback_params = [
"braintrust_host",
"slack_webhook_url",
"lunary_public_key",
"dd_api_key",
"dd_site",
"dd_agent_host",
"dd_agent_port",
]
_request_blocked_callback_params = {
"gcs_bucket_name",
"gcs_path_service_account",
"dd_api_key",
"dd_site",
"dd_agent_host",
"dd_agent_port",
}

40
litellm/litellm_core_utils/litellm_logging.py
View file

@ -376,13 +376,14 @@ class Logging(LiteLLMLoggingBaseClass):
List[Union[str, Callable, CustomLogger]]
] = dynamic_async_failure_callbacks
# Process dynamic callbacks
self.process_dynamic_callbacks()
## DYNAMIC LANGFUSE / GCS / logging callback KEYS ##
self.standard_callback_dynamic_params: StandardCallbackDynamicParams = (
self.initialize_standard_callback_dynamic_params(kwargs)
)
# Process dynamic callbacks (after standard_callback_dynamic_params is initialized,
# so team-scoped credentials are available for callback initialization)
self.process_dynamic_callbacks()
self.standard_built_in_tools_params: StandardBuiltInToolsParams = (
self.initialize_standard_built_in_tools_params(kwargs)
)
@ -477,8 +478,21 @@ class Logging(LiteLLMLoggingBaseClass):
isinstance(callback, str)
and callback in litellm._known_custom_logger_compatible_callbacks
):
# For callbacks that support team-scoped credentials (e.g. datadog),
# pass only the relevant dynamic params as custom_logger_init_args.
_custom_logger_init_args: Optional[dict] = None
if callback == "datadog":
_custom_logger_init_args = {
k: v
for k, v in self.standard_callback_dynamic_params.items()
if k.startswith("dd_")
}
callback_class = _init_custom_logger_compatible_class(
callback, internal_usage_cache=None, llm_router=None # type: ignore
callback, # type: ignore[arg-type]
internal_usage_cache=None,
llm_router=None, # type: ignore
custom_logger_init_args=_custom_logger_init_args,
)
if callback_class is not None:
processed_list.append(callback_class)
@ -3890,6 +3904,24 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
_in_memory_loggers.append(_prometheus_logger)
return _prometheus_logger # type: ignore
elif logging_integration == "datadog":
# Check if team-scoped credentials are provided
_dd_api_key = custom_logger_init_args.get("dd_api_key")
_dd_site = custom_logger_init_args.get("dd_site")
_dd_agent_host = custom_logger_init_args.get("dd_agent_host")
_dd_agent_port = custom_logger_init_args.get("dd_agent_port")
if _dd_api_key or _dd_site or _dd_agent_host:
# Team-scoped credentials: use DynamicLoggingCache for per-credential isolation
from litellm.integrations.datadog.datadog_team_handler import (
DataDogHandler,
)
return DataDogHandler.get_datadog_logger_for_request(
standard_callback_dynamic_params=custom_logger_init_args, # type: ignore
in_memory_dynamic_logger_cache=in_memory_dynamic_logger_cache,
)
# Global (env-var based): reuse cached instance
for callback in _in_memory_loggers:
if isinstance(callback, DataDogLogger):
return callback # type: ignore

7
litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py
View file

@ -6,6 +6,7 @@ from typing import Any, Dict, List, Literal, Optional, Tuple
import litellm
from litellm.constants import OPENAI_FILE_SEARCH_COST_PER_1K_CALLS
from litellm.litellm_core_utils.llm_cost_calc.utils import _get_web_search_requests
from litellm.types.llms.openai import (
FileSearchTool,
ResponsesAPIResponse,
@ -339,8 +340,7 @@ class StandardBuiltInToolCostTracking:
# and _handle_web_search_cost() is never called.
if (
hasattr(usage, "server_tool_use")
and usage.server_tool_use is not None
and usage.server_tool_use.web_search_requests is not None
and _get_web_search_requests(usage.server_tool_use) is not None
):
return True
return False
@ -352,8 +352,7 @@ class StandardBuiltInToolCostTracking:
elif usage is not None:
if (
hasattr(usage, "server_tool_use")
and usage.server_tool_use is not None
and usage.server_tool_use.web_search_requests is not None
and _get_web_search_requests(usage.server_tool_use) is not None
):
return True
elif (

22
litellm/litellm_core_utils/llm_cost_calc/utils.py
View file

@ -1,7 +1,7 @@
# What is this?
## Helper utilities for cost_per_token()
from typing import Literal, Optional, Tuple, TypedDict, cast
from typing import Any, Literal, Optional, Tuple, TypedDict, cast
import litellm
from litellm._logging import verbose_logger
@ -42,6 +42,26 @@ def _get_token_detail_value(details: object, key: str) -> Optional[int]:
return value if isinstance(value, int) else None
def _get_web_search_requests(server_tool_use: Any) -> Optional[int]:
"""
Tolerantly read ``web_search_requests`` from a ``server_tool_use`` value
that may be ``None``, a ``dict``, a ``ServerToolUse`` pydantic instance,
or any other object supporting attribute access.
Returns ``None`` when the value cannot be resolved — callers can
distinguish "absent" from "zero" using ``is None``.
See https://github.com/BerriAI/litellm/issues/26153 — ``stream_chunk_builder``
historically left this as a plain ``dict``, which broke direct attribute
access in cost calculation.
"""
if server_tool_use is None:
return None
if isinstance(server_tool_use, dict):
return server_tool_use.get("web_search_requests")
return getattr(server_tool_use, "web_search_requests", None)
def _is_above_128k(tokens: float) -> bool:
if tokens > 128000:
return True

13
litellm/litellm_core_utils/streaming_chunk_builder_utils.py
View file

@ -588,7 +588,18 @@ class ChunkProcessor:
hasattr(usage_chunk, "server_tool_use")
and usage_chunk.server_tool_use is not None
):
server_tool_use = usage_chunk.server_tool_use
# Coerce dict to ServerToolUse so downstream cost-calc code
# (which accesses .web_search_requests as an attribute)
# doesn't raise AttributeError. Some providers / streaming
# paths leave server_tool_use as a plain dict on the chunk.
if isinstance(usage_chunk.server_tool_use, dict):
server_tool_use = ServerToolUse(**usage_chunk.server_tool_use)
elif isinstance(usage_chunk.server_tool_use, ServerToolUse):
server_tool_use = usage_chunk.server_tool_use
else:
server_tool_use = ServerToolUse.model_validate(
usage_chunk.server_tool_use
)
if (
usage_chunk_dict["prompt_tokens_details"] is not None
and getattr(

14
litellm/llms/anthropic/cost_calculation.py
View file

@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Optional, Tuple
from litellm.litellm_core_utils.llm_cost_calc.utils import (
_get_token_base_cost,
_get_web_search_requests,
_parse_prompt_tokens_details,
calculate_cache_writing_cost,
generic_cost_per_token,
@ -110,11 +111,12 @@ def get_cost_for_anthropic_web_search(
if model_info is None:
return 0.0
if (
usage is None
or usage.server_tool_use is None
or usage.server_tool_use.web_search_requests is None
):
if usage is None:
return 0.0
web_search_requests = _get_web_search_requests(
getattr(usage, "server_tool_use", None)
)
if web_search_requests is None:
return 0.0
## Get the cost per web search request
@ -128,5 +130,5 @@ def get_cost_for_anthropic_web_search(
return 0.0
## Calculate the total cost
total_cost = cost_per_web_search_request * usage.server_tool_use.web_search_requests
total_cost = cost_per_web_search_request * web_search_requests
return total_cost

20
litellm/llms/base_llm/responses/transformation.py
View file

@ -62,6 +62,26 @@ class BaseResponsesAPIConfig(ABC):
"""
return False
def sign_request(
self,
headers: dict,
optional_params: dict,
request_data: dict,
api_base: str,
api_key: Optional[str] = None,
model: Optional[str] = None,
stream: Optional[bool] = None,
fake_stream: Optional[bool] = None,
) -> Tuple[dict, Optional[bytes]]:
"""Sign the request after the body is finalized.
Default is a no-op (returns headers unchanged, no signed body). Providers
whose endpoint requires request signing (e.g. Bedrock Mantle SigV4)
override this and return the signed body bytes so the handler sends those
exact bytes.
"""
return headers, None
@abstractmethod
def get_supported_openai_params(self, model: str) -> list:
pass

120
litellm/llms/bedrock_mantle/responses/transformation.py
View file

@ -4,14 +4,26 @@ Amazon Bedrock Mantle - Responses API backend.
gpt-5.5 / gpt-5.4 on Mantle are exposed ONLY on the `/openai/v1/responses`
path (not the standard `/v1/responses`). Payloads and SSE follow the OpenAI
Responses spec, so this config inherits OpenAIResponsesAPIConfig and overrides
only the endpoint URL and Bearer authentication.
only the endpoint URL and authentication.
Auth: AWS Bedrock API key as Bearer token (BEDROCK_MANTLE_API_KEY or the
standard AWS_BEARER_TOKEN_BEDROCK), NOT SigV4.
Auth: Bearer token (BEDROCK_MANTLE_API_KEY or the standard
AWS_BEARER_TOKEN_BEDROCK, or litellm_params.api_key) when present; otherwise
AWS SigV4 (service name "bedrock") using the standard credential chain (IAM
role / access key / profile / web identity), signed via the shared
BaseAWSLLM._sign_request after the request body is finalized.
"""
from typing import Optional
import re
from typing import Optional, Tuple
from botocore.exceptions import (
CredentialRetrievalError,
NoCredentialsError,
PartialCredentialsError,
ProfileNotFound,
)
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
from litellm.secret_managers.main import get_secret_str
from litellm.types.router import GenericLiteLLMParams
@ -29,22 +41,44 @@ _BASE_SUFFIXES_TO_STRIP = (
"/v1",
)
# Standard Mantle host: https://bedrock-mantle.<region>.api.aws (group 1 = region).
_MANTLE_HOST_RE = re.compile(
r"^https?://bedrock-mantle\.([^/.]+)\.api\.aws", re.IGNORECASE
)
class BedrockMantleResponsesAPIConfig(OpenAIResponsesAPIConfig):
def __init__(self, aws_signer: Optional[BaseAWSLLM] = None):
super().__init__()
self._aws_signer = aws_signer or BaseAWSLLM()
@property
def custom_llm_provider(self) -> LlmProviders:
return LlmProviders.BEDROCK_MANTLE
@staticmethod
def _resolve_region(params: dict) -> str:
region = params.get("aws_region_name")
if region:
return region
base = params.get("api_base") or get_secret_str("BEDROCK_MANTLE_API_BASE")
if base:
match = _MANTLE_HOST_RE.match(base.rstrip("/"))
if match:
return match.group(1)
return (
get_secret_str("BEDROCK_MANTLE_REGION")
or get_secret_str("AWS_REGION_NAME")
or get_secret_str("AWS_REGION")
or BEDROCK_MANTLE_DEFAULT_REGION
)
def get_complete_url(
self,
api_base: Optional[str],
litellm_params: dict,
) -> str:
region = (
get_secret_str("BEDROCK_MANTLE_REGION")
or get_secret_str("AWS_REGION")
or BEDROCK_MANTLE_DEFAULT_REGION
)
region = self._resolve_region({**litellm_params, "api_base": api_base})
base = (
api_base
or get_secret_str("BEDROCK_MANTLE_API_BASE")
@ -55,6 +89,11 @@ class BedrockMantleResponsesAPIConfig(OpenAIResponsesAPIConfig):
if base.endswith(suffix):
base = base[: -len(suffix)]
break
# For the standard Mantle host (including the default-region base that
# responses/main.py auto-injects into litellm_params.api_base), pin to the
# single resolved region so aws_region_name wins; preserve custom proxy hosts.
if _MANTLE_HOST_RE.match(base):
base = f"https://bedrock-mantle.{region}.api.aws"
return f"{base}/openai/v1/responses"
def validate_environment(
@ -66,12 +105,8 @@ class BedrockMantleResponsesAPIConfig(OpenAIResponsesAPIConfig):
or get_secret_str("BEDROCK_MANTLE_API_KEY")
or get_secret_str("AWS_BEARER_TOKEN_BEDROCK")
)
if not api_key:
raise ValueError(
"Bedrock Mantle API key is required. Set BEDROCK_MANTLE_API_KEY "
"(or AWS_BEARER_TOKEN_BEDROCK) or pass api_key."
)
headers["Authorization"] = f"Bearer {api_key}"
if api_key:
headers["Authorization"] = f"Bearer {api_key}"
return headers
def supports_native_file_search(self) -> bool:
@ -79,3 +114,58 @@ class BedrockMantleResponsesAPIConfig(OpenAIResponsesAPIConfig):
def supports_native_websocket(self) -> bool:
return False
def sign_request(
self,
headers: dict,
optional_params: dict,
request_data: dict,
api_base: str,
api_key: Optional[str] = None,
model: Optional[str] = None,
stream: Optional[bool] = None,
fake_stream: Optional[bool] = None,
) -> Tuple[dict, Optional[bytes]]:
bearer = (
api_key
or get_secret_str("BEDROCK_MANTLE_API_KEY")
or get_secret_str("AWS_BEARER_TOKEN_BEDROCK")
)
if not bearer:
# SigV4 path. Pin the credential-scope region to the region of the actual
# signing URL (api_base, already region-resolved by get_complete_url) so the
# SigV4 scope and the URL host can never disagree. Resolve from api_base first,
# then fall back to the regular precedence. Also drop any caller Authorization
# so _sign_request's restore-original-Authorization step cannot override the
# SigV4 header.
optional_params = {
**optional_params,
"aws_region_name": self._resolve_region(
{**optional_params, "api_base": api_base}
),
}
headers = {k: v for k, v in headers.items() if k.lower() != "authorization"}
try:
return self._aws_signer._sign_request(
service_name="bedrock",
headers=headers,
optional_params=optional_params,
request_data=request_data,
api_base=api_base,
api_key=bearer,
model=model,
stream=stream,
fake_stream=fake_stream,
)
except (
NoCredentialsError,
PartialCredentialsError,
ProfileNotFound,
CredentialRetrievalError,
) as e:
raise ValueError(
"Bedrock Mantle auth failed: no Bearer token and no usable AWS "
"credentials. Set BEDROCK_MANTLE_API_KEY (or AWS_BEARER_TOKEN_BEDROCK) "
"or pass api_key for Bearer auth, or provide AWS credentials "
"(IAM role / access key / profile / web identity) for SigV4."
) from e

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

@ -2318,6 +2318,31 @@ class BaseLLMHTTPHandler:
# but never included in the outbound provider payload.
request_context["litellm_params"] = dict(litellm_params)
is_stream_request = bool(stream)
if is_stream_request and fake_stream is True:
stream, data = self._prepare_fake_stream_request(
stream=stream,
data=data,
fake_stream=fake_stream,
)
# Sign after the body is final (post-transform/normalize/extra_body and post
# fake-stream prep) so signed bytes match what we send. No-op for providers
# that inherit the default sign_request.
headers, signed_body = responses_api_provider_config.sign_request(
headers=headers,
optional_params=dict(litellm_params),
request_data=data,
api_base=api_base,
api_key=litellm_params.api_key,
model=model,
stream=stream,
fake_stream=fake_stream,
)
body_kwargs: Dict[str, Any] = (
{"data": signed_body} if signed_body is not None else {"json": data}
)
## LOGGING
logging_obj.pre_call(
input=input,
@ -2330,22 +2355,14 @@ class BaseLLMHTTPHandler:
)
try:
if stream:
# For streaming, use stream=True in the request
if fake_stream is True:
stream, data = self._prepare_fake_stream_request(
stream=stream,
data=data,
fake_stream=fake_stream,
)
if is_stream_request:
response = sync_httpx_client.post(
url=api_base,
headers=headers,
json=data,
timeout=timeout
or float(response_api_optional_request_params.get("timeout", 0)),
stream=stream,
**body_kwargs,
)
if fake_stream is True:
return MockResponsesAPIStreamingIterator(
@ -2370,13 +2387,12 @@ class BaseLLMHTTPHandler:
call_type=CallTypes.responses.value,
)
else:
# For non-streaming requests
response = sync_httpx_client.post(
url=api_base,
headers=headers,
json=data,
timeout=timeout
or float(response_api_optional_request_params.get("timeout", 0)),
**body_kwargs,
)
except Exception as e:
raise self._handle_error(
@ -2464,6 +2480,28 @@ class BaseLLMHTTPHandler:
# but never included in the outbound provider payload.
request_context["litellm_params"] = dict(litellm_params)
is_stream_request = bool(stream)
if is_stream_request and fake_stream is True:
stream, data = self._prepare_fake_stream_request(
stream=stream,
data=data,
fake_stream=fake_stream,
)
headers, signed_body = responses_api_provider_config.sign_request(
headers=headers,
optional_params=dict(litellm_params),
request_data=data,
api_base=api_base,
api_key=litellm_params.api_key,
model=model,
stream=stream,
fake_stream=fake_stream,
)
body_kwargs: Dict[str, Any] = (
{"data": signed_body} if signed_body is not None else {"json": data}
)
## LOGGING
logging_obj.pre_call(
input=input,
@ -2476,22 +2514,14 @@ class BaseLLMHTTPHandler:
)
try:
if stream:
# For streaming, we need to use stream=True in the request
if fake_stream is True:
stream, data = self._prepare_fake_stream_request(
stream=stream,
data=data,
fake_stream=fake_stream,
)
if is_stream_request:
response = await async_httpx_client.post(
url=api_base,
headers=headers,
json=data,
timeout=timeout
or float(response_api_optional_request_params.get("timeout", 0)),
stream=stream,
**body_kwargs,
)
if fake_stream is True:
@ -2518,13 +2548,12 @@ class BaseLLMHTTPHandler:
call_type=CallTypes.responses.value,
)
else:
# For non-streaming, proceed as before
response = await async_httpx_client.post(
url=api_base,
headers=headers,
json=data,
timeout=timeout
or float(response_api_optional_request_params.get("timeout", 0)),
**body_kwargs,
)
except Exception as e:
@ -4005,6 +4034,18 @@ class BaseLLMHTTPHandler:
)
data = BaseResponsesAPIConfig.normalize_responses_api_request_dict(data)
headers, signed_body = responses_api_provider_config.sign_request(
headers=headers,
optional_params=dict(litellm_params),
request_data=data,
api_base=url,
api_key=litellm_params.api_key,
model=model,
)
body_kwargs: Dict[str, Any] = (
{"data": signed_body} if signed_body is not None else {"json": data}
)
## LOGGING
logging_obj.pre_call(
input=input,
@ -4018,7 +4059,7 @@ class BaseLLMHTTPHandler:
try:
response = sync_httpx_client.post(
url=url, headers=headers, json=data, timeout=timeout
url=url, headers=headers, timeout=timeout, **body_kwargs
)
except Exception as e:
@ -4088,6 +4129,18 @@ class BaseLLMHTTPHandler:
)
data = BaseResponsesAPIConfig.normalize_responses_api_request_dict(data)
headers, signed_body = responses_api_provider_config.sign_request(
headers=headers,
optional_params=dict(litellm_params),
request_data=data,
api_base=url,
api_key=litellm_params.api_key,
model=model,
)
body_kwargs: Dict[str, Any] = (
{"data": signed_body} if signed_body is not None else {"json": data}
)
## LOGGING
logging_obj.pre_call(
input=input,
@ -4101,7 +4154,7 @@ class BaseLLMHTTPHandler:
try:
response = await async_httpx_client.post(
url=url, headers=headers, json=data, timeout=timeout
url=url, headers=headers, timeout=timeout, **body_kwargs
)
except Exception as e:

28
litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py
View file

@ -105,6 +105,16 @@ def _extract_text_from_content(content: object) -> str:
return ""
def _merge_metadata_bags(request_data: Mapping[str, Any]) -> Optional[dict[str, Any]]:
merged: dict[str, Any] = {}
present = False
for bag in (request_data.get("metadata"), request_data.get("litellm_metadata")):
if isinstance(bag, Mapping):
present = True
merged.update(bag)
return merged if present else None
class CrowdStrikeAIDRHandler(CustomGuardrail):
"""
CrowdStrike AIDR AI Guardrail handler to interact with the CrowdStrike AIDR
@ -312,11 +322,27 @@ class CrowdStrikeAIDRHandler(CustomGuardrail):
event_type = "output"
hook_name = "apply_guardrail (response)"
ai_guard_payload = {
ai_guard_payload: dict[str, Any] = {
"guard_input": guard_input.model_dump(mode="json"),
"event_type": event_type,
}
model = inputs.get("model")
if model:
ai_guard_payload["model"] = model
metadata = _merge_metadata_bags(request_data)
if metadata is not None:
user_id = metadata.get("user_api_key_user_id")
if user_id:
ai_guard_payload["user_id"] = user_id
extra_info: dict[str, str] = {}
user_email = metadata.get("user_api_key_user_email")
if user_email:
extra_info["user_name"] = user_email
ai_guard_payload["extra_info"] = extra_info
ai_guard_response = await self._call_crowdstrike_aidr_guard(
ai_guard_payload, hook_name
)

29
litellm/proxy/hooks/batch_rate_limiter.py
View file

@ -518,11 +518,17 @@ class _PROXY_BatchRateLimiter(CustomLogger):
# Check if this is a managed file (base64 encoded unified file ID)
from litellm.proxy.openai_files_endpoints.common_utils import (
_is_base64_encoded_unified_file_id,
get_models_from_unified_file_id,
)
# Managed files require bypassing the HTTP endpoint (which runs access-check hooks)
# and calling the managed files hook directly with the user's credentials.
is_managed_file = _is_base64_encoded_unified_file_id(file_id)
target_model_names = (
get_models_from_unified_file_id(is_managed_file)
if is_managed_file
else []
)
if is_managed_file and user_api_key_dict is not None:
file_content = await self._fetch_managed_file_content(
file_id=file_id,
@ -560,6 +566,7 @@ class _PROXY_BatchRateLimiter(CustomLogger):
await self._enforce_batch_file_model_access(
user_api_key_dict=user_api_key_dict,
file_content_as_dict=file_content_as_dict,
target_model_names=target_model_names or None,
)
input_file_usage = _get_batch_job_input_file_usage(
@ -595,9 +602,13 @@ class _PROXY_BatchRateLimiter(CustomLogger):
self,
user_api_key_dict: UserAPIKeyAuth,
file_content_as_dict: List[dict],
target_model_names: Optional[List[str]] = None,
) -> None:
"""Reject the batch if the caller is not authorized for every
``body.model`` named inside the JSONL.
"""Reject the batch if the caller is not authorized for the upload target.
For managed files, ``target_model_names`` (from the unified file id) is
the proxy alias the file was uploaded for and is used directly for auth.
For legacy/non-managed files, falls back to ``body.model`` values in the JSONL.
Reuses standard auth helpers so the same model access rules the proxy
enforces on `/chat/completions` apply here.
@ -614,9 +625,12 @@ class _PROXY_BatchRateLimiter(CustomLogger):
from litellm.proxy.proxy_server import proxy_logging_obj
from litellm.proxy.proxy_server import user_api_key_cache
models = _get_models_from_batch_input_file_content(file_content_as_dict)
if not models:
return
if target_model_names:
models = target_model_names
else:
models = _get_models_from_batch_input_file_content(file_content_as_dict)
if not models:
return
team_object = None
if (
@ -647,12 +661,7 @@ class _PROXY_BatchRateLimiter(CustomLogger):
llm_model_list = llm_router.model_list if llm_router is not None else None
for model in models:
# body.model may be the provider id after replace_model_in_jsonl; map to proxy model_name for auth.
model_to_check = model
if llm_router is not None:
proxy_model_name = llm_router.resolve_model_name_from_model_id(model)
if proxy_model_name is not None:
model_to_check = proxy_model_name
try:
if team_object is not None:
try:

11
litellm/types/utils.py
View file

@ -1571,7 +1571,7 @@ class Usage(SafeAttributeModel, CompletionUsage):
completion_tokens_details: Optional[
Union[CompletionTokensDetailsWrapper, dict]
] = None,
server_tool_use: Optional[ServerToolUse] = None,
server_tool_use: Optional[Union[ServerToolUse, dict]] = None,
cost: Optional[float] = None,
**params,
):
@ -1672,6 +1672,9 @@ class Usage(SafeAttributeModel, CompletionUsage):
prompt_tokens_details=_prompt_tokens_details or None,
)
if isinstance(server_tool_use, dict):
server_tool_use = ServerToolUse(**server_tool_use)
if server_tool_use is not None:
self.server_tool_use = server_tool_use
else: # maintain openai compatibility in usage object if possible
@ -3026,6 +3029,12 @@ class StandardCallbackDynamicParams(TypedDict, total=False):
wandb_api_key: Optional[str]
weave_project_id: Optional[str]
# Datadog dynamic params
dd_api_key: Optional[str]
dd_site: Optional[str]
dd_agent_host: Optional[str]
dd_agent_port: Optional[str]
# Logging settings
turn_off_message_logging: Optional[bool] # when true will not log messages
litellm_disabled_callbacks: Optional[List[str]]

194
tests/test_litellm/integrations/datadog/test_datadog_team_handler.py Normal file
View file

@ -0,0 +1,194 @@
"""
Tests for team-scoped Datadog callback support.
Verifies that DataDogLogger can be instantiated with per-team credentials
(dd_api_key, dd_site) instead of relying solely on environment variables,
and that the DataDogHandler correctly resolves and caches per-team loggers.
"""
from unittest.mock import patch
import pytest
from litellm.integrations.datadog.datadog import DataDogLogger
from litellm.integrations.datadog.datadog_team_handler import (
DataDogHandler,
DatadogLoggingConfig,
)
from litellm.litellm_core_utils.specialty_caches.dynamic_logging_cache import (
DynamicLoggingCache,
)
from litellm.types.utils import StandardCallbackDynamicParams
@pytest.fixture
def datadog_env(monkeypatch):
"""Set global DD env vars for the default/global logger."""
monkeypatch.setenv("DD_API_KEY", "global_api_key")
monkeypatch.setenv("DD_SITE", "us1.datadoghq.com")
class TestDataDogLoggerCredentialKwargs:
"""Test that DataDogLogger accepts credentials as kwargs."""
def test_init_with_explicit_credentials(self):
"""Logger should use explicit kwargs instead of env vars."""
with patch("asyncio.create_task"):
logger = DataDogLogger(
dd_api_key="team_api_key",
dd_site="eu1.datadoghq.com",
)
assert logger.DD_API_KEY == "team_api_key"
assert "eu1.datadoghq.com" in logger.intake_url
def test_init_falls_back_to_env_vars(self, datadog_env):
"""Logger should fall back to env vars when no kwargs provided."""
with patch("asyncio.create_task"):
logger = DataDogLogger()
assert logger.DD_API_KEY == "global_api_key"
assert "us1.datadoghq.com" in logger.intake_url
def test_init_kwargs_override_env_vars(self, datadog_env):
"""Explicit kwargs should take precedence over env vars."""
with patch("asyncio.create_task"):
logger = DataDogLogger(
dd_api_key="override_key",
dd_site="ap1.datadoghq.com",
)
assert logger.DD_API_KEY == "override_key"
assert "ap1.datadoghq.com" in logger.intake_url
def test_init_with_agent_credentials(self):
"""Logger should use agent mode when dd_agent_host is provided."""
with patch("asyncio.create_task"):
logger = DataDogLogger(
dd_agent_host="dd-agent.local",
dd_agent_port="8125",
dd_api_key="agent_api_key",
)
assert "dd-agent.local:8125" in logger.intake_url
assert logger.DD_API_KEY == "agent_api_key"
def test_init_raises_without_credentials(self, monkeypatch):
"""Logger should raise if no credentials are available."""
monkeypatch.delenv("DD_API_KEY", raising=False)
monkeypatch.delenv("DD_SITE", raising=False)
monkeypatch.delenv("LITELLM_DD_AGENT_HOST", raising=False)
with pytest.raises(Exception, match="DD_API_KEY"):
with patch("asyncio.create_task"):
DataDogLogger()
class TestDataDogHandler:
"""Test that DataDogHandler resolves the correct logger per team."""
def test_creates_team_logger_with_dynamic_credentials(self, datadog_env):
"""Should create a new logger when team credentials are provided."""
cache = DynamicLoggingCache()
params = StandardCallbackDynamicParams(
dd_api_key="team_a_key",
dd_site="eu1.datadoghq.com",
)
with patch("asyncio.create_task"):
result = DataDogHandler.get_datadog_logger_for_request(
standard_callback_dynamic_params=params,
in_memory_dynamic_logger_cache=cache,
)
assert result.DD_API_KEY == "team_a_key"
assert "eu1.datadoghq.com" in result.intake_url
def test_caches_team_logger(self, datadog_env):
"""Same team credentials should return the same cached logger instance."""
cache = DynamicLoggingCache()
params = StandardCallbackDynamicParams(
dd_api_key="team_b_key",
dd_site="us5.datadoghq.com",
)
with patch("asyncio.create_task"):
result1 = DataDogHandler.get_datadog_logger_for_request(
standard_callback_dynamic_params=params,
in_memory_dynamic_logger_cache=cache,
)
result2 = DataDogHandler.get_datadog_logger_for_request(
standard_callback_dynamic_params=params,
in_memory_dynamic_logger_cache=cache,
)
assert result1 is result2
def test_different_teams_get_different_loggers(self, datadog_env):
"""Different team credentials should create separate logger instances."""
cache = DynamicLoggingCache()
params_a = StandardCallbackDynamicParams(
dd_api_key="team_a_key",
dd_site="us1.datadoghq.com",
)
params_b = StandardCallbackDynamicParams(
dd_api_key="team_b_key",
dd_site="eu1.datadoghq.com",
)
with patch("asyncio.create_task"):
result_a = DataDogHandler.get_datadog_logger_for_request(
standard_callback_dynamic_params=params_a,
in_memory_dynamic_logger_cache=cache,
)
result_b = DataDogHandler.get_datadog_logger_for_request(
standard_callback_dynamic_params=params_b,
in_memory_dynamic_logger_cache=cache,
)
assert result_a is not result_b
assert result_a.DD_API_KEY == "team_a_key"
assert result_b.DD_API_KEY == "team_b_key"
def test_request_blocked_callback_params_includes_dd(self):
"""DD params should be blocked from request-level metadata (security)."""
from litellm.litellm_core_utils.initialize_dynamic_callback_params import (
_request_blocked_callback_params,
)
assert "dd_api_key" in _request_blocked_callback_params
assert "dd_site" in _request_blocked_callback_params
assert "dd_agent_host" in _request_blocked_callback_params
assert "dd_agent_port" in _request_blocked_callback_params
class TestDynamicCredentialDetection:
"""Test that _dynamic_datadog_credentials_are_passed works correctly."""
def test_no_credentials(self):
params = StandardCallbackDynamicParams()
assert DataDogHandler._dynamic_datadog_credentials_are_passed(params) is False
def test_dd_api_key_only(self):
params = StandardCallbackDynamicParams(dd_api_key="key")
assert DataDogHandler._dynamic_datadog_credentials_are_passed(params) is True
def test_dd_site_only(self):
params = StandardCallbackDynamicParams(dd_site="site")
assert DataDogHandler._dynamic_datadog_credentials_are_passed(params) is True
def test_dd_agent_host_only(self):
params = StandardCallbackDynamicParams(dd_agent_host="host")
assert DataDogHandler._dynamic_datadog_credentials_are_passed(params) is True
class TestStandardCallbackDynamicParamsIncludesDatadog:
"""Verify that Datadog params are in the allow-list."""
def test_dd_params_in_annotations(self):
annotations = StandardCallbackDynamicParams.__annotations__
assert "dd_api_key" in annotations
assert "dd_site" in annotations
assert "dd_agent_host" in annotations
assert "dd_agent_port" in annotations

88
tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking_dict_safety.py Normal file
View file

@ -0,0 +1,88 @@
"""
Tests that the cost-tracking call sites tolerate ``server_tool_use`` being
either a ``dict`` or a ``ServerToolUse`` pydantic instance.
See https://github.com/BerriAI/litellm/issues/26153.
"""
import os
import sys
import pytest
sys.path.insert(0, os.path.abspath("../../../.."))
from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import (
StandardBuiltInToolCostTracking,
_get_web_search_requests,
)
from litellm.types.utils import ModelResponse, ServerToolUse, Usage
class _UsageWithDictServerToolUse:
"""
Tiny stand-in that mimics the broken streaming-rebuild shape:
``server_tool_use`` is a plain dict.
"""
def __init__(self, server_tool_use):
self.server_tool_use = server_tool_use
self.prompt_tokens_details = None
def test_get_web_search_requests_handles_none():
assert _get_web_search_requests(None) is None
def test_get_web_search_requests_handles_dict():
assert _get_web_search_requests({"web_search_requests": 5}) == 5
def test_get_web_search_requests_handles_dict_missing_key():
assert _get_web_search_requests({}) is None
def test_get_web_search_requests_handles_pydantic():
stu = ServerToolUse(web_search_requests=7)
assert _get_web_search_requests(stu) == 7
def test_get_web_search_requests_handles_pydantic_with_none_value():
stu = ServerToolUse()
assert _get_web_search_requests(stu) is None
def test_response_object_includes_web_search_call_with_dict_server_tool_use():
"""
The exact bug: ``usage.server_tool_use`` is a dict and the check in
``response_object_includes_web_search_call`` used to crash with
``AttributeError``.
"""
response = ModelResponse()
usage = _UsageWithDictServerToolUse({"web_search_requests": 2})
# Must not raise — and must correctly detect the web search call.
result = StandardBuiltInToolCostTracking.response_object_includes_web_search_call(
response_object=response, usage=usage # type: ignore[arg-type]
)
assert result is True
def test_response_object_includes_web_search_call_with_pydantic_server_tool_use():
response = ModelResponse()
usage = _UsageWithDictServerToolUse(ServerToolUse(web_search_requests=2))
result = StandardBuiltInToolCostTracking.response_object_includes_web_search_call(
response_object=response, usage=usage # type: ignore[arg-type]
)
assert result is True
def test_response_object_includes_web_search_call_with_none_server_tool_use():
response = ModelResponse()
usage = _UsageWithDictServerToolUse(None)
result = StandardBuiltInToolCostTracking.response_object_includes_web_search_call(
response_object=response, usage=usage # type: ignore[arg-type]
)
assert result is False

130
tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_server_tool_use.py Normal file
View file

@ -0,0 +1,130 @@
"""
Regression tests for https://github.com/BerriAI/litellm/issues/26153
``stream_chunk_builder`` used to leave ``usage.server_tool_use`` as a plain
``dict`` when reconstructing a streaming response. Downstream cost-calculation
code (``StandardBuiltInToolCostTracking.response_object_includes_web_search_call``
and ``get_cost_for_anthropic_web_search``) accesses
``usage.server_tool_use.web_search_requests`` as an attribute, which raised
``AttributeError: 'dict' object has no attribute 'web_search_requests'``.
These tests reconstruct streaming chunks for an Anthropic-style web_search
response and assert:
1. ``stream_chunk_builder`` returns ``ServerToolUse`` (not ``dict``) for
``usage.server_tool_use``.
2. ``completion_cost`` runs end-to-end on the rebuilt response without
raising ``AttributeError``.
"""
import os
import sys
import pytest
sys.path.insert(0, os.path.abspath("../../.."))
from litellm import completion_cost, stream_chunk_builder
from litellm.types.utils import (
Delta,
ModelResponseStream,
ServerToolUse,
StreamingChoices,
Usage,
)
def _make_text_chunk(text: str) -> ModelResponseStream:
return ModelResponseStream(
id="chatcmpl-test-26153",
created=1700000000,
model="claude-3-haiku-20240307",
object="chat.completion.chunk",
choices=[
StreamingChoices(
finish_reason=None,
index=0,
delta=Delta(role="assistant", content=text),
)
],
)
def _make_finish_chunk_with_usage_dict_server_tool_use() -> ModelResponseStream:
"""Final chunk where server_tool_use is a *dict* — reproduces the bug shape."""
return ModelResponseStream(
id="chatcmpl-test-26153",
created=1700000000,
model="claude-3-haiku-20240307",
object="chat.completion.chunk",
choices=[
StreamingChoices(
finish_reason="stop",
index=0,
delta=Delta(),
)
],
usage=Usage(
prompt_tokens=42,
completion_tokens=11,
total_tokens=53,
# NOTE: passed as a dict on purpose — this is the shape that
# historically slipped through stream_chunk_builder unchanged.
server_tool_use={"web_search_requests": 3},
),
)
def test_stream_chunk_builder_coerces_server_tool_use_to_pydantic():
"""
Regression: stream_chunk_builder must produce ServerToolUse, not dict.
"""
chunks = [
_make_text_chunk("Otters "),
_make_text_chunk("are great."),
_make_finish_chunk_with_usage_dict_server_tool_use(),
]
rebuilt = stream_chunk_builder(chunks)
assert rebuilt is not None
assert rebuilt.usage is not None # type: ignore[attr-defined]
server_tool_use = rebuilt.usage.server_tool_use # type: ignore[attr-defined]
assert (
server_tool_use is not None
), "server_tool_use should be carried through from the final chunk"
assert isinstance(server_tool_use, ServerToolUse), (
f"expected ServerToolUse, got {type(server_tool_use).__name__}: "
f"{server_tool_use!r}"
)
# Attribute access must not raise (this is exactly what was broken).
assert server_tool_use.web_search_requests == 3
def test_completion_cost_does_not_raise_on_streaming_web_search_response():
"""
Regression: completion_cost(...) must not raise AttributeError when the
response was reconstructed by stream_chunk_builder from a streaming
Anthropic web_search call.
"""
chunks = [
_make_text_chunk("hello"),
_make_finish_chunk_with_usage_dict_server_tool_use(),
]
rebuilt = stream_chunk_builder(chunks)
assert rebuilt is not None
# The exact dollar amount depends on the model-pricing table; what matters
# for this regression is that it does NOT raise AttributeError on
# `dict has no attribute 'web_search_requests'`.
try:
cost = completion_cost(completion_response=rebuilt)
except AttributeError as e: # pragma: no cover - regression guard
pytest.fail(
"completion_cost raised AttributeError after stream_chunk_builder "
f"(issue #26153 regression): {e}"
)
assert isinstance(cost, (int, float))

5
tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py
View file

@ -520,7 +520,10 @@ def test_stream_chunk_builder_anthropic_web_search():
assert usage.prompt_tokens == 50
assert usage.completion_tokens == 27
assert usage.total_tokens == 77
assert usage.server_tool_use["web_search_requests"] == 2
# server_tool_use must be a ServerToolUse pydantic so downstream cost-calc
# (which uses attribute access) works. See issue #26153.
assert isinstance(usage.server_tool_use, ServerToolUse)
assert usage.server_tool_use.web_search_requests == 2
def test_sort_chunks_handles_dict_hidden_params_created_at():

94
tests/test_litellm/llms/anthropic/test_cost_calculation_dict_safety.py Normal file
View file

@ -0,0 +1,94 @@
"""
Tests that ``get_cost_for_anthropic_web_search`` tolerates ``server_tool_use``
being either a ``dict`` or a ``ServerToolUse`` pydantic instance.
See https://github.com/BerriAI/litellm/issues/26153.
"""
import os
import sys
import pytest
sys.path.insert(0, os.path.abspath("../../../.."))
from litellm.llms.anthropic.cost_calculation import (
_get_web_search_requests,
get_cost_for_anthropic_web_search,
)
from litellm.types.utils import ModelInfo, ServerToolUse
class _UsageWithServerToolUse:
def __init__(self, server_tool_use):
self.server_tool_use = server_tool_use
def _make_model_info(cost_per_query: float = 0.01) -> ModelInfo:
info: ModelInfo = { # type: ignore[typeddict-item]
"search_context_cost_per_query": {
"search_context_size_low": cost_per_query,
"search_context_size_medium": cost_per_query,
"search_context_size_high": cost_per_query,
}
}
return info
def test_get_web_search_requests_handles_none():
assert _get_web_search_requests(None) is None
def test_get_web_search_requests_handles_dict():
assert _get_web_search_requests({"web_search_requests": 4}) == 4
def test_get_web_search_requests_handles_dict_missing_key():
assert _get_web_search_requests({}) is None
def test_get_web_search_requests_handles_pydantic():
assert _get_web_search_requests(ServerToolUse(web_search_requests=2)) == 2
def test_get_cost_for_anthropic_web_search_with_dict_server_tool_use():
"""
Regression: ``server_tool_use`` was a dict from ``stream_chunk_builder`` and
direct attribute access on it raised ``AttributeError``.
"""
usage = _UsageWithServerToolUse({"web_search_requests": 3})
info = _make_model_info(cost_per_query=0.01)
cost = get_cost_for_anthropic_web_search(
model_info=info, usage=usage # type: ignore[arg-type]
)
assert cost == pytest.approx(0.03)
def test_get_cost_for_anthropic_web_search_with_pydantic_server_tool_use():
usage = _UsageWithServerToolUse(ServerToolUse(web_search_requests=3))
info = _make_model_info(cost_per_query=0.01)
cost = get_cost_for_anthropic_web_search(
model_info=info, usage=usage # type: ignore[arg-type]
)
assert cost == pytest.approx(0.03)
def test_get_cost_for_anthropic_web_search_with_none_server_tool_use():
usage = _UsageWithServerToolUse(None)
info = _make_model_info(cost_per_query=0.01)
cost = get_cost_for_anthropic_web_search(
model_info=info, usage=usage # type: ignore[arg-type]
)
assert cost == 0.0
def test_get_cost_for_anthropic_web_search_with_no_usage():
info = _make_model_info(cost_per_query=0.01)
cost = get_cost_for_anthropic_web_search(model_info=info, usage=None)
assert cost == 0.0

398
tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py
View file

@ -12,6 +12,11 @@ import sys
sys.path.insert(0, os.path.abspath("../../../../.."))
import pytest
from botocore.exceptions import (
ConnectTimeoutError,
PartialCredentialsError,
ProfileNotFound,
)
import litellm
from litellm.llms.bedrock_mantle.responses.transformation import (
@ -114,16 +119,15 @@ class TestBedrockMantleResponsesAuth:
)
assert headers["Authorization"] == "Bearer bearer-key"
def test_missing_key_raises(self, monkeypatch):
def test_missing_bearer_does_not_raise_in_validate_environment(self, monkeypatch):
# SigV4 may still apply, so validate_environment must defer instead of raising.
monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False)
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
cfg = BedrockMantleResponsesAPIConfig()
with pytest.raises(ValueError, match="Bedrock Mantle API key"):
cfg.validate_environment(
headers={},
model="openai.gpt-5.5",
litellm_params=GenericLiteLLMParams(),
)
headers = cfg.validate_environment(
headers={}, model="openai.gpt-5.5", litellm_params=GenericLiteLLMParams()
)
assert "Authorization" not in headers
def test_custom_llm_provider(self):
cfg = BedrockMantleResponsesAPIConfig()
@ -261,6 +265,386 @@ def local_cost_map(monkeypatch):
litellm.get_model_info.cache_clear()
class TestBedrockMantleResponsesSigV4:
def test_bearer_short_circuits_without_credentials(self, monkeypatch):
from unittest.mock import MagicMock
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False)
signer = BaseAWSLLM()
signer.get_credentials = MagicMock(
side_effect=AssertionError("get_credentials must not run for bearer auth")
)
cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer)
headers, signed_body = cfg.sign_request(
headers={},
optional_params={},
request_data={"input": "hi"},
api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses",
api_key="bearer-from-config",
)
assert headers["Authorization"] == "Bearer bearer-from-config"
assert signed_body == b'{"input": "hi"}'
signer.get_credentials.assert_not_called()
def test_bearer_resolved_from_mantle_env_key(self, monkeypatch):
from unittest.mock import MagicMock
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
monkeypatch.setenv("BEDROCK_MANTLE_API_KEY", "env-bearer")
signer = BaseAWSLLM()
signer.get_credentials = MagicMock(
side_effect=AssertionError("get_credentials must not run for bearer auth")
)
cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer)
headers, _ = cfg.sign_request(
headers={},
optional_params={},
request_data={"input": "hi"},
api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses",
api_key=None,
)
assert headers["Authorization"] == "Bearer env-bearer"
def test_bearer_arg_takes_priority_over_mantle_env_key(self, monkeypatch):
# The passed api_key (e.g. litellm_params.api_key) must win over the env
# bearer; a reordered precedence chain would silently use the wrong token.
from unittest.mock import MagicMock
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
monkeypatch.setenv("BEDROCK_MANTLE_API_KEY", "env-bearer")
signer = BaseAWSLLM()
signer.get_credentials = MagicMock(
side_effect=AssertionError("get_credentials must not run for bearer auth")
)
cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer)
headers, _ = cfg.sign_request(
headers={},
optional_params={},
request_data={"input": "hi"},
api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses",
api_key="arg-bearer",
)
assert headers["Authorization"] == "Bearer arg-bearer"
signer.get_credentials.assert_not_called()
def test_access_key_produces_sigv4_headers(self, monkeypatch):
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False)
cfg = BedrockMantleResponsesAPIConfig(aws_signer=BaseAWSLLM())
headers, signed_body = cfg.sign_request(
headers={},
optional_params={
"aws_access_key_id": "AKIAEXAMPLE",
"aws_secret_access_key": "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0",
"aws_session_token": "session-token-test",
"aws_region_name": "us-east-2",
},
request_data={"input": "hi"},
api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses",
api_key=None,
)
assert headers["Authorization"].startswith("AWS4-HMAC-SHA256")
assert "Credential=AKIAEXAMPLE/" in headers["Authorization"]
assert "/us-east-2/bedrock/aws4_request" in headers["Authorization"]
assert "X-Amz-Date" in headers
assert headers["X-Amz-Security-Token"] == "session-token-test"
assert signed_body == b'{"input": "hi"}'
def test_assume_role_path_produces_sigv4_headers(self, monkeypatch):
from unittest.mock import MagicMock
from botocore.credentials import Credentials
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False)
signer = BaseAWSLLM()
signer.get_credentials = MagicMock(
return_value=Credentials(
access_key="ASIAEXAMPLE",
secret_key="YXNzdW1lZC1yb2xlLXNlY3JldC1hc3N1bWVk",
token="assumed-session-token",
)
)
cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer)
headers, _ = cfg.sign_request(
headers={},
optional_params={
"aws_role_name": "arn:aws:iam::000000000000:role/test-role",
"aws_session_name": "litellm-test",
"aws_region_name": "us-east-2",
},
request_data={"input": "hi"},
api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses",
api_key=None,
)
signer.get_credentials.assert_called_once()
call = signer.get_credentials.call_args.kwargs
assert call["aws_role_name"] == "arn:aws:iam::000000000000:role/test-role"
assert call["aws_session_name"] == "litellm-test"
assert headers["Authorization"].startswith("AWS4-HMAC-SHA256")
assert "/us-east-2/bedrock/aws4_request" in headers["Authorization"]
def test_signed_body_matches_final_data_after_normalize(self, monkeypatch):
"""Core regression: the signed bytes must equal the bytes actually sent.
Sign the *final* data dict and assert the returned signed_body decodes to
exactly that dict, so a later change to the data would break the SigV4 hash.
"""
import json
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False)
final_data = {"model": "openai.gpt-5.5", "input": "hi", "max_output_tokens": 16}
cfg = BedrockMantleResponsesAPIConfig(aws_signer=BaseAWSLLM())
_, signed_body = cfg.sign_request(
headers={},
optional_params={
"aws_access_key_id": "AKIAEXAMPLE",
"aws_secret_access_key": "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0",
"aws_region_name": "us-east-2",
},
request_data=final_data,
api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses",
api_key=None,
)
assert signed_body is not None
assert json.loads(signed_body) == final_data
def test_region_comes_from_optional_params(self, monkeypatch):
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False)
monkeypatch.delenv("AWS_REGION", raising=False)
monkeypatch.delenv("AWS_REGION_NAME", raising=False)
cfg = BedrockMantleResponsesAPIConfig(aws_signer=BaseAWSLLM())
headers, _ = cfg.sign_request(
headers={},
optional_params={
"aws_access_key_id": "AKIAEXAMPLE",
"aws_secret_access_key": "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0",
"aws_region_name": "eu-west-1",
},
request_data={"input": "hi"},
api_base="https://bedrock-mantle.eu-west-1.api.aws/openai/v1/responses",
api_key=None,
)
assert "/eu-west-1/bedrock/aws4_request" in headers["Authorization"]
def test_url_region_and_sigv4_region_agree_from_litellm_params(self, monkeypatch):
"""Adversarial-review regression: a caller-supplied aws_region_name (no region
env set) must shape BOTH the URL host and the SigV4 credential scope, or the
request is signed for one region and sent to another -> 401.
"""
monkeypatch.delenv("BEDROCK_MANTLE_REGION", raising=False)
monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False)
monkeypatch.delenv("AWS_REGION", raising=False)
monkeypatch.delenv("AWS_REGION_NAME", raising=False)
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False)
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
params = {
"aws_region_name": "ap-southeast-2",
"aws_access_key_id": "AKIAEXAMPLE",
"aws_secret_access_key": "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0",
}
cfg = BedrockMantleResponsesAPIConfig(aws_signer=BaseAWSLLM())
url = cfg.get_complete_url(api_base=None, litellm_params=params)
assert (
url == "https://bedrock-mantle.ap-southeast-2.api.aws/openai/v1/responses"
)
headers, _ = cfg.sign_request(
headers={},
optional_params=params,
request_data={"input": "hi"},
api_base=url,
api_key=None,
)
assert "/ap-southeast-2/bedrock/aws4_request" in headers["Authorization"]
def test_injected_default_region_base_does_not_override_aws_region_name(
self, monkeypatch
):
"""2nd-round adversarial regression: responses/main.py auto-injects
litellm_params.api_base = https://bedrock-mantle.<DEFAULT>.api.aws/v1 (default
region, ignoring aws_region_name). The config must still pin BOTH the URL host
and the SigV4 scope to aws_region_name, or the IAM deployment 401s. A naive
'resolve region only when api_base is None' fix would fail this test.
"""
monkeypatch.delenv("BEDROCK_MANTLE_REGION", raising=False)
monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False)
monkeypatch.delenv("AWS_REGION", raising=False)
monkeypatch.delenv("AWS_REGION_NAME", raising=False)
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False)
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
injected_base = "https://bedrock-mantle.us-east-1.api.aws/v1" # default region
params = {
"aws_region_name": "us-east-2", # what the caller actually wants
"api_base": injected_base,
"aws_access_key_id": "AKIAEXAMPLE",
"aws_secret_access_key": "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0",
}
cfg = BedrockMantleResponsesAPIConfig(aws_signer=BaseAWSLLM())
url = cfg.get_complete_url(api_base=injected_base, litellm_params=params)
assert url == "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses"
headers, _ = cfg.sign_request(
headers={},
optional_params=params,
request_data={"input": "hi"},
api_base=url,
api_key=None,
)
assert "/us-east-2/bedrock/aws4_request" in headers["Authorization"]
assert "us-east-1" not in headers["Authorization"]
def test_custom_proxy_host_is_preserved(self, monkeypatch):
"""A genuinely custom (non-Mantle) api_base host must be preserved, not rewritten
to a bedrock-mantle host. Only standard Mantle hosts are region-pinned.
"""
monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False)
cfg = BedrockMantleResponsesAPIConfig()
url = cfg.get_complete_url(
api_base="https://mantle-proxy.internal.example/openai/v1",
litellm_params={"aws_region_name": "us-east-2"},
)
assert url == "https://mantle-proxy.internal.example/openai/v1/responses"
def test_caller_authorization_does_not_override_sigv4(self, monkeypatch):
"""Adversarial-review regression: a caller-supplied Authorization header (e.g.
from extra_headers, surviving the relaxed validate_environment) must not clobber
the SigV4 Authorization that _sign_request would otherwise restore.
"""
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False)
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
cfg = BedrockMantleResponsesAPIConfig(aws_signer=BaseAWSLLM())
headers, _ = cfg.sign_request(
headers={"Authorization": "Bearer stale-caller-token"},
optional_params={
"aws_access_key_id": "AKIAEXAMPLE",
"aws_secret_access_key": "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0",
"aws_region_name": "us-east-2",
},
request_data={"input": "hi"},
api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses",
api_key=None,
)
assert headers["Authorization"].startswith("AWS4-HMAC-SHA256")
assert "Bearer stale-caller-token" not in headers["Authorization"]
def test_no_bearer_and_no_credentials_raises_both_paths(self, monkeypatch):
from unittest.mock import MagicMock
from botocore.exceptions import NoCredentialsError
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False)
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
signer = BaseAWSLLM()
signer.get_credentials = MagicMock(side_effect=NoCredentialsError())
cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer)
with pytest.raises(ValueError) as exc:
cfg.sign_request(
headers={},
optional_params={"aws_region_name": "us-east-2"},
request_data={"input": "hi"},
api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses",
api_key=None,
)
msg = str(exc.value)
assert "Bearer" in msg
assert "SigV4" in msg or "IAM" in msg
@pytest.mark.parametrize(
"cred_error",
[
PartialCredentialsError(provider="env", cred_var="aws_secret_access_key"),
ProfileNotFound(profile="missing-profile"),
],
)
def test_partial_credentials_raises_both_paths(self, monkeypatch, cred_error):
from unittest.mock import MagicMock
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False)
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
signer = BaseAWSLLM()
signer.get_credentials = MagicMock(side_effect=cred_error)
cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer)
with pytest.raises(ValueError) as exc:
cfg.sign_request(
headers={},
optional_params={"aws_region_name": "us-east-2"},
request_data={"input": "hi"},
api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses",
api_key=None,
)
msg = str(exc.value)
assert "Bearer" in msg
assert "SigV4" in msg or "IAM" in msg
def test_sts_transport_error_is_not_masked_as_credentials(self, monkeypatch):
# An AssumeRole / web-identity flow hits STS over the network, so a transient
# connection error must surface as itself, not be rewritten into the
# "no usable AWS credentials" message that would send the user to fix the
# wrong thing.
from unittest.mock import MagicMock
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False)
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
signer = BaseAWSLLM()
signer.get_credentials = MagicMock(
side_effect=ConnectTimeoutError(
endpoint_url="https://sts.us-east-2.amazonaws.com"
)
)
cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer)
with pytest.raises(ConnectTimeoutError):
cfg.sign_request(
headers={},
optional_params={
"aws_role_name": "arn:aws:iam::000000000000:role/test-role",
"aws_region_name": "us-east-2",
},
request_data={"input": "hi"},
api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses",
api_key=None,
)
class TestBedrockMantleResponsesPricing:
def test_gpt_5_5_pricing_and_mode(self, local_cost_map):
info = litellm.get_model_info("bedrock_mantle/openai.gpt-5.5")

238
tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py
View file

@ -742,3 +742,241 @@ async def test_anthropic_post_retry_reserializes_mutated_body():
assert first_sent == prebuilt # attempt 0 used prebuilt
assert second_sent == _json.dumps(request_body) # attempt 1 re-serialized
assert "MUTATED" in second_sent # ... the mutated body
def test_base_responses_config_sign_request_is_noop_by_default():
"""Default responses sign_request must be a no-op: unchanged headers, no signed body.
Guards the 15 existing responses providers from accidental signing when the
handler starts calling sign_request.
"""
from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
cfg = OpenAIResponsesAPIConfig()
headers = {"Authorization": "Bearer sk-existing"}
out_headers, signed_body = cfg.sign_request(
headers=headers,
optional_params={},
request_data={"input": "hi"},
api_base="https://api.openai.com/v1/responses",
)
assert out_headers == {"Authorization": "Bearer sk-existing"}
assert signed_body is None
def _make_responses_handler_call(signed_body):
"""Drive BaseLLMHTTPHandler.response_api_handler with a fully mocked provider
config + sync client, returning the kwargs the client.post was called with.
signed_body=None simulates a no-op (non-signing) provider; bytes simulates a
signing provider (e.g. Bedrock Mantle).
"""
from unittest.mock import MagicMock
from litellm.llms.custom_httpx.http_handler import HTTPHandler
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.types.router import GenericLiteLLMParams
provider_config = MagicMock()
provider_config.validate_environment.return_value = {}
provider_config.get_complete_url.return_value = (
"https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses"
)
provider_config.transform_responses_api_request.return_value = {"input": "hi"}
provider_config.should_fake_stream.return_value = False
provider_config.sign_request.return_value = ({"X-Signed": "1"}, signed_body)
mock_client = MagicMock(spec=HTTPHandler)
mock_client.post.return_value = MagicMock()
handler = BaseLLMHTTPHandler()
handler.response_api_handler(
model="openai.gpt-5.5",
input="hi",
responses_api_provider_config=provider_config,
response_api_optional_request_params={},
custom_llm_provider="bedrock_mantle",
litellm_params=GenericLiteLLMParams(aws_region_name="us-east-2"),
logging_obj=MagicMock(),
client=mock_client,
_is_async=False,
)
return mock_client.post.call_args.kwargs
def test_responses_handler_sends_json_when_not_signed():
"""No-op provider (signed_body is None) -> handler posts json=data, no data= bytes."""
kwargs = _make_responses_handler_call(signed_body=None)
assert kwargs.get("json") == {"input": "hi"}
assert "data" not in kwargs
def test_responses_handler_sends_signed_bytes_when_signed():
"""Signing provider -> handler posts the exact signed bytes via data=, not json=."""
kwargs = _make_responses_handler_call(signed_body=b'{"input": "hi"}')
assert kwargs.get("data") == b'{"input": "hi"}'
assert "json" not in kwargs
assert kwargs["headers"] == {"X-Signed": "1"}
def test_responses_handler_signs_after_fake_stream_prep_strips_stream():
"""Fake-stream signing-order invariant: the bytes SIGNED must equal the bytes SENT.
In the streaming + fake-stream path the handler first runs
_prepare_fake_stream_request, which pops "stream" out of the body, and only
then calls sign_request. If signing ran before that pop, the signed body
would still carry "stream" while the body sent over the wire would not,
producing a SigV4 payload-hash mismatch (401) for a real Mantle deployment.
We snapshot request_data at sign time and assert "stream" is already gone.
"""
from unittest.mock import MagicMock
from litellm.llms.custom_httpx.http_handler import HTTPHandler
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.types.llms.openai import ResponsesAPIResponse
from litellm.types.router import GenericLiteLLMParams
provider_config = MagicMock()
provider_config.validate_environment.return_value = {}
provider_config.get_complete_url.return_value = (
"https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses"
)
provider_config.transform_responses_api_request.return_value = {
"input": "hi",
"stream": True,
}
provider_config.should_fake_stream.return_value = True
provider_config.transform_response_api_response.return_value = ResponsesAPIResponse(
id="resp_1",
created_at=0,
output=[],
status="completed",
model="openai.gpt-5.5",
)
captured = {}
def _capture_sign(**kwargs):
captured["request_data"] = dict(kwargs["request_data"])
return ({"X-Signed": "1"}, b'{"input": "hi"}')
provider_config.sign_request.side_effect = _capture_sign
mock_client = MagicMock(spec=HTTPHandler)
mock_client.post.return_value = MagicMock()
handler = BaseLLMHTTPHandler()
handler.response_api_handler(
model="openai.gpt-5.5",
input="hi",
responses_api_provider_config=provider_config,
response_api_optional_request_params={"stream": True},
custom_llm_provider="bedrock_mantle",
litellm_params=GenericLiteLLMParams(aws_region_name="us-east-2"),
logging_obj=MagicMock(),
client=mock_client,
_is_async=False,
fake_stream=True,
)
assert "stream" not in captured["request_data"]
assert "input" in captured["request_data"]
post_kwargs = mock_client.post.call_args.kwargs
assert post_kwargs.get("data") == b'{"input": "hi"}'
assert "json" not in post_kwargs
assert "stream" in post_kwargs
def _make_compact_handler_call(signed_body, is_async):
"""Drive (async_)compact_response_api_handler with a fully mocked provider config
+ client, returning the kwargs the client.post was called with.
signed_body=None simulates a no-op (non-signing) provider; bytes simulates a
signing provider (e.g. Bedrock Mantle SigV4 / bearer).
"""
from unittest.mock import MagicMock
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.types.router import GenericLiteLLMParams
compact_url = "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses/compact"
provider_config = MagicMock()
provider_config.validate_environment.return_value = {}
provider_config.get_complete_url.return_value = (
"https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses"
)
provider_config.transform_compact_response_api_request.return_value = (
compact_url,
{"model": "openai.gpt-5.5", "input": "hi"},
)
provider_config.sign_request.return_value = ({"X-Signed": "1"}, signed_body)
provider_config.transform_compact_response_api_response.return_value = "ok"
spec = AsyncHTTPHandler if is_async else HTTPHandler
mock_client = MagicMock(spec=spec)
if is_async:
mock_client.post = AsyncMock(return_value=MagicMock())
else:
mock_client.post.return_value = MagicMock()
handler = BaseLLMHTTPHandler()
result = handler.compact_response_api_handler(
model="openai.gpt-5.5",
input="hi",
responses_api_provider_config=provider_config,
response_api_optional_request_params={},
custom_llm_provider="bedrock_mantle",
litellm_params=GenericLiteLLMParams(aws_region_name="us-east-2"),
logging_obj=MagicMock(),
client=mock_client,
_is_async=is_async,
)
if is_async:
asyncio.run(result)
return provider_config, mock_client.post.call_args.kwargs
def test_compact_handler_sends_json_when_not_signed():
"""No-op provider on compact (signed_body is None) -> posts json=data, no data= bytes."""
provider_config, kwargs = _make_compact_handler_call(
signed_body=None, is_async=False
)
provider_config.sign_request.assert_called_once()
assert kwargs.get("json") == {"model": "openai.gpt-5.5", "input": "hi"}
assert "data" not in kwargs
def test_compact_handler_sends_signed_bytes_when_signed():
"""Signing provider on compact -> posts the signed bytes via data=, not json=.
Regression for the adversarial-review finding that /responses/compact bypassed
the SigV4 signing hook, so IAM-only Mantle callers sent unsigned bodies.
"""
provider_config, kwargs = _make_compact_handler_call(
signed_body=b'{"model": "openai.gpt-5.5", "input": "hi"}', is_async=False
)
assert kwargs.get("data") == b'{"model": "openai.gpt-5.5", "input": "hi"}'
assert "json" not in kwargs
assert kwargs["headers"] == {"X-Signed": "1"}
# signing must use the compact endpoint as api_base, not the create URL
assert provider_config.sign_request.call_args.kwargs["api_base"].endswith(
"/openai/v1/responses/compact"
)
def test_async_compact_handler_sends_signed_bytes_when_signed():
"""Async compact must sign identically to sync (same omission in the async twin)."""
provider_config, kwargs = _make_compact_handler_call(
signed_body=b'{"model": "openai.gpt-5.5", "input": "hi"}', is_async=True
)
assert kwargs.get("data") == b'{"model": "openai.gpt-5.5", "input": "hi"}'
assert "json" not in kwargs
assert kwargs["headers"] == {"X-Signed": "1"}
def test_async_compact_handler_sends_json_when_not_signed():
"""Async no-op provider on compact -> posts json=data, no data= bytes."""
_provider_config, kwargs = _make_compact_handler_call(
signed_body=None, is_async=True
)
assert kwargs.get("json") == {"model": "openai.gpt-5.5", "input": "hi"}
assert "data" not in kwargs

194
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py
View file

@ -478,3 +478,197 @@ async def test_apply_guardrail_request_skipped_messages_stay_aligned(
assert result["texts"][1] == ""
assert result["texts"][2] == "Here is my SSN: <US_SSN>"
assert result["structured_messages"] == inputs["structured_messages"]
@pytest.mark.asyncio
async def test_apply_guardrail_sends_user_id_model_and_extra_info(
crowdstrike_aidr_guardrail: CrowdStrikeAIDRHandler,
) -> None:
inputs: GenericGuardrailAPIInputs = {
"texts": ["Hello"],
"structured_messages": [{"role": "user", "content": "Hello"}],
"model": "gpt-4o",
}
request_data = {
"messages": inputs["structured_messages"],
"model": "gpt-4o",
"litellm_metadata": {
"user_api_key_user_id": "uid-abc",
"user_api_key_user_email": "alice@example.com",
},
}
guardrail_endpoint = (
f"{crowdstrike_aidr_guardrail.api_base}/v1/guard_chat_completions"
)
with patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
return_value=httpx.Response(
status_code=200,
json={"result": {"blocked": False, "transformed": False}},
request=httpx.Request(method="POST", url=guardrail_endpoint),
),
) as mock_method:
await crowdstrike_aidr_guardrail.apply_guardrail(
inputs=inputs,
request_data=request_data,
input_type="request",
)
payload = mock_method.call_args.kwargs["json"]
assert payload["user_id"] == "uid-abc"
assert payload["model"] == "gpt-4o"
assert payload["extra_info"] == {"user_name": "alice@example.com"}
@pytest.mark.asyncio
async def test_apply_guardrail_empty_extra_info_when_no_email(
crowdstrike_aidr_guardrail: CrowdStrikeAIDRHandler,
) -> None:
inputs: GenericGuardrailAPIInputs = {
"texts": ["Hello"],
"structured_messages": [{"role": "user", "content": "Hello"}],
"model": "gemini-flash",
}
request_data = {
"messages": inputs["structured_messages"],
"model": "gemini-flash",
"litellm_metadata": {
"user_api_key_user_id": "uid-no-email",
"user_api_key_user_email": None,
},
}
guardrail_endpoint = (
f"{crowdstrike_aidr_guardrail.api_base}/v1/guard_chat_completions"
)
with patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
return_value=httpx.Response(
status_code=200,
json={"result": {"blocked": False, "transformed": False}},
request=httpx.Request(method="POST", url=guardrail_endpoint),
),
) as mock_method:
await crowdstrike_aidr_guardrail.apply_guardrail(
inputs=inputs,
request_data=request_data,
input_type="request",
)
payload = mock_method.call_args.kwargs["json"]
assert payload["user_id"] == "uid-no-email"
assert payload["model"] == "gemini-flash"
assert payload["extra_info"] == {}
@pytest.mark.asyncio
async def test_apply_guardrail_no_metadata_skips_user_fields(
crowdstrike_aidr_guardrail: CrowdStrikeAIDRHandler,
) -> None:
inputs: GenericGuardrailAPIInputs = {
"texts": ["Hello"],
"structured_messages": [{"role": "user", "content": "Hello"}],
}
request_data = {"messages": inputs["structured_messages"]}
guardrail_endpoint = (
f"{crowdstrike_aidr_guardrail.api_base}/v1/guard_chat_completions"
)
with patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
return_value=httpx.Response(
status_code=200,
json={"result": {"blocked": False, "transformed": False}},
request=httpx.Request(method="POST", url=guardrail_endpoint),
),
) as mock_method:
await crowdstrike_aidr_guardrail.apply_guardrail(
inputs=inputs,
request_data=request_data,
input_type="request",
)
payload = mock_method.call_args.kwargs["json"]
assert "user_id" not in payload
assert "model" not in payload
assert "extra_info" not in payload
@pytest.mark.asyncio
@pytest.mark.parametrize(
"litellm_metadata, metadata",
[
(
None,
{
"user_api_key_user_id": "uid-abc",
"user_api_key_user_email": "alice@example.com",
},
),
(
{"trace_id": "t1"},
{
"user_api_key_user_id": "uid-abc",
"user_api_key_user_email": "alice@example.com",
},
),
(
["unexpected"],
{
"user_api_key_user_id": "uid-abc",
"user_api_key_user_email": "alice@example.com",
},
),
(
{
"user_api_key_user_id": "uid-abc",
"user_api_key_user_email": "alice@example.com",
},
{"trace_id": "t1"},
),
],
ids=[
"identity_in_metadata_llm_none",
"identity_in_metadata_llm_user_dict",
"identity_in_metadata_llm_non_mapping",
"identity_in_litellm_metadata",
],
)
async def test_apply_guardrail_reads_identity_from_either_metadata_bag(
crowdstrike_aidr_guardrail: CrowdStrikeAIDRHandler,
litellm_metadata,
metadata,
) -> None:
inputs: GenericGuardrailAPIInputs = {
"texts": ["Hello"],
"structured_messages": [{"role": "user", "content": "Hello"}],
"model": "gpt-4o",
}
request_data = {
"messages": inputs["structured_messages"],
"model": "gpt-4o",
"litellm_metadata": litellm_metadata,
"metadata": metadata,
}
guardrail_endpoint = (
f"{crowdstrike_aidr_guardrail.api_base}/v1/guard_chat_completions"
)
with patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
return_value=httpx.Response(
status_code=200,
json={"result": {"blocked": False, "transformed": False}},
request=httpx.Request(method="POST", url=guardrail_endpoint),
),
) as mock_method:
await crowdstrike_aidr_guardrail.apply_guardrail(
inputs=inputs,
request_data=request_data,
input_type="request",
)
payload = mock_method.call_args.kwargs["json"]
assert payload["user_id"] == "uid-abc"
assert payload["extra_info"] == {"user_name": "alice@example.com"}

99
tests/test_litellm/proxy/hooks/test_batch_file_validation.py
View file

@ -713,7 +713,8 @@ async def test_count_input_file_usage_decodes_model_embedded_file_id():
@pytest.mark.asyncio
async def test_pre_call_allows_stripped_provider_model_when_key_has_proxy_alias():
"""After replace_model_in_jsonl, body.model is the provider id (e.g. gpt-5.5).
Auth must check the proxy model_name the key was granted, not the stripped id."""
Auth must check target_model_names from the unified file id, not reverse-map
the stripped id."""
from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter
rate_limiter = _PROXY_BatchRateLimiter(
@ -732,7 +733,6 @@ async def test_pre_call_allows_stripped_provider_model_when_key_has_proxy_alias(
)
mock_router = MagicMock()
mock_router.model_list = []
mock_router.resolve_model_name_from_model_id.return_value = proxy_alias
can_key_call_model = AsyncMock(return_value=True)
with (
@ -745,10 +745,105 @@ async def test_pre_call_allows_stripped_provider_model_when_key_has_proxy_alias(
await rate_limiter._enforce_batch_file_model_access(
user_api_key_dict=user,
file_content_as_dict=file_dict,
target_model_names=[proxy_alias],
)
can_key_call_model.assert_awaited_once()
assert can_key_call_model.await_args.kwargs["model"] == proxy_alias
mock_router.resolve_model_name_from_model_id.assert_not_called()
@pytest.mark.asyncio
@pytest.mark.parametrize(
"model_list_order",
[
[
"openai/openai/gpt-5.5",
"openai/openai/gpt-5.5-batch",
"us/azure/openai/gpt-5.5",
],
[
"us/azure/openai/gpt-5.5",
"openai/openai/gpt-5.5",
"openai/openai/gpt-5.5-batch",
],
[
"openai/openai/gpt-5.5-batch",
"us/azure/openai/gpt-5.5",
"openai/openai/gpt-5.5",
],
],
)
async def test_pre_call_uses_target_model_names_not_stripped_reverse_lookup(
model_list_order,
):
"""LIT-3593: three deployments strip to gpt-5.5; auth must use the upload
target alias from target_model_names, not first-match reverse lookup."""
from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter
rate_limiter = _PROXY_BatchRateLimiter(
internal_usage_cache=MagicMock(),
parallel_request_limiter=MagicMock(),
)
batch_alias = "openai/openai/gpt-5.5-batch"
deployment_templates = {
"openai/openai/gpt-5.5": {
"model_name": "openai/openai/gpt-5.5",
"litellm_params": {"model": "openai/gpt-5.5"},
"model_info": {"id": "openai/openai/gpt-5.5", "mode": "chat"},
},
"openai/openai/gpt-5.5-batch": {
"model_name": "openai/openai/gpt-5.5-batch",
"litellm_params": {"model": "openai/gpt-5.5"},
"model_info": {"id": "openai/openai/gpt-5.5-batch", "mode": "batch"},
},
"us/azure/openai/gpt-5.5": {
"model_name": "us/azure/openai/gpt-5.5",
"litellm_params": {"model": "azure/gpt-5.5"},
"model_info": {"id": "openai/openai/gpt-5.5", "mode": "chat"},
},
}
mock_router = MagicMock()
mock_router.model_list = [deployment_templates[name] for name in model_list_order]
def _resolve(model_id):
for deployment in mock_router.model_list:
actual_model = deployment.get("litellm_params", {}).get("model")
if actual_model == model_id or (
actual_model and actual_model.endswith(f"/{model_id}")
):
return deployment.get("model_name")
return None
mock_router.resolve_model_name_from_model_id.side_effect = _resolve
file_dict = [
{"body": {"model": "gpt-5.5", "messages": [{"role": "user", "content": "x"}]}}
]
user = UserAPIKeyAuth(
api_key="sk-ok",
user_id="alice",
models=[batch_alias],
user_role=LitellmUserRoles.INTERNAL_USER.value,
)
can_key_call_model = AsyncMock(return_value=True)
with (
patch(
"litellm.proxy.auth.auth_checks.can_key_call_model",
new=can_key_call_model,
),
patch("litellm.proxy.proxy_server.llm_router", mock_router),
):
await rate_limiter._enforce_batch_file_model_access(
user_api_key_dict=user,
file_content_as_dict=file_dict,
target_model_names=[batch_alias],
)
can_key_call_model.assert_awaited_once()
assert can_key_call_model.await_args.kwargs["model"] == batch_alias
mock_router.resolve_model_name_from_model_id.assert_not_called()
@pytest.mark.asyncio