scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-08 22:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_lit_6342_claude_subagent_router

Browse source
This commit is contained in:
mateo-berri 2026-09-02 12:10:37 -07:00
commit 01b20b78c0
103 changed files with 8187 additions and 627 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
.github/workflows/image-scan.yml vendored
View file

@ -80,7 +80,7 @@ jobs:
LITELLM_IMAGE: litellm-image-scan:${{ github.sha }}
run: |
python -m pip install "pytest==9.0.3"
python -m pytest tests/proxy_migration_tests/test_offline_image_migration.py -v
python -m pytest tests/proxy_migration_tests/test_offline_image_migration.py tests/proxy_migration_tests/test_image_bedrock_realtime_extra.py -v
# Scans the whole shipped artifact: OS/apk plus every language package
# baked into the image, including ones no lockfile declares (e.g. prisma's
@ -124,7 +124,7 @@ jobs:
LITELLM_IMAGE: litellm-runtime-scan:${{ github.sha }}
run: |
python -m pip install "pytest==9.0.3"
python -m pytest tests/proxy_migration_tests/test_offline_image_migration.py -v
python -m pytest tests/proxy_migration_tests/test_offline_image_migration.py tests/proxy_migration_tests/test_image_bedrock_realtime_extra.py -v
migrations-image:
name: migrations-image
@ -185,7 +185,7 @@ jobs:
LITELLM_COMPONENT_PORT: "4000"
run: |
python -m pip install "pytest==9.0.3"
python -m pytest tests/proxy_migration_tests/test_component_image_serves_offline.py -v
python -m pytest tests/proxy_migration_tests/test_component_image_serves_offline.py tests/proxy_migration_tests/test_image_bedrock_realtime_extra.py -v
ui-image:
name: ui-image

2
Dockerfile
View file

@ -66,6 +66,7 @@ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-gr
--extra extra_proxy \
--extra semantic-router \
--extra saml \
--extra bedrock-realtime \
--python python3.13
# Copy full source tree
@ -87,6 +88,7 @@ RUN uv sync --frozen --no-default-groups --no-editable \
--extra extra_proxy \
--extra semantic-router \
--extra saml \
--extra bedrock-realtime \
--python python3.13
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \

2
backend/Dockerfile
View file

@ -46,6 +46,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
--extra proxy-runtime \
--extra extra_proxy \
--extra semantic-router \
--extra saml \
--python python3.13
# Stage 2 — copy source and install the project + workspace members.
@ -57,6 +58,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
--extra proxy-runtime \
--extra extra_proxy \
--extra semantic-router \
--extra saml \
--python python3.13
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \

2
basedpyright-code-budget.json
View file

@ -117,7 +117,7 @@
"limit": 111
},
"reportUnnecessaryComparison": {
"limit": 695
"limit": 692
},
"reportUnnecessaryContains": {
"limit": 5

2
docker/Dockerfile.database
View file

@ -64,6 +64,7 @@ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-gr
--extra extra_proxy \
--extra semantic-router \
--extra saml \
--extra bedrock-realtime \
--python python3.13
# Copy full source tree
@ -85,6 +86,7 @@ RUN uv sync --frozen --no-default-groups --no-editable \
--extra extra_proxy \
--extra semantic-router \
--extra saml \
--extra bedrock-realtime \
--python python3.13
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \

3
docker/Dockerfile.non_root
View file

@ -70,6 +70,7 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \
--extra extra_proxy \
--extra semantic-router \
--extra saml \
--extra bedrock-realtime \
--python python3.13
# Copy full source tree
@ -97,6 +98,7 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \
--extra extra_proxy \
--extra semantic-router \
--extra saml \
--extra bedrock-realtime \
--python python3.13 \
--no-sources-package litellm-proxy-extras; \
else \
@ -106,6 +108,7 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \
--extra extra_proxy \
--extra semantic-router \
--extra saml \
--extra bedrock-realtime \
--python python3.13; \
fi

2
helm/litellm-helm/Chart.yaml
View file

@ -18,7 +18,7 @@ type: application
# This is the chart version. This version number should be incremented each time you make changes
# to the chart and its templates, including the app version.
# Versions are expected to follow Semantic Versioning (https://semver.org/)
version: 1.1.2
version: 1.1.3
# This is the version number of the application being deployed. This version number should be
# incremented each time you make changes to the application. Versions are not expected to

4
helm/litellm-helm/README.md
View file

@ -26,7 +26,7 @@ If `db.useStackgresOperator` is used (not yet implemented):
| `replicaCount` | The number of LiteLLM Proxy pods to be deployed | `1` |
| `masterkeySecretName` | The name of the Kubernetes Secret that contains the Master API Key for LiteLLM. If not specified, use the generated secret name. | N/A |
| `masterkeySecretKey` | The key within the Kubernetes Secret that contains the Master API Key for LiteLLM. If not specified, use `masterkey` as the key. | N/A |
| `masterkey` | The Master API Key for LiteLLM. If not specified, a random key in the `sk-...` format is generated. | N/A |
| `masterkey` | The Master API Key for LiteLLM. If not specified, a random key in the `sk-...` format is generated on first install and reused on upgrades. | N/A |
| `environmentSecrets` | An optional array of Secret object names. The keys and values in these secrets will be presented to the LiteLLM proxy pod as environment variables. See below for an example Secret object. | `[]` |
| `environmentConfigMaps` | An optional array of ConfigMap object names. The keys and values in these configmaps will be presented to the LiteLLM proxy pod as environment variables. See below for an example Secret object. | `[]` |
| `image.repository` | LiteLLM Proxy image repository | `ghcr.io/berriai/litellm` |
@ -212,6 +212,8 @@ service, the **Proxy Endpoint** should be set to `http://<RELEASE>-litellm:4000`
The **Proxy Key** is the value specified for `masterkey` or, if a `masterkey`
was not provided to the helm command line, the `masterkey` is a randomly
generated string in the `sk-...` format stored in the `<RELEASE>-litellm-masterkey` Kubernetes Secret.
The key is generated once on the first install; later `helm upgrade` runs reuse the
value already in that Secret, so upgrading never rotates the master key.
```bash
kubectl -n litellm get secret <RELEASE>-litellm-masterkey -o jsonpath="{.data.masterkey}"

6
helm/litellm-helm/templates/secret-masterkey.yaml
View file

@ -1,9 +1,11 @@
{{- if not .Values.masterkeySecretName }}
{{ $masterkey := (.Values.masterkey | default (printf "sk-%s" (randAlphaNum 18))) }}
{{- $secretName := printf "%s-masterkey" (include "litellm.fullname" .) }}
{{- $existing := lookup "v1" "Secret" .Release.Namespace $secretName }}
{{- $masterkey := .Values.masterkey | default (dig "data" "masterkey" "" $existing | b64dec) | default (printf "sk-%s" (randAlphaNum 18)) }}
apiVersion: v1
kind: Secret
metadata:
name: {{ include "litellm.fullname" . }}-masterkey
name: {{ $secretName }}
data:
masterkey: {{ $masterkey | b64enc }}
type: Opaque

42
helm/litellm-helm/tests/hpa_tests.yaml
View file

@ -1,4 +1,4 @@
suite: "hpa with behavior"
suite: "hpa"
templates:
- hpa.yaml
tests:
@ -23,14 +23,44 @@ tests:
- equal: { path: spec.behavior.scaleUp.stabilizationWindowSeconds, value: 60 }
- equal: { path: spec.behavior.scaleDown.stabilizationWindowSeconds, value: 90 }
---
suite: "hpa without behavior"
templates:
- hpa.yaml
tests:
- it: "does not render behavior when not set"
set:
autoscaling.enabled: true
asserts:
- isKind: { of: HorizontalPodAutoscaler }
- isNull: { path: spec.behavior }
- it: "scales on cpu at the documented 60 percent by default"
set:
autoscaling.enabled: true
asserts:
- isKind: { of: HorizontalPodAutoscaler }
- equal: { path: "spec.metrics[0].resource.name", value: cpu }
- equal: { path: "spec.metrics[0].resource.target.type", value: Utilization }
- equal: { path: "spec.metrics[0].resource.target.averageUtilization", value: 60 }
- it: "does not scale on memory by default"
set:
autoscaling.enabled: true
asserts:
- lengthEqual: { path: spec.metrics, count: 1 }
- it: "honours an explicit cpu target override"
set:
autoscaling.enabled: true
autoscaling.targetCPUUtilizationPercentage: 75
asserts:
- equal: { path: "spec.metrics[0].resource.target.averageUtilization", value: 75 }
- it: "renders a memory metric only when a memory target is set"
set:
autoscaling.enabled: true
autoscaling.targetMemoryUtilizationPercentage: 80
asserts:
- lengthEqual: { path: spec.metrics, count: 2 }
- equal: { path: "spec.metrics[1].resource.name", value: memory }
- equal: { path: "spec.metrics[1].resource.target.averageUtilization", value: 80 }
- it: "renders no hpa when autoscaling is disabled"
asserts:
- hasDocuments: { count: 0 }

47
helm/litellm-helm/tests/masterkey-secret_tests.yaml
View file

@ -15,6 +15,53 @@ tests:
# Note: The masterkey is generated as "sk-<18-random-chars>" in plain text,
# but stored as base64 encoded in Kubernetes secret (requirement).
# "sk-" base64 encodes to "c2st", so we check for "^c2st" pattern.
- it: should reuse the master key already stored in the cluster instead of generating a new one on upgrade
template: secret-masterkey.yaml
set:
masterkeySecretName: ""
kubernetesProvider:
scheme:
"v1/Secret":
gvr:
version: "v1"
resource: "secrets"
namespaced: true
objects:
- kind: Secret
apiVersion: v1
metadata:
name: RELEASE-NAME-litellm-masterkey
namespace: NAMESPACE
data:
masterkey: c2stZXhpc3Rpbmcta2V5
asserts:
- equal:
path: data.masterkey
value: c2stZXhpc3Rpbmcta2V5
- it: should let an explicit masterkey value override the one already stored in the cluster
template: secret-masterkey.yaml
set:
masterkeySecretName: ""
masterkey: sk-explicit
kubernetesProvider:
scheme:
"v1/Secret":
gvr:
version: "v1"
resource: "secrets"
namespaced: true
objects:
- kind: Secret
apiVersion: v1
metadata:
name: RELEASE-NAME-litellm-masterkey
namespace: NAMESPACE
data:
masterkey: c2stZXhpc3Rpbmcta2V5
asserts:
- equal:
path: data.masterkey
value: c2stZXhwbGljaXQ=
- it: should not create a secret if masterkeySecretName is set
template: secret-masterkey.yaml
set:

11
helm/litellm-helm/values.yaml
View file

@ -200,7 +200,16 @@ autoscaling:
enabled: false
minReplicas: 1
maxReplicas: 100
targetCPUUtilizationPercentage: 80
# 60 is the documented recommendation. See "Recommended Machine Specifications"
# in https://docs.litellm.ai/docs/proxy/prod. A new replica clears the startupProbe
# above only after up to failureThreshold x periodSeconds = 300 seconds, so a target
# high enough to trip near saturation adds capacity minutes after it was needed.
targetCPUUtilizationPercentage: 60
# Deliberately left unset rather than given a value. The prisma query engine's
# resident memory is a high-water mark that ratchets to the pod's worst-ever write
# and is never returned, so a memory target reads the largest write a pod ever did
# rather than what it is doing now, and replicas ratchet up without scaling back in.
# Memory is a floor to provision under 'resources', not a signal to scale on.
# targetMemoryUtilizationPercentage: 80
# behavior: {}

65
litellm/integrations/custom_guardrail.py
View file

@ -1,4 +1,5 @@
import contextvars
import copy
import hashlib
import os
import secrets
@ -39,6 +40,7 @@ except ImportError:
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
dc: Final = DualCache()
@ -852,6 +854,69 @@ class CustomGuardrail(CustomLogger):
return result
async def async_logging_hook(
self,
kwargs: dict, # mutable-ok: CustomLogger.async_logging_hook contract
result: object,
call_type: str,
) -> tuple[dict, object]: # mutable-ok: CustomLogger.async_logging_hook contract
"""logging_only: run apply_guardrail on copies of the logged request/response and record the verdict."""
from litellm.llms import get_guardrail_translation_mapping
if not self.uses_apply_guardrail_interface() or self.use_native_lifecycle_hooks:
return kwargs, result
try:
translation: Final = get_guardrail_translation_mapping(CallTypes(call_type))()
except ValueError:
verbose_logger.debug(
"Guardrail %s: no guardrail translation for call_type=%s, skipping logging_only scan",
self.guardrail_name,
call_type,
)
return kwargs, result
litellm_params: Final = kwargs.get("litellm_params") or {}
scratch_metadata: Final = {
key: value
for key, value in (litellm_params.get("metadata") or {}).items()
if key != "standard_logging_guardrail_information"
}
try:
await self._scan_logged_call(kwargs, result, translation, scratch_metadata)
except Exception as e:
verbose_logger.warning("Guardrail %s: logging_only scan raised: %s", self.guardrail_name, e)
recorded: Final = scratch_metadata.get("standard_logging_guardrail_information")
standard_logging_object: Final = kwargs.get("standard_logging_object")
if not recorded or not isinstance(standard_logging_object, dict):
return kwargs, result
entries: Final = recorded if isinstance(recorded, list) else [recorded]
existing: Final = standard_logging_object.get("guardrail_information") or []
return {
**kwargs,
"standard_logging_object": {**standard_logging_object, "guardrail_information": [*existing, *entries]},
}, result
async def _scan_logged_call(
self,
kwargs: dict, # mutable-ok: CustomLogger.async_logging_hook contract
result: object,
translation: "BaseTranslation",
scratch_metadata: dict, # mutable-ok: apply_guardrail records its verdict into request metadata
) -> None:
optional_params: Final = kwargs.get("optional_params") or {}
scratch_input: Final = copy.deepcopy(kwargs.get("messages") or kwargs.get("input"))
scratch_request: Final = {
"model": kwargs.get("model"),
"messages": scratch_input,
"input": scratch_input,
"tools": copy.deepcopy(optional_params.get("tools")),
"litellm_call_id": kwargs.get("litellm_call_id"),
"metadata": scratch_metadata,
}
await translation.process_input_messages(data=scratch_request, guardrail_to_apply=self)
await translation.process_output_response(
response=copy.deepcopy(result), guardrail_to_apply=self, request_data=scratch_request
)
def supports_scan_only_tool_results(self) -> bool:
"""Whether this guardrail can scan tool-result content.

393
litellm/integrations/datadog/datadog_llm_obs.py
View file

@ -11,6 +11,7 @@ import json
import os
from collections.abc import Mapping, Sequence
from datetime import datetime
from types import MappingProxyType
from typing import Any, Final, Literal
import httpx
@ -30,12 +31,16 @@ from litellm.integrations.datadog.datadog_mock_client import (
)
from litellm.litellm_core_utils.dd_tracing import tracer
from litellm.litellm_core_utils.prompt_templates.common_utils import (
convert_content_list_to_str,
handle_any_messages_to_chat_completion_str_messages_conversion,
)
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.proxy.spend_tracking.savings import extract_cache_creation_tokens, extract_cache_read_tokens
from litellm.types.integrations.datadog_llm_obs import *
from litellm.types.utils import (
CallTypes,
@ -44,6 +49,189 @@ from litellm.types.utils import (
StandardLoggingPayloadErrorInformation,
)
_EMPTY_MAPPING: Final[Mapping[str, Any]] = MappingProxyType({})
_EMPTY_MESSAGE: Final[Message] = {"role": "", "content": ""}
_MAX_PARSED_TOOL_ARGUMENT_CHARS: Final = 256 * 1024
def _mapping_field(source: Mapping[str, Any], key: str) -> Mapping[str, Any]:
"""The value at `key` when it is a mapping, else an empty one."""
value: Final = source.get(key)
return value if isinstance(value, dict) else _EMPTY_MAPPING
def _content_blocks(message: Mapping[str, Any]) -> tuple[Mapping[str, Any], ...]:
content: Final = message.get("content")
if not isinstance(content, list):
return ()
return tuple(block for block in content if isinstance(block, dict))
def _to_dd_arguments(raw_arguments: object) -> dict[str, Any] | str:
"""
Arguments as the object LLM Obs types them as, or the raw string when they are not one.
Strings past the size bound ship unparsed: decoding multiplies memory on hostile compact
JSON, and the raw string is what the intake receives either way.
"""
if not isinstance(raw_arguments, str):
return raw_arguments if isinstance(raw_arguments, dict) else str(raw_arguments)
if len(raw_arguments) > _MAX_PARSED_TOOL_ARGUMENT_CHARS:
return raw_arguments
parsed: Final = safe_json_loads(raw_arguments)
return parsed if isinstance(parsed, dict) else raw_arguments
def _to_dd_tool_calls(message: Mapping[str, Any]) -> tuple[ToolCall, ...]:
"""
The tool calls a message carries, in LLM Obs' ToolCall schema, from either dialect.
OpenAI puts them in `tool_calls` with the callee nested under `function` and `arguments`
serialized; Anthropic puts them in `content` as `tool_use` blocks with `input` already an
object. LLM Obs reads `name` / `arguments` / `tool_id` either way.
"""
raw_tool_calls: Final = message.get("tool_calls")
openai_calls: Final = tuple(
ToolCall(
name=function.get("name", ""),
arguments=_to_dd_arguments(function.get("arguments", "")),
tool_id=tool_call.get("id", ""),
type=tool_call.get("type", "function"),
)
for tool_call in (raw_tool_calls if isinstance(raw_tool_calls, list) else ())
if isinstance(tool_call, dict)
for function in [_mapping_field(tool_call, "function")]
)
anthropic_calls: Final = tuple(
ToolCall(
name=block.get("name", ""),
arguments=_to_dd_arguments(block.get("input") or {}),
tool_id=block.get("id", ""),
type="tool_use",
)
for block in _content_blocks(message)
if block.get("type") == "tool_use"
)
return openai_calls + anthropic_calls
def _to_dd_tool_results(message: Mapping[str, Any], tool_call_names: Mapping[str, str]) -> tuple[ToolResult, ...]:
"""
The tool results a message carries, linked back to the call each answers.
OpenAI models a result as a whole `role: "tool"` message keyed by `tool_call_id`;
Anthropic nests `tool_result` blocks inside a user message, keyed by `tool_use_id`.
"""
def to_result(tool_id: str, result: object) -> ToolResult:
return ToolResult(
name=tool_call_names.get(tool_id, ""),
result=result if isinstance(result, str) else safe_dumps(result),
tool_id=tool_id,
type="function",
)
if message.get("role") == "tool":
return (to_result(str(message.get("tool_call_id", "")), message.get("content") or ""),)
return tuple(
to_result(str(block.get("tool_use_id", "")), block.get("content") or "")
for block in _content_blocks(message)
if block.get("type") == "tool_result"
)
def _tool_call_names_by_id(messages: Sequence[object]) -> Mapping[str, str]:
"""Ids to tool names for result linking; reads names structurally and parses nothing."""
openai_pairs: Final = tuple(
(tool_call.get("id"), function.get("name", ""))
for message in messages
if isinstance(message, dict) and isinstance(message.get("tool_calls"), list)
for tool_call in message["tool_calls"]
if isinstance(tool_call, dict)
for function in [_mapping_field(tool_call, "function")]
)
anthropic_pairs: Final = tuple(
(block.get("id"), block.get("name", ""))
for message in messages
if isinstance(message, dict)
for block in _content_blocks(message)
if block.get("type") == "tool_use"
)
return MappingProxyType({str(tool_id): str(name) for tool_id, name in openai_pairs + anthropic_pairs if tool_id})
def _to_dd_message(message: object, tool_call_names: Mapping[str, str]) -> Message:
"""
Map one chat message onto LLM Obs' Message schema, adding fields and never destroying content.
Content collapses to its text only when it has text; a content list with none (tool blocks,
images) rides along unchanged so nothing the caller logged is lost. Tool calls and results
move into the fields the LLM Obs Tools panel reads, from both the OpenAI and Anthropic shapes.
"""
if not isinstance(message, dict):
converted: Final = handle_any_messages_to_chat_completion_str_messages_conversion(message)
return converted[0] if converted else _EMPTY_MESSAGE
text: Final = convert_content_list_to_str(message) # pyright: ignore[reportArgumentType] # caller-supplied dict
original_content: Final = message.get("content")
content: Final = (
text if text or not isinstance(original_content, list) or not original_content else original_content
)
reasoning: Final = message.get("reasoning_content")
tool_calls: Final = _to_dd_tool_calls(message)
tool_results: Final = _to_dd_tool_results(message, tool_call_names)
dd_message: Final[Message] = {
"role": message.get("role", ""),
"content": content,
**({"reasoning_content": reasoning} if reasoning is not None else {}),
**({"tool_calls": tool_calls} if tool_calls else {}),
**({"tool_results": tool_results} if tool_results else {}),
}
return dd_message
def _to_dd_messages(messages: object) -> tuple[Message, ...]:
"""Map a whole conversation, resolving each tool result against the calls that precede it."""
if messages is None:
return ()
if not isinstance(messages, list):
return tuple(handle_any_messages_to_chat_completion_str_messages_conversion(messages))
tool_call_names: Final = _tool_call_names_by_id(messages)
return tuple(_to_dd_message(message, tool_call_names) for message in messages)
def _to_dd_tool_definition(entry: Mapping[str, Any]) -> ToolDefinition | None:
function: Final = entry.get("function")
declared: Final[Mapping[str, Any]] = function if isinstance(function, dict) else entry
name: Final = declared.get("name")
if not name:
return None
schema: Final = declared.get("parameters") or declared.get("input_schema")
description: Final = declared.get("description", "")
if not isinstance(schema, dict):
return ToolDefinition(name=name, description=description)
return ToolDefinition(name=name, description=description, schema=schema)
def _to_dd_tool_definitions(model_parameters: object) -> tuple[ToolDefinition, ...]:
"""
Map the request's declared tools onto LLM Obs' ToolDefinition schema.
Handles the wrapped chat-completions shape and the bare shape the Anthropic and
Responses surfaces use, since both reach this logger through `model_parameters`.
"""
if not isinstance(model_parameters, dict):
return ()
raw_tools: Final = model_parameters.get("tools") or model_parameters.get("functions")
if not isinstance(raw_tools, list):
return ()
return tuple(
definition
for entry in raw_tools
if isinstance(entry, dict)
if (definition := _to_dd_tool_definition(entry)) is not None
)
class DataDogLLMObsLogger(CustomBatchLogger):
def __init__(self, **kwargs):
@ -222,12 +410,9 @@ class DataDogLLMObsLogger(CustomBatchLogger):
if standard_logging_payload is None:
raise Exception("DataDogLLMObs: standard_logging_object is not set")
messages = standard_logging_payload["messages"]
messages = self._ensure_string_content(messages=messages)
metadata: Final = kwargs.get("litellm_params", {}).get("metadata", {})
input_meta: Final = InputMeta(messages=handle_any_messages_to_chat_completion_str_messages_conversion(messages))
input_meta: Final = InputMeta(messages=_to_dd_messages(standard_logging_payload["messages"]))
output_meta: Final = OutputMeta(
messages=self._get_response_messages(
standard_logging_payload=standard_logging_payload,
@ -241,22 +426,20 @@ class DataDogLLMObsLogger(CustomBatchLogger):
if isinstance(metadata, dict):
metadata_parent_id = metadata.get("parent_id")
meta: Final = Meta(
kind=self._get_datadog_span_kind(standard_logging_payload.get("call_type"), metadata_parent_id),
input=input_meta,
output=output_meta,
metadata=self._get_dd_llm_obs_payload_metadata(standard_logging_payload),
error=error_info,
)
tool_definitions: Final = _to_dd_tool_definitions(standard_logging_payload.get("model_parameters"))
span_kind: Final = self._get_datadog_span_kind(standard_logging_payload.get("call_type"), metadata_parent_id)
payload_metadata: Final = self._get_dd_llm_obs_payload_metadata(standard_logging_payload)
# Calculate metrics (you may need to adjust these based on available data)
metrics: Final = LLMMetrics(
input_tokens=float(standard_logging_payload.get("prompt_tokens", 0)),
output_tokens=float(standard_logging_payload.get("completion_tokens", 0)),
total_tokens=float(standard_logging_payload.get("total_tokens", 0)),
total_cost=float(standard_logging_payload.get("response_cost", 0)),
time_to_first_token=self._get_time_to_first_token_seconds(standard_logging_payload),
)
meta: Final[Meta] = {
"kind": span_kind,
"input": input_meta,
"output": output_meta,
"metadata": payload_metadata,
"error": error_info,
**({"tool_definitions": tool_definitions} if tool_definitions else {}),
}
metrics: Final = self._assemble_metrics(standard_logging_payload)
payload: Final[LLMObsPayload] = LLMObsPayload(
parent_id=metadata_parent_id if metadata_parent_id else "undefined",
@ -314,6 +497,45 @@ class DataDogLLMObsLogger(CustomBatchLogger):
)
return error_info
def _assemble_metrics(self, standard_logging_payload: StandardLoggingPayload) -> LLMMetrics:
"""
Build the span metrics, including the prompt-cache counts LLM Obs charts cache savings from.
Cache counts resolve through the same owners the savings dashboard uses, so every provider
spelling is covered, and `non_cached_input_tokens` subtracts BOTH cache categories because
litellm's normalized prompt count includes both (the invariant the cost calculator's custom
pricing helper documents). A zero residual on a fully cached request is real data and is
emitted; a zero read or write count is absence and is not.
"""
prompt_tokens: Final = float(standard_logging_payload.get("prompt_tokens", 0))
completion_tokens: Final = float(standard_logging_payload.get("completion_tokens", 0))
total_tokens: Final = float(standard_logging_payload.get("total_tokens", 0))
total_cost: Final = float(standard_logging_payload.get("response_cost", 0))
time_to_first_token: Final = self._get_time_to_first_token_seconds(standard_logging_payload)
raw_usage: Final = (standard_logging_payload.get("metadata") or {}).get("usage_object")
usage_object: Final = raw_usage if isinstance(raw_usage, dict) else None
cache_read: Final = float(extract_cache_read_tokens(usage_object))
cache_write: Final = float(extract_cache_creation_tokens(usage_object))
metrics: Final[LLMMetrics] = {
"input_tokens": prompt_tokens,
"output_tokens": completion_tokens,
"total_tokens": total_tokens,
"total_cost": total_cost,
"time_to_first_token": time_to_first_token,
**(
{
**({"cache_read_input_tokens": cache_read} if cache_read else {}),
**({"cache_write_input_tokens": cache_write} if cache_write else {}),
"non_cached_input_tokens": max(prompt_tokens - cache_read - cache_write, 0.0),
}
if cache_read or cache_write
else {}
),
}
return metrics
def _get_time_to_first_token_seconds(self, standard_logging_payload: StandardLoggingPayload) -> float:
"""
Get the time to first token in seconds
@ -335,7 +557,7 @@ class DataDogLLMObsLogger(CustomBatchLogger):
def _get_response_messages(
self, standard_logging_payload: StandardLoggingPayload, call_type: str | None
) -> list[object]:
) -> tuple[Message, ...]:
"""
Get the messages from the response object
@ -344,7 +566,7 @@ class DataDogLLMObsLogger(CustomBatchLogger):
response_obj = standard_logging_payload.get("response")
if response_obj is None:
return []
return ()
# edge case: handle response_obj is a string representation of a dict
if isinstance(response_obj, str):
@ -357,7 +579,7 @@ class DataDogLLMObsLogger(CustomBatchLogger):
# fallback to json parsing
response_obj = json.loads(str(response_obj))
except json.JSONDecodeError:
return []
return ()
if call_type in [
CallTypes.completion.value,
@ -375,12 +597,12 @@ class DataDogLLMObsLogger(CustomBatchLogger):
if isinstance(response_obj, dict) and "choices" in response_obj:
choices: Final = response_obj["choices"]
if choices and len(choices) > 0 and "message" in choices[0]:
return [choices[0]["message"]]
return []
return _to_dd_messages([choices[0]["message"]])
return ()
except (KeyError, IndexError, TypeError):
# In case of any error accessing the response structure, return empty list
return []
return []
return ()
return ()
def _get_datadog_span_kind(
self, call_type: str | None, parent_id: str | None = None
@ -485,17 +707,6 @@ class DataDogLLMObsLogger(CustomBatchLogger):
# Default fallback for unknown or passthrough operations
return "llm"
def _ensure_string_content(self, messages: str | Sequence[object] | Mapping[object, object] | None) -> list[object]:
if messages is None:
return []
if isinstance(messages, str):
return [messages]
elif isinstance(messages, list):
return [message for message in messages]
elif isinstance(messages, dict):
return [str(messages.get("content", ""))]
return []
def _get_dd_llm_obs_payload_metadata(self, standard_logging_payload: StandardLoggingPayload) -> dict[str, object]:
"""
Fields to track in DD LLM Observability metadata from litellm standard logging payload
@ -524,10 +735,6 @@ class DataDogLLMObsLogger(CustomBatchLogger):
spend_metrics: Final = self._get_spend_metrics(standard_logging_payload)
_metadata.update({"spend_metrics": dict(spend_metrics)})
## extract tool calls and add to metadata
tool_call_metadata: Final = self._extract_tool_call_metadata(standard_logging_payload)
_metadata.update(tool_call_metadata)
_standard_logging_metadata: Final[dict] = dict(standard_logging_payload.get("metadata", {})) or {}
_metadata.update(_standard_logging_metadata)
return _metadata
@ -647,107 +854,3 @@ class DataDogLLMObsLogger(CustomBatchLogger):
verbose_logger.debug("Original value: %s", user_api_key_budget_reset_at)
return spend_metrics
def _process_input_messages_preserving_tool_calls(self, messages: Sequence[object]) -> list[dict[str, object]]:
"""
Process input messages while preserving tool_calls and tool message types.
This bypasses the lossy string conversion when tool calls are present,
allowing complex nested tool_calls objects to be preserved for Datadog.
"""
processed: Final = []
for msg in messages:
if isinstance(msg, dict):
# Preserve messages with tool_calls or tool role as-is
if "tool_calls" in msg or msg.get("role") == "tool":
processed.append(msg)
else:
# For regular messages, still apply string conversion
converted = handle_any_messages_to_chat_completion_str_messages_conversion([msg])
processed.extend(converted)
else:
# For non-dict messages, apply string conversion
converted = handle_any_messages_to_chat_completion_str_messages_conversion([msg])
processed.extend(converted)
return processed
@staticmethod
def _tool_calls_kv_pair(tool_calls: list[dict[str, Any]]) -> dict[str, object]:
"""
Extract tool call information into key-value pairs for Datadog metadata.
Similar to OpenTelemetry's implementation but adapted for Datadog's format.
"""
kv_pairs: Final[dict[str, object]] = {}
for idx, tool_call in enumerate(tool_calls):
try:
# Extract tool call ID
tool_id = tool_call.get("id")
if tool_id:
kv_pairs[f"tool_calls.{idx}.id"] = tool_id
# Extract tool call type
tool_type = tool_call.get("type")
if tool_type:
kv_pairs[f"tool_calls.{idx}.type"] = tool_type
# Extract function information
function = tool_call.get("function")
if function:
function_name = function.get("name")
if function_name:
kv_pairs[f"tool_calls.{idx}.function.name"] = function_name
function_arguments = function.get("arguments")
if function_arguments:
# Store arguments as JSON string for Datadog
if isinstance(function_arguments, str):
kv_pairs[f"tool_calls.{idx}.function.arguments"] = function_arguments
else:
import json
kv_pairs[f"tool_calls.{idx}.function.arguments"] = json.dumps(function_arguments)
except (KeyError, TypeError, ValueError) as e:
verbose_logger.debug("DataDogLLMObs: Error processing tool call %s: %s", idx, e)
continue
return kv_pairs
def _extract_tool_call_metadata(self, standard_logging_payload: StandardLoggingPayload) -> dict[str, object]:
"""
Extract tool call information from both input messages and response for Datadog metadata.
"""
tool_call_metadata: Final[dict[str, object]] = {}
try:
# Extract tool calls from input messages
messages: Final = standard_logging_payload.get("messages", [])
if messages and isinstance(messages, list):
for message in messages:
if isinstance(message, dict) and "tool_calls" in message:
tool_calls = message.get("tool_calls")
if tool_calls:
input_tool_calls_kv = self._tool_calls_kv_pair(tool_calls)
# Prefix with "input_" to distinguish from response tool calls
for key, value in input_tool_calls_kv.items():
tool_call_metadata[f"input_{key}"] = value
# Extract tool calls from response
response_obj: Final = standard_logging_payload.get("response")
if response_obj and isinstance(response_obj, dict):
choices: Final = response_obj.get("choices", [])
for choice in choices:
if isinstance(choice, dict):
message = choice.get("message")
if message and isinstance(message, dict):
tool_calls = message.get("tool_calls")
if tool_calls:
response_tool_calls_kv = self._tool_calls_kv_pair(tool_calls)
# Prefix with "output_" to distinguish from input tool calls
for key, value in response_tool_calls_kv.items():
tool_call_metadata[f"output_{key}"] = value
except Exception as e:
verbose_logger.debug("DataDogLLMObs: Error extracting tool call metadata: %s", e)
return tool_call_metadata

43
litellm/integrations/otel/model/payloads.py
View file

@ -6,6 +6,7 @@ import json
from collections.abc import Mapping
from dataclasses import dataclass, field
from enum import Enum
from types import MappingProxyType
from typing import TYPE_CHECKING, ClassVar, Final, cast
from urllib.parse import urlsplit
@ -62,6 +63,31 @@ if TYPE_CHECKING:
# --- typed sub-structures ---------------------------------------------------- #
def _cache_token_value(*values: object) -> int | None:
explicit_zero = False
invalid_before_zero = False
for raw_value in values:
if raw_value is None:
continue
if isinstance(raw_value, bool):
parsed = None
else:
try:
parsed = as_int(raw_value)
except (OverflowError, ValueError):
parsed = None
if parsed is None:
if not explicit_zero:
invalid_before_zero = True
elif parsed > 0:
return parsed
elif parsed == 0:
explicit_zero = True
elif not explicit_zero:
invalid_before_zero = True
return 0 if explicit_zero and not invalid_before_zero else None
@dataclass(frozen=True)
class LLMRequestParams:
temperature: float | None = None
@ -104,12 +130,25 @@ class LLMUsage:
metadata: Final[Mapping[str, object]] = payload.get("metadata") or {}
raw_usage: Final = metadata.get("usage_object")
usage_object: Final[Mapping[str, object]] = raw_usage if isinstance(raw_usage, Mapping) else {}
raw_details: Final = usage_object.get("prompt_tokens_details")
prompt_details: Final[Mapping[str, object]] = (
raw_details if isinstance(raw_details, Mapping) else MappingProxyType({})
)
return cls(
input_tokens=as_int(payload.get("prompt_tokens")),
output_tokens=as_int(payload.get("completion_tokens")),
total_tokens=as_int(payload.get("total_tokens")),
cache_creation_input_tokens=as_int(usage_object.get("cache_creation_input_tokens")),
cache_read_input_tokens=as_int(usage_object.get("cache_read_input_tokens")),
cache_creation_input_tokens=_cache_token_value(
usage_object.get("cache_creation_input_tokens"),
prompt_details.get("cache_write_tokens"),
prompt_details.get("cache_creation_tokens"),
prompt_details.get("cache_creation_input_tokens"),
),
cache_read_input_tokens=_cache_token_value(
usage_object.get("cache_read_input_tokens"),
prompt_details.get("cached_tokens"),
usage_object.get("prompt_cache_hit_tokens"),
),
)

120
litellm/integrations/prometheus.py
View file

@ -8,6 +8,7 @@ import math
import os
import sys
from collections.abc import Awaitable, Callable, Mapping, Sequence
from dataclasses import replace
from datetime import datetime, timedelta
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, cast
@ -58,6 +59,7 @@ from litellm.types.utils import (
if TYPE_CHECKING:
from apscheduler.schedulers.asyncio import AsyncIOScheduler
from prometheus_client import Gauge
from prometheus_client.metrics import MetricWrapperBase
from litellm.router import Router
@ -476,6 +478,30 @@ class PrometheusLogger(CustomLogger):
labelnames=self.get_labels_for_metric("litellm_remaining_api_key_tokens_for_model"),
)
self.litellm_api_key_rate_limit_allowed_metric = self._gauge_factory(
"litellm_api_key_rate_limit_allowed_metric",
"Configured rate limit for the API Key in the current window (rpm_limit / tpm_limit), by rate_limit_type",
labelnames=self.get_labels_for_metric("litellm_api_key_rate_limit_allowed_metric"),
)
self.litellm_api_key_rate_limit_used_metric = self._gauge_factory(
"litellm_api_key_rate_limit_used_metric",
"Requests or tokens the API Key has consumed in the current rate limit window, by rate_limit_type",
labelnames=self.get_labels_for_metric("litellm_api_key_rate_limit_used_metric"),
)
self.litellm_team_rate_limit_allowed_metric = self._gauge_factory(
"litellm_team_rate_limit_allowed_metric",
"Configured rate limit for the Team in the current window (team rpm_limit / tpm_limit), by rate_limit_type",
labelnames=self.get_labels_for_metric("litellm_team_rate_limit_allowed_metric"),
)
self.litellm_team_rate_limit_used_metric = self._gauge_factory(
"litellm_team_rate_limit_used_metric",
"Requests or tokens the Team has consumed in the current rate limit window, by rate_limit_type",
labelnames=self.get_labels_for_metric("litellm_team_rate_limit_used_metric"),
)
########################################
# LLM API Deployment Metrics / analytics
########################################
@ -1475,6 +1501,11 @@ class PrometheusLogger(CustomLogger):
model_id=enum_values.model_id,
)
self._set_key_and_team_rate_limit_metrics(
standard_logging_payload=standard_logging_payload, # pyright: ignore[reportArgumentType] # isinstance(dict) above narrows the TypedDict to dict[Unknown, Unknown]
enum_values=enum_values,
)
# set latency metrics
self._set_latency_metrics(
kwargs=kwargs,
@ -2002,17 +2033,102 @@ class PrometheusLogger(CustomLogger):
"""
if standard_logging_payload is None:
return None
return PrometheusLogger._get_int_from_v3_rate_limit_headers(
standard_logging_payload=standard_logging_payload,
header_name=f"x-ratelimit-model_per_key-remaining-{rate_limit_type}",
)
@staticmethod
def _get_int_from_v3_rate_limit_headers(
standard_logging_payload: StandardLoggingPayload,
header_name: str,
) -> int | None:
hidden_params: Final = standard_logging_payload.get("hidden_params")
if hidden_params is None:
return None
additional_headers: Final = hidden_params.get("additional_headers")
additional_headers: Final[Mapping[str, object] | None] = hidden_params.get("additional_headers")
if additional_headers is None:
return None
value: Final = dict(additional_headers).get(f"x-ratelimit-model_per_key-remaining-{rate_limit_type}")
value: Final = additional_headers.get(header_name)
if isinstance(value, bool) or not isinstance(value, int):
return None
return value
def _set_key_and_team_rate_limit_metrics(
self,
standard_logging_payload: StandardLoggingPayload,
enum_values: UserAPIKeyLabelValues,
) -> None:
"""
Export the key-level and team-level RPM / TPM limit and current window
usage from the ``x-ratelimit-{api_key,team}-{limit,remaining}-*``
headers the v3 rate limiter mirrors into the logging payload. The
limiter already read these counters (from Redis when configured) on
the request path, so no extra store lookup happens here. Descriptors
without a configured limit emit no header, so their series is removed
rather than left at the value from before the limit was dropped.
"""
descriptor_gauges: Final[
tuple[tuple[Literal["api_key", "team"], DEFINED_PROMETHEUS_METRICS, Gauge, Gauge], ...]
] = (
(
"api_key",
"litellm_api_key_rate_limit_allowed_metric",
self.litellm_api_key_rate_limit_allowed_metric,
self.litellm_api_key_rate_limit_used_metric,
),
(
"team",
"litellm_team_rate_limit_allowed_metric",
self.litellm_team_rate_limit_allowed_metric,
self.litellm_team_rate_limit_used_metric,
),
)
for descriptor_key, metric_name, allowed_gauge, used_gauge in descriptor_gauges:
for rate_limit_type in ("requests", "tokens"):
self._set_rate_limit_allowed_and_used_gauges(
standard_logging_payload=standard_logging_payload,
enum_values=enum_values,
descriptor_key=descriptor_key,
metric_name=metric_name,
allowed_gauge=allowed_gauge,
used_gauge=used_gauge,
rate_limit_type=rate_limit_type,
)
def _set_rate_limit_allowed_and_used_gauges(
self,
standard_logging_payload: StandardLoggingPayload,
enum_values: UserAPIKeyLabelValues,
descriptor_key: Literal["api_key", "team"],
metric_name: DEFINED_PROMETHEUS_METRICS,
allowed_gauge: Gauge,
used_gauge: Gauge,
rate_limit_type: Literal["requests", "tokens"],
) -> None:
limit: Final = self._get_int_from_v3_rate_limit_headers(
standard_logging_payload=standard_logging_payload,
header_name=f"x-ratelimit-{descriptor_key}-limit-{rate_limit_type}",
)
remaining: Final = self._get_int_from_v3_rate_limit_headers(
standard_logging_payload=standard_logging_payload,
header_name=f"x-ratelimit-{descriptor_key}-remaining-{rate_limit_type}",
)
labelled_values: Final = replace(enum_values, rate_limit_type=rate_limit_type)
labelnames: Final = self.get_labels_for_metric(metric_name)
labels: Final = prometheus_label_factory(
supported_enum_labels=labelnames,
enum_values=labelled_values,
label_context=PrometheusLabelFactoryContext(labelled_values),
)
if limit is None or remaining is None:
label_values: Final = tuple(labels.get(label) for label in labelnames)
self._bounded_prometheus_series_tracker.remove_series(allowed_gauge, label_values)
self._bounded_prometheus_series_tracker.remove_series(used_gauge, label_values)
return
allowed_gauge.labels(**labels).set(limit)
used_gauge.labels(**labels).set(limit - remaining)
def _set_virtual_key_rate_limit_metrics(
self,
user_api_key: str | None,

4
litellm/integrations/prometheus_helpers/bounded_prometheus_series_tracker.py
View file

@ -60,6 +60,10 @@ class BoundedPrometheusSeriesTracker:
break
del series[tracked_label_values]
def remove_series(self, metric: object, label_values: tuple[str | None, ...]) -> bool:
"""Drop one child series, True when it is gone (removed or never existed)."""
return self._remove_metric_child(metric, label_values)
def _should_run_ttl_cleanup(
self,
metric_name: str,

126
litellm/litellm_core_utils/get_model_cost_map.py
View file

@ -12,6 +12,7 @@ import asyncio
import json
import os
import random
import time
from collections.abc import Awaitable, Callable
from dataclasses import dataclass
from datetime import datetime, timezone
@ -154,18 +155,6 @@ class GetModelCostMap:
return True
@staticmethod
def fetch_remote_model_cost_map(url: str, timeout: int = 5) -> dict:
"""
Fetch the model cost map from a remote URL.
Returns the parsed JSON dict. Raises on network/parse errors
(caller is expected to handle).
"""
response: Final = httpx.get(url, timeout=timeout)
response.raise_for_status()
return response.json()
RETRYABLE_FETCH_STATUS_CODES: Final = frozenset({429, 500, 502, 503, 504})
MODEL_COST_MAP_FETCH_MAX_ATTEMPTS: Final = 3
@ -212,6 +201,13 @@ class _AsyncGetClient(Protocol):
def get(self, url: str, *, timeout: float | None = None) -> Awaitable[httpx.Response]: ...
class _SyncGetClient(Protocol):
def get(self, url: str, *, timeout: float | None = None) -> httpx.Response: ...
_FetchAttemptOutcome = ModelCostMapReloaded | ModelCostMapReloadUnavailable | _FetchAttemptRetryable
def _default_reload_client() -> _AsyncGetClient:
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.types.llms.custom_http import httpxSpecialProvider
@ -219,13 +215,30 @@ def _default_reload_client() -> _AsyncGetClient:
return get_async_httpx_client(llm_provider=httpxSpecialProvider.ModelCostMap)
async def _attempt_fetch(
client: _AsyncGetClient, url: str, timeout: int
) -> ModelCostMapReloaded | ModelCostMapReloadUnavailable | _FetchAttemptRetryable:
def _classify_fetch_error(error: httpx.HTTPError | httpx.InvalidURL, url: str) -> _FetchAttemptOutcome:
reason: Final = f"{type(error).__name__} fetching {url}: {error}"
if isinstance(error, (httpx.InvalidURL, httpx.UnsupportedProtocol)):
return ModelCostMapReloadUnavailable(reason=reason)
return _FetchAttemptRetryable(reason=reason, retry_after_seconds=None)
async def _attempt_fetch(client: _AsyncGetClient, url: str, timeout: int) -> _FetchAttemptOutcome:
try:
response: Final = await client.get(url, timeout=timeout)
except httpx.HTTPError as e:
return _FetchAttemptRetryable(reason=f"{type(e).__name__} fetching {url}: {e}", retry_after_seconds=None)
except (httpx.HTTPError, httpx.InvalidURL) as e:
return _classify_fetch_error(e, url)
return _classify_fetch_response(response, url)
def _attempt_fetch_sync(client: _SyncGetClient, url: str, timeout: int) -> _FetchAttemptOutcome:
try:
response: Final = client.get(url, timeout=timeout)
except (httpx.HTTPError, httpx.InvalidURL) as e:
return _classify_fetch_error(e, url)
return _classify_fetch_response(response, url)
def _classify_fetch_response(response: httpx.Response, url: str) -> _FetchAttemptOutcome:
if response.status_code in RETRYABLE_FETCH_STATUS_CODES:
return _FetchAttemptRetryable(
reason=f"HTTP {response.status_code} from {url}",
@ -242,6 +255,22 @@ async def _attempt_fetch(
return ModelCostMapReloaded(model_cost_map=parsed)
def _next_retry_wait(
outcome: _FetchAttemptRetryable, attempt: int, max_attempts: int, rng: random.Random
) -> float | ModelCostMapReloadUnavailable:
if attempt == max_attempts:
return ModelCostMapReloadUnavailable(reason=f"{outcome.reason} (after {max_attempts} attempts)")
wait_seconds: Final = _retry_wait_seconds(outcome=outcome, attempt=attempt, rng=rng)
verbose_logger.warning(
"LiteLLM: model cost map fetch attempt %d/%d failed (%s); retrying in %.1fs",
attempt,
max_attempts,
outcome.reason,
wait_seconds,
)
return wait_seconds
async def _fetch_remote_model_cost_map_with_retry(
url: str,
timeout: int,
@ -254,20 +283,32 @@ async def _fetch_remote_model_cost_map_with_retry(
outcome = await _attempt_fetch(client=client, url=url, timeout=timeout)
if not isinstance(outcome, _FetchAttemptRetryable):
return outcome
if attempt == max_attempts:
return ModelCostMapReloadUnavailable(reason=f"{outcome.reason} (after {max_attempts} attempts)")
wait_seconds = _retry_wait_seconds(outcome=outcome, attempt=attempt, rng=rng)
verbose_logger.warning(
"LiteLLM: model cost map fetch attempt %d/%d failed (%s); retrying in %.1fs",
attempt,
max_attempts,
outcome.reason,
wait_seconds,
)
wait_seconds = _next_retry_wait(outcome=outcome, attempt=attempt, max_attempts=max_attempts, rng=rng)
if isinstance(wait_seconds, ModelCostMapReloadUnavailable):
return wait_seconds
await sleep(wait_seconds)
return ModelCostMapReloadUnavailable(reason="model cost map fetch failed")
def _fetch_remote_model_cost_map_with_retry_sync(
url: str,
timeout: int,
max_attempts: int,
sleep: Callable[[float], None],
rng: random.Random,
client: _SyncGetClient,
) -> ModelCostMapReloadResult:
for attempt in range(1, max_attempts + 1):
outcome = _attempt_fetch_sync(client=client, url=url, timeout=timeout)
if not isinstance(outcome, _FetchAttemptRetryable):
return outcome
wait_seconds = _next_retry_wait(outcome=outcome, attempt=attempt, max_attempts=max_attempts, rng=rng)
if isinstance(wait_seconds, ModelCostMapReloadUnavailable):
return wait_seconds
sleep(wait_seconds)
return ModelCostMapReloadUnavailable(reason="model cost map fetch failed")
async def refetch_model_cost_map(
url: str,
timeout: int = 5,
@ -423,13 +464,21 @@ def _finalize_model_cost_map(model_cost: dict) -> dict:
return _expand_model_aliases(model_cost)
def get_model_cost_map(url: str) -> dict:
def get_model_cost_map(
url: str,
timeout: int = 5,
max_attempts: int = MODEL_COST_MAP_FETCH_MAX_ATTEMPTS,
sleep: Callable[[float], None] = time.sleep,
rng: random.Random | None = None,
client: "_SyncGetClient | None" = None,
) -> dict:
"""
Public entry point — returns the model cost map dict.
1. If ``LITELLM_LOCAL_MODEL_COST_MAP`` is set, uses the local backup only.
2. Otherwise fetches from ``url``, validates integrity, and falls back
to the local backup on any failure.
2. Otherwise fetches from ``url``, retrying transient HTTP errors
(429/5xx/transport) with Retry-After-aware backoff, validates
integrity, and falls back to the local backup on any failure.
Only the backup model count is cached (a single int) for validation.
The full backup dict is only parsed when it must be *returned* as a
@ -448,17 +497,24 @@ def get_model_cost_map(url: str) -> dict:
_cost_map_source_info.url = url
_cost_map_source_info.is_env_forced = False
try:
content: Final = GetModelCostMap.fetch_remote_model_cost_map(url)
except Exception as e:
result: Final = _fetch_remote_model_cost_map_with_retry_sync(
url=url,
timeout=timeout,
max_attempts=max_attempts,
sleep=sleep,
rng=rng if rng is not None else random.Random(),
client=client if client is not None else httpx,
)
if isinstance(result, ModelCostMapReloadUnavailable):
verbose_logger.warning(
"LiteLLM: Failed to fetch remote model cost map from %s: %s. Falling back to local backup.",
url,
str(e),
result.reason,
)
_cost_map_source_info.source = "local"
_cost_map_source_info.fallback_reason = f"Remote fetch failed: {e}"
_cost_map_source_info.fallback_reason = f"Remote fetch failed: {result.reason}"
return _finalize_model_cost_map(GetModelCostMap.load_local_model_cost_map())
content: Final = result.model_cost_map
# Validate using cached count (cheap int comparison, no file I/O)
if not GetModelCostMap.validate_model_cost_map(

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

@ -4957,10 +4957,13 @@ def make_valid_bedrock_tool_name(input_tool_name: str) -> str:
def add_cache_point_tool_block(tool: dict, model: str | None = None) -> BedrockToolBlock | None:
from litellm.llms.bedrock.common_utils import is_claude_4_5_on_bedrock
from litellm.llms.bedrock.common_utils import (
bedrock_model_accepts_cache_points,
is_claude_4_5_on_bedrock,
)
cache_control: Final = tool.get("cache_control", None)
if cache_control is not None:
if cache_control is not None and bedrock_model_accepts_cache_points(model):
cache_point: Final = cache_control.get("type", "ephemeral")
if cache_point == "ephemeral":
cache_point_block: Final[CachePointBlock] = {"type": "default"}

8
litellm/litellm_core_utils/streaming_chunk_builder_utils.py
View file

@ -36,6 +36,8 @@ from litellm.types.utils import (
from litellm.utils import print_verbose, token_counter
if TYPE_CHECKING:
from openai.types.completion_usage import CompletionUsage
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.types.litellm_core_utils.streaming_chunk_builder_utils import (
UsagePerChunk,
@ -794,7 +796,7 @@ class ChunkProcessor:
@staticmethod
def _extract_usage_chunk(chunk: "_UsageBearingChunk | ModelResponse | ModelResponseStream") -> Usage | None:
usage_chunk: Usage | None = None
usage_chunk: Usage | CompletionUsage | None = None
if hasattr(chunk, "usage") and chunk.usage is not None:
usage_chunk = chunk.usage
elif "usage" in chunk:
@ -806,7 +808,9 @@ class ChunkProcessor:
if isinstance(usage_chunk, dict):
return Usage(**usage_chunk)
return usage_chunk
if usage_chunk is None or isinstance(usage_chunk, Usage):
return usage_chunk
return Usage(**usage_chunk.model_dump())
def _calculate_usage_per_chunk(
self,

17
litellm/llms/anthropic/common_utils.py
View file

@ -1378,31 +1378,38 @@ def process_anthropic_headers(headers: httpx.Headers | dict) -> dict:
return additional_headers
def _anthropic_model_entry(model: ModelInfoResponse, created_at: str) -> Mapping[str, object]:
def _anthropic_model_entry(
model: ModelInfoResponse, created_at: str, display_names: Mapping[str, str]
) -> Mapping[str, object]:
return { # mutable-ok: JSON response body, serialized by the route and never mutated
"type": "model",
"id": model["id"],
"display_name": model["id"],
"display_name": display_names.get(model["id"], model["id"]),
"created_at": created_at,
"max_input_tokens": model.get("max_input_tokens"),
"max_tokens": model.get("max_output_tokens"),
}
def create_anthropic_model_list_response(models: Sequence[ModelInfoResponse]) -> Mapping[str, object]:
def create_anthropic_model_list_response(
models: Sequence[ModelInfoResponse],
display_names: Mapping[str, str] = MappingProxyType({}),
) -> Mapping[str, object]:
"""Build the Anthropic-native /v1/models envelope.
Clients that send an anthropic-version header parse the Anthropic Models API
shape (type/display_name/created_at plus has_more/first_id/last_id) and filter
the list themselves, so every model is returned here. The token limits carry
over from the OpenAI-shaped listing, named as the Messages API names them, and
are always present because the vendor shape declares them nullable, not optional
are always present because the vendor shape declares them nullable, not optional.
display_names maps a listed model id to a configured human-readable name; ids
without an entry fall back to the id itself, matching the vendor behavior
"""
created_at: Final = (
datetime.fromtimestamp(DEFAULT_MODEL_CREATED_AT_TIME, tz=timezone.utc).isoformat().replace("+00:00", "Z")
)
data: Final = [ # mutable-ok: JSON response body, serialized by the route and never mutated
_anthropic_model_entry(model, created_at) for model in models
_anthropic_model_entry(model, created_at, display_names) for model in models
]
return { # mutable-ok: JSON response body, serialized by the route and never mutated
"data": data,

4
litellm/llms/base_llm/guardrail_translation/base_translation.py
View file

@ -155,8 +155,8 @@ class BaseTranslation(ABC):
self,
exc: "ModifyResponseException",
stream_started: bool = False,
responses_so_far: list[Any] | None = None,
) -> list[bytes] | None:
responses_so_far: Sequence[Any] | None = None,
) -> Sequence[bytes] | None:
"""
Build the streaming chunks that deliver a guardrail block message and
cleanly terminate the stream in this provider's wire format.

55
litellm/llms/base_llm/guardrail_translation/utils.py
View file

@ -124,6 +124,61 @@ def blocked_responses_api_usage(original_response: object) -> ResponseAPIUsage:
)
def stream_item_field(item: object, field: str) -> object | None:
if isinstance(item, dict):
return item.get(field)
return getattr(item, field, None)
def blocked_chat_stream_usage(original_response: object) -> tuple[int, int]:
"""
``(prompt_tokens, completion_tokens)`` for a synthetic guardrail-blocked
chat completions stream.
A mid-stream block carries the chunks received so far as a list; real usage
rides on the final chunk when the upstream sent one
(``stream_options.include_usage``). Non-list originals defer to
``blocked_response_usage``.
"""
if not isinstance(original_response, list):
usage: Final = blocked_response_usage(original_response)
return usage.get("input_tokens", 0), usage.get("output_tokens", 0)
usage_obj: Final = next(
(
chunk_usage
for item in reversed(original_response)
if (chunk_usage := stream_item_field(item, "usage")) is not None
),
None,
)
return (
_usage_tokens(usage_obj, "prompt_tokens", "input_tokens"),
_usage_tokens(usage_obj, "completion_tokens", "output_tokens"),
)
def blocked_responses_stream_usage(original_response: object) -> ResponseAPIUsage:
"""
``ResponseAPIUsage`` for a synthetic guardrail-blocked /v1/responses stream.
A mid-stream block carries the events received so far as a list; real usage
rides on the ``response.completed`` event's response when the upstream sent
one. Non-list originals defer to ``blocked_responses_api_usage``.
"""
if not isinstance(original_response, list):
return blocked_responses_api_usage(original_response)
completed: Final = next(
(
response
for item in reversed(original_response)
if stream_item_field(item, "type") == "response.completed"
and (response := stream_item_field(item, "response")) is not None
),
None,
)
return blocked_responses_api_usage(completed)
def effective_skip_system_message_for_guardrail(guardrail_to_apply: Any) -> bool:
per: Final = getattr(guardrail_to_apply, "skip_system_message_in_guardrail", None)
if per is not None:

6
litellm/llms/bedrock/base_aws_llm.py
View file

@ -1442,7 +1442,7 @@ class BaseAWSLLM:
@tracer.wrap()
def get_request_headers(
self,
credentials: Credentials,
credentials: Credentials | None,
aws_region_name: str,
extra_headers: dict | None,
endpoint_url: str,
@ -1469,9 +1469,13 @@ class BaseAWSLLM:
try:
from botocore.auth import SigV4Auth
from botocore.awsrequest import AWSRequest
from botocore.exceptions import NoCredentialsError
except ImportError:
raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
if credentials is None:
raise NoCredentialsError()
# Filter headers for AWS signature calculation
# AWS SigV4 only includes specific headers in signature calculation
aws_signature_headers: Final = self._filter_headers_for_aws_signature(headers)

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

@ -1,4 +1,6 @@
import json
from collections.abc import Mapping
from types import MappingProxyType
from typing import Any, Final
import httpx
@ -24,6 +26,22 @@ from ..common_utils import BedrockError, _get_all_bedrock_regions
from .invoke_handler import AWSEventStreamDecoder, MockResponseIterator, make_call
def _sigv4_principal(credentials: Credentials | None) -> Mapping[str, str]:
if credentials is None:
return MappingProxyType({})
return MappingProxyType(
{
key: value
for key, value in (
("aws_access_key_id", credentials.access_key),
("aws_secret_access_key", credentials.secret_key),
("aws_session_token", credentials.token),
)
if value is not None
}
)
def make_sync_call(
client: HTTPHandler | None,
api_base: str,
@ -95,7 +113,7 @@ class BedrockConverseLLM(BaseAWSLLM):
stream,
optional_params: dict,
litellm_params: dict,
credentials: Credentials,
credentials: Credentials | None,
logger_fn=None,
headers={},
client: AsyncHTTPHandler | None = None,
@ -167,7 +185,7 @@ class BedrockConverseLLM(BaseAWSLLM):
stream,
optional_params: dict,
litellm_params: dict,
credentials: Credentials,
credentials: Credentials | None,
logger_fn=None,
headers: dict = {},
client: AsyncHTTPHandler | None = None,
@ -331,7 +349,7 @@ class BedrockConverseLLM(BaseAWSLLM):
litellm_params["aws_region_name"] = aws_region_name # [DO NOT DELETE] important for async calls
credentials: Final[Credentials] = self.get_credentials(
credentials: Final[Credentials | None] = self.get_credentials(
aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
aws_session_token=aws_session_token,
@ -368,19 +386,13 @@ class BedrockConverseLLM(BaseAWSLLM):
# The Rust core owns the whole call for the subset it accepts. Ask
# before transforming so whichever path runs emits pre_call once, and
# hand down the credentials, region and endpoint this handler already
# resolved so both paths sign as the same principal.
# resolved so both paths sign as the same principal. Bearer-token auth
# resolves no SigV4 principal at all, and each path reads that token
# itself.
rust_optional_params: Final = { # mutable-ok: json.dumps in the bridge rejects a mappingproxy
**optional_params,
**{ # mutable-ok: merged into its mutable parent above
key: value
for key, value in (
("aws_access_key_id", credentials.access_key),
("aws_secret_access_key", credentials.secret_key),
("aws_session_token", credentials.token),
("aws_region_name", aws_region_name),
)
if value is not None
},
**_sigv4_principal(credentials),
"aws_region_name": aws_region_name,
}
serves_via_rust: Final = rust_chat_completions_accepts(
model=model,

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

@ -87,6 +87,7 @@ from ..common_utils import (
BedrockError,
BedrockModelInfo,
bedrock_converse_supports_parallel_tool_use_config,
bedrock_model_accepts_cache_points,
get_anthropic_beta_from_headers,
get_bedrock_tool_name,
is_bedrock_application_inference_profile_arn,
@ -1149,7 +1150,7 @@ class AmazonConverseConfig(BaseConfig):
model: str | None = None,
) -> SystemContentBlock | ContentBlock | None:
cache_control: Final = message_block.get("cache_control", None)
if cache_control is None:
if cache_control is None or not bedrock_model_accepts_cache_points(model):
return None
cache_point: Final = self._build_cache_point_block(cache_control, model)
@ -1613,7 +1614,7 @@ class AmazonConverseConfig(BaseConfig):
# Append cachePoint to tools if cache_control_injection_points has tool_config
cache_injection_points: Final = additional_request_params.pop("cache_control_injection_points", None)
if cache_injection_points and len(bedrock_tools) > 0:
if cache_injection_points and len(bedrock_tools) > 0 and bedrock_model_accepts_cache_points(model):
for point in cache_injection_points:
if point.get("location") == "tool_config":
cache_point = self._build_cache_point_block(point.get("control"), model)

24
litellm/llms/bedrock/common_utils.py
View file

@ -816,6 +816,30 @@ def bedrock_converse_supports_parallel_tool_use_config(model: str) -> bool:
)
def bedrock_model_accepts_cache_points(model: str | None) -> bool:
"""
Whether Converse ``cachePoint`` blocks may be sent to this model.
Bedrock rejects requests carrying cachePoint blocks for models without prompt
caching support ("You invoked an unsupported model or your request did not allow
prompt caching"), so a model whose cost-map entry does not declare
``supports_prompt_caching`` must not receive them. A model absent from the map
(an application inference profile ARN, a model newer than the map) keeps emitting
so existing caching setups never silently degrade. ``litellm.utils.supports_prompt_caching``
is not reusable here: it returns False for unmapped models, the opposite polarity.
"""
if model is None:
return True
entries: Final = tuple(
entry
for candidate in (model, get_bedrock_base_model(model))
if (entry := litellm.model_cost.get(candidate)) is not None
)
if not entries:
return True
return any(entry.get("supports_prompt_caching") is True for entry in entries)
def is_claude_4_5_on_bedrock(model: str) -> bool:
"""
Check if the model supports Bedrock prompt caching with an extended '1h' TTL

13
litellm/llms/custom_httpx/aiohttp_transport.py
View file

@ -3,6 +3,7 @@ import concurrent.futures
import contextlib
import os
import ssl
import sys
import typing
import urllib.request
from collections.abc import Callable, Generator
@ -75,10 +76,22 @@ except ImportError:
pass
def _current_task_is_cancelling() -> bool:
task: Final = asyncio.current_task()
if task is None or sys.version_info < (3, 11):
return True
return task.cancelling() > 0
@contextlib.contextmanager
def map_aiohttp_exceptions() -> Generator[None, None, None]:
try:
yield
except asyncio.CancelledError as exc:
# a closing connector cancels its shielded DNS task; that surfaces here without the request task being cancelled
if _current_task_is_cancelling():
raise
raise httpx.ConnectError("aiohttp transport cancelled the request internally") from exc
except Exception as exc:
mapped_exc: type[Exception] | None = None

138
litellm/llms/openai/chat/guardrail_translation/handler.py
View file

@ -14,9 +14,14 @@ Pattern Overview:
This pattern can be replicated for other message formats (e.g., Anthropic).
"""
import json
import time
import uuid
from collections.abc import Sequence
from typing import TYPE_CHECKING, Any, Final, Union, cast
from typing_extensions import NotRequired, ReadOnly, TypedDict
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.llms.base_llm.guardrail_translation.base_translation import (
@ -24,6 +29,7 @@ from litellm.llms.base_llm.guardrail_translation.base_translation import (
StreamTransformSink,
)
from litellm.llms.base_llm.guardrail_translation.utils import (
blocked_chat_stream_usage,
effective_scan_only_tool_results_for_guardrail,
effective_skip_system_message_for_guardrail,
effective_skip_tool_message_for_guardrail,
@ -32,6 +38,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import (
openai_tool_name,
role_out_of_guardrail_scope,
scoped_structured_message_indices,
stream_item_field,
)
from litellm.main import stream_chunk_builder
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
@ -49,7 +56,10 @@ from litellm.types.utils import (
if TYPE_CHECKING:
from fastapi import HTTPException
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.integrations.custom_guardrail import (
CustomGuardrail,
ModifyResponseException,
)
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy._types import UserAPIKeyAuth
@ -1005,3 +1015,129 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
else:
# Subsequent chunks - clear the text
content_item["text"] = ""
def _check_streaming_has_ended(self, responses_so_far: Sequence[object]) -> bool:
"""
True once any relayed chunk carries a non-null ``finish_reason``.
The unified guardrail's ``end_of_stream_only`` streaming path probes
this via ``hasattr`` to withhold the terminal chunks until
end-of-stream moderation runs, so a block can replace the finish
instead of trailing after a ``finish_reason`` the client already saw.
"""
return any(
stream_item_field(choice, "finish_reason") is not None
for item in responses_so_far
for choice in _stream_chunk_choices(item)
)
def build_block_sse_chunks(
self,
exc: "ModifyResponseException",
stream_started: bool = False,
responses_so_far: Sequence[object] | None = None,
) -> Sequence[bytes]:
"""
Build OpenAI chat-completions SSE chunks that deliver the guardrail
block message and terminate the stream cleanly, mirroring the
non-streaming block response: ``finish_reason`` ``content_filter`` plus
the real usage the upstream call consumed.
- ``stream_started`` False (buffered / pre-stream): nothing has been
sent, so open a standalone completion with a ``role`` delta.
- ``stream_started`` True (sampling / mid-stream): chunks already
reached the client, so continue the in-progress completion (reuse its
id/created/model, content-only delta).
The proxy's data generator appends ``data: [DONE]`` itself.
"""
chunk_id, created, model = _blocked_stream_identity(exc, responses_so_far or ())
prompt_tokens, completion_tokens = blocked_chat_stream_usage(exc.original_response)
continuation_delta: Final[_BlockedChunkDelta] = {"content": exc.message}
standalone_delta: Final[_BlockedChunkDelta] = {"role": "assistant", "content": exc.message}
message_chunk: Final[_BlockedChunk] = {
"id": chunk_id,
"object": "chat.completion.chunk",
"created": created,
"model": model,
"choices": (
{
"index": 0,
"delta": continuation_delta if stream_started else standalone_delta,
"finish_reason": None,
},
),
}
final_chunk: Final[_BlockedChunk] = {
"id": chunk_id,
"object": "chat.completion.chunk",
"created": created,
"model": model,
"choices": ({"index": 0, "delta": {}, "finish_reason": "content_filter"},),
"usage": {
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
"total_tokens": prompt_tokens + completion_tokens,
},
}
return _chat_sse_chunk(message_chunk), _chat_sse_chunk(final_chunk)
class _BlockedChunkDelta(TypedDict, total=False):
role: ReadOnly[str]
content: ReadOnly[str]
class _BlockedChunkChoice(TypedDict):
index: ReadOnly[int]
delta: ReadOnly[_BlockedChunkDelta]
finish_reason: ReadOnly[str | None]
class _BlockedChunkUsage(TypedDict):
prompt_tokens: ReadOnly[int]
completion_tokens: ReadOnly[int]
total_tokens: ReadOnly[int]
class _BlockedChunk(TypedDict):
id: ReadOnly[str]
object: ReadOnly[str]
created: ReadOnly[int]
model: ReadOnly[str]
choices: ReadOnly[tuple[_BlockedChunkChoice, ...]]
usage: NotRequired[ReadOnly[_BlockedChunkUsage]]
def _chat_sse_chunk(payload: _BlockedChunk) -> bytes:
return f"data: {json.dumps(payload)}\n\n".encode()
def _stream_chunk_choices(item: object) -> Sequence[object]:
choices: Final = stream_item_field(item, "choices")
if isinstance(choices, Sequence) and not isinstance(choices, (str, bytes)):
return choices
return ()
def _blocked_stream_identity(
exc: "ModifyResponseException", responses_so_far: Sequence[object]
) -> tuple[str, int, str]:
identified: Final = next(
(
(chunk_id, item)
for item in responses_so_far
if isinstance(chunk_id := stream_item_field(item, "id"), str) and chunk_id
),
None,
)
if identified is None:
return f"chatcmpl-{uuid.uuid4()}", int(time.time()), exc.model
chunk_id, source = identified
created: Final = stream_item_field(source, "created")
model: Final = stream_item_field(source, "model")
return (
chunk_id,
created if isinstance(created, int) else int(time.time()),
model if isinstance(model, str) and model else exc.model,
)

358
litellm/llms/openai/responses/guardrail_translation/handler.py
View file

@ -28,12 +28,16 @@ Output: response.output is List[GenericResponseOutputItem] where each has:
- text: str
"""
from collections.abc import Sequence
import time
import uuid
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Union, cast
from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall
from openai.types.responses.tool_param import FunctionToolParam
from pydantic import BaseModel
from pydantic import BaseModel, TypeAdapter
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_proxy_logger
@ -41,17 +45,33 @@ from litellm.completion_extras.litellm_responses_transformation.transformation i
OpenAiResponsesToChatCompletionStreamIterator,
)
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
from litellm.llms.base_llm.guardrail_translation.utils import (
blocked_responses_stream_usage,
stream_item_field,
)
from litellm.responses.litellm_completion_transformation.transformation import (
LiteLLMCompletionResponsesConfig,
)
from litellm.types.llms.openai import (
AllMessageValues,
BaseLiteLLMOpenAIResponseObject,
ChatCompletionToolCallChunk,
ChatCompletionToolParam,
ContentPartAddedEvent,
ContentPartDoneEvent,
ContentPartDonePartOutputText,
ErrorEvent,
ErrorEventError,
OpenAIMcpServerTool,
OutputItemAddedEvent,
OutputItemDoneEvent,
OutputTextDeltaEvent,
OutputTextDoneEvent,
ResponseAPIUsage,
ResponseCompletedEvent,
ResponsesAPIResponse,
ResponsesAPIStreamEvents,
ResponsesAPIStreamingResponse,
)
from litellm.types.responses.main import (
GenericResponseOutputItem,
@ -63,11 +83,13 @@ from litellm.types.utils import GenericGuardrailAPIInputs
if TYPE_CHECKING:
from fastapi import HTTPException
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.integrations.custom_guardrail import (
CustomGuardrail,
ModifyResponseException,
)
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.llms.openai import ResponseInputParam
from litellm.types.utils import ResponsesAPIResponse
class ResponseOutputEnvelope(TypedDict, total=False):
@ -865,3 +887,331 @@ class OpenAIResponsesHandler(BaseTranslation):
content[content_idx]["text"] = guardrail_response
elif hasattr(content[content_idx], "text"):
content[content_idx].text = guardrail_response
def build_block_sse_chunks(
self,
exc: "ModifyResponseException",
stream_started: bool = False,
responses_so_far: Sequence[object] | None = None,
) -> Sequence[bytes]:
"""
Build Responses API SSE events that deliver the guardrail block message
and terminate the stream cleanly, mirroring the non-streaming block
response: a completed response whose only output is the violation text,
with the real usage the upstream call consumed.
- ``stream_started`` False (buffered / pre-stream): nothing has been
sent, so emit the full synthetic sequence (``response.created``
through ``response.completed``).
- ``stream_started`` True (sampling / mid-stream): events already
reached the client, so continue the in-progress response: close the
output item still open on the wire, deliver the block message as a
new output item under the same response id, and close with a
``response.completed`` carrying only the replacement item.
The proxy's data generator appends ``data: [DONE]`` itself.
"""
events: Final = (
self._block_continuation_events(exc, responses_so_far or ())
if stream_started
else self._standalone_block_events(exc)
)
return tuple(
f"data: {event.model_dump_json(exclude_none=True, exclude_unset=True, serialize_as_any=True)}\n\n".encode()
for event in events
)
@staticmethod
def _standalone_block_events(exc: "ModifyResponseException") -> Sequence[ResponsesAPIStreamingResponse]:
from litellm.responses.streaming_iterator import build_synthetic_response_events
return build_synthetic_response_events(
transformed=_blocked_response(exc, response_id=f"resp_{uuid.uuid4()}", model=exc.model),
logging_obj=None,
chunk_size=max(len(exc.message), 1),
)
@staticmethod
def _block_continuation_events(
exc: "ModifyResponseException", responses_so_far: Sequence[object]
) -> Sequence[ResponsesAPIStreamingResponse]:
response_id, model, output_index = _continuation_identity(exc, responses_so_far)
item: Final = _blocked_output_item(exc)
item_id: Final = item.id
part: Final[_BlockedContentPart] = {"type": "output_text", "text": exc.message, "annotations": ()}
done_part: Final[_BlockedDoneContentPart] = {
"type": "output_text",
"text": exc.message,
"annotations": (),
"logprobs": None,
}
return (
*_open_item_closing_events(responses_so_far),
OutputItemAddedEvent(
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED,
output_index=output_index,
item=item,
),
ContentPartAddedEvent(
type=ResponsesAPIStreamEvents.CONTENT_PART_ADDED,
item_id=item_id,
output_index=output_index,
content_index=0,
part=BaseLiteLLMOpenAIResponseObject.model_validate(part),
),
OutputTextDeltaEvent(
type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA,
item_id=item_id,
output_index=output_index,
content_index=0,
delta=exc.message,
),
OutputTextDoneEvent(
type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE,
item_id=item_id,
output_index=output_index,
content_index=0,
text=exc.message,
),
ContentPartDoneEvent(
type=ResponsesAPIStreamEvents.CONTENT_PART_DONE,
item_id=item_id,
output_index=output_index,
content_index=0,
part=ContentPartDonePartOutputText.model_validate(done_part),
),
OutputItemDoneEvent(
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE,
output_index=output_index,
item=item,
),
ResponseCompletedEvent(
type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
response=_blocked_response(exc, response_id=response_id, model=model, output_item=item),
),
)
class _BlockedContentPart(TypedDict):
type: ReadOnly[str]
text: ReadOnly[str]
annotations: ReadOnly[tuple[object, ...]]
class _BlockedDoneContentPart(TypedDict):
type: ReadOnly[str]
text: ReadOnly[str]
annotations: ReadOnly[tuple[object, ...]]
logprobs: ReadOnly[None]
class _BlockedItemPayload(TypedDict):
type: ReadOnly[str]
id: ReadOnly[str]
status: ReadOnly[str]
role: ReadOnly[str]
content: ReadOnly[tuple[_BlockedContentPart, ...]]
class _BlockedResponsePayload(TypedDict):
id: ReadOnly[str]
object: ReadOnly[str]
created_at: ReadOnly[int]
model: ReadOnly[str]
output: ReadOnly[tuple[GenericResponseOutputItem, ...]]
status: ReadOnly[str]
usage: ReadOnly[ResponseAPIUsage]
def _blocked_output_item(exc: "ModifyResponseException") -> GenericResponseOutputItem:
payload: Final[_BlockedItemPayload] = {
"type": "message",
"id": f"msg_{uuid.uuid4()}",
"status": "completed",
"role": "assistant",
"content": ({"type": "output_text", "text": exc.message, "annotations": ()},),
}
return GenericResponseOutputItem.model_validate(payload)
def _blocked_response(
exc: "ModifyResponseException",
response_id: str,
model: str,
output_item: GenericResponseOutputItem | None = None,
) -> ResponsesAPIResponse:
payload: Final[_BlockedResponsePayload] = {
"id": response_id,
"object": "response",
"created_at": int(time.time()),
"model": model,
"output": (output_item if output_item is not None else _blocked_output_item(exc),),
"status": "completed",
"usage": blocked_responses_stream_usage(exc.original_response),
}
return ResponsesAPIResponse.model_validate(payload)
def _continuation_identity(exc: "ModifyResponseException", responses_so_far: Sequence[object]) -> tuple[str, str, int]:
responses: Final = tuple(
response for item in responses_so_far if (response := stream_item_field(item, "response")) is not None
)
response_id: Final = next(
(rid for response in responses if isinstance(rid := stream_item_field(response, "id"), str) and rid),
f"resp_{uuid.uuid4()}",
)
model: Final = next(
(m for response in responses if isinstance(m := stream_item_field(response, "model"), str) and m),
exc.model,
)
indices: Final = tuple(
index for item in responses_so_far if isinstance(index := stream_item_field(item, "output_index"), int)
)
return response_id, model, max(indices) + 1 if indices else 0
@dataclass(frozen=True, slots=True)
class _OpenItemState:
item_id: str
item_type: str
role: str
output_index: int
content_index: int
text: str
part_open: bool
payload: object
def _open_item_state(responses_so_far: Sequence[object]) -> _OpenItemState | None:
typed: Final = tuple((stream_item_field(event, "type"), event) for event in responses_so_far)
added: Final = tuple(
(added_index, stream_item_field(event, "item"))
for event_type, event in typed
if event_type == "response.output_item.added"
and isinstance(added_index := stream_item_field(event, "output_index"), int)
)
done_indices: Final = frozenset(
done_index
for event_type, event in typed
if event_type == "response.output_item.done"
and isinstance(done_index := stream_item_field(event, "output_index"), int)
)
open_added: Final = tuple((index, payload) for index, payload in added if index not in done_indices)
if not open_added:
return None
output_index, item_payload = open_added[-1]
if item_payload is None:
return None
item_id: Final = stream_item_field(item_payload, "id")
if not isinstance(item_id, str) or not item_id:
return None
raw_type: Final = stream_item_field(item_payload, "type")
raw_role: Final = stream_item_field(item_payload, "role")
part_added: Final = tuple(
part_index
for event_type, event in typed
if event_type == "response.content_part.added"
and stream_item_field(event, "item_id") == item_id
and isinstance(part_index := stream_item_field(event, "content_index"), int)
)
part_done: Final = frozenset(
part_done_index
for event_type, event in typed
if event_type == "response.content_part.done"
and stream_item_field(event, "item_id") == item_id
and isinstance(part_done_index := stream_item_field(event, "content_index"), int)
)
open_parts: Final = tuple(index for index in part_added if index not in part_done)
text: Final = "".join(
delta
for event_type, event in typed
if event_type == "response.output_text.delta"
and stream_item_field(event, "item_id") == item_id
and isinstance(delta := stream_item_field(event, "delta"), str)
)
return _OpenItemState(
item_id=item_id,
item_type=raw_type if isinstance(raw_type, str) and raw_type else "message",
role=raw_role if isinstance(raw_role, str) and raw_role else "assistant",
output_index=output_index,
content_index=open_parts[-1] if open_parts else 0,
text=text,
part_open=bool(open_parts),
payload=item_payload,
)
_item_fields_adapter: Final = TypeAdapter(Mapping[str, object])
_no_item_fields: Final[Mapping[str, object]] = MappingProxyType({})
def _incomplete_item_fields(payload: object) -> Mapping[str, object]:
raw: Final = payload.model_dump() if isinstance(payload, BaseModel) else payload
if not isinstance(raw, dict):
return _no_item_fields
return _item_fields_adapter.validate_python(raw)
def _open_item_closing_events(responses_so_far: Sequence[object]) -> Sequence[ResponsesAPIStreamingResponse]:
"""Close the output item still in progress on the relayed stream before the
block item is appended: strict Responses clients reject a
``response.completed`` that arrives while an earlier ``output_item.added``
was never closed. A message item closes ``completed`` with exactly the text
the client has received so far; any other item type (a function call the
guardrail rejected, for instance) closes ``incomplete`` so the synthetic
done event can never authorize acting on it."""
open_item: Final = _open_item_state(responses_so_far)
if open_item is None:
return ()
if open_item.item_type != "message":
return (
OutputItemDoneEvent(
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE,
output_index=open_item.output_index,
item=BaseLiteLLMOpenAIResponseObject.model_validate(
MappingProxyType({**_incomplete_item_fields(open_item.payload), "status": "incomplete"})
),
),
)
partial_part: Final[_BlockedContentPart] = {
"type": "output_text",
"text": open_item.text,
"annotations": (),
}
closed_payload: Final[_BlockedItemPayload] = {
"type": open_item.item_type,
"id": open_item.item_id,
"status": "completed",
"role": open_item.role,
"content": (partial_part,),
}
item_done: Final = OutputItemDoneEvent(
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE,
output_index=open_item.output_index,
item=GenericResponseOutputItem.model_validate(closed_payload),
)
if not open_item.part_open:
return (item_done,)
partial_done_part: Final[_BlockedDoneContentPart] = {
"type": "output_text",
"text": open_item.text,
"annotations": (),
"logprobs": None,
}
return (
OutputTextDoneEvent(
type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE,
item_id=open_item.item_id,
output_index=open_item.output_index,
content_index=open_item.content_index,
text=open_item.text,
),
ContentPartDoneEvent(
type=ResponsesAPIStreamEvents.CONTENT_PART_DONE,
item_id=open_item.item_id,
output_index=open_item.output_index,
content_index=open_item.content_index,
part=ContentPartDonePartOutputText.model_validate(partial_done_part),
),
item_done,
)

90
litellm/llms/parallel_ai/search/cost_calculator.py Normal file
View file

@ -0,0 +1,90 @@
from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import Final
from pydantic import TypeAdapter, ValidationError
from litellm.utils import get_model_info
PARALLEL_AI_DEFAULT_RESULTS: Final = 10
PARALLEL_AI_ADDITIONAL_RESULT_COST: Final = 0.001
PARALLEL_AI_USAGE_PARAM: Final = "_parallel_ai_usage"
PARALLEL_AI_STANDARD_SEARCH_MODEL: Final = "parallel_ai/search"
PARALLEL_AI_FAST_SEARCH_MODEL: Final = "parallel_ai/search-fast"
PARALLEL_AI_TURBO_SEARCH_MODEL: Final = "parallel_ai/search-turbo"
PARALLEL_AI_PRICING_MODEL_BY_MODE: Final[Mapping[str, str]] = MappingProxyType(
{
"fast": PARALLEL_AI_FAST_SEARCH_MODEL,
"turbo": PARALLEL_AI_TURBO_SEARCH_MODEL,
}
)
ADVANCED_SETTINGS_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object])
def _non_negative_int(value: object) -> int | None:
if isinstance(value, bool) or not isinstance(value, int) or value < 0:
return None
return value
def _usage_count(usage: Sequence[Mapping[str, object]], sku: str) -> int | None:
counts: Final = tuple(
count
for item in usage
if item.get("name") == sku
if (count := _non_negative_int(item.get("count"))) is not None
)
return sum(counts) if counts else None
def _effective_mode(optional_params: Mapping[str, object]) -> str:
mode: Final = optional_params.get("mode")
if isinstance(mode, str):
return mode
processor: Final = optional_params.get("processor")
if processor == "pro":
return "advanced"
return "basic"
def _effective_max_results(optional_params: Mapping[str, object]) -> int:
try:
advanced_settings: Final = ADVANCED_SETTINGS_ADAPTER.validate_python(optional_params.get("advanced_settings"))
advanced_max_results: Final = _non_negative_int(advanced_settings.get("max_results"))
if advanced_max_results is not None:
return advanced_max_results
except ValidationError:
pass
max_results: Final = _non_negative_int(optional_params.get("max_results"))
return max_results if max_results is not None else PARALLEL_AI_DEFAULT_RESULTS
def _request_cost(mode: str) -> float:
pricing_model: Final = PARALLEL_AI_PRICING_MODEL_BY_MODE.get(mode, PARALLEL_AI_STANDARD_SEARCH_MODEL)
model_info: Final = get_model_info(model=pricing_model, custom_llm_provider="parallel_ai")
return float(model_info.get("input_cost_per_query") or 0.0)
def _additional_results(
optional_params: Mapping[str, object],
usage: Sequence[Mapping[str, object]] | None,
) -> int:
usage_count: Final = _usage_count(usage, "sku_search_additional_results") if usage is not None else None
if usage_count is not None:
return usage_count
if usage is not None:
return 0
return max(_effective_max_results(optional_params) - PARALLEL_AI_DEFAULT_RESULTS, 0)
def parallel_ai_search_cost(
optional_params: Mapping[str, object],
usage: Sequence[Mapping[str, object]] | None,
) -> float:
request_cost: Final = _request_cost(_effective_mode(optional_params))
request_count_from_usage: Final = _usage_count(usage, "sku_search") if usage is not None else None
request_count: Final = request_count_from_usage if request_count_from_usage is not None else 1
additional_results: Final = _additional_results(optional_params, usage)
return request_count * request_cost + additional_results * PARALLEL_AI_ADDITIONAL_RESULT_COST

143
litellm/llms/parallel_ai/search/transformation.py
View file

@ -4,9 +4,13 @@ Calls Parallel AI's /v1/search endpoint to search the web.
Parallel AI API Reference: https://docs.parallel.ai/api-reference/search/search
"""
from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import Final, TypedDict
import httpx
from pydantic import BaseModel, ConfigDict
from typing_extensions import ReadOnly
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.search.transformation import (
@ -14,9 +18,29 @@ from litellm.llms.base_llm.search.transformation import (
SearchResponse,
SearchResult,
)
from litellm.llms.parallel_ai.search.cost_calculator import PARALLEL_AI_USAGE_PARAM
from litellm.secret_managers.main import get_secret_str
class _ParallelAIV1SearchResult(BaseModel):
model_config = ConfigDict(extra="ignore")
url: str | None = None
title: str | None = None
publish_date: str | None = None
excerpts: Sequence[str] | None = None
class _ParallelAIV1SearchResponse(BaseModel):
model_config = ConfigDict(extra="ignore")
search_id: str | None = None
session_id: str | None = None
results: Sequence[_ParallelAIV1SearchResult] = ()
usage: Sequence[Mapping[str, object]] | None = None
warnings: Sequence[Mapping[str, object]] | None = None
class _ParallelAISourcePolicy(TypedDict, total=False):
include_domains: list[str]
exclude_domains: list[str]
@ -27,10 +51,16 @@ class _ParallelAIExcerptSettings(TypedDict, total=False):
max_chars_per_result: int
class _ParallelAIFetchPolicy(TypedDict, total=False):
max_age_seconds: ReadOnly[int]
timeout_seconds: ReadOnly[float]
disable_cache_fallback: ReadOnly[bool]
class _ParallelAIAdvancedSettings(TypedDict, total=False):
source_policy: _ParallelAISourcePolicy
excerpt_settings: _ParallelAIExcerptSettings
fetch_policy: dict
fetch_policy: _ParallelAIFetchPolicy
location: str
max_results: int
@ -43,14 +73,14 @@ class ParallelAISearchRequest(TypedDict, total=False):
search_queries: list[str] # Required - at least one keyword search query
objective: str # Optional - natural-language description of search goal
mode: str # Optional - 'turbo', 'basic', or 'advanced' (default 'advanced')
mode: str # Optional - 'turbo', 'fast', 'basic', or 'advanced' (default 'advanced')
max_chars_total: int # Optional - upper bound on total excerpt characters
session_id: str # Optional - tracks calls across search/extract requests
client_model: str # Optional - model consuming the results
advanced_settings: _ParallelAIAdvancedSettings
LEGACY_PROCESSOR_TO_MODE: Final = {"base": "basic", "pro": "advanced"}
LEGACY_PROCESSOR_TO_MODE: Final = MappingProxyType({"base": "basic", "pro": "advanced"})
class ParallelAISearchConfig(BaseSearchConfig):
@ -67,16 +97,16 @@ class ParallelAISearchConfig(BaseSearchConfig):
api_base: str | None = None,
**kwargs,
) -> dict:
api_key = self.resolve_server_api_key(
resolved_api_key: Final = self.resolve_server_api_key(
caller_api_key=api_key,
caller_api_base=api_base,
key_env_vars=("PARALLEL_AI_API_KEY", "PARALLEL_API_KEY"),
base_env_var="PARALLEL_AI_API_BASE",
default_api_base=self.PARALLEL_AI_API_BASE,
)
if not api_key:
if not resolved_api_key:
raise ValueError("PARALLEL_API_KEY is not set. Set `PARALLEL_API_KEY` environment variable.")
headers["x-api-key"] = api_key
headers["x-api-key"] = resolved_api_key
headers["Content-Type"] = "application/json"
return headers
@ -87,13 +117,12 @@ class ParallelAISearchConfig(BaseSearchConfig):
data: dict | list[dict] | None = None,
**kwargs,
) -> str:
api_base = api_base or get_secret_str("PARALLEL_AI_API_BASE") or self.PARALLEL_AI_API_BASE
resolved_api_base: Final = api_base or get_secret_str("PARALLEL_AI_API_BASE") or self.PARALLEL_AI_API_BASE
api_base = api_base.rstrip("/")
if not api_base.endswith("/v1/search"):
api_base = f"{api_base.removesuffix('/v1')}/v1/search"
return api_base
trimmed: Final = resolved_api_base.rstrip("/")
if trimmed.endswith("/v1/search"):
return trimmed
return f"{trimmed.removesuffix('/v1')}/v1/search"
def transform_search_request(
self,
@ -109,14 +138,17 @@ class ParallelAISearchConfig(BaseSearchConfig):
- If string: maps to `search_queries` (single item) and `objective`
- If list: maps to `search_queries` (keyword queries)
optional_params: Optional parameters for the request
- mode: Search mode ('turbo', 'basic', 'advanced'); defaults to 'basic'
- mode: Search mode ('turbo', 'fast', 'basic', 'advanced'); defaults to 'basic'
- processor: Legacy v1beta param; 'base' maps to mode 'basic', 'pro' to 'advanced'
- max_results: Maximum number of search results -> `advanced_settings.max_results`
- search_domain_filter: Domains to include -> `advanced_settings.source_policy.include_domains`
- search_domain_filter / include_domains: Domains to include -> `advanced_settings.source_policy.include_domains`
- exclude_domains: Domains to exclude -> `advanced_settings.source_policy.exclude_domains`
- country: ISO 3166-1 alpha-2 code -> `advanced_settings.location`
- after_date: RFC 3339 date (YYYY-MM-DD) -> `advanced_settings.source_policy.after_date`
- country / location: ISO 3166-1 alpha-2 code -> `advanced_settings.location`
- max_chars_per_result: -> `advanced_settings.excerpt_settings.max_chars_per_result`
- Any other params are passed through to the request body as-is
- fetch_policy: Cache vs live-fetch policy -> `advanced_settings.fetch_policy`
- Any other params (objective, max_chars_total, session_id, client_model, ...)
are passed through to the request body as-is
Returns:
Dict with request data following the v1 search request spec
@ -137,7 +169,7 @@ class ParallelAISearchConfig(BaseSearchConfig):
mode = LEGACY_PROCESSOR_TO_MODE.get(processor, processor)
# the v1 API defaults to 'advanced' when mode is omitted; default to 'basic'
# instead to keep v1beta's default tier (processor 'base') and litellm's
# $0.004/query cost map entry for `parallel_ai/search` accurate
# cost map entry for `parallel_ai/search` accurate
request_data["mode"] = mode or "basic"
advanced_settings: Final[_ParallelAIAdvancedSettings] = {}
@ -148,17 +180,29 @@ class ParallelAISearchConfig(BaseSearchConfig):
if "country" in params:
advanced_settings["location"] = params.pop("country")
if "location" in params:
advanced_settings["location"] = params.pop("location")
if "max_chars_per_result" in params:
advanced_settings["excerpt_settings"] = {"max_chars_per_result": params.pop("max_chars_per_result")}
if "fetch_policy" in params:
advanced_settings["fetch_policy"] = params.pop("fetch_policy")
source_policy: Final[_ParallelAISourcePolicy] = {}
if "search_domain_filter" in params:
source_policy["include_domains"] = params.pop("search_domain_filter")
if "include_domains" in params:
source_policy["include_domains"] = params.pop("include_domains")
if "exclude_domains" in params:
source_policy["exclude_domains"] = params.pop("exclude_domains")
if "after_date" in params:
source_policy["after_date"] = params.pop("after_date")
if source_policy:
advanced_settings["source_policy"] = source_policy
@ -170,9 +214,11 @@ class ParallelAISearchConfig(BaseSearchConfig):
# unified-spec param with no v1 equivalent
params.pop("max_tokens_per_page", None)
result_data: Final[dict] = dict(request_data)
result_data.update(params)
return result_data
# reserved for the provider's own reported usage, which prices the request;
# a caller-supplied value would otherwise set its own cost
params.pop(PARALLEL_AI_USAGE_PARAM, None)
return {**request_data, **params}
def transform_search_response(
self,
@ -186,26 +232,49 @@ class ParallelAISearchConfig(BaseSearchConfig):
Parallel AI -> LiteLLM mappings:
- results[].title -> SearchResult.title
- results[].url -> SearchResult.url
- results[].excerpts (array) -> SearchResult.snippet (joined string)
- results[].excerpts (array) -> SearchResult.snippet (joined string); the raw
array is preserved as an extra `excerpts` field on each result
- results[].publish_date -> SearchResult.date
- search_id / session_id / warnings are preserved as extra fields on the
response; usage is preserved as `parallel_usage` (the `usage` name is
reserved for LiteLLM's token-usage object)
"""
response_json: Final = raw_response.json()
parsed: Final = _ParallelAIV1SearchResponse.model_validate(raw_response.json())
results: Final = []
for result in response_json.get("results", []):
excerpts = result.get("excerpts") or []
snippet = " ... ".join(excerpts) if excerpts else ""
# written unconditionally: leaving a caller-supplied value in place when the
# provider reports no usage would let the caller price its own request
logging_obj.optional_params = {
**logging_obj.optional_params,
PARALLEL_AI_USAGE_PARAM: parsed.usage,
}
search_result = SearchResult(
title=result.get("title") or "",
url=result.get("url") or "",
snippet=snippet,
date=result.get("publish_date"),
last_updated=None,
results: Final = tuple(
SearchResult.model_validate(
MappingProxyType(
{
"title": result.title or "",
"url": result.url or "",
"snippet": " ... ".join(result.excerpts or ()),
"date": result.publish_date,
"last_updated": None,
"excerpts": result.excerpts or (),
}
)
)
results.append(search_result)
return SearchResponse(
results=results,
object="search",
for result in parsed.results
)
extra_fields: Final = MappingProxyType(
{
key: value
for key, value in (
("search_id", parsed.search_id),
("session_id", parsed.session_id),
("parallel_usage", parsed.usage),
("warnings", parsed.warnings),
)
if value is not None
}
)
return SearchResponse.model_validate(MappingProxyType({"results": results, "object": "search", **extra_fields}))

2
litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
View file

@ -949,7 +949,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
# For Gemini 3+ models, use thinkingLevel instead of thinkingBudget
if model and VertexGeminiConfig._is_gemini_3_or_newer(model):
if thinking_enabled:
if thinking_budget is None or thinking_budget == 0:
if thinking_budget == 0:
params["includeThoughts"] = False
else:
params["includeThoughts"] = True

3
litellm/llms/vertex_ai/ocr/deepseek_transformation.py
View file

@ -177,8 +177,9 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig):
content_item = {"type": "image_url", "image_url": document_url}
# Build DeepSeek OCR request
provider_model: Final = model if model.startswith("deepseek-ai/") else f"deepseek-ai/{model}"
data: Final = {
"model": "deepseek-ai/" + model,
"model": provider_model,
"messages": [{"role": "user", "content": [content_item]}],
}

22
litellm/main.py
View file

@ -8637,6 +8637,16 @@ def _set_stream_builder_response_cost(response: ModelResponse, logging_obj: Opti
hidden_params["response_cost"] = response_cost
def _stamp_streaming_usage_cost(usage: Usage, response: ModelResponse, logging_obj: Optional["Logging"]) -> None:
if logging_obj is None:
return
if isinstance(getattr(usage, "cost", None), (int, float)):
return
computed_cost: Final = logging_obj._response_cost_calculator(result=response)
if isinstance(computed_cost, (int, float)) and computed_cost > 0:
setattr(usage, "cost", computed_cost)
def stream_chunk_builder(
chunks: list,
messages: list | None = None,
@ -8731,12 +8741,7 @@ def stream_chunk_builder(
)
break
if litellm.include_cost_in_streaming_usage and logging_obj is not None:
setattr(
usage,
"cost",
logging_obj._response_cost_calculator(result=response),
)
_stamp_streaming_usage_cost(usage, response, logging_obj)
_set_stream_builder_response_cost(response, logging_obj)
processor.apply_provider_assembled_streaming_metadata(response, chunks, logging_obj)
@ -8915,10 +8920,7 @@ def stream_chunk_builder(
)
break
# Add cost to usage object if include_cost_in_streaming_usage is True
if litellm.include_cost_in_streaming_usage and logging_obj is not None:
setattr(usage, "cost", logging_obj._response_cost_calculator(result=response))
_stamp_streaming_usage_cost(usage, response, logging_obj)
_set_stream_builder_response_cost(response, logging_obj)
processor.apply_provider_assembled_streaming_metadata(response, chunks, logging_obj)

1005
litellm/model_prices_and_context_window_backup.json
View file

File diff suppressed because it is too large Load diff

10
litellm/models/user.py
View file

@ -7,7 +7,7 @@ Canonical definition for ``litellm_usertable``. Re-exported from
from datetime import datetime
from pydantic import ConfigDict, Field, model_validator
from pydantic import BaseModel, ConfigDict, Field, model_validator
from litellm.models.object_permission import LiteLLM_ObjectPermissionTable
from litellm.models.organization_membership import (
@ -67,3 +67,11 @@ class LiteLLM_UserTable(LiteLLMPydanticObjectBase):
if not self.models:
return True
return model_name in self.models
class SCIMPlaceholder(BaseModel):
"""A user row keyed by a value that names another account by SSO identity or email."""
placeholder_user_id: str
resolved_user_ids: tuple[str, ...]
team_ids: tuple[str, ...]

71
litellm/proxy/_experimental/mcp_server/rest_endpoints.py
View file

@ -1,7 +1,8 @@
import asyncio
import importlib
from collections.abc import Awaitable, Callable, Mapping
from collections.abc import Awaitable, Callable, Mapping, Sequence
from datetime import datetime
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal
import anyio
@ -20,8 +21,11 @@ from litellm.proxy._experimental.mcp_server.exceptions import (
MCPUpstreamAuthError,
)
from litellm.proxy._experimental.mcp_server.faults.list_outcomes import (
ServerListOk,
ServerOutcome,
classify_list_exception,
list_fault_http_status,
outcome_wire_value,
)
from litellm.proxy._experimental.mcp_server.ui_session_utils import (
acting_user_auth,
@ -99,6 +103,7 @@ if MCP_AVAILABLE:
ListMCPToolsRestAPIResponseObject,
MCPInfo,
MCPServer,
_aggregate_server_key, # pyright: ignore[reportPrivateUsage] # same per-server key as the tools/list _meta outcomes
_apply_toolset_scope,
_fire_mcp_tool_call_logging,
execute_mcp_tool,
@ -803,9 +808,6 @@ if MCP_AVAILABLE:
list(allowed_server_ids_set), _rest_client_ip
)
list_tools_result: Final = []
error_message = None
# If server_id is specified, only query that specific server
if server_id:
return await _list_tools_for_single_server(
@ -849,22 +851,19 @@ if MCP_AVAILABLE:
else {}
)
# Query all servers the user has access to
errors: Final = []
for allowed_server_id in allowed_server_ids:
server = global_mcp_server_manager.get_mcp_server_by_id(allowed_server_id)
if server is None:
continue
server_auth_header = _get_server_auth_header(server, mcp_server_auth_headers, mcp_auth_header)
user_oauth_extra_headers = await _get_user_oauth_extra_headers(
async def list_server(
server: MCPServer,
) -> tuple[Sequence[ListMCPToolsRestAPIResponseObject], ServerOutcome]:
server_auth_header: Final = _get_server_auth_header(
server, mcp_server_auth_headers, mcp_auth_header
)
user_oauth_extra_headers: Final = await _get_user_oauth_extra_headers(
server,
user_api_key_dict,
prefetched_creds=prefetched_oauth_creds,
)
try:
tools_result = await _get_tools_for_single_server(
tools_result: Final = await _get_tools_for_single_server(
server,
server_auth_header,
raw_headers_from_request,
@ -872,24 +871,36 @@ if MCP_AVAILABLE:
extra_headers=user_oauth_extra_headers,
apply_tool_filters=apply_tool_filters,
)
list_tools_result.extend(tools_result)
except Exception as e:
verbose_logger.exception("Error getting tools from %s: %s", server.name, e)
errors.append(
f"{get_server_prefix(server)}: {classify_list_exception(e).tag}"
if isinstance(e, (MCPServerListError, MCPUpstreamAuthError))
else f"{get_server_prefix(server)}: {e}"
)
continue
return (), classify_list_exception(e)
return tools_result, ServerListOk(tool_count=len(tools_result))
if errors and not list_tools_result:
error_message = "Failed to get tools from servers: " + "; ".join(errors)
return {
"tools": list_tools_result,
"error": "partial_failure" if error_message else None,
"message": (error_message if error_message else "Successfully retrieved tools"),
}
# Query all servers the user has access to
queried_servers: Final = tuple(
server
for server in map(global_mcp_server_manager.get_mcp_server_by_id, allowed_server_ids)
if server is not None
)
listings: Final = tuple([await list_server(server) for server in queried_servers])
list_tools_result: Final = [tool for tools, _ in listings for tool in tools]
server_outcomes: Final = MappingProxyType(
{_aggregate_server_key(server): outcome for server, (_, outcome) in zip(queried_servers, listings)}
)
errors: Final = tuple(
f"{key}: {outcome.tag}" for key, outcome in server_outcomes.items() if outcome.tag != "ok"
)
error_message: Final = (
"Failed to get tools from servers: " + "; ".join(errors)
if errors and not list_tools_result
else None
)
return {
"tools": list_tools_result,
"error": "partial_failure" if error_message else None,
"message": (error_message if error_message else "Successfully retrieved tools"),
"server_outcomes": {key: outcome_wire_value(outcome) for key, outcome in server_outcomes.items()},
}
except MCPUpstreamAuthError as e:
# Surface upstream pass-through 401/403 challenges to the client so

179
litellm/proxy/_lazy_openapi_snapshot.json
View file

@ -32002,6 +32002,62 @@
"title": "SCIMPatchOperation",
"type": "object"
},
"SCIMPlaceholder": {
"description": "A user row keyed by a value that names another account by SSO identity or email.",
"properties": {
"placeholder_user_id": {
"title": "Placeholder User Id",
"type": "string"
},
"resolved_user_ids": {
"items": {
"type": "string"
},
"title": "Resolved User Ids",
"type": "array"
},
"team_ids": {
"items": {
"type": "string"
},
"title": "Team Ids",
"type": "array"
}
},
"required": [
"placeholder_user_id",
"resolved_user_ids",
"team_ids"
],
"title": "SCIMPlaceholder",
"type": "object"
},
"SCIMPlaceholderMergeResult": {
"properties": {
"merged_into_user_id": {
"title": "Merged Into User Id",
"type": "string"
},
"placeholder_user_id": {
"title": "Placeholder User Id",
"type": "string"
},
"team_ids": {
"items": {
"type": "string"
},
"title": "Team Ids",
"type": "array"
}
},
"required": [
"placeholder_user_id",
"merged_into_user_id",
"team_ids"
],
"title": "SCIMPlaceholderMergeResult",
"type": "object"
},
"SCIMServiceProviderConfig": {
"properties": {
"authenticationSchemes": {
@ -33641,6 +33697,129 @@
"scim"
]
}
},
"/scim/v2/placeholders": {
"get": {
"description": "List user rows whose id is another account's SSO identity or email.\n\nAn earlier release provisioned a group member it could not match as a user keyed\nby the raw member value, and that row now shadows the account the value really\nnames, so every push of that member is refused. This lists those rows so an\noperator can fold each one into the account it shadows with\n``POST /scim/v2/placeholders/{user_id}/merge``. A row that has an SSO identity of\nits own or owns virtual keys is left out: someone uses that account.",
"operationId": "list_placeholders_scim_v2_placeholders_get",
"parameters": [
{
"in": "query",
"name": "feature",
"required": false,
"schema": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Feature"
}
}
],
"responses": {
"200": {
"content": {
"application/json": {
"schema": {
"items": {
"$ref": "#/components/schemas/SCIMPlaceholder"
},
"title": "Response List Placeholders Scim V2 Placeholders Get",
"type": "array"
}
}
},
"description": "Successful Response"
},
"422": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/HTTPValidationError"
}
}
},
"description": "Validation Error"
}
},
"security": [
{
"APIKeyHeader": []
}
],
"summary": "List Placeholders",
"tags": [
"scim"
]
}
},
"/scim/v2/placeholders/{user_id}/merge": {
"post": {
"description": "Fold a placeholder user into the one account its id names by SSO identity or email.\n\nThe account is added to every team the placeholder is on, then the placeholder is\ndeleted the way ``DELETE /scim/v2/Users/{id}`` deletes a user, so the next group\npush resolves the member value to the real account. Refused with 409 when the row\nhas an SSO identity of its own, owns virtual keys, or names no account or several.",
"operationId": "merge_placeholder_scim_v2_placeholders__user_id__merge_post",
"parameters": [
{
"in": "path",
"name": "user_id",
"required": true,
"schema": {
"title": "User ID",
"type": "string"
}
},
{
"in": "query",
"name": "feature",
"required": false,
"schema": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Feature"
}
}
],
"responses": {
"200": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/SCIMPlaceholderMergeResult"
}
}
},
"description": "Successful Response"
},
"422": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/HTTPValidationError"
}
}
},
"description": "Validation Error"
}
},
"security": [
{
"APIKeyHeader": []
}
],
"summary": "Merge Placeholder",
"tags": [
"scim"
]
}
}
}
},

25
litellm/proxy/common_utils/model_listing_utils.py
View file

@ -10,13 +10,36 @@ legacy internal names with `general_settings.use_team_public_model_name: false`.
from __future__ import annotations
from collections.abc import Mapping
from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, cast
if TYPE_CHECKING:
from litellm.router import Router
def configured_display_names(
entries: Sequence[tuple[str, str]],
llm_router: Router | None,
) -> Mapping[str, str]:
"""response_id -> configured `model_info.display_name` for the listing entries
that have one.
Metadata is looked up by each entry's internal lookup id (so team-scoped rows
resolve), while the returned map is keyed by the public response id the
Anthropic-shaped listing is built from. Entries without a configured name are
omitted so the listing falls back to the id itself.
"""
if llm_router is None:
return MappingProxyType({})
resolved: Final = (
(response_id, llm_router.get_configured_display_name(lookup_id)) for response_id, lookup_id in entries
)
return MappingProxyType(
{response_id: display_name for response_id, display_name in resolved if display_name is not None}
)
class TeamModelNameTranslator:
"""Translates internal team routing keys to their public names for the model
listing/retrieve responses. Stateless; the live router and general_settings

124
litellm/proxy/management_endpoints/scim/scim_v2.py
View file

@ -29,6 +29,7 @@ import litellm
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.models.user import SCIMPlaceholder
from litellm.proxy._types import (
LiteLLM_TeamTable,
LiteLLM_UserTable,
@ -585,6 +586,37 @@ async def _users_named_by_member_value(
return tuple(dict.fromkeys(row.user_id for row in rows))
async def _accounts_named_by_member_value(value: str, prisma_client: PrismaClient) -> tuple[str, ...]:
"""Every user id this member value names, by user id, SSO identity or email.
Classification needs to know whether the value is one account's ``user_id`` and
whether it names any other account, so all three fields are read in one pass. The
id is compared exactly and unstripped, as a primary key lookup would; the
identities compare as ``_users_named_by_member_value`` describes. Two rows are
enough to tell one account from several, so the read stops there. Only a full
read that lacks the row keyed by the value leaves that row's existence open, and
only then is the id read on its own.
"""
subject: Final = value.strip()
email: Final[_CaseInsensitiveMatch] = {"equals": subject, "mode": "insensitive"}
users: Final = _table(UserRepository(prisma_client))
rows: Final = await users.find_many(
where={ # mutable-ok: Prisma filter
"OR": [ # mutable-ok: Prisma filter
{"user_id": value}, # mutable-ok: Prisma filter
{"sso_user_id": subject}, # mutable-ok: Prisma filter
{"user_email": email}, # mutable-ok: Prisma filter
],
},
take=2,
)
named: Final = tuple(dict.fromkeys(row.user_id for row in rows))
if len(named) < 2 or value in named:
return named
keyed: Final = await users.find_unique(where={"user_id": value})
return named if keyed is None else (value, *named)
async def _classify_group_member(member: SCIMMember, prisma_client: PrismaClient) -> _ClassifiedGroupMember:
"""
Decide what a single SCIM group member refers to.
@ -627,11 +659,9 @@ async def _classify_group_member(member: SCIMMember, prisma_client: PrismaClient
if member_type == "group":
return _SkippedGroupMember(value=value, reason="nested_group")
user: Final = await _table(UserRepository(prisma_client)).find_unique(where={"user_id": value})
if user is not None:
shared_with: Final = tuple(
other for other in await _users_named_by_member_value(value, prisma_client) if other != value
)
named: Final = await _accounts_named_by_member_value(value, prisma_client)
if value in named:
shared_with: Final = tuple(other for other in named if other != value)
if shared_with:
verbose_proxy_logger.warning(
"SCIM: group member '%s' is one account's user id and is also account '%s' by SSO identity or email, "
@ -651,7 +681,6 @@ async def _classify_group_member(member: SCIMMember, prisma_client: PrismaClient
if team is not None and _team_metadata_has_scim_provenance(team.metadata):
return _SkippedGroupMember(value=value, reason="existing_team")
named: Final = await _users_named_by_member_value(value, prisma_client)
if len(named) == 1:
verbose_proxy_logger.info(
"SCIM: group member '%s' matched user_id '%s' by SSO identity or email",
@ -1834,6 +1863,89 @@ async def delete_user(
raise handle_exception_on_proxy(e)
@scim_router.get(
"/placeholders",
response_model=tuple[SCIMPlaceholder, ...],
dependencies=(Depends(user_api_key_auth),),
)
async def list_placeholders() -> tuple[SCIMPlaceholder, ...]:
"""
List user rows whose id is another account's SSO identity or email.
An earlier release provisioned a group member it could not match as a user keyed
by the raw member value, and that row now shadows the account the value really
names, so every push of that member is refused. This lists those rows so an
operator can fold each one into the account it shadows with
``POST /scim/v2/placeholders/{user_id}/merge``. A row that has an SSO identity of
its own or owns virtual keys is left out: someone uses that account.
"""
try:
prisma_client: Final = await _get_prisma_client_or_raise_exception()
async with prisma_client.tx() as tx:
return await UserRepository(prisma_client).find_shadowing_placeholders(tx)
except Exception as e:
raise handle_exception_on_proxy(e)
def _placeholder_rejection(placeholder: LiteLLM_UserTable, resolved: tuple[str, ...], key_count: int) -> str | None:
if placeholder.sso_user_id is not None:
return f"User '{placeholder.user_id}' has an SSO identity of its own, so it is an account someone signs in to"
if key_count:
return f"User '{placeholder.user_id}' owns {key_count} virtual keys. Move or delete them before merging it"
if not resolved:
return f"User '{placeholder.user_id}' shadows no account: no other user has that id as SSO identity or email"
if len(resolved) > 1:
return (
f"User '{placeholder.user_id}' names {len(resolved)} accounts ({', '.join(resolved)}). Resolve that first"
)
return None
@scim_router.post(
"/placeholders/{user_id}/merge",
response_model=SCIMPlaceholderMergeResult,
dependencies=(Depends(user_api_key_auth),),
)
async def merge_placeholder(
user_id: str = Path(..., title="User ID"),
) -> SCIMPlaceholderMergeResult:
"""
Fold a placeholder user into the one account its id names by SSO identity or email.
The account is added to every team the placeholder is on, then the placeholder is
deleted the way ``DELETE /scim/v2/Users/{id}`` deletes a user, so the next group
push resolves the member value to the real account. Refused with 409 when the row
has an SSO identity of its own, owns virtual keys, or names no account or several.
"""
try:
prisma_client: Final = await _get_prisma_client_or_raise_exception()
placeholder: Final = await _check_user_exists(user_id)
resolved: Final = tuple(
other for other in await _users_named_by_member_value(user_id, prisma_client, take=None) if other != user_id
)
owned_keys: Final[_UserIdWhere] = {"user_id": user_id}
keys: Final = await _table(VerificationTokenRepository(prisma_client)).find_many(where=owned_keys)
rejection: Final = _placeholder_rejection(placeholder, resolved, len(keys))
if rejection is not None:
detail: Final[_ScimErrorDetail] = {"error": rejection}
raise HTTPException(status_code=409, detail=detail)
target_user_id: Final = resolved[0]
team_ids: Final = tuple(placeholder.teams)
for team_id in team_ids:
await _add_user_to_team(user_id=target_user_id, team_id=team_id)
await delete_user(user_id=user_id)
await _recompute_scim_member_roles(prisma_client, (target_user_id,))
verbose_proxy_logger.info(
"SCIM: merged placeholder user '%s' into '%s', moving teams %s", user_id, target_user_id, team_ids
)
return SCIMPlaceholderMergeResult(
placeholder_user_id=user_id, merged_into_user_id=target_user_id, team_ids=team_ids
)
except Exception as e:
raise handle_exception_on_proxy(e)
def _parse_member_entry(entry: object) -> SCIMMember | None:
"""Parse one entry of a SCIM patch value, or None when it carries no id."""
if isinstance(entry, str):

6
litellm/proxy/management_helpers/access_group_key_sync.py
View file

@ -38,6 +38,7 @@ from litellm.proxy._types import (
from litellm.proxy.auth.auth_checks import (
_delete_cache_access_object, # pyright: ignore[reportPrivateUsage] # the access-group endpoints reach for this same cache primitive
)
from litellm.proxy.db.routing_prisma_wrapper import WriterPinnedClient
from litellm.repositories.table_repositories import AccessGroupRepository
@ -72,8 +73,9 @@ _REPOINT_KEY_SQL: Final = (
def _raw_executor(prisma_client: object) -> _RawExecutor:
"""Narrow the untyped Prisma client down to the raw-query call this module makes."""
return AccessGroupRepository(prisma_client).prisma_client.db # pyright: ignore[reportAny] # untyped Prisma client
"""Narrow the untyped Prisma client down to the raw-query call this module makes, pinned to the writer."""
db: Final = AccessGroupRepository(prisma_client).prisma_client.db # pyright: ignore[reportAny] # untyped Prisma client
return WriterPinnedClient(db).db # pyright: ignore[reportAny, reportReturnType] # untyped Prisma client behind the pin
async def _invalidate_access_group_cache(access_group_id: str) -> None:

2
litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
View file

@ -812,6 +812,8 @@ def _resolve_team_callback_wiring(
else { # mutable-ok: Logging arg
**callback_vars,
TRUSTED_CALLBACK_VARS_FIELD: callback_vars,
"metadata": {}, # mutable-ok: Logging arg
"model_info": {}, # mutable-ok: Logging arg
}
)
return _TeamCallbackWiring(

21
litellm/proxy/proxy_server.py
View file

@ -352,7 +352,10 @@ from litellm.proxy.common_utils.load_config_utils import (
get_file_contents_from_s3,
)
from litellm.proxy.common_utils.model_deprecation import collect_model_deprecations
from litellm.proxy.common_utils.model_listing_utils import TeamModelNameTranslator
from litellm.proxy.common_utils.model_listing_utils import (
TeamModelNameTranslator,
configured_display_names,
)
from litellm.proxy.common_utils.openai_endpoint_utils import (
remove_sensitive_info_from_deployment,
)
@ -10223,7 +10226,8 @@ async def model_list(
# The internal routing key drives the metadata/fallback lookup, while the
# public name is what the client sees as the model id.
model_data = []
for response_id, lookup_id in TeamModelNameTranslator.listing_entries(all_models, llm_router, settings):
admin_entries: Final = TeamModelNameTranslator.listing_entries(all_models, llm_router, settings)
for response_id, lookup_id in admin_entries:
model_info = create_model_info_response(
model_id=lookup_id,
provider="openai",
@ -10236,7 +10240,10 @@ async def model_list(
if wants_anthropic_format:
admin_listing: Final = cast(Sequence[ModelInfoResponse], model_data) # cast-ok: rows built above
return create_anthropic_model_list_response(admin_listing)
return create_anthropic_model_list_response(
admin_listing,
display_names=configured_display_names(admin_entries, llm_router),
)
return dict(
data=model_data,
@ -10267,7 +10274,8 @@ async def model_list(
# The internal routing key drives the metadata/fallback lookup, while the
# public name is what the client sees as the model id.
model_data = []
for response_id, lookup_id in TeamModelNameTranslator.listing_entries(all_models, llm_router, settings):
entries: Final = TeamModelNameTranslator.listing_entries(all_models, llm_router, settings)
for response_id, lookup_id in entries:
model_info = create_model_info_response(
model_id=lookup_id,
provider="openai",
@ -10280,7 +10288,10 @@ async def model_list(
if wants_anthropic_format:
listing: Final = cast(Sequence[ModelInfoResponse], model_data) # cast-ok: rows built above
return create_anthropic_model_list_response(listing)
return create_anthropic_model_list_response(
listing,
display_names=configured_display_names(entries, llm_router),
)
return dict(
data=model_data,

26
litellm/repositories/user_repository.py
View file

@ -6,15 +6,34 @@ import json
from collections.abc import Mapping
from typing import TYPE_CHECKING, Final
from litellm.models.user import LiteLLM_UserTable
from pydantic import TypeAdapter
from litellm.models.user import LiteLLM_UserTable, SCIMPlaceholder
from litellm.repositories.base_repository import BaseRepository, DbRecord, record_to_dict
from litellm.repositories.prisma_protocols import TableActions
if TYPE_CHECKING:
from prisma import Prisma
from prisma import models as prisma_models
_JSON_ENCODED_COLUMNS: Final = frozenset({"metadata", "model_spend", "model_max_budget"})
_SHADOWING_PLACEHOLDERS_SQL: Final = """
SELECT p.user_id AS placeholder_user_id,
array_agg(r.user_id ORDER BY r.user_id) AS resolved_user_ids,
p.teams AS team_ids
FROM "LiteLLM_UserTable" p
JOIN "LiteLLM_UserTable" r
ON r.user_id <> p.user_id
AND (r.sso_user_id = p.user_id OR LOWER(r.user_email) = LOWER(p.user_id))
WHERE p.sso_user_id IS NULL
AND NOT EXISTS (SELECT 1 FROM "LiteLLM_VerificationToken" k WHERE k.user_id = p.user_id)
GROUP BY p.user_id, p.teams
ORDER BY p.user_id
"""
_PLACEHOLDER_ROWS_ADAPTER: Final = TypeAdapter(tuple[SCIMPlaceholder, ...])
class UserRepository(BaseRepository[LiteLLM_UserTable]):
"""Repository for user database operations."""
@ -59,6 +78,11 @@ class UserRepository(BaseRepository[LiteLLM_UserTable]):
"""Find all users in a team."""
return await self.find_many(where={"teams": {"has": team_id}})
async def find_shadowing_placeholders(self, tx: "Prisma") -> tuple[SCIMPlaceholder, ...]:
"""Users with no SSO id and no virtual keys whose id is another user's SSO id or email."""
rows: Final = await tx.query_raw(_SHADOWING_PLACEHOLDERS_SQL)
return _PLACEHOLDER_ROWS_ADAPTER.validate_python(rows)
async def count_billable_users(self) -> int:
"""Number of users that count toward the license seat limit.

29
litellm/rerank_api/main.py
View file

@ -6,6 +6,7 @@ from typing import Any, Final, Literal
import litellm
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig
from litellm.llms.bedrock.rerank.handler import BedrockRerankHandler
@ -43,10 +44,23 @@ async def arerank(
"""
Async: Reranks a list of documents based on their relevance to the query
"""
_custom_llm_provider: str | None = (
None # rebind-ok: set by the declared-provider guard or the get_llm_provider unpack; read in the except
)
try:
loop: Final = asyncio.get_event_loop()
kwargs["arerank"] = True
declared_provider: Final = declared_authenticating_provider(model, custom_llm_provider)
if declared_provider is not None:
_custom_llm_provider = declared_provider # rebind-ok: see pre-declaration above
else:
_, _custom_llm_provider, _, _ = litellm.get_llm_provider( # rebind-ok: see pre-declaration above
model=model,
custom_llm_provider=custom_llm_provider,
api_base=kwargs.get("api_base", None),
)
func: Final = partial(
rerank,
model,
@ -70,7 +84,11 @@ async def arerank(
response = init_response
return response
except Exception as e:
raise e
raise exception_type(
model=model,
custom_llm_provider=_custom_llm_provider or custom_llm_provider,
original_exception=e,
)
@client
@ -115,6 +133,7 @@ def rerank(
model_info: Final = kwargs.get("model_info", None)
user: Final = kwargs.get("user", None)
client: Final = kwargs.get("client", None)
_custom_llm_provider: str | None = None # rebind-ok: set by the get_llm_provider unpack; read in the except
try:
_is_async: Final = kwargs.pop("arerank", False) is True
optional_params: Final = GenericLiteLLMParams(**kwargs)
@ -127,7 +146,7 @@ def rerank(
(
model,
_custom_llm_provider,
_custom_llm_provider, # rebind-ok: see pre-declaration above
dynamic_api_key,
dynamic_api_base,
) = litellm.get_llm_provider(
@ -538,4 +557,8 @@ def rerank(
return response
except Exception as e:
verbose_logger.error("Error in rerank: %s", e)
raise exception_type(model=model, custom_llm_provider=custom_llm_provider, original_exception=e)
raise exception_type(
model=model,
custom_llm_provider=_custom_llm_provider or custom_llm_provider,
original_exception=e,
)

10
litellm/responses/litellm_completion_transformation/streaming_iterator.py
View file

@ -1169,16 +1169,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
def _emit_response_completed_event(self, litellm_model_response: ModelResponse) -> ResponseCompletedEvent | None:
if litellm_model_response:
# Add cost to usage object if include_cost_in_streaming_usage is True
if litellm.include_cost_in_streaming_usage and self.litellm_logging_obj is not None:
usage: Final[object] = getattr(litellm_model_response, "usage", None)
if usage is not None:
setattr(
usage,
"cost",
self.litellm_logging_obj._response_cost_calculator(result=litellm_model_response),
)
# Transform the response
responses_api_response: Final = (
LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response(

54
litellm/responses/streaming_iterator.py
View file

@ -407,23 +407,7 @@ class BaseResponsesAPIStreamingIterator:
openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED,
):
self.completed_response = openai_responses_api_chunk
# Add cost to usage object if include_cost_in_streaming_usage is True
if litellm.include_cost_in_streaming_usage and self.logging_obj is not None:
response_obj: Final[ResponsesAPIResponse | None] = getattr(
openai_responses_api_chunk, "response", None
)
if response_obj:
usage_obj: Final[ResponseAPIUsage | None] = getattr(response_obj, "usage", None)
if usage_obj is not None:
try:
cost: Final[float | None] = self.logging_obj._response_cost_calculator(
result=response_obj
)
if cost is not None:
setattr(usage_obj, "cost", cost)
except Exception:
# Best-effort usage cost annotation should not break stream replay.
pass
_stamp_responses_usage_cost(getattr(openai_responses_api_chunk, "response", None), self.logging_obj)
if _chunk_type == openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED:
self._handle_logging_failed_response()
@ -1023,7 +1007,7 @@ class MockResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
transformed: ResponsesAPIResponse,
logging_obj: LiteLLMLoggingObj,
) -> None:
self._events: list[ResponsesAPIStreamingResponse] = _build_synthetic_response_events(
self._events: Sequence[ResponsesAPIStreamingResponse] = build_synthetic_response_events(
transformed=transformed,
logging_obj=logging_obj,
chunk_size=self.CHUNK_SIZE,
@ -1090,7 +1074,7 @@ class CachedResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
transformed: ResponsesAPIResponse,
logging_obj: LiteLLMLoggingObj,
) -> None:
self._events = _build_synthetic_response_events(
self._events = build_synthetic_response_events(
transformed=transformed,
logging_obj=logging_obj,
chunk_size=MockResponsesAPIStreamingIterator.CHUNK_SIZE,
@ -1274,22 +1258,32 @@ def _add_text_like_part_events(
)
def _build_synthetic_response_events(
def _stamp_responses_usage_cost(
response_obj: ResponsesAPIResponse | None, logging_obj: LiteLLMLoggingObj | None
) -> None:
if response_obj is None or logging_obj is None:
return
usage_obj: Final[ResponseAPIUsage | None] = getattr(response_obj, "usage", None)
if usage_obj is None:
return
if isinstance(getattr(usage_obj, "cost", None), (int, float)):
return
try:
cost: Final[float | None] = logging_obj._response_cost_calculator(result=response_obj)
except Exception:
return
if isinstance(cost, (int, float)) and cost > 0:
setattr(usage_obj, "cost", cost)
def build_synthetic_response_events(
*,
transformed: ResponsesAPIResponse,
logging_obj: LiteLLMLoggingObj,
logging_obj: LiteLLMLoggingObj | None,
chunk_size: int,
) -> list[ResponsesAPIStreamingResponse]:
openai_types: Final = _get_openai_response_types()
if litellm.include_cost_in_streaming_usage and logging_obj is not None:
usage_obj: Final = transformed.usage if hasattr(transformed, "usage") else None
if usage_obj is not None:
try:
cost: Final[float | None] = logging_obj._response_cost_calculator(result=transformed)
if cost is not None:
setattr(usage_obj, "cost", cost)
except Exception:
pass
_stamp_responses_usage_cost(transformed, logging_obj)
events: Final[list[ResponsesAPIStreamingResponse]] = [
_build_response_status_event(openai_types.ResponsesAPIStreamEvents.RESPONSE_CREATED, transformed),

46
litellm/router.py
View file

@ -22,7 +22,7 @@ import traceback
import weakref
from collections import defaultdict
from collections.abc import AsyncGenerator, AsyncIterator, Callable, Generator, Mapping, Sequence
from functools import lru_cache
from functools import lru_cache, partial
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias, TypeVar, Union, cast
@ -831,6 +831,7 @@ class Router:
self._zero_cost_cache: dict[str, bool] = {}
self._routing_group_rows: tuple[DeploymentTypedDict, ...] | None = None
self._init_routing_groups(None)
self._provider_unresolved_deployments: tuple[Callable[[], Deployment | None], ...] = ()
self.deployment_affinity_ttl_seconds = deployment_affinity_ttl_seconds
self.model_group_affinity_config = model_group_affinity_config
@ -8478,6 +8479,19 @@ class Router:
return deployment
except Exception as e:
if self.ignore_invalid_deployments:
if isinstance(e, litellm.BadRequestError):
self._provider_unresolved_deployments = (
*self._provider_unresolved_deployments,
partial(
self._create_deployment,
deployment_info=deployment_info,
_model_name=_model_name,
_litellm_params=_litellm_params,
_model_info=_model_info,
declared_id=declared_id,
duplicate_ids=duplicate_ids,
),
)
verbose_router_logger.exception(
"Error creating deployment: %s, ignoring and continuing with other deployments.", e
)
@ -8907,6 +8921,7 @@ class Router:
self.quality_routers = {}
self.complexity_routers = {}
self.auto_routers = {}
self._provider_unresolved_deployments = ()
self._invalidate_model_group_info_cache()
self._invalidate_access_groups_cache()
# we add api_base/api_key each model so load balancing between azure/gpt on api_base1 and api_base2 works
@ -9529,8 +9544,12 @@ class Router:
"""Re-assert this router's deployments onto a freshly fetched catalog.
Reads ``model_list`` at call time, so only deployments the router still
serves are restored.
serves are restored, plus any config deployment the fresh catalog now resolves.
"""
provider_unresolved: Final = self._provider_unresolved_deployments
self._provider_unresolved_deployments = ()
for create_deployment in provider_unresolved:
create_deployment()
for entry in tuple(self.model_list):
try:
deployment = entry if isinstance(entry, Deployment) else Deployment(**entry)
@ -9733,6 +9752,26 @@ class Router:
coerce_token_limit(model_info.get("max_output_tokens")),
)
def get_configured_display_name(self, model_name: str) -> "str | None":
"""
Return the display_name explicitly configured in a concrete deployment's
model_info for model_name, via O(1) index lookup.
Returns None for wildcard-expanded or unknown names, and treats a
non-string or empty configured value as absent rather than failing the
listing. Like get_configured_token_limits, this never triggers pattern
matching or deep copies, so it is safe to call per listed model on the
/v1/models hot path.
"""
deployment: Final = self.get_deployment_by_model_group_name(model_group_name=model_name)
if deployment is None:
return None
display_name: Final = deployment.model_info.get("display_name")
if isinstance(display_name, str) and display_name.strip():
return display_name
return None
def get_deployment_credentials_with_provider(
self, model_id: str, team_id: str | None = None
) -> dict[str, Any] | None:
@ -12587,8 +12626,7 @@ class Router:
await self._claude_code_session_router_cache.async_delete_cache(key=cache_key)
except Exception as e: # noqa: BLE001 # cache cleanup must not fail an otherwise routable request
verbose_router_logger.warning(
"Failed to delete Claude Code session router binding; "
"the binding may remain until its TTL expires: %s",
"Failed to delete Claude Code session router binding; the binding may remain until its TTL expires: %s",
e,
)

12
litellm/router_utils/search_api_router.py
View file

@ -9,6 +9,7 @@ import random
import traceback
from collections.abc import Callable
from functools import partial
from types import MappingProxyType
from typing import Any, Final
from litellm._logging import verbose_router_logger
@ -214,6 +215,15 @@ class SearchAPIRouter:
api_key, api_base = SearchAPIRouter._resolve_search_provider_credentials(
tool_litellm_params=litellm_params,
)
protected_params: Final = frozenset(("search_provider", "api_key", "api_base"))
search_params: Final = MappingProxyType(
{
key: value
for params in (litellm_params, kwargs)
for key, value in params.items()
if key not in protected_params and value is not None
}
)
verbose_router_logger.debug("Selected search tool with provider: %s", search_provider)
@ -222,7 +232,7 @@ class SearchAPIRouter:
search_provider=search_provider,
api_key=api_key,
api_base=api_base,
**kwargs,
**search_params,
)
return response

35
litellm/search/cost_calculator.py
View file

@ -2,16 +2,37 @@
Cost calculation for search providers.
"""
from collections.abc import Mapping
from types import MappingProxyType
from typing import Final
from pydantic import TypeAdapter, ValidationError
from litellm.utils import get_model_info
PROVIDER_USAGE_ADAPTER: Final[TypeAdapter[tuple[Mapping[str, object], ...]]] = TypeAdapter(
tuple[Mapping[str, object], ...]
)
EMPTY_OPTIONAL_PARAMS: Final[Mapping[str, object]] = MappingProxyType({})
def _provider_usage(
optional_params: Mapping[str, object] | None,
usage_param: str,
) -> tuple[Mapping[str, object], ...] | None:
params: Final = optional_params if optional_params is not None else EMPTY_OPTIONAL_PARAMS
raw_usage: Final[object] = params.get(usage_param)
try:
return PROVIDER_USAGE_ADAPTER.validate_python(raw_usage)
except ValidationError:
return None
def search_provider_cost_per_query(
model: str,
custom_llm_provider: str | None = None,
number_of_queries: int = 1,
optional_params: dict | None = None,
optional_params: Mapping[str, object] | None = None,
) -> tuple[float, float]:
"""
Calculate cost for search-only providers.
@ -28,6 +49,18 @@ def search_provider_cost_per_query(
Returns:
Tuple of (input_cost, output_cost) where output_cost is always 0.0
"""
if custom_llm_provider == "parallel_ai":
from litellm.llms.parallel_ai.search.cost_calculator import (
PARALLEL_AI_USAGE_PARAM,
parallel_ai_search_cost,
)
input_cost: Final = parallel_ai_search_cost(
optional_params=optional_params if optional_params is not None else EMPTY_OPTIONAL_PARAMS,
usage=_provider_usage(optional_params, PARALLEL_AI_USAGE_PARAM),
)
return (input_cost, 0.0)
model_info: Final = get_model_info(model=model, custom_llm_provider=custom_llm_provider)
# Check for tiered pricing (e.g., Exa AI based on max_results)

49
litellm/types/integrations/datadog_llm_obs.py
View file

@ -4,21 +4,58 @@ Payloads for Datadog LLM Observability Service (LLMObs)
API Reference: https://docs.datadoghq.com/llm_observability/setup/api/?tab=example#api-standards
"""
from collections.abc import Sequence
from typing import Any, Literal
from typing_extensions import TypedDict
from typing_extensions import ReadOnly, TypedDict
from litellm.types.integrations.custom_logger import StandardCustomLoggerInitParams
class ToolCall(TypedDict, total=False):
"""A tool call on a message, as LLM Obs names its fields."""
name: ReadOnly[str]
arguments: ReadOnly[dict[str, Any] | str] # parsed object, or the raw string when it will not parse to one
tool_id: ReadOnly[str]
type: ReadOnly[str]
class ToolResult(TypedDict, total=False):
"""The result of a tool call, as LLM Obs names its fields."""
name: ReadOnly[str]
result: ReadOnly[str]
tool_id: ReadOnly[str]
type: ReadOnly[str]
class ToolDefinition(TypedDict, total=False):
"""A tool the model was offered on the request."""
name: ReadOnly[str]
description: ReadOnly[str]
schema: ReadOnly[dict[str, Any]]
class Message(TypedDict, total=False):
"""A message on a span, as LLM Obs names its fields."""
content: ReadOnly[str]
role: ReadOnly[str]
reasoning_content: ReadOnly[str]
tool_calls: ReadOnly[Sequence[ToolCall]]
tool_results: ReadOnly[Sequence[ToolResult]]
class InputMeta(TypedDict):
messages: list[
dict[str, Any] # changed to fit with tool calls
messages: Sequence[
Message | dict[str, Any] # changed to fit with tool calls
] # Relevant Issue: https://github.com/BerriAI/litellm/issues/9494
class OutputMeta(TypedDict):
messages: list[Any]
messages: Sequence[Any]
class DDLLMObsError(TypedDict, total=False):
@ -36,6 +73,7 @@ class Meta(TypedDict, total=False):
output: OutputMeta # The span's output information.
metadata: dict[str, Any]
error: DDLLMObsError | None # Error information on the span
tool_definitions: ReadOnly[Sequence[ToolDefinition]] # The tools offered to the model on this request
class LLMMetrics(TypedDict, total=False):
@ -45,6 +83,9 @@ class LLMMetrics(TypedDict, total=False):
time_to_first_token: float
time_per_output_token: float
total_cost: float
cache_read_input_tokens: ReadOnly[float]
cache_write_input_tokens: ReadOnly[float]
non_cached_input_tokens: ReadOnly[float]
class LLMObsPayload(TypedDict, total=False):

20
litellm/types/integrations/prometheus.py
View file

@ -270,6 +270,10 @@ DEFINED_PROMETHEUS_METRICS = Literal[
"litellm_deployment_rpm_limit",
"litellm_remaining_api_key_requests_for_model",
"litellm_remaining_api_key_tokens_for_model",
"litellm_api_key_rate_limit_allowed_metric",
"litellm_api_key_rate_limit_used_metric",
"litellm_team_rate_limit_allowed_metric",
"litellm_team_rate_limit_used_metric",
"litellm_llm_api_failed_requests_metric",
"litellm_callback_logging_failures_metric",
"litellm_in_flight_requests",
@ -775,6 +779,22 @@ class PrometheusMetricLabels:
UserAPIKeyLabelNames.MODEL_ID.value,
]
litellm_api_key_rate_limit_allowed_metric: ClassVar[tuple[str, ...]] = (
UserAPIKeyLabelNames.API_KEY_HASH.value,
UserAPIKeyLabelNames.API_KEY_ALIAS.value,
UserAPIKeyLabelNames.RATE_LIMIT_TYPE.value,
)
litellm_api_key_rate_limit_used_metric = litellm_api_key_rate_limit_allowed_metric
litellm_team_rate_limit_allowed_metric: ClassVar[tuple[str, ...]] = (
UserAPIKeyLabelNames.TEAM.value,
UserAPIKeyLabelNames.TEAM_ALIAS.value,
UserAPIKeyLabelNames.RATE_LIMIT_TYPE.value,
)
litellm_team_rate_limit_used_metric = litellm_team_rate_limit_allowed_metric
litellm_llm_api_failed_requests_metric = [
UserAPIKeyLabelNames.END_USER.value,
UserAPIKeyLabelNames.API_KEY_HASH.value,

6
litellm/types/proxy/management_endpoints/scim_v2.py
View file

@ -150,6 +150,12 @@ class SCIMGroup(SCIMResource):
members: list[SCIMMember] | None = None
class SCIMPlaceholderMergeResult(BaseModel):
placeholder_user_id: str
merged_into_user_id: str
team_ids: tuple[str, ...]
# SCIM List Response Models
class SCIMListResponse(BaseModel):
schemas: list[str] = ["urn:ietf:params:scim:api:messages:2.0:ListResponse"]

2
litellm/types/utils.py
View file

@ -3636,6 +3636,8 @@ all_litellm_params = (
"client",
"rpm",
"tpm",
"default_api_key_rpm_limit",
"default_api_key_tpm_limit",
"itpm",
"otpm",
"max_parallel_requests",

1005
model_prices_and_context_window.json
View file

File diff suppressed because it is too large Load diff

45
tests/e2e/test_junit_properties.py
View file

@ -24,25 +24,10 @@ from junit_properties import (
)
class FakeMarker:
def __init__(self, name: str, *args: object) -> None:
self.name = name
self.args = args
class FakeItem:
"""The three attributes junit_properties reads off a pytest Item."""
def __init__(
self, nodeid: str, location: tuple[str, int | None, str], markers: tuple[FakeMarker, ...] = ()
) -> None:
self.nodeid = nodeid
self.location = location
self.user_properties: list[tuple[str, str]] = []
self._markers = markers
def iter_markers(self, name: str):
return (marker for marker in self._markers if marker.name == name)
def collected_item(request: pytest.FixtureRequest, name: str) -> pytest.Item:
"""The Item pytest collected for test ``name`` in this file: the real nodeid,
location and marker machinery the collection hook reads, as pytest built it."""
return next(item for item in request.session.items if item.path == request.path and item.name == name)
def repo_root() -> Path | None:
@ -109,22 +94,22 @@ class TestSourceFromLocation:
class TestResultProperties:
def test_every_test_carries_package_covers_and_source(self) -> None:
item = FakeItem(
"logging/test_x.py::TestFoo::test_bar",
("logging/test_x.py", 40, "TestFoo.test_bar"),
(FakeMarker("covers", "LOG-1", "LOG-2"),),
)
assert result_properties(item) == (
("package", "logging"),
def test_every_test_carries_package_covers_and_source(self, request: pytest.FixtureRequest) -> None:
"""Read off this test's own collected Item, so the nodeid and location are
whatever pytest reports for the launch shape in use, and the marker is added
at run time so the coverage registry's collect-only pass never sees it."""
test = type(self).test_every_test_carries_package_covers_and_source
request.applymarker(pytest.mark.covers("LOG-1", "LOG-2"))
assert result_properties(collected_item(request, test.__name__)) == (
("package", "root"),
("covers", "LOG-1,LOG-2"),
("source", "tests/e2e/logging/test_x.py:41"),
("source", f"tests/e2e/test_junit_properties.py:{test.__code__.co_firstlineno}"),
)
def test_attach_is_idempotent(self) -> None:
def test_attach_is_idempotent(self, request: pytest.FixtureRequest) -> None:
"""Collection can run the hook more than once; a second pass must not
double the <property> entries in the report."""
item = FakeItem("logging/test_x.py::test_bar", ("logging/test_x.py", 40, "test_bar"))
item = collected_item(request, type(self).test_attach_is_idempotent.__name__)
attach_result_properties(item)
attach_result_properties(item)
assert [name for name, _ in item.user_properties] == ["package", "covers", "source"]

2
tests/llm_responses_api_testing/test_responses_hooks.py
View file

@ -841,7 +841,7 @@ def test_build_synthetic_response_events_covers_annotations_function_calls_and_r
)
try:
events = streaming_module._build_synthetic_response_events(
events = streaming_module.build_synthetic_response_events(
transformed=transformed,
logging_obj=logging_obj,
chunk_size=5,

20
tests/ocr_tests/test_ocr_vertex_ai.py
View file

@ -5,9 +5,11 @@ Note: Vertex AI OCR automatically converts URLs to base64 data URIs since
the Vertex AI endpoint doesn't have internet access.
"""
import os
import json
import os
import tempfile
from typing import Final
import pytest
from base_ocr_unit_tests import BaseOCRTest
@ -139,3 +141,19 @@ def test_vertex_ai_ocr_routing():
assert isinstance(
deepseek_variant, VertexAIDeepSeekOCRConfig
), "DeepSeek variant should route to VertexAIDeepSeekOCRConfig"
@pytest.mark.parametrize("model", ("deepseek-ocr-maas", "deepseek-ai/deepseek-ocr-maas"))
def test_deepseek_request_uses_single_provider_namespace(model: str) -> None:
from litellm.llms.vertex_ai.ocr.deepseek_transformation import (
VertexAIDeepSeekOCRConfig,
)
request: Final = VertexAIDeepSeekOCRConfig().transform_ocr_request(
model=model,
document={"type": "image_url", "image_url": "data:image/png;base64,AA=="},
optional_params={},
headers={},
)
assert request.data["model"] == "deepseek-ai/deepseek-ocr-maas"

58
tests/proxy_migration_tests/test_image_bedrock_realtime_extra.py Normal file
View file

@ -0,0 +1,58 @@
"""Image-level check that the built proxy image can import the Bedrock realtime SDK.
Bedrock Nova Sonic (`/v1/realtime`) imports `aws_sdk_bedrock_runtime` lazily on the
first session, so an image whose `uv sync` stages skip the `bedrock-realtime` extra
boots, passes health checks, and then fails every Nova Sonic session with
"Missing aws_sdk_bedrock_runtime". Importing inside the built image is what catches
that class of regression (missing extra, lockfile drift, a stage that syncs a
different set of extras), which a static Dockerfile check cannot.
Gated on LITELLM_IMAGE like the other image checks in this directory; exercised
where an image has been built (the image-scan workflow). Requires a working docker CLI.
"""
import os
import shutil
import subprocess
from typing import Final
import pytest
IMAGE: Final = os.getenv("LITELLM_IMAGE")
NON_ROOT_UID: Final = "12345:0"
IMPORT_PROBE: Final = "import aws_sdk_bedrock_runtime, smithy_aws_core; print('bedrock-realtime ok')"
pytestmark = [
pytest.mark.skipif(IMAGE is None, reason="requires a built image (set LITELLM_IMAGE)"),
pytest.mark.skipif(shutil.which("docker") is None, reason="requires the docker CLI"),
]
def test_image_imports_bedrock_realtime_sdk():
assert IMAGE is not None
probe: Final = subprocess.run(
[
"docker",
"run",
"--rm",
"--network",
"none",
"--user",
NON_ROOT_UID,
"--entrypoint",
"python",
IMAGE,
"-c",
IMPORT_PROBE,
],
capture_output=True,
text=True,
check=False,
)
assert probe.returncode == 0 and "bedrock-realtime ok" in probe.stdout, (
f"{IMAGE} cannot import aws_sdk_bedrock_runtime as uid {NON_ROOT_UID}, so Bedrock Nova Sonic "
"/v1/realtime sessions fail with 'Missing aws_sdk_bedrock_runtime'. Is `--extra bedrock-realtime` "
f"passed to every `uv sync` in its Dockerfile?\nstdout:\n{probe.stdout}\nstderr:\n{probe.stderr}"
)

3
tests/router_unit_tests/test_router_helper_utils.py
View file

@ -2294,6 +2294,7 @@ def search_tools():
"search_provider": "perplexity",
"api_key": "test-api-key",
"api_base": "https://api.perplexity.ai",
"mode": "turbo",
},
},
{
@ -2302,6 +2303,7 @@ def search_tools():
"search_provider": "perplexity",
"api_key": "test-api-key-2",
"api_base": "https://api.perplexity.ai",
"mode": "turbo",
},
},
]
@ -2393,6 +2395,7 @@ async def test_asearch_with_fallbacks_helper(search_tools):
assert "search_provider" in kwargs
assert kwargs["search_provider"] == "perplexity"
assert "api_key" in kwargs
assert kwargs["mode"] == "turbo"
assert kwargs["query"] == "helper test query"
return mock_response

469
tests/test_litellm/integrations/datadog/test_datadog_llm_obs.py Normal file
View file

@ -0,0 +1,469 @@
"""
Regression tests for the Datadog LLM Observability payload schema (issue #35786).
Datadog renders tool calls, tool results and prompt-cache savings only from the fields its
own schema names. These assert on the payload `create_llm_obs_payload` actually hands the
intake, so a regression that moves data back into `meta.metadata` fails here.
Fixtures mirror what a live proxy run recorded on the callback, including the provider
spelling of prompt-cache counts (`prompt_tokens_details.cached_tokens`).
"""
import json
import os
from datetime import datetime, timedelta
from typing import Any
from unittest.mock import patch
import pytest
from litellm.integrations.datadog.datadog_llm_obs import DataDogLLMObsLogger
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
TOOL_DEFINITION: dict[str, Any] = {
"type": "function",
"function": {
"name": "get_weather",
"description": "Get current weather for a city",
"parameters": {
"type": "object",
"properties": {"city": {"type": "string"}},
"required": ["city"],
},
},
}
ASSISTANT_TOOL_CALL: dict[str, Any] = {
"id": "call_abc123",
"type": "function",
"function": {"name": "get_weather", "arguments": '{"city":"Paris","unit":"c"}'},
}
@pytest.fixture
def logger() -> DataDogLLMObsLogger:
with patch.dict(os.environ, {"DD_API_KEY": "k", "DD_SITE": "us5.datadoghq.com"}, clear=True):
with patch("asyncio.create_task"):
return DataDogLLMObsLogger()
NOT_GIVEN: Any = object()
def build_payload(
messages: Any = NOT_GIVEN,
response_message: dict[str, Any] | None = None,
usage_object: dict[str, Any] | None = None,
model_parameters: dict[str, Any] | None = None,
prompt_tokens: int = 4447,
) -> dict[str, Any]:
return {
"standard_logging_object": {
"call_type": "acompletion",
"messages": [{"role": "user", "content": "hi"}] if messages is NOT_GIVEN else messages,
"response": {"choices": [{"message": response_message or {"role": "assistant", "content": "hello"}}]},
"model_parameters": model_parameters or {},
"metadata": {"usage_object": usage_object} if usage_object is not None else {},
"prompt_tokens": prompt_tokens,
"completion_tokens": 507,
"total_tokens": prompt_tokens + 507,
"response_cost": 0.02,
"status": "success",
},
"litellm_params": {"metadata": {}},
}
def build(logger: DataDogLLMObsLogger, **kwargs: Any) -> dict[str, Any]:
"""Build a span and read it back as the JSON the intake receives, not as Python objects."""
start = datetime(2026, 9, 1, 12, 0, 0)
payload = logger.create_llm_obs_payload(build_payload(**kwargs), start, start + timedelta(seconds=2))
return json.loads(safe_dumps(payload))
def test_output_tool_calls_use_the_datadog_tool_call_schema(logger: DataDogLLMObsLogger) -> None:
"""Datadog reads name/arguments/tool_id off the tool call; OpenAI nests them under `function`."""
payload = build(
logger,
response_message={"role": "assistant", "content": None, "tool_calls": [ASSISTANT_TOOL_CALL]},
)
message = payload["meta"]["output"]["messages"][0]
assert message["tool_calls"] == [
{
"name": "get_weather",
"arguments": {"city": "Paris", "unit": "c"},
"tool_id": "call_abc123",
"type": "function",
}
]
assert "function" not in message["tool_calls"][0]
def test_tool_calls_are_not_duplicated_into_metadata(logger: DataDogLLMObsLogger) -> None:
"""The flat `output_tool_calls.*` keys were a second copy of a fact that now has its own field."""
payload = build(
logger,
response_message={"role": "assistant", "content": None, "tool_calls": [ASSISTANT_TOOL_CALL]},
)
assert [key for key in payload["meta"]["metadata"] if "tool_calls." in key] == []
def test_tool_result_message_links_back_to_its_tool_call(logger: DataDogLLMObsLogger) -> None:
"""Datadog pairs a result with its call through tool_id, and names the tool from the call."""
payload = build(
logger,
messages=[
{"role": "user", "content": "Weather in Paris?"},
{"role": "assistant", "content": None, "tool_calls": [ASSISTANT_TOOL_CALL]},
{"role": "tool", "tool_call_id": "call_abc123", "content": '{"temp_c": 18}'},
],
)
tool_message = payload["meta"]["input"]["messages"][2]
assert tool_message["tool_results"] == [
{"name": "get_weather", "result": '{"temp_c": 18}', "tool_id": "call_abc123", "type": "function"}
]
def test_tool_result_without_a_matching_call_still_reports_its_id(logger: DataDogLLMObsLogger) -> None:
"""A truncated conversation loses the call, so the name is unknown but the link must survive."""
payload = build(
logger,
messages=[{"role": "tool", "tool_call_id": "call_orphan", "content": "42"}],
)
assert payload["meta"]["input"]["messages"][0]["tool_results"] == [
{"name": "", "result": "42", "tool_id": "call_orphan", "type": "function"}
]
def test_cache_tokens_are_reported_as_span_metrics(logger: DataDogLLMObsLogger) -> None:
"""
Datadog charts cache savings from span metrics; nested usage_object is not read for it.
litellm's normalized prompt count includes both cache categories, so the three cache
metrics must partition input_tokens: read + write + non_cached == input.
"""
payload = build(
logger,
usage_object={"prompt_tokens_details": {"cached_tokens": 4300, "cache_write_tokens": 95}},
)
metrics = payload["metrics"]
assert metrics["cache_read_input_tokens"] == 4300.0
assert metrics["cache_write_input_tokens"] == 95.0
assert metrics["non_cached_input_tokens"] == 4447.0 - 4300.0 - 95.0
assert (
metrics["cache_read_input_tokens"] + metrics["cache_write_input_tokens"] + metrics["non_cached_input_tokens"]
== metrics["input_tokens"]
)
def test_cache_write_tokens_are_not_counted_as_non_cached(logger: DataDogLLMObsLogger) -> None:
"""A cache-priming request must not report its primed prefix as full-price uncached input."""
payload = build(logger, usage_object={"prompt_tokens_details": {"cache_write_tokens": 4000}})
assert payload["metrics"]["cache_write_input_tokens"] == 4000.0
assert payload["metrics"]["non_cached_input_tokens"] == 4447.0 - 4000.0
assert "cache_read_input_tokens" not in payload["metrics"]
def test_a_fully_cached_request_reports_a_zero_non_cached_count(logger: DataDogLLMObsLogger) -> None:
"""Zero residual is real data: everything was served from cache. Inconsistent counts clamp to it."""
payload = build(
logger,
usage_object={"prompt_tokens_details": {"cached_tokens": 4352, "cache_write_tokens": 95}},
)
assert payload["metrics"]["non_cached_input_tokens"] == 0.0
def test_anthropic_top_level_cache_keys_are_read(logger: DataDogLLMObsLogger) -> None:
"""A raw Anthropic usage dict records the counts top level, not under prompt_tokens_details."""
payload = build(
logger,
usage_object={"cache_read_input_tokens": 4300, "cache_creation_input_tokens": 95},
)
metrics = payload["metrics"]
assert metrics["cache_read_input_tokens"] == 4300.0
assert metrics["cache_write_input_tokens"] == 95.0
assert metrics["non_cached_input_tokens"] == 4447.0 - 4300.0 - 95.0
def test_cache_metrics_come_from_the_normalized_field_not_the_anthropic_one(logger: DataDogLLMObsLogger) -> None:
"""
litellm normalizes every provider's cache counters into prompt_tokens_details.
A real cached request from a non-Anthropic provider carries only `cached_tokens`, so
reading the Anthropic-specific `cache_read_input_tokens` key reports nothing for it.
"""
payload = build(
logger,
usage_object={"prompt_tokens_details": {"audio_tokens": None, "cached_tokens": 4096}},
prompt_tokens=4335,
)
assert payload["metrics"]["cache_read_input_tokens"] == 4096.0
assert payload["metrics"]["non_cached_input_tokens"] == 4335.0 - 4096.0
@pytest.mark.parametrize(
"usage_object",
[
{"prompt_tokens_details": {"cache_write_tokens": 95}},
{"prompt_tokens_details": {"cache_creation_tokens": 95}},
{"cache_creation_input_tokens": 95},
],
)
def test_every_spelling_of_cache_write_tokens_is_read(
logger: DataDogLLMObsLogger, usage_object: dict[str, Any]
) -> None:
"""A raw usage dict that bypassed litellm's normalizer can carry any provider's spelling."""
payload = build(logger, usage_object=usage_object)
assert payload["metrics"]["cache_write_input_tokens"] == 95.0
def test_a_cache_read_does_not_emit_a_zero_cache_write(logger: DataDogLLMObsLogger) -> None:
"""A zero write on every cache-read span would drag Datadog's cache-write average to nothing."""
payload = build(logger, usage_object={"prompt_tokens_details": {"cached_tokens": 4096}})
assert payload["metrics"]["cache_read_input_tokens"] == 4096.0
assert "cache_write_input_tokens" not in payload["metrics"]
def test_no_cache_keys_when_the_provider_reports_no_caching(logger: DataDogLLMObsLogger) -> None:
"""An uncached request must not gain zero-valued cache metrics that dilute cache dashboards."""
payload = build(logger, usage_object={"prompt_tokens_details": None})
assert "cache_read_input_tokens" not in payload["metrics"]
assert "cache_write_input_tokens" not in payload["metrics"]
assert "non_cached_input_tokens" not in payload["metrics"]
def test_tool_definitions_are_sent_on_meta(logger: DataDogLLMObsLogger) -> None:
payload = build(logger, model_parameters={"tools": [TOOL_DEFINITION]})
assert payload["meta"]["tool_definitions"] == [
{
"name": "get_weather",
"description": "Get current weather for a city",
"schema": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]},
}
]
def test_tool_definitions_accept_the_bare_anthropic_shape(logger: DataDogLLMObsLogger) -> None:
"""The Anthropic surface declares tools unwrapped, with input_schema instead of parameters."""
payload = build(
logger,
model_parameters={"tools": [{"name": "get_weather", "description": "d", "input_schema": {"type": "object"}}]},
)
assert payload["meta"]["tool_definitions"] == [
{"name": "get_weather", "description": "d", "schema": {"type": "object"}}
]
def test_meta_omits_tool_definitions_when_no_tools_were_offered(logger: DataDogLLMObsLogger) -> None:
assert "tool_definitions" not in build(logger)["meta"]
def test_unparseable_tool_arguments_are_preserved_rather_than_dropped(logger: DataDogLLMObsLogger) -> None:
"""A truncated argument string is still the only record of what the model tried to call."""
payload = build(
logger,
response_message={
"role": "assistant",
"content": None,
"tool_calls": [{"id": "c1", "type": "function", "function": {"name": "f", "arguments": '{"city":'}}],
},
)
assert payload["meta"]["output"]["messages"][0]["tool_calls"][0]["arguments"] == '{"city":'
def test_oversized_tool_arguments_ship_unparsed(logger: DataDogLLMObsLogger) -> None:
"""
Decoding attacker-sized compact JSON multiplies memory for a span that is only logging.
This payload is perfectly valid JSON, so the only reason it arrives as a string is the
size bound; a smaller copy of the same shape comes back as an object below.
"""
oversized = '{"a":"' + "x" * 300_000 + '"}'
payload = build(
logger,
response_message={
"role": "assistant",
"content": None,
"tool_calls": [{"id": "c1", "type": "function", "function": {"name": "f", "arguments": oversized}}],
},
)
assert payload["meta"]["output"]["messages"][0]["tool_calls"][0]["arguments"] == oversized
def test_valid_arguments_below_the_bound_still_parse(logger: DataDogLLMObsLogger) -> None:
"""The size bound must not swallow ordinary arguments; this is the oversized test's control."""
payload = build(
logger,
response_message={
"role": "assistant",
"content": None,
"tool_calls": [
{"id": "c1", "type": "function", "function": {"name": "f", "arguments": '{"a":"' + "x" * 64 + '"}'}}
],
},
)
assert payload["meta"]["output"]["messages"][0]["tool_calls"][0]["arguments"] == {"a": "x" * 64}
def test_a_result_is_named_even_when_its_call_had_unparseable_arguments(logger: DataDogLLMObsLogger) -> None:
"""Correlating a result to its call reads ids and names, so bad arguments cannot break linking."""
payload = build(
logger,
messages=[
{
"role": "assistant",
"content": None,
"tool_calls": [
{"id": "call_abc123", "type": "function", "function": {"name": "get_weather", "arguments": "{"}}
],
},
{"role": "tool", "tool_call_id": "call_abc123", "content": "18C"},
],
)
assert payload["meta"]["input"]["messages"][1]["tool_results"] == [
{"name": "get_weather", "result": "18C", "tool_id": "call_abc123", "type": "function"}
]
def test_deeply_nested_tool_arguments_do_not_drop_the_span(logger: DataDogLLMObsLogger) -> None:
"""json.loads raises RecursionError, not JSONDecodeError, on hostile nesting."""
payload = build(
logger,
response_message={
"role": "assistant",
"content": None,
"tool_calls": [{"id": "c1", "type": "function", "function": {"name": "f", "arguments": "[" * 50_000}}],
},
)
assert payload["meta"]["output"]["messages"][0]["tool_calls"][0]["arguments"] == "[" * 50_000
def test_tool_arguments_that_parse_to_a_non_object_stay_a_string(logger: DataDogLLMObsLogger) -> None:
"""Datadog types arguments as an object, so a bare JSON scalar must not land there as one."""
payload = build(
logger,
response_message={
"role": "assistant",
"content": None,
"tool_calls": [{"id": "c1", "type": "function", "function": {"name": "f", "arguments": "42"}}],
},
)
assert payload["meta"]["output"]["messages"][0]["tool_calls"][0]["arguments"] == "42"
def test_a_tool_without_a_name_is_not_offered_as_a_definition(logger: DataDogLLMObsLogger) -> None:
"""A nameless tool cannot be matched to a call, so it is dropped rather than sent blank."""
payload = build(logger, model_parameters={"tools": [{"function": {"description": "no name"}}, TOOL_DEFINITION]})
assert [tool["name"] for tool in payload["meta"]["tool_definitions"]] == ["get_weather"]
def test_a_tool_definition_without_a_schema_omits_the_field(logger: DataDogLLMObsLogger) -> None:
"""An empty schema object would read as a tool that takes no arguments, which is a different claim."""
payload = build(logger, model_parameters={"tools": [{"name": "ping", "description": "d"}]})
assert payload["meta"]["tool_definitions"] == [{"name": "ping", "description": "d"}]
def test_a_non_dict_message_still_reaches_datadog(logger: DataDogLLMObsLogger) -> None:
"""Callers can log arbitrary message payloads, and dropping the span over one loses the request."""
payload = build(logger, messages=["just a bare string"])
assert payload["meta"]["input"]["messages"] == [{"input": "just a bare string"}]
def test_messages_logged_as_a_bare_string_still_reach_datadog(logger: DataDogLLMObsLogger) -> None:
payload = build(logger, messages="the whole prompt as one string")
assert payload["meta"]["input"]["messages"] == [{"input": "the whole prompt as one string"}]
def test_non_chat_call_types_log_an_empty_input(logger: DataDogLLMObsLogger) -> None:
"""Embedding and image calls carry no messages; fabricating an "None" turn misreads in Datadog."""
payload = build(logger, messages=None)
assert payload["meta"]["input"]["messages"] == []
def test_anthropic_tool_blocks_map_to_tool_calls_and_results(logger: DataDogLLMObsLogger) -> None:
"""/v1/messages carries tool traffic as content blocks, not OpenAI fields."""
payload = build(
logger,
messages=[
{"role": "user", "content": [{"type": "text", "text": "Weather in Tokyo?"}]},
{
"role": "assistant",
"content": [{"type": "tool_use", "id": "toolu_1", "name": "get_weather", "input": {"city": "Tokyo"}}],
},
{"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_1", "content": "18C"}]},
],
)
assistant, result_turn = payload["meta"]["input"]["messages"][1:3]
assert assistant["tool_calls"] == [
{"name": "get_weather", "arguments": {"city": "Tokyo"}, "tool_id": "toolu_1", "type": "tool_use"}
]
assert result_turn["tool_results"] == [
{"name": "get_weather", "result": "18C", "tool_id": "toolu_1", "type": "function"}
]
def test_content_with_no_text_parts_is_preserved_not_blanked(logger: DataDogLLMObsLogger) -> None:
"""A content list the mapper does not understand must ride along, not be erased."""
blocks = [{"type": "image_url", "image_url": {"url": "https://example.com/x.png"}}]
payload = build(logger, messages=[{"role": "user", "content": blocks}])
assert payload["meta"]["input"]["messages"][0]["content"] == blocks
def test_multimodal_content_parts_are_flattened_to_text(logger: DataDogLLMObsLogger) -> None:
"""Datadog types Message.content as a string, so content lists collapse to their text."""
payload = build(
logger,
messages=[
{"role": "user", "content": [{"type": "text", "text": "describe "}, {"type": "text", "text": "this"}]}
],
)
assert payload["meta"]["input"]["messages"][0]["content"] == "describe this"
def test_mapping_input_messages_does_not_mutate_the_shared_payload(logger: DataDogLLMObsLogger) -> None:
"""Sibling callbacks read the same messages list, so flattening must not write through it."""
messages: list[dict[str, Any]] = [{"role": "user", "content": [{"type": "text", "text": "hi"}]}]
kwargs = build_payload(messages=messages)
start = datetime(2026, 9, 1, 12, 0, 0)
logger.create_llm_obs_payload(kwargs, start, start + timedelta(seconds=1))
assert messages[0]["content"] == [{"type": "text", "text": "hi"}]
def test_reasoning_content_survives_the_mapping(logger: DataDogLLMObsLogger) -> None:
payload = build(
logger,
response_message={"role": "assistant", "content": "answer", "reasoning_content": "thinking"},
)
assert payload["meta"]["output"]["messages"][0]["reasoning_content"] == "thinking"

68
tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py
View file

@ -541,6 +541,74 @@ def test_llm_call_adapter_extracts_cache_tokens_from_usage_object():
assert data.usage.cache_read_input_tokens == 3
def test_llm_call_adapter_normalizes_nested_cache_tokens():
cases: Final = (
({"prompt_tokens_details": {"cached_tokens": 3}}, 3, None),
({"prompt_cache_hit_tokens": 11}, 11, None),
({"prompt_tokens_details": {"cache_write_tokens": 7}}, None, 7),
({"prompt_tokens_details": {"cache_creation_tokens": 13}}, None, 13),
({"prompt_tokens_details": {"cache_creation_input_tokens": 17}}, None, 17),
)
for usage_object, expected_read, expected_creation in cases:
case_payload = _sample_payload(metadata={"usage_object": usage_object})
data = LLMCallSpanData.from_standard_logging_payload(case_payload)
assert data.usage.cache_read_input_tokens == expected_read
assert data.usage.cache_creation_input_tokens == expected_creation
def test_llm_call_adapter_prefers_nested_count_over_zero_top_level():
payload = _sample_payload(
metadata={
"usage_object": {
"cache_read_input_tokens": 0,
"cache_creation_input_tokens": 0,
"prompt_tokens_details": {"cached_tokens": 5, "cache_write_tokens": 7},
}
}
)
data = LLMCallSpanData.from_standard_logging_payload(payload)
assert data.usage.cache_read_input_tokens == 5
assert data.usage.cache_creation_input_tokens == 7
def test_llm_call_adapter_ignores_invalid_cache_values_before_valid_fallbacks():
payload = _sample_payload(
metadata={
"usage_object": {
"cache_read_input_tokens": -1,
"cache_creation_input_tokens": "5.0",
"prompt_tokens_details": {"cached_tokens": 5, "cache_write_tokens": 7},
}
}
)
data = LLMCallSpanData.from_standard_logging_payload(payload)
assert data.usage.cache_read_input_tokens == 5
assert data.usage.cache_creation_input_tokens == 7
def test_llm_call_adapter_ignores_non_finite_cache_values():
payload = _sample_payload(
metadata={
"usage_object": {
"prompt_tokens_details": {"cached_tokens": float("nan")},
}
}
)
data = LLMCallSpanData.from_standard_logging_payload(payload)
assert data.usage.cache_read_input_tokens is None
def test_llm_call_adapter_preserves_explicit_zero_and_omits_missing_cache_tokens():
for usage_object, expected_read, expected_creation in (
({"prompt_tokens_details": {"cached_tokens": 0}}, 0, None),
({}, None, None),
):
case_payload = _sample_payload(metadata={"usage_object": usage_object})
data = LLMCallSpanData.from_standard_logging_payload(case_payload)
assert data.usage.cache_read_input_tokens == expected_read
assert data.usage.cache_creation_input_tokens == expected_creation
def test_llm_call_adapter_cache_tokens_none_without_usage_object():
data = LLMCallSpanData.from_standard_logging_payload(_sample_payload())
assert data.usage.cache_creation_input_tokens is None

199
tests/test_litellm/integrations/test_custom_guardrail.py
View file

@ -2237,3 +2237,202 @@ class TestRecordsOwnGuardrailInformation:
)
assert _guardrail_entries(request_data) == []
class _ApplyOnlyObserver(CustomGuardrail):
"""Overrides only apply_guardrail, like panw_prisma_airs; inherits async_logging_hook."""
def __init__(self, block: bool = False):
from litellm.types.guardrails import GuardrailEventHooks
super().__init__(guardrail_name="apply-only-observer", event_hook=GuardrailEventHooks.logging_only)
self.block = block
self.calls: list = []
@log_guardrail_information
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
from fastapi import HTTPException
self.calls.append((input_type, list(inputs.get("texts") or [])))
if self.block:
raise HTTPException(status_code=400, detail={"error": "flagged"})
return GenericGuardrailAPIInputs(texts=["[MASKED]" for _ in inputs.get("texts") or []])
def _logged_call(messages: list | str) -> tuple[dict, object]:
from litellm.types.utils import Choices, Message, ModelResponse
response = ModelResponse(choices=[Choices(message=Message(role="assistant", content="general kenobi"))])
kwargs = {
"model": "gpt-5.4-mini",
"messages": messages,
"litellm_call_id": "call-1",
"litellm_params": {"metadata": {"user_api_key_user_id": "u1"}},
"optional_params": {},
"standard_logging_object": {"guardrail_information": None},
}
return kwargs, response
class TestLoggingOnlyApplyGuardrail:
"""LIT-4876 regression: a guardrail in mode logging_only that implements only
apply_guardrail must still run against the logged request and response and
record guardrail_information, instead of inheriting the CustomLogger no-op."""
@pytest.mark.asyncio
async def test_runs_apply_guardrail_observe_only_and_records_verdict(self):
guardrail = _ApplyOnlyObserver()
messages = [{"role": "user", "content": "hello there"}]
kwargs, response = _logged_call(messages)
out_kwargs, out_response = await guardrail.async_logging_hook(kwargs, response, CallTypes.acompletion.value)
assert guardrail.calls == [("request", ["hello there"]), ("response", ["general kenobi"])]
assert out_kwargs["messages"] == [{"role": "user", "content": "hello there"}]
assert out_response.choices[0].message.content == "general kenobi"
entries = out_kwargs["standard_logging_object"]["guardrail_information"]
assert [e["guardrail_name"] for e in entries] == ["apply-only-observer", "apply-only-observer"]
assert {e["guardrail_mode"] for e in entries} == {"logging_only"}
assert {e["guardrail_status"] for e in entries} == {"success"}
assert "standard_logging_guardrail_information" not in kwargs["litellm_params"]["metadata"]
assert kwargs["standard_logging_object"] == {"guardrail_information": None}
@pytest.mark.asyncio
async def test_appends_to_pre_call_verdicts_without_duplicating_them(self):
guardrail = _ApplyOnlyObserver()
kwargs, response = _logged_call([{"role": "user", "content": "hello there"}])
pre_call_entry = {"guardrail_name": "pii-blocker", "guardrail_mode": "pre_call", "guardrail_status": "success"}
kwargs["litellm_params"]["metadata"]["standard_logging_guardrail_information"] = [pre_call_entry]
kwargs["standard_logging_object"]["guardrail_information"] = [pre_call_entry]
out_kwargs, _ = await guardrail.async_logging_hook(kwargs, response, CallTypes.acompletion.value)
entries = out_kwargs["standard_logging_object"]["guardrail_information"]
assert [e["guardrail_name"] for e in entries] == ["pii-blocker", "apply-only-observer", "apply-only-observer"]
assert kwargs["litellm_params"]["metadata"]["standard_logging_guardrail_information"] == [pre_call_entry]
@pytest.mark.asyncio
async def test_request_copy_failure_is_swallowed(self):
import threading
guardrail = _ApplyOnlyObserver()
kwargs, response = _logged_call([{"role": "user", "content": "hello there", "lock": threading.Lock()}])
out_kwargs, out_response = await guardrail.async_logging_hook(kwargs, response, CallTypes.acompletion.value)
assert guardrail.calls == []
assert out_kwargs is kwargs
assert out_response is response
@pytest.mark.asyncio
async def test_block_verdict_is_recorded_without_raising(self):
guardrail = _ApplyOnlyObserver(block=True)
kwargs, response = _logged_call([{"role": "user", "content": "flagged content"}])
out_kwargs, _ = await guardrail.async_logging_hook(kwargs, response, CallTypes.acompletion.value)
assert guardrail.calls == [("request", ["flagged content"])]
entries = out_kwargs["standard_logging_object"]["guardrail_information"]
assert [e["guardrail_status"] for e in entries] == ["guardrail_intervened"]
@pytest.mark.asyncio
async def test_call_type_without_translation_is_skipped(self):
guardrail = _ApplyOnlyObserver()
kwargs, response = _logged_call([{"role": "user", "content": "hello there"}])
out_kwargs, _ = await guardrail.async_logging_hook(kwargs, response, CallTypes.amoderation.value)
assert guardrail.calls == []
assert out_kwargs["standard_logging_object"]["guardrail_information"] is None
@pytest.mark.asyncio
async def test_aembedding_scans_logged_input(self):
from litellm.types.utils import EmbeddingResponse
guardrail = _ApplyOnlyObserver()
kwargs, _ = _logged_call("hello there")
response = EmbeddingResponse(data=[{"embedding": [0.1], "index": 0, "object": "embedding"}])
out_kwargs, out_response = await guardrail.async_logging_hook(kwargs, response, CallTypes.aembedding.value)
assert guardrail.calls == [("request", ["hello there"])]
assert out_kwargs["messages"] == "hello there"
assert out_response is response
entries = out_kwargs["standard_logging_object"]["guardrail_information"]
assert [e["guardrail_status"] for e in entries] == ["success"]
@pytest.mark.asyncio
async def test_native_lifecycle_hook_guardrail_is_left_alone(self):
class _NativeHooks(_ApplyOnlyObserver):
use_native_lifecycle_hooks = True
guardrail = _NativeHooks()
kwargs, response = _logged_call([{"role": "user", "content": "hello there"}])
out_kwargs, out_response = await guardrail.async_logging_hook(kwargs, response, CallTypes.acompletion.value)
assert guardrail.calls == []
assert out_kwargs is kwargs
assert out_response is response
@pytest.mark.asyncio
async def test_aresponses_scans_logged_messages_when_input_is_cleared(self):
from litellm.types.llms.openai import ResponsesAPIResponse
guardrail = _ApplyOnlyObserver()
kwargs, _ = _logged_call([{"role": "user", "content": "hello there"}])
kwargs["input"] = None
response = ResponsesAPIResponse(
id="resp_1",
created_at=1,
model="gpt-5.4-mini",
object="response",
status="completed",
output=[
{
"type": "message",
"id": "msg_1",
"status": "completed",
"role": "assistant",
"content": [{"type": "output_text", "text": "general kenobi"}],
}
],
)
out_kwargs, _ = await guardrail.async_logging_hook(kwargs, response, CallTypes.aresponses.value)
assert guardrail.calls == [("request", ["hello there"]), ("response", ["general kenobi"])]
entries = out_kwargs["standard_logging_object"]["guardrail_information"]
assert [e["guardrail_status"] for e in entries] == ["success", "success"]
@pytest.mark.asyncio
async def test_async_success_handler_records_verdict_in_standard_logging_object(self):
import datetime as dt
from litellm.litellm_core_utils.litellm_logging import Logging
guardrail = _ApplyOnlyObserver()
guardrail.default_on = True
messages = [{"role": "user", "content": "hello there"}]
_, response = _logged_call(messages)
logging_obj = Logging(
model="gpt-5.4-mini",
messages=messages,
stream=False,
call_type=CallTypes.acompletion.value,
start_time=dt.datetime.now(),
litellm_call_id="call-1",
function_id="fn-1",
dynamic_async_success_callbacks=[guardrail],
)
logging_obj.update_environment_variables(
litellm_params={"metadata": {}}, optional_params={}, model="gpt-5.4-mini", custom_llm_provider="openai"
)
await logging_obj.async_success_handler(
result=response, start_time=dt.datetime.now(), end_time=dt.datetime.now()
)
assert guardrail.calls == [("request", ["hello there"]), ("response", ["general kenobi"])]
entries = logging_obj.model_call_details["standard_logging_object"]["guardrail_information"]
assert [e["guardrail_status"] for e in entries] == ["success", "success"]

1
tests/test_litellm/integrations/test_prometheus_client_ip_user_agent.py
View file

@ -93,6 +93,7 @@ async def test_async_post_call_success_hook_includes_client_ip_user_agent():
logger._increment_token_metrics = MagicMock()
logger._increment_remaining_budget_metrics = AsyncMock()
logger._set_virtual_key_rate_limit_metrics = MagicMock()
logger._set_key_and_team_rate_limit_metrics = MagicMock()
logger._set_latency_metrics = MagicMock()
logger.set_llm_deployment_success_metrics = MagicMock()
logger._increment_cache_metrics = MagicMock()

252
tests/test_litellm/integrations/test_prometheus_rate_limit_labels.py
View file

@ -13,6 +13,7 @@ Covers two follow-up gaps to the unified rate-limit error work:
429s don't silently break when the new class lands.
"""
from collections.abc import Mapping
from unittest.mock import MagicMock, patch
import pytest
@ -471,3 +472,254 @@ def test_should_ignore_non_int_v3_header_values(bad_value):
logger.litellm_remaining_api_key_tokens_for_model.labels.return_value.set.assert_called_once_with(
sys.maxsize
)
KEY_AND_TEAM_RATE_LIMIT_METRICS = (
"litellm_api_key_rate_limit_allowed_metric",
"litellm_api_key_rate_limit_used_metric",
"litellm_team_rate_limit_allowed_metric",
"litellm_team_rate_limit_used_metric",
)
def _clear_prometheus_registry() -> None:
from prometheus_client import REGISTRY
for collector in list(REGISTRY._collector_to_names.keys()):
try:
REGISTRY.unregister(collector)
except Exception:
pass
def _collected_samples(metric_name: str) -> dict[tuple[tuple[str, str], ...], float]:
from prometheus_client import REGISTRY
return {
tuple(sorted(sample.labels.items())): sample.value
for metric in REGISTRY.collect()
for sample in metric.samples
if sample.name == metric_name
}
def _success_kwargs_with_rate_limit_headers(additional_headers: Mapping[str, object] | None) -> dict[str, object]:
return {
"model": "claude-haiku-4-5",
"litellm_params": {"metadata": {}},
"standard_logging_object": {
"id": "t",
"call_type": "completion",
"response_cost": 0.001,
"status": "success",
"total_tokens": 20,
"prompt_tokens": 15,
"completion_tokens": 5,
"startTime": 1.0,
"endTime": 2.0,
"completionStartTime": 1.5,
"model": "claude-haiku-4-5",
"model_id": "model-123",
"model_group": "anthropic-haiku-4-5",
"api_base": "https://api.anthropic.com",
"custom_llm_provider": "anthropic",
"request_tags": [],
"end_user": None,
"cache_hit": False,
"stream": False,
"response": None,
"model_parameters": None,
"metadata": {
"user_api_key_hash": "key-hash",
"user_api_key_alias": "key-alias",
"user_api_key_team_id": "team-id",
"user_api_key_team_alias": "team-alias",
"user_api_key_user_id": "u",
"user_api_key_user_email": "e@x.com",
"user_api_key_org_id": None,
"user_api_key_org_alias": None,
"requester_metadata": None,
"user_api_key_end_user_id": None,
"usage_object": None,
},
"hidden_params": {
"litellm_overhead_time_ms": None,
"additional_headers": additional_headers,
},
},
}
async def _run_success_event(
additional_headers: Mapping[str, object] | None, logger: PrometheusLogger | None = None
) -> None:
import datetime
now = datetime.datetime.now()
await (logger or PrometheusLogger()).async_log_success_event(
_success_kwargs_with_rate_limit_headers(additional_headers), None, now, now
)
@pytest.mark.asyncio
async def test_should_emit_key_and_team_rate_limit_allowed_and_used_from_v3_headers():
"""
LIT-1672: the v3 limiter mirrors ``x-ratelimit-{api_key,team}-{limit,remaining}-*``
into the logging payload. The gauges must expose the configured limit as-is
and the window consumption as ``limit - remaining`` for each key / team
dimension, split by ``rate_limit_type``.
"""
_clear_prometheus_registry()
try:
await _run_success_event(
{
"x-ratelimit-api_key-limit-requests": 10,
"x-ratelimit-api_key-remaining-requests": 7,
"x-ratelimit-api_key-limit-tokens": 20000,
"x-ratelimit-api_key-remaining-tokens": 19947,
"x-ratelimit-team-limit-requests": 50,
"x-ratelimit-team-remaining-requests": 47,
"x-ratelimit-team-limit-tokens": 40000,
"x-ratelimit-team-remaining-tokens": 39960,
"x-ratelimit-model_per_key-limit-requests": 5,
"x-ratelimit-model_per_key-remaining-requests": 1,
}
)
key_requests = (
("api_key_alias", "key-alias"),
("hashed_api_key", "key-hash"),
("rate_limit_type", "requests"),
)
key_tokens = (
("api_key_alias", "key-alias"),
("hashed_api_key", "key-hash"),
("rate_limit_type", "tokens"),
)
team_requests = (
("rate_limit_type", "requests"),
("team", "team-id"),
("team_alias", "team-alias"),
)
team_tokens = (
("rate_limit_type", "tokens"),
("team", "team-id"),
("team_alias", "team-alias"),
)
assert _collected_samples("litellm_api_key_rate_limit_allowed_metric") == {
key_requests: 10,
key_tokens: 20000,
}
assert _collected_samples("litellm_api_key_rate_limit_used_metric") == {
key_requests: 3,
key_tokens: 53,
}
assert _collected_samples("litellm_team_rate_limit_allowed_metric") == {
team_requests: 50,
team_tokens: 40000,
}
assert _collected_samples("litellm_team_rate_limit_used_metric") == {
team_requests: 3,
team_tokens: 40,
}
finally:
_clear_prometheus_registry()
@pytest.mark.asyncio
async def test_should_emit_only_the_dimensions_the_limiter_enforced():
"""
A key with only ``rpm_limit`` set and no team limits produces only the
key/requests headers, so no tokens series and no team series may appear
(a phantom 0 or sys.maxsize series would misreport an unlimited dimension).
"""
_clear_prometheus_registry()
try:
await _run_success_event(
{
"x-ratelimit-api_key-limit-requests": 10,
"x-ratelimit-api_key-remaining-requests": 10,
}
)
key_requests = (
("api_key_alias", "key-alias"),
("hashed_api_key", "key-hash"),
("rate_limit_type", "requests"),
)
assert _collected_samples("litellm_api_key_rate_limit_allowed_metric") == {key_requests: 10}
assert _collected_samples("litellm_api_key_rate_limit_used_metric") == {key_requests: 0}
assert _collected_samples("litellm_team_rate_limit_allowed_metric") == {}
assert _collected_samples("litellm_team_rate_limit_used_metric") == {}
finally:
_clear_prometheus_registry()
@pytest.mark.asyncio
async def test_should_drop_key_and_team_series_once_the_limiter_stops_reporting_a_limit():
"""
Removing a key's ``rpm_limit`` / ``tpm_limit`` (or a team's ``tpm_limit``)
makes the v3 limiter stop emitting that descriptor's headers on later
requests. The old allowed/used samples must disappear instead of keeping
a limit that no longer exists on the scrape.
"""
_clear_prometheus_registry()
try:
logger = PrometheusLogger()
await _run_success_event(
{
"x-ratelimit-api_key-limit-requests": 10,
"x-ratelimit-api_key-remaining-requests": 7,
"x-ratelimit-api_key-limit-tokens": 20000,
"x-ratelimit-api_key-remaining-tokens": 19947,
"x-ratelimit-team-limit-requests": 50,
"x-ratelimit-team-remaining-requests": 47,
"x-ratelimit-team-limit-tokens": 40000,
"x-ratelimit-team-remaining-tokens": 39960,
},
logger=logger,
)
await _run_success_event(
{
"x-ratelimit-team-limit-requests": 50,
"x-ratelimit-team-remaining-requests": 46,
},
logger=logger,
)
team_requests = (
("rate_limit_type", "requests"),
("team", "team-id"),
("team_alias", "team-alias"),
)
assert _collected_samples("litellm_api_key_rate_limit_allowed_metric") == {}
assert _collected_samples("litellm_api_key_rate_limit_used_metric") == {}
assert _collected_samples("litellm_team_rate_limit_allowed_metric") == {team_requests: 50}
assert _collected_samples("litellm_team_rate_limit_used_metric") == {team_requests: 4}
finally:
_clear_prometheus_registry()
@pytest.mark.asyncio
@pytest.mark.parametrize(
"additional_headers",
[
None,
{"x-ratelimit-model_per_key-remaining-requests": 42},
{"x-ratelimit-api_key-limit-requests": 10},
{"x-ratelimit-api_key-limit-requests": "10", "x-ratelimit-api_key-remaining-requests": "7"},
{"x-ratelimit-team-limit-tokens": True, "x-ratelimit-team-remaining-tokens": 5},
],
)
async def test_should_emit_no_key_or_team_rate_limit_series_without_a_complete_int_pair(
additional_headers,
):
_clear_prometheus_registry()
try:
await _run_success_event(additional_headers)
for metric_name in KEY_AND_TEAM_RATE_LIMIT_METRICS:
assert _collected_samples(metric_name) == {}, metric_name
finally:
_clear_prometheus_registry()

86
tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py
View file

@ -1522,7 +1522,7 @@ def test_gpt_5_6_alias_prices_match_sol(local_model_cost_map):
sol = litellm.model_cost["gpt-5.6-sol"]
cost_fields = sorted(field for field in sol if "cost" in field)
assert len(cost_fields) == 23
assert len(cost_fields) == 27
for field in cost_fields:
assert alias.get(field) == sol.get(field), field
@ -4039,8 +4039,8 @@ def test_fast_service_tier_matches_priority_above_the_context_threshold(_local_m
)
assert fast == priority
assert fast[0] == pytest.approx(300_000 * 8e-06, rel=1e-9)
assert fast[1] == pytest.approx(1_000 * 3e-05, rel=1e-9)
assert fast[0] == pytest.approx(300_000 * 1.6e-05, rel=1e-9)
assert fast[1] == pytest.approx(1_000 * 6e-05, rel=1e-9)
def test_priority_reasoning_tokens_bill_at_the_priority_output_rate(_local_model_cost_map):
@ -4200,6 +4200,86 @@ def test_generic_cost_per_token_gemini_37_flash(_local_model_cost_map):
assert completion_cost == pytest.approx(0.001875)
GEMINI_38_FLASH_LAUNCH_PRICING = [
("gemini-3.8-flash", 7.5e-07, 3.75e-06, 7.5e-08),
("gemini/gemini-3.8-flash", 7.5e-07, 3.75e-06, 7.5e-08),
("vertex_ai/gemini-3.8-flash", 7.5e-07, 3.75e-06, 7.5e-08),
]
@pytest.mark.parametrize("model,input_cost,output_cost,cache_read_cost", GEMINI_38_FLASH_LAUNCH_PRICING)
def test_gemini_38_flash_launch_pricing(model, input_cost, output_cost, cache_read_cost, _local_model_cost_map):
model_cost_map = litellm.model_cost[model]
assert model_cost_map["input_cost_per_token"] == input_cost
assert model_cost_map["output_cost_per_token"] == output_cost
assert model_cost_map["output_cost_per_reasoning_token"] == output_cost
assert model_cost_map["cache_read_input_token_cost"] == cache_read_cost
assert model_cost_map["mode"] == "chat"
assert model_cost_map["supports_reasoning"] is True
assert model_cost_map["supports_function_calling"] is True
assert model_cost_map["max_input_tokens"] == 1048576
GEMINI_38_FLASH_FIELDS_SHARED_WITH_37_FLASH = (
"input_cost_per_token",
"output_cost_per_token",
"output_cost_per_reasoning_token",
"cache_read_input_token_cost",
"input_cost_per_token_batches",
"output_cost_per_token_batches",
"input_cost_per_token_flex",
"output_cost_per_token_flex",
"cache_read_input_token_cost_flex",
"input_cost_per_token_priority",
"output_cost_per_token_priority",
"cache_read_input_token_cost_priority",
"search_context_cost_per_query",
"google_maps_grounding_cost_per_query",
"prompt_cache_min_tokens",
"max_input_tokens",
"max_output_tokens",
"supports_reasoning",
"supports_function_calling",
"supports_prompt_caching",
"supports_vision",
"supports_pdf_input",
"supports_audio_input",
"supports_video_input",
"supports_response_schema",
"supports_tool_choice",
"supports_web_search",
"supports_url_context",
)
@pytest.mark.parametrize("prefix", ["", "gemini/", "vertex_ai/"])
def test_gemini_38_flash_matches_37_flash_promotional_pricing(prefix, _local_model_cost_map):
new_model = litellm.model_cost[f"{prefix}gemini-3.8-flash"]
old_model = litellm.model_cost[f"{prefix}gemini-3.7-flash"]
for field in GEMINI_38_FLASH_FIELDS_SHARED_WITH_37_FLASH:
assert new_model[field] == old_model[field], field
def test_generic_cost_per_token_gemini_38_flash(_local_model_cost_map):
usage = Usage(
prompt_tokens=1000,
completion_tokens=500,
total_tokens=1500,
completion_tokens_details=CompletionTokensDetailsWrapper(
reasoning_tokens=200,
text_tokens=300,
),
prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=1000),
)
prompt_cost, completion_cost = generic_cost_per_token(
model="gemini-3.8-flash",
usage=usage,
custom_llm_provider="gemini",
)
assert prompt_cost == pytest.approx(0.00075)
assert completion_cost == pytest.approx(0.001875)
def test_grok_46_launch_pricing(_local_model_cost_map):
model_cost_map = litellm.model_cost["xai/grok-4.6"]
assert model_cost_map["input_cost_per_token"] == 2e-06

22
tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py
View file

@ -2932,6 +2932,28 @@ def test_add_cache_point_tool_block_passes_ttl_for_claude_4_5(monkeypatch):
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", old_env)
def test_add_cache_point_tool_block_stands_down_for_model_without_prompt_caching(monkeypatch):
"""A tool carrying cache_control must not become a cachePoint for a Bedrock model
whose cost-map entry lacks prompt caching support, since Bedrock rejects the whole
request. An unmapped id keeps emitting so ARN deployments do not lose caching."""
from litellm.litellm_core_utils.prompt_templates.factory import (
add_cache_point_tool_block,
)
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
tool = {"cache_control": {"type": "ephemeral"}}
assert add_cache_point_tool_block(tool, model="nvidia.nemotron-super-3-120b") is None
assert add_cache_point_tool_block(tool, model="us.nvidia.nemotron-super-3-120b") is None
assert add_cache_point_tool_block(
tool, model="arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc123"
) == {"cachePoint": {"type": "default"}}
assert add_cache_point_tool_block(tool, model="us.anthropic.claude-sonnet-4-5-20250929-v1:0") == {
"cachePoint": {"type": "default"}
}
def test_bedrock_tools_pt_passes_ttl_for_claude_4_5(monkeypatch):
"""
End-to-end: _bedrock_tools_pt should produce cachePoint blocks with ttl

106
tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py
View file

@ -256,14 +256,12 @@ def test_get_model_cost_map_stamps_loaded_at(monkeypatch):
from litellm.litellm_core_utils import get_model_cost_map as module
monkeypatch.setattr(module._cost_map_source_info, "loaded_at", None)
monkeypatch.setattr(
module.GetModelCostMap,
"fetch_remote_model_cost_map",
staticmethod(lambda url, timeout=5: _load_root_cost_map()),
client, _calls = _mock_client(
[httpx.Response(200, content=_real_map_bytes())], client_cls=httpx.Client
)
before = datetime.now(timezone.utc)
module.get_model_cost_map(url="https://example.invalid/cost_map.json")
module.get_model_cost_map(url="https://example.invalid/cost_map.json", client=client)
loaded_at = module.get_model_cost_map_loaded_at()
assert loaded_at is not None
@ -308,7 +306,7 @@ def _unset_local_cost_map_env(monkeypatch):
monkeypatch.delenv("LITELLM_LOCAL_MODEL_COST_MAP", raising=False)
def _mock_client(outcomes):
def _mock_client(outcomes, client_cls=httpx.AsyncClient):
"""httpx client over a MockTransport serving one outcome per request; an exception instance is raised."""
calls = {"count": 0}
@ -320,7 +318,7 @@ def _mock_client(outcomes):
raise outcome
return outcome
return httpx.AsyncClient(transport=httpx.MockTransport(handler)), calls
return client_cls(transport=httpx.MockTransport(handler)), calls
@pytest.mark.asyncio
@ -450,3 +448,97 @@ async def test_refetch_respects_local_env_override(monkeypatch):
)
assert isinstance(result, ModelCostMapReloaded)
assert len(result.model_cost_map) > 100
# ---------------------------------------------------------------------------
# get_model_cost_map: the boot-time load retries transient failures like a reload does
# ---------------------------------------------------------------------------
from litellm.litellm_core_utils.get_model_cost_map import (
get_model_cost_map,
get_model_cost_map_source_info,
)
class _SyncSleepRecorder:
"""Injected in place of time.sleep so the boot path's waits are asserted without delay."""
def __init__(self):
self.waits = []
def __call__(self, seconds: float) -> None:
self.waits.append(seconds)
def test_boot_load_retries_transient_failures_instead_of_falling_back():
"""A refused connection then a 503 at pod boot used to pin the process to the bundled
backup for its lifetime; both are transient and must be retried before giving up."""
client, calls = _mock_client(
[
httpx.ConnectError("connection refused"),
httpx.Response(503),
httpx.Response(200, content=_real_map_bytes()),
],
client_cls=httpx.Client,
)
sleeper = _SyncSleepRecorder()
cost_map = get_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client)
assert calls["count"] == 3
assert len(sleeper.waits) == 2
assert 2.0 <= sleeper.waits[0] < 3.0
assert 4.0 <= sleeper.waits[1] < 5.0
source = get_model_cost_map_source_info()
assert source["source"] == "remote"
assert source["fallback_reason"] is None
assert cost_map.keys() >= _load_root_cost_map().keys() - {"sample_spec", FALLBACK_GENERALIZATIONS_KEY}
def test_boot_load_honors_retry_after_then_falls_back_after_max_attempts():
"""An outage longer than the retry budget still ends on the bundled backup, and the
recorded fallback reason says how many attempts were spent so operators can tell."""
client, calls = _mock_client(
[httpx.Response(429, headers={"Retry-After": "7"})], client_cls=httpx.Client
)
sleeper = _SyncSleepRecorder()
cost_map = get_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client)
assert calls["count"] == 3
assert sleeper.waits == [7.0, 7.0]
source = get_model_cost_map_source_info()
assert source["source"] == "local"
assert "after 3 attempts" in source["fallback_reason"]
assert len(cost_map) > 100
def test_boot_load_does_not_retry_permanent_failures():
"""A 404 or a malformed URL cannot heal by waiting: one attempt, no sleeps, backup."""
client, calls = _mock_client([httpx.Response(404)], client_cls=httpx.Client)
sleeper = _SyncSleepRecorder()
get_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client)
assert calls["count"] == 1
assert sleeper.waits == []
assert get_model_cost_map_source_info()["source"] == "local"
get_model_cost_map(url="not a url", sleep=sleeper, rng=random.Random(0))
assert sleeper.waits == []
assert get_model_cost_map_source_info()["source"] == "local"
def test_boot_load_respects_local_env_override(monkeypatch):
"""LITELLM_LOCAL_MODEL_COST_MAP=True still short-circuits to the backup with zero HTTP."""
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
def _fail(request):
raise AssertionError("no HTTP request should be made when local map is forced")
cost_map = get_model_cost_map(
url=_URL,
sleep=_SyncSleepRecorder(),
client=httpx.Client(transport=httpx.MockTransport(_fail)),
)
assert len(cost_map) > 100
assert get_model_cost_map_source_info()["is_env_forced"] is True

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

@ -592,6 +592,59 @@ def test_stream_chunk_builder_litellm_usage_chunks():
assert usage.total_tokens == 77
def test_calculate_usage_honors_openai_sdk_completion_usage_chunks():
from openai.types.completion_usage import CompletionUsage
content_chunk = ModelResponseStream(
id="chatcmpl-sdk-usage-1",
created=1745513206,
model="mantle-claude",
object="chat.completion.chunk",
system_fingerprint=None,
choices=[
StreamingChoices(
finish_reason="stop",
index=0,
delta=Delta(
provider_specific_fields=None,
content="ok",
role=None,
function_call=None,
tool_calls=None,
audio=None,
),
logprobs=None,
)
],
provider_specific_fields=None,
stream_options={"include_usage": True},
)
usage_chunk = ModelResponseStream(
id="chatcmpl-sdk-usage-1",
created=1745513207,
model="mantle-claude",
object="chat.completion.chunk",
system_fingerprint=None,
choices=[],
provider_specific_fields=None,
stream_options={"include_usage": True},
)
usage_chunk.usage = CompletionUsage(
prompt_tokens=20, completion_tokens=60, total_tokens=80, cost=0.000704
)
assert type(usage_chunk.usage) is CompletionUsage
chunks = [content_chunk, usage_chunk]
usage = ChunkProcessor(chunks=chunks).calculate_usage(
chunks=chunks, model="mantle-claude", completion_output=""
)
assert usage.prompt_tokens == 20
assert usage.completion_tokens == 60
assert usage.total_tokens == 80
assert getattr(usage, "cost", None) == pytest.approx(0.000704)
def test_get_model_from_chunks_azure_model_router():
"""
Test that _get_model_from_chunks finds the actual model from Azure Model Router chunks.

39
tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py
View file

@ -3,11 +3,13 @@ import threading
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
import litellm
from litellm.constants import RESPONSE_FORMAT_TOOL_NAME
from litellm.llms.anthropic.chat.handler import ModelResponseIterator, make_call
from litellm.llms.custom_httpx.http_handler import HTTPHandler
from litellm.types.llms.openai import (
ChatCompletionToolCallChunk,
ChatCompletionToolCallFunctionChunk,
@ -46,6 +48,43 @@ async def test_make_call_passes_logging_obj_to_client_post():
assert call_kwargs.get("logging_obj") is logging_obj
def test_anthropic_completion_does_not_send_deployment_default_limits():
captured_requests: list[httpx.Request] = []
def respond(request: httpx.Request) -> httpx.Response:
captured_requests.append(request)
return httpx.Response(
200,
json={
"id": "msg_default_limits",
"type": "message",
"role": "assistant",
"model": "claude-3-5-haiku-20241022",
"content": [{"type": "text", "text": "Hello"}],
"stop_reason": "end_turn",
"stop_sequence": None,
"usage": {"input_tokens": 1, "output_tokens": 1},
},
)
client = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond)))
try:
litellm.completion(
model="anthropic/claude-3-5-haiku-20241022",
messages=[{"role": "user", "content": "Hello"}],
api_key="test-key",
client=client,
default_api_key_rpm_limit=60,
default_api_key_tpm_limit=5000000,
)
finally:
client.close()
request_body = json.loads(captured_requests[0].content)
assert "default_api_key_rpm_limit" not in request_body
assert "default_api_key_tpm_limit" not in request_body
def test_redacted_thinking_content_block_delta():
chunk = {
"type": "content_block_start",

36
tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py
View file

@ -96,10 +96,8 @@ def _completion_kwargs(**overrides):
return kwargs
def _run(**overrides):
with patch.object(
BedrockConverseLLM, "get_credentials", return_value=RESOLVED_CREDENTIALS
):
def _run(*, credentials: Credentials | None = RESOLVED_CREDENTIALS, **overrides):
with patch.object(BedrockConverseLLM, "get_credentials", return_value=credentials):
return BedrockConverseLLM().completion(**_completion_kwargs(**overrides))
@ -360,7 +358,7 @@ async def test_async_completion_logs_pre_call_by_default():
def _sync_client_returning_converse_response():
client = MagicMock()
client.post = lambda **_kwargs: httpx.Response(
client.post.side_effect = lambda **_kwargs: httpx.Response(
200,
json=CONVERSE_RESPONSE,
request=httpx.Request("POST", "https://bedrock-runtime.us-west-2.amazonaws.com"),
@ -487,3 +485,31 @@ def test_post_call_is_not_logged_twice_when_the_sync_rust_call_declines():
assert response.choices[0].message.content == "hi"
assert len(calls["post_call"]) == 1
assert "hi" in calls["post_call"][0]["original_response"]
def test_bearer_token_auth_serves_when_boto3_resolves_no_sigv4_credentials(monkeypatch):
"""With only `AWS_BEARER_TOKEN_BEDROCK` configured boto3 resolves no
credentials at all. Preparing the Rust handoff must not dereference that
None: the bearer token signs the request on its own."""
monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "bedrock-bearer-token")
client = _sync_client_returning_converse_response()
response = _run(credentials=None, litellm_params={}, client=client)
assert response.choices[0].message.content == "hi"
sent_headers = client.post.call_args.kwargs["headers"]
assert sent_headers["Authorization"] == "Bearer bedrock-bearer-token"
def test_the_rust_opt_in_needs_no_sigv4_principal():
"""The core resolves the bearer token itself, so a bearer-only deployment
keeps its opt-in and the gate sees no aws_* credential keys to sign with."""
seen = _inject()
response = _run(credentials=None, api_key="bedrock-bearer-token")
assert response.choices[0].message.content == "hello from rust"
params = seen["call"][0]["optional_params"]
assert not {"aws_access_key_id", "aws_secret_access_key", "aws_session_token"} & params.keys()
assert params["aws_region_name"] == "us-east-1"
assert seen["call"][0]["api_key"] == "bedrock-bearer-token"

80
tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py
View file

@ -5248,6 +5248,84 @@ def test_cache_control_injection_tool_config_drops_ttl_for_unsupported_model():
assert tools[-1] == {"cachePoint": {"type": "default"}}
@pytest.mark.parametrize(
("model", "expects_cache_points"),
[
pytest.param("nvidia.nemotron-super-3-120b", False, id="mapped-model-without-prompt-caching"),
pytest.param("us.nvidia.nemotron-super-3-120b", False, id="regional-prefix-resolves-through-base-model"),
pytest.param(
"us.anthropic.claude-3-5-sonnet-20240620-v1:0", False, id="claude-named-but-not-caching-on-bedrock"
),
pytest.param("us.anthropic.claude-sonnet-4-5-20250929-v1:0", True, id="mapped-model-with-prompt-caching"),
pytest.param(
"arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc123",
True,
id="unmapped-arn-keeps-emitting",
),
],
)
def test_cache_points_emitted_only_for_models_that_support_prompt_caching(model, expects_cache_points, monkeypatch):
"""Bedrock rejects cachePoint blocks for models without prompt caching support
("You invoked an unsupported model or your request did not allow prompt caching"),
and clients like Claude Code attach cache_control to every request, so a map-known
model without the capability must not receive them. Unmapped ids (application
inference profile ARNs, models newer than the map) keep emitting so existing
caching setups never silently degrade."""
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
body = AmazonConverseConfig().transform_request(
model=model,
messages=[
{"role": "system", "content": [{"type": "text", "text": "sys", "cache_control": {"type": "ephemeral"}}]},
{"role": "user", "content": [{"type": "text", "text": "hi", "cache_control": {"type": "ephemeral"}}]},
],
optional_params={},
litellm_params={},
headers={},
)
assert ("cachePoint" in json.dumps(body)) is expects_cache_points
assert body["system"][0]["text"] == "sys"
assert body["messages"][0]["content"][0]["text"] == "hi"
def test_tool_config_cachepoint_not_placed_or_credited_for_model_without_prompt_caching(monkeypatch):
"""The tool_config injection point must stand down with the rest of the cachePoint
emission when the model cannot cache, and spend attribution must not credit the
gateway for a breakpoint that was never placed."""
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
bucket: dict = {"user_api_key": "sk-test"}
data = AmazonConverseConfig()._transform_request_helper(
model="nvidia.nemotron-super-3-120b",
system_content_blocks=[],
optional_params={
"tools": [
{
"type": "function",
"function": {
"name": "get_weather",
"description": "Get weather",
"parameters": {
"type": "object",
"properties": {"location": {"type": "string"}},
"required": ["location"],
},
},
}
],
"cache_control_injection_points": [{"location": "tool_config"}],
},
messages=[{"role": "user", "content": "hi"}],
litellm_params={"metadata": bucket, "litellm_metadata": None, "model_info": {"id": "dep-bedrock"}},
)
assert "cachePoint" not in json.dumps(data.get("toolConfig", {}))
assert "litellm_gateway_injected_cache" not in bucket
def test_translate_response_format_json_schema_still_injects_tool():
"""
response_format with an explicit json_schema should still use the
@ -6211,7 +6289,7 @@ def test_message_level_cache_control_drops_ttl_for_unsupported_model(ttl_target)
result = _bedrock_converse_messages_pt(
messages=_agentic_messages_with_ttl(ttl_target),
model="anthropic.claude-3-5-sonnet-20240620-v1:0",
model="anthropic.claude-3-5-sonnet-20241022-v2:0",
llm_provider="bedrock_converse",
)

18
tests/test_litellm/llms/bedrock/test_base_aws_llm.py
View file

@ -15,6 +15,7 @@ from unittest.mock import MagicMock, patch
from botocore.awsrequest import AWSPreparedRequest, AWSRequest
from botocore.auth import SigV4Auth
from botocore.credentials import Credentials
from botocore.exceptions import NoCredentialsError
import litellm
from litellm.llms.bedrock.base_aws_llm import (
@ -801,6 +802,23 @@ def test_get_request_headers_with_sigv4():
assert result == mock_request.prepare.return_value
def test_get_request_headers_without_credentials_or_bearer_token_raises_no_credentials():
"""Bearer-token auth needs no SigV4 principal, so `credentials` may be None.
Reaching the SigV4 branch with neither must fail the way botocore always
has instead of signing with a missing principal."""
llm = BaseAWSLLM()
with patch.dict(os.environ, {}, clear=True), pytest.raises(NoCredentialsError):
llm.get_request_headers(
credentials=None,
aws_region_name="us-west-2",
extra_headers=None,
endpoint_url="https://api.example.com",
data='{"prompt": "test"}',
headers={"Content-Type": "application/json"},
)
def test_sigv4_matches_rust_golden_vector():
request = AWSRequest(
method="POST",

56
tests/test_litellm/llms/custom_httpx/test_aiohttp_transport.py
View file

@ -1,7 +1,11 @@
import asyncio
import concurrent.futures
import socket
import sys
from typing import Final
import aiohttp
import aiohttp.abc
import aiohttp.client_exceptions
import aiohttp.http_exceptions
import httpx
@ -1140,3 +1144,55 @@ async def test_stopped_loop_session_disposed_synchronously_on_recycle():
finally:
await new_session.close()
result["loop"].close()
class _CancellingResolver(aiohttp.abc.AbstractResolver):
"""Cancels the given task (or, by default, aiohttp's shielded DNS child task) mid-lookup."""
def __init__(self, task_to_cancel: "asyncio.Task[object] | None" = None):
self._task_to_cancel: Final = task_to_cancel
async def resolve(
self, host: str, port: int = 0, family: socket.AddressFamily = socket.AF_INET
) -> list[aiohttp.abc.ResolveResult]:
target: Final = self._task_to_cancel or asyncio.current_task()
assert target is not None
target.cancel()
await asyncio.sleep(0)
raise OSError("resolver finished after the task was cancelled")
async def close(self) -> None:
return None
@pytest.mark.asyncio
@pytest.mark.skipif(
sys.version_info < (3, 11), reason="Task.cancelling() is needed to tell the two cancellations apart"
)
async def test_internal_dns_cancellation_maps_to_connect_error():
"""A CancelledError the request task never asked for must surface as a mapped httpx transport error."""
session = aiohttp.ClientSession(connector=aiohttp.TCPConnector(resolver=_CancellingResolver()))
transport = LiteLLMAiohttpTransport(client=session)
try:
with pytest.raises(httpx.ConnectError):
await transport.handle_async_request(httpx.Request("GET", "http://example.invalid/"))
current = asyncio.current_task()
assert current is not None and current.cancelling() == 0
finally:
await transport.aclose()
@pytest.mark.asyncio
async def test_genuine_request_cancellation_still_propagates():
"""Cancelling the request task itself (client disconnect, shutdown) must still propagate unmapped."""
current = asyncio.current_task()
assert current is not None
session = aiohttp.ClientSession(connector=aiohttp.TCPConnector(resolver=_CancellingResolver(current)))
transport = LiteLLMAiohttpTransport(client=session)
try:
with pytest.raises(asyncio.CancelledError):
await transport.handle_async_request(httpx.Request("GET", "http://example.invalid/"))
finally:
if sys.version_info >= (3, 11):
current.uncancel()
await transport.aclose()

84
tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py
View file

@ -1559,3 +1559,87 @@ class TestScanOnlyToolResults:
assert data["messages"][3]["content"] == "page says [BLOCKED] here"
assert data["messages"][3]["tool_call_id"] == "call_1"
assert data["messages"][4]["content"] == "and then?"
class TestBuildBlockSseChunks:
"""build_block_sse_chunks turns a streaming ModifyResponseException into 200 SSE chunks"""
def _exc(self, original_response=None):
from litellm.exceptions import ModifyResponseException
return ModifyResponseException(
message="Blocked by policy.",
model="gpt-5.4-mini",
request_data={},
guardrail_name="test",
original_response=original_response,
)
def _payloads(self, chunks):
return [json.loads(chunk.decode().removeprefix("data: ").strip()) for chunk in chunks]
def test_standalone_block_uses_fresh_identity_and_zero_usage(self):
handler = OpenAIChatCompletionsHandler()
first, final = self._payloads(handler.build_block_sse_chunks(self._exc(), stream_started=False))
assert first["id"].startswith("chatcmpl-")
assert first["model"] == "gpt-5.4-mini"
assert first["choices"][0]["delta"] == {"role": "assistant", "content": "Blocked by policy."}
assert first["choices"][0]["finish_reason"] is None
assert final["choices"][0]["delta"] == {}
assert final["choices"][0]["finish_reason"] == "content_filter"
assert final["usage"] == {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
def test_continuation_reuses_stream_identity_and_real_usage(self):
handler = OpenAIChatCompletionsHandler()
yielded = [
{"id": "chatcmpl-live", "created": 1724900000, "model": "gpt-5.4-mini-2026-01-01"},
]
original = yielded + [
{"id": "chatcmpl-live", "usage": {"prompt_tokens": 11, "completion_tokens": 5}},
]
first, final = self._payloads(
handler.build_block_sse_chunks(
self._exc(original_response=original), stream_started=True, responses_so_far=yielded
)
)
assert (first["id"], first["created"], first["model"]) == (
"chatcmpl-live",
1724900000,
"gpt-5.4-mini-2026-01-01",
)
assert first["choices"][0]["delta"] == {"content": "Blocked by policy."}
assert final["id"] == "chatcmpl-live"
assert final["usage"] == {"prompt_tokens": 11, "completion_tokens": 5, "total_tokens": 16}
class TestCheckStreamingHasEnded:
"""_check_streaming_has_ended lets end_of_stream_only withhold the finish chunk until moderation"""
def test_empty_and_content_only_chunks_are_not_ended(self):
handler = OpenAIChatCompletionsHandler()
assert handler._check_streaming_has_ended([]) is False
content_only = [
{"id": "chatcmpl-live", "choices": [{"index": 0, "delta": {"content": "hi"}, "finish_reason": None}]},
{"id": "chatcmpl-live", "choices": []},
{"id": "chatcmpl-live", "usage": {"prompt_tokens": 1, "completion_tokens": 1}},
]
assert handler._check_streaming_has_ended(content_only) is False
def test_dict_finish_chunk_marks_stream_ended(self):
handler = OpenAIChatCompletionsHandler()
chunks = [
{"id": "chatcmpl-live", "choices": [{"index": 0, "delta": {"content": "hi"}, "finish_reason": None}]},
{"id": "chatcmpl-live", "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]},
]
assert handler._check_streaming_has_ended(chunks) is True
def test_object_finish_chunk_marks_stream_ended(self):
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
handler = OpenAIChatCompletionsHandler()
chunks = [
ModelResponseStream(
choices=[StreamingChoices(index=0, delta=Delta(content=None), finish_reason="stop")]
)
]
assert handler._check_streaming_has_ended(chunks) is True

216
tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py
View file

@ -1321,3 +1321,219 @@ class TestOpenAIResponsesHandlerToolInjection:
names = [t.get("name") for t in result["tools"]]
assert "get_weather" in names
assert "injected_tool" in names
class TestBuildBlockSseChunks:
"""build_block_sse_chunks turns a streaming ModifyResponseException into 200 SSE events"""
def _exc(self, original_response=None):
from litellm.exceptions import ModifyResponseException
return ModifyResponseException(
message="Blocked by policy.",
model="gpt-5.4-mini",
request_data={},
guardrail_name="test",
original_response=original_response,
)
def _payloads(self, chunks):
import json
return [json.loads(chunk.decode().removeprefix("data: ").strip()) for chunk in chunks]
def test_standalone_block_emits_complete_synthetic_stream(self):
handler = OpenAIResponsesHandler()
payloads = self._payloads(handler.build_block_sse_chunks(self._exc(), stream_started=False))
types = [payload["type"] for payload in payloads]
assert types[0] == "response.created"
assert types[-1] == "response.completed"
completed = payloads[-1]["response"]
assert completed["id"].startswith("resp_")
assert completed["model"] == "gpt-5.4-mini"
assert completed["output"][0]["content"][0]["text"] == "Blocked by policy."
def test_continuation_appends_item_at_next_output_index_with_real_usage(self):
handler = OpenAIResponsesHandler()
yielded = [
{"type": "response.created", "response": {"id": "resp_live", "model": "gpt-5.4-mini-2026-01-01"}},
{"type": "response.output_item.added", "output_index": 2, "item": {"id": "msg_orig"}},
]
original = yielded + [
{
"type": "response.completed",
"response": {
"id": "resp_live",
"model": "gpt-5.4-mini-2026-01-01",
"output": [],
"usage": {"input_tokens": 7, "output_tokens": 21, "total_tokens": 28},
},
}
]
payloads = self._payloads(
handler.build_block_sse_chunks(
self._exc(original_response=original), stream_started=True, responses_so_far=yielded
)
)
types = [payload["type"] for payload in payloads]
assert "response.created" not in types
assert types[0] == "response.output_item.done"
assert payloads[0]["output_index"] == 2
assert payloads[0]["item"]["id"] == "msg_orig"
assert payloads[0]["item"]["status"] == "completed"
assert types[1] == "response.output_item.added"
assert payloads[1]["output_index"] == 3
completed = payloads[-1]["response"]
assert completed["id"] == "resp_live"
assert completed["model"] == "gpt-5.4-mini-2026-01-01"
assert completed["output"][0]["content"][0]["text"] == "Blocked by policy."
assert completed["usage"] == {"input_tokens": 7, "output_tokens": 21, "total_tokens": 28}
def test_continuation_reads_usage_from_typed_completed_event(self):
from litellm.types.llms.openai import (
ResponseCompletedEvent,
ResponsesAPIResponse,
ResponsesAPIStreamEvents,
)
handler = OpenAIResponsesHandler()
original = [
ResponseCompletedEvent(
type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
response=ResponsesAPIResponse.model_validate(
{
"id": "resp_live",
"created_at": 1,
"model": "gpt-5.4-mini",
"output": [],
"usage": {"input_tokens": 7, "output_tokens": 21, "total_tokens": 28},
}
),
)
]
payloads = self._payloads(
handler.build_block_sse_chunks(
self._exc(original_response=original), stream_started=True, responses_so_far=[]
)
)
completed = payloads[-1]["response"]
assert completed["usage"]["input_tokens"] == 7
assert completed["usage"]["output_tokens"] == 21
assert completed["usage"]["total_tokens"] == 28
def test_continuation_closes_open_item_given_pydantic_events_with_enum_types(self):
from litellm.types.llms.openai import (
BaseLiteLLMOpenAIResponseObject,
ContentPartAddedEvent,
OutputItemAddedEvent,
OutputTextDeltaEvent,
ResponsesAPIStreamEvents,
)
handler = OpenAIResponsesHandler()
open_item = GenericResponseOutputItem.model_validate(
{"type": "message", "id": "msg_live", "status": "in_progress", "role": "assistant", "content": []}
)
yielded = [
OutputItemAddedEvent(
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, output_index=0, item=open_item
),
ContentPartAddedEvent(
type=ResponsesAPIStreamEvents.CONTENT_PART_ADDED,
item_id="msg_live",
output_index=0,
content_index=0,
part=BaseLiteLLMOpenAIResponseObject.model_validate(
{"type": "output_text", "text": "", "annotations": []}
),
),
OutputTextDeltaEvent(
type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA,
item_id="msg_live",
output_index=0,
content_index=0,
delta="partial ",
),
OutputTextDeltaEvent(
type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA,
item_id="msg_live",
output_index=0,
content_index=0,
delta="text",
),
]
payloads = self._payloads(
handler.build_block_sse_chunks(
self._exc(original_response=yielded), stream_started=True, responses_so_far=yielded
)
)
types = [payload["type"] for payload in payloads]
assert types[:3] == [
"response.output_text.done",
"response.content_part.done",
"response.output_item.done",
]
assert payloads[0]["text"] == "partial text"
assert payloads[2]["item"]["id"] == "msg_live"
assert payloads[2]["item"]["status"] == "completed"
assert payloads[2]["item"]["content"][0]["text"] == "partial text"
assert types[3] == "response.output_item.added"
assert payloads[3]["output_index"] == 1
def test_continuation_closes_open_function_call_as_incomplete(self):
handler = OpenAIResponsesHandler()
yielded = [
{"type": "response.created", "response": {"id": "resp_live", "model": "gpt-5.4-mini"}},
{
"type": "response.output_item.added",
"output_index": 0,
"item": {
"id": "fc_live",
"type": "function_call",
"status": "in_progress",
"call_id": "call_1",
"name": "run_payment",
"arguments": "",
},
},
{
"type": "response.function_call_arguments.delta",
"item_id": "fc_live",
"output_index": 0,
"delta": '{"amount": 100}',
},
]
payloads = self._payloads(
handler.build_block_sse_chunks(
self._exc(original_response=yielded), stream_started=True, responses_so_far=yielded
)
)
types = [payload["type"] for payload in payloads]
assert types[0] == "response.output_item.done"
closed = payloads[0]["item"]
assert closed["id"] == "fc_live"
assert closed["type"] == "function_call"
assert closed["status"] == "incomplete"
assert closed["name"] == "run_payment"
assert "content" not in closed
assert types[1] == "response.output_item.added"
assert payloads[1]["output_index"] == 1
assert types[-1] == "response.completed"
def test_continuation_without_open_item_emits_no_closing_events(self):
handler = OpenAIResponsesHandler()
yielded = [
{"type": "response.created", "response": {"id": "resp_live", "model": "gpt-5.4-mini"}},
{"type": "response.in_progress", "response": {"id": "resp_live"}},
]
payloads = self._payloads(
handler.build_block_sse_chunks(
self._exc(original_response=yielded), stream_started=True, responses_so_far=yielded
)
)
types = [payload["type"] for payload in payloads]
assert types[0] == "response.output_item.added"
assert types[-1] == "response.completed"
dones = [payload for payload in payloads if payload["type"] == "response.output_item.done"]
assert len(dones) == 1
assert dones[0]["item"]["content"][0]["text"] == "Blocked by policy."

190
tests/test_litellm/llms/parallel_ai/test_parallel_ai_search.py
View file

@ -2,6 +2,7 @@
Tests for Parallel AI Search API integration (v1 endpoint).
"""
import json
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@ -30,13 +31,41 @@ MOCK_V1_RESPONSE = {
}
def _mock_response():
def _mock_response(payload=None):
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = MOCK_V1_RESPONSE
mock_response.json.return_value = payload if payload is not None else MOCK_V1_RESPONSE
return mock_response
@pytest.fixture
def httpx_transport(monkeypatch):
monkeypatch.setattr( # test-quality-ok: respx needs HTTPX enabled to fake the provider HTTP boundary.
litellm,
"disable_aiohttp_transport",
True,
)
litellm.in_memory_llm_clients_cache.flush_cache()
yield
litellm.in_memory_llm_clients_cache.flush_cache()
@pytest.fixture
def bundled_cost_map(monkeypatch):
"""Price lookups against the bundled cost map.
litellm caches model-info lookups, so swapping ``model_cost`` only takes
effect once those caches are invalidated -- on the way in and back out.
"""
from litellm.utils import _invalidate_model_cost_lowercase_map
monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
_invalidate_model_cost_lowercase_map()
yield
monkeypatch.undo()
_invalidate_model_cost_lowercase_map()
class TestParallelAISearch:
@pytest.fixture(autouse=True)
def _set_api_key(self, monkeypatch):
@ -135,9 +164,7 @@ class TestParallelAISearch:
json_data = mock_post.call_args.kwargs.get("json")
assert json_data["mode"] == "basic"
@pytest.mark.parametrize(
"processor,expected_mode", [("base", "basic"), ("pro", "advanced")]
)
@pytest.mark.parametrize("processor,expected_mode", [("base", "basic"), ("pro", "advanced")])
@pytest.mark.asyncio
async def test_legacy_processor_maps_to_mode(self, processor, expected_mode):
with patch(
@ -222,9 +249,7 @@ class TestParallelAISearch:
"arxiv.org",
"nature.com",
]
assert advanced_settings["source_policy"]["exclude_domains"] == [
"reddit.com"
]
assert advanced_settings["source_policy"]["exclude_domains"] == ["reddit.com"]
assert advanced_settings["excerpt_settings"]["max_chars_per_result"] == 1500
assert "max_results" not in json_data
@ -306,10 +331,7 @@ class TestParallelAISearch:
)
call_args = mock_post.call_args
assert (
call_args.kwargs["url"]
== "https://proxy.internal.example.com/v1/search"
)
assert call_args.kwargs["url"] == "https://proxy.internal.example.com/v1/search"
@pytest.mark.asyncio
async def test_caller_api_base_without_key_is_refused(self, monkeypatch):
@ -338,3 +360,147 @@ class TestParallelAISearch:
query="AI developments",
search_provider="parallel_ai",
)
@pytest.mark.asyncio
async def test_flat_source_and_fetch_params_nest_under_advanced_settings(self, respx_mock, httpx_transport):
route = respx_mock.post("https://api.parallel.ai/v1/search").respond(json=MOCK_V1_RESPONSE)
await litellm.asearch(
query="AI developments",
search_provider="parallel_ai",
objective="find peer-reviewed AI research",
include_domains=["arxiv.org"],
after_date="2026-01-01",
location="gb",
fetch_policy={"max_age_seconds": 600, "disable_cache_fallback": True},
client_model="claude-fable-5",
)
json_data = json.loads(route.calls[0].request.content)
assert json_data["objective"] == "find peer-reviewed AI research"
assert json_data["client_model"] == "claude-fable-5"
advanced_settings = json_data["advanced_settings"]
assert advanced_settings["location"] == "gb"
assert advanced_settings["fetch_policy"] == {
"max_age_seconds": 600,
"disable_cache_fallback": True,
}
assert advanced_settings["source_policy"]["include_domains"] == ["arxiv.org"]
assert advanced_settings["source_policy"]["after_date"] == "2026-01-01"
assert "include_domains" not in json_data
assert "after_date" not in json_data
assert "location" not in json_data
assert "fetch_policy" not in json_data
@pytest.mark.asyncio
async def test_response_preserves_raw_parallel_fields(self, respx_mock, httpx_transport):
respx_mock.post("https://api.parallel.ai/v1/search").respond(json=MOCK_V1_RESPONSE)
response = await litellm.asearch(
query="AI developments",
search_provider="parallel_ai",
)
dumped = response.model_dump()
assert dumped["search_id"] == "search_abc123"
assert dumped["session_id"] == "session_xyz"
assert dumped["parallel_usage"] == [{"name": "search_advanced", "count": 1}]
first = response.results[0].model_dump()
assert first["excerpts"] == ["First excerpt.", "Second excerpt."]
@pytest.mark.asyncio
async def test_response_normalizes_null_result_fields(self, respx_mock, httpx_transport):
response_payload = {
**MOCK_V1_RESPONSE,
"results": [{"url": None, "title": None, "publish_date": None, "excerpts": None}],
}
respx_mock.post("https://api.parallel.ai/v1/search").respond(json=response_payload)
response = await litellm.asearch(
query="AI developments",
search_provider="parallel_ai",
)
assert len(response.results) == 1
result = response.results[0]
assert result.url == ""
assert result.title == ""
assert result.snippet == ""
assert result.date is None
assert result.model_dump()["excerpts"] == ()
@pytest.mark.parametrize(
"mode,usage,max_results,expected_cost",
[
("turbo", [{"name": "sku_search", "count": 1}], None, 0.001),
("fast", [{"name": "sku_search", "count": 1}], None, 0.001),
("basic", [{"name": "sku_search", "count": 1}], None, 0.005),
("advanced", [{"name": "sku_search", "count": 1}], None, 0.005),
(
"basic",
[
{"name": "sku_search", "count": 1},
{"name": "sku_search_additional_results", "count": 2},
],
20,
0.007,
),
("basic", None, 20, 0.015),
],
)
@pytest.mark.asyncio
async def test_search_cost_uses_mode_and_provider_usage(
self, mode, usage, max_results, expected_cost, bundled_cost_map, respx_mock, httpx_transport
):
response_payload = {**MOCK_V1_RESPONSE, "usage": usage}
respx_mock.post("https://api.parallel.ai/v1/search").respond(json=response_payload)
response = await litellm.asearch(
query="AI developments",
search_provider="parallel_ai",
mode=mode,
max_results=max_results,
)
assert response._hidden_params["response_cost"] == pytest.approx(expected_cost)
@pytest.mark.asyncio
async def test_search_cost_treats_keyword_queries_as_one_request(
self, bundled_cost_map, respx_mock, httpx_transport
):
response_payload = {
**MOCK_V1_RESPONSE,
"usage": [{"name": "sku_search", "count": 1}],
}
respx_mock.post("https://api.parallel.ai/v1/search").respond(json=response_payload)
response = await litellm.asearch(
query=["AI developments", "machine learning trends"],
search_provider="parallel_ai",
mode="basic",
)
assert response._hidden_params["response_cost"] == pytest.approx(0.005)
@pytest.mark.asyncio
async def test_caller_cannot_supply_provider_usage(self, bundled_cost_map, respx_mock, httpx_transport):
"""`_parallel_ai_usage` prices the request, so a caller must not be able to set it.
The provider reports no usage here, which is the case where a caller-supplied
value would otherwise survive into the cost calculation.
"""
response_payload = {k: v for k, v in MOCK_V1_RESPONSE.items() if k != "usage"}
route = respx_mock.post("https://api.parallel.ai/v1/search").respond(json=response_payload)
response = await litellm.asearch(
query="AI developments",
search_provider="parallel_ai",
mode="basic",
_parallel_ai_usage=[{"name": "sku_search", "count": 0}],
)
assert response._hidden_params["response_cost"] == pytest.approx(0.005)
assert "_parallel_ai_usage" not in json.loads(route.calls[0].request.content)

191
tests/test_litellm/llms/parallel_ai/test_parallel_ai_search_gateway.py Normal file
View file

@ -0,0 +1,191 @@
"""Gateway coverage for Parallel AI Search."""
from __future__ import annotations
from collections.abc import Iterator
from typing import Final
from unittest.mock import AsyncMock
import httpx
import pytest
from fastapi.testclient import TestClient
import litellm
from litellm import Router
from litellm.integrations.websearch_interception.handler import (
WebSearchInterceptionLogger,
)
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.proxy import proxy_server
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.types.utils import LlmProviders
PARALLEL_SEARCH_URL: Final = "https://api.parallel.ai/v1/search"
@pytest.fixture
def client() -> TestClient:
return TestClient(proxy_server.app, raise_server_exceptions=False)
@pytest.fixture
def auth_as() -> Iterator[None]:
async def _authorized_request() -> UserAPIKeyAuth:
return UserAPIKeyAuth(
api_key="hashed-sk-test",
user_id="parallel-test-user",
)
previous: Final = proxy_server.app.dependency_overrides.get(user_api_key_auth)
proxy_server.app.dependency_overrides[user_api_key_auth] = _authorized_request
try:
yield
finally:
if previous is None:
proxy_server.app.dependency_overrides.pop(user_api_key_auth, None)
else:
proxy_server.app.dependency_overrides[user_api_key_auth] = previous
def _parallel_search_body() -> dict[str, object]:
return {
"search_id": "search_parallel_gateway",
"results": [
{
"url": "https://example.com/parallel",
"title": "Parallel result",
"publish_date": "2026-08-13",
"excerpts": ["First excerpt", "Second excerpt"],
}
],
"usage": [{"name": "sku_search", "count": 1}],
}
def _parallel_router(mode: str = "turbo") -> Router:
return Router(
model_list=[],
search_tools=[
{
"search_tool_name": "parallel-search",
"litellm_params": {
"search_provider": "parallel_ai",
"api_key": "parallel-search-key",
"mode": mode,
},
}
],
num_retries=0,
)
def _mock_async_post(
monkeypatch,
*,
url: str,
response_body: dict[str, object],
) -> AsyncMock:
response = httpx.Response(
status_code=200,
json=response_body,
request=httpx.Request("POST", url),
)
mock_post = AsyncMock(return_value=response)
monkeypatch.setattr(AsyncHTTPHandler, "post", mock_post)
return mock_post
def test_parallel_search_gateway_route(client, auth_as, monkeypatch):
"""The named search route selects its configured Parallel Search tool.
The tool-level `mode` must survive the router hop, so the upstream request
is sent as `turbo` rather than falling back to the adapter default.
"""
monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
monkeypatch.setattr(proxy_server, "llm_router", _parallel_router())
mock_post = _mock_async_post(
monkeypatch,
url=PARALLEL_SEARCH_URL,
response_body=_parallel_search_body(),
)
response = client.post(
"/v1/search/parallel-search",
json={"query": "Parallel AI news", "max_results": 3},
)
assert response.status_code == 200, response.text
assert response.json()["results"] == [
{
"title": "Parallel result",
"url": "https://example.com/parallel",
"snippet": "First excerpt ... Second excerpt",
"date": "2026-08-13",
"last_updated": None,
"excerpts": ["First excerpt", "Second excerpt"],
}
]
request_kwargs = mock_post.await_args.kwargs
assert request_kwargs["url"] == PARALLEL_SEARCH_URL
assert request_kwargs["headers"]["x-api-key"] == "parallel-search-key"
assert request_kwargs["json"] == {
"objective": "Parallel AI news",
"search_queries": ["Parallel AI news"],
"mode": "turbo",
"advanced_settings": {"max_results": 3},
}
@pytest.mark.asyncio
async def test_web_search_interception_executes_parallel_search(monkeypatch):
"""An intercepted web-search call uses the configured Parallel Search tool."""
monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
monkeypatch.setattr(proxy_server, "llm_router", _parallel_router(mode="fast"))
mock_post = _mock_async_post(
monkeypatch,
url=PARALLEL_SEARCH_URL,
response_body=_parallel_search_body(),
)
logger = WebSearchInterceptionLogger(
enabled_providers=[LlmProviders.OPENAI],
search_tool_name="parallel-search",
)
plan = await logger.async_build_responses_agentic_loop_plan(
tools={
"tool_calls": [
{
"id": "fc_parallel",
"call_id": "fc_parallel",
"type": "function_call",
"name": "litellm_web_search",
"arguments": '{"query":"Parallel AI news"}',
"input": {"query": "Parallel AI news"},
}
]
},
model="gpt-5",
messages=[{"role": "user", "content": "Research Parallel"}],
response=None,
optional_params={"tools": [{"type": "function", "name": "litellm_web_search"}]},
logging_obj=None,
stream=False,
kwargs={"custom_llm_provider": "openai"},
)
assert plan.run_agentic_loop is True
assert plan.request_patch is not None
assert plan.request_patch.messages[-1] == {
"type": "function_call_output",
"call_id": "fc_parallel",
"output": (
"Title: Parallel result\nURL: https://example.com/parallel\nSnippet: First excerpt ... Second excerpt"
),
}
request_kwargs = mock_post.await_args.kwargs
assert request_kwargs["url"] == PARALLEL_SEARCH_URL
assert request_kwargs["headers"]["x-api-key"] == "parallel-search-key"
assert request_kwargs["json"]["mode"] == "fast"

3
tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py
View file

@ -1096,10 +1096,13 @@ def test_natively_signed_parallel_turn_never_carries_a_placeholder(model):
"gemini-3.5-flash",
"gemini-3.6-flash",
"gemini-3.7-flash",
"gemini-3.8-flash",
"vertex_ai/gemini-3.5-flash",
"vertex_ai/gemini-3.7-flash",
"vertex_ai/gemini-3.8-flash",
"gemini/gemini-3.5-flash",
"gemini/gemini-3.7-flash",
"gemini/gemini-3.8-flash",
],
)
def test_placeholder_scoped_to_first_call_across_gemini_3_variants(model):

12
tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py
View file

@ -1185,6 +1185,18 @@ def test_vertex_ai_map_thinking_param_with_budget_tokens_0():
}
def test_vertex_ai_map_thinking_param_without_budget_tokens_for_gemini_3():
v = VertexGeminiConfig()
result = v.map_openai_params(
non_default_params={"thinking": {"type": "enabled"}},
optional_params={},
model="gemini-3.5-flash",
drop_params=False,
)
assert result["thinkingConfig"] == {"includeThoughts": True}
def test_vertex_ai_map_tools():
v = VertexGeminiConfig()
optional_params = {}

11
tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py
View file

@ -1431,7 +1431,11 @@ class TestListToolsRestAPI:
async def test_aggregate_list_absorbs_one_server_auth_failure(self, monkeypatch):
"""The multi-server aggregate listing degrades a server whose upstream
rejects auth to an empty contribution and still returns the healthy
server's tools with a 200, rather than surfacing a 401."""
server's tools with a 200, rather than surfacing a 401. The absorbed
server must still show up as a classified per-server outcome so a REST
caller can tell "needs upstream auth" apart from "has no tools"."""
from pydantic import TypeAdapter
from litellm.proxy._experimental.mcp_server.exceptions import (
MCPUpstreamAuthError,
)
@ -1497,6 +1501,11 @@ class TestListToolsRestAPI:
assert result["tools"] == ["good-tool"]
assert result["error"] is None
wire_body = json.loads(TypeAdapter(dict).dump_json(result))
assert wire_body["server_outcomes"] == {
"good": {"status": "ok", "tool_count": 1},
"bad": {"status": "auth_required", "http_status": 401},
}
async def test_name_resolution_finds_server_by_uuid(self, monkeypatch):
"""When server_id is a name string, it should be resolved to its UUID

7
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py
View file

@ -5524,7 +5524,9 @@ async def test_streaming_end_of_stream_block_emits_error_frame_instead_of_trunca
"""Regression for PR #38722: a topicPolicy DENY caught by the end-of-stream
scan used to raise after SSE headers were flushed, so the client saw a
silently truncated stream. The unified hook must emit the chat in-stream
error frame instead."""
error frame instead. The finish chunk is withheld while the end-of-stream
scan runs, so on a block it is dropped rather than relayed before the
frame."""
from litellm.llms import load_guardrail_translation_mappings
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail import (
unified_guardrail as unified_module,
@ -5582,8 +5584,9 @@ async def test_streaming_end_of_stream_block_emits_error_frame_instead_of_trunca
finally:
unified_module.endpoint_guardrail_translation_mappings = None
assert len(out) == 3
assert len(out) == 2
assert isinstance(out[0], ModelResponseStream)
assert out[0].choices[0].finish_reason is None
frame = out[-1]
assert isinstance(frame, bytes)
payload = json.loads(frame.decode()[len("data: ") :])

327
tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_openai_streaming_block.py Normal file
View file

@ -0,0 +1,327 @@
"""
Regression tests for blocking an OpenAI-format streaming response from the
unified guardrail post-call streaming iterator hook.
When a guardrail's ``apply_guardrail`` raises ``ModifyResponseException``
while (or at the end of) a chat completions or Responses API stream is being
relayed, the hook must emit a well-formed SSE termination sequence carrying
the block message - NOT a bare ``data: {"error": ...}`` blob that surfaces as
an HTTP 500 error frame and truncates the stream.
"""
import json
from typing import Any, AsyncGenerator, Dict, Literal, Optional, Tuple, Union
import pytest
from litellm.integrations.custom_guardrail import (
CustomGuardrail,
ModifyResponseException,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
UnifiedLLMGuardrails,
)
from litellm.types.utils import (
Delta,
GenericGuardrailAPIInputs,
ModelResponseStream,
StreamingChoices,
)
BLOCK_MESSAGE = "This response was replaced by policy."
JsonPayload = Dict[str, object]
StreamChunk = Union[ModelResponseStream, JsonPayload, bytes]
class _BlockingGuardrail(CustomGuardrail):
"""Mock guardrail that always blocks response scans by raising ModifyResponseException."""
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict,
input_type: Literal["request", "response"],
logging_obj: Optional[Any] = None,
) -> GenericGuardrailAPIInputs:
raise ModifyResponseException(
message=BLOCK_MESSAGE,
model="gpt-5.4-mini",
request_data=request_data,
guardrail_name=self.guardrail_name,
)
class _PassingGuardrail(CustomGuardrail):
"""Mock guardrail that always lets response scans through unchanged."""
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict,
input_type: Literal["request", "response"],
logging_obj: Optional[Any] = None,
) -> GenericGuardrailAPIInputs:
return inputs
def _chat_chunk(delta: Delta, finish_reason: Optional[str] = None) -> ModelResponseStream:
return ModelResponseStream(
id="chatcmpl-live",
created=1724900000,
model="gpt-5.4-mini",
choices=[StreamingChoices(index=0, delta=delta, finish_reason=finish_reason)],
)
async def _chat_stream(end: bool) -> AsyncGenerator[ModelResponseStream, None]:
yield _chat_chunk(Delta(role="assistant", content="This "))
for text in ["is ", "the ", "original ", "answer."]:
yield _chat_chunk(Delta(content=text))
if end:
yield _chat_chunk(Delta(), finish_reason="stop")
async def _responses_stream(end: bool) -> AsyncGenerator[JsonPayload, None]:
original_text = "This is the original answer."
response_envelope = {"id": "resp_live", "model": "gpt-5.4-mini", "status": "in_progress", "output": []}
yield {"type": "response.created", "response": response_envelope}
yield {"type": "response.in_progress", "response": response_envelope}
yield {
"type": "response.output_item.added",
"output_index": 0,
"item": {"id": "msg_orig", "type": "message", "role": "assistant", "content": []},
}
yield {
"type": "response.content_part.added",
"item_id": "msg_orig",
"output_index": 0,
"content_index": 0,
"part": {"type": "output_text", "text": "", "annotations": []},
}
for delta in ["This ", "is ", "the ", "original ", "answer."]:
yield {
"type": "response.output_text.delta",
"item_id": "msg_orig",
"output_index": 0,
"content_index": 0,
"delta": delta,
}
yield {
"type": "response.output_text.done",
"item_id": "msg_orig",
"output_index": 0,
"content_index": 0,
"text": original_text,
}
if end:
yield {
"type": "response.completed",
"response": {
"id": "resp_live",
"model": "gpt-5.4-mini",
"status": "completed",
"output": [
{
"id": "msg_orig",
"type": "message",
"role": "assistant",
"status": "completed",
"content": [{"type": "output_text", "text": original_text, "annotations": []}],
}
],
"usage": {"input_tokens": 7, "output_tokens": 21, "total_tokens": 28},
},
}
async def _run_hook(
route: str,
stream: AsyncGenerator[Union[ModelResponseStream, JsonPayload], None],
sampling_rate: int = 1,
end_of_stream_only: bool = False,
buffer_until_moderated: bool = False,
blocks: bool = True,
) -> Tuple[StreamChunk, ...]:
guardrail = (
_BlockingGuardrail(guardrail_name="test-blocking-guardrail", event_hook="post_call")
if blocks
else _PassingGuardrail(guardrail_name="test-passing-guardrail", event_hook="post_call")
)
guardrail.streaming_sampling_rate = sampling_rate
guardrail.streaming_end_of_stream_only = end_of_stream_only
guardrail.streaming_buffer_until_moderated = buffer_until_moderated
unified_guardrail = UnifiedLLMGuardrails()
user_api_key_dict = UserAPIKeyAuth(api_key="test", request_route=route)
request_data = {
"messages": [{"role": "user", "content": "hi"}],
"guardrail_to_apply": guardrail,
"metadata": {"guardrails": [guardrail.guardrail_name]},
}
return tuple(
[
chunk
async for chunk in unified_guardrail.async_post_call_streaming_iterator_hook(
user_api_key_dict=user_api_key_dict,
response=stream,
request_data=request_data,
)
]
)
def _sse_payloads(collected: Tuple[StreamChunk, ...]) -> Tuple[JsonPayload, ...]:
return tuple(
json.loads(line[len("data:") :].strip())
for chunk in collected
if isinstance(chunk, bytes)
for block in chunk.decode().split("\n\n")
for line in block.strip().split("\n")
if line.startswith("data:")
)
def _assert_no_error_frame(collected: Tuple[StreamChunk, ...]) -> None:
raw = "".join(chunk.decode() for chunk in collected if isinstance(chunk, bytes))
assert '"error"' not in raw, f"unexpected error blob in stream: {raw!r}"
@pytest.mark.asyncio
async def test_chat_pre_stream_block_emits_standalone_completion():
"""Block on the first chunk: a standalone completion opens with a role delta
and ends with finish_reason content_filter."""
collected = await _run_hook("/v1/chat/completions", _chat_stream(end=False))
_assert_no_error_frame(collected)
payloads = _sse_payloads(collected)
assert payloads, "no block SSE chunks were emitted"
assert payloads[0]["choices"][0]["delta"] == {"role": "assistant", "content": BLOCK_MESSAGE}
assert payloads[-1]["choices"][0]["finish_reason"] == "content_filter"
@pytest.mark.asyncio
async def test_chat_mid_stream_block_continues_the_completion():
"""Regression for the LIT-6496 500 error frame: after chunks were already
forwarded, the block continues the same completion id and terminates with
finish_reason content_filter instead of raising into an error blob."""
collected = await _run_hook("/v1/chat/completions", _chat_stream(end=False), sampling_rate=5)
_assert_no_error_frame(collected)
forwarded = [chunk for chunk in collected if isinstance(chunk, ModelResponseStream)]
assert forwarded, "original chunks should have streamed before the block"
payloads = _sse_payloads(collected)
assert payloads, "no block SSE chunks were emitted"
assert all(payload["id"] == "chatcmpl-live" for payload in payloads), (
"block chunks must continue the in-progress completion, not start a new one"
)
assert payloads[0]["choices"][0]["delta"] == {"content": BLOCK_MESSAGE}
assert payloads[-1]["choices"][0]["finish_reason"] == "content_filter"
@pytest.mark.asyncio
async def test_chat_end_of_stream_block_terminates_cleanly():
"""Regression for bugbot's finish-ordering finding: in end_of_stream_only
mode the original finish chunk must be withheld until moderation decides,
so a block's content_filter finish is the only stream terminator a client
ever sees - never policy text trailing after finish_reason stop."""
collected = await _run_hook("/v1/chat/completions", _chat_stream(end=True), end_of_stream_only=True)
_assert_no_error_frame(collected)
forwarded = [chunk for chunk in collected if isinstance(chunk, ModelResponseStream)]
assert forwarded, "content chunks still stream to the client before end-of-stream moderation"
assert all(choice.finish_reason is None for chunk in forwarded for choice in chunk.choices), (
"the original finish chunk must be withheld until moderation decides"
)
payloads = _sse_payloads(collected)
assert BLOCK_MESSAGE in json.dumps(payloads)
assert payloads[-1]["choices"][0]["finish_reason"] == "content_filter"
@pytest.mark.asyncio
async def test_chat_end_of_stream_pass_releases_withheld_finish_chunk():
"""When end-of-stream moderation passes, the withheld finish chunk is
released so a clean stream still terminates normally."""
collected = await _run_hook(
"/v1/chat/completions", _chat_stream(end=True), end_of_stream_only=True, blocks=False
)
assert not [chunk for chunk in collected if isinstance(chunk, bytes)], (
"a clean stream must carry no synthetic block frames"
)
forwarded = [chunk for chunk in collected if isinstance(chunk, ModelResponseStream)]
finish_reasons = [choice.finish_reason for chunk in forwarded for choice in chunk.choices]
assert finish_reasons[-1] == "stop", "the withheld finish chunk must be released after moderation passes"
assert all(reason is None for reason in finish_reasons[:-1])
@pytest.mark.asyncio
async def test_responses_buffered_block_emits_full_event_sequence():
"""Buffered moderation blocks before anything streams: a complete synthetic
Responses stream from response.created through response.completed carrying
the block message, with the original content never released."""
collected = await _run_hook("/v1/responses", _responses_stream(end=True), buffer_until_moderated=True)
_assert_no_error_frame(collected)
assert not [chunk for chunk in collected if isinstance(chunk, dict)], (
"buffered original chunks must never be released after a block"
)
payloads = _sse_payloads(collected)
event_types = [payload["type"] for payload in payloads]
assert event_types[0] == "response.created"
assert "response.output_text.delta" in event_types
assert event_types[-1] == "response.completed"
completed = payloads[-1]["response"]
assert completed["status"] == "completed"
assert completed["output"][0]["content"][0]["text"] == BLOCK_MESSAGE
assert completed["usage"] == {"input_tokens": 7, "output_tokens": 21, "total_tokens": 28}
assert "original answer" not in json.dumps(payloads)
@pytest.mark.asyncio
async def test_responses_mid_stream_block_continues_the_response():
"""Regression for the LIT-6496 500 error frame and bugbot's unclosed-item
finding: after events were already forwarded, the block first closes the
output item still open on the wire, then appends the replacement item under
the same response id, and closes with response.completed - never a second
response.created and never a completed response with an item left open."""
collected = await _run_hook("/v1/responses", _responses_stream(end=False))
_assert_no_error_frame(collected)
forwarded = [chunk for chunk in collected if isinstance(chunk, dict)]
forwarded_types = [chunk["type"] for chunk in forwarded]
assert "response.created" in forwarded_types, "original events should have streamed before the block"
payloads = _sse_payloads(collected)
assert payloads, "no block SSE chunks were emitted"
block_types = [payload["type"] for payload in payloads]
assert "response.created" not in block_types, "a mid-stream block must not restart the response"
assert block_types[-1] == "response.completed"
all_events = forwarded + list(payloads)
opened = sorted(event["output_index"] for event in all_events if event["type"] == "response.output_item.added")
closed = sorted(event["output_index"] for event in all_events if event["type"] == "response.output_item.done")
assert opened == closed, "every output item opened on the stream must be closed before response.completed"
original_done_position = block_types.index("response.output_item.done")
block_item_position = block_types.index("response.output_item.added")
assert original_done_position < block_item_position, (
"the in-progress original item must be closed before the block item is appended"
)
assert payloads[original_done_position]["item"]["id"] == "msg_orig"
assert payloads[block_item_position]["output_index"] == 1, (
"the block item must continue after the original output item"
)
completed = payloads[-1]["response"]
assert completed["id"] == "resp_live"
assert completed["output"][0]["content"][0]["text"] == BLOCK_MESSAGE
@pytest.mark.asyncio
async def test_responses_end_of_stream_block_reports_original_usage():
collected = await _run_hook("/v1/responses", _responses_stream(end=True), end_of_stream_only=True)
_assert_no_error_frame(collected)
forwarded_types = [chunk["type"] for chunk in collected if isinstance(chunk, dict)]
assert "response.completed" not in forwarded_types, (
"the original terminal event must be withheld and replaced by the block sequence"
)
payloads = _sse_payloads(collected)
completed = payloads[-1]["response"]
assert payloads[-1]["type"] == "response.completed"
assert completed["id"] == "resp_live"
assert completed["output"][0]["content"][0]["text"] == BLOCK_MESSAGE
assert completed["usage"] == {"input_tokens": 7, "output_tokens": 21, "total_tokens": 28}

3
tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py
View file

@ -1844,7 +1844,8 @@ class TestStreamingHttpErrorFrames:
out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks)
assert out[:2] == chunks
assert out[0] == chunks[0]
assert chunks[1] not in out
frame = out[-1]
assert isinstance(frame, bytes)
text = frame.decode()

386
tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py
View file

@ -1,7 +1,8 @@
import logging
import time
from collections.abc import Mapping
from collections.abc import Callable, Mapping, Sequence
from itertools import chain
from types import MappingProxyType
from typing import Final
from unittest.mock import AsyncMock, MagicMock, call
@ -38,6 +39,7 @@ from litellm.proxy.management_endpoints.scim.scim_v2 import (
get_groups,
get_users,
get_service_provider_config,
merge_placeholder,
patch_group,
patch_team_membership,
patch_user,
@ -52,6 +54,7 @@ from litellm.types.proxy.management_endpoints.scim_v2 import (
SCIMMember,
SCIMPatchOp,
SCIMPatchOperation,
SCIMPlaceholderMergeResult,
SCIMServiceProviderConfig,
SCIMUser,
SCIMUserEmail,
@ -778,13 +781,17 @@ async def test_handle_existing_user_by_email_without_teams_preserves_memberships
"litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user",
AsyncMock(return_value=None),
)
mock_team_member_add = mocker.patch( # test-quality-ok: roster helpers are module-level, not injectable into the helper
"litellm.proxy.management_endpoints.scim.scim_v2.team_member_add",
AsyncMock(),
mock_team_member_add = (
mocker.patch( # test-quality-ok: roster helpers are module-level, not injectable into the helper
"litellm.proxy.management_endpoints.scim.scim_v2.team_member_add",
AsyncMock(),
)
)
mock_team_member_delete = mocker.patch( # test-quality-ok: roster helpers are module-level, not injectable into the helper
"litellm.proxy.management_endpoints.scim.scim_v2.team_member_delete",
AsyncMock(),
mock_team_member_delete = (
mocker.patch( # test-quality-ok: roster helpers are module-level, not injectable into the helper
"litellm.proxy.management_endpoints.scim.scim_v2.team_member_delete",
AsyncMock(),
)
)
new_user_request = NewUserRequest(
@ -1645,6 +1652,25 @@ async def test_update_group_e2e(mocker):
ScimTransformations.transform_litellm_team_to_scim_group.assert_called_once_with(updated_team)
def _rows_by_exact_id(
user_row: Callable[[Mapping[str, str]], LiteLLM_UserTable | MagicMock | None],
) -> Callable[..., tuple[LiteLLM_UserTable | MagicMock, ...]]:
"""``find_many`` stand-in for the classifier's cross-field read on a table where a
member value only ever matches as an exact ``user_id``."""
def rows(where: Mapping[str, object], take: int | None = None) -> tuple[LiteLLM_UserTable | MagicMock, ...]:
clauses: Final = where["OR"]
assert isinstance(clauses, list)
found: Final = tuple(user_row(clause) for clause in clauses if "user_id" in clause)
return tuple(row for row in found if row is not None)
return rows
def _user_row_for(where: Mapping[str, str]) -> LiteLLM_UserTable:
return LiteLLM_UserTable(user_id=where["user_id"])
@pytest.mark.asyncio
async def test_create_group_with_nonexistent_users_rejects(mocker, monkeypatch):
"""
@ -1696,9 +1722,8 @@ async def test_create_group_with_nonexistent_users_rejects(mocker, monkeypatch):
return mock_user
return None # new-user-1 and new-user-2 don't exist
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=mock_user_lookup)
mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None)
mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=[])
mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(side_effect=_rows_by_exact_id(mock_user_lookup))
# Mock dependencies
mocker.patch(
@ -1782,9 +1807,8 @@ async def test_update_group_with_nonexistent_users_rejects(mocker, monkeypatch):
return mock_user
return None # new-user-3 and new-user-4 don't exist
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=mock_user_lookup)
mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None)
mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=[])
mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(side_effect=_rows_by_exact_id(mock_user_lookup))
# Mock dependencies
mocker.patch(
@ -1853,9 +1877,8 @@ async def test_create_group_with_nonexistent_users_creates_when_flag_true(mocker
return mock_user
return None # new-user-1 and new-user-2 don't exist
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=mock_user_lookup)
mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None)
mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=[])
mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(side_effect=_rows_by_exact_id(mock_user_lookup))
# Mock user creation
created_user_1 = NewUserResponse(user_id="new-user-1", key="test-key-1")
@ -1943,9 +1966,8 @@ async def test_extract_group_member_ids_with_flag_true_creates_users(mocker, mon
return mock_user
return None # new-user-1 doesn't exist
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=mock_user_lookup)
mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None)
mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=[])
mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(side_effect=_rows_by_exact_id(mock_user_lookup))
# Mock user creation
created_user = NewUserResponse(user_id="new-user-1", key="test-key-1")
@ -2013,9 +2035,8 @@ async def test_extract_group_member_ids_with_flag_false_rejects(mocker, monkeypa
return mock_user
return None # new-user-1 doesn't exist
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=mock_user_lookup)
mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None)
mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=[])
mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(side_effect=_rows_by_exact_id(mock_user_lookup))
# Mock dependencies
mocker.patch(
@ -3121,8 +3142,7 @@ async def test_process_group_patch_operations_add_retains_existing_members(mocke
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
# new-user already exists in the DB
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock(user_id="new-user"))
mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=())
mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=(mocker.MagicMock(user_id="new-user"),))
_, final_members, _ = await _process_group_patch_operations(
patch_ops=patch_ops,
@ -3415,8 +3435,7 @@ async def test_patch_group_add_applies_delta_and_keeps_concurrent_add(mocker):
)
mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=final_team)
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock())
mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=())
mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(side_effect=_rows_by_exact_id(_user_row_for))
mocker.patch(
"litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception",
@ -3509,8 +3528,7 @@ async def test_patch_group_replace_stays_absolute_against_concurrent_roster(mock
)
mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=final_team)
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock())
mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=())
mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(side_effect=_rows_by_exact_id(_user_row_for))
mocker.patch(
"litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception",
@ -3640,8 +3658,7 @@ async def test_process_group_patch_add_filtered_path_without_value(mocker):
prisma_client = mocker.MagicMock()
prisma_client.db = mocker.MagicMock()
prisma_client.db.litellm_usertable = mocker.MagicMock()
prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=LiteLLM_UserTable(user_id="user-3"))
prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=())
prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=(LiteLLM_UserTable(user_id="user-3"),))
_, final_members, _ = await _process_group_patch_operations(
patch_ops=patch_ops,
@ -3733,12 +3750,14 @@ def _member_resolution_prisma(
starts folding it, fails here instead of passing.
A caller that must know which accounts match rather than merely how many
passes take=None, so an unbounded read returns every match.
passes take=None, so an unbounded read returns every match. The row keyed by
the value comes last, the order a bounded read is least prepared for, since
the database promises no order at all.
"""
clauses: Final = where["OR"]
assert isinstance(clauses, list)
fields: Final = tuple(next(iter(clause)) for clause in clauses)
assert fields == ("sso_user_id", "user_email"), fields
assert fields in (("user_id", "sso_user_id", "user_email"), ("sso_user_id", "user_email")), fields
def comparison(clause: Mapping[str, object]) -> tuple[str, bool]:
"""The needle and whether production asked for a case-insensitive compare,
@ -3749,8 +3768,9 @@ def _member_resolution_prisma(
assert isinstance(criterion, dict), criterion
return criterion["equals"], criterion.get("mode") == "insensitive"
sso_needle, sso_insensitive = comparison(clauses[0])
email_needle, email_insensitive = comparison(clauses[1])
by_field: Final = dict(zip(fields, (comparison(clause) for clause in clauses)))
sso_needle, sso_insensitive = by_field["sso_user_id"]
email_needle, email_insensitive = by_field["user_email"]
def same(stored: str, needle: str, insensitive: bool) -> bool:
return stored.casefold() == needle.casefold() if insensitive else stored == needle
@ -3768,6 +3788,11 @@ def _member_resolution_prisma(
if same(email, email_needle, email_insensitive)
for user_id in user_ids
),
(
user_id
for user_id in users
if "user_id" in by_field and same(user_id, by_field["user_id"][0], by_field["user_id"][1])
),
)
)
found: Final = tuple(dict.fromkeys(matched))
@ -4452,9 +4477,11 @@ async def test_create_group_applies_default_team_params(
"litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception",
AsyncMock(return_value=_member_resolution_prisma(mocker, users=set(), teams=set())),
)
new_team_mock = mocker.patch( # test-quality-ok: endpoint collaborators are module-level, not injectable into create_group
"litellm.proxy.management_endpoints.scim.scim_v2.new_team",
AsyncMock(return_value=mocker.MagicMock()),
new_team_mock = (
mocker.patch( # test-quality-ok: endpoint collaborators are module-level, not injectable into create_group
"litellm.proxy.management_endpoints.scim.scim_v2.new_team",
AsyncMock(return_value=mocker.MagicMock()),
)
)
mocker.patch( # test-quality-ok: endpoint collaborators are module-level, not injectable into create_group
"litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_team_to_scim_group",
@ -4611,9 +4638,15 @@ async def test_resolve_group_member_ids_dedupes_repeated_member(mocker, scim_ups
def _identity_lookup(value: str) -> object:
"""The single cross-field lookup the classifier is expected to issue."""
"""The single cross-field lookup the classifier is expected to issue per member."""
return call(
where={"OR": [{"sso_user_id": value}, {"user_email": {"equals": value, "mode": "insensitive"}}]},
where={
"OR": [
{"user_id": value},
{"sso_user_id": value},
{"user_email": {"equals": value, "mode": "insensitive"}},
]
},
take=2,
)
@ -4903,9 +4936,7 @@ async def test_process_group_patch_remove_by_the_id_the_directory_added_with(
@pytest.mark.asyncio
async def test_process_group_patch_remove_still_drops_a_placeholder_by_its_literal_id(
mocker, scim_upsert_user_enabled
):
async def test_process_group_patch_remove_still_drops_a_placeholder_by_its_literal_id(mocker, scim_upsert_user_enabled):
"""An earlier release put unmatched ids on the roster verbatim, so a remove has to
keep clearing the id as written even once it also resolves."""
patch_ops = SCIMPatchOp(
@ -4916,7 +4947,10 @@ async def test_process_group_patch_remove_still_drops_a_placeholder_by_its_liter
team_id="parent-group",
team_alias="Parent Group",
members=[],
members_with_roles=[Member(user_id="legacy@example.com", role="user"), Member(user_id="keep-user", role="user")],
members_with_roles=[
Member(user_id="legacy@example.com", role="user"),
Member(user_id="keep-user", role="user"),
],
)
_, final_members, _ = await _process_group_patch_operations(
@ -5081,11 +5115,8 @@ async def test_process_group_patch_remove_refuses_when_two_members_share_the_id(
assert "more than one member of this group" in str(exc_info.value.detail)
@pytest.mark.asyncio
async def test_resolve_group_member_ids_exact_user_id_wins_when_it_names_nobody_else(
mocker, scim_upsert_user_enabled
):
async def test_resolve_group_member_ids_exact_user_id_wins_when_it_names_nobody_else(mocker, scim_upsert_user_enabled):
"""The canonical user id stays authoritative, including when the same account also
holds that value as its email, which is how a SCIM-provisioned account is keyed."""
prisma_client = _member_resolution_prisma(
@ -5147,10 +5178,79 @@ async def test_resolve_group_member_ids_refuses_a_user_id_that_names_another_acc
assert exc_info.value.status_code == 400
assert "member-id" in str(exc_info.value.detail)
create_user_mock.assert_not_called()
assert any(
record.levelno >= logging.WARNING and "someone-else" in record.getMessage() for record in caplog.records
assert any(record.levelno >= logging.WARNING and "someone-else" in record.getMessage() for record in caplog.records)
@pytest.mark.asyncio
async def test_resolve_group_member_ids_reads_the_exact_id_when_two_other_accounts_fill_the_lookup(
mocker, scim_upsert_user_enabled
):
"""A value that is one account's id and two other accounts' identities fills the
bounded lookup with the other two. The account keyed by the value must still be
found, or the id would lose its precedence and a non-canonical type would skip
a member that names a real user."""
prisma_client = _member_resolution_prisma(
mocker,
users={"shared"},
teams=set(),
sso_user_id_to_user_id={"shared": "by-sso"},
email_to_user_id={"shared": "by-email"},
)
create_user_mock = mocker.patch( # test-quality-ok: user creation is module-level, not injectable into the resolver
"litellm.proxy.management_endpoints.scim.scim_v2._create_user_if_not_exists",
AsyncMock(return_value=None),
)
with pytest.raises(HTTPException) as exc_info:
await _resolve_group_member_ids(
members=[SCIMMember(value="shared", type="direct")],
created_via="scim_group_membership",
prisma_client=prisma_client,
)
assert exc_info.value.status_code == 400
assert "shared" in str(exc_info.value.detail)
create_user_mock.assert_not_called()
assert prisma_client.db.litellm_usertable.find_many.await_args_list == [_identity_lookup("shared")]
prisma_client.db.litellm_usertable.find_unique.assert_awaited_once_with(where={"user_id": "shared"})
@pytest.mark.asyncio
async def test_resolve_group_member_ids_reads_the_user_table_once_per_member(mocker, scim_upsert_user_enabled):
"""Every member costs one read of the user table, however it resolves: by its exact
id (which still outranks a non-canonical type), by identity, as a SCIM team, or not
at all. Looking the exact id up on its own before the identity read doubled the
reads of a push, and the identity read is a scan."""
prisma_client = _member_resolution_prisma(
mocker,
users={"by-id"},
teams={"by-team"},
email_to_user_id={"by-email@example.com": "email-user"},
)
mocker.patch( # test-quality-ok: user creation is module-level, not injectable into the resolver
"litellm.proxy.management_endpoints.scim.scim_v2._create_user_if_not_exists",
AsyncMock(return_value=NewUserResponse(user_id="nobody", key="key")),
)
result = await _resolve_group_member_ids(
members=[
SCIMMember(value="by-id", type="direct"),
SCIMMember(value="by-email@example.com"),
SCIMMember(value="by-team"),
SCIMMember(value="nobody"),
],
created_via="scim_group_membership",
prisma_client=prisma_client,
)
assert result.all_member_ids == ["by-id", "email-user", "nobody"]
prisma_client.db.litellm_usertable.find_unique.assert_not_awaited()
assert prisma_client.db.litellm_usertable.find_many.await_args_list == [
_identity_lookup("by-id"),
_identity_lookup("by-email@example.com"),
_identity_lookup("by-team"),
_identity_lookup("nobody"),
]
@pytest.mark.asyncio
@ -5536,10 +5636,7 @@ async def test_resolve_group_member_ids_admits_member_created_concurrently(mocke
the member is still admitted: the id resolves to a real user row, so failing
or dropping it would be wrong either way."""
prisma_client = _member_resolution_prisma(mocker, users=set(), teams=set())
prisma_client.db.litellm_usertable.find_unique = AsyncMock(
side_effect=[None, LiteLLM_UserTable(user_id="raced-user")]
)
prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=())
prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=LiteLLM_UserTable(user_id="raced-user"))
mocker.patch(
"litellm.proxy.management_endpoints.scim.scim_v2._create_user_if_not_exists",
AsyncMock(return_value=None),
@ -5619,3 +5716,196 @@ async def test_patch_group_404s_when_team_deleted_mid_request(mocker):
assert exc_info.value.code == "404"
assert f"Group not found with ID: {group_id}" in exc_info.value.message
_SHADOW_MEMBER_VALUE: Final = "00u1shadow"
_SHADOWED_ACCOUNT: Final = "real-1"
_SHADOWED_GROUP: Final = "grp-eng"
def _shadowed_tenant_rows() -> tuple[LiteLLM_UserTable, ...]:
"""A placeholder keyed by the raw member value, and the real account that value names by SSO id."""
return (
LiteLLM_UserTable(user_id=_SHADOW_MEMBER_VALUE, user_email=_SHADOW_MEMBER_VALUE, teams=[_SHADOWED_GROUP]),
LiteLLM_UserTable(user_id=_SHADOWED_ACCOUNT, user_email="alice@example.com", sso_user_id=_SHADOW_MEMBER_VALUE),
)
def _shadow_tenant_prisma(
mocker: MockerFixture,
*,
rows: Sequence[LiteLLM_UserTable],
keys_owned_by: Mapping[str, int] = MappingProxyType({}),
) -> MagicMock:
"""Prisma fake whose user rows are live: deleting one removes it from every later lookup."""
users: Final[dict[str, LiteLLM_UserTable]] = {row.user_id: row for row in rows}
team: Final = LiteLLM_TeamTable(
team_id=_SHADOWED_GROUP,
members=[_SHADOW_MEMBER_VALUE],
members_with_roles=[Member(user_id=_SHADOW_MEMBER_VALUE, role="user")],
metadata={SCIM_MANAGED_TEAM_METADATA_KEY: True},
)
async def find_unique(where: Mapping[str, str]) -> LiteLLM_UserTable | None:
return users.get(where["user_id"])
def clause_matches(row: LiteLLM_UserTable, clause: Mapping[str, object]) -> bool:
if "user_id" in clause:
return row.user_id == clause["user_id"]
if "sso_user_id" in clause:
return row.sso_user_id == clause["sso_user_id"]
email_filter: Final = clause["user_email"]
assert isinstance(email_filter, dict)
return (row.user_email or "").casefold() == str(email_filter["equals"]).casefold()
async def identity_rows(where: Mapping[str, object], take: int | None = None) -> tuple[LiteLLM_UserTable, ...]:
clauses: Final = where["OR"]
assert isinstance(clauses, list)
matched: Final = tuple(row for row in users.values() if any(clause_matches(row, clause) for clause in clauses))
return matched[:take] if take else matched
async def delete(where: Mapping[str, str]) -> LiteLLM_UserTable | None:
return users.pop(where["user_id"], None)
async def keys_for(where: Mapping[str, object]) -> tuple[MagicMock, ...]:
return tuple(mocker.MagicMock() for _ in range(keys_owned_by.get(str(where["user_id"]), 0)))
async def team_lookup(where: Mapping[str, str]) -> LiteLLM_TeamTable | None:
return team if where["team_id"] == team.team_id else None
prisma_client = mocker.MagicMock()
prisma_client.db = mocker.MagicMock()
prisma_client.db.litellm_usertable = mocker.MagicMock()
prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=find_unique)
prisma_client.db.litellm_usertable.find_many = AsyncMock(side_effect=identity_rows)
prisma_client.db.litellm_usertable.delete = AsyncMock(side_effect=delete)
prisma_client.db.litellm_teamtable = mocker.MagicMock()
prisma_client.db.litellm_teamtable.find_unique = AsyncMock(side_effect=team_lookup)
prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=team)
prisma_client.db.litellm_verificationtoken = mocker.MagicMock()
prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(side_effect=keys_for)
prisma_client.db.litellm_invitationlink = mocker.MagicMock(delete_many=AsyncMock(return_value=0))
prisma_client.db.litellm_organizationmembership = mocker.MagicMock(delete_many=AsyncMock(return_value=0))
prisma_client.db.litellm_teammembership = mocker.MagicMock(delete_many=AsyncMock(return_value=0))
return prisma_client
@pytest.fixture
def shadowed_tenant(mocker, monkeypatch, scim_upsert_user_enabled) -> MagicMock:
from litellm.proxy import proxy_server
prisma_client: Final = _shadow_tenant_prisma(mocker, rows=_shadowed_tenant_rows())
monkeypatch.setattr(proxy_server, "prisma_client", prisma_client)
return prisma_client
async def _push_shadow_member(prisma_client: MagicMock):
return await _resolve_group_member_ids(
members=[SCIMMember(value=_SHADOW_MEMBER_VALUE)],
created_via="scim_group_membership",
prisma_client=prisma_client,
)
@pytest.mark.asyncio
async def test_merge_placeholder_hands_the_group_to_the_shadowed_account(mocker, shadowed_tenant):
"""Every group push of the shadowing value is refused until the placeholder is folded into
the real account; after the merge the same push resolves to that account."""
team_member_add_mock = (
mocker.patch( # test-quality-ok: roster helpers are module-level, not injectable into the endpoint
"litellm.proxy.management_endpoints.scim.scim_v2.team_member_add", AsyncMock()
)
)
team_member_delete_mock = (
mocker.patch( # test-quality-ok: roster helpers are module-level, not injectable into the endpoint
"litellm.proxy.management_endpoints.scim.scim_v2.team_member_delete", AsyncMock()
)
)
with pytest.raises(HTTPException) as before:
await _push_shadow_member(shadowed_tenant)
assert before.value.status_code == 400
result: Final = await merge_placeholder(user_id=_SHADOW_MEMBER_VALUE)
assert result == SCIMPlaceholderMergeResult(
placeholder_user_id=_SHADOW_MEMBER_VALUE,
merged_into_user_id=_SHADOWED_ACCOUNT,
team_ids=(_SHADOWED_GROUP,),
)
added: Final = team_member_add_mock.call_args.kwargs["data"]
assert (added.team_id, added.member.user_id) == (_SHADOWED_GROUP, _SHADOWED_ACCOUNT)
dropped: Final = team_member_delete_mock.call_args.kwargs["data"]
assert (dropped.team_id, dropped.user_id) == (_SHADOWED_GROUP, _SHADOW_MEMBER_VALUE)
shadowed_tenant.db.litellm_teammembership.delete_many.assert_awaited_once_with(
where={"user_id": _SHADOW_MEMBER_VALUE}
)
shadowed_tenant.db.litellm_usertable.delete.assert_awaited_once_with(where={"user_id": _SHADOW_MEMBER_VALUE})
after: Final = await _push_shadow_member(shadowed_tenant)
assert after.all_member_ids == [_SHADOWED_ACCOUNT]
assert after.created_users == []
@pytest.mark.asyncio
async def test_merge_placeholder_keeps_the_placeholder_when_the_roster_write_fails(mocker, shadowed_tenant):
"""If the real account cannot join the team, the placeholder stays on it, or the membership is gone
from both accounts."""
mocker.patch( # test-quality-ok: roster helpers are module-level, not injectable into the endpoint
"litellm.proxy.management_endpoints.scim.scim_v2.team_member_add",
AsyncMock(side_effect=Exception("database connection lost")),
)
team_member_delete_mock = (
mocker.patch( # test-quality-ok: roster helpers are module-level, not injectable into the endpoint
"litellm.proxy.management_endpoints.scim.scim_v2.team_member_delete", AsyncMock()
)
)
with pytest.raises(ProxyException):
await merge_placeholder(user_id=_SHADOW_MEMBER_VALUE)
team_member_delete_mock.assert_not_awaited()
shadowed_tenant.db.litellm_usertable.delete.assert_not_awaited()
assert await shadowed_tenant.db.litellm_usertable.find_unique(where={"user_id": _SHADOW_MEMBER_VALUE}) is not None
@pytest.mark.parametrize(
("rows", "keys_owned_by", "merged", "reason"),
[
pytest.param(_shadowed_tenant_rows(), {}, _SHADOWED_ACCOUNT, "SSO identity of its own", id="real-account"),
pytest.param(
_shadowed_tenant_rows(), {_SHADOW_MEMBER_VALUE: 2}, _SHADOW_MEMBER_VALUE, "2 virtual keys", id="owns-keys"
),
pytest.param(_shadowed_tenant_rows()[:1], {}, _SHADOW_MEMBER_VALUE, "shadows no account", id="names-nobody"),
pytest.param(
(*_shadowed_tenant_rows(), LiteLLM_UserTable(user_id="real-2", user_email=_SHADOW_MEMBER_VALUE.upper())),
{},
_SHADOW_MEMBER_VALUE,
"names 2 accounts (real-1, real-2)",
id="names-two-accounts",
),
],
)
@pytest.mark.asyncio
async def test_merge_placeholder_refuses_rows_that_are_not_a_lone_placeholder(
mocker, monkeypatch, scim_upsert_user_enabled, rows, keys_owned_by, merged, reason
):
"""Only a row with no SSO identity and no keys whose id names exactly one other account is folded;
anything else could move memberships to the wrong person, so nothing is written."""
from litellm.proxy import proxy_server
prisma_client: Final = _shadow_tenant_prisma(mocker, rows=rows, keys_owned_by=keys_owned_by)
monkeypatch.setattr(proxy_server, "prisma_client", prisma_client)
team_member_add_mock = (
mocker.patch( # test-quality-ok: roster helpers are module-level, not injectable into the endpoint
"litellm.proxy.management_endpoints.scim.scim_v2.team_member_add", AsyncMock()
)
)
with pytest.raises(ProxyException) as exc_info:
await merge_placeholder(user_id=merged)
assert int(exc_info.value.code) == 409
assert reason in str(exc_info.value.message)
team_member_add_mock.assert_not_awaited()
prisma_client.db.litellm_usertable.delete.assert_not_awaited()

57
tests/test_litellm/proxy/management_helpers/test_access_group_key_sync.py Normal file
View file

@ -0,0 +1,57 @@
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock
import pytest
from litellm.proxy.db.prisma_client import PrismaWrapper
from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper
from litellm.proxy.management_helpers.access_group_key_sync import (
sync_key_access_group_membership,
sync_key_regeneration_access_group_membership,
)
def _routed_prisma_client():
writer_inner = MagicMock(name="writer_prisma")
reader_inner = MagicMock(name="reader_prisma")
writer_inner.query_raw = AsyncMock(return_value=[])
reader_inner.query_raw = AsyncMock(return_value=[])
writer = PrismaWrapper(original_prisma=writer_inner, iam_token_db_auth=False)
reader = PrismaWrapper(original_prisma=reader_inner, iam_token_db_auth=False)
routing = RoutingPrismaWrapper(writer=writer, reader=reader)
return SimpleNamespace(db=routing), writer_inner, reader_inner
@pytest.mark.asyncio
async def test_regeneration_repoint_update_runs_on_the_writer():
prisma_client, writer_inner, reader_inner = _routed_prisma_client()
await sync_key_regeneration_access_group_membership(
prisma_client=prisma_client,
previous_key_token="old-token",
new_key_token="new-token",
data=None,
existing_key_row=MagicMock(),
)
writer_inner.query_raw.assert_awaited_once()
assert writer_inner.query_raw.await_args.args[0].startswith('UPDATE "LiteLLM_AccessGroupTable"')
reader_inner.query_raw.assert_not_awaited()
@pytest.mark.asyncio
async def test_membership_attach_and_detach_updates_run_on_the_writer():
prisma_client, writer_inner, reader_inner = _routed_prisma_client()
await sync_key_access_group_membership(
prisma_client=prisma_client,
key_token="token",
previous_access_group_ids=["ag-old"],
updated_access_group_ids=["ag-new"],
)
assert writer_inner.query_raw.await_count == 2
assert all(
call.args[0].startswith('UPDATE "LiteLLM_AccessGroupTable"') for call in writer_inner.query_raw.await_args_list
)
reader_inner.query_raw.assert_not_awaited()

111
tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py
View file

@ -2,6 +2,7 @@ import asyncio
import json
import logging
import os
from collections.abc import Callable
from contextlib import ExitStack, contextmanager
from io import BytesIO
from types import SimpleNamespace
@ -29,6 +30,7 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
websocket_passthrough_request,
)
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY,
@ -5464,7 +5466,10 @@ def test_the_marker_check_distinguishes_the_two_route_kinds():
assert request_dispatched_to_pass_through_endpoint(builtin) is False
async def _drive_passthrough_request_and_capture_logging(user_api_key_dict: UserAPIKeyAuth) -> tuple[int, object]:
async def _drive_passthrough_request_and_capture_logging(
user_api_key_dict: UserAPIKeyAuth,
on_pre_call: Callable[[LiteLLMLoggingObj | None], None] | None = None,
) -> tuple[int, LiteLLMLoggingObj | None]:
import litellm
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.types.llms.custom_http import httpxSpecialProvider
@ -5487,10 +5492,12 @@ async def _drive_passthrough_request_and_capture_logging(user_api_key_dict: User
mock_request.query_params = QueryParams({})
mock_request.body = AsyncMock(return_value=b'{"model": "gemini-2.0-flash"}')
captured_data: dict = {}
captured_data: dict = {} # mutable-ok: the pre-call hook records the request data into it
async def capture_pre_call_hook(user_api_key_dict, data, call_type):
captured_data.update(data)
if on_pre_call is not None:
on_pre_call(data.get("litellm_logging_obj"))
return data
mock_proxy_logging = MagicMock()
@ -5623,3 +5630,103 @@ async def test_resolve_team_callback_wiring_fails_open_on_operational_error():
assert wiring.success_callbacks is None
assert wiring.failure_callbacks is None
assert wiring.logging_kwargs is None
@pytest.mark.asyncio
async def test_pass_through_request_leaves_guardrail_readable_metadata():
"""A pre-call guardrail reads the request headers off the passthrough logging
params without raising."""
from litellm.proxy.guardrails.guardrail_hooks.hiddenlayer.hiddenlayer import (
_logged_request_headers,
)
user_api_key_dict = UserAPIKeyAuth(
api_key="test-key",
team_id="test-team",
team_metadata={
"logging": [
{
"callback_name": "langfuse",
"callback_type": "success_and_failure",
"callback_vars": {
"langfuse_public_key": "pk_test",
"langfuse_secret_key": "sk_test",
},
}
]
},
)
observed: dict[str, dict[str, str] | BaseException] = {} # mutable-ok: the pre-call hook records into it
def read_headers_the_way_a_guardrail_does(logging_obj: LiteLLMLoggingObj | None) -> None:
assert logging_obj is not None
try:
observed["headers"] = _logged_request_headers(logging_obj)
except Exception as exc: # noqa: BLE001 - the regression is that this used to raise
observed["headers"] = exc
status_code, logging_obj = await _drive_passthrough_request_and_capture_logging(
user_api_key_dict, on_pre_call=read_headers_the_way_a_guardrail_does
)
assert "headers" in observed, "the pre-call hook never ran, so nothing was observed"
assert observed["headers"] == {}, f"guardrail header read failed: {observed['headers']!r}"
assert status_code == 200
assert logging_obj is not None
assert logging_obj.dynamic_success_callbacks, "team success callbacks must stay wired"
assert logging_obj.standard_callback_dynamic_params.get("langfuse_public_key") == "pk_test"
@pytest.mark.asyncio
async def test_pass_through_request_leaves_cost_router_logger_working():
"""The cost router's logger reads the deployment id off the passthrough logging
params without raising. least_busy shares the read but swallows the exception,
so this is the strategy where the break is observable."""
from litellm._logging import verbose_logger
from litellm.caching.caching import DualCache
from litellm.router_strategy.lowest_cost import LowestCostLoggingHandler
handler = LowestCostLoggingHandler(router_cache=DualCache())
user_api_key_dict = UserAPIKeyAuth(
api_key="test-key",
team_id="test-team",
team_metadata={
"logging": [
{
"callback_name": "langfuse",
"callback_type": "success_and_failure",
"callback_vars": {
"langfuse_public_key": "pk_test",
"langfuse_secret_key": "sk_test",
},
}
]
},
)
status_code, logging_obj = await _drive_passthrough_request_and_capture_logging(user_api_key_dict)
assert status_code == 200
assert logging_obj is not None
raised: list[logging.LogRecord] = [] # mutable-ok: logging.Handler records into it
class _RecordTracebacks(logging.Handler):
def emit(self, record: logging.LogRecord) -> None:
if record.exc_info is not None:
raised.append(record)
recorder = _RecordTracebacks()
verbose_logger.addHandler(recorder)
try:
await handler.async_log_success_event(
kwargs=logging_obj.model_call_details,
response_obj=None,
start_time=None,
end_time=None,
)
finally:
verbose_logger.removeHandler(recorder)
assert not raised, f"cost router logger raised on the passthrough logging params: {raised[0].exc_info}"

78
tests/test_litellm/proxy/proxy_server/test_routes_models.py
View file

@ -45,6 +45,7 @@ def patched_models(monkeypatch):
deployment = MagicMock()
deployment.litellm_params.model = "gpt-4"
router.get_deployment_by_model_group_name = MagicMock(return_value=deployment)
router.get_configured_display_name = MagicMock(return_value=None)
monkeypatch.setattr(proxy_server, "llm_router", router)
monkeypatch.setattr(proxy_server, "prisma_client", MagicMock())
@ -187,6 +188,83 @@ def test_anthropic_format_carries_router_configured_token_limits(client, auth_as
assert (claude["max_input_tokens"], claude["max_tokens"]) == (500000, 4096)
@pytest.mark.parametrize("path", ["/v1/models", "/models"])
def test_anthropic_format_uses_configured_display_name(client, auth_as, patched_models, path):
"""A deployment's ``model_info.display_name`` becomes the Anthropic-native
``display_name`` so Claude Code's picker shows a clean name while the id keeps
routing; models without one keep the id fallback, and the OpenAI-shaped
listing carries no display_name either way."""
def _configured(model_name):
return "Kimi K3" if model_name == "gpt-4" else None
patched_models.get_configured_display_name = MagicMock(side_effect=_configured)
with auth_as():
anthropic_response = client.get(path, headers={"anthropic-version": "2023-06-01"})
openai_response = client.get(path)
assert anthropic_response.status_code == 200
gpt_4, claude = anthropic_response.json()["data"]
assert (gpt_4["id"], gpt_4["display_name"]) == ("gpt-4", "Kimi K3")
assert (claude["id"], claude["display_name"]) == ("claude-sonnet", "claude-sonnet")
assert openai_response.status_code == 200
openai_models = openai_response.json()["data"]
assert [m["id"] for m in openai_models] == ["gpt-4", "claude-sonnet"]
assert all("display_name" not in m for m in openai_models)
@pytest.mark.parametrize("params", [{}, {"scope": "expand"}])
def test_anthropic_display_name_resolved_via_internal_team_key(
client, auth_as, patched_models, monkeypatch, params
):
"""For a team-scoped row the configured display name must be looked up by the
internal routing key while the entry itself is keyed by the public name, so
the clean name lands on the id the client actually sees."""
from litellm.proxy import utils as proxy_utils
from litellm.proxy.auth import model_checks
internal_name = "model_name_team-1_c0ffee"
patched_models.get_model_list = MagicMock(
return_value=[
{
"model_name": internal_name,
"model_info": {
"team_id": "team-1",
"team_public_model_name": "gpt-4-team",
},
}
]
)
patched_models.get_model_names = MagicMock(return_value=[internal_name])
patched_models.get_configured_display_name = MagicMock(
side_effect=lambda model_name: "Team GPT" if model_name == internal_name else None
)
async def _fake_get_available_models_for_user(**kwargs):
return [internal_name]
monkeypatch.setattr(
proxy_utils,
"get_available_models_for_user",
_fake_get_available_models_for_user,
)
monkeypatch.setattr(
model_checks, "get_complete_model_list", lambda **kwargs: [internal_name]
)
with auth_as():
response = client.get(
"/v1/models", params=params, headers={"anthropic-version": "2023-06-01"}
)
assert response.status_code == 200
(entry,) = response.json()["data"]
assert (entry["id"], entry["display_name"]) == ("gpt-4-team", "Team GPT")
@pytest.mark.parametrize("path", ["/v1/models", "/models"])
def test_get_models_invalid_scope_returns_400(client, auth_as, patched_models, path):
"""Pins: ``GET /v1/models``, ``GET /models`` (error path: invalid scope)."""

26
tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py
View file

@ -19,7 +19,10 @@ from litellm.proxy._types import (
LitellmUserRoles,
UserAPIKeyAuth,
)
from litellm.proxy.common_utils.model_listing_utils import TeamModelNameTranslator
from litellm.proxy.common_utils.model_listing_utils import (
TeamModelNameTranslator,
configured_display_names,
)
from litellm.proxy.proxy_server import (
_get_proxy_model_info,
_translate_model_name_for_response,
@ -1391,6 +1394,27 @@ def test_resolve_public_name_respects_legacy_flag():
)
def test_configured_display_names_keyed_by_response_id():
"""The map is keyed by the public response id while the router lookup uses
the internal routing key, and entries without a configured name are omitted."""
router = MagicMock()
router.get_configured_display_name = MagicMock(
side_effect=lambda model_name: "Team Sonnet" if model_name == "model_name_team-abc-123_4a6b8" else None
)
assert configured_display_names(
entries=[
("team-claude-sonnet", "model_name_team-abc-123_4a6b8"),
("gpt-4o", "gpt-4o"),
],
llm_router=router,
) == {"team-claude-sonnet": "Team Sonnet"}
def test_configured_display_names_empty_without_router():
assert configured_display_names(entries=[("gpt-4o", "gpt-4o")], llm_router=None) == {}
@pytest.mark.asyncio
async def test_retrieve_model_by_public_name_returns_200(monkeypatch):
"""Regression: `GET /v1/models/{public_name}` must NOT 404. The listing

93
tests/test_litellm/rerank_api/test_main.py
View file

@ -111,6 +111,99 @@ def test_together_rerank_honors_api_base(respx_mock: respx.MockRouter):
assert mock_route.calls[0].request.headers["authorization"] == "Bearer fake-together-key"
DASHSCOPE_404_BODY = {
"error": {
"message": "The model `does-not-exist` does not exist or you do not have access to it.",
"type": "invalid_request_error",
"param": None,
"code": "model_not_found",
},
"request_id": "mock-request-id",
}
def test_rerank_error_names_provider_and_keeps_body(respx_mock: respx.MockRouter, monkeypatch):
"""Regression for the rerank error path mapping with the unresolved provider param:
a provider 404 surfaced as 'None - ' instead of naming the provider and its error body."""
monkeypatch.delenv("DASHSCOPE_API_BASE", raising=False)
monkeypatch.delenv("DASHSCOPE_API_BASE_RERANK", raising=False)
mock_route = respx_mock.post("https://dashscope.example/v1/reranks")
mock_route.return_value = httpx.Response(404, json=DASHSCOPE_404_BODY)
with pytest.raises(litellm.NotFoundError) as exc_info:
litellm.rerank(
model="dashscope/does-not-exist",
query=MARKER_QUERY,
documents=[MARKER_DOC],
api_key="fake-dashscope-key",
api_base="https://dashscope.example/v1",
)
assert mock_route.called
assert "DashscopeException" in str(exc_info.value)
assert "does not exist or you do not have access to it" in str(exc_info.value)
assert "None - " not in str(exc_info.value)
@pytest.mark.asyncio
async def test_arerank_error_is_mapped_to_litellm_exception(respx_mock: respx.MockRouter, monkeypatch):
"""Regression for arerank's bare re-raise: provider errors escaped as raw
provider exception classes instead of the mapped litellm exception contract."""
monkeypatch.delenv("DASHSCOPE_API_BASE", raising=False)
monkeypatch.delenv("DASHSCOPE_API_BASE_RERANK", raising=False)
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
mock_route = respx_mock.post("https://dashscope.example/v1/reranks")
mock_route.return_value = httpx.Response(404, json=DASHSCOPE_404_BODY)
with pytest.raises(litellm.NotFoundError) as exc_info:
await litellm.arerank(
model="dashscope/does-not-exist",
query=MARKER_QUERY,
documents=[MARKER_DOC],
api_key="fake-dashscope-key",
api_base="https://dashscope.example/v1",
)
assert mock_route.called
assert "DashscopeException" in str(exc_info.value)
assert "does not exist or you do not have access to it" in str(exc_info.value)
assert "None - " not in str(exc_info.value)
@pytest.mark.asyncio
async def test_arerank_declared_authenticating_provider_skips_resolution(monkeypatch):
"""Regression for the event-loop hazard in arerank's provider pre-resolution:
get_llm_provider runs the blocking OAuth device flow for github_copilot/chatgpt,
so arerank must adopt the declared provider instead of resolving it, while the
except path still maps with that declared provider."""
from litellm.llms.base_llm.chat.transformation import BaseLLMException
resolution_calls = []
def record_resolution(*args, **kwargs):
resolution_calls.append((args, kwargs))
return "gpt-4o", "github_copilot", None, None
def rerank_raises_provider_error(*args, **kwargs):
raise BaseLLMException(status_code=401, message='{"error":"bad key"}')
monkeypatch.setattr(litellm, "get_llm_provider", record_resolution)
monkeypatch.setattr("litellm.rerank_api.main.rerank", rerank_raises_provider_error)
with pytest.raises(litellm.AuthenticationError) as exc_info:
await litellm.arerank(
model="github_copilot/gpt-4o",
query=MARKER_QUERY,
documents=[MARKER_DOC],
)
assert resolution_calls == []
assert "Github_copilotException" in str(exc_info.value)
assert "None - " not in str(exc_info.value)
@pytest.mark.asyncio
async def test_together_rerank_async_honors_env_api_base(respx_mock: respx.MockRouter, monkeypatch):
"""Regression: TOGETHER_AI_API_BASE was honored by chat but ignored by rerank."""

52
tests/test_litellm/responses/test_streaming_iterator.py
View file

@ -326,3 +326,55 @@ def test_run_post_success_hooks_does_not_report_generation_time_as_overhead():
assert iterator.completed_response._hidden_params["_response_ms"] == 10000.0
assert "litellm_overhead_time_ms" not in iterator.completed_response._hidden_params
def _responses_api_response_with_usage() -> ResponsesAPIResponse:
from litellm.types.llms.openai import ResponseAPIUsage
return ResponsesAPIResponse(
id="resp_lit6427",
created_at=int(datetime(2025, 1, 1).timestamp()),
status="completed",
model="mantle-claude",
object="response",
output=[],
usage=ResponseAPIUsage(input_tokens=20, output_tokens=60, total_tokens=80),
)
def test_stamp_responses_usage_cost_stamps_computed_cost():
from litellm.responses.streaming_iterator import _stamp_responses_usage_cost
response = _responses_api_response_with_usage()
logging_obj = Mock(spec=LiteLLMLoggingObj)
logging_obj._response_cost_calculator.return_value = 0.000704
_stamp_responses_usage_cost(response, logging_obj)
assert getattr(response.usage, "cost", None) == pytest.approx(0.000704)
logging_obj._response_cost_calculator.assert_called_once_with(result=response)
def test_stamp_responses_usage_cost_keeps_provider_reported_cost():
from litellm.responses.streaming_iterator import _stamp_responses_usage_cost
response = _responses_api_response_with_usage()
setattr(response.usage, "cost", 0.5)
logging_obj = Mock(spec=LiteLLMLoggingObj)
_stamp_responses_usage_cost(response, logging_obj)
assert getattr(response.usage, "cost", None) == pytest.approx(0.5)
logging_obj._response_cost_calculator.assert_not_called()
def test_stamp_responses_usage_cost_survives_calculator_failure():
from litellm.responses.streaming_iterator import _stamp_responses_usage_cost
response = _responses_api_response_with_usage()
logging_obj = Mock(spec=LiteLLMLoggingObj)
logging_obj._response_cost_calculator.side_effect = RuntimeError("cost map unavailable")
_stamp_responses_usage_cost(response, logging_obj)
assert getattr(response.usage, "cost", None) is None

209
tests/test_litellm/test_bedrock_usgov_pricing.py
View file

@ -26,9 +26,7 @@ import pytest
@pytest.fixture(scope="module")
def model_data():
json_path = os.path.join(
os.path.dirname(__file__), "../../model_prices_and_context_window.json"
)
json_path = os.path.join(os.path.dirname(__file__), "../../model_prices_and_context_window.json")
with open(json_path) as f:
return json.load(f)
@ -51,21 +49,14 @@ def test_usgov_sonnet_4_5_pricing(model_data, model_key):
info = model_data[model_key]
assert info["input_cost_per_token"] == 3.6e-06, (
f"{model_key}: input_cost_per_token should be $3.60/MTok "
f"(got {info['input_cost_per_token']})"
f"{model_key}: input_cost_per_token should be $3.60/MTok (got {info['input_cost_per_token']})"
)
assert (
info["output_cost_per_token"] == 1.8e-05
), f"{model_key}: output_cost_per_token should be $18.00/MTok"
assert (
info["cache_creation_input_token_cost"] == 4.5e-06
), f"{model_key}: 5m cache write should be $4.50/MTok"
assert (
info["cache_creation_input_token_cost_above_1hr"] == 7.2e-06
), f"{model_key}: 1h cache write should be $7.20/MTok"
assert (
info["cache_read_input_token_cost"] == 3.6e-07
), f"{model_key}: cache read should be $0.36/MTok"
assert info["output_cost_per_token"] == 1.8e-05, f"{model_key}: output_cost_per_token should be $18.00/MTok"
assert info["cache_creation_input_token_cost"] == 4.5e-06, f"{model_key}: 5m cache write should be $4.50/MTok"
assert info["cache_creation_input_token_cost_above_1hr"] == 7.2e-06, (
f"{model_key}: 1h cache write should be $7.20/MTok"
)
assert info["cache_read_input_token_cost"] == 3.6e-07, f"{model_key}: cache read should be $0.36/MTok"
def test_usgov_carries_20_percent_premium_over_global(model_data):
@ -84,9 +75,7 @@ def test_usgov_carries_20_percent_premium_over_global(model_data):
"cache_read_input_token_cost",
):
ratio = usgov_info[field] / global_info[field]
assert (
abs(ratio - 1.2) < 1e-9
), f"{field}: us-gov / global ratio is {ratio}, expected 1.2"
assert abs(ratio - 1.2) < 1e-9, f"{field}: us-gov / global ratio is {ratio}, expected 1.2"
# The us-gov.anthropic.* cross-region inference profile is the only us-gov
@ -112,9 +101,7 @@ def test_usgov_cross_region_above_200k_carries_gov_premium(model_data, field, ex
"""
info = model_data[USGOV_CROSS_REGION_KEY]
assert field in info, f"{USGOV_CROSS_REGION_KEY}: missing field {field}"
assert (
info[field] == expected
), f"{USGOV_CROSS_REGION_KEY}: {field} should be {expected} (got {info[field]})"
assert info[field] == expected, f"{USGOV_CROSS_REGION_KEY}: {field} should be {expected} (got {info[field]})"
def test_usgov_cross_region_above_200k_ratio_to_global(model_data):
@ -127,6 +114,176 @@ def test_usgov_cross_region_above_200k_ratio_to_global(model_data):
usgov_info = model_data[USGOV_CROSS_REGION_KEY]
for field in EXPECTED_USGOV_ABOVE_200K:
ratio = usgov_info[field] / global_info[field]
assert (
abs(ratio - 1.2) < 1e-9
), f"{field}: us-gov / global ratio is {ratio}, expected 1.2"
assert abs(ratio - 1.2) < 1e-9, f"{field}: us-gov / global ratio is {ratio}, expected 1.2"
CLAUDE_GOV_EXPECTED = {
"anthropic.claude-sonnet-5": {
"input_cost_per_token": 2.4e-06,
"output_cost_per_token": 1.2e-05,
"cache_creation_input_token_cost": 3e-06,
"cache_creation_input_token_cost_above_1hr": 4.8e-06,
"cache_read_input_token_cost": 2.4e-07,
},
"anthropic.claude-opus-4-8": {
"input_cost_per_token": 6e-06,
"output_cost_per_token": 3e-05,
"cache_creation_input_token_cost": 7.5e-06,
"cache_creation_input_token_cost_above_1hr": 1.2e-05,
"cache_read_input_token_cost": 6e-07,
},
}
USGOV_CLAUDE_KEY_TEMPLATES = {
"bedrock/us-gov-east-1/{base_key}": "bedrock",
"bedrock/us-gov-west-1/{base_key}": "bedrock",
"us-gov.{base_key}": "bedrock_converse",
}
@pytest.mark.parametrize("base_key", CLAUDE_GOV_EXPECTED)
@pytest.mark.parametrize("key_template,expected_provider", USGOV_CLAUDE_KEY_TEMPLATES.items())
def test_usgov_claude_sonnet5_opus48_pricing(model_data, key_template, expected_provider, base_key):
"""Sonnet 5 and Opus 4.8 gov entries, both in-region keys and the us-gov.
geo inference profile the model cards list for GovCloud, must match the
rates AWS publishes on the Bedrock pricing page (1.2x global).
"""
gov_key = key_template.format(base_key=base_key)
assert gov_key in model_data, f"Missing model entry: {gov_key}"
info = model_data[gov_key]
assert info["litellm_provider"] == expected_provider
for field, expected in CLAUDE_GOV_EXPECTED[base_key].items():
assert info[field] == expected, f"{gov_key}: {field} should be {expected} (got {info[field]})"
ratio = info[field] / model_data[base_key][field]
assert abs(ratio - 1.2) < 1e-9, f"{gov_key}: {field} gov/global ratio is {ratio}, expected 1.2"
CONVERSE_GOV_EXPECTED = {
"nvidia.nemotron-nano-3-30b": (7.2e-08, 2.88e-07),
"nvidia.nemotron-nano-12b-v2": (2.4e-07, 7.2e-07),
"nvidia.nemotron-super-3-120b": (1.8e-07, 7.8e-07),
"openai.gpt-oss-20b-1:0": (8.4e-08, 3.6e-07),
"openai.gpt-oss-120b-1:0": (1.8e-07, 7.2e-07),
}
@pytest.mark.parametrize("base_key", CONVERSE_GOV_EXPECTED)
@pytest.mark.parametrize("region", ["us-gov-east-1", "us-gov-west-1"])
def test_usgov_converse_model_pricing(model_data, region, base_key):
"""Nemotron and gpt-oss gov entries must match the AWS Bedrock offer file,
which prices both GovCloud regions identically at 1.2x commercial.
"""
gov_key = f"bedrock/{region}/{base_key}"
assert gov_key in model_data, f"Missing model entry: {gov_key}"
info = model_data[gov_key]
expected_input, expected_output = CONVERSE_GOV_EXPECTED[base_key]
assert info["input_cost_per_token"] == expected_input
assert info["output_cost_per_token"] == expected_output
assert info["litellm_provider"] == "bedrock"
base = model_data[base_key]
assert abs(info["input_cost_per_token"] / base["input_cost_per_token"] - 1.2) < 1e-9
assert abs(info["output_cost_per_token"] / base["output_cost_per_token"] - 1.2) < 1e-9
def test_usgov_west_llama3_8b_output_price_fixed(model_data):
"""The us-gov-west-1 llama3-8b entry carried the 70B output rate ($2.65/MTok);
the AWS Bedrock offer file prices output at $0.60/MTok. AWS lists the model
in us-gov-west-1 only, so there is no east entry to check.
"""
info = model_data["bedrock/us-gov-west-1/meta.llama3-8b-instruct-v1:0"]
assert info["input_cost_per_token"] == 3e-07
assert info["output_cost_per_token"] == 6e-07
MANTLE_GOV_TIERED_EXPECTED = {
"openai.gpt-5.6-luna": {
"input_cost_per_token": 2.64e-07,
"input_cost_per_token_above_272k_tokens": 5.28e-07,
"cache_creation_input_token_cost": 3.3e-07,
"cache_creation_input_token_cost_above_272k_tokens": 6.6e-07,
"cache_read_input_token_cost": 2.64e-08,
"cache_read_input_token_cost_above_272k_tokens": 5.28e-08,
"output_cost_per_token": 1.584e-06,
"output_cost_per_token_above_272k_tokens": 2.376e-06,
},
"openai.gpt-5.6-terra": {
"input_cost_per_token": 2.64e-06,
"input_cost_per_token_above_272k_tokens": 5.28e-06,
"cache_creation_input_token_cost": 3.3e-06,
"cache_creation_input_token_cost_above_272k_tokens": 6.6e-06,
"cache_read_input_token_cost": 2.64e-07,
"cache_read_input_token_cost_above_272k_tokens": 5.28e-07,
"output_cost_per_token": 1.584e-05,
"output_cost_per_token_above_272k_tokens": 2.376e-05,
},
}
@pytest.mark.parametrize("model", MANTLE_GOV_TIERED_EXPECTED)
def test_usgov_west_mantle_terra_luna_pricing(model_data, model):
"""Terra and Luna carry 1.2x commercial across every tier in the
us-gov-west-1 offer file; the us-gov-east-1 offer file has no SKUs for them.
"""
gov_key = f"bedrock_mantle/us-gov-west-1/{model}"
assert gov_key in model_data, f"Missing model entry: {gov_key}"
info = model_data[gov_key]
for field, expected in MANTLE_GOV_TIERED_EXPECTED[model].items():
assert info[field] == expected, f"{gov_key}: {field} should be {expected} (got {info[field]})"
assert info["litellm_provider"] == "bedrock_mantle"
assert f"bedrock_mantle/us-gov-east-1/{model}" not in model_data
@pytest.mark.parametrize("region", ["us-gov-east-1", "us-gov-west-1"])
def test_usgov_mantle_gpt_5_4_pricing_has_no_long_context_tier(model_data, region):
"""gpt-5.4 gov rates come from the offer file, which publishes only the
standard tier in GovCloud: no long-context SKUs exist there, unlike commercial.
"""
gov_key = f"bedrock_mantle/{region}/openai.gpt-5.4"
assert gov_key in model_data, f"Missing model entry: {gov_key}"
info = model_data[gov_key]
assert info["input_cost_per_token"] == 3.3e-06
assert info["cache_read_input_token_cost"] == 3.3e-07
assert info["output_cost_per_token"] == 1.98e-05
assert not any(field.endswith("_above_272k_tokens") for field in info)
def test_usgov_mantle_grok_4_3_west_only(model_data):
"""grok-4.3 is priced in the us-gov-west-1 offer file only; the east offer
file carries grok-4.6 instead.
"""
info = model_data["bedrock_mantle/us-gov-west-1/xai.grok-4.3"]
assert info["input_cost_per_token"] == 1.5e-06
assert info["output_cost_per_token"] == 3e-06
assert info["cache_read_input_token_cost"] == 2.4e-07
assert "bedrock_mantle/us-gov-east-1/xai.grok-4.3" not in model_data
AZURE_GOV_EXPECTED = {
"azure/us-gov/gpt-5.1": {
"input_cost_per_token": 1.71875e-06,
"cache_read_input_token_cost": 1.71875e-07,
"output_cost_per_token": 1.375e-05,
},
"azure/us-gov/o3-mini": {
"input_cost_per_token": 1.513e-06,
"cache_read_input_token_cost": 7.57e-07,
"output_cost_per_token": 6.05e-06,
},
"azure/us-gov/text-embedding-3-large": {"input_cost_per_token": 1.63e-07},
"azure/us-gov/text-embedding-3-small": {"input_cost_per_token": 2.5e-08},
}
@pytest.mark.parametrize("gov_key", AZURE_GOV_EXPECTED)
def test_azure_usgov_pricing(model_data, gov_key):
"""Azure Government meters from the Azure retail prices API
(usgovvirginia/usgovarizona, serviceName 'Foundry Models'). No Government
retirement schedule is published, so these entries carry no deprecation_date.
"""
assert gov_key in model_data, f"Missing model entry: {gov_key}"
info = model_data[gov_key]
for field, expected in AZURE_GOV_EXPECTED[gov_key].items():
assert info[field] == expected, f"{gov_key}: {field} should be {expected} (got {info[field]})"
assert info["litellm_provider"] == "azure"
assert "deprecation_date" not in info

16
tests/test_litellm/test_cloudflare_workers_ai_model_metadata.py
View file

@ -75,6 +75,22 @@ def test_additional_current_models_are_present():
assert entry["output_cost_per_token"] > 0
@pytest.mark.parametrize(
"key, published_price_per_audio_minute",
[
("cloudflare/@cf/openai/whisper", 0.00045),
("cloudflare/@cf/openai/whisper-large-v3-turbo", 0.00051),
],
)
def test_whisper_transcription_pricing_is_stored_per_second(key, published_price_per_audio_minute):
entry = litellm.model_cost[key]
assert entry["litellm_provider"] == "cloudflare"
assert entry["mode"] == "audio_transcription"
assert entry["supported_endpoints"] == ["/v1/audio/transcriptions"]
assert entry["output_cost_per_second"] == 0.0
assert entry["input_cost_per_second"] == pytest.approx(published_price_per_audio_minute / 60)
def test_root_and_backup_have_identical_cloudflare_keys():
if not os.path.exists(ROOT_MAP):
pytest.skip("root cost map only ships in source checkouts")

56
tests/test_litellm/test_dockerfile_bedrock_realtime_extra.py Normal file
View file

@ -0,0 +1,56 @@
"""
Static checks that every proxy Docker image installs the `bedrock-realtime` extra.
Bedrock Nova Sonic speech-to-speech (`/v1/realtime`) needs `aws-sdk-bedrock-runtime`,
which only ships in the `bedrock-realtime` extra. An image whose `uv sync` stages
omit the extra fails every Nova Sonic realtime session with
"Missing aws_sdk_bedrock_runtime. Install with: pip install aws-sdk-bedrock-runtime".
"""
import os
import re
from typing import Final
import pytest
REPO_ROOT: Final = os.path.join(os.path.dirname(__file__), "..", "..")
PROXY_DOCKERFILES: Final = (
"Dockerfile",
os.path.join("docker", "Dockerfile.non_root"),
os.path.join("docker", "Dockerfile.database"),
os.path.join("gateway", "Dockerfile"),
)
CONTINUED_LINE_RE: Final = re.compile(r"(?:\\\n|[^\n])+")
UV_SYNC_BOUNDARY_RE: Final = re.compile(r"(?=uv sync)")
def _uv_sync_invocations(dockerfile_text: str) -> tuple[str, ...]:
"""Return each `uv sync ...` command, split apart when one RUN holds several (if/else branches)."""
return tuple(
part
for line in CONTINUED_LINE_RE.finditer(dockerfile_text)
for part in UV_SYNC_BOUNDARY_RE.split(line.group(0))
if part.startswith("uv sync")
)
@pytest.mark.parametrize("relative_path", PROXY_DOCKERFILES)
def test_every_uv_sync_installs_bedrock_realtime_extra(relative_path: str):
dockerfile_path: Final = os.path.join(REPO_ROOT, relative_path)
if not os.path.exists(dockerfile_path):
pytest.skip(f"{relative_path} not present in this checkout")
with open(dockerfile_path, "r", encoding="utf-8") as f:
contents: Final = f.read()
invocations: Final = _uv_sync_invocations(contents)
assert invocations, f"{relative_path} has no `uv sync` invocation"
missing: Final = tuple(invocation for invocation in invocations if "--extra bedrock-realtime" not in invocation)
assert not missing, (
f"{relative_path}: {len(missing)} of {len(invocations)} `uv sync` invocations omit "
"`--extra bedrock-realtime`, so aws-sdk-bedrock-runtime is absent and Bedrock Nova Sonic "
"/v1/realtime sessions fail with 'Missing aws_sdk_bedrock_runtime'"
)

66
tests/test_litellm/test_main.py
View file

@ -3150,8 +3150,8 @@ def _stream_builder_logging_obj() -> LiteLLMLogging:
return logging_obj
def test_stream_chunk_builder_reports_streaming_usage_cost_when_enabled(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True)
def test_stream_chunk_builder_stamps_streaming_usage_cost_by_default(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", False)
chunks: Final = [
_stream_builder_text_chunk("gpt-4o", "Hello "),
_stream_builder_text_chunk("gpt-4o", "world.", finish_reason="stop"),
@ -3168,11 +3168,45 @@ def test_stream_chunk_builder_reports_streaming_usage_cost_when_enabled(monkeypa
assert response._hidden_params["response_cost"] == pytest.approx(usage_cost)
def test_stream_chunk_builder_defers_cost_to_logging_obj_when_usage_cost_absent(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", False)
def test_stream_chunk_builder_skips_stamp_when_cost_is_unpriceable():
import time as time_module
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
logging_obj: Final = LiteLLMLogging(
model="us.anthropic.claude-opus-5",
messages=[{"role": "user", "content": "hi"}],
stream=True,
call_type="completion",
start_time=time_module.time(),
litellm_call_id="stream-builder-alias-unpriceable",
function_id="1",
)
logging_obj.model_call_details["custom_llm_provider"] = "bedrock"
logging_obj.optional_params = {}
usage_chunk: Final = _stream_builder_text_chunk("bedrock-claude-opus-5", "")
usage_chunk.usage = Usage(prompt_tokens=40, completion_tokens=5, total_tokens=45)
chunks: Final = [
_stream_builder_text_chunk("bedrock-claude-opus-5", "Hello ", finish_reason="stop"),
usage_chunk,
]
response: Final = litellm.stream_chunk_builder(
chunks=chunks, messages=[{"role": "user", "content": "hi"}], logging_obj=logging_obj
)
assert response is not None
assert getattr(response.usage, "cost", None) is None
assert response._hidden_params.get("response_cost") is None
def test_stream_chunk_builder_keeps_provider_reported_usage_cost():
usage_chunk: Final = _stream_builder_text_chunk("gpt-4o", "")
usage_chunk.usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15, cost=0.5)
chunks: Final = [
_stream_builder_text_chunk("gpt-4o", "Hello "),
_stream_builder_text_chunk("gpt-4o", "world.", finish_reason="stop"),
usage_chunk,
]
response: Final = litellm.stream_chunk_builder(
@ -3180,4 +3214,26 @@ def test_stream_chunk_builder_defers_cost_to_logging_obj_when_usage_cost_absent(
)
assert response is not None
assert response._hidden_params.get("response_cost") is None
assert getattr(response.usage, "cost", None) == pytest.approx(0.5)
assert response._hidden_params["response_cost"] == pytest.approx(0.5)
def test_stream_chunk_builder_prices_alias_from_openai_sdk_usage_chunk():
from openai.types.completion_usage import CompletionUsage
usage_chunk: Final = _stream_builder_text_chunk("mantle-claude", "")
usage_chunk.usage = CompletionUsage(prompt_tokens=20, completion_tokens=60, total_tokens=80, cost=0.000704)
assert type(usage_chunk.usage) is CompletionUsage
chunks: Final = [
_stream_builder_text_chunk("mantle-claude", "Hello "),
_stream_builder_text_chunk("mantle-claude", "world.", finish_reason="stop"),
usage_chunk,
]
response: Final = litellm.stream_chunk_builder(chunks=chunks, messages=[{"role": "user", "content": "hi"}])
assert response is not None
assert response.usage.prompt_tokens == 20
assert response.usage.completion_tokens == 60
assert getattr(response.usage, "cost", None) == pytest.approx(0.000704)
assert response._hidden_params["response_cost"] == pytest.approx(0.000704)

156
tests/test_litellm/test_openai_service_tier_long_context_pricing.py Normal file
View file

@ -0,0 +1,156 @@
import json
from functools import lru_cache
from pathlib import Path
import pytest
import litellm
REPO_ROOT = Path(__file__).parents[2]
MAIN_PATH = REPO_ROOT / "model_prices_and_context_window.json"
BACKUP_PATH = REPO_ROOT / "litellm" / "model_prices_and_context_window_backup.json"
FLEX_LONG_CONTEXT = {
"gpt-5.4": {
"input_cost_per_token_above_272k_tokens_flex": 2.5e-06,
"output_cost_per_token_above_272k_tokens_flex": 1.125e-05,
"cache_read_input_token_cost_above_272k_tokens_flex": 2.5e-07,
},
"gpt-5.4-pro": {
"input_cost_per_token_above_272k_tokens_flex": 3e-05,
"output_cost_per_token_above_272k_tokens_flex": 0.000135,
},
"gpt-5.5": {
"input_cost_per_token_above_272k_tokens_flex": 5e-06,
"output_cost_per_token_above_272k_tokens_flex": 2.25e-05,
"cache_read_input_token_cost_above_272k_tokens_flex": 5e-07,
},
}
PRIORITY_LONG_CONTEXT = {
"gpt-5.6": {
"input_cost_per_token_above_272k_tokens_priority": 1.6e-05,
"output_cost_per_token_above_272k_tokens_priority": 6e-05,
"cache_read_input_token_cost_above_272k_tokens_priority": 1.6e-06,
"cache_creation_input_token_cost_above_272k_tokens_priority": 2e-05,
},
"gpt-5.6-sol": {
"input_cost_per_token_above_272k_tokens_priority": 1.6e-05,
"output_cost_per_token_above_272k_tokens_priority": 6e-05,
"cache_read_input_token_cost_above_272k_tokens_priority": 1.6e-06,
"cache_creation_input_token_cost_above_272k_tokens_priority": 2e-05,
},
"gpt-5.6-terra": {
"input_cost_per_token_above_272k_tokens_priority": 8e-06,
"output_cost_per_token_above_272k_tokens_priority": 3.6e-05,
"cache_read_input_token_cost_above_272k_tokens_priority": 8e-07,
"cache_creation_input_token_cost_above_272k_tokens_priority": 1e-05,
},
"gpt-5.6-luna": {
"input_cost_per_token_above_272k_tokens_priority": 8e-07,
"output_cost_per_token_above_272k_tokens_priority": 3.6e-06,
"cache_read_input_token_cost_above_272k_tokens_priority": 8e-08,
"cache_creation_input_token_cost_above_272k_tokens_priority": 1e-06,
},
}
EXPECTED = {**FLEX_LONG_CONTEXT, **PRIORITY_LONG_CONTEXT}
NO_PUBLISHED_PRIORITY_LONG_CONTEXT = ("gpt-5.4", "gpt-5.5")
@pytest.fixture(autouse=True)
def _local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
@lru_cache(maxsize=2)
def _load(path: Path) -> dict[str, dict[str, object]]:
with open(path) as f:
return json.load(f)
@pytest.mark.parametrize("path", [MAIN_PATH, BACKUP_PATH], ids=["main", "backup"])
@pytest.mark.parametrize("model", sorted(EXPECTED))
def test_service_tier_long_context_rates_are_published(model: str, path: Path) -> None:
"""Each tier must carry its own above-272K rates, in both price files."""
info = _load(path).get(model)
assert info is not None, f"{model} not found in {path.name}"
for key, expected in EXPECTED[model].items():
assert info.get(key) == pytest.approx(expected), f"{model}.{key} is {info.get(key)!r}, expected {expected!r}"
@pytest.mark.parametrize("model", sorted(EXPECTED))
def test_tier_long_context_rate_is_half_or_double_the_standard(model: str) -> None:
"""Flex is half the standard long-context rate; priority is double it."""
info = _load(MAIN_PATH)[model]
tier = "flex" if model in FLEX_LONG_CONTEXT else "priority"
ratio = 0.5 if tier == "flex" else 2.0
for base in ("input_cost_per_token", "output_cost_per_token"):
standard = info[f"{base}_above_272k_tokens"]
tiered = info[f"{base}_above_272k_tokens_{tier}"]
assert tiered == pytest.approx(standard * ratio), (
f"{model}.{base}_above_272k_tokens_{tier} is {tiered!r}, "
f"expected {ratio}x the standard long-context rate {standard!r}"
)
@pytest.mark.parametrize("model", NO_PUBLISHED_PRIORITY_LONG_CONTEXT)
def test_no_priority_long_context_rates_where_openai_publishes_none(model: str) -> None:
"""Guard against back-filling a rate OpenAI does not publish."""
info = _load(MAIN_PATH)[model]
assert "input_cost_per_token_above_272k_tokens_priority" not in info
LONG_CONTEXT_PROMPT_TOKENS = 300_000
COMPLETION_TOKENS = 1_000
TIERED_COST_CASES = [
("gpt-5.4", "flex", 2.5e-06, 1.125e-05),
("gpt-5.4-pro", "flex", 3e-05, 0.000135),
("gpt-5.5", "flex", 5e-06, 2.25e-05),
("gpt-5.6", "priority", 1.6e-05, 6e-05),
("gpt-5.6-sol", "priority", 1.6e-05, 6e-05),
("gpt-5.6-terra", "priority", 8e-06, 3.6e-05),
("gpt-5.6-luna", "priority", 8e-07, 3.6e-06),
]
@pytest.mark.parametrize("model,tier,input_rate,output_rate", TIERED_COST_CASES)
def test_cost_per_token_bills_long_context_at_the_tier_rate(
model: str, tier: str, input_rate: float, output_rate: float
) -> None:
"""A prompt over 272K on flex or priority must bill at that tier's long-context rate."""
input_cost, output_cost = litellm.cost_per_token(
model=model,
prompt_tokens=LONG_CONTEXT_PROMPT_TOKENS,
completion_tokens=COMPLETION_TOKENS,
service_tier=tier,
)
assert input_cost == pytest.approx(LONG_CONTEXT_PROMPT_TOKENS * input_rate)
assert output_cost == pytest.approx(COMPLETION_TOKENS * output_rate)
@pytest.mark.parametrize("model,tier,input_rate,output_rate", TIERED_COST_CASES)
def test_cost_per_token_tier_differs_from_the_standard_long_context_cost(
model: str, tier: str, input_rate: float, output_rate: float
) -> None:
"""Flex halves the standard long-context bill and priority doubles it."""
ratio = 0.5 if tier == "flex" else 2.0
standard = sum(
litellm.cost_per_token(
model=model,
prompt_tokens=LONG_CONTEXT_PROMPT_TOKENS,
completion_tokens=COMPLETION_TOKENS,
)
)
tiered = sum(
litellm.cost_per_token(
model=model,
prompt_tokens=LONG_CONTEXT_PROMPT_TOKENS,
completion_tokens=COMPLETION_TOKENS,
service_tier=tier,
)
)
assert tiered == pytest.approx(standard * ratio)

65
tests/test_litellm/test_router.py
View file

@ -7271,6 +7271,71 @@ def test_get_configured_token_limits_coerces_numeric_strings():
assert router.get_configured_token_limits("quoted-limits-model") == (32000, 8000)
def test_get_configured_display_name_reads_deployment_model_info():
router = litellm.Router(
model_list=[
{
"model_name": "Kimi K3-claude-compatible",
"litellm_params": {"model": "openai/some-unmapped-model"},
"model_info": {"display_name": "Kimi K3"},
}
]
)
assert router.get_configured_display_name("Kimi K3-claude-compatible") == "Kimi K3"
def test_get_configured_display_name_returns_none_for_unset_or_unknown():
router = litellm.Router(
model_list=[
{
"model_name": "no-display-model",
"litellm_params": {"model": "openai/some-unmapped-model"},
}
]
)
assert router.get_configured_display_name("no-display-model") is None
assert router.get_configured_display_name("not-a-real-model") is None
def test_get_configured_display_name_skips_wildcard_pattern_matching():
router = litellm.Router(
model_list=[
{
"model_name": "bedrock/*",
"litellm_params": {"model": "bedrock/*"},
"model_info": {"display_name": "Bedrock"},
}
]
)
with patch.object(
router.pattern_router, "route", side_effect=AssertionError("pattern route called")
):
assert (
router.get_configured_display_name("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0")
is None
)
def test_get_configured_display_name_treats_malformed_values_as_absent():
malformed = ["", " ", 12345, ["Kimi K3"], {"name": "Kimi K3"}, True]
router = litellm.Router(
model_list=[
{
"model_name": f"bad-display-{i}",
"litellm_params": {"model": "openai/some-unmapped-model"},
"model_info": {"display_name": bad},
}
for i, bad in enumerate(malformed)
]
)
for i in range(len(malformed)):
assert router.get_configured_display_name(f"bad-display-{i}") is None
@pytest.mark.asyncio
async def test_acreate_batch_disable_fallbacks_surfaces_owning_provider_error():
router = litellm.Router(

80
tests/test_litellm/test_router_model_cost_isolation.py
View file

@ -2281,3 +2281,83 @@ def test_every_declaring_deployment_is_named(caplog):
assert "azure-ptu-east" in warnings[0]
assert "azure-ptu-west" in warnings[0]
assert "plain-gpt-4o" not in warnings[0]
def _simulate_price_data_reload_with_provider_sets(monkeypatch, fetched_catalog):
"""Like `_simulate_price_data_reload`, plus the provider model-set refresh the proxy's
`_swap_in_model_cost_map` does before replaying, so bare names in the new catalog resolve."""
monkeypatch.setattr(litellm, "model_cost", fetched_catalog)
_invalidate_model_cost_lowercase_map()
litellm.add_known_models(model_cost_map=fetched_catalog)
reapply_runtime_model_cost_registrations()
def test_a_config_deployment_dropped_by_a_stale_cost_map_comes_back_on_reload(monkeypatch):
"""
Booting on the bundled backup, a bare model that only the remote catalog knows
cannot be provider-resolved, so the proxy router (ignore_invalid_deployments) drops
it. Once a reload brings in a catalog that knows the model, the deployment must be
served again with its access groups, and exactly once however many reloads follow.
"""
backend = "lit-5766-only-in-remote-catalog"
try:
router = Router(
model_list=[
{
"model_name": "new-model",
"litellm_params": {"model": backend, "api_key": "k"},
"model_info": {"id": "new-id", "access_groups": ["team-models"]},
},
{
"model_name": "control-model",
"litellm_params": {"model": "hosted_vllm/control-backend", "api_key": "k"},
"model_info": {"id": "control-id", "access_groups": ["team-models"]},
},
],
ignore_invalid_deployments=True,
)
assert router.get_model_names() == ["control-model"]
assert router.get_model_access_groups(model_name="new-model") == {}
fresh_catalog = {**litellm.model_cost, backend: {"litellm_provider": "openai", "mode": "chat"}}
_simulate_price_data_reload_with_provider_sets(monkeypatch, fresh_catalog)
_simulate_price_data_reload_with_provider_sets(monkeypatch, fresh_catalog)
assert sorted(router.get_model_names()) == ["control-model", "new-model"]
assert router.get_model_access_groups(model_name="new-model") == {"team-models": ["new-model"]}
assert [d["model_info"]["id"] for d in router.model_list] == ["control-id", "new-id"]
assert "new-id" in litellm.model_cost
finally:
litellm.open_ai_chat_completion_models.discard(backend)
litellm.models_by_provider["openai"].discard(backend)
def test_a_config_deployment_dropped_for_a_permanent_reason_is_not_retried_on_reload(monkeypatch):
"""
Only provider-resolution drops can be healed by a fresh catalog. A deployment that
fails after its provider resolved (here a pass-through vertex entry with no project)
has already touched router state, so replaying it on every reload would leak into
`deployment_names` each time.
"""
router = Router(
model_list=[
{
"model_name": "vertex-passthrough",
"litellm_params": {"model": "vertex_ai/gemini-2.5-flash", "use_in_pass_through": True},
"model_info": {"id": "vertex-id"},
},
{
"model_name": "control-model",
"litellm_params": {"model": "hosted_vllm/control-backend", "api_key": "k"},
"model_info": {"id": "control-id"},
},
],
ignore_invalid_deployments=True,
)
assert router.get_model_names() == ["control-model"]
names_after_boot = list(router.deployment_names)
_simulate_price_data_reload_with_provider_sets(monkeypatch, dict(litellm.model_cost))
assert router.get_model_names() == ["control-model"]
assert router.deployment_names == names_after_boot

1
tests/test_litellm/test_utils.py
View file

@ -4655,6 +4655,7 @@ GEMINI_4096_CACHE_MIN_MODELS: Final = tuple(
"gemini-3.5-flash",
"gemini-3.6-flash",
"gemini-3.7-flash",
"gemini-3.8-flash",
"gemini-3.1-pro-preview",
"gemini-3.1-pro-preview-customtools",
)

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