scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-31 06:56:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 245

fix(guardrails): retain assistant scope with structured rewrites

Browse source
This commit is contained in:
CHX606 2026-10-18 01:00:19 -07:00
commit f571e39d96
281 changed files with 7645 additions and 314 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
enterprise/pyproject.toml
View file

@ -1,6 +1,6 @@
[project]
name = "litellm-enterprise"
version = "0.1.77"
version = "0.1.78"
description = "Package for LiteLLM Enterprise features"
readme = "README.md"
requires-python = ">=3.9"
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.1.77"
version = "0.1.78"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-enterprise==",

4
litellm-proxy-extras/pyproject.toml
View file

@ -1,6 +1,6 @@
[project]
name = "litellm-proxy-extras"
version = "0.4.109"
version = "0.4.110"
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
readme = "README.md"
requires-python = ">=3.9"
@ -30,7 +30,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.4.109"
version = "0.4.110"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-proxy-extras==",

11
litellm-rust/crates/llms-types/src/providers/anthropic/beta.rs
View file

@ -195,9 +195,10 @@ impl AnthropicBeta {
fn rejected_by(&self) -> &'static [BetaProvider] {
match self {
Self::AdvancedToolUse20251120 | Self::ContextManagement20250627 => {
&[BetaProvider::BedrockConverse]
Self::AdvancedToolUse20251120 => {
&[BetaProvider::BedrockConverse, BetaProvider::Databricks]
}
Self::ContextManagement20250627 => &[BetaProvider::BedrockConverse],
Self::AdvisorTool20260301 | Self::Compact20260904 => &[
BetaProvider::AzureAi,
BetaProvider::BedrockConverse,
@ -231,12 +232,11 @@ impl AnthropicBeta {
BetaProvider::Bedrock,
BetaProvider::BedrockMantle,
BetaProvider::VertexAi,
BetaProvider::Databricks,
],
Self::Compact20260112 => &[BetaProvider::AzureAi, BetaProvider::BedrockConverse],
Self::ComputerUse20250124 | Self::ComputerUse20251124 | Self::Context1m20250807 => &[],
Self::DangerousToolUse20260903 => {
&[BetaProvider::BedrockConverse, BetaProvider::Databricks]
}
Self::DangerousToolUse20260903 => &[BetaProvider::BedrockConverse],
Self::Effort20251124 => &[BetaProvider::VertexAi],
Self::FastMode20260201 | Self::Oauth20250420 => &[
BetaProvider::AzureAi,
@ -244,6 +244,7 @@ impl AnthropicBeta {
BetaProvider::Bedrock,
BetaProvider::BedrockMantle,
BetaProvider::VertexAi,
BetaProvider::Databricks,
],
Self::FineGrainedToolStreaming20250514 => {
&[BetaProvider::AzureAi, BetaProvider::VertexAi]

2
litellm-rust/crates/llms/src/anthropic/beta_headers.rs
View file

@ -76,7 +76,7 @@ mod tests {
#[case::vertex_ai(BetaProvider::VertexAi, Some("tool-search-tool-2025-10-19"))]
#[case::anthropic(BetaProvider::Anthropic, Some("advanced-tool-use-2025-11-20"))]
#[case::azure_ai(BetaProvider::AzureAi, Some("advanced-tool-use-2025-11-20"))]
#[case::databricks(BetaProvider::Databricks, Some("advanced-tool-use-2025-11-20"))]
#[case::databricks(BetaProvider::Databricks, None)]
#[case::bedrock_converse(BetaProvider::BedrockConverse, None)]
fn filter_renames_per_host(
#[case] provider: BetaProvider,

22
litellm/anthropic_beta_headers_config.json
View file

@ -220,36 +220,36 @@
},
"databricks": {
"advisor-tool-2026-03-01": null,
"advanced-tool-use-2025-11-20": "advanced-tool-use-2025-11-20",
"advanced-tool-use-2025-11-20": null,
"bash_20241022": null,
"bash_20250124": null,
"code-execution-2025-08-25": "code-execution-2025-08-25",
"code-execution-2025-08-25": null,
"compact-2026-01-12": "compact-2026-01-12",
"computer-use-2025-01-24": "computer-use-2025-01-24",
"computer-use-2025-11-24": "computer-use-2025-11-24",
"context-1m-2025-08-07": "context-1m-2025-08-07",
"context-management-2025-06-27": "context-management-2025-06-27",
"dangerous-tool-use-2026-09-03": null,
"dangerous-tool-use-2026-09-03": "dangerous-tool-use-2026-09-03",
"effort-2025-11-24": "effort-2025-11-24",
"fast-mode-2026-02-01": "fast-mode-2026-02-01",
"files-api-2025-04-14": "files-api-2025-04-14",
"fast-mode-2026-02-01": null,
"files-api-2025-04-14": null,
"structured-output-2024-03-01": null,
"fine-grained-tool-streaming-2025-05-14": "fine-grained-tool-streaming-2025-05-14",
"interleaved-thinking-2025-05-14": "interleaved-thinking-2025-05-14",
"mcp-client-2025-11-20": "mcp-client-2025-11-20",
"mcp-client-2025-04-04": "mcp-client-2025-04-04",
"mcp-client-2025-11-20": null,
"mcp-client-2025-04-04": null,
"mcp-servers-2025-12-04": null,
"oauth-2025-04-20": "oauth-2025-04-20",
"oauth-2025-04-20": null,
"output-128k-2025-02-19": "output-128k-2025-02-19",
"per-turn-control-2026-07-01": null,
"prompt-caching-scope-2026-01-05": "prompt-caching-scope-2026-01-05",
"skills-2025-10-02": "skills-2025-10-02",
"prompt-caching-scope-2026-01-05": null,
"skills-2025-10-02": null,
"structured-outputs-2025-11-13": "structured-outputs-2025-11-13",
"text_editor_20241022": null,
"text_editor_20250124": null,
"thinking-binding-controls-2026-08-01": "thinking-binding-controls-2026-08-01",
"token-efficient-tools-2025-02-19": "token-efficient-tools-2025-02-19",
"web-fetch-2025-09-10": "web-fetch-2025-09-10",
"web-fetch-2025-09-10": null,
"web-search-2025-03-05": "web-search-2025-03-05"
}
}

11
litellm/integrations/custom_guardrail.py
View file

@ -1139,6 +1139,17 @@ class CustomGuardrail(CustomLogger):
"""
return False
def supports_structured_messages(self) -> bool:
"""Whether this guardrail scans every row of ``structured_messages`` and returns texts that agree with them.
Guardrails that scan only part of the rows (the latest user turn, the
last row, the turns after the last assistant turn, or no tool results),
derive their returned texts from the rows, or cannot rewrite rows whose
content is a list of parts, override this to return False. Translation handlers that can rewrite the request from ``texts``
alone, like Bedrock Converse, then hand them only ``texts``.
"""
return True
def _suppressed_by_auto_router_compression(self) -> bool:
"""True when an auto router's own compression policy suppresses this guardrail.

10
litellm/litellm_core_utils/get_provider_specific_headers.py
View file

@ -9,6 +9,7 @@ class ProviderSpecificHeaderUtils:
def get_provider_specific_headers(
provider_specific_header: ProviderSpecificHeader | Sequence[ProviderSpecificHeader] | None,
custom_llm_provider: str | None,
native_anthropic_messages: bool = False,
) -> Mapping[str, str]:
"""
Get the provider specific headers for the given custom llm provider.
@ -17,7 +18,8 @@ class ProviderSpecificHeaderUtils:
carries its own comma-separated provider list, so headers that are safe for several
providers and headers that are safe for exactly one can travel on the same request
without sharing a scope. Entries whose provider list does not contain
`custom_llm_provider` contribute nothing.
`custom_llm_provider` contribute nothing, and entries marked
`native_anthropic_messages_only` contribute only to a provider's native Messages route.
Returns:
Dict: The provider specific headers for the given custom llm provider
@ -31,6 +33,12 @@ class ProviderSpecificHeaderUtils:
matched_headers: Final = {}
for scoped_header in scoped_headers:
if (
"native_anthropic_messages_only" in scoped_header
and scoped_header["native_anthropic_messages_only"]
and not native_anthropic_messages
):
continue
stored_providers = scoped_header.get("custom_llm_provider", "")
provider_list = [p.strip() for p in stored_providers.split(",")]
if custom_llm_provider in provider_list:

5
litellm/litellm_core_utils/prompt_templates/factory.py
View file

@ -5533,6 +5533,11 @@ def _bedrock_tools_pt(
return tool_block_list
def bedrock_tool_blocks(tools: Sequence[Mapping[str, object]], model: str | None) -> tuple[BedrockToolBlock, ...]:
"""Bedrock Converse ``toolConfig.tools`` entries for OpenAI or Anthropic tool definitions."""
return tuple(_bedrock_tools_pt([dict(tool) for tool in tools], model=model))
def _is_system_message(message: Mapping[str, object]) -> bool:
role: Final = message.get("role")
return isinstance(role, str) and "system" in role

66
litellm/llms/anthropic/pass_through/messages/handler.py
View file

@ -7,13 +7,14 @@
import asyncio
import contextvars
from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence
from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine, Iterator, Mapping, Sequence
from functools import partial
from typing import Final, cast
from pydantic import TypeAdapter
import litellm
from litellm.anthropic_beta_headers_manager import aload_beta_headers_config
from litellm.litellm_core_utils.exception_mapping_utils import exception_type
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.anthropic.common_utils import (
@ -39,8 +40,7 @@ from ..responses_adapters.handler import LiteLLMMessagesToResponsesAPIHandler
from ..safeguards import (
SafeguardsEvaluator,
native_route_classifies,
read_classifier_model_setting,
requested_dangerous_tool_use,
proxy_serves_safeguards,
with_safeguard_results,
)
from .interceptors import get_messages_interceptors
@ -106,11 +106,17 @@ def _should_route_to_responses_api(
async def _with_native_safeguard_results(
pending: Coroutine[object, object, AnthropicMessagesResponse | AsyncIterator[bytes]],
evaluator: SafeguardsEvaluator,
proxy_classifies: Callable[[], Awaitable[bool]],
build_evaluator: Callable[[], SafeguardsEvaluator | None],
) -> AnthropicMessagesResponse | AsyncIterator[bytes]:
from .streaming_iterator import AnthropicMessagesStreamingResponse, SafeguardedAnthropicSSEStream
response: Final = await pending
if not await proxy_classifies():
return response
evaluator: Final = await asyncio.to_thread(build_evaluator)
if evaluator is None:
return response
if isinstance(response, AnthropicMessagesStreamingResponse):
return AnthropicMessagesStreamingResponse(
completion_stream=SafeguardedAnthropicSSEStream(response, evaluator),
@ -121,6 +127,10 @@ async def _with_native_safeguard_results(
return SafeguardedAnthropicSSEStream(response, evaluator)
def _native_safeguards_evaluator(kwargs: Mapping[str, object], messages: object) -> SafeguardsEvaluator | None:
return request_safeguards_evaluator(kwargs, _SAFEGUARD_MESSAGES.validate_python(messages))
def _deployment_passes_through_anthropic_messages(model_info: object) -> bool:
"""Whether the deployment opted into forwarding /v1/messages untranslated.
@ -139,7 +149,7 @@ def _deployment_supports_cache_control_ttl(model_info: object) -> bool:
def _native_messages_config(
model: str, custom_llm_provider: str | None, model_info: object
model: str, custom_llm_provider: str | None, model_info: object, api_base: str | None
) -> BaseAnthropicMessagesConfig | None:
provider_config: Final = (
ProviderConfigManager.get_provider_anthropic_messages_config(
@ -149,8 +159,11 @@ def _native_messages_config(
if custom_llm_provider is not None and custom_llm_provider in [provider.value for provider in LlmProviders]
else None
)
if provider_config is not None or not _deployment_passes_through_anthropic_messages(model_info):
return provider_config
native_config: Final = (
provider_config if provider_config is not None and provider_config.serves_api_base(model, api_base) else None
)
if native_config is not None or not _deployment_passes_through_anthropic_messages(model_info):
return native_config
from litellm.llms.openai_like.messages.transformation import (
OpenAILikeAnthropicMessagesConfig,
)
@ -158,12 +171,26 @@ def _native_messages_config(
return OpenAILikeAnthropicMessagesConfig(cache_control_ttl=_deployment_supports_cache_control_ttl(model_info))
def proxy_classifies_tool_uses(model: str, custom_llm_provider: str | None, kwargs: Mapping[str, object]) -> bool:
async def _proxy_classifies(model: str, safeguards: object, native_config: BaseAnthropicMessagesConfig | None) -> bool:
if not proxy_serves_safeguards(safeguards):
return False
if native_config is None:
return True
if native_config.should_filter_anthropic_beta_headers():
await aload_beta_headers_config()
return not native_config.supports_safeguards() or native_route_classifies(model, safeguards)
async def proxy_classifies_tool_uses(model: str, custom_llm_provider: str | None, kwargs: Mapping[str, object]) -> bool:
"""Whether this proxy, rather than the upstream, answers ``safeguards`` for a request routed this way."""
safeguards: Final = kwargs.get("safeguards")
if _native_messages_config(model, custom_llm_provider, kwargs.get("model_info")) is not None:
return native_route_classifies(model, safeguards)
return requested_dangerous_tool_use(safeguards) is not None and read_classifier_model_setting() is not None
api_base: Final = kwargs.get("api_base")
return await _proxy_classifies(
model,
kwargs.get("safeguards"),
_native_messages_config(
model, custom_llm_provider, kwargs.get("model_info"), api_base if isinstance(api_base, str) else None
),
)
####### ENVIRONMENT VARIABLES ###################
@ -608,7 +635,7 @@ def anthropic_messages_handler(
)
anthropic_messages_provider_config: Final = _native_messages_config(
model, custom_llm_provider, kwargs.get("model_info")
model, custom_llm_provider, kwargs.get("model_info"), api_base
)
if anthropic_messages_provider_config is None:
# Route to Responses API for OpenAI / Azure, chat/completions for everything else.
@ -691,11 +718,6 @@ def anthropic_messages_handler(
safeguard_kwargs: Final = (
_SAFEGUARD_KWARGS.validate_python(kwargs) if is_async is True and "safeguards" in kwargs else None
)
safeguards_evaluator: Final = (
request_safeguards_evaluator(safeguard_kwargs, _SAFEGUARD_MESSAGES.validate_python(messages))
if safeguard_kwargs is not None and native_route_classifies(model, safeguard_kwargs.get("safeguards"))
else None
)
native_response: Final = base_llm_http_handler.anthropic_messages_handler(
model=model,
messages=cast( # cast-ok: anthropic messages are dicts at runtime
@ -716,6 +738,10 @@ def anthropic_messages_handler(
stream=stream,
kwargs=kwargs,
)
if safeguards_evaluator is None or not isinstance(native_response, Coroutine):
if safeguard_kwargs is None or not isinstance(native_response, Coroutine):
return native_response
return _with_native_safeguard_results(native_response, safeguards_evaluator)
return _with_native_safeguard_results(
native_response,
partial(_proxy_classifies, model, safeguard_kwargs.get("safeguards"), anthropic_messages_provider_config),
partial(_native_safeguards_evaluator, safeguard_kwargs, messages),
)

2
litellm/llms/anthropic/pass_through/messages/mcp_handler.py
View file

@ -136,7 +136,7 @@ async def anthropic_messages_with_mcp(
)
custom_llm_provider: Final = kwargs.get("custom_llm_provider")
classified_here: Final = proxy_classifies_tool_uses(
classified_here: Final = await proxy_classifies_tool_uses(
model, custom_llm_provider if isinstance(custom_llm_provider, str) else None, kwargs
)
loop_call_args: Final[Mapping[str, object]] = (

61
litellm/llms/anthropic/pass_through/messages/transformation.py
View file

@ -3,10 +3,12 @@ from typing import ClassVar, Final
import httpx
import litellm
from litellm._logging import verbose_logger
from litellm.anthropic_beta_headers_manager import is_beta_header_supported_by_bundled_config
from litellm.anthropic_beta_headers_manager import get_unsupported_headers, is_beta_header_supported_by_bundled_config
from litellm.exceptions import AuthenticationError
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.anthropic.pass_through.safeguards import proxy_serves_safeguards
from litellm.llms.base_llm.anthropic_messages.transformation import (
BaseAnthropicMessagesConfig,
)
@ -110,6 +112,27 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
def should_filter_anthropic_beta_headers(self) -> bool:
return self._resolved_provider != "anthropic"
def supports_safeguards(self) -> bool:
"""
Whether the upstream platform answers the ``safeguards`` body field itself. A platform
whose beta header map column marks the dangerous-tool-use beta unsupported does not, and
a platform the map does not list is trusted to unless its config says otherwise.
"""
return not self.should_filter_anthropic_beta_headers() or (
ANTHROPIC_BETA_HEADER_VALUES.DANGEROUS_TOOL_USE_2026_09_03.value
not in get_unsupported_headers(self._resolved_provider)
)
def drops_safeguards(self, safeguards: object, drop_params: bool) -> bool:
"""
``safeguards`` and the dangerous-tool-use beta leave together for a platform that does not
answer them, when the proxy classifies in its place or drop_params is set. Otherwise both
are forwarded as sent and the platform answers its own 400, like any unsupported param.
"""
return not self.supports_safeguards() and (
drop_params or litellm.drop_params is True or proxy_serves_safeguards(safeguards)
)
def _remove_scope_from_cache_control(self, anthropic_messages_request: dict) -> None:
"""
Remove `scope` field from cache_control blocks.
@ -208,7 +231,8 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
ids were found without tool_result blocks immediately after") while
consecutive user turns merge upstream.
Billing-header system blocks are stripped from the top-level ``system``
field regardless of whether anything was hoisted.
field when ``should_strip_billing_metadata`` says so, regardless of
whether anything was hoisted.
Subclasses whose upstream rejects the role opt in by calling this from
their ``transform_anthropic_messages_request``; the first-party Anthropic
@ -242,7 +266,11 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
)
for block in as_system_content_blocks(source)
]
filtered_system: Final = self._filter_billing_headers_from_system(system_content)
filtered_system: Final = (
self._filter_billing_headers_from_system(system_content)
if self.should_strip_billing_metadata()
else system_content
)
if filtered_system:
anthropic_messages_request["system"] = filtered_system
else:
@ -592,6 +620,12 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
status_code=400,
)
if "safeguards" in anthropic_messages_optional_request_params and self.drops_safeguards(
anthropic_messages_optional_request_params["safeguards"], # pyright: ignore[reportUnknownArgumentType] # the request params arrive as a bare dict
litellm_params.drop_params is True,
):
anthropic_messages_optional_request_params.pop("safeguards") # pyright: ignore[reportUnknownMemberType] # the request params arrive as a bare dict
self._translate_reasoning_effort_to_anthropic(
model=model,
optional_params=anthropic_messages_optional_request_params,
@ -684,9 +718,9 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
self,
model: str,
httpx_response: httpx.Response,
request_body: dict,
request_body: dict[str, object], # mutable-ok: mirrors the base signature
litellm_logging_obj: LiteLLMLoggingObj,
) -> AsyncIterator:
) -> AsyncIterator[bytes]:
"""Helper function to handle Anthropic streaming responses using the existing logging handlers"""
from litellm.llms.anthropic.pass_through.messages.streaming_iterator import (
BaseAnthropicMessagesStreamingIterator,
@ -703,13 +737,14 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
litellm_logging_obj=litellm_logging_obj,
)
@staticmethod
def _update_headers_with_anthropic_beta(
self,
headers: dict,
optional_params: dict,
custom_llm_provider: str = "anthropic",
messages: Sequence[object] = (),
beta_map_provider: str | None = None,
drop_params: bool = False,
) -> dict:
"""
Auto-inject anthropic-beta headers based on features used.
@ -727,6 +762,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
custom_llm_provider: Provider name for looking up correct tool search header
messages: Request messages, scanned for per-message output_config
beta_map_provider: Provider whose beta map column gates safeguards, defaults to custom_llm_provider
drop_params: Whether the deployment or request set drop_params
"""
beta_values: Final[set[str]] = set()
@ -805,9 +841,20 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
else ()
)
all_beta_values: Final = beta_values.union(thinking_display_betas, tool_change_betas)
dangerous_tool_use: Final = ANTHROPIC_BETA_HEADER_VALUES.DANGEROUS_TOOL_USE_2026_09_03.value
forwarded_beta_values: Final = ( # pyright: ignore[reportUnknownVariableType] # derived from the bare beta set
all_beta_values - {dangerous_tool_use}
if dangerous_tool_use in all_beta_values
and self.drops_safeguards(
optional_params.get("safeguards"), # pyright: ignore[reportUnknownArgumentType] # the optional params arrive as a bare dict
drop_params,
)
else all_beta_values
)
if not all_beta_values:
return headers
merged: Final = {key: value for key, value in headers.items() if key.lower() != "anthropic-beta"}
merged["anthropic-beta"] = ",".join(sorted(all_beta_values))
if forwarded_beta_values:
merged["anthropic-beta"] = ",".join(sorted(forwarded_beta_values))
return merged

10
litellm/llms/anthropic/pass_through/safeguards.py
View file

@ -273,12 +273,12 @@ def read_classifier_model_setting() -> str | None:
return configured if isinstance(configured, str) and configured else None
def proxy_serves_safeguards(safeguards: object) -> bool:
return requested_dangerous_tool_use(safeguards) is not None and read_classifier_model_setting() is not None
def native_route_classifies(model: str, safeguards: object) -> bool:
return (
"claude" not in model.lower()
and requested_dangerous_tool_use(safeguards) is not None
and read_classifier_model_setting() is not None
)
return "claude" not in model.lower() and proxy_serves_safeguards(safeguards)
def build_safeguards_evaluator(

9
litellm/llms/base_llm/anthropic_messages/transformation.py
View file

@ -138,6 +138,15 @@ class BaseAnthropicMessagesConfig(ABC):
"""
return True
def supports_safeguards(self) -> bool:
return True
def serves_api_base(self, model: str, api_base: str | None) -> bool:
return True
def is_native_messages_route(self) -> bool:
return True
def handles_web_search_natively(self) -> bool:
"""
Whether the upstream this config routes to executes ``web_search`` tools

353
litellm/llms/bedrock/chat/converse_to_openai.py Normal file
View file

@ -0,0 +1,353 @@
import json
from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
from dataclasses import dataclass
from itertools import accumulate, chain, groupby, repeat, takewhile
from typing import Final, TypeAlias, TypeVar
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError
from litellm._logging import verbose_proxy_logger
from litellm.types.llms.openai import (
AllMessageValues,
ChatCompletionAssistantMessage,
ChatCompletionAssistantToolCall,
ChatCompletionSystemMessage,
ChatCompletionTextObject,
ChatCompletionToolCallFunctionChunk,
ChatCompletionToolMessage,
ChatCompletionToolParam,
ChatCompletionToolParamFunctionChunk,
ChatCompletionUserMessage,
)
JsonPath: TypeAlias = tuple[str | int, ...]
_Node: Final = TypeVar("_Node")
_OBJECT: Final = TypeAdapter(Mapping[str, object])
_ITEMS: Final = TypeAdapter(Sequence[object])
class _ConverseModel(BaseModel):
model_config = ConfigDict(extra="ignore", frozen=True)
class _GuardText(_ConverseModel):
text: str = ""
class _GuardContent(_ConverseModel):
text: _GuardText | None = None
class _ToolUse(_ConverseModel):
tool_use_id: str = Field(default="", alias="toolUseId")
name: str = ""
tool_input: object = Field(default=None, alias="input")
class _ToolResultContent(_ConverseModel):
text: str = ""
json_value: object = Field(default=None, alias="json")
class _ToolResult(_ConverseModel):
tool_use_id: str = Field(default="", alias="toolUseId")
content: tuple[_ToolResultContent, ...] = ()
class _ContentBlock(_ConverseModel):
text: str = ""
guard_content: _GuardContent | None = Field(default=None, alias="guardContent")
tool_use: _ToolUse | None = Field(default=None, alias="toolUse")
tool_result: _ToolResult | None = Field(default=None, alias="toolResult")
class _Message(_ConverseModel):
role: str = ""
content: tuple[_ContentBlock, ...] = ()
class _ToolInputSchema(_ConverseModel):
json_schema: Mapping[str, object] | None = Field(default=None, alias="json")
class _ToolSpec(_ConverseModel):
name: str = ""
description: str = ""
input_schema: _ToolInputSchema | None = Field(default=None, alias="inputSchema")
class _Tool(_ConverseModel):
tool_spec: _ToolSpec | None = Field(default=None, alias="toolSpec")
class _ToolConfig(_ConverseModel):
tools: tuple[_Tool, ...] = ()
class ConverseRequest(_ConverseModel):
system: tuple[_ContentBlock, ...] = ()
messages: tuple[_Message, ...] = ()
tool_config: _ToolConfig | None = Field(default=None, alias="toolConfig")
additional_model_request_fields: object = Field(default=None, alias="additionalModelRequestFields")
@dataclass(frozen=True, slots=True)
class TextSlot:
path: JsonPath
text: str
@dataclass(frozen=True, slots=True)
class JsonSlot:
path: JsonPath
value: object
Slot: TypeAlias = TextSlot | JsonSlot
@dataclass(frozen=True, slots=True)
class ConverseRow:
message: AllMessageValues
content: tuple[Slot, ...]
calls: tuple[JsonSlot, ...] = ()
def parse_converse_request(body: Mapping[str, object]) -> ConverseRequest | None:
try:
return ConverseRequest.model_validate(body)
except ValidationError:
verbose_proxy_logger.debug("Converse body not convertible to OpenAI messages")
return None
def as_json_object(value: object) -> Mapping[str, object] | None:
try:
return _OBJECT.validate_python(value)
except ValidationError:
return None
def as_json_array(value: object) -> Sequence[object] | None:
try:
return _ITEMS.validate_python(value)
except ValidationError:
return None
def breadth_first(root: _Node, children: Callable[[_Node], tuple[_Node, ...]]) -> Iterator[_Node]:
def next_level(level: tuple[_Node, ...], _: None) -> tuple[_Node, ...]:
return tuple(chain.from_iterable(map(children, level)))
return chain.from_iterable(takewhile(bool, accumulate(repeat(None), next_level, initial=(root,))))
def _json_children(node: tuple[JsonPath, object]) -> tuple[tuple[JsonPath, object], ...]:
path, value = node
fields: Final = as_json_object(value)
if fields is not None:
return tuple(((*path, key), child) for key, child in fields.items())
items: Final = as_json_array(value)
return () if items is None else tuple(((*path, index), child) for index, child in enumerate(items))
def _text_leaf(node: tuple[JsonPath, object]) -> TextSlot | None:
path, value = node
return TextSlot(path, value) if isinstance(value, str) and value else None
def _flat_text_order(root: tuple[JsonPath, object]) -> Iterator[tuple[JsonPath, object]]:
"""
The nodes under ``root`` in the order the flat texts list their strings: each container's own strings in field
order, then its nested containers, the last one first.
"""
pending: Final = [root]
while pending:
children = _json_children(pending.pop())
yield from (child for child in children if isinstance(child[1], str))
pending.extend(child for child in children if not isinstance(child[1], str))
def _string_leaves(path: JsonPath, value: object) -> tuple[TextSlot, ...]:
leaves: Final = map(_text_leaf, _flat_text_order((path, value)))
return tuple(leaf for leaf in leaves if leaf is not None)
def _block_slot(path: JsonPath, block: _ContentBlock) -> TextSlot | None:
if block.text:
return TextSlot((*path, "text"), block.text)
guard_text: Final = block.guard_content.text if block.guard_content else None
return (
TextSlot((*path, "guardContent", "text", "text"), guard_text.text) if guard_text and guard_text.text else None
)
def _text_slots(path: JsonPath, blocks: Iterable[tuple[int, _ContentBlock]]) -> tuple[TextSlot, ...]:
slots: Final = (_block_slot((*path, index), block) for index, block in blocks)
return tuple(slot for slot in slots if slot is not None)
def _text_parts(slots: Sequence[TextSlot]) -> tuple[ChatCompletionTextObject, ...]:
return tuple(ChatCompletionTextObject(type="text", text=slot.text) for slot in slots)
def _slot_text(slot: Slot) -> str:
match slot:
case TextSlot():
return slot.text
case JsonSlot():
return json.dumps(slot.value)
def _tool_part_slot(path: JsonPath, part: _ToolResultContent) -> Slot | None:
if part.text:
return TextSlot((*path, "text"), part.text)
return None if part.json_value is None else JsonSlot((*path, "json"), part.json_value)
def _tool_result_row(path: JsonPath, result: _ToolResult) -> ConverseRow:
slots: Final = tuple(
slot
for index, part in enumerate(result.content)
if (slot := _tool_part_slot((*path, "content", index), part)) is not None
)
message: Final = ChatCompletionToolMessage(
role="tool", tool_call_id=result.tool_use_id, content="\n".join(map(_slot_text, slots))
)
return ConverseRow(message, slots)
def _user_group(
path: JsonPath, is_tool_result: bool, group: Sequence[tuple[int, _ContentBlock]]
) -> tuple[ConverseRow, ...]:
if is_tool_result:
return tuple(
_tool_result_row((*path, index, "toolResult"), result)
for index, block in group
if (result := block.tool_result) is not None
)
slots: Final = _text_slots(path, group)
return (
(ConverseRow(ChatCompletionUserMessage(role="user", content=list(_text_parts(slots))), slots),) if slots else ()
)
def _user_rows(path: JsonPath, blocks: Sequence[_ContentBlock]) -> tuple[ConverseRow, ...]:
groups: Final = groupby(enumerate(blocks), key=lambda item: item[1].tool_result is not None)
return tuple(
chain.from_iterable(_user_group(path, is_tool_result, tuple(group)) for is_tool_result, group in groups)
)
def _tool_call(tool_use: _ToolUse) -> ChatCompletionAssistantToolCall:
return ChatCompletionAssistantToolCall(
id=tool_use.tool_use_id,
type="function",
function=ChatCompletionToolCallFunctionChunk(
name=tool_use.name,
arguments="{}" if tool_use.tool_input is None else json.dumps(tool_use.tool_input),
),
)
def _assistant_rows(
path: JsonPath, blocks: Sequence[_ContentBlock], include_tool_calls: bool
) -> tuple[ConverseRow, ...]:
slots: Final = _text_slots(path, enumerate(blocks))
uses: Final = (
tuple((index, tool_use) for index, block in enumerate(blocks) if (tool_use := block.tool_use) is not None)
if include_tool_calls
else ()
)
if uses:
message: Final = ChatCompletionAssistantMessage(
role="assistant",
content=list(_text_parts(slots)) or None,
tool_calls=[_tool_call(tool_use) for _, tool_use in uses],
)
calls: Final = tuple(
JsonSlot((*path, index, "toolUse", "input"), tool_use.tool_input) for index, tool_use in uses
)
return (ConverseRow(message, slots, calls),)
if not slots:
return ()
return (ConverseRow(ChatCompletionAssistantMessage(role="assistant", content=list(_text_parts(slots))), slots),)
def _message_rows(index: int, message: _Message, include_tool_calls: bool) -> tuple[ConverseRow, ...]:
path: Final = ("messages", index, "content")
if message.role == "assistant":
return _assistant_rows(path, message.content, include_tool_calls)
return _user_rows(path, message.content)
def _system_rows(blocks: Sequence[_ContentBlock]) -> tuple[ConverseRow, ...]:
slots: Final = _text_slots(("system",), enumerate(blocks))
return (
(ConverseRow(ChatCompletionSystemMessage(role="system", content=list(_text_parts(slots))), slots),)
if slots
else ()
)
def _user_row(slots: Sequence[Slot]) -> ConverseRow:
parts: Final = [ChatCompletionTextObject(type="text", text=_slot_text(slot)) for slot in slots]
return ConverseRow(ChatCompletionUserMessage(role="user", content=parts), tuple(slots))
def _with_request_fields(rows: tuple[ConverseRow, ...], value: object, user_turn: int) -> tuple[ConverseRow, ...]:
"""
``rows`` with the ``additionalModelRequestFields`` strings after the text of the latest user row, so a guardrail
that judges only the latest user message reads both. When no user row has text they get a user row of their own
at ``user_turn``, where the rows of the latest user turn that is not only tool results start.
"""
slots: Final = _string_leaves(("additionalModelRequestFields",), value)
if not slots:
return rows
latest: Final = next((index for index in reversed(range(len(rows))) if rows[index].message["role"] == "user"), None)
if latest is None:
return (*rows[:user_turn], _user_row(slots), *rows[user_turn:])
return (*rows[:latest], _user_row((*rows[latest].content, *slots)), *rows[latest + 1 :])
def _is_user_turn(message: _Message) -> bool:
return message.role != "assistant" and any(block.tool_result is None for block in message.content)
def converse_rows(request: ConverseRequest, include_tool_calls: bool) -> tuple[ConverseRow, ...]:
system: Final = _system_rows(request.system)
messages: Final = tuple(
_message_rows(index, message, include_tool_calls) for index, message in enumerate(request.messages)
)
starts: Final = tuple(accumulate(map(len, messages), initial=len(system)))
rows: Final = (*system, *chain.from_iterable(messages))
user_turn: Final = next(
(starts[index] for index in reversed(range(len(messages))) if _is_user_turn(request.messages[index])),
len(rows),
)
return _with_request_fields(rows, request.additional_model_request_fields, user_turn)
@dataclass(frozen=True, slots=True)
class ConverseTool:
index: int
param: ChatCompletionToolParam
def converse_tools(request: ConverseRequest) -> tuple[ConverseTool, ...]:
tools: Final = request.tool_config.tools if request.tool_config else ()
return tuple(
ConverseTool(
index,
ChatCompletionToolParam(
type="function",
function=ChatCompletionToolParamFunctionChunk(
name=spec.name,
description=spec.description,
parameters=dict(spec.input_schema.json_schema),
),
),
)
for index, tool in enumerate(tools)
if (spec := tool.tool_spec) and spec.name and spec.input_schema and spec.input_schema.json_schema is not None
)

3
litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py
View file

@ -98,6 +98,9 @@ class AmazonAnthropicClaudeMessagesConfig(
def beta_headers_provider(self) -> str:
return self.custom_llm_provider or "bedrock"
def should_strip_billing_metadata(self) -> bool:
return True
def drops_dangerous_tool_use_beta_without_safeguards(self) -> bool:
return True

4
litellm/llms/bedrock/messages/mantle_transformation.py
View file

@ -176,9 +176,9 @@ class AmazonMantleMessagesConfig(AmazonAnthropicClaudeMessagesConfig):
self,
model: str,
httpx_response: httpx.Response,
request_body: dict,
request_body: dict[str, object], # mutable-ok: mirrors the base signature
litellm_logging_obj: LiteLLMLoggingObj,
) -> AsyncIterator:
) -> AsyncIterator[bytes]:
return AnthropicMessagesConfig.get_async_streaming_response_iterator(
self,
model=model,

448
litellm/llms/bedrock/passthrough/guardrail_translation/handler.py
View file

@ -5,10 +5,12 @@ import struct
from binascii import crc32
from collections.abc import AsyncGenerator, AsyncIterable, Callable, Mapping, Sequence
from dataclasses import dataclass
from functools import partial, reduce
from itertools import accumulate, chain
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Optional, Protocol, TypeAlias, TypeVar
from pydantic import TypeAdapter, ValidationError
from pydantic import InstanceOf, TypeAdapter, ValidationError
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_proxy_logger
@ -19,11 +21,32 @@ from litellm.llms.base_llm.guardrail_translation.base_translation import (
StreamTransformSink,
)
from litellm.llms.base_llm.guardrail_translation.utils import (
effective_scan_only_tool_results_for_guardrail,
effective_skip_assistant_message_for_guardrail,
effective_skip_system_message_for_guardrail,
effective_skip_tool_message_for_guardrail,
scoped_structured_message_indices,
unappliable_request_rewrite,
)
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.bedrock.chat.converse_to_openai import (
ConverseRow,
JsonPath,
as_json_array,
as_json_object,
converse_rows,
converse_tools,
parse_converse_request,
)
from litellm.llms.bedrock.passthrough.guardrail_translation.structured_edits import (
Edit,
ReplacedTools,
ReturnedTool,
ToolResultText,
structured_edits,
tool_edits,
)
from litellm.proxy.guardrails.auto_router_compression import is_compression_guardrail
from litellm.types.utils import GenericGuardrailAPIInputs
if TYPE_CHECKING:
@ -63,7 +86,32 @@ def _generic_passthrough_handler() -> "PassThroughEndpointHandler":
return PassThroughEndpointHandler()
def _reads_flat_texts_only(guardrail: "CustomGuardrail") -> bool:
"""Whether `guardrail` gets only the flat texts on Converse, without rows or tools."""
return is_compression_guardrail(guardrail) or not guardrail.supports_structured_messages()
_StringHolder = tuple[Any, str | int]
_LeafIdentity: TypeAlias = tuple[int, str | int]
@dataclass(frozen=True, slots=True)
class _Assign:
holder: _StringHolder
value: object
@dataclass(frozen=True, slots=True)
class _Remove:
holder: _StringHolder
_Write: TypeAlias = _Assign
_ContentHolder: TypeAlias = tuple[object, str]
_BlankDrop: TypeAlias = tuple[_ContentHolder, frozenset[int]]
_BODY_BY_REFERENCE: Final = TypeAdapter(InstanceOf[dict[str, object]])
_OBJECT_BY_REFERENCE: Final = TypeAdapter(InstanceOf[Mapping[str, object]])
_ARRAY_BY_REFERENCE: Final = TypeAdapter(InstanceOf[list[object]])
def _collect_strings(node: object, holders: list[_StringHolder]) -> None:
@ -99,10 +147,17 @@ def _collect_block_text(block: dict, holders: list[_StringHolder]) -> None:
holders.append((block, "text"))
def _request_field_holders(body: Mapping[str, object]) -> tuple[_StringHolder, ...]:
holders: Final = list[_StringHolder]()
_collect_strings(body.get("additionalModelRequestFields"), holders)
return tuple(holders)
def _extract_converse_texts(
body: dict,
skip_system: bool,
skip_tool: bool,
scan_only_tool_results: bool = False,
skip_assistant: bool = False,
) -> tuple[list[str], list[_StringHolder]]:
"""
@ -123,22 +178,25 @@ def _extract_converse_texts(
inflating guardrail usage for a request whose only prompt is one user
message. No other guardrail translation handler puts tool definitions in
``texts``; the chat and messages handlers carry them in the structured
``tools`` input instead, which this handler does not populate because a
Bedrock ``toolSpec`` is not the OpenAI tool shape those consumers expect.
``tools`` input instead, which this handler also populates by converting
each Bedrock ``toolSpec`` to the OpenAI tool shape.
``additionalModelRequestFields`` is treated differently on purpose. Bedrock
gives ``toolConfig.tools`` a fixed schema whose contents are tool metadata by
contract, while ``additionalModelRequestFields`` is free-form and defined by
the target model, so what it carries cannot be classified without knowing
that model. Scanning it stays the fail-closed default.
that model. Scanning it stays the fail-closed default. Its strings come
last, in the same order as they follow the latest user message in the rows.
"""
holders: Final[list[_StringHolder]] = []
if not skip_system:
system_holders: Final = list[_StringHolder]()
if not skip_system and not scan_only_tool_results:
for block in body.get("system") or []:
if isinstance(block, dict):
_collect_block_text(block, holders)
_collect_block_text(block, system_holders)
request_field_holders: Final = () if scan_only_tool_results else _request_field_holders(body)
message_holders: Final = list[_StringHolder]()
for message in body.get("messages") or []:
if not isinstance(message, dict):
continue
@ -149,23 +207,286 @@ def _extract_converse_texts(
continue
if skip_tool and ("toolUse" in block or "toolResult" in block):
continue
_collect_block_text(block, holders)
if scan_only_tool_results and "toolResult" not in block:
continue
_collect_block_text(block, message_holders)
tool_use = block.get("toolUse")
if isinstance(tool_use, dict):
_collect_strings(tool_use.get("input"), holders)
_collect_strings(tool_use.get("input"), message_holders)
tool_result = block.get("toolResult")
if isinstance(tool_result, dict):
for inner in tool_result.get("content") or []:
if isinstance(inner, dict):
_collect_block_text(inner, holders)
_collect_strings(inner.get("json"), holders)
_collect_strings(body.get("additionalModelRequestFields"), holders)
_collect_block_text(inner, message_holders)
_collect_strings(inner.get("json"), message_holders)
holders: Final = [*system_holders, *message_holders, *request_field_holders]
texts: Final = [container[key] for container, key in holders]
return texts, holders
@dataclass(frozen=True, slots=True)
class _Scope:
skip_system: bool
skip_tool: bool
scan_only_tool_results: bool
skip_assistant: bool = False
def _texts_applied_rows(body: Mapping[str, object], texts: Sequence[str], scope: _Scope) -> tuple[ConverseRow, ...]:
rewritten: Final = _BODY_BY_REFERENCE.validate_python(json.loads(json.dumps(body)))
_, holders = _extract_converse_texts(
rewritten,
scope.skip_system,
scope.skip_tool,
scope.scan_only_tool_results,
skip_assistant=scope.skip_assistant,
)
_write_back_texts(texts, holders)
request: Final = parse_converse_request(rewritten)
return converse_rows(request, include_tool_calls=not scope.skip_tool) if request else ()
def _texts_follow_rows(
body: Mapping[str, object],
rows: Sequence[ConverseRow],
texts: Sequence[str],
guardrailed_texts: Sequence[str],
changed_rows: frozenset[int],
scope: _Scope,
) -> bool:
"""Whether the returned texts sit in the flat layout and only touch rows the vendor also changed."""
if len(guardrailed_texts) != len(texts):
return False
view: Final = _texts_applied_rows(body, guardrailed_texts, scope)
if len(view) != len(rows):
return False
rows_changed_by_texts: Final = frozenset(
index for index, (row, viewed) in enumerate(zip(rows, view)) if row.message != viewed.message
)
return rows_changed_by_texts <= changed_rows
def _child(node: object, key: str | int) -> object:
try:
match key:
case str():
return _OBJECT_BY_REFERENCE.validate_python(node).get(key)
case int():
items: Final = _ARRAY_BY_REFERENCE.validate_python(node)
return items[key] if 0 <= key < len(items) else None
except ValidationError:
return None
def _holder(body: object, path: JsonPath) -> _StringHolder | None:
container: Final = reduce(_child, path[:-1], body)
return (container, path[-1]) if isinstance(_child(container, path[-1]), str) else None
def _leaf_identity(holder: _StringHolder) -> _LeafIdentity:
container, key = holder
return id(container), key
def _changed_leaves(
holders: Sequence[_StringHolder], texts: Sequence[str], guardrailed_texts: Sequence[str]
) -> frozenset[_LeafIdentity]:
return frozenset(
_leaf_identity(holder) for holder, before, after in zip(holders, texts, guardrailed_texts) if after != before
)
def _tool_result_write(body: Mapping[str, object], edit: ToolResultText) -> _Write | None:
container: Final = reduce(_child, edit.content[:-1], body)
blocks: Final = as_json_array(_child(container, edit.content[-1]))
if blocks is None:
return None
dropped: Final = frozenset(edit.indices[1:])
return _Assign(
(container, edit.content[-1]),
[
{"text": edit.text} if index == edit.indices[0] else block
for index, block in enumerate(blocks)
if index not in dropped
],
)
def _is_cache_point(entry: object) -> bool:
fields: Final = as_json_object(entry)
return fields is not None and "cachePoint" in fields
def _tool_name(entry: object) -> object:
return _child(_child(entry, "toolSpec"), "name")
def _entries_after(unsent: Sequence[tuple[int, object]], anchor: int) -> tuple[object, ...]:
return tuple(entry for at, entry in unsent if at == anchor)
def _sent_anchor(sent: frozenset[int], anchor: int, index: int) -> int:
return index if index in sent else anchor
def _replaced(entry: ReturnedTool) -> int | None:
return entry if isinstance(entry, int) else entry.replaces
def _replaced_tools_write(body: Mapping[str, object], edit: ReplacedTools) -> _Write | None:
"""
Rebuild ``toolConfig.tools`` in the order the guardrail returned. Entries it was never sent (cache points, system
tools) stay right after the tool they followed, and a cache point goes away with its tool. Refuses when
``toolChoice`` names a tool that is no longer there, and creates ``toolConfig`` when the request had none. A
``toolConfig`` left with no tool stays until the request is sent, so later guardrails still see the tool history.
"""
tool_config: Final = _child(body, "toolConfig")
tools: Final = () if tool_config is None else as_json_array(_child(tool_config, "tools"))
if tools is None:
return None
anchors: Final = tuple(accumulate(range(len(tools)), partial(_sent_anchor, edit.sent), initial=-1))[1:]
unsent: Final = tuple(
(anchor, tool) for index, (anchor, tool) in enumerate(zip(anchors, tools)) if index not in edit.sent
)
replaced: Final = frozenset(_replaced(entry) for entry in edit.tools)
placed: Final = chain.from_iterable(
(tools[entry], *_entries_after(unsent, entry))
if isinstance(entry, int)
else (*entry.blocks, *(() if entry.replaces is None else _entries_after(unsent, entry.replaces)))
for entry in edit.tools
)
orphans: Final = tuple(tool for at, tool in unsent if at not in replaced and at >= 0 and not _is_cache_point(tool))
rebuilt: Final = [*_entries_after(unsent, -1), *placed, *orphans]
chosen: Final = _child(_child(_child(tool_config, "toolChoice"), "tool"), "name")
if isinstance(chosen, str) and all(_tool_name(tool) != chosen for tool in rebuilt):
return None
if tool_config is None:
return None if all(map(_is_cache_point, rebuilt)) else _Assign((body, "toolConfig"), {"tools": rebuilt})
return _Assign((tool_config, "tools"), rebuilt)
def _resolved(body: Mapping[str, object], edit: Edit) -> _Write | None:
match edit:
case ToolResultText():
return _tool_result_write(body, edit)
case ReplacedTools():
return _replaced_tools_write(body, edit)
case (path, text):
holder: Final = _holder(body, path)
return None if holder is None else _Assign(holder, text)
def _kept_write(write: _Write, flat_changes: frozenset[_LeafIdentity]) -> bool:
return _leaf_identity(write.holder) not in flat_changes
def _holder_edits(
body: Mapping[str, object], edits: Sequence[Edit], flat_changes: frozenset[_LeafIdentity]
) -> tuple[_Write, ...] | None:
resolved: Final = tuple(_resolved(body, edit) for edit in edits)
if any(write is None for write in resolved):
return None
return tuple(write for write in resolved if write is not None and _kept_write(write, flat_changes))
def _apply(writes: Sequence[_Assign | _Remove]) -> None:
for write in writes:
match write:
case _Assign(holder=(container, key), value=value):
container[key] = value
case _Remove(holder=(container, key)):
del container[key]
def _block_text_holder(block: object) -> _ContentHolder | None:
if isinstance(_child(block, "text"), str):
return block, "text"
guarded: Final = _child(_child(block, "guardContent"), "text")
return (guarded, "text") if isinstance(_child(guarded, "text"), str) else None
def _final_texts(
holders: Sequence[_StringHolder], applied_texts: Sequence[str], writes: Sequence[_Write]
) -> Mapping[_LeafIdentity, object]:
flat: Final = {_leaf_identity(holder): text for holder, text in zip(holders, applied_texts)}
structured: Final = {_leaf_identity(write.holder): write.value for write in writes}
return MappingProxyType({**flat, **structured})
def _blanked(final: Mapping[_LeafIdentity, object], block: object) -> bool:
holder: Final = _block_text_holder(block)
if holder is None:
return False
container, key = holder
value: Final = final.get((id(container), key))
return isinstance(value, str) and value != _child(container, key) and not value.strip()
def _blank_blocks(final: Mapping[_LeafIdentity, object], holder: _ContentHolder) -> frozenset[int]:
container, key = holder
return frozenset(id(block) for block in as_json_array(_child(container, key)) or () if _blanked(final, block))
def _empties_turn(drop: _BlankDrop) -> bool:
(container, key), dropped = drop
blocks: Final = as_json_array(_child(container, key)) or ()
return bool(dropped) and all(id(block) in dropped or _is_cache_point(block) for block in blocks)
def _turn_drop(final: Mapping[_LeafIdentity, object], message: object) -> _BlankDrop:
return (message, "content"), _blank_blocks(final, (message, "content"))
def _blank_drops(body: Mapping[str, object], final: Mapping[_LeafIdentity, object]) -> tuple[_BlankDrop, ...]:
"""The system and message text blocks the guardrail left blank, except in a turn they would leave empty."""
turns: Final = (_turn_drop(final, message) for message in as_json_array(body.get("messages")) or ())
system: Final = ((body, "system"), _blank_blocks(final, (body, "system")))
return tuple(drop for drop in (system, *(turn for turn in turns if not _empties_turn(turn))) if drop[1])
def _blank_drop_write(drop: _BlankDrop) -> _Assign | _Remove:
holder, dropped = drop
container, key = holder
kept: Final = [block for block in as_json_array(_child(container, key)) or () if id(block) not in dropped]
return _Assign(holder, kept) if not all(map(_is_cache_point, kept)) else _Remove(holder)
def _edits_drop_a_text_change(
body: Mapping[str, object],
edits: Sequence[Edit],
texts: Sequence[str],
guardrailed_texts: Sequence[str],
scope: _Scope,
) -> bool:
"""
Whether the row and tool edits alone leave out a change the guardrail made in its returned texts. Texts of the
length that was sent are read position by position, so a changed string whose field the edits leave at its
original value is dropped. Texts of another length cannot be placed, so what they mask cannot be checked either,
while an empty list leaves every text as it was.
"""
if not guardrailed_texts:
return False
if len(guardrailed_texts) != len(texts):
return True
rewritten: Final = _BODY_BY_REFERENCE.validate_python(json.loads(json.dumps(body)))
_, holders = _extract_converse_texts(
rewritten,
scope.skip_system,
scope.skip_tool,
scope.scan_only_tool_results,
skip_assistant=scope.skip_assistant,
)
writes: Final = _holder_edits(rewritten, edits, frozenset())
if writes is None:
return True
_apply(writes)
edited: Final = [container[key] for container, key in holders]
return any(
returned != before and after == before for before, returned, after in zip(texts, guardrailed_texts, edited)
)
def _extract_converse_output_texts(
content_blocks: Sequence[object],
) -> tuple[list[str], list[_StringHolder]]:
@ -200,8 +521,8 @@ def _extract_converse_output_texts(
def _write_back_texts(
guardrailed_texts: list[str],
holders: list[_StringHolder],
guardrailed_texts: Sequence[str],
holders: Sequence[_StringHolder],
) -> None:
if len(guardrailed_texts) < len(holders):
verbose_proxy_logger.warning(
@ -986,17 +1307,42 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation):
if not isinstance(body, dict) or not isinstance(body.get("messages"), list):
return data
request_body: Final = _BODY_BY_REFERENCE.validate_python(body)
skip_system: Final = effective_skip_system_message_for_guardrail(guardrail_to_apply)
skip_tool: Final = effective_skip_tool_message_for_guardrail(guardrail_to_apply)
skip_assistant: Final = effective_skip_assistant_message_for_guardrail(guardrail_to_apply)
reads_rows: Final = not _reads_flat_texts_only(guardrail_to_apply)
scope: Final = _Scope(
skip_system=effective_skip_system_message_for_guardrail(guardrail_to_apply),
skip_tool=effective_skip_tool_message_for_guardrail(guardrail_to_apply),
scan_only_tool_results=effective_scan_only_tool_results_for_guardrail(guardrail_to_apply),
skip_assistant=effective_skip_assistant_message_for_guardrail(guardrail_to_apply),
)
texts, holders = _extract_converse_texts(
request_body,
scope.skip_system,
scope.skip_tool,
scope.scan_only_tool_results,
skip_assistant=scope.skip_assistant,
)
request: Final = parse_converse_request(request_body) if reads_rows else None
rows: Final = converse_rows(request, include_tool_calls=not scope.skip_tool) if request else ()
scoped_indices: Final = scoped_structured_message_indices(
tuple(row.message for row in rows),
scan_only_tool_results=scope.scan_only_tool_results,
skip_system=scope.skip_system,
skip_tool=scope.skip_tool,
skip_assistant=scope.skip_assistant,
)
texts, holders = _extract_converse_texts(body, skip_system, skip_tool, skip_assistant)
if not texts:
if not texts and not scoped_indices:
return data
inputs: Final = GenericGuardrailAPIInputs(texts=texts)
sent_messages: Final = [rows[index].message for index in scoped_indices]
inputs: Final = GenericGuardrailAPIInputs(texts=list(texts))
if sent_messages:
inputs["structured_messages"] = json.loads(json.dumps(sent_messages))
tools: Final = converse_tools(request) if request and not scope.scan_only_tool_results else ()
if tools:
inputs["tools"] = copy.deepcopy([tool.param for tool in tools])
model: Final = data.get("model")
if model:
inputs["model"] = model
@ -1009,8 +1355,60 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation):
)
guardrailed_texts: Final = guardrailed_inputs.get("texts", [])
if guardrailed_texts:
_write_back_texts(guardrailed_texts, holders)
guardrailed_messages: Final = guardrailed_inputs.get("structured_messages")
rewritten_rows: Final = (
guardrailed_messages
if sent_messages and isinstance(guardrailed_messages, list) and guardrailed_messages != sent_messages
else None
)
changed_rows: Final[frozenset[int]] = (
frozenset()
if rewritten_rows is None
else frozenset(
index
for index, sent, returned in zip(scoped_indices, sent_messages, rewritten_rows)
if sent != returned
)
)
texts_follow_rows: Final = rewritten_rows is None or _texts_follow_rows(
request_body, rows, texts, guardrailed_texts, changed_rows, scope
)
applied_texts: Final = guardrailed_texts if texts_follow_rows else texts
row_edits: Final = (
()
if rewritten_rows is None
else structured_edits(
rows, scoped_indices, rewritten_rows, partial(_texts_applied_rows, request_body, applied_texts, scope)
)
)
returned_tools: Final = guardrailed_inputs.get("tools")
tool_changes: Final = (
()
if returned_tools is None or not reads_rows
else tool_edits(tools, returned_tools, model if isinstance(model, str) else None)
)
holder_edits: Final = (
None
if row_edits is None
or tool_changes is None
or (
not texts_follow_rows
and _edits_drop_a_text_change(
request_body, (*row_edits, *tool_changes), texts, guardrailed_texts, scope
)
)
else _holder_edits(
request_body, (*row_edits, *tool_changes), _changed_leaves(holders, texts, applied_texts)
)
)
if holder_edits is None:
raise unappliable_request_rewrite(guardrail_to_apply.guardrail_name)
if holder_edits or (applied_texts and list(applied_texts) != texts):
blank_drops: Final = _blank_drops(request_body, _final_texts(holders, applied_texts, holder_edits))
if applied_texts:
_write_back_texts(applied_texts, holders)
_apply(holder_edits)
_apply(tuple(_blank_drop_write(drop) for drop in blank_drops))
return data

337
litellm/llms/bedrock/passthrough/guardrail_translation/structured_edits.py Normal file
View file

@ -0,0 +1,337 @@
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from itertools import chain
from typing import Final, TypeAlias
from pydantic import TypeAdapter, ValidationError
from litellm.litellm_core_utils.prompt_templates.factory import bedrock_tool_blocks
from litellm.llms.bedrock.chat.converse_to_openai import (
ConverseRow,
ConverseTool,
JsonPath,
JsonSlot,
Slot,
TextSlot,
as_json_array,
as_json_object,
breadth_first,
)
from litellm.types.llms.bedrock import BedrockToolBlock
from litellm.types.llms.openai import ChatCompletionTextObject
@dataclass(frozen=True, slots=True)
class ToolResultText:
content: JsonPath
indices: tuple[int, ...]
text: str
@dataclass(frozen=True, slots=True)
class ConvertedTool:
replaces: int | None
blocks: tuple[BedrockToolBlock, ...]
ReturnedTool: TypeAlias = int | ConvertedTool
@dataclass(frozen=True, slots=True)
class ReplacedTools:
sent: frozenset[int]
tools: tuple[ReturnedTool, ...]
TextEdit: TypeAlias = tuple[JsonPath, str]
Edit: TypeAlias = TextEdit | ToolResultText | ReplacedTools
_JsonPair: TypeAlias = tuple[JsonPath, object, object]
_TEXT_PARTS: Final = TypeAdapter(tuple[ChatCompletionTextObject, ...])
_PARSED_JSON: Final = TypeAdapter(object)
def _without(mapping: Mapping[str, object], key: str) -> tuple[tuple[str, object], ...]:
return tuple(sorted(((name, value) for name, value in mapping.items() if name != key), key=lambda item: item[0]))
def _without_content(message: Mapping[str, object]) -> tuple[tuple[str, object], ...]:
return _without(message, "content")
def _text_parts(message: Mapping[str, object]) -> tuple[str, ...] | None:
content: Final = message.get("content")
if isinstance(content, str):
return (content,)
try:
parts: Final = _TEXT_PARTS.validate_python(content)
except ValidationError:
return None
return tuple(part["text"] for part in parts)
def _merged(children: Sequence[tuple[Edit, ...] | None]) -> tuple[Edit, ...] | None:
if any(child is None for child in children):
return None
return tuple(chain.from_iterable(child for child in children if child is not None))
def _paired_children(path: JsonPath, original: object, returned: object) -> tuple[_JsonPair, ...] | None:
original_object: Final = as_json_object(original)
returned_object: Final = as_json_object(returned)
if original_object is not None and returned_object is not None:
if original_object.keys() != returned_object.keys():
return None
return tuple(((*path, key), value, returned_object[key]) for key, value in original_object.items())
original_items: Final = as_json_array(original)
returned_items: Final = as_json_array(returned)
if original_items is None or returned_items is None or len(original_items) != len(returned_items):
return None
return tuple(((*path, index), value, returned_items[index]) for index, value in enumerate(original_items))
def _pair_children(pair: _JsonPair | None) -> tuple[_JsonPair | None, ...]:
if pair is None:
return ()
path, original, returned = pair
if original == returned or (isinstance(original, str) and isinstance(returned, str)):
return ()
children: Final = _paired_children(path, original, returned)
return (None,) if children is None else children
def _leaf_edit(pair: _JsonPair | None) -> TextEdit | None:
match pair:
case (path, str() as original, str() as returned) if original != returned:
return path, returned
case _:
return None
def _json_edits(path: JsonPath, original: object, returned: object) -> tuple[TextEdit, ...] | None:
pairs: Final = tuple(breadth_first((path, original, returned), _pair_children))
if any(pair is None for pair in pairs):
return None
return tuple(edit for pair in pairs if (edit := _leaf_edit(pair)) is not None)
def _slot_edits(slot: Slot, text: str) -> tuple[Edit, ...] | None:
match slot:
case TextSlot():
return () if text == slot.text else ((slot.path, text),)
case JsonSlot():
try:
parsed: Final = _PARSED_JSON.validate_json(text)
except ValidationError:
return None
return _json_edits(slot.path, slot.value, parsed)
def _call_function(call: object) -> Mapping[str, object] | None:
fields: Final = as_json_object(call)
return as_json_object(fields.get("function")) if fields is not None else None
def _call_edits(slot: JsonSlot, original: object, returned: object) -> tuple[Edit, ...] | None:
if original == returned:
return ()
original_call: Final = as_json_object(original)
returned_call: Final = as_json_object(returned)
original_function: Final = _call_function(original)
returned_function: Final = _call_function(returned)
if original_call is None or returned_call is None or original_function is None or returned_function is None:
return None
arguments: Final = returned_function.get("arguments")
if (
not isinstance(arguments, str)
or _without(returned_call, "function") != _without(original_call, "function")
or _without(returned_function, "arguments") != _without(original_function, "arguments")
):
return None
return _slot_edits(slot, arguments)
def _calls_edits(slots: Sequence[JsonSlot], original: object, returned: object) -> tuple[Edit, ...] | None:
original_calls: Final = as_json_array(original)
returned_calls: Final = as_json_array(returned)
if original_calls is None or returned_calls is None or len(original_calls) != len(returned_calls):
return None
return _merged(tuple(_call_edits(*call) for call in zip(slots, original_calls, returned_calls)))
def _carries_no_text(content: object) -> bool:
match content:
case None | "" | []:
return True
case _:
return False
def _content_texts(row: ConverseRow, returned: Mapping[str, object]) -> tuple[str, ...] | None:
content: Final = returned.get("content")
sent: Final = _text_parts(row.message)
if (
isinstance(content, str)
and len(row.content) > 1
and sent is not None
and all("\n" not in text for text in sent)
):
return tuple(content.split("\n"))
return _text_parts(returned)
def _content_edits(row: ConverseRow, original: object, returned: Mapping[str, object]) -> tuple[Edit, ...] | None:
if returned.get("content") == original:
return ()
if not row.content and _carries_no_text(returned.get("content")):
return ()
texts: Final = _content_texts(row, returned)
if texts is None or len(texts) != len(row.content):
return None
return _merged(tuple(_slot_edits(slot, text) for slot, text in zip(row.content, texts)))
def _field(message: Mapping[str, object], key: str) -> object:
return message.get(key)
def _row_edits(row: ConverseRow, returned: Mapping[str, object]) -> tuple[Edit, ...] | None:
if returned == row.message:
return ()
calls: Final = _field(row.message, "tool_calls")
call_edits: Final = _calls_edits(row.calls, calls, returned.get("tool_calls")) if row.calls else ()
comparable: Final = {**returned, "tool_calls": calls} if row.calls else returned
if call_edits is None or _without_content(comparable) != _without_content(row.message):
return None
return _merged((call_edits, _content_edits(row, _field(row.message, "content"), comparable)))
def _same_content(expected: Mapping[str, object], returned: Mapping[str, object]) -> bool:
if expected.get("content") == returned.get("content"):
return True
expected_parts: Final = _text_parts(expected)
returned_parts: Final = _text_parts(returned)
return (
expected_parts is not None
and returned_parts is not None
and "\n".join(expected_parts) == "\n".join(returned_parts)
)
def _reproduces(expected: ConverseRow, returned: Mapping[str, object] | None) -> bool:
return (
returned is not None
and _without_content(returned) == _without_content(expected.message)
and _same_content(expected.message, returned)
)
def _tool_result_text(row: ConverseRow, returned: Mapping[str, object]) -> ToolResultText | None:
texts: Final = _text_parts(returned)
indices: Final = tuple(index for slot in row.content if isinstance(index := slot.path[-2], int))
if (
row.message["role"] != "tool"
or not indices
or texts is None
or _without_content(returned) != _without_content(row.message)
):
return None
return ToolResultText(row.content[0].path[:-2], indices, "\n".join(texts))
def _unmapped_edits(
row: ConverseRow, view: ConverseRow, returned: Mapping[str, object] | None
) -> tuple[Edit, ...] | None:
if _reproduces(view, returned):
return ()
rewrite: Final = _tool_result_text(row, returned) if returned is not None and view == row else None
return None if rewrite is None else (rewrite,)
def structured_edits(
rows: Sequence[ConverseRow],
scoped_indices: Sequence[int],
returned: Sequence[object],
texts_applied_rows: Callable[[], tuple[ConverseRow, ...]],
) -> tuple[Edit, ...] | None:
if len(returned) != len(scoped_indices):
return None
returned_rows: Final = tuple(as_json_object(message) for message in returned)
proposed: Final = tuple(
_row_edits(rows[index], row) if row is not None else None for index, row in zip(scoped_indices, returned_rows)
)
unmapped: Final = tuple(
(index, row) for index, row, edits in zip(scoped_indices, returned_rows, proposed) if edits is None
)
view: Final = texts_applied_rows() if unmapped else ()
if unmapped and len(view) != len(rows):
return None
fallbacks: Final = tuple(_unmapped_edits(rows[index], view[index], row) for index, row in unmapped)
return _merged((*(edits for edits in proposed if edits is not None), *fallbacks))
def _spec_edit(index: int, edit: TextEdit) -> TextEdit | None:
path, text = edit
match path:
case ("function", "description"):
return ("toolConfig", "tools", index, "toolSpec", "description"), text
case ("function", "parameters", *schema_path):
return ("toolConfig", "tools", index, "toolSpec", "inputSchema", "json", *schema_path), text
case _:
return None
def _tool_edits(tool: ConverseTool | None, returned: object) -> tuple[Edit, ...] | None:
edits: Final = None if tool is None else _json_edits((), tool.param, returned)
if tool is None or edits is None:
return None
spec_edits: Final = tuple(_spec_edit(tool.index, edit) for edit in edits)
return None if None in spec_edits else tuple(edit for edit in spec_edits if edit is not None)
def _tool_name(tool: object) -> object:
function: Final = _call_function(tool)
return None if function is None else function.get("name")
def _named(first_by_name: Mapping[str, ConverseTool], returned: object) -> ConverseTool | None:
name: Final = _tool_name(returned)
return first_by_name.get(name) if isinstance(name, str) else None
def _bedrock_blocks(returned: object, model: str | None) -> tuple[BedrockToolBlock, ...] | None:
tool: Final = as_json_object(returned)
function: Final = _call_function(returned)
if (
tool is None
or function is None
or not isinstance(function.get("name"), str)
or ("parameters" in function and as_json_object(function["parameters"]) is None)
):
return None
return bedrock_tool_blocks((tool,), model)
def _returned_tool(
tool: ConverseTool | None, returned: object, model: str | None
) -> tuple[ReturnedTool, tuple[Edit, ...]] | None:
edits: Final = _tool_edits(tool, returned)
if tool is not None and edits is not None:
return tool.index, edits
blocks: Final = _bedrock_blocks(returned, model)
return None if blocks is None else (ConvertedTool(None if tool is None else tool.index, blocks), ())
def tool_edits(tools: Sequence[ConverseTool], returned: object, model: str | None) -> tuple[Edit, ...] | None:
returned_tools: Final = as_json_array(returned)
if returned_tools is None:
return None
first_by_name: Final = {tool.param["function"]["name"]: tool for tool in reversed(tools)}
planned: Final = tuple(_returned_tool(_named(first_by_name, item), item, model) for item in returned_tools)
placed: Final = tuple(plan for plan in planned if plan is not None)
if len(placed) != len(planned):
return None
entries: Final = tuple(entry for entry, _ in placed)
edits: Final = tuple(chain.from_iterable(edits for _, edits in placed))
if entries == tuple(tool.index for tool in tools):
return edits
return (*edits, ReplacedTools(frozenset(tool.index for tool in tools), entries))

72
litellm/llms/bedrock/passthrough/transformation.py
View file

@ -13,6 +13,7 @@ from litellm._logging import verbose_logger
from litellm.constants import AWS_ERROR_TYPE_HEADER
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig, PassthroughStreamCollector
from litellm.llms.bedrock.chat.converse_to_openai import as_json_array, as_json_object
from litellm.types.llms.openai import ChatCompletionSystemMessage, ChatCompletionUserMessage
from litellm.types.utils import ModelResponseStream
@ -122,6 +123,71 @@ def _with_invoke_request_metadata(
return merge_bedrock_invoke_headers(dict(headers), (), metadata_headers, owned_names), body
def _is_cache_point(entry: object) -> bool:
fields: Final = as_json_object(entry)
return fields is not None and "cachePoint" in fields
def _tool_result_part(item: object) -> str | None:
fields: Final = as_json_object(item)
if fields is None:
return None
text: Final = fields.get("text")
if isinstance(text, str):
return text
return json.dumps(fields["json"]) if "json" in fields else None
def _tool_result_text(tool_result: Mapping[str, object]) -> str:
parts: Final = as_json_array(tool_result.get("content")) or ()
rendered: Final = "\n".join(part for item in parts if (part := _tool_result_part(item)) is not None).strip()
return rendered or "<non-text tool result omitted>"
def _block_field(block: object, key: str) -> Mapping[str, object] | None:
fields: Final = as_json_object(block)
return None if fields is None else as_json_object(fields.get(key))
def _as_text_block(block: object) -> object:
tool_use: Final = _block_field(block, "toolUse")
if tool_use is not None:
arguments: Final = json.dumps(tool_use.get("input"))
return {"text": f"[tool call {tool_use.get('toolUseId')}: {tool_use.get('name')}({arguments})]"}
tool_result: Final = _block_field(block, "toolResult")
if tool_result is None:
return block
return {"text": f"[tool result for {tool_result.get('toolUseId')}: {_tool_result_text(tool_result)}]"}
def _has_tool_block(content: Sequence[object]) -> bool:
return any(
_block_field(block, "toolUse") is not None or _block_field(block, "toolResult") is not None for block in content
)
def _as_text_message(message: object) -> object:
fields: Final = as_json_object(message)
content: Final = None if fields is None else as_json_array(fields.get("content"))
if fields is None or content is None or not _has_tool_block(content):
return message
return {**fields, "content": [_as_text_block(block) for block in content]}
def _converse_body_to_send(body: Mapping[str, object]) -> Mapping[str, object]:
"""
``body`` without a Converse ``toolConfig`` that holds no tool, with every ``toolUse`` and ``toolResult`` turned into
text, since Bedrock rejects tool blocks without a tool. Any other ``body`` comes back as it is.
"""
tool_config: Final = as_json_object(body.get("toolConfig"))
tools: Final = None if tool_config is None else as_json_array(tool_config.get("tools"))
if tools is None or not all(map(_is_cache_point, tools)):
return body
messages: Final = as_json_array(body.get("messages"))
rest: Final = {key: value for key, value in body.items() if key != "toolConfig"}
return rest if messages is None else {**rest, "messages": [_as_text_message(message) for message in messages]}
_TEXT_ONLY_DELTA_FIELDS: Final = frozenset({"content", "role"})
@ -378,11 +444,13 @@ class BedrockPassthroughConfig(BaseAWSLLM, BedrockModelInfo, BedrockEventStreamD
api_base: str,
model: str | None = None,
) -> tuple[dict, bytes | None]:
action: Final = _bedrock_endpoint_action(api_base)
body: Final[Mapping[str, object]] = request_data or {}
optional_params: Final = litellm_params.copy()
signed_headers, signed_body = _with_request_metadata(
action=_bedrock_endpoint_action(api_base) if bedrock_request_metadata_is_owned() else None,
action=action if bedrock_request_metadata_is_owned() else None,
headers=headers,
body=request_data or {},
body=_converse_body_to_send(body) if action in _BEDROCK_CONVERSE_ACTIONS else body,
litellm_params=litellm_params,
)
return self._sign_request(

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

@ -2314,6 +2314,7 @@ class BaseLLMHTTPHandler:
provider_specific_headers: Final = ProviderSpecificHeaderUtils.get_provider_specific_headers(
provider_specific_header=provider_specific_header,
custom_llm_provider=custom_llm_provider,
native_anthropic_messages=anthropic_messages_provider_config.is_native_messages_route(),
)
forwarded_headers: Final = kwargs.get("headers", None)
# Also check for extra_headers in kwargs (from config or direct calls)

2
litellm/llms/databricks/chat/transformation.py
View file

@ -278,7 +278,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
litellm_params: dict,
stream: bool | None = None,
) -> str:
use_ai_gateway: Final = model.removeprefix("databricks/").count(".") >= 2
use_ai_gateway: Final = self.uses_ai_gateway(model)
api_base = self._get_api_base(api_base, use_ai_gateway=use_ai_gateway)
url_base: Final = api_base.rstrip("/") if use_ai_gateway else api_base
complete_url: Final = f"{url_base}/chat/completions"

81
litellm/llms/databricks/common_utils.py
View file

@ -12,12 +12,16 @@ Authentication priority:
import os
import re
from collections.abc import Mapping
from types import MappingProxyType
from typing import Final, Literal
from urllib.parse import urlsplit, urlunsplit
from pydantic import BaseModel, ConfigDict
import litellm
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.secret_managers.main import get_secret_str
_DEFAULT_M2M_TOKEN_LIFETIME_SECONDS: Final = 3600
@ -26,6 +30,28 @@ class DatabricksException(BaseLLMException):
pass
DATABRICKS_ANTHROPIC_MESSAGES_PATH: Final = "/anthropic/v1/messages"
DATABRICKS_SERVING_ENDPOINTS_MESSAGES_PATH: Final = f"/serving-endpoints{DATABRICKS_ANTHROPIC_MESSAGES_PATH}"
DATABRICKS_AI_GATEWAY_ANTHROPIC_PATH: Final = "/ai-gateway/anthropic"
DATABRICKS_AI_GATEWAY_MESSAGES_PATH: Final = f"{DATABRICKS_AI_GATEWAY_ANTHROPIC_PATH}/v1/messages"
DATABRICKS_WORKSPACE_ROUTE_PREFIXES: Final = (
"/serving-endpoints",
"/serving-endpoints/anthropic",
DATABRICKS_SERVING_ENDPOINTS_MESSAGES_PATH,
"/ai-gateway/mlflow/v1",
DATABRICKS_AI_GATEWAY_ANTHROPIC_PATH,
DATABRICKS_AI_GATEWAY_MESSAGES_PATH,
)
DATABRICKS_NATIVE_MESSAGES_MODEL_PREFIXES: Final = ("databricks-claude-", "system.ai.claude-")
DATABRICKS_ENDPOINT_PATHS: Final = MappingProxyType(
{
"chat_completions": "/chat/completions",
"embeddings": "/embeddings",
"anthropic_messages": DATABRICKS_ANTHROPIC_MESSAGES_PATH,
}
)
class OAuthM2MToken(BaseModel):
model_config = ConfigDict(frozen=True)
@ -144,6 +170,47 @@ class DatabricksBase:
return redacted
@staticmethod
def resolve_api_key(api_key: str | None) -> str | None:
return api_key or litellm.api_key or litellm.databricks_key or get_secret_str("DATABRICKS_API_KEY")
@staticmethod
def resolve_api_base(api_base: str | None) -> str | None:
return api_base or litellm.api_base or os.getenv("DATABRICKS_API_BASE")
@staticmethod
def serves_native_anthropic_messages(model: str) -> bool:
return model.removeprefix("databricks/").lower().startswith(DATABRICKS_NATIVE_MESSAGES_MODEL_PREFIXES)
@staticmethod
def uses_ai_gateway(model: str) -> bool:
return model.removeprefix("databricks/").count(".") >= 2
@staticmethod
def recognizes_workspace_route(api_base: str, model: str) -> bool:
route: Final = urlsplit(api_base).path.rstrip("/")
return route.endswith(DATABRICKS_WORKSPACE_ROUTE_PREFIXES) or (
route == "" and DatabricksBase.uses_ai_gateway(model)
)
@staticmethod
def workspace_root(api_base: str) -> str:
normalized: Final = api_base.rstrip("/")
route_prefix: Final = next(
(prefix for prefix in DATABRICKS_WORKSPACE_ROUTE_PREFIXES if normalized.endswith(prefix)), ""
)
return normalized.removesuffix(route_prefix)
@staticmethod
def configured_user_agent(litellm_params: Mapping[str, object]) -> str | None:
configured: Final = (
litellm_params.get("user_agent")
or litellm_params.get("databricks_user_agent")
or os.getenv("LITELLM_USER_AGENT")
or os.getenv("DATABRICKS_USER_AGENT")
)
return configured if isinstance(configured, str) else None
@staticmethod
def _build_user_agent(custom_user_agent: str | None = None) -> str:
"""
@ -327,11 +394,11 @@ class DatabricksBase:
self,
api_key: str | None,
api_base: str | None,
endpoint_type: Literal["chat_completions", "embeddings"],
endpoint_type: Literal["chat_completions", "embeddings", "anthropic_messages"],
custom_endpoint: bool | None,
headers: dict | None,
custom_user_agent: str | None = None,
) -> tuple[str, dict]:
) -> tuple[str, dict[str, str]]: # mutable-ok: becomes the outbound headers
"""
Validate and configure the Databricks environment.
@ -343,7 +410,7 @@ class DatabricksBase:
Args:
api_key: Personal access token (PAT)
api_base: Databricks workspace URL with /serving-endpoints
endpoint_type: Type of endpoint (chat_completions or embeddings)
endpoint_type: Type of endpoint (chat_completions, embeddings or anthropic_messages)
custom_endpoint: Whether using a custom endpoint URL
headers: Existing headers dict
custom_user_agent: Optional custom user agent to prefix
@ -410,9 +477,5 @@ class DatabricksBase:
# Debug logging with redaction (never log actual tokens)
verbose_logger.debug("Databricks request headers: %s", self.redact_headers_for_logging(headers))
if endpoint_type == "chat_completions" and custom_endpoint is not True:
api_base = f"{api_base}/chat/completions"
elif endpoint_type == "embeddings" and custom_endpoint is not True:
api_base = f"{api_base}/embeddings"
return api_base, headers
endpoint_path: Final = "" if custom_endpoint is True else DATABRICKS_ENDPOINT_PATHS[endpoint_type]
return f"{api_base}{endpoint_path}", headers

194
litellm/llms/databricks/messages/transformation.py Normal file
View file

@ -0,0 +1,194 @@
import json
from collections.abc import AsyncIterator, Mapping
from typing import Final
import httpx
import litellm
from litellm.litellm_core_utils.core_helpers import as_str_mapping, iter_items
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.anthropic.pass_through.messages.transformation import (
AnthropicMessagesConfig,
)
from litellm.llms.databricks.common_utils import (
DATABRICKS_AI_GATEWAY_MESSAGES_PATH,
DATABRICKS_ANTHROPIC_MESSAGES_PATH,
DATABRICKS_SERVING_ENDPOINTS_MESSAGES_PATH,
DatabricksBase,
)
from litellm.types.llms.anthropic import ANTHROPIC_API_HEADERS
from litellm.types.router import GenericLiteLLMParams
DATABRICKS_UNSUPPORTED_SERVER_TOOL_PREFIXES: Final = ("web_search_", "web_fetch_", "code_execution_")
DROP_CONTAINER_WARNING: Final = (
"Dropping `container` for model=%s. The Databricks Anthropic Messages route rejects it, "
"and the chat completions bridge it replaces never forwarded it."
)
BEDROCK_INVOCATION_METRICS_FIELD: Final = "amazon-bedrock-invocationMetrics"
def _is_unsupported_server_tool(tool: object) -> bool:
tool_type: Final = (as_str_mapping(tool) or {}).get("type")
return isinstance(tool_type, str) and tool_type.startswith(DATABRICKS_UNSUPPORTED_SERVER_TOOL_PREFIXES)
def _tool_name(tool: object) -> object:
return (as_str_mapping(tool) or {}).get("name")
def _without_unsupported_server_tools(params: Mapping[str, object]) -> dict[str, object]:
if not isinstance(params.get("tools"), list):
return dict(params)
sent_tools: Final = tuple(iter_items(params.get("tools")))
kept_tools: Final = [tool for tool in sent_tools if not _is_unsupported_server_tool(tool)]
dropped_names: Final = frozenset(_tool_name(tool) for tool in sent_tools if _is_unsupported_server_tool(tool))
forced_tool: Final = _tool_name(params.get("tool_choice"))
keeps_tool_choice: Final = (
"tool_choice" in params and bool(kept_tools) and (forced_tool is None or forced_tool not in dropped_names)
)
remaining: Final = {key: value for key, value in params.items() if key not in ("tools", "tool_choice")}
return (
remaining
| ({"tools": kept_tools} if kept_tools else {})
| ({"tool_choice": params["tool_choice"]} if keeps_tool_choice else {})
)
def _without_invocation_metrics(sse_line: bytes) -> bytes:
if not sse_line.startswith(b"data:") or BEDROCK_INVOCATION_METRICS_FIELD.encode() not in sse_line:
return sse_line
event: Final = as_str_mapping(json.loads(sse_line.removeprefix(b"data:")))
if event is None:
return sse_line
kept_fields: Final = {key: value for key, value in event.items() if key != BEDROCK_INVOCATION_METRICS_FIELD}
return b"data: " + json.dumps(kept_fields, separators=(",", ":")).encode()
async def _strip_invocation_metrics(sse_stream: AsyncIterator[bytes]) -> AsyncIterator[bytes]:
partial_line: Final = bytearray()
async for chunk in sse_stream:
partial_line.extend(chunk)
if complete_end := partial_line.rfind(b"\n") + 1:
yield b"\n".join(
_without_invocation_metrics(line) for line in bytes(partial_line[:complete_end]).split(b"\n")
)
del partial_line[:complete_end]
if partial_line:
yield _without_invocation_metrics(bytes(partial_line))
class DatabricksAnthropicMessagesConfig(DatabricksBase, AnthropicMessagesConfig):
"""
Forwards /v1/messages for Claude models to the workspace's native Anthropic
route (``/serving-endpoints/anthropic/v1/messages`` for a serving endpoint,
``/ai-gateway/anthropic/v1/messages`` for a Unity Catalog model service), so
Anthropic-only request fields and response fields (``safeguards`` and
``safeguard_results``) survive instead of being bridged through chat completions.
"""
@property
def custom_llm_provider(self) -> str | None:
return "databricks"
def handles_web_search_natively(self) -> bool:
return False
def drops_safeguards(self, safeguards: object, drop_params: bool) -> bool:
return not self.supports_safeguards()
def serves_api_base(self, model: str, api_base: str | None) -> bool:
resolved_api_base: Final = self.resolve_api_base(api_base)
return resolved_api_base is None or self.recognizes_workspace_route(resolved_api_base, model)
def get_async_streaming_response_iterator(
self,
model: str,
httpx_response: httpx.Response,
request_body: dict[str, object], # mutable-ok: mirrors the base signature
litellm_logging_obj: LiteLLMLoggingObj,
) -> AsyncIterator[bytes]:
return _strip_invocation_metrics(
super().get_async_streaming_response_iterator(
model=model,
httpx_response=httpx_response,
request_body=request_body,
litellm_logging_obj=litellm_logging_obj,
)
)
def validate_anthropic_messages_environment(
self,
headers: dict[str, str], # mutable-ok: mirrors the base signature
model: str,
messages: list[object], # mutable-ok: mirrors the base signature
optional_params: dict[str, object], # mutable-ok: mirrors the base signature
litellm_params: dict[str, object], # mutable-ok: mirrors the base signature
api_key: str | None = None,
api_base: str | None = None,
) -> tuple[dict[str, str], str | None]: # mutable-ok: mirrors the base signature
client_anthropic_headers: Final = {
key: value for key, value in headers.items() if key.lower() in ANTHROPIC_API_HEADERS
}
messages_api_base, auth_headers = self.databricks_validate_environment( # pyright: ignore[reportUnknownMemberType] # the shared helper takes a bare dict
api_key=self.resolve_api_key(api_key),
api_base=self.resolve_api_base(api_base),
endpoint_type="anthropic_messages",
custom_endpoint=False,
headers={key: value for key, value in headers.items() if key not in client_anthropic_headers},
custom_user_agent=self.configured_user_agent(litellm_params),
)
workspace_api_base: Final = messages_api_base.removesuffix(DATABRICKS_ANTHROPIC_MESSAGES_PATH)
normalized_headers: Final = {
("content-type" if key.lower() == "content-type" else key): value
for key, value in {**auth_headers, **client_anthropic_headers}.items()
}
finalized_headers: Final[dict[str, str]] = self._finalize_messages_headers( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the shared helper returns a bare dict # mutable-ok: becomes the outbound headers
headers=normalized_headers,
optional_params=optional_params,
messages=messages,
)
return finalized_headers, workspace_api_base
def get_complete_url(
self,
api_base: str | None,
api_key: str | None,
model: str,
optional_params: dict[str, object], # mutable-ok: mirrors the base signature
litellm_params: dict[str, object], # mutable-ok: mirrors the base signature
stream: bool | None = None,
) -> str:
messages_path: Final = (
DATABRICKS_AI_GATEWAY_MESSAGES_PATH
if self.uses_ai_gateway(model)
else DATABRICKS_SERVING_ENDPOINTS_MESSAGES_PATH
)
return f"{self.workspace_root(self._get_api_base(api_base))}{messages_path}"
def transform_anthropic_messages_request(
self,
model: str,
messages: list[dict[str, object]], # mutable-ok: mirrors the base signature
anthropic_messages_optional_request_params: dict[str, object], # mutable-ok: mirrors the base signature
litellm_params: GenericLiteLLMParams,
headers: dict[str, str], # mutable-ok: mirrors the base signature
) -> dict[str, object]: # mutable-ok: mirrors the base signature
if "container" in anthropic_messages_optional_request_params:
litellm.verbose_logger.warning(DROP_CONTAINER_WARNING, model)
workspace_params: Final = {
key: value for key, value in anthropic_messages_optional_request_params.items() if key != "container"
}
drops_unsupported_params: Final = litellm.drop_params is True or litellm_params.drop_params is True
request_params: Final = (
_without_unsupported_server_tools(workspace_params) if drops_unsupported_params else workspace_params
)
anthropic_messages_request: Final[dict[str, object]] = super().transform_anthropic_messages_request( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the shared base returns a bare dict # mutable-ok: becomes the outbound request body
model=model,
messages=messages,
anthropic_messages_optional_request_params=request_params,
litellm_params=litellm_params,
headers=headers,
)
self._normalize_system_role_messages(anthropic_messages_request, model=model) # pyright: ignore[reportUnknownMemberType] # the shared helper takes a bare dict
self._remove_scope_from_cache_control(anthropic_messages_request) # pyright: ignore[reportUnknownMemberType] # the shared helper takes a bare dict
return anthropic_messages_request

7
litellm/llms/openai_like/messages/transformation.py
View file

@ -33,6 +33,9 @@ class OpenAILikeAnthropicMessagesConfig(AnthropicMessagesConfig):
super().__init__()
self._cache_control_ttl: Final = cache_control_ttl
def is_native_messages_route(self) -> bool:
return False
def validate_anthropic_messages_environment(
self,
headers: dict[str, str],
@ -59,6 +62,7 @@ class OpenAILikeAnthropicMessagesConfig(AnthropicMessagesConfig):
optional_params=optional_params,
messages=messages,
beta_map_provider=LlmProviders.OPENAI_LIKE.value,
drop_params=litellm_params.get("drop_params") is True, # pyright: ignore[reportUnknownMemberType] # the litellm params arrive as a bare dict
)
return merged, api_base
@ -132,6 +136,9 @@ class JSONProviderAnthropicMessagesConfig(OpenAILikeAnthropicMessagesConfig):
def should_strip_billing_metadata(self) -> bool:
return True
def supports_safeguards(self) -> bool:
return False
def _resolve_api_key(self, api_key: str | None) -> str | None:
return api_key or get_secret_str(self._provider.api_key_env) or litellm.api_key

21
litellm/main.py
View file

@ -239,6 +239,7 @@ from .llms.cohere.embed import handler as cohere_embed
from .llms.custom_httpx.aiohttp_handler import BaseLLMAIOHTTPHandler
from .llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from .llms.custom_llm import CustomLLM, custom_chat_llm_router
from .llms.databricks.common_utils import DatabricksBase
from .llms.databricks.embed.handler import DatabricksEmbeddingHandler
from .llms.deprecated_providers import aleph_alpha
from .llms.gdc.chat.transformation import GDCGeminiConfig
@ -3361,8 +3362,8 @@ def _complete_oobabooga(ctx: CompletionDispatchContext) -> _CompletionDispatchRe
def _complete_databricks(ctx: CompletionDispatchContext) -> _CompletionDispatchResult:
acompletion: Final = ctx.acompletion
api_base = ctx.api_base
api_key = ctx.api_key
api_base: Final = DatabricksBase.resolve_api_base(ctx.api_base)
api_key: Final = DatabricksBase.resolve_api_key(ctx.api_key)
client: Final = _dispatch_client_http(ctx)
headers = ctx.headers
litellm_params: Final = ctx.litellm_params
@ -3374,20 +3375,6 @@ def _complete_databricks(ctx: CompletionDispatchContext) -> _CompletionDispatchR
stream: Final = ctx.stream
timeout: Final = ctx.timeout
api_base = (
api_base # for databricks we check in get_llm_provider and pass in the api base from there
or litellm.api_base
or os.getenv("DATABRICKS_API_BASE")
)
# set API KEY
api_key = (
api_key
or litellm.api_key # for databricks we check in get_llm_provider and pass in the api key from there
or litellm.databricks_key
or get_secret("DATABRICKS_API_KEY")
)
headers = headers or litellm.headers
## COMPLETION CALL
@ -6832,7 +6819,7 @@ def embedding(
api_base = api_base or litellm.api_base or get_secret("DATABRICKS_API_BASE")
# set API KEY
api_key = api_key or litellm.api_key or litellm.databricks_key or get_secret("DATABRICKS_API_KEY")
api_key = DatabricksBase.resolve_api_key(api_key)
## EMBEDDING CALL
response = databricks_embedding.embedding(

25
litellm/model_prices_and_context_window_backup.json
View file

@ -73794,14 +73794,14 @@
"input_cost_per_token_flex": 1e-06,
"input_cost_per_token_priority": 3.5e-06,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 500000,
"max_output_tokens": 500000,
"max_tokens": 500000,
"max_input_tokens": 524288,
"max_output_tokens": 524288,
"max_tokens": 524288,
"mode": "chat",
"output_cost_per_token": 6e-06,
"output_cost_per_token_flex": 3e-06,
"output_cost_per_token_priority": 1.05e-05,
"source": "https://b0.p.awsstatic.com/pricing/2.0/meteredUnitMaps/bedrock/USD/current/bedrock.json",
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-xai-grok-4-7.html",
"supports_function_calling": true,
"supports_prompt_caching": false,
"supports_reasoning": true,
@ -73820,14 +73820,14 @@
"input_cost_per_token_flex": 1.1e-06,
"input_cost_per_token_priority": 3.85e-06,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 500000,
"max_output_tokens": 500000,
"max_tokens": 500000,
"max_input_tokens": 524288,
"max_output_tokens": 524288,
"max_tokens": 524288,
"mode": "chat",
"output_cost_per_token": 6.6e-06,
"output_cost_per_token_flex": 3.3e-06,
"output_cost_per_token_priority": 1.155e-05,
"source": "https://b0.p.awsstatic.com/pricing/2.0/meteredUnitMaps/bedrock/USD/current/bedrock.json",
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-xai-grok-4-7.html",
"supports_function_calling": true,
"supports_prompt_caching": false,
"supports_reasoning": true,
@ -75198,31 +75198,34 @@
"us.twelvelabs.pegasus-1-5-v1:0": {
"input_cost_per_video_per_second": 0.00049,
"litellm_provider": "bedrock",
"max_input_tokens": 261120,
"max_output_tokens": 98304,
"max_tokens": 98304,
"mode": "chat",
"output_cost_per_token": 7.5e-06,
"source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json",
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-twelvelabs-pegasus-v1-5.html",
"supports_video_input": true
},
"global.twelvelabs.pegasus-1-5-v1:0": {
"input_cost_per_video_per_second": 0.00049,
"litellm_provider": "bedrock",
"max_input_tokens": 261120,
"max_output_tokens": 98304,
"max_tokens": 98304,
"mode": "chat",
"output_cost_per_token": 7.5e-06,
"source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json",
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-twelvelabs-pegasus-v1-5.html",
"supports_video_input": true
},
"twelvelabs.pegasus-1-5-v1:0": {
"input_cost_per_video_per_second": 0.00049,
"litellm_provider": "bedrock",
"max_input_tokens": 261120,
"max_output_tokens": 98304,
"max_tokens": 98304,
"mode": "chat",
"output_cost_per_token": 7.5e-06,
"source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json",
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-twelvelabs-pegasus-v1-5.html",
"supports_video_input": true
},
"azure_ai/Microsoft-Decision-1": {

30
litellm/proxy/auth/auth_checks.py
View file

@ -4803,6 +4803,24 @@ def check_model_access_helper(
_check_model_access_helper: Final = check_model_access_helper
def _owning_team_names_for_team_deployment(
model: str,
llm_router: Router | None,
team_id: str | None,
team_model_aliases: Mapping[str, str] | None,
) -> tuple[str, ...]:
if llm_router is None or team_id is None or not model.startswith(f"model_name_{team_id}_"):
return ()
deployment: Final = llm_router.get_deployment_by_model_group_name(model_group_name=model)
public_name: Final = (
deployment.model_info.team_public_model_name
if deployment is not None and deployment.model_info.team_id == team_id
else None
)
legacy_aliases: Final = tuple(alias for alias, target in (team_model_aliases or {}).items() if target == model)
return (*((public_name,) if public_name else ()), *legacy_aliases)
def _can_object_call_model(
model: str | list[str],
llm_router: Router | None,
@ -4874,6 +4892,13 @@ def _can_object_call_model(
*((global_or_router_alias_target,) if global_or_router_alias_target else ()),
)
)
owning_team_names: Final = (
()
if key_alias_applied
else _owning_team_names_for_team_deployment(
model=model, llm_router=llm_router, team_id=team_id, team_model_aliases=team_model_aliases
)
)
## check model access for alias + underlying model - allow if either is in allowed models
for m in potential_models:
@ -4885,6 +4910,11 @@ def _can_object_call_model(
team_id=team_id,
):
return True
if any(
check_model_access_helper(model=name, llm_router=llm_router, models=models, team_id=team_id)
for name in owning_team_names
):
return True
internal_message: Final = (
f"{object_type} not allowed to access model. This {object_type} can only access models={models}. "

3
litellm/proxy/guardrails/guardrail_hooks/airia/airia.py
View file

@ -128,6 +128,9 @@ class AiriaGuardrail(CustomGuardrail):
)
verbose_proxy_logger.info("AiriaGuardrail initialized with gateway: %s", self.api_base)
def supports_structured_messages(self) -> bool:
return False
@log_guardrail_information
async def apply_guardrail(
self,

4
litellm/proxy/guardrails/guardrail_hooks/akto/akto.py
View file

@ -794,6 +794,10 @@ class AktoGuardrail(CustomGuardrail):
raise failure
return main_task.result()
@override
def supports_structured_messages(self) -> bool:
return False
@override
@log_guardrail_information
async def apply_guardrail(

3
litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py
View file

@ -586,6 +586,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
def supports_scan_only_tool_results(self) -> bool:
return self.experimental_use_latest_role_message_only is not True
def supports_structured_messages(self) -> bool:
return False
def _prepare_guardrail_messages_for_role(
self,
messages: list[AllMessageValues] | None,

3
litellm/proxy/guardrails/guardrail_hooks/conduct/conduct.py
View file

@ -125,6 +125,9 @@ else:
def get_config_model() -> type[ConductGuardrailConfigModel]:
return ConductGuardrailConfigModel
def supports_structured_messages(self) -> bool:
return False
@log_guardrail_information
async def apply_guardrail(
self,

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

@ -474,6 +474,10 @@ class CrowdStrikeAIDRHandler(CustomGuardrail):
or effective_skip_assistant_message_for_guardrail(self)
)
@override
def supports_structured_messages(self) -> bool:
return False
def _writeback_messages(
self,
structured_messages: list[AllMessageValues],

6
litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py
View file

@ -262,6 +262,9 @@ class HiddenlayerGuardrail(CustomGuardrail):
self._http_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback)
super().__init__(**kwargs)
def supports_structured_messages(self) -> bool:
return False
@log_guardrail_information
async def apply_guardrail(
self,
@ -472,6 +475,9 @@ class HiddenlayerGuardrailV2(CustomGuardrail):
self._http_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback)
super().__init__(**kwargs)
def supports_structured_messages(self) -> bool:
return False
@log_guardrail_information
async def apply_guardrail(
self,

3
litellm/proxy/guardrails/guardrail_hooks/llm_as_a_judge/__init__.py
View file

@ -256,6 +256,9 @@ class LLMAsAJudgeGuardrail(CustomGuardrail):
raw: Final = response.choices[0].message.content or "{}"
return _parse_judge_verdict(raw)
def supports_structured_messages(self) -> bool:
return False
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,

3
litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py
View file

@ -200,6 +200,9 @@ class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail):
},
)
def supports_structured_messages(self) -> bool:
return False
@log_guardrail_information
async def apply_guardrail(
self,

3
litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py
View file

@ -1744,6 +1744,9 @@ class PanwPrismaAirsHandler(CustomGuardrail):
def supports_scan_only_tool_results(self) -> bool:
return False
def supports_structured_messages(self) -> bool:
return False
@classmethod
def _get_scannable_text_indices(
cls,

3
litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py
View file

@ -170,6 +170,9 @@ class PromptSecurityGuardrail(CustomGuardrail):
def supports_scan_only_tool_results(self) -> bool:
return self.check_tool_results
def supports_structured_messages(self) -> bool:
return False
@log_guardrail_information
async def apply_guardrail(
self,

3
litellm/proxy/guardrails/guardrail_hooks/promptguard/promptguard.py
View file

@ -119,6 +119,9 @@ class PromptGuardGuardrail(CustomGuardrail):
GuardrailEventHooks.post_call,
]
def supports_structured_messages(self) -> bool:
return False
@log_guardrail_information
async def apply_guardrail(
self,

3
litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py
View file

@ -420,6 +420,9 @@ class QualifireGuardrail(CustomGuardrail):
verbose_proxy_logger.exception("Qualifire Guardrail error: %s", e)
raise
def supports_structured_messages(self) -> bool:
return False
@log_guardrail_information
async def apply_guardrail(
self,

3
litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py
View file

@ -402,6 +402,9 @@ class SingulrGuardrail(CustomGuardrail):
)
return inputs
def supports_structured_messages(self) -> bool:
return False
@log_guardrail_information
async def apply_guardrail(
self,

24
litellm/proxy/litellm_pre_call_utils.py
View file

@ -4089,13 +4089,25 @@ def add_provider_specific_headers_to_request(
}
scoped_headers: Final = [
ProviderSpecificHeader(custom_llm_provider=providers, extra_headers=extra_headers)
for providers, extra_headers in (
(_ANTHROPIC_API_HEADER_PROVIDERS, anthropic_api_headers),
(_ANTHROPIC_OAUTH_CREDENTIAL_PROVIDERS, anthropic_oauth_credential_headers),
(LlmProviders.TYPESAFE.value, typesafe_danger_zone_headers),
scoped_header
for scoped_header in (
ProviderSpecificHeader(
custom_llm_provider=_ANTHROPIC_API_HEADER_PROVIDERS, extra_headers=anthropic_api_headers
),
ProviderSpecificHeader(
custom_llm_provider=LlmProviders.DATABRICKS.value,
extra_headers=anthropic_api_headers,
native_anthropic_messages_only=True,
),
ProviderSpecificHeader(
custom_llm_provider=_ANTHROPIC_OAUTH_CREDENTIAL_PROVIDERS,
extra_headers=anthropic_oauth_credential_headers,
),
ProviderSpecificHeader(
custom_llm_provider=LlmProviders.TYPESAFE.value, extra_headers=typesafe_danger_zone_headers
),
)
if extra_headers
if scoped_header["extra_headers"]
]
if scoped_headers:

14
litellm/proxy/policy_engine/pipeline_executor.py
View file

@ -26,6 +26,7 @@ from litellm.litellm_core_utils.core_helpers import (
get_or_create_metadata_bucket,
independent_snapshot,
)
from litellm.llms.base_llm.guardrail_translation.utils import UnappliableRequestRewrite
from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
UnifiedLLMGuardrails,
@ -410,7 +411,7 @@ class PipelineExecutor:
duration = time.perf_counter() - start_time
action = _pipeline_action_for_outcome(step, outcome)
action = _pipeline_action_for_step(step, outcome, original_exception)
step_result = PipelineStepResult(
guardrail_name=step.guardrail,
@ -748,6 +749,17 @@ def _carry_working_guardrail_information(
_append_guardrail_information(request_data=request_data, entries=[e for e in recorded if e not in existing])
def _pipeline_action_for_step(step: PipelineStep, outcome: str, original_exception: Exception | None) -> str:
"""
The action a step takes: block when its guardrail returned a rewrite the endpoint
cannot apply, whatever on_fail and on_error say, else the action configured for
the outcome.
"""
if isinstance(original_exception, UnappliableRequestRewrite):
return "block"
return _pipeline_action_for_outcome(step, outcome)
def _pipeline_action_for_outcome(step: PipelineStep, outcome: str) -> str:
"""
Map pipeline step outcome to the configured action.

1
litellm/types/utils.py
View file

@ -4538,6 +4538,7 @@ class HttpHandlerRequestFields(TypedDict, total=False):
class ProviderSpecificHeader(TypedDict):
custom_llm_provider: str
extra_headers: dict
native_anthropic_messages_only: NotRequired[ReadOnly[bool]]
class SelectTokenizerResponse(TypedDict):

9
litellm/utils.py
View file

@ -9367,6 +9367,15 @@ class ProviderConfigManager:
)
return AzureAnthropicMessagesConfig()
elif litellm.LlmProviders.DATABRICKS == provider:
from litellm.llms.databricks.common_utils import DatabricksBase
if DatabricksBase.serves_native_anthropic_messages(model):
from litellm.llms.databricks.messages.transformation import (
DatabricksAnthropicMessagesConfig,
)
return DatabricksAnthropicMessagesConfig()
elif litellm.LlmProviders.MINIMAX == provider:
from litellm.llms.minimax.messages.transformation import (
MinimaxMessagesConfig,

25
model_prices_and_context_window.json
View file

@ -73794,14 +73794,14 @@
"input_cost_per_token_flex": 1e-06,
"input_cost_per_token_priority": 3.5e-06,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 500000,
"max_output_tokens": 500000,
"max_tokens": 500000,
"max_input_tokens": 524288,
"max_output_tokens": 524288,
"max_tokens": 524288,
"mode": "chat",
"output_cost_per_token": 6e-06,
"output_cost_per_token_flex": 3e-06,
"output_cost_per_token_priority": 1.05e-05,
"source": "https://b0.p.awsstatic.com/pricing/2.0/meteredUnitMaps/bedrock/USD/current/bedrock.json",
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-xai-grok-4-7.html",
"supports_function_calling": true,
"supports_prompt_caching": false,
"supports_reasoning": true,
@ -73820,14 +73820,14 @@
"input_cost_per_token_flex": 1.1e-06,
"input_cost_per_token_priority": 3.85e-06,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 500000,
"max_output_tokens": 500000,
"max_tokens": 500000,
"max_input_tokens": 524288,
"max_output_tokens": 524288,
"max_tokens": 524288,
"mode": "chat",
"output_cost_per_token": 6.6e-06,
"output_cost_per_token_flex": 3.3e-06,
"output_cost_per_token_priority": 1.155e-05,
"source": "https://b0.p.awsstatic.com/pricing/2.0/meteredUnitMaps/bedrock/USD/current/bedrock.json",
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-xai-grok-4-7.html",
"supports_function_calling": true,
"supports_prompt_caching": false,
"supports_reasoning": true,
@ -75198,31 +75198,34 @@
"us.twelvelabs.pegasus-1-5-v1:0": {
"input_cost_per_video_per_second": 0.00049,
"litellm_provider": "bedrock",
"max_input_tokens": 261120,
"max_output_tokens": 98304,
"max_tokens": 98304,
"mode": "chat",
"output_cost_per_token": 7.5e-06,
"source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json",
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-twelvelabs-pegasus-v1-5.html",
"supports_video_input": true
},
"global.twelvelabs.pegasus-1-5-v1:0": {
"input_cost_per_video_per_second": 0.00049,
"litellm_provider": "bedrock",
"max_input_tokens": 261120,
"max_output_tokens": 98304,
"max_tokens": 98304,
"mode": "chat",
"output_cost_per_token": 7.5e-06,
"source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json",
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-twelvelabs-pegasus-v1-5.html",
"supports_video_input": true
},
"twelvelabs.pegasus-1-5-v1:0": {
"input_cost_per_video_per_second": 0.00049,
"litellm_provider": "bedrock",
"max_input_tokens": 261120,
"max_output_tokens": 98304,
"max_tokens": 98304,
"mode": "chat",
"output_cost_per_token": 7.5e-06,
"source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json",
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-twelvelabs-pegasus-v1-5.html",
"supports_video_input": true
},
"azure_ai/Microsoft-Decision-1": {

4
pyproject.toml
View file

@ -77,8 +77,8 @@ proxy = [
"mcp>=2.2.0,<3",
"httpx2>=2.5.0,<3",
"pydantic>=2.12.0,<3",
"litellm-proxy-extras==0.4.109",
"litellm-enterprise==0.1.77",
"litellm-proxy-extras==0.4.110",
"litellm-enterprise==0.1.78",
"RestrictedPython>=8.5,<9.0",
"rich>=13.9.4,<14.0",
"InquirerPy>=0.3.4,<1.0",

92
tests/e2e/perf/UI_QUERIES.md Normal file
View file

@ -0,0 +1,92 @@
# SpendLog UI query performance
Run `test_ui_queries.py` against a dedicated proxy and a disposable database restored from the metadata-only performance seed. These tests do not call an LLM or measure browser rendering
The suite measures 42 request shapes across 20 routes used by current dashboard screens. It replays direct HTTP requests against the existing bulk seed, with the seed anchor replacing the wall clock and a UTC browser timezone. It does not run a browser or claim to cover every dashboard request
## UI callers
All paths below are relative to `ui/litellm-dashboard/src/`. Each case requires an authenticated admin session with access to that screen, and selected resources must exist in the fixture
| Cases | UI action and source |
| --- | --- |
| `logs_*` | Logs > Request Logs, change date range, filter, sort or next page. `components/logs/request/useLogFilterLogic.ts` and `RequestLogsFilters.tsx` |
| `usage_*` | Usage overview or entity tab, open key activity, search keys, expand a model group. `app/(dashboard)/usage/_components/components/UsagePageView.tsx`, `EntityUsage/EntityUsage.tsx`, `EntityUsage/entityFetchFns.ts`, `components/UsagePage/components/KeyActivityPanel.tsx` |
| `team_users_*` | Team Usage. `EntityUsage/EntityUsage.tsx` passes selected teams or all available non-dashboard teams to `TeamUserSpendCard.tsx` |
| `cache_activity_*` | Caching > Cache Analytics, optionally filter aliases/models. `app/(dashboard)/hooks/caching/useCacheActivity.ts` |
| `cache_leakage_keys` | Cost Optimization, switch cache leakage to Keys. `app/(dashboard)/cost-optimization/_components/useCacheLeakageKeys.ts` |
Request serialization lives in `components/networking.tsx`. Logs use page size 25 and session grouping, including the next-session cursor returned by page one. Individual Logs filters use a 30-day window to exercise larger matching datasets. Their case names end in `_30d` so the measurements do not share history with the old 24-hour filters. The synthetic cursorless page-three request is excluded
Usage list, search and model drilldown requests do not inject an `api_key` filter. Team activity excludes `litellm-dashboard`. Cache leakage supplies `include_current_utc_day=true`. Dates and filter values are deterministic fixture values, not captured production traffic
Old Usage is excluded because its expensive-query gate normally hides those calls at this database size. LiteAdmin's agent-selected reports are excluded from this deterministic dashboard set. Missing-token alias recovery is backend functional coverage, not a separate UI action, and its artificial requests and orphan fixtures are removed
This mapping is source-verified, not browser-recorded. Recheck the named callers when changing cases. Neither browser capture nor an automatic UI-drift checker is installed by this suite
## Dataset scope
Logs, team-user spend and cache activity query the bulk SpendLogs data. User, team, tag and customer usage, plus cache leakage, query daily tables aggregated from that dataset. These are sequential large-dataset queries, not concurrent traffic or saturation tests
A read-only restore of `db-perf-v1-105-1tb` gave exact request totals from daily aggregates: 9,442,606 across 24 hours, 66,077,435 across seven days and 283,189,684 across 30 days. The retained 30-day key filter matches 1,134,559 requests, user filter 11,316,454, team filter 113,298,119 and model filter 94,395,473. These describe matching requests, not rows scanned or returned by the API, and are documentation rather than performance assertions
PostgreSQL table-wide estimates (`pg_class.reltuples`) are approximately 2.70 million daily user rows, 2.70 million daily team rows, 5.40 million daily tag rows and 78.27 million daily end-user rows. These are not exact counts within each Usage case's seven-day window
The individual end-user filter is excluded because it matched only 104 requests in 24 hours and 8,616 across all dates. User/end-user dropdown requests are also excluded because their queries cap the input at 10,000 matching SpendLogs before extracting distinct values. Neither is a broad matching-data workload in this seed
Excluded because they only had tiny matching fixtures or individual-record lookups: organization/agent usage, guardrail/tool logs, request-ID lookup, payload/session drawers, session/search filters for `perf-ui-*`, the OAuth-token filter, and all prompt-caching request cases. The seed has response-cache hits but no prompt-cache usage or injection metadata, so the prompt-caching cases only matched the 320 added rows. No UI-specific rows are inserted anymore
| Remaining cases | Count | Variations |
| --- | ---: | --- |
| Logs | 21 | 24h/7d/30d, page 2 with session cursor, 12 individual filters over 30 days, five additional descending sorts |
| User/team/tag/customer usage | 16 | Aggregated totals, keys, key search, model top keys for each |
| Team-user spend | 2 | All 100 seeded teams or `perf-team-0` |
| Cache activity | 2 | All or selected aliases/model |
| Cache leakage keys | 1 | 30-day window |
## Deployment prerequisites
Use a fresh proxy against the disposable restore, never the active seed source or a shared test database. Apply the target migrations and wait for readiness first. The source seed's `litellm_perf_seed.checkpoint` must exist and contain rows
Configure `litellm_settings.service_callback: [prometheus_system]` and expose metrics for every proxy worker/pod. Set `E2E_DATABASE_METRICS_URLS` to their comma-separated URLs, not a load-balanced endpoint that can miss a pod. Absent failure counters count as zero only when the endpoint exposes a positive PostgreSQL request counter. Missing instrumentation stops the run
No expensive-query override or legacy view creation is needed for these cases. Keep the normal large-database UI restrictions
The cases use explicit dates from the seed anchor where the UI supplies dates. A missing nonempty fixture result records a contract failure rather than reporting a fast empty query as success
## Run
Set `DATABASE_URL` and `LITELLM_MASTER_KEY` securely for this disposable deployment. The URL must name the same database used by the proxy, with hostname `db-perf-perf-<run>.litellm.svc`. The explicit disposable confirmation permits cancellation of leftover queries in this database. Setup reads the existing seed checkpoint without adding UI-specific rows. Destroy the disposable database after the run
```bash
export LITELLM_PROXY_URL=http://localhost:4000
export E2E_DATABASE_METRICS_URLS=http://localhost:4000/metrics
export E2E_PERF_TESTS=1
export E2E_UI_PERF_DISPOSABLE=1
export E2E_RESET_SPEND_LOGS=0
export E2E_ARTIFACT_DIR="$PWD/ui-query-results"
mkdir -p "$E2E_ARTIFACT_DIR"
uv run pytest tests/e2e/perf/test_ui_queries.py --reruns 0 --junitxml="$E2E_ARTIFACT_DIR/junit.xml" -q
```
Do not enable SpendLog retention or run other traffic on this stack. The `quiet_stack` marker coordinates local pytest callers but cannot stop external traffic. Use one pytest worker for comparable timings
`E2E_UI_QUERY_REPEATS` defaults to 9, giving ten measured requests including the first, and `E2E_UI_QUERY_TIMEOUT_SECONDS` defaults to 120. Every request is issued once without transport retries. A failed first request stops that case's repeated samples. Cursor prerequisites are recorded separately, and failed prerequisites stop that case's measurements
## Reading results
Each case writes `ui-query-<name>.json` with its exact method, path, parameters and body, response status counts, elapsed times, row counts, errors and before/after PostgreSQL failure counters. Response bodies and authentication headers are not saved
The dashboard's `perf_ui_seconds_p50`, `_p90` and `_max` include all successful measured requests, including the first. `perf_ui_samples` is their sample count. p50 is the median and p90 uses nearest rank, the ninth sorted value with ten successes. Failures and cursor prerequisites contribute no latency. Pass percentage is `(perf_ui_requests - perf_ui_failures) / perf_ui_requests * 100`, including cursor prerequisites in the request counts
The first-request and repeated-request metrics are retained for existing consumers. This is not a cold-cache claim: previous cases and previous runs can warm database buffers or application caches. Restart the disposable proxy between independent runs when comparing independent runs
Repeated p50 is reported for successful samples, p95 only with at least 20 and p99 only with at least 100. First-request latency is omitted when that request fails. HTTP, JSON, response-shape, nonempty-result, cursor failures are recorded, never asserted. Endpoint failures, database counter increases and latency do not fail CI
Each JUnit testcase `test_ui_query_performance[<case>]` carries `perf_ui_requests`, `perf_ui_failures`, `perf_ui_repeated_samples`, and nonzero `perf_ui_error_<code>` counts. Codes distinguish `http_400`, `http_500`, `http_503`, other HTTP statuses, `timeout`, `network`, `invalid_json` and `contract`. Successful latency properties are `perf_ui_first_seconds` and `perf_ui_repeated_seconds_p50`, `_p95`, `_p99`, `_max`. A failed cursor prerequisite is counted once and contributes no latency
Before every case the runner requires three idle database observations and stable failure counters. After each case it cancels leftover SQL, terminates idle transactions and waits for idle and stable counters again. `ui-isolation-<case>.jsonl` retains active SQL, waits, elapsed times, raw samples and counter observations. `perf_ui_database_failures` measures the increase before cleanup, while `perf_ui_cleanup_database_failures` measures the increase observed afterward, including cancellations and potentially delayed request failures. Isolation, instrumentation or setup failures still stop the run rather than publish contaminated measurements. Disk and process caches are not reset between cases
The Buildkite `ui_queries` lane restores the pinned snapshot independently, waits for migrations and indexes to complete, runs these cases with no retries, uploads JUnit and JSON artifacts, and removes the temporary resources. It follows the nightly regular perf and database-upgrade builds using the same LiteLLM SHA. Shinkiro ingestion is a separate change
Validate a small restore first. Small-database numbers validate the runner only, not full-size performance. Run the same cases against the completed full-size snapshot before publishing large-database results

227
tests/e2e/perf/setup/ui_queries.py Normal file
View file

@ -0,0 +1,227 @@
from __future__ import annotations
import math
import statistics
import time
from collections.abc import Callable
from dataclasses import dataclass
from itertools import chain
from typing import Final, Literal
from urllib.parse import urlencode
from e2e_http import (
NetworkError,
RateLimitedError,
Result,
Success,
UnauthorizedError,
UnknownApiError,
ValidationError,
without_retries,
)
from e2e_metadata import step
from prometheus_client.parser import text_string_to_metric_families
from pydantic import BaseModel, ConfigDict, JsonValue, RootModel
from transport import HttpTransport
class Query(BaseModel):
model_config = ConfigDict(frozen=True)
class Document(RootModel[JsonValue]):
model_config = ConfigDict(frozen=True)
@dataclass(frozen=True, slots=True)
class Contract:
collection: tuple[str, ...] = ()
row_fields: tuple[str, ...] = ()
positive: tuple[str, ...] = ()
object_fields: tuple[str, ...] = ()
excluded_ids: frozenset[str] = frozenset()
def count(self, document: Document) -> int:
value: Final = descend(document.root, self.collection)
if self.object_fields:
assert isinstance(value, dict), "Expected a JSON object"
assert set(self.object_fields) <= value.keys(), f"Missing fields {self.object_fields}"
return 1
assert isinstance(value, list) and value, f"Empty/missing collection {self.collection}, check the fixture"
for row in value:
if self.excluded_ids:
assert isinstance(row, dict) and isinstance(row_id := row.get("request_id"), str), "Missing request_id"
assert row_id not in self.excluded_ids, "Repeated page-one row"
if self.row_fields or self.positive:
assert isinstance(row, dict), "Expected object rows"
assert set(self.row_fields) <= row.keys(), f"Missing row fields {self.row_fields}"
for name in self.positive:
assert any(tuple(_positive(row, name) for row in value)), f"Expected positive {name}"
return len(value)
def _positive(row: JsonValue, name: str) -> bool:
assert isinstance(row, dict), "Expected object row"
number: Final = row.get(name)
assert isinstance(number, (int, float)) and not isinstance(number, bool), f"Expected numeric {name}"
assert math.isfinite(number), f"Expected finite {name}"
return number > 0
def descend(value: JsonValue, path: tuple[str, ...]) -> JsonValue:
if not path:
return value
assert isinstance(value, dict) and path[0] in value, f"Missing response path {path}"
return descend(value[path[0]], path[1:])
@dataclass(frozen=True, slots=True)
class UiQueryCase:
name: str
path: str
params: BaseModel
contract: Contract
surface: Literal["current"] = "current"
body: BaseModel | None = None
cursor: Literal["none", "session"] = "none"
class Sample(BaseModel):
model_config = ConfigDict(frozen=True)
seconds: float
status: int
rows: int = 0
error: str = ""
class Timings(BaseModel):
model_config = ConfigDict(frozen=True)
first_seconds: float | None
repeated_samples: int
repeated_p50_seconds: float | None
repeated_p95_seconds: float | None
repeated_p99_seconds: float | None
repeated_max_seconds: float | None
def summarize(samples: tuple[Sample, ...]) -> Timings:
assert samples, "No UI requests were measured"
repeated: Final = tuple(sorted(sample.seconds for sample in samples[1:] if not sample.error))
return Timings(
first_seconds=samples[0].seconds if not samples[0].error else None,
repeated_samples=len(repeated),
repeated_p50_seconds=statistics.median(repeated) if repeated else None,
repeated_p95_seconds=repeated[math.ceil(len(repeated) * 0.95) - 1] if len(repeated) >= 20 else None,
repeated_p99_seconds=repeated[math.ceil(len(repeated) * 0.99) - 1] if len(repeated) >= 100 else None,
repeated_max_seconds=max(repeated) if repeated else None,
)
def error_code(sample: Sample) -> str:
if not sample.error:
return ""
if sample.status > 0 and sample.status != 200:
return f"http_{sample.status}"
if sample.error in ("timeout", "network"):
return sample.error
if sample.error == "Response is not valid JSON":
return "invalid_json"
return "contract"
def ui_junit_properties(
samples: tuple[Sample, ...], prerequisite: Sample | None = None
) -> tuple[tuple[str, float], ...]:
requests: Final = (prerequisite, *samples) if prerequisite is not None and not prerequisite.error else samples
codes: Final = tuple(error_code(sample) for sample in requests if sample.error)
timings: Final = summarize(samples)
successful: Final = tuple(sorted(sample.seconds for sample in samples if not sample.error))
latency: Final = (
("seconds_p50", statistics.median(successful) if successful else None),
("seconds_p90", successful[math.ceil(len(successful) * 0.90) - 1] if successful else None),
("seconds_max", max(successful) if successful else None),
("first_seconds", timings.first_seconds),
("repeated_seconds_p50", timings.repeated_p50_seconds),
("repeated_seconds_p95", timings.repeated_p95_seconds),
("repeated_seconds_p99", timings.repeated_p99_seconds),
("repeated_seconds_max", timings.repeated_max_seconds),
)
return (
("perf_ui_requests", len(requests)),
("perf_ui_failures", len(codes)),
("perf_ui_samples", len(successful)),
("perf_ui_repeated_samples", timings.repeated_samples),
*((f"perf_ui_error_{code}", codes.count(code)) for code in sorted(frozenset(codes))),
*((f"perf_ui_{name}", value) for name, value in latency if value is not None),
)
def _failure(result: Result[Document], seconds: float) -> Sample:
match result:
case Success(status_code=status):
return Sample(seconds=seconds, status=status, error=f"Unexpected HTTP {status}")
case UnknownApiError(status_code=status):
return Sample(seconds=seconds, status=status, error=f"HTTP {status}")
case UnauthorizedError():
return Sample(seconds=seconds, status=401, error="Unauthorized")
case RateLimitedError():
return Sample(seconds=seconds, status=429, error="Rate limited")
case NetworkError(message=message):
return Sample(seconds=seconds, status=-1, error="timeout" if "timed out" in message.lower() else "network")
case ValidationError():
return Sample(seconds=seconds, status=-1, error="Response is not valid JSON")
def measure(
issue: Callable[[], Result[Document]],
contract: Contract,
*,
clock: Callable[[], float] = time.monotonic,
) -> tuple[Sample, Document | None]:
started: Final = clock()
with without_retries():
result: Final = issue()
seconds: Final = clock() - started
match result:
case Success(status_code=200, data=document):
try:
count: Final = contract.count(document)
except AssertionError as error:
return Sample(seconds=seconds, status=200, error=str(error)), document
return Sample(seconds=seconds, status=200, rows=count), document
case _:
return _failure(result, seconds), None
def query_path(case: UiQueryCase) -> str:
query: Final = urlencode(case.params.model_dump(by_alias=True, exclude_none=True), doseq=True)
return f"{case.path}?{query}" if query else case.path
@step("Measure UI query {case.name} at {case.path}")
def issue_query(transport: HttpTransport, case: UiQueryCase, *, timeout: float) -> Result[Document]:
if case.body is not None:
return transport.post(
case.path, headers=transport.master, json=case.body, response_type=Document, timeout=timeout
)
return transport.get(
query_path(case), headers=transport.master, params=Query(), response_type=Document, timeout=timeout
)
def database_failures(text: str) -> float:
families: Final = tuple(text_string_to_metric_families(text))
required: Final = frozenset(("litellm_postgres_failed_requests", "litellm_batch_write_to_db_failed_requests"))
missing: Final = required - frozenset(family.name for family in families)
all_samples: Final = tuple(chain.from_iterable(family.samples for family in families))
has_activity: Final = any(
sample.name == "litellm_postgres_total_requests_total" and math.isfinite(sample.value) and sample.value > 0
for sample in all_samples
)
assert not missing or has_activity, f"Missing PostgreSQL failure counters {missing}"
names: Final = frozenset(f"{name}_total" for name in required)
values: Final = tuple(sample.value for sample in all_samples if sample.name in names)
assert all(math.isfinite(value) and value >= 0 for value in values), "Invalid PostgreSQL failure counter"
return sum(values)

206
tests/e2e/perf/setup/ui_query_cases.py Normal file
View file

@ -0,0 +1,206 @@
from __future__ import annotations
from dataclasses import replace
from datetime import datetime, timedelta
from itertools import chain
from typing import Final, Literal
from pydantic import ConfigDict
from setup.ui_queries import Contract, Document, Query, UiQueryCase, descend
class Params(Query):
model_config = ConfigDict(frozen=True, extra="forbid", populate_by_name=True)
start_date: str | None = None
timezone: str | None = None
offset: int | None = None
model_group: str | None = None
by_model_group: str | None = None
end_date: str | None = None
page: int | None = None
page_size: int | None = None
sort_by: str | None = None
sort_order: str | None = None
group_by_session: Literal["true"] | None = None
session_cursor: str | None = None
api_key: str | None = None
team_id: str | None = None
team_ids: str | None = None
user_id: str | None = None
model: str | None = None
model_id: str | None = None
key_alias: str | None = None
search: str | None = None
status_filter: str | None = None
cache_hit_filter: str | None = None
error_code: str | None = None
error_message: str | None = None
span_type: str | None = None
exclude_internal_health_checks: Literal["true"] | None = None
key_aliases: tuple[str, ...] | None = None
models: tuple[str, ...] | None = None
limit: int | None = None
tags: str | None = None
exclude_team_ids: str | None = None
include_current_utc_day: Literal["true"] | None = None
def _timestamp(value: datetime) -> str:
return value.strftime("%Y-%m-%d %H:%M:%S")
def next_page(case: UiQueryCase, document: Document) -> UiQueryCase:
params: Final = Params.model_validate(case.params.model_dump())
rows: Final = descend(document.root, case.contract.collection)
assert isinstance(rows, list), "Missing page-one rows"
ids: Final = frozenset(str(row["request_id"]) for row in rows if isinstance(row, dict))
contract: Final = replace(case.contract, excluded_ids=ids)
cursor: Final = descend(document.root, ("next_session_cursor",))
assert isinstance(cursor, str) and cursor, "Missing session cursor"
return replace(case, params=params.model_copy(update={"page": 2, "session_cursor": cursor}), contract=contract)
_LOGS: Final = Contract(("data",), ("request_id", "startTime", "metadata"))
def log_cases(anchor: datetime, api_key: str) -> tuple[UiQueryCase, ...]:
params: Final = Params(
start_date=_timestamp(anchor - timedelta(days=1)),
end_date=_timestamp(anchor),
page=1,
page_size=25,
sort_by="startTime",
sort_order="desc",
group_by_session="true",
)
base: Final = UiQueryCase("logs_default_24h", "/spend/logs/ui", params, _LOGS)
filtered: Final = params.model_copy(update={"start_date": _timestamp(anchor - timedelta(days=30))})
filters: Final = (
("api_key", api_key),
("team_id", "perf-team-0"),
("user_id", "perf-user-0"),
("model", "perf-chat"),
("model_id", "perf-model-0"),
("key_alias", "perf-key-0"),
("status_filter", "failure"),
("cache_hit_filter", "hit"),
("error_code", "429"),
("error_message", "Synthetic"),
("span_type", "llm"),
("exclude_internal_health_checks", "true"),
)
return (
base,
replace(base, name="logs_cursor_page_2", params=params.model_copy(update={"page": 2}), cursor="session"),
replace(
base,
name="logs_7d",
params=params.model_copy(update={"start_date": _timestamp(anchor - timedelta(days=7))}),
),
replace(
base,
name="logs_30d",
params=params.model_copy(update={"start_date": _timestamp(anchor - timedelta(days=30))}),
),
*(
replace(base, name=f"logs_{name}_30d", params=filtered.model_copy(update={name: value}))
for name, value in filters
),
*(
replace(base, name=f"logs_sort_{name}", params=params.model_copy(update={"sort_by": name}))
for name in ("spend", "total_tokens", "request_duration_ms", "model", "ttft_ms")
),
)
def analytics_cases(anchor: datetime) -> tuple[UiQueryCase, ...]:
dates: Final = Params(
start_date=(anchor - timedelta(days=7)).date().isoformat(), end_date=anchor.date().isoformat()
)
cache: Final = UiQueryCase(
"cache_activity_7d", "/global/activity/cache_hits", dates, Contract(("groups",), ("call_type",))
)
return (
UiQueryCase(
"team_users_all_teams",
"/team/spend/by_user",
dates.model_copy(update={"team_ids": ",".join(f"perf-team-{i}" for i in range(100))}),
Contract(("results",), ("team_id", "user_id", "spend"), ("api_requests",)),
),
UiQueryCase(
"team_users_filtered",
"/team/spend/by_user",
dates.model_copy(update={"team_ids": "perf-team-0"}),
Contract(("results",), ("team_id", "user_id", "spend"), ("api_requests",)),
),
cache,
replace(
cache,
name="cache_activity_filtered",
params=dates.model_copy(update={"key_aliases": ("perf-key-0", "perf-key-1"), "models": ("perf-chat",)}),
),
)
def _entity_activity_cases(anchor: datetime, entity: str) -> tuple[UiQueryCase, ...]:
dates: Final = Params(
start_date=(anchor - timedelta(days=7)).date().isoformat(),
end_date=anchor.date().isoformat(),
timezone="0",
exclude_team_ids="litellm-dashboard" if entity == "team" else None,
)
routes: Final = (
("aggregated", dates, Contract(("results",), ("date", "breakdown"))),
(
"aggregated/keys",
dates.model_copy(update={"offset": 0, "limit": 50}),
Contract(("api_keys",), ("api_key", "metrics", "metadata")),
),
(
"aggregated/search",
dates.model_copy(update={"search": "perf-key"}),
Contract(("api_keys",), ("api_key", "metrics", "metadata")),
),
(
"aggregated/model_top_keys",
dates.model_copy(update={"model_group": "perf-chat", "by_model_group": "true"}),
Contract(("api_keys",), ("api_key", "metrics", "metadata")),
),
)
return tuple(
UiQueryCase(
f"usage_{entity}_{route.replace('/', '_')}",
f"/{entity}/daily/activity/{route}",
params,
contract,
)
for route, params, contract in routes
)
def activity_cases(anchor: datetime) -> tuple[UiQueryCase, ...]:
groups: Final = tuple(_entity_activity_cases(anchor, entity) for entity in ("user", "team", "tag", "customer"))
return (
*chain.from_iterable(groups),
UiQueryCase(
"cache_leakage_keys",
"/user/daily/activity/aggregated/cache_leakage_keys",
Params(
start_date=(anchor - timedelta(days=30)).date().isoformat(),
end_date=anchor.date().isoformat(),
timezone="0",
include_current_utc_day="true",
),
Contract(("api_keys",), ("api_key", "metrics", "metadata")),
),
)
def cases(anchor: datetime, api_key: str) -> tuple[UiQueryCase, ...]:
return (
*log_cases(anchor, api_key),
*analytics_cases(anchor),
*activity_cases(anchor),
)

53
tests/e2e/perf/setup/ui_query_fixture.py Normal file
View file

@ -0,0 +1,53 @@
from __future__ import annotations
import hashlib
import os
import time
from datetime import datetime
from typing import Final
import psycopg
from pydantic import AwareDatetime, BaseModel, ConfigDict, TypeAdapter
from setup.ui_query_isolation import validate_target
from setup.upgrade_worker import read_state, release_indexes
class UiFixture(BaseModel):
model_config = ConfigDict(frozen=True)
anchor: AwareDatetime
api_key: str
def prepare_fixture() -> UiFixture:
validate_target(os.environ["DATABASE_URL"])
_finish_upgrade()
with psycopg.connect(os.environ["DATABASE_URL"], connect_timeout=10) as connection:
_ = connection.execute("SET statement_timeout='30s'")
row: Final = connection.execute("SELECT spec->>'anchor', rows FROM litellm_perf_seed.checkpoint").fetchone()
assert row is not None and TypeAdapter(int).validate_python(row[1]) > 0, "Missing completed seed rows"
fixture: Final = UiFixture(
anchor=TypeAdapter(datetime).validate_python(row[0]),
api_key=hashlib.md5(b"perf-key:0", usedforsecurity=False).hexdigest()
+ hashlib.md5(b"perf-key-tail:0", usedforsecurity=False).hexdigest(),
)
return fixture
def _finish_upgrade() -> None:
if os.environ.get("E2E_DATABASE_UPGRADE") != "1":
return
with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True, connect_timeout=10) as connection:
release_indexes(connection)
deadline: Final = time.monotonic() + 6 * 3600
while time.monotonic() < deadline:
match read_state(connection).phase:
case "failed":
raise AssertionError("Snapshot upgrade failed")
case "complete":
return
case _:
pass
time.sleep(5)
raise AssertionError("Snapshot indexes did not finish before UI measurements")

110
tests/e2e/perf/setup/ui_query_isolation.py Normal file
View file

@ -0,0 +1,110 @@
from __future__ import annotations
import json
import os
import re
import time
from collections.abc import Callable, Generator
from contextlib import contextmanager
from pathlib import Path
from typing import Final
from urllib.parse import urlparse
import psycopg
import pytest
from psycopg.rows import dict_row
from pydantic import BaseModel, ConfigDict, TypeAdapter
class ActiveQuery(BaseModel):
model_config = ConfigDict(frozen=True)
pid: int
leader_pid: int | None
application_name: str
state: str | None
wait_event_type: str | None
wait_event: str | None
seconds: float | None
query: str
def validate_target(url: str) -> None:
assert os.environ.get("E2E_UI_PERF_DISPOSABLE") == "1", "Disposable database confirmation required"
assert os.environ.get("E2E_RESET_SPEND_LOGS") != "1", "Never truncate the snapshot baseline"
assert re.fullmatch(r"db-perf-perf-[a-z0-9-]+\.litellm\.svc", urlparse(url).hostname or ""), (
"SQL cancellation is restricted to the disposable perf-lane restore"
)
def active(connection: psycopg.Connection[dict[str, object]]) -> tuple[ActiveQuery, ...]:
rows: Final = connection.execute("""SELECT pid, leader_pid, application_name, state,
wait_event_type, wait_event, extract(epoch from clock_timestamp()-query_start)::float AS seconds,
left(query, 1500) AS query FROM pg_stat_activity
WHERE datname = current_database() AND pid <> pg_backend_pid()
AND backend_type IN ('client backend', 'parallel worker')
AND state IS DISTINCT FROM 'idle' ORDER BY pid""").fetchall()
return TypeAdapter(tuple[ActiveQuery, ...]).validate_python(rows)
def wait_idle(connection: psycopg.Connection[dict[str, object]], evidence: Path, *, cancel: bool) -> None:
consecutive = 0
for queries in (active(connection) for _ in range(30)):
with evidence.open("a", encoding="utf-8") as output:
_ = output.write(
json.dumps(
{"time": time.time(), "cancel": cancel, "active": tuple(query.model_dump() for query in queries)}
)
+ "\n"
)
if queries:
consecutive = 0 # rebind-ok: consecutive idle observations reset on activity
for query in queries:
if cancel and query.leader_pid is None:
_ = connection.execute(
"SELECT pg_cancel_backend(%s)" if query.state == "active" else "SELECT pg_terminate_backend(%s)",
(query.pid,),
)
else:
consecutive += 1 # rebind-ok: consecutive idle observations across polling iterations
if consecutive == 3:
return
time.sleep(1)
pytest.exit("Database did not become idle. No further UI measurements will run", returncode=2)
def settled_counter(read: Callable[[], float], evidence: Path) -> float:
previous: float | None = None
consecutive = 0
for value in (read() for _ in range(10)):
with evidence.open("a", encoding="utf-8") as output:
_ = output.write(json.dumps({"time": time.time(), "database_failures": value}) + "\n")
consecutive = consecutive + 1 if value == previous else 1 # rebind-ok: polling state
if consecutive == 3:
return value
previous = value # rebind-ok: last observed counter for stabilization
time.sleep(1)
pytest.exit("Database counters did not settle. No further UI measurements will run", returncode=2)
@contextmanager
def isolated_database(name: str) -> Generator[None]:
url: Final = os.environ["DATABASE_URL"]
validate_target(url)
evidence: Final = Path(os.environ["E2E_ARTIFACT_DIR"]) / f"ui-isolation-{name}.jsonl"
with psycopg.Connection[dict[str, object]].connect(
url,
autocommit=True,
connect_timeout=10,
application_name="ui-isolation-monitor",
options="-c statement_timeout=5000",
row_factory=dict_row,
) as connection:
try:
wait_idle(connection, evidence, cancel=False)
yield
finally:
try:
wait_idle(connection, evidence, cancel=True)
except Exception as error:
pytest.exit(f"SQL cleanup failed ({type(error).__name__}). Stopping UI measurements", returncode=2)

159
tests/e2e/perf/test_ui_queries.py Normal file
View file

@ -0,0 +1,159 @@
from __future__ import annotations
import json
import os
from dataclasses import dataclass, replace
from datetime import UTC, datetime
from pathlib import Path
from typing import Final
import pytest
from e2e_config import CONTROL_PLANE_BASE_URL, MASTER_KEY
from e2e_http import send_text_external, without_retries
from setup.ui_queries import (
Sample,
UiQueryCase,
database_failures,
error_code,
issue_query,
measure,
summarize,
ui_junit_properties,
)
from setup.ui_query_cases import Params, cases, next_page
from setup.ui_query_fixture import UiFixture, prepare_fixture
from setup.ui_query_isolation import isolated_database, settled_counter
from transport import HttpTransport
pytestmark = [
pytest.mark.perf,
pytest.mark.quiet_stack,
pytest.mark.skipif(os.environ.get("E2E_UI_PERF_DISPOSABLE") != "1", reason="Dedicated disposable database only"),
]
def pytest_generate_tests(metafunc: pytest.Metafunc) -> None:
if "ui_case_name" not in metafunc.fixturenames:
return
if os.environ.get("E2E_UI_PERF_DISPOSABLE") == "1" and os.environ.get("E2E_RESET_SPEND_LOGS") == "1":
raise pytest.UsageError("UI performance tests cannot run with E2E_RESET_SPEND_LOGS=1")
names: Final = tuple(case.name for case in cases(datetime(2000, 1, 1, tzinfo=UTC), "collection-only"))
metafunc.parametrize("ui_case_name", names, ids=names)
@pytest.fixture(scope="module")
def ui_fixture() -> UiFixture:
return prepare_fixture()
def _metrics(transport: HttpTransport) -> float:
urls: Final = tuple(url.strip() for url in os.environ["E2E_DATABASE_METRICS_URLS"].split(",") if url.strip())
try:
assert urls, "Provide a metrics URL for every proxy worker/pod"
readings: Final = tuple(send_text_external("GET", url, headers=transport.master, timeout=10) for url in urls)
assert all(reading.status_code == 200 for reading in readings), "Cannot read database failure metrics"
return sum(database_failures(reading.body) for reading in readings)
except Exception as error:
pytest.exit(f"Database metrics unavailable ({type(error).__name__}). Stopping UI measurements", returncode=2)
@dataclass(frozen=True, slots=True)
class Prepared:
case: UiQueryCase
prerequisite: Sample | None = None
def _page_two(transport: HttpTransport, case: UiQueryCase, timeout: float) -> Prepared:
if case.cursor == "none":
return Prepared(case)
params: Final = Params.model_validate(case.params.model_dump()).model_copy(update={"page": 1})
first: Final = replace(case, params=params)
sample, document = measure(lambda: issue_query(transport, first, timeout=timeout), case.contract)
if sample.error or document is None:
return Prepared(case, sample)
try:
return Prepared(next_page(case, document), sample)
except AssertionError as error:
return Prepared(case, sample.model_copy(update={"error": str(error)}))
def _samples(transport: HttpTransport, case: UiQueryCase, timeout: float, repeats: int) -> tuple[Sample, ...]:
measured, _ = measure(lambda: issue_query(transport, case, timeout=timeout), case.contract)
if measured.error:
return (measured,)
repeated: Final = tuple(
measure(lambda: issue_query(transport, case, timeout=timeout), case.contract)[0] for _ in range(repeats)
)
return (measured, *repeated)
def test_ui_query_performance(ui_case_name: str, ui_fixture: UiFixture, request: pytest.FixtureRequest) -> None:
case: Final = next(case for case in cases(ui_fixture.anchor, ui_fixture.api_key) if case.name == ui_case_name)
transport: Final = HttpTransport(base_url=CONTROL_PLANE_BASE_URL, master_key=MASTER_KEY)
timeout: Final = float(os.environ.get("E2E_UI_QUERY_TIMEOUT_SECONDS", "120"))
repeats: Final = int(os.environ.get("E2E_UI_QUERY_REPEATS", "9"))
assert 0 < timeout <= 600 and 1 <= repeats <= 1000, "Invalid UI performance sample settings"
target: Final = Path(os.environ["E2E_ARTIFACT_DIR"])
assert target.is_dir(), "E2E_ARTIFACT_DIR must exist"
evidence: Final = target / f"ui-isolation-{case.name}.jsonl"
with without_retries():
with isolated_database(case.name):
before: Final = settled_counter(lambda: _metrics(transport), evidence)
paginated: Final = _page_two(transport, case, timeout)
samples: Final = (
(paginated.prerequisite,)
if paginated.prerequisite is not None and paginated.prerequisite.error
else _samples(transport, paginated.case, timeout, repeats)
)
with evidence.open("a", encoding="utf-8") as sample_output:
_ = sample_output.write(
json.dumps(
{
"samples": tuple(sample.model_dump() for sample in samples),
"prerequisite": paginated.prerequisite.model_dump()
if paginated.prerequisite is not None
else None,
}
)
+ "\n"
)
after_requests: Final = _metrics(transport)
item: Final = request.node # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # pytest: untyped
assert isinstance(item, pytest.Item)
item.user_properties.extend(ui_junit_properties(samples, paginated.prerequisite))
after: Final = settled_counter(lambda: _metrics(transport), evidence)
item.user_properties.extend(
(
("perf_ui_database_failures", after_requests - before),
("perf_ui_cleanup_database_failures", after - after_requests),
)
)
timings: Final = summarize(samples)
artifact: Final = {
"name": case.name,
"surface": case.surface,
"path": case.path,
"method": "POST" if case.body is not None else "GET",
"params": paginated.case.params.model_dump(mode="json", by_alias=True, exclude_none=True),
"prerequisite": paginated.prerequisite.model_dump() if paginated.prerequisite is not None else None,
"prerequisite_failed": paginated.prerequisite is not None and bool(paginated.prerequisite.error),
"body": case.body.model_dump(mode="json", exclude_none=True) if case.body is not None else None,
"samples": [sample.model_dump() for sample in samples],
"error_codes": [error_code(sample) for sample in samples],
"timings": timings.model_dump(),
"timeouts": sum(sample.error == "timeout" for sample in samples),
"failures": sum(bool(sample.error) for sample in samples),
"status_counts": {
str(status): sum(sample.status == status for sample in samples)
for status in frozenset(sample.status for sample in samples)
},
"validation_failures": sum(
bool(sample.error) and (sample.status == 200 or sample.error == "Response is not valid JSON")
for sample in samples
),
"database_failures_before": before,
"database_failures_after_requests": after_requests,
"database_failures_after_cleanup": after,
}
output: Final = target / f"ui-query-{case.name}.json"
_ = output.write_text(json.dumps(artifact, indent=2), encoding="utf-8")

257
tests/e2e_harness/perf/setup/test_ui_queries.py Normal file
View file

@ -0,0 +1,257 @@
from __future__ import annotations
from datetime import UTC, datetime
from typing import Final
from urllib.parse import parse_qs, urlsplit
import pytest
from e2e_http import NetworkError, Result, Success, UnknownApiError
from setup.ui_queries import (
Contract,
Document,
Sample,
database_failures,
measure,
query_path,
summarize,
ui_junit_properties,
)
from setup.ui_query_cases import Params, cases, next_page
def test_query_string_preserves_ui_arrays_and_boolean_strings() -> None:
requests: Final = cases(datetime(2026, 1, 1, tzinfo=UTC), "key")
cache: Final = next(case for case in requests if case.name == "cache_activity_filtered")
assert parse_qs(urlsplit(query_path(cache)).query) == {
"start_date": ["2025-12-25"],
"end_date": ["2026-01-01"],
"key_aliases": ["perf-key-0", "perf-key-1"],
"models": ["perf-chat"],
}
health: Final = next(case for case in requests if case.name == "logs_exclude_internal_health_checks_30d")
query: Final = parse_qs(urlsplit(query_path(health)).query)
assert query["exclude_internal_health_checks"] == ["true"]
assert query["group_by_session"] == ["true"]
@pytest.mark.parametrize("entity", ("user", "team", "tag", "customer"))
@pytest.mark.parametrize(
("route", "extra"),
(
("aggregated", {}),
("aggregated/keys", {"offset": 0, "limit": 50}),
("aggregated/search", {"search": "perf-key"}),
("aggregated/model_top_keys", {"model_group": "perf-chat", "by_model_group": "true"}),
),
)
def test_usage_requests_list_keys_without_a_single_key_filter(
entity: str, route: str, extra: dict[str, str | int]
) -> None:
case: Final = next(
case
for case in cases(datetime(2026, 1, 1, tzinfo=UTC), "key")
if case.path == f"/{entity}/daily/activity/{route}"
)
assert case.params.model_dump(exclude_none=True) == {
"start_date": "2025-12-25",
"end_date": "2026-01-01",
"timezone": "0",
**({"exclude_team_ids": "litellm-dashboard"} if entity == "team" else {}),
**extra,
}
@pytest.mark.parametrize(("name", "value"), (("api_key", "key"), ("key_alias", "perf-key-0")))
def test_key_filters_request_30_days_of_logs(name: str, value: str) -> None:
case: Final = next(
case for case in cases(datetime(2026, 1, 1, tzinfo=UTC), "key") if case.name == f"logs_{name}_30d"
)
assert parse_qs(urlsplit(query_path(case)).query) == {
"start_date": ["2025-12-02 00:00:00"],
"end_date": ["2026-01-01 00:00:00"],
"page_size": ["25"],
"page": ["1"],
"sort_by": ["startTime"],
"sort_order": ["desc"],
"group_by_session": ["true"],
name: [value],
}
@pytest.mark.parametrize("invalid", ("not-a-number", True, float("nan"), float("inf")))
def test_positive_row_does_not_hide_a_malformed_later_row(invalid: str | bool | float) -> None:
with pytest.raises(AssertionError, match=r"Expected (numeric|finite) spend"):
Contract(positive=("spend",)).count(Document([{"spend": 1}, {"spend": invalid}]))
def test_session_cursor_is_reused_and_repeated_rows_rejected() -> None:
case: Final = next(case for case in cases(datetime(2026, 1, 1, tzinfo=UTC), "key") if case.cursor == "session")
first: Final = Document.model_validate_json('{"data":[{"request_id":"one"}],"next_session_cursor":"opaque"}')
second: Final = next_page(case, first)
assert Params.model_validate(second.params).session_cursor == "opaque"
assert Params.model_validate(second.params).page == 2
with pytest.raises(AssertionError, match="Repeated page-one row"):
_ = second.contract.count(first)
def test_first_query_is_separate_from_repeated_latency() -> None:
samples: Final = (Sample(seconds=80, status=200), *(Sample(seconds=i, status=200) for i in range(1, 101)))
result: Final = summarize(samples)
assert result.first_seconds == 80
assert result.repeated_samples == 100
assert result.repeated_p50_seconds == 50.5
assert result.repeated_p95_seconds == 95
assert result.repeated_p99_seconds == 99
assert result.repeated_max_seconds == 100
def test_failed_repeats_do_not_look_like_fast_successes() -> None:
result: Final = summarize((Sample(seconds=10, status=200), Sample(seconds=0.01, status=503, error="HTTP 503")))
assert result.repeated_samples == 0
assert result.repeated_p50_seconds is None
assert result.repeated_p95_seconds is None
assert result.repeated_p99_seconds is None
assert result.repeated_max_seconds is None
def test_small_sample_does_not_report_tail_percentiles() -> None:
result: Final = summarize(tuple(Sample(seconds=1, status=200) for _ in range(6)))
assert result.repeated_samples == 5
assert result.repeated_p50_seconds == 1
assert result.repeated_p95_seconds is None
assert result.repeated_p99_seconds is None
def test_failed_first_request_emits_error_code_without_latency() -> None:
sample: Final = Sample(seconds=30, status=500, error="HTTP 500")
assert summarize((sample,)).first_seconds is None
assert dict(ui_junit_properties((sample,))) == {
"perf_ui_requests": 1,
"perf_ui_failures": 1,
"perf_ui_samples": 0,
"perf_ui_repeated_samples": 0,
"perf_ui_error_http_500": 1,
}
def test_junit_counts_failures_but_only_measures_successful_latency() -> None:
samples: Final = (
Sample(seconds=2, status=200),
Sample(seconds=3, status=200),
Sample(seconds=30, status=-1, error="timeout"),
Sample(seconds=30, status=503, error="HTTP 503"),
Sample(seconds=1, status=200, error="Empty collection"),
)
assert dict(ui_junit_properties(samples, Sample(seconds=4, status=200))) == {
"perf_ui_requests": 6,
"perf_ui_failures": 3,
"perf_ui_samples": 2,
"perf_ui_seconds_p50": 2.5,
"perf_ui_seconds_p90": 3,
"perf_ui_seconds_max": 3,
"perf_ui_repeated_samples": 1,
"perf_ui_error_timeout": 1,
"perf_ui_error_http_503": 1,
"perf_ui_error_contract": 1,
"perf_ui_first_seconds": 2,
"perf_ui_repeated_seconds_p50": 3,
"perf_ui_repeated_seconds_max": 3,
}
def test_failed_cursor_prerequisite_is_counted_once() -> None:
sample: Final = Sample(seconds=30, status=503, error="HTTP 503")
assert dict(ui_junit_properties((sample,), sample)) == dict(ui_junit_properties((sample,)))
def test_all_request_percentiles_include_first_and_use_nearest_rank_p90() -> None:
samples: Final = tuple(Sample(seconds=seconds, status=200) for seconds in (10, 1, 9, 2, 8, 3, 7, 4, 6, 5))
metrics: Final = dict(ui_junit_properties(samples))
assert {key: value for key, value in metrics.items() if key.startswith("perf_ui_seconds_")} == {
"perf_ui_seconds_p50": 5.5,
"perf_ui_seconds_p90": 9,
"perf_ui_seconds_max": 10,
}
assert metrics["perf_ui_samples"] == 10
@pytest.mark.parametrize("body", ("{}", '{"data": []}', '{"data": [{"spend": 1}]}', '{"data": [null]}'))
def test_empty_or_malformed_rows_fail_even_with_http_200(body: str) -> None:
ticks: Final = iter((10.0, 12.5))
result, _ = measure(
lambda: Success(status_code=200, data=Document.model_validate_json(body)),
Contract(("data",), ("request_id",)),
clock=lambda: next(ticks),
)
assert result.status == 200
assert result.seconds == 2.5
assert result.rows == 0
assert result.error
def test_valid_response_records_elapsed_seconds_and_rows() -> None:
ticks: Final = iter((5.0, 5.125))
result, document = measure(
lambda: Success(status_code=200, data=Document.model_validate_json('{"data":[{"request_id":"one"}]}')),
Contract(("data",), ("request_id",)),
clock=lambda: next(ticks),
)
assert result == Sample(seconds=0.125, status=200, rows=1)
assert document is not None
def test_zero_spend_rows_are_allowed_alongside_nonzero_fixture_data() -> None:
contract: Final = Contract(positive=("spend",))
assert contract.count(Document.model_validate_json('[{"spend":0},{"spend":2}]')) == 2
with pytest.raises(AssertionError, match="positive spend"):
_ = contract.count(Document.model_validate_json('[{"spend":0}]'))
@pytest.mark.parametrize(
"failure,error,status",
(
(UnknownApiError(status_code=500, body="sensitive error detail"), "HTTP 500", 500),
(NetworkError(message="Read timed out"), "timeout", -1),
(NetworkError(message="connection refused"), "network", -1),
),
)
def test_failed_requests_are_recorded_without_response_bodies(
failure: Result[Document],
error: str,
status: int,
) -> None:
ticks: Final = iter((1.0, 2.0))
result, document = measure(lambda: failure, Contract(), clock=lambda: next(ticks))
assert result == Sample(seconds=1, status=status, error=error)
assert document is None
def _metrics(postgres: str, batch: str) -> str:
return (
"# TYPE litellm_postgres_failed_requests_total counter\n"
f'litellm_postgres_failed_requests_total{{call_type="query_raw"}} {postgres}\n'
"# TYPE litellm_batch_write_to_db_failed_requests_total counter\n"
f'litellm_batch_write_to_db_failed_requests_total{{call_type="batch"}} {batch}\n'
)
def test_postgres_failures_include_background_batch_writes() -> None:
assert database_failures(_metrics("0", "0")) == 0
assert database_failures(_metrics("2", "3")) == 5
def test_lazy_failure_counters_require_evidence_of_database_activity() -> None:
assert (
database_failures(
"# TYPE litellm_postgres_total_requests_total counter\nlitellm_postgres_total_requests_total 3\n"
)
== 0
)
@pytest.mark.parametrize(
"text", ("", "# TYPE unrelated counter\nunrelated_total 0\n", _metrics("NaN", "0"), _metrics("-1", "0"))
)
def test_missing_or_invalid_metrics_cannot_pass_as_zero(text: str) -> None:
with pytest.raises(AssertionError):
_ = database_failures(text)

28
tests/e2e_harness/perf/setup/test_ui_query_isolation.py Normal file
View file

@ -0,0 +1,28 @@
from __future__ import annotations
from pathlib import Path
from typing import Final
import pytest
from setup.ui_query_isolation import settled_counter, validate_target
@pytest.mark.parametrize("host", ("production.example.com", "db-perf-source.litellm.svc", "localhost"))
def test_query_cancellation_rejects_non_lane_databases(host: str, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("E2E_UI_PERF_DISPOSABLE", "1")
monkeypatch.setenv("E2E_RESET_SPEND_LOGS", "0")
with pytest.raises(AssertionError, match="restricted"):
validate_target(f"postgresql://user@{host}/litellm")
def test_query_cancellation_requires_explicit_disposable_opt_in(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.delenv("E2E_UI_PERF_DISPOSABLE", raising=False)
with pytest.raises(AssertionError, match="confirmation"):
validate_target("postgresql://user@db-perf-perf-test-1.litellm.svc/litellm")
def test_counter_stabilization_waits_for_delayed_failures(tmp_path: Path) -> None:
readings: Final = iter((0.0, 0.0, 2.0, 2.0, 2.0))
evidence: Final = tmp_path / "isolation.jsonl"
assert settled_counter(lambda: next(readings), evidence) == 2
assert len(evidence.read_text().splitlines()) == 5

4
tests/integration/_support/actors.py
View file

@ -15,7 +15,7 @@ from pydantic import JsonValue
from litellm.proxy.utils import PrismaClient, hash_token
from tests.integration._support.client import Gateway, Scenario, string_value
from tests.integration._support.prisma_clients import connected_proxy_prisma
from tests.integration._support.process import owned_proxy
from tests.integration._support.process import OwnedState, owned_proxy
TEAM_ALPHA: Final = "behavior-pin-team-alpha"
TEAM_BETA: Final = "behavior-pin-team-beta"
@ -124,7 +124,7 @@ def owned_settings_client(
config: Final = {**base, "litellm_settings": {**base["litellm_settings"], **litellm_settings}}
path: Final = directory / "owned-settings.yaml"
path.write_text(yaml.safe_dump(config))
with owned_proxy(client.gateway, directory, {}, config=path) as candidate:
with owned_proxy(client.gateway, directory, {}, config=path, state=OwnedState.SHARED_WITH_WORKER) as candidate:
yield GatewayClient(candidate, client.replacements)

3
tests/integration/_support/bedrock_discovery.py
View file

@ -23,7 +23,7 @@ from integration._support.aws_control_plane import (
verified_scope,
)
from integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, held, object_value, string_value
from integration._support.process import OwnedProxy, owned_proxy_process
from integration._support.process import OwnedProxy, OwnedState, owned_proxy_process
from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
from pydantic import JsonValue
@ -115,6 +115,7 @@ def discovery_proxy(
config=discovery_config(parent, directory, check_provider_endpoint=check_provider_endpoint),
remove_environment=aws_environment_names(),
workers=workers,
state=OwnedState.SHARED_WITH_WORKER,
) as owned:
wait_for_every_worker(owned.log, workers)
yield owned

28
tests/integration/_support/client.py
View file

@ -35,9 +35,9 @@ def list_value(value: JsonValue) -> Sequence[JsonValue]:
def delete_key_if_present(candidate: Gateway, key: str) -> None:
digest: Final = sha256(key.encode()).hexdigest()
if read_rows('SELECT token FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)):
if candidate.rows('SELECT token FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)):
candidate.post("/key/delete", {"keys": [key]})
assert read_rows('SELECT token FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)) == []
assert candidate.rows('SELECT token FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)) == []
def eventually(
@ -95,6 +95,10 @@ class Gateway:
client: httpx.Client
key: str
upstream_url: str
database_url: str | None = None
def rows(self, query: str, parameters: tuple[str, ...]) -> list[dict[str, JsonValue]]:
return read_rows(query, parameters, database_url=self.database_url)
def request(
self,
@ -175,7 +179,7 @@ class Scenario:
def delete_team(self, identity: str) -> None:
self.gateway.post("/team/delete", {"team_ids": [identity]})
assert read_rows('SELECT team_id FROM "LiteLLM_TeamTable" WHERE team_id = %s', (identity,)) == []
assert self.gateway.rows('SELECT team_id FROM "LiteLLM_TeamTable" WHERE team_id = %s', (identity,)) == []
def project(self, team_id: str, **fields: JsonValue) -> str:
created: Final = self.gateway.post(
@ -188,7 +192,9 @@ class Scenario:
def delete_project(self, identity: str) -> None:
response: Final = self.gateway.request("DELETE", "/project/delete", {"project_ids": [identity]})
assert response.status_code == 200, response.text
assert read_rows('SELECT project_id FROM "LiteLLM_ProjectTable" WHERE project_id = %s', (identity,)) == []
assert (
self.gateway.rows('SELECT project_id FROM "LiteLLM_ProjectTable" WHERE project_id = %s', (identity,)) == []
)
def budget(self, **fields: JsonValue) -> str:
created: Final = self.gateway.post("/budget/new", fields)
@ -198,7 +204,7 @@ class Scenario:
def delete_budget(self, identity: str) -> None:
self.gateway.post("/budget/delete", {"id": identity})
assert read_rows('SELECT budget_id FROM "LiteLLM_BudgetTable" WHERE budget_id = %s', (identity,)) == []
assert self.gateway.rows('SELECT budget_id FROM "LiteLLM_BudgetTable" WHERE budget_id = %s', (identity,)) == []
def organization(self, **fields: JsonValue) -> str:
created: Final = self.gateway.post(
@ -212,7 +218,9 @@ class Scenario:
response: Final = self.gateway.request("DELETE", "/organization/delete", {"organization_ids": [identity]})
assert response.status_code == 200, response.text
assert (
read_rows('SELECT organization_id FROM "LiteLLM_OrganizationTable" WHERE organization_id = %s', (identity,))
self.gateway.rows(
'SELECT organization_id FROM "LiteLLM_OrganizationTable" WHERE organization_id = %s', (identity,)
)
== []
)
self.delete_budget(budget_id)
@ -236,7 +244,7 @@ class Scenario:
def delete_user(self, identity: str) -> None:
response: Final = self.gateway.request("POST", "/user/delete", {"user_ids": [identity]})
assert response.status_code == 200 and response.json() == 1, response.text
assert read_rows('SELECT user_id FROM "LiteLLM_UserTable" WHERE user_id = %s', (identity,)) == []
assert self.gateway.rows('SELECT user_id FROM "LiteLLM_UserTable" WHERE user_id = %s', (identity,)) == []
def member(self, team_id: str, role: str = "user") -> str:
"""Create an internal user and add them to ``team_id``; deleting the user later removes the membership."""
@ -247,7 +255,7 @@ class Scenario:
def delete_key(self, token: str) -> None:
self.gateway.post("/key/delete", {"keys": [token]})
hashed: Final = sha256(token.encode()).hexdigest()
assert read_rows('SELECT token FROM "LiteLLM_VerificationToken" WHERE token = %s', (hashed,)) == []
assert self.gateway.rows('SELECT token FROM "LiteLLM_VerificationToken" WHERE token = %s', (hashed,)) == []
info: Final = object_value(self.gateway.get("/key/info", {"key": hashed})["info"])
assert info["status"] == "deleted", f"Deleted key still served as live: {info['status']}"
@ -265,7 +273,9 @@ class Scenario:
lambda served: identity not in served,
seconds=2 * PROXY_CONFIG_RELOAD_INTERVAL_SECONDS,
)
assert read_rows('SELECT model_id FROM "LiteLLM_ProxyModelTable" WHERE model_id = %s', (identity,)) == []
assert (
self.gateway.rows('SELECT model_id FROM "LiteLLM_ProxyModelTable" WHERE model_id = %s', (identity,)) == []
)
def model(self, *, model_info: Mapping[str, JsonValue] | None = None, **parameters: JsonValue) -> str:
name: Final = f"integration-{uuid.uuid4().hex}"

205
tests/integration/_support/process.py
View file

@ -7,9 +7,9 @@ import subprocess
import sys
import time
import uuid
from collections.abc import Generator, Iterable, Iterator, Mapping
from collections.abc import Callable, Generator, Iterable, Iterator, Mapping
from concurrent.futures import ThreadPoolExecutor
from contextlib import contextmanager, suppress
from contextlib import AbstractContextManager, contextmanager, suppress
from dataclasses import dataclass
from enum import Enum
from datetime import UTC, datetime, timedelta
@ -17,7 +17,7 @@ from ipaddress import IPv4Address
from itertools import accumulate
from pathlib import Path
from types import MappingProxyType
from typing import TYPE_CHECKING, Final
from typing import TYPE_CHECKING, Final, TypeAlias
from urllib.parse import urlsplit
import httpx
@ -33,6 +33,8 @@ from tests.integration._support.budget import import_seconds, readiness_for, rea
from tests.integration._support.client import GATEWAY_LIMITS, Gateway, gateway_from_environment
from tests.integration._support.database import scratch_database
from tests.integration._support.ports import allocate_port
from tests.integration._support.redis_process import owned_redis
from tests.integration._support.workers import gateway_profile, replica_urls
if TYPE_CHECKING:
import pytest
@ -43,6 +45,75 @@ LEGACY_MIGRATE_DEPLOY: Final = ("--use_legacy_migration_resolver",)
CONFIGURATION: Final = TypeAdapter(dict[str, object])
_DEPLOYMENT_IDS: Final = TypeAdapter(dict[str, list[str]])
MCP_COVERAGE_CONFIG: Final = Path(__file__).resolve().parents[1] / "mcp_coverage.toml"
OWNED_STATE_VARIABLE: Final = "INTEGRATION_OWNED_STATE"
class OwnedState(Enum):
PRIVATE = "private"
SHARED_WITH_WORKER = "shared_with_worker"
@dataclass(frozen=True, slots=True)
class PrivateState:
database_url: str
environment: Mapping[str, str]
PrivateStateFactory: TypeAlias = Callable[[Path, Mapping[str, str]], AbstractContextManager[PrivateState]]
def boots_private_state(
state: OwnedState | None,
overrides: Mapping[str, str],
remove_environment: tuple[str, ...],
inherited: Mapping[str, str],
) -> bool:
chosen: Final = state or OwnedState(inherited.get(OWNED_STATE_VARIABLE, OwnedState.PRIVATE.value))
return chosen is OwnedState.PRIVATE and "DATABASE_URL" not in overrides and "DATABASE_URL" not in remove_environment
@contextmanager
def private_state(directory: Path, inherited: Mapping[str, str]) -> Iterator[PrivateState]:
with scratch_database("owned") as database, owned_redis(directory, inherited=inherited) as cache:
replica: Final = replica_urls(database, inherited) if gateway_profile(inherited).replica else None
yield PrivateState(
database,
MappingProxyType(
{
"DATABASE_URL": database if replica is None else replica["INTEGRATION_PROXY_DATABASE_URL"],
**(
{}
if replica is None
else {"DATABASE_URL_READ_REPLICA": replica["INTEGRATION_PROXY_READ_REPLICA_URL"]}
),
"REDIS_HOST": cache.host,
"REDIS_PORT": str(cache.port),
}
),
)
@dataclass(frozen=True, slots=True)
class OwnedEnvironment:
overrides: Mapping[str, str]
database_url: str | None
@contextmanager
def owned_environment(
directory: Path,
overrides: Mapping[str, str],
remove_environment: tuple[str, ...],
state: OwnedState | None,
inherited: Mapping[str, str],
private: PrivateStateFactory = private_state,
) -> Iterator[OwnedEnvironment]:
if not boots_private_state(state, overrides, remove_environment, inherited):
yield OwnedEnvironment(overrides, None)
return
with private(directory, inherited) as booted:
kept: Final = {name: value for name, value in booted.environment.items() if name not in remove_environment}
yield OwnedEnvironment(MappingProxyType({**kept, **overrides}), booted.database_url)
def proxy_database_environment(inherited: Mapping[str, str] = os.environ) -> Mapping[str, str]:
@ -368,6 +439,7 @@ def owned_proxy(
database_setup: tuple[str, ...] = DB_PUSH,
coverage_file: Path | None = None,
inherited: Mapping[str, str] = os.environ,
state: OwnedState | None = None,
) -> Iterator[Gateway]:
with owned_proxy_process(
gateway,
@ -379,20 +451,21 @@ def owned_proxy(
database_setup=database_setup,
coverage_file=coverage_file,
inherited=inherited,
state=state,
) as owned:
yield owned.gateway
@contextmanager
def isolated_gateway(directory: Path, environment: "pytest.MonkeyPatch") -> Iterator[Gateway]:
with gateway_from_environment() as shared, scratch_database() as database:
environment.setenv("DATABASE_URL", database)
environment.setenv("INTEGRATION_PROXY_DATABASE_URL", database)
with gateway_from_environment() as shared, owned_proxy(shared, directory, {}, state=OwnedState.PRIVATE) as owned:
assert owned.database_url is not None
environment.setenv("DATABASE_URL", owned.database_url)
environment.setenv("INTEGRATION_PROXY_DATABASE_URL", owned.database_url)
environment.delenv("INTEGRATION_PROXY_READ_REPLICA_URL", raising=False)
environment.delenv("DATABASE_URL_READ_REPLICA", raising=False)
with owned_proxy(shared, directory, {}) as owned:
owned.post("/config/update", {"router_settings": {"num_retries": 0}})
yield owned
owned.post("/config/update", {"router_settings": {"num_retries": 0}})
yield owned
def _stop(process: subprocess.Popen[bytes]) -> None:
@ -634,6 +707,32 @@ def _proxy_root() -> Path:
return Path(os.environ.get("INTEGRATION_PROXY_ROOT") or Path(__file__).resolve().parents[3])
def _configuration_copy(config: Path, directory: Path, prefix: str) -> Path:
checkouts: Final = (_harness_root(), _proxy_root().resolve())
inside_checkout: Final = any(config.resolve().is_relative_to(checkout) for checkout in checkouts)
target: Final = directory if inside_checkout else config.parent
target.mkdir(parents=True, exist_ok=True)
return target / f"{prefix}-{uuid.uuid4().hex}.yaml"
def private_configuration(config: Path, directory: Path) -> Path:
loaded: Final[object] = yaml.safe_load(config.read_text())
if not isinstance(loaded, dict):
return config
settings: Final = CONFIGURATION.validate_python(loaded)
router: Final = settings.get("router_settings") or {}
if not isinstance(router, dict) or "num_retries" in router:
return config
destination: Final = _configuration_copy(config, directory, "private-config")
destination.write_text(
yaml.safe_dump(
{**settings, "router_settings": {**CONFIGURATION.validate_python(router), "num_retries": 0}},
sort_keys=False,
)
)
return destination
def worker_configuration(config: Path, directory: Path, inherited: Mapping[str, str] = os.environ) -> Path:
if inherited.get("INTEGRATION_WORKER_ISOLATION") != "1":
return config
@ -644,11 +743,7 @@ def worker_configuration(config: Path, directory: Path, inherited: Mapping[str,
return f"http://127.0.0.1:{inherited['INTEGRATION_PROVIDER_PORT']}"
contents: Final = re.sub(r"http://127[.]0[.]0[.]1:(8190|8191)(?![0-9])", endpoint, config.read_text())
checkouts: Final = (_harness_root(), _proxy_root().resolve())
inside_checkout: Final = any(config.resolve().is_relative_to(checkout) for checkout in checkouts)
target: Final = directory if inside_checkout else config.parent
target.mkdir(parents=True, exist_ok=True)
destination: Final = target / f"worker-config-{uuid.uuid4().hex}.yaml"
destination: Final = _configuration_copy(config, directory, "worker-config")
loaded: Final[object] = yaml.safe_load(contents)
if not isinstance(loaded, dict):
destination.write_text(contents)
@ -753,42 +848,58 @@ def owned_proxy_process(
healthcheck: WorkerHealthcheck = WorkerHealthcheck.STRICT,
inherited: Mapping[str, str] = os.environ,
client_timeout: float = 15,
state: OwnedState | None = None,
) -> Iterator[OwnedProxy]:
assert "--timeout_worker_healthcheck" not in extra_arguments, "Choose the worker healthcheck with `healthcheck`"
root: Final = _proxy_root()
environment: Final = _proxy_environment(
gateway,
{**overrides, **({"COVERAGE_FILE": str(coverage_file)} if coverage_file else {})},
remove_environment,
inherited,
)
output: Final = Path(inherited.get("INTEGRATION_RESULTS_DIR", str(directory)))
output.mkdir(parents=True, exist_ok=True)
command: Final = (
sys.executable,
"-m",
*(("coverage", "run", f"--rcfile={MCP_COVERAGE_CONFIG}", "-m") if coverage_file else ()),
"integration._support.proxy",
"--config",
str(worker_configuration(config or root / "tests/integration/proxy_config.yaml", directory, inherited)),
"--host",
"127.0.0.1",
"--num_workers",
str(workers),
"--timeout_worker_healthcheck",
str(healthcheck_seconds(healthcheck)),
*database_setup,
*extra_arguments,
)
launch: Final = _launch_until_ready(command, root, environment, output, workers)
process: Final = launch.process
try:
with httpx.Client(
base_url=f"http://127.0.0.1:{launch.port}", timeout=client_timeout, trust_env=False, limits=GATEWAY_LIMITS
) as client:
yield OwnedProxy(Gateway(client, gateway.key, gateway.upstream_url), process, launch.log, healthcheck)
finally:
_stop_launch(launch)
with owned_environment(directory, overrides, remove_environment, state, inherited) as owned:
configured: Final = worker_configuration(
config or root / "tests/integration/proxy_config.yaml", directory, inherited
)
owned_configuration: Final = (
configured if owned.database_url is None else private_configuration(configured, directory)
)
environment: Final = _proxy_environment(
gateway,
{**owned.overrides, **({"COVERAGE_FILE": str(coverage_file)} if coverage_file else {})},
remove_environment,
inherited,
)
output: Final = Path(inherited.get("INTEGRATION_RESULTS_DIR", str(directory)))
output.mkdir(parents=True, exist_ok=True)
command: Final = (
sys.executable,
"-m",
*(("coverage", "run", f"--rcfile={MCP_COVERAGE_CONFIG}", "-m") if coverage_file else ()),
"integration._support.proxy",
"--config",
str(owned_configuration),
"--host",
"127.0.0.1",
"--num_workers",
str(workers),
"--timeout_worker_healthcheck",
str(healthcheck_seconds(healthcheck)),
*database_setup,
*extra_arguments,
)
launch: Final = _launch_until_ready(command, root, environment, output, workers)
process: Final = launch.process
try:
with httpx.Client(
base_url=f"http://127.0.0.1:{launch.port}",
timeout=client_timeout,
trust_env=False,
limits=GATEWAY_LIMITS,
) as client:
yield OwnedProxy(
Gateway(client, gateway.key, gateway.upstream_url, owned.database_url),
process,
launch.log,
healthcheck,
)
finally:
_stop_launch(launch)
@contextmanager

3
tests/integration/_support/prometheus_series.py
View file

@ -19,7 +19,7 @@ import httpx
import yaml
from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment, object_value
from integration._support.database import read_rows
from integration._support.process import OwnedProxy, owned_proxy_process
from integration._support.process import OwnedProxy, OwnedState, owned_proxy_process
from integration._support.responses_vendor import ResponsesVendor, same_response
from integration._support.wire import Reply, Request, Wire, wire_server
from prometheus_client.mmap_dict import MmapedDict
@ -207,6 +207,7 @@ def series_cap_rig(
config=config,
remove_environment=() if shared_samples else ("PROMETHEUS_MULTIPROC_DIR",),
workers=workers,
state=OwnedState.SHARED_WITH_WORKER,
)
)
scenario: Final = stack.enter_context(owned.gateway.scenario())

13
tests/integration/_support/safeguards_rig.py
View file

@ -30,7 +30,7 @@ from integration._support.client import (
string_value,
)
from integration._support.database import read_rows
from integration._support.process import OwnedProxy, owned_proxy_process
from integration._support.process import OwnedProxy, OwnedState, owned_proxy_process
from integration._support.wire import Reply, Request, Wire, wire_server
from pydantic import JsonValue, TypeAdapter
@ -287,6 +287,7 @@ def safeguards_rig(
directory: Path,
*,
environment: Callable[[str], Mapping[str, str]] = no_environment,
remove_environment: tuple[str, ...] = (),
workers: int = 2,
classifier_model: JsonValue | Omitted = CLASSIFIER_MODEL,
general_settings: Mapping[str, JsonValue] | None = None,
@ -305,7 +306,15 @@ def safeguards_rig(
litellm_settings=litellm_settings,
top_level=top_level,
)
with owned_proxy_process(shared, directory, dict(environment(wire.url)), config=path, workers=workers) as owned:
with owned_proxy_process(
shared,
directory,
dict(environment(wire.url)),
config=path,
remove_environment=remove_environment,
workers=workers,
state=OwnedState.SHARED_WITH_WORKER,
) as owned:
yield Rig(wire, scripts, owned)

8
tests/integration/_support/workers.py
View file

@ -86,9 +86,9 @@ def worker_endpoints() -> Mapping[str, str]:
)
def _replica_urls(database: str) -> Mapping[str, str]:
writer: Final = urlsplit(os.environ["INTEGRATION_PROXY_DATABASE_URL"])
reader: Final = urlsplit(os.environ["INTEGRATION_PROXY_READ_REPLICA_URL"])
def replica_urls(database: str, inherited: Mapping[str, str] = os.environ) -> Mapping[str, str]:
writer: Final = urlsplit(inherited["INTEGRATION_PROXY_DATABASE_URL"])
reader: Final = urlsplit(inherited["INTEGRATION_PROXY_READ_REPLICA_URL"])
database_path: Final = urlsplit(database).path
assert writer.username is not None and reader.username is not None
with psycopg.connect(database, autocommit=True) as admin:
@ -164,7 +164,7 @@ def worker_services(
"INTEGRATION_PROXY_READ_REPLICA_URL": "",
"REDIS_HOST": cache.host,
"REDIS_PORT": str(cache.port),
**(_replica_urls(database) if replica else {}),
**(replica_urls(database, inherited) if replica else {}),
}
)

2
tests/integration/authorization/test_access_group_budget.py
View file

@ -22,6 +22,8 @@ from integration._support.process import owned_proxy
from integration._support.wire import Reply, Request, Wire, wire_server
from pydantic import BaseModel, TypeAdapter
pytestmark: Final = pytest.mark.shares_worker_state
_INPUT_RATE: Final = 0.000001
_OUTPUT_RATE: Final = 0.00001
_MAX_BUDGET: Final = 0.0001

2
tests/integration/authorization/test_alias_chain_budget_waiver.py
View file

@ -62,7 +62,7 @@ from integration.authorization._hidden_alias_budget import (
)
from pydantic import JsonValue, TypeAdapter
pytestmark: Final = pytest.mark.timeout(240)
pytestmark: Final = [pytest.mark.timeout(240), pytest.mark.shares_worker_state]
PAID_INPUT_COST_PER_TOKEN: Final = 0.001
PAID_OUTPUT_COST_PER_TOKEN: Final = 0.002

2
tests/integration/authorization/test_claude_code_gateway_oauth.py
View file

@ -23,6 +23,8 @@ from integration._support.process import owned_proxy_process
from integration._support.wire import Reply, Request, Wire, wire_server
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter
pytestmark: Final = pytest.mark.shares_worker_state
_CLIENT_ID: Final = "scripted-cli-client"
_CLIENT_SECRET: Final = "scripted-cli-secret"
_AUTHORIZATION_CODE: Final = "scripted-auth-code"

2
tests/integration/authorization/test_cli_sso_login_ui_disabled.py
View file

@ -34,7 +34,7 @@ from tests.integration._support.process import OwnedProxy, group_members, owned_
from tests.integration._support.provider import SharedProvider
from tests.integration._support.wire import Reply, Request, Wire, wire_server
pytestmark: Final = pytest.mark.timeout(300)
pytestmark: Final = [pytest.mark.timeout(300), pytest.mark.shares_worker_state]
CLI_SOURCE: Final = "litellm-cli"
CLI_STATE_PREFIX: Final = "litellm-session-token"

2
tests/integration/authorization/test_end_user_block.py
View file

@ -12,6 +12,8 @@ from integration._support.database import read_rows
from integration._support.process import owned_proxy
from pydantic import BaseModel, JsonValue
pytestmark: Final = pytest.mark.shares_worker_state
class _BlockedCustomerInfo(BaseModel):
user_id: str

2
tests/integration/authorization/test_hidden_alias_budget_bypass.py
View file

@ -62,7 +62,7 @@ from integration.authorization._hidden_alias_budget import (
)
from pydantic import JsonValue
pytestmark: Final = pytest.mark.timeout(240)
pytestmark: Final = [pytest.mark.timeout(240), pytest.mark.shares_worker_state]
# $0 comes from litellm's own cost-map row (model_prices_and_context_window.json, added 2025-10-14 in 900a57f5b1f):
# Lemonade is a local inference server, so there is no vendor price for the row to drift from. A row that gains a

2
tests/integration/authorization/test_hidden_alias_budget_bypass_chaos.py
View file

@ -32,7 +32,7 @@ from integration.authorization._hidden_alias_budget import (
upstream_requests,
)
pytestmark: Final = pytest.mark.timeout(240)
pytestmark: Final = [pytest.mark.timeout(240), pytest.mark.shares_worker_state]
_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]")
_WAVE: Final = 8

2
tests/integration/authorization/test_jwt_auto_register_map_existing_key.py
View file

@ -24,7 +24,7 @@ from tests.integration._support.wire import Reply, Request, wire_server
KEY_ID: Final = "integration-jwt-map-existing-key"
MAPPING_INSERT: Final = b'INSERT INTO "public"."LiteLLM_JWTKeyMapping"'
pytestmark = pytest.mark.timeout(240)
pytestmark: Final = [pytest.mark.timeout(240), pytest.mark.shares_worker_state]
def _hash(key: str) -> str:

2
tests/integration/authorization/test_jwt_default_team_provisioning.py
View file

@ -14,6 +14,8 @@ from tests.integration._support.database import read_rows
from tests.integration._support.process import owned_proxy
from tests.integration._support.wire import Reply, Request, wire_server
pytestmark: Final = pytest.mark.shares_worker_state
KEY_ID: Final = "integration-jwt-signing-key"
TEAM_BUDGET: Final = 25.0

2
tests/integration/authorization/test_jwt_key_mapping_lifecycle.py
View file

@ -36,7 +36,7 @@ _KEY_ID: Final = "integration-signing-key"
_UPSTREAM_KEY: Final = "synthetic-openai-key"
_CHANNEL: Final = "litellm_proxy.auth_cache_invalidation"
_CHAT_BODY: Final = TypeAdapter(dict[str, JsonValue])
pytestmark: Final = pytest.mark.timeout(180)
pytestmark: Final = [pytest.mark.timeout(180), pytest.mark.shares_worker_state]
class _Usage(BaseModel):

2
tests/integration/authorization/test_jwt_mapped_key_email_backfill.py
View file

@ -14,6 +14,8 @@ from integration._support.process import owned_proxy
from integration._support.wire import Reply, Request, wire_server
from jwt.algorithms import RSAAlgorithm
pytestmark: Final = pytest.mark.shares_worker_state
AUDIENCE: Final = "litellm-integration"
KEY_ID: Final = "integration-signing-key"
CLIENT_CLAIM: Final = "client_id"

3
tests/integration/authorization/test_key_guardrail_opt_out.py
View file

@ -3,6 +3,7 @@ from pathlib import Path
from typing import Final
import httpx
import pytest
import yaml
from integration._support.client import Gateway, Scenario, object_value, string_value
@ -16,6 +17,8 @@ from integration.authorization._guardrail_opt_out import (
stored_metadata,
)
pytestmark: Final = pytest.mark.shares_worker_state
_KEY_ROUTES: Final = ["/key/generate", "/key/update", "/key/regenerate", "/v1/chat/completions"]

3
tests/integration/authorization/test_key_guardrail_opt_out_chaos.py
View file

@ -9,6 +9,7 @@ from queue import SimpleQueue
from typing import Final
import httpx
import pytest
from integration._support.client import Gateway, eventually, string_value
from integration._support.database import read_rows
@ -24,6 +25,8 @@ from integration.authorization._guardrail_opt_out import (
upstream_observations,
)
pytestmark: Final = pytest.mark.shares_worker_state
class _GuardrailSink:
"""Test-owned guardrail endpoint that can be stopped and restarted on the same port."""

3
tests/integration/authorization/test_key_guardrail_opt_out_runtime.py
View file

@ -5,6 +5,7 @@ from pathlib import Path
from typing import Final
import httpx
import pytest
from anthropic import Anthropic
from openai import AsyncOpenAI, OpenAI
@ -20,6 +21,8 @@ from integration.authorization._guardrail_opt_out import (
upstream_hits,
)
pytestmark: Final = pytest.mark.shares_worker_state
def _wire_hits(wire: Wire, marker: str) -> int:
return sum(1 for request in wire.drain() if marker.encode() in request.body)

2
tests/integration/authorization/test_rag_query_vector_store_allowlist.py
View file

@ -25,6 +25,8 @@ from openai.types.chat import ChatCompletion
from pydantic import JsonValue, TypeAdapter
from redis import Redis
pytestmark: Final = pytest.mark.shares_worker_state
CONFIG_STORE_ID: Final = "vs_integration_config_store"
PROXY_CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml"
REMOVE_OPENAI_API_BASE: Final = ("OPENAI_API_BASE",)

2
tests/integration/authorization/test_rag_team_provider_credentials.py
View file

@ -89,7 +89,7 @@ AWS_AMBIENT_CREDENTIALS: Final = (
"AWS_WEB_IDENTITY_TOKEN_FILE",
)
pytestmark: Final = pytest.mark.timeout(OWNED_PROXY_CELL_SECONDS)
pytestmark: Final = [pytest.mark.timeout(OWNED_PROXY_CELL_SECONDS), pytest.mark.shares_worker_state]
Signature = tuple[str, str, str, str | None]

2
tests/integration/authorization/test_responses_compact_id_ownership.py
View file

@ -50,6 +50,8 @@ from integration.authorization._responses_compact import (
)
from pydantic import JsonValue
pytestmark: Final = pytest.mark.shares_worker_state
_CLIENTS: Final = ("sync", "async", "httpx")

2
tests/integration/authorization/test_responses_compact_id_ownership_chaos.py
View file

@ -40,6 +40,8 @@ from integration.authorization._responses_compact import (
vendor_id,
)
pytestmark: Final = pytest.mark.shares_worker_state
_Kind = Literal["owner", "other_user", "other_team", "unmanaged"]
_KINDS: Final[tuple[_Kind, ...]] = ("owner", "other_user", "other_team", "unmanaged")
_REFUSING: Final[tuple[_Kind, ...]] = ("other_user", "other_team", "unmanaged")

2
tests/integration/authorization/test_responses_id_ownership.py
View file

@ -19,6 +19,8 @@ from integration._support.worker_readiness import wait_for_worker_readiness
from openai.types.responses import ResponseInputParam
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter
pytestmark: Final = pytest.mark.shares_worker_state
_MODEL: Final = "gpt-5"
_API_KEY: Final = "synthetic-responses-key"
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])

2
tests/integration/authorization/test_sso_blank_identity.py
View file

@ -30,7 +30,7 @@ from tests.integration._support.database_relay import dropped_connection_relay
from tests.integration._support.process import OwnedProxy, owned_proxy_process
from tests.integration._support.wire import Reply, Request, Wire, wire_server
pytestmark: Final = pytest.mark.timeout(300)
pytestmark: Final = [pytest.mark.timeout(300), pytest.mark.shares_worker_state]
CLIENT_ID: Final = "integration-entra-client"
CLIENT_SECRET: Final = "integration-entra-secret"

2
tests/integration/authorization/test_tool_search_replay_guardrails.py
View file

@ -18,7 +18,7 @@ from pydantic import JsonValue
from tests.integration.caching.response_cache_case import await_entry, redis_store
pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
pytestmark: Final = [pytest.mark.timeout(2 * graceful_stop_seconds() + 120), pytest.mark.shares_worker_state]
_POLICY: Final = "integration-tool-policy"
_BLOCK: Final = "tool-permission-block"

2
tests/integration/caching/test_cache_max_messages.py
View file

@ -19,6 +19,8 @@ from pydantic import JsonValue, TypeAdapter
from redis import Redis
from redis.exceptions import ResponseError
pytestmark: Final = pytest.mark.shares_worker_state
_Turns = Callable[[str], tuple[list[JsonValue], list[JsonValue]]]
_INDEXES: Final = TypeAdapter(tuple[bytes, ...])
_COUNT: Final = TypeAdapter(tuple[int])

2
tests/integration/caching/test_response_cache_chaos.py
View file

@ -61,6 +61,8 @@ from tests.integration.caching.response_cache_case import (
slowed,
)
pytestmark: Final = pytest.mark.shares_worker_state
RecordProperty = Callable[[str, object], None]
WORKERS: Final = 2

2
tests/integration/compatibility/_status_code_audit.py
View file

@ -248,7 +248,7 @@ def owned_gateway(
gateway_from_environment() as gateway,
owned_proxy_process(gateway, directory, environment or {}, config=path, workers=workers) as owned,
):
yield Gateway(owned.gateway.client, owned.gateway.key, gateway.upstream_url)
yield owned.gateway
def stream_finished(endpoint: str, lines: Sequence[str]) -> bool:

2
tests/integration/compatibility/test_codex_model_catalog.py
View file

@ -32,6 +32,8 @@ from tests.integration._support.database import read_rows
from tests.integration._support.process import owned_proxy
from tests.integration._support.wire import Reply, Request, Wire, wire_server
pytestmark: Final = pytest.mark.shares_worker_state
T = TypeVar("T")
CLIENT_VERSION: Final = "0.159.3"

2
tests/integration/compatibility/test_missing_body_param_status.py
View file

@ -43,6 +43,8 @@ from tests.integration.cost_calculation.cost_tracking_case import (
SseResponse,
)
pytestmark: Final = pytest.mark.shares_worker_state
_Route = tuple[str, str, tuple[str, ...], dict[str, JsonValue], str]
_EVAL_DATA_SOURCE_CONFIG: Final[dict[str, JsonValue]] = {
"type": "custom",

2
tests/integration/compatibility/test_openai_sdk_client_paths_chaos.py
View file

@ -26,6 +26,8 @@ from pydantic import JsonValue
from tests.integration.cost_calculation.cost_tracking_case import JsonResponse, RoutedResponse, SseResponse
pytestmark: Final = pytest.mark.shares_worker_state
_Kind: TypeAlias = Literal["chat", "chat_stream", "messages", "responses", "completions", "moderations"]
_ALL_KINDS: Final[tuple[_Kind, ...]] = ("chat", "chat_stream", "messages", "responses", "completions", "moderations")

2
tests/integration/compatibility/test_persisted_toolsets.py
View file

@ -12,6 +12,8 @@ from integration._support.database import read_rows
from integration._support.mcp import call_tool, mcp_peer, register_mcp, tool_names
from integration._support.process import owned_proxy
pytestmark: Final = pytest.mark.shares_worker_state
@pytest.mark.covers("other.compatibility.mcp.persisted_tool_names_survive_candidate_startup")
def test_existing_toolset_format_loads_before_start_and_keeps_sibling_denied(gateway: Gateway, tmp_path: Path) -> None:

2
tests/integration/compatibility/test_status_code_followups.py
View file

@ -207,7 +207,7 @@ def router_default_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[G
},
)
with owned_proxy_process(gateway, directory, {}, config=config) as owned:
yield Gateway(owned.gateway.client, owned.gateway.key, gateway.upstream_url)
yield owned.gateway
def test_ocr_multipart_without_file_returns_400(gateway: Gateway) -> None:

2
tests/integration/compatibility/test_status_code_followups_chaos.py
View file

@ -41,7 +41,7 @@ from pydantic import JsonValue
from tests.integration.cost_calculation.cost_tracking_case import StoredResponse
pytestmark: Final = pytest.mark.timeout(480)
pytestmark: Final = [pytest.mark.timeout(480), pytest.mark.shares_worker_state]
_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]")
_MARKER_HEADER: Final = "x-litellm-spend-logs-metadata"

2
tests/integration/configuration/test_codex_catalog_boot.py
View file

@ -16,6 +16,8 @@ from pydantic import JsonValue
from tests.integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, object_value, string_value
from tests.integration._support.process import owned_proxy, owned_proxy_process
pytestmark: Final = pytest.mark.shares_worker_state
CLIENT_VERSION: Final = "0.159.3"
FRESH_CONNECTION: Final = {"Connection": "close"}
CONSISTENT_READS: Final = 8

2
tests/integration/configuration/test_effective_settings.py
View file

@ -23,6 +23,8 @@ from tests.integration._support.database import read_rows, scratch_database
from tests.integration._support.process import owned_proxy
from tests.integration._support.wire import Reply, Request, Wire, wire_server
pytestmark: Final = pytest.mark.shares_worker_state
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])

30
tests/integration/conftest.py
View file

@ -5,7 +5,7 @@ import json
import os
import re
from collections.abc import AsyncIterator, Generator, Iterator, Mapping, Sequence
from contextlib import ExitStack
from contextlib import ExitStack, contextmanager
from copy import copy
from functools import partial
from importlib.metadata import version
@ -23,6 +23,7 @@ from tests.integration._support.client import Gateway, Scenario, eventually, gat
from tests.integration._support.generation import LIFECYCLE_SETTINGS
from tests.integration._support.manifest import OWNED_DIRECTORIES
from tests.integration._support.ordering import collected_case, fixture_groups, scheduling_groups
from tests.integration._support.process import OWNED_STATE_VARIABLE, OwnedState
from tests.integration._support.provider import SharedProvider, shared_provider
from tests.integration._support.routing import RoutingPlugin
from tests.integration._support.sdk_models import build_sdk_models
@ -69,6 +70,7 @@ def pytest_addoption(parser: pytest.Parser) -> None:
def pytest_configure(config: pytest.Config) -> None:
config.addinivalue_line("markers", "integration: owned real-service integration contracts")
config.addinivalue_line("markers", "covers(*ids): legacy contract IDs kept for existing tests, not enforced")
config.addinivalue_line("markers", "shares_worker_state: owned proxies here share the worker DB and Redis")
config.stash[REPORTS] = []
config.stash[SHARD_FILES] = _shard_files(SHARD_PATH.validate_python(config.getoption("integration_shard_files")))
config.pluginmanager.register(IntegrationReportPlugin(config))
@ -295,6 +297,32 @@ def pytest_sessionfinish(session: pytest.Session, exitstatus: int) -> None:
session.exitstatus = pytest.ExitCode.OK
@contextmanager
def _owned_state(item: pytest.Item) -> Iterator[None]:
with pytest.MonkeyPatch.context() as environment:
if item.get_closest_marker("shares_worker_state") is not None:
environment.setenv(OWNED_STATE_VARIABLE, OwnedState.SHARED_WITH_WORKER.value)
yield
@pytest.hookimpl(wrapper=True)
def pytest_runtest_setup(item: pytest.Item) -> Generator[None, None, None]:
with _owned_state(item):
return (yield)
@pytest.hookimpl(wrapper=True)
def pytest_runtest_call(item: pytest.Item) -> Generator[None, None, None]:
with _owned_state(item):
return (yield)
@pytest.hookimpl(wrapper=True)
def pytest_runtest_teardown(item: pytest.Item, nextitem: pytest.Item | None) -> Generator[None, None, None]:
with _owned_state(item):
return (yield)
@pytest.fixture
def gateway() -> Iterator[Gateway]:
with gateway_from_environment() as value:

2
tests/integration/cost_calculation/test_base_model_cross_provider_spend.py
View file

@ -16,7 +16,7 @@ from integration._support.process import OwnedProxy, owned_proxy_process
from integration._support.wire import Reply, Request, Wire, wire_server
from pydantic import JsonValue
pytestmark: Final = pytest.mark.timeout(240)
pytestmark: Final = [pytest.mark.timeout(240), pytest.mark.shares_worker_state]
_MODEL: Final = "cc-subscription-sol"
_INPUT_TOKENS: Final = 1000

2
tests/integration/cost_calculation/test_bridged_previous_response_id_spend.py
View file

@ -17,7 +17,7 @@ from integration._support.workers import cost_map_environment
from integration.cost_calculation.cost_tracking_case import COST_MAP
from pydantic import JsonValue, TypeAdapter
pytestmark: Final = pytest.mark.timeout(240)
pytestmark: Final = [pytest.mark.timeout(240), pytest.mark.shares_worker_state]
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
_MESSAGES: Final = TypeAdapter(list[dict[str, JsonValue]])

2
tests/integration/cost_calculation/test_responses_websocket_guardrail_spend.py
View file

@ -21,7 +21,7 @@ from integration._support.responses_ws_peer import PROVIDER_MODEL, USAGE, create
from integration._support.tls import write_self_signed_cert
from pydantic import JsonValue
pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
pytestmark: Final = [pytest.mark.timeout(2 * graceful_stop_seconds() + 120), pytest.mark.shares_worker_state]
_INPUT_COST_PER_TOKEN: Final = 0.002
_OUTPUT_COST_PER_TOKEN: Final = 0.004

2
tests/integration/database/test_reader_writer_regeneration.py
View file

@ -17,6 +17,8 @@ from integration._support.wire import Reply, Request, wire_server
from psycopg import sql
from pydantic import BaseModel, JsonValue, TypeAdapter
pytestmark: Final = pytest.mark.shares_worker_state
@pytest.mark.covers("other.database.regeneration.writer_updates_dependent_grants")
def test_key_regeneration_uses_writer_with_a_real_readonly_reader(gateway: Gateway, tmp_path: Path) -> None:

2
tests/integration/management/test_batch_file_usage_caps.py
View file

@ -59,7 +59,7 @@ from integration.management._batch_file_caps import (
from pydantic import JsonValue
from redis import Redis
pytestmark = pytest.mark.timeout(600)
pytestmark: Final = [pytest.mark.timeout(600), pytest.mark.shares_worker_state]
YAML_RECORDS: Final = 3
YAML_UPLOADS: Final = 4

2
tests/integration/management/test_batch_file_usage_caps_chaos.py
View file

@ -46,7 +46,7 @@ from integration.management._batch_file_caps import (
)
from pydantic import JsonValue
pytestmark = pytest.mark.timeout(600)
pytestmark: Final = [pytest.mark.timeout(600), pytest.mark.shares_worker_state]
WORKERS: Final = 2
UPLOAD_CAP: Final = 5

Some files were not shown because too many files have changed in this diff Show more