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_containers_route_model_list_creds

Browse source
This commit is contained in:
Devin AI 2026-09-02 19:10:45 +00:00
commit 4a09c35d94
148 changed files with 11429 additions and 825 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: {}

32
litellm/constants.py
View file

@ -9,6 +9,38 @@ DEFAULT_HEALTH_CHECK_PROMPT: Final = str(os.getenv("DEFAULT_HEALTH_CHECK_PROMPT"
AZURE_DEFAULT_RESPONSES_API_VERSION: Final = str(os.getenv("AZURE_DEFAULT_RESPONSES_API_VERSION", "preview"))
ROUTER_MAX_FALLBACKS: Final = int(os.getenv("ROUTER_MAX_FALLBACKS", 5))
ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS: Final = 2000
RUNTIME_UPDATABLE_ROUTER_SETTINGS: Final[frozenset[str]] = frozenset(
{
"routing_strategy_args",
"routing_strategy",
"routing_groups",
"allowed_fails",
"cooldown_time",
"num_retries",
"timeout",
"max_retries",
"retry_after",
"fallbacks",
"context_window_fallbacks",
"retry_policy",
"model_group_retry_policy",
"model_group_alias",
"enable_weighted_failover",
"enable_tag_filtering",
"tag_routing_prefix",
"optional_pre_call_checks",
}
)
ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG: Final[frozenset[str]] = frozenset(
{
"model_list",
"search_tools",
"assistants_config",
"router_general_settings",
"ignore_invalid_deployments",
"fallback_access_check",
}
)
DEFAULT_BATCH_SIZE: Final = int(os.getenv("DEFAULT_BATCH_SIZE", 512))
DEFAULT_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_FLUSH_INTERVAL_SECONDS", 5))
DEFAULT_S3_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_S3_FLUSH_INTERVAL_SECONDS", 10))

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,

25
litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py
View file

@ -86,22 +86,41 @@ def _decoded_sse_data_line(line: bytes) -> object | None:
return None
def _anthropic_error_event_payload(chunk: object) -> Mapping[str, object] | None:
def _anthropic_event_payload(chunk: object, event_type: str) -> Mapping[str, object] | None:
if isinstance(chunk, dict):
return chunk if chunk.get("type") == "error" else None
return chunk if chunk.get("type") == event_type else None
if isinstance(chunk, (bytes, bytearray)):
decoded_lines: Final = (_decoded_sse_data_line(line) for line in chunk.splitlines())
return next(
(
candidate
for candidate in decoded_lines
if isinstance(candidate, dict) and candidate.get("type") == "error"
if isinstance(candidate, dict) and candidate.get("type") == event_type
),
None,
)
return None
def _anthropic_error_event_payload(chunk: object) -> Mapping[str, object] | None:
return _anthropic_event_payload(chunk, "error")
def parse_anthropic_refusal_stop_details(chunk: object) -> Mapping[str, object] | None:
"""
Return the ``stop_details`` object of an Anthropic SSE ``message_delta``
chunk whose delta carries ``stop_reason: "refusal"`` (a safeguard refusal:
https://platform.claude.com/docs/en/build-with-claude/refusals-and-fallback),
or None for any other chunk, a plain refusal without ``stop_details`` included.
"""
payload: Final = _anthropic_event_payload(chunk, "message_delta")
delta: Final = payload.get("delta") if payload is not None else None
if not isinstance(delta, dict) or delta.get("stop_reason") != "refusal":
return None
stop_details: Final = delta.get("stop_details")
return stop_details if isinstance(stop_details, dict) else None
def _anthropic_error_body(chunk: object) -> Mapping[str, object] | None:
"""Return the ``error`` object of an Anthropic SSE ``event: error`` chunk, or None."""
payload: Final = _anthropic_error_event_payload(chunk)

49
litellm/llms/anthropic/experimental_pass_through/messages/utils.py
View file

@ -1,11 +1,40 @@
from collections.abc import Mapping
from functools import lru_cache
from typing import Any, Final, cast, get_type_hints
from typing import TYPE_CHECKING, Any, Final, cast, get_type_hints
from litellm.types.llms.anthropic import AnthropicMessagesRequestOptionalParams
from litellm.types.llms.anthropic_messages.anthropic_response import (
AnthropicMessagesResponse,
)
if TYPE_CHECKING:
from litellm.exceptions import ContentPolicyViolationError
def get_safeguard_refusal_stop_details(response: object) -> Mapping[str, Any] | None:
"""
Return the ``stop_details`` of an Anthropic Messages response refused by a
safeguard (``stop_reason: "refusal"`` carrying ``stop_details``:
https://platform.claude.com/docs/en/build-with-claude/refusals-and-fallback),
or None for any other response, a plain refusal without ``stop_details`` included.
"""
if not isinstance(response, dict) or response.get("stop_reason") != "refusal":
return None
stop_details: Final = response.get("stop_details")
return stop_details if isinstance(stop_details, dict) else None
def safeguard_refusal_error(model: str, stop_details: Mapping[str, object]) -> "ContentPolicyViolationError":
"""The exception a safeguard-refused Anthropic response converts into so the
content-policy fallback chain can re-dispatch it."""
from litellm.exceptions import ContentPolicyViolationError
return ContentPolicyViolationError(
message=f"Anthropic safeguard refusal (category: {stop_details.get('category')}).",
model=model,
llm_provider="anthropic",
)
@lru_cache(maxsize=1)
def _anthropic_messages_optional_param_keys() -> frozenset[str]:
@ -100,14 +129,12 @@ def mock_response(
model=model,
)
return AnthropicMessagesResponse(
**{
"content": [{"text": mock_response, "type": "text"}],
"id": "msg_013Zva2CMHLNnXjNJJKqJ2EF",
"model": "claude-sonnet-4-20250514",
"role": "assistant",
"stop_reason": "end_turn",
"stop_sequence": None,
"type": "message",
"usage": {"input_tokens": 2095, "output_tokens": 503},
}
content=[{"text": mock_response, "type": "text"}],
id="msg_013Zva2CMHLNnXjNJJKqJ2EF",
model="claude-sonnet-4-20250514",
role="assistant",
stop_reason="end_turn",
stop_sequence=None,
type="message",
usage={"input_tokens": 2095, "output_tokens": 503},
)

2
litellm/llms/azure_ai/vector_stores/transformation.py
View file

@ -19,6 +19,7 @@ from litellm.types.vector_stores import (
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
from litellm.router import Router
LiteLLMLoggingObj = _LiteLLMLoggingObj
else:
@ -115,6 +116,7 @@ class AzureAIVectorStoreConfig(BaseVectorStoreConfig, BaseAzureLLM):
litellm_logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
extra_body: dict[str, Any] | None = None,
router: "Router | None" = None,
) -> tuple[str, dict[str, Any]]:
"""
Transform search request for Azure AI Search API

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:

5
litellm/llms/base_llm/vector_store/transformation.py
View file

@ -17,6 +17,7 @@ from litellm.types.vector_stores import (
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
from litellm.router import Router
from ..chat.transformation import BaseLLMException as _BaseLLMException
@ -57,6 +58,7 @@ class BaseVectorStoreConfig:
litellm_logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
extra_body: dict[str, Any] | None = None,
router: "Router | None" = None,
) -> tuple[str, dict]:
pass
@ -69,6 +71,7 @@ class BaseVectorStoreConfig:
litellm_logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
extra_body: dict[str, Any] | None = None,
router: "Router | None" = None,
) -> tuple[str, dict]:
"""
Optional async version of transform_search_vector_store_request.
@ -84,6 +87,7 @@ class BaseVectorStoreConfig:
litellm_logging_obj=litellm_logging_obj,
litellm_params=litellm_params,
extra_body=extra_body,
router=router,
)
@abstractmethod
@ -197,6 +201,7 @@ class BaseDirectVectorStoreConfig(BaseVectorStoreConfig):
litellm_logging_obj: LiteLLMLoggingObj,
litellm_params: Mapping[str, object],
extra_body: Mapping[str, object] | None = None,
router: "Router | None" = None,
) -> NoReturn:
raise NotImplementedError("Direct vector store providers execute the search themselves; no HTTP request shape")

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

2
litellm/llms/bedrock/vector_stores/transformation.py
View file

@ -27,6 +27,7 @@ from litellm.types.vector_stores import (
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.router import Router
else:
LiteLLMLoggingObj = Any
@ -196,6 +197,7 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM):
litellm_logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
extra_body: dict[str, Any] | None = None,
router: "Router | None" = None,
) -> tuple[str, dict]:
if isinstance(query, list):
query = " ".join(query)

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

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

@ -178,6 +178,7 @@ if TYPE_CHECKING:
AnthropicMessagesStreamingResponse,
)
from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig
from litellm.router import Router
from litellm.types.llms.openai_evals import (
CancelEvalResponse,
CancelRunResponse,
@ -2923,7 +2924,7 @@ class BaseLLMHTTPHandler:
final_response: Final = await self._call_agentic_completion_hooks(
response=initial_response,
model=model,
messages=(input if isinstance(input, list) else [{"role": "user", "content": input}]),
messages=(input if isinstance(input, list) else [{"role": "user", "content": input}]), # pyright: ignore[reportArgumentType] # pre-existing mismatch surfaced by the Router import; the hook accepts response input items at runtime
anthropic_messages_provider_config=responses_api_provider_config,
anthropic_messages_optional_request_params=response_api_optional_request_params,
logging_obj=logging_obj,
@ -5415,7 +5416,7 @@ class BaseLLMHTTPHandler:
try:
response: ResponsesAPIResponse | BaseResponsesAPIStreamingIterator = await litellm.aresponses(
model=patch.model or model,
input=patch.messages,
input=patch.messages, # pyright: ignore[reportArgumentType] # pre-existing mismatch surfaced by the Router import; patch messages are valid response input at runtime
**optional_params,
**kwargs_for_followup,
)
@ -9688,6 +9689,7 @@ class BaseLLMHTTPHandler:
timeout: float | httpx.Timeout | None = None,
client: HTTPHandler | AsyncHTTPHandler | None = None,
_is_async: bool = False,
router: "Router | None" = None,
) -> VectorStoreSearchResponse:
if isinstance(vector_store_provider_config, BaseDirectVectorStoreConfig):
self._pre_call_direct_vector_store_search(
@ -9738,6 +9740,7 @@ class BaseLLMHTTPHandler:
litellm_logging_obj=logging_obj,
litellm_params=dict(litellm_params),
extra_body=extra_body,
router=router,
)
else:
(
@ -9751,6 +9754,7 @@ class BaseLLMHTTPHandler:
litellm_logging_obj=logging_obj,
litellm_params=dict(litellm_params),
extra_body=extra_body,
router=router,
)
all_optional_params: Final[dict[str, object]] = dict(litellm_params)
all_optional_params.update(vector_store_search_optional_params or {})
@ -9802,6 +9806,7 @@ class BaseLLMHTTPHandler:
timeout: float | httpx.Timeout | None = None,
client: HTTPHandler | AsyncHTTPHandler | None = None,
_is_async: bool = False,
router: "Router | None" = None,
) -> VectorStoreSearchResponse | Coroutine[object, object, VectorStoreSearchResponse]:
if _is_async:
return self.async_vector_store_search_handler(
@ -9816,6 +9821,7 @@ class BaseLLMHTTPHandler:
extra_body=extra_body,
timeout=timeout,
client=client,
router=router,
)
if isinstance(vector_store_provider_config, BaseDirectVectorStoreConfig):
@ -9862,6 +9868,7 @@ class BaseLLMHTTPHandler:
litellm_logging_obj=logging_obj,
litellm_params=dict(litellm_params),
extra_body=extra_body,
router=router,
)
all_optional_params: Final[dict[str, object]] = dict(litellm_params)

2
litellm/llms/gemini/vector_stores/transformation.py
View file

@ -33,6 +33,7 @@ from litellm.types.vector_stores import (
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.router import Router
else:
LiteLLMLoggingObj = Any
@ -168,6 +169,7 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig):
litellm_logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
extra_body: Mapping[str, object] | None = None,
router: "Router | None" = None,
) -> tuple[str, dict]:
"""
Transform search request to Gemini's generateContent format.

2
litellm/llms/milvus/vector_stores/transformation.py
View file

@ -19,6 +19,7 @@ from litellm.types.vector_stores import (
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
from litellm.router import Router
LiteLLMLoggingObj = _LiteLLMLoggingObj
else:
@ -123,6 +124,7 @@ class MilvusVectorStoreConfig(BaseVectorStoreConfig):
litellm_logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
extra_body: dict[str, Any] | None = None,
router: "Router | None" = None,
) -> tuple[str, dict[str, Any]]:
"""
Transform search request for Azure AI Search API

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

2
litellm/llms/openai/vector_stores/transformation.py
View file

@ -21,6 +21,7 @@ from litellm.utils import add_openai_metadata
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
from litellm.router import Router
LiteLLMLoggingObj = _LiteLLMLoggingObj
else:
@ -99,6 +100,7 @@ class OpenAIVectorStoreConfig(BaseVectorStoreConfig):
litellm_logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
extra_body: dict[str, Any] | None = None,
router: "Router | None" = None,
) -> tuple[str, dict]:
encoded_vector_store_id: Final = encode_url_path_segment(vector_store_id, field_name="vector_store_id")
url: Final = f"{api_base}/{encoded_vector_store_id}/search"

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/pg_vector/vector_stores/transformation.py
View file

@ -8,6 +8,7 @@ from litellm.types.vector_stores import VectorStoreSearchOptionalRequestParams
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.router import Router
else:
LiteLLMLoggingObj = Any
@ -80,6 +81,7 @@ class PGVectorStoreConfig(OpenAIVectorStoreConfig):
litellm_logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
extra_body: dict[str, Any] | None = None,
router: "Router | None" = None,
) -> tuple[str, dict]:
encoded_vector_store_id: Final = encode_url_path_segment(vector_store_id, field_name="vector_store_id")
url: Final = f"{api_base}/{encoded_vector_store_id}/search"

2
litellm/llms/ragflow/vector_stores/transformation.py
View file

@ -17,6 +17,7 @@ from litellm.types.vector_stores import (
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.router import Router
else:
LiteLLMLoggingObj = Any
@ -92,6 +93,7 @@ class RAGFlowVectorStoreConfig(BaseVectorStoreConfig):
litellm_logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
extra_body: dict[str, Any] | None = None,
router: "Router | None" = None,
) -> tuple[str, dict]:
"""RAGFlow vector stores are management-only, search is not supported."""
raise NotImplementedError("RAGFlow vector stores support dataset management only, not search/retrieval")

38
litellm/llms/s3_vectors/vector_stores/transformation.py
View file

@ -1,8 +1,8 @@
import re
from typing import TYPE_CHECKING, Any, Final
import httpx
from litellm.caching._embedding_router import resolve_embedding_router
from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.types.router import GenericLiteLLMParams
@ -18,6 +18,7 @@ from litellm.types.vector_stores import (
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.router import Router
else:
LiteLLMLoggingObj = Any
@ -58,13 +59,20 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM):
return headers
def get_complete_url(self, api_base: str | None, litellm_params: dict) -> str:
aws_region_name: Final = litellm_params.get("aws_region_name")
if not aws_region_name:
raise ValueError("aws_region_name is required for S3 Vectors")
if not re.match(r"^[a-z][a-z0-9-]*$", aws_region_name):
raise ValueError("Invalid aws_region_name format")
# Resolve region the same way the ingestion path does:
# dynamic param -> AWS_REGION_NAME -> AWS_REGION -> default (us-west-2)
aws_region_name: Final = self.get_aws_region_name_for_non_llm_api_calls(litellm_params.get("aws_region_name"))
return f"https://s3vectors.{aws_region_name}.api.aws"
def _resolve_query_embedding_router(self, embedding_model: str, router: "Router | None") -> "Router | None":
"""Return the router iff it serves ``embedding_model`` as a deployment."""
if router is None:
return None
model_list: Final = [
dict(m) for m in (router.get_model_list() or ())
] # mutable-ok: resolve_embedding_router requires list[dict]
return resolve_embedding_router(embedding_model=embedding_model, llm_router=router, llm_model_list=model_list)
def transform_search_vector_store_request(
self,
vector_store_id: str,
@ -74,6 +82,7 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM):
litellm_logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
extra_body: dict[str, Any] | None = None,
router: "Router | None" = None,
) -> tuple[str, dict]:
"""Sync version - generates embedding synchronously."""
# For S3 Vectors, vector_store_id should be in format: bucket_name:index_name
@ -99,10 +108,16 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM):
# Generate embedding for the query
embedding_model: Final = litellm_params.get("embedding_model", "text-embedding-3-small")
embedding_router: Final = self._resolve_query_embedding_router(embedding_model=embedding_model, router=router)
import litellm as litellm_module
embedding_response: Final = litellm_module.embedding(model=embedding_model, input=[query])
embedding_input: Final = [query] # mutable-ok: the embedding API takes list input
embedding_response: Final = (
embedding_router.embedding(model=embedding_model, input=embedding_input)
if embedding_router is not None
else litellm_module.embedding(model=embedding_model, input=embedding_input)
)
query_embedding: Final = embedding_response.data[0]["embedding"]
url: Final = f"{api_base}/QueryVectors"
@ -128,6 +143,7 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM):
litellm_logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
extra_body: dict[str, Any] | None = None,
router: "Router | None" = None,
) -> tuple[str, dict]:
"""Async version - generates embedding asynchronously."""
# For S3 Vectors, vector_store_id should be in format: bucket_name:index_name
@ -153,10 +169,16 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM):
# Generate embedding for the query asynchronously
embedding_model: Final = litellm_params.get("embedding_model", "text-embedding-3-small")
embedding_router: Final = self._resolve_query_embedding_router(embedding_model=embedding_model, router=router)
import litellm as litellm_module
embedding_response: Final = await litellm_module.aembedding(model=embedding_model, input=[query])
embedding_input: Final = [query] # mutable-ok: the embedding API takes list input
embedding_response: Final = (
await embedding_router.aembedding(model=embedding_model, input=embedding_input)
if embedding_router is not None
else await litellm_module.aembedding(model=embedding_model, input=embedding_input)
)
query_embedding: Final = embedding_response.data[0]["embedding"]
url: Final = f"{api_base}/QueryVectors"

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]}],
}

2
litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py
View file

@ -21,6 +21,7 @@ from litellm.types.vector_stores import (
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
from litellm.router import Router
LiteLLMLoggingObj = _LiteLLMLoggingObj
else:
@ -161,6 +162,7 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase):
litellm_logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
extra_body: Mapping[str, object] | None = None,
router: "Router | None" = None,
) -> tuple[str, dict[str, object]]:
"""
Transform search request for Vertex AI RAG API

2
litellm/llms/vertex_ai/vector_stores/search_api/transformation.py
View file

@ -25,6 +25,7 @@ from litellm.types.vector_stores import (
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
from litellm.router import Router
LiteLLMLoggingObj = _LiteLLMLoggingObj
else:
@ -245,6 +246,7 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase):
litellm_logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
extra_body: Mapping[str, object] | None = None,
router: "Router | None" = None,
) -> tuple[str, dict[str, object]]:
"""
Transform a search request for the Vertex AI Search (Discovery Engine) API.

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)

1007
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"
]
}
}
}
},

174
litellm/proxy/anthropic_endpoints/streaming_model_restamp.py Normal file
View file

@ -0,0 +1,174 @@
"""
Restamp the public ``model`` on the Anthropic Messages ``message_start`` event, the only
stream event carrying a model, so streamed responses report the requested model like
non-streaming ones do.
Chunks reach the serializer either as already-encoded SSE frames (``bytes``/``str``, the
provider passthrough path) or as event dicts (fake-stream and agentic paths).
"""
import json
import re
from collections.abc import Mapping
from typing import Final
from pydantic import TypeAdapter, ValidationError
_MESSAGE_START_EVENT: Final = "message_start"
_MESSAGE_START_MARKER: Final = b"message_start"
_SSE_DATA_FIELD: Final = "data:"
_SSE_FRAME_END_PATTERN: Final = re.compile(rb"\r\n\r\n|\r\r|\n\n")
_MAX_HELD_BYTES: Final = 65536
_PING_MARKERS: Final = (b"event: ping", b'"type": "ping"', b'"type":"ping"')
_EVENT_ADAPTER: Final = TypeAdapter(Mapping[str, object])
def _restamped_event(event: Mapping[str, object], requested_model: str) -> Mapping[str, object] | None:
message: Final = event.get("message")
if event.get("type") != _MESSAGE_START_EVENT or not isinstance(message, dict):
return None
if message.get("model") == requested_model:
return None
return {**event, "message": {**message, "model": requested_model}} # mutable-ok: SSE payload, re-serialized as is
def _restamped_data_line(line: str, requested_model: str) -> str | None:
stripped: Final = line.strip()
if not stripped.startswith(_SSE_DATA_FIELD):
return None
payload: Final = stripped[len(_SSE_DATA_FIELD) :].strip()
if not payload or payload == "[DONE]":
return None
try:
event: Final = _EVENT_ADAPTER.validate_json(payload)
except ValidationError:
return None
restamped: Final = _restamped_event(event, requested_model)
if restamped is None:
return None
terminator: Final = line[len(line.rstrip("\r\n")) :]
return f"data: {json.dumps(restamped, separators=(',', ':'))}{terminator}"
def _restamped_frame(frame: str, requested_model: str) -> str | None:
lines: Final = frame.splitlines(keepends=True)
restamped: Final = tuple(_restamped_data_line(line, requested_model) for line in lines)
if all(line is None for line in restamped):
return None
return "".join(new if new is not None else old for new, old in zip(restamped, lines))
def restamp_anthropic_stream_chunk_model(chunk: object, requested_model: str) -> object:
"""
Return ``chunk`` with the ``message_start`` model replaced by ``requested_model``.
Chunks that carry no model are returned unchanged.
"""
if isinstance(chunk, dict):
try:
event: Final = _EVENT_ADAPTER.validate_python(chunk)
except ValidationError:
return chunk
return _restamped_event(event, requested_model) or chunk
if isinstance(chunk, (bytes, bytearray)):
if _MESSAGE_START_EVENT.encode() not in chunk:
return chunk
restamped_bytes: Final = _restamped_frame(chunk.decode("utf-8", errors="ignore"), requested_model)
return chunk if restamped_bytes is None else restamped_bytes.encode("utf-8")
if isinstance(chunk, str):
if _MESSAGE_START_EVENT not in chunk:
return chunk
restamped_text: Final = _restamped_frame(chunk, requested_model)
return chunk if restamped_text is None else restamped_text
return chunk
def _is_ping_frame(frame: bytes) -> bool:
return any(marker in frame for marker in _PING_MARKERS)
class AnthropicStreamModelRestamper:
"""
Per-stream restamper for the encoded passthrough path, where chunks are raw
transport reads: the ``message_start`` SSE frame can arrive split across
chunks or coalesced with later frames. Complete frames (``\\n\\n``,
``\\r\\n\\r\\n``, or ``\\r\\r`` terminated) are emitted as their terminator
closes them and an incomplete tail is held until it completes, so the
restamp never misses a torn frame; ``flush`` returns whatever is still held
when the stream ends so no bytes are swallowed. Once ``message_start`` has
been handled, or the first real event proves the stream carries none, every
later chunk passes through untouched.
"""
def __init__(self, requested_model: str) -> None:
self._requested_model: Final = requested_model
self._held = b""
self._armed = True
def process(self, chunk: object) -> object:
if not self._armed:
return chunk
if isinstance(chunk, (bytes, bytearray)):
return self._process_encoded(bytes(chunk))
if isinstance(chunk, str):
return self._process_encoded(chunk.encode("utf-8"))
restamped: Final = restamp_anthropic_stream_chunk_model(chunk, self._requested_model)
if isinstance(chunk, dict) and chunk.get("type") not in (None, "ping"):
self._armed = False
return restamped
def flush(self) -> bytes:
held: Final = self._held
self._held = b""
self._armed = False
if not held:
return b""
restamped: Final = restamp_anthropic_stream_chunk_model(held, self._requested_model)
return restamped if isinstance(restamped, bytes) else held
def _process_encoded(self, data: bytes) -> bytes:
combined: Final = self._held + data
boundaries: Final = tuple(match.end() for match in _SSE_FRAME_END_PATTERN.finditer(combined))
if not boundaries:
if len(combined) > _MAX_HELD_BYTES:
self._held = b""
self._armed = False
return combined
self._held = combined
return b""
emitted: Final = self._restamped_closed_block(combined[: boundaries[-1]])
tail: Final = combined[boundaries[-1] :]
if not self._armed:
self._held = b""
return emitted + tail
self._held = tail
return emitted
def _restamped_closed_block(self, closed: bytes) -> bytes:
boundaries: Final = tuple(match.end() for match in _SSE_FRAME_END_PATTERN.finditer(closed))
frames: Final = tuple(closed[start:end] for start, end in zip((0, *boundaries[:-1]), boundaries))
decider: Final = next(
(
index
for index, frame in enumerate(frames)
if _MESSAGE_START_MARKER in frame or (b"data:" in frame and not _is_ping_frame(frame))
),
None,
)
if decider is None:
return closed
self._armed = False
if _MESSAGE_START_MARKER not in frames[decider]:
return closed
restamped_text: Final = _restamped_frame(
frames[decider].decode("utf-8", errors="ignore"), self._requested_model
)
if restamped_text is None:
return closed
return b"".join(
restamped_text.encode("utf-8") if index == decider else frame for index, frame in enumerate(frames)
)

36
litellm/proxy/common_request_processing.py
View file

@ -176,6 +176,9 @@ if TYPE_CHECKING:
ProxyConfig = _ProxyConfig
else:
ProxyConfig = Any
from litellm.proxy.anthropic_endpoints.streaming_model_restamp import (
AnthropicStreamModelRestamper,
)
from litellm.proxy.litellm_pre_call_utils import (
add_litellm_data_to_request,
refresh_proxy_server_request_body_snapshot,
@ -2490,6 +2493,9 @@ class ProxyBaseLLMRequestProcessing:
request_data=self.data,
proxy_logging_obj=proxy_logging_obj,
request=request,
restamp_model=(
None if _should_return_raw_model_name(self.data) else requested_model_from_client
),
)
return await create_response(
generator=wrap_sse_stream_with_keepalive_pings(
@ -3442,6 +3448,16 @@ class ProxyBaseLLMRequestProcessing:
else:
return chunk
@staticmethod
def _sse_chunk_serializer(restamper: AnthropicStreamModelRestamper | None) -> StreamChunkSerializer:
if restamper is None:
return ProxyBaseLLMRequestProcessing.return_sse_chunk
def serialize(chunk: object) -> str:
return ProxyBaseLLMRequestProcessing.return_sse_chunk(restamper.process(chunk))
return serialize
@staticmethod
async def _finalize_streaming_generator_cleanup(
request: Request | None,
@ -3502,11 +3518,16 @@ class ProxyBaseLLMRequestProcessing:
serialize_chunk: StreamChunkSerializer,
serialize_error: StreamErrorSerializer,
request: Request | None = None,
flush_tail: Callable[[], bytes] | None = None,
) -> AsyncGenerator[str, None]:
"""
Shared streaming data generator: runs proxy iterator hook, per-chunk hook,
cost injection, then yields chunks via serialize_chunk; on exception runs
failure hook and yields via serialize_error. Use for SSE or NDJSON.
``flush_tail`` runs once after the upstream iterator completes cleanly and
its non-empty result is yielded, so a serializer that buffers bytes across
chunks can emit anything still held at end of stream.
"""
verbose_proxy_logger.debug("inside generator")
# Resolve per-stream (not per-chunk) whether the heavy per-chunk path
@ -3569,6 +3590,9 @@ class ProxyBaseLLMRequestProcessing:
# so it must not suppress that refund.
delivered_chunk = delivered_chunk or chunk != STREAM_SSE_KEEPALIVE_PING_BYTES
yield serialize_chunk(chunk)
held_tail: Final = flush_tail() if flush_tail is not None else b""
if held_tail:
yield serialize_chunk(held_tail)
stream_completed = True
except (asyncio.CancelledError, GeneratorExit):
# Client disconnected mid-stream. CancelledError / GeneratorExit
@ -3579,8 +3603,7 @@ class ProxyBaseLLMRequestProcessing:
# billing and release exactly once. This is the outermost generator
# Starlette closes on disconnect, so the nested iterator hook (which
# only sees GeneratorExit on GC) cannot own the refund.
if not stream_completed:
client_disconnected = True
client_disconnected = not stream_completed
if not delivered_chunk and not _withheld_provider_output(response):
from litellm.proxy.spend_tracking.budget_reservation import (
release_budget_reservation_on_cancel,
@ -3634,6 +3657,7 @@ class ProxyBaseLLMRequestProcessing:
request_data: dict,
proxy_logging_obj: ProxyLogging,
request: Request | None = None,
restamp_model: str | None = None,
) -> AsyncGenerator[str, None]:
"""
Anthropic /messages and Google /generateContent streaming data generator require SSE events.
@ -3642,17 +3666,23 @@ class ProxyBaseLLMRequestProcessing:
SSE serializers directly (rather than re-wrapping it in another
``async for: yield`` trampoline), so a streamed chunk traverses one
fewer async-generator layer / coroutine resume on the hot path.
``restamp_model`` publishes that name on the Anthropic ``message_start``
event in place of the provider's model, matching what the non-streaming
response reports.
"""
restamper: Final = AnthropicStreamModelRestamper(restamp_model) if restamp_model else None
return ProxyBaseLLMRequestProcessing.async_streaming_data_generator(
response=response,
user_api_key_dict=user_api_key_dict,
request_data=request_data,
proxy_logging_obj=proxy_logging_obj,
serialize_chunk=ProxyBaseLLMRequestProcessing.return_sse_chunk,
serialize_chunk=ProxyBaseLLMRequestProcessing._sse_chunk_serializer(restamper),
serialize_error=lambda proxy_exc: (
f"{STREAM_SSE_DATA_PREFIX}{json.dumps({'error': proxy_exc.to_dict()})}\n\n"
),
request=request,
flush_tail=None if restamper is None else restamper.flush,
)
@overload

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

120
litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py
View file

@ -10,6 +10,7 @@ from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, TypeGuard
import httpx
from fastapi import HTTPException
from httpx import Response as HttpxResponse
from pydantic import TypeAdapter
import litellm
from litellm._logging import verbose_proxy_logger
@ -52,6 +53,10 @@ BYPASS_HEADER: Final = "x-headroom-bypass"
HEADROOM_RETRIEVE_TOOL_NAME: Final = "headroom_retrieve"
_HASH_PATTERN: Final = re.compile(r"hash=([a-f0-9]{24})")
_HASH_CACHE_TTL_SECONDS: Final = 15 * 60
# Narrows the base class's bare-dict ``request_data`` at the boundary so its
# untranslated messages can be read with concrete types (values pass through by
# reference, so this is a shallow top-level reconstruction).
_REQUEST_DATA_ADAPTER: Final = TypeAdapter(dict[str, object])
def _is_str_object_dict(value: object) -> TypeGuard[dict[str, object]]: # guard-ok: isinstance narrows correctly; predicate is trivially correct # fmt: skip
@ -116,16 +121,119 @@ def _restore_content_shapes(
return restored
def _protected_indices(messages: Sequence[Mapping[str, object]]) -> frozenset[int]:
def _tool_call_name(tool_call: Mapping[str, object]) -> str | None:
function: Final = tool_call.get("function")
if not _is_str_object_dict(function):
return None
name: Final = function.get("name")
return name if isinstance(name, str) else None
def _is_retrieve_tool_name(name: str | None) -> bool:
"""Match the retrieve tool whether called directly or via the MCP gateway.
Server-side the tool is ``headroom_retrieve``; exposed through LiteLLM's MCP
gateway a client calls it as ``mcp__<server>__headroom_retrieve``.
"""
return name is not None and (
name == HEADROOM_RETRIEVE_TOOL_NAME or name.endswith(f"__{HEADROOM_RETRIEVE_TOOL_NAME}")
)
def _retrieve_call_ids_in_message(message: Mapping[str, object]) -> frozenset[str]:
if message.get("role") != "assistant":
return frozenset()
tool_calls: Final = message.get("tool_calls")
if not _is_object_list(tool_calls):
return frozenset()
return frozenset(
str(tool_call["id"])
for tool_call in tool_calls
if _is_str_object_dict(tool_call) and tool_call.get("id") and _is_retrieve_tool_name(_tool_call_name(tool_call))
)
def _anthropic_tool_use_retrieve_id(block: object) -> str | None:
if not _is_str_object_dict(block) or block.get("type") != "tool_use":
return None
name: Final = block.get("name")
call_id: Final = block.get("id")
if isinstance(name, str) and call_id is not None and _is_retrieve_tool_name(name):
return str(call_id)
return None
def _anthropic_retrieve_ids_in_message(message: Mapping[str, object]) -> frozenset[str]:
content: Final = message.get("content")
if not _is_object_list(content):
return frozenset()
return frozenset(call_id for block in content if (call_id := _anthropic_tool_use_retrieve_id(block)) is not None)
def _raw_retrieve_call_ids(messages: object) -> frozenset[str]:
"""Retrieve-tool call ids read from the request's own, untranslated messages.
The guardrail otherwise scans an OpenAI-translated view where a tool name
over 64 chars is truncated to ``{prefix}_{hash}``, which drops the
``__headroom_retrieve`` suffix a long ``mcp__<server>__`` prefix pushes past
the limit. Tool-call ids are never truncated, so pairing the tool result to
an id read from the original request keeps the match intact. Both wire
shapes are handled: OpenAI ``tool_calls`` and Anthropic ``tool_use`` blocks.
"""
if not _is_object_list(messages):
return frozenset()
return frozenset(
call_id
for message in messages
if _is_str_object_dict(message)
for call_id in _retrieve_call_ids_in_message(message) | _anthropic_retrieve_ids_in_message(message)
)
def _retrieval_result_indices(
messages: Sequence[Mapping[str, object]], extra_retrieve_call_ids: frozenset[str] = frozenset()
) -> frozenset[int]:
"""Indices of tool-result rows that carry ``headroom_retrieve`` output.
When the retrieve tool is exposed to a client that runs its own tool loop
(the LiteLLM MCP gateway path), the client executes the call and sends the
recovered original content back as a tool result on the next turn. That
content is exactly what a prior compression stubbed, so compressing it again
re-derives the identical content hash: a no-op that strands the model on the
marker and loops the agent. Hold those rows back so the expansion survives.
``extra_retrieve_call_ids`` carries ids recovered from the untruncated
request so the pairing survives tool-name truncation (see
``_raw_retrieve_call_ids``).
"""
retrieve_call_ids: Final = extra_retrieve_call_ids | frozenset(
call_id for message in messages for call_id in _retrieve_call_ids_in_message(message)
)
if not retrieve_call_ids:
return frozenset()
return frozenset(
index
for index, message in enumerate(messages)
if message.get("role") in ("tool", "function") and str(message.get("tool_call_id")) in retrieve_call_ids
)
def _protected_indices(
messages: Sequence[Mapping[str, object]], extra_retrieve_call_ids: frozenset[str] = frozenset()
) -> frozenset[int]:
"""Indices headroom must not send to the compression service.
``get_protected_indices`` is litellm's own compression policy: the system
rows, the last user row, the last assistant row. It is expanded over whole
rows, the last user row, the last assistant row. Rows carrying just-retrieved
``headroom_retrieve`` output are added so re-compression can't collapse them
back to the marker they were expanded from. The union is expanded over whole
tool exchanges the way ``compress()`` expands it, so a protected assistant
tool call cannot end up answered by a marker standing in for the result the
model just asked for.
"""
protected: Final = frozenset(get_protected_indices(messages))
protected: Final = frozenset(get_protected_indices(messages)) | _retrieval_result_indices(
messages, extra_retrieve_call_ids
)
return protected | frozenset(
index
for group in group_tool_exchanges(messages)
@ -634,7 +742,11 @@ class HeadroomGuardrail(CustomGuardrail):
# /v1/compress grows a field for sending the live turn as the retrieval
# query without compressing it: query-aware compression reads the newest
# user message, so it is withheld here at some cost to history ranking.
protected_indices: Final = _protected_indices(messages)
# request_data is a bare dict on the base signature; narrow it before
# reading the untranslated messages so long tool names can be recovered.
raw_messages: Final = _REQUEST_DATA_ADAPTER.validate_python(request_data).get("messages")
raw_retrieve_call_ids: Final = _raw_retrieve_call_ids(raw_messages)
protected_indices: Final = _protected_indices(messages, raw_retrieve_call_ids)
compressible: Final = [m for i, m in enumerate(messages) if i not in protected_indices]
if not compressible:
return inputs

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(

70
litellm/proxy/proxy_server.py
View file

@ -39,7 +39,7 @@ from typing import (
import anyio
import websockets
import websockets.exceptions
from pydantic import BaseModel, Json, JsonValue, ValidationError
from pydantic import BaseModel, Json, JsonValue, TypeAdapter, ValidationError
from typing_extensions import NotRequired, ReadOnly, assert_never
from litellm._uuid import uuid
@ -60,6 +60,7 @@ from litellm.constants import (
LITELLM_SETTINGS_SAFE_DB_OVERRIDES,
LITELLM_UI_ALLOW_HEADERS,
LITELLM_UI_SESSION_DURATION,
RUNTIME_UPDATABLE_ROUTER_SETTINGS,
)
from litellm.litellm_core_utils.litellm_logging import (
_init_custom_logger_compatible_class,
@ -253,6 +254,7 @@ from litellm.constants import (
PROXY_BUDGET_RESCHEDULER_MAX_TIME,
PROXY_BUDGET_RESCHEDULER_MIN_TIME,
PROXY_CONFIG_RELOAD_INTERVAL_SECONDS,
ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG,
USER_SPEND_ALERTS_JOB_ID,
WEEKLY_SPEND_REPORT_JOB_ID,
)
@ -352,7 +354,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,
)
@ -5710,13 +5715,9 @@ class ProxyConfig:
router_settings: Final = config.get("router_settings", None)
if router_settings and isinstance(router_settings, dict):
# model list and search_tools already set
exclude_args: Final = {
"model_list",
"search_tools",
}
available_args: Final = [x for x in litellm.Router.get_valid_args() if x not in exclude_args]
available_args: Final = [
x for x in litellm.Router.get_valid_args() if x not in ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG
]
for k, v in router_settings.items():
if k in available_args:
@ -10223,7 +10224,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 +10238,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 +10272,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 +10286,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,
@ -16207,6 +16216,7 @@ async def invitation_delete(
)
async def update_config(
config_info: ConfigYAML,
request: Request,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
@ -16222,6 +16232,26 @@ async def update_config(
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
raise HTTPException(status_code=403, detail="Only proxy admins can update config")
request_body: Final[Mapping[str, JsonValue]] = TypeAdapter(Mapping[str, JsonValue]).validate_python(
await request.json()
)
raw_router_settings: Final = request_body.get("router_settings")
if isinstance(raw_router_settings, dict):
supported_router_settings: Final = RUNTIME_UPDATABLE_ROUTER_SETTINGS | (
frozenset(litellm.Router.get_valid_args()) - ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG
)
unsupported_router_settings: Final = sorted(set(raw_router_settings) - supported_router_settings)
if unsupported_router_settings:
raise HTTPException(
status_code=400,
detail={
"error": (
f"Unsupported router settings: {', '.join(unsupported_router_settings)} "
"are not valid router settings"
)
},
)
if prisma_client is None:
raise Exception("No DB Connected")
@ -16323,11 +16353,19 @@ async def update_config(
)
# router_settings: merge existing + request, request wins.
if config_info.router_settings is not None:
if isinstance(raw_router_settings, dict):
existing = await _read_section("router_settings")
before_router_settings: Final = copy.deepcopy(existing)
updates = config_info.router_settings.dict(exclude_none=True)
new_router_settings: Final = {**existing, **updates}
typed_router_settings: Final = (
config_info.router_settings.dict(exclude_none=True) if config_info.router_settings is not None else {}
)
raw_router_settings_without_none: Final = {
key: value
for key, value in raw_router_settings.items()
if key not in typed_router_settings and value is not None
}
router_settings_updates: Final = {**typed_router_settings, **raw_router_settings_without_none}
new_router_settings: Final = {**existing, **router_settings_updates}
await _upsert_section("router_settings", new_router_settings)
asyncio.create_task(
create_config_audit_log(

56
litellm/proxy/rag_endpoints/endpoints.py
View file

@ -9,6 +9,7 @@ Provides:
import base64
import json
from collections.abc import Mapping
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final
import orjson
@ -19,6 +20,9 @@ from starlette.datastructures import UploadFile
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import (
LiteLLM_ManagedVectorStore,
)
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm.proxy._types import *
from litellm.proxy.auth.auth_utils import is_request_body_safe
@ -36,6 +40,10 @@ from litellm.proxy.rag_endpoints.upload_security import (
RejectedUpload,
validate_upload,
)
from litellm.proxy.vector_store_endpoints.endpoints import (
build_request_data_from_managed_vector_store,
reject_caller_embedding_selection_params,
)
from litellm.proxy.vector_store_endpoints.utils import (
assert_user_can_access_vector_store_id,
)
@ -120,12 +128,21 @@ def _collect_vector_store_ids_from_payload(payload: object) -> set[str]:
async def _authorize_nested_vector_store_ids(
payload: object,
user_api_key_dict: UserAPIKeyAuth,
) -> None:
for vector_store_id in sorted(_collect_vector_store_ids_from_payload(payload)):
await assert_user_can_access_vector_store_id(
vector_store_id=vector_store_id,
user_api_key_dict=user_api_key_dict,
)
) -> Mapping[str, LiteLLM_ManagedVectorStore]:
"""Authorize every nested vector store id and return the managed stores it resolved."""
return MappingProxyType(
{
vector_store_id: store
for vector_store_id in sorted(_collect_vector_store_ids_from_payload(payload))
if (
store := await assert_user_can_access_vector_store_id(
vector_store_id=vector_store_id,
user_api_key_dict=user_api_key_dict,
)
)
is not None
}
)
def _build_file_metadata_entry(
@ -700,11 +717,27 @@ async def rag_query(
status_code=400,
detail={"error": "retrieval_config must contain 'vector_store_id'"},
)
await _authorize_nested_vector_store_ids(
reject_caller_embedding_selection_params(payload=retrieval_config, source="retrieval_config")
resolved_stores: Final = await _authorize_nested_vector_store_ids(
payload=retrieval_config,
user_api_key_dict=user_api_key_dict,
)
# Merge litellm-managed vector store params (provider, region, embedding
# model, credentials, ...) from the registry: the same source the direct
# /vector_stores/{id}/search endpoint uses. Store-managed keys win on
# conflict so callers cannot override the store's provider or credentials.
managed_store: Final = resolved_stores.get(retrieval_config["vector_store_id"])
store_data: Final = (
await build_request_data_from_managed_vector_store(managed_store)
if managed_store is not None
else MappingProxyType({})
)
merged_retrieval_config: Final = {
**retrieval_config,
**store_data,
} # mutable-ok: litellm.aquery requires a plain dict payload
# Add litellm data
request_data: dict[str, object] = {}
request_data = await add_litellm_data_to_request(
@ -716,13 +749,18 @@ async def rag_query(
proxy_config=proxy_config,
)
verbose_proxy_logger.debug("RAG Query - model: %s, retrieval_config: %s", model, retrieval_config)
verbose_proxy_logger.debug(
"RAG Query - model: %s, vector_store_id: %s, custom_llm_provider: %s",
model,
retrieval_config["vector_store_id"],
merged_retrieval_config.get("custom_llm_provider"),
)
# Call query
response: Final = await litellm.aquery(
model=model,
messages=messages,
retrieval_config=retrieval_config,
retrieval_config=merged_retrieval_config,
rerank=rerank,
stream=stream,
router=llm_router,

110
litellm/proxy/vector_store_endpoints/endpoints.py
View file

@ -1,3 +1,5 @@
from collections.abc import Mapping
from types import MappingProxyType
from typing import (
Annotated,
Any, # noqa: TID251 # jsonify_object in proxy/utils.py is annotated with a bare dict
@ -27,11 +29,69 @@ from litellm.types.vector_stores import IndexCreateRequest, IndexListResponse
from litellm.vector_stores.vector_store_registry import VectorStoreIndexRegistry
router: Final = APIRouter()
BLOCKED_QUERY_EMBEDDING_SELECTION_PARAMS: Final = frozenset(
{
"embedding_model",
"litellm_embedding_model",
"litellm_embedding_config",
"litellm_credential_name",
}
)
def reject_caller_embedding_selection_params(payload: Mapping[str, object], source: str) -> None:
blocked: Final = sorted(BLOCKED_QUERY_EMBEDDING_SELECTION_PARAMS & payload.keys())
if blocked:
raise HTTPException(
status_code=400,
detail={
"error": f"'{blocked[0]}' cannot be set in {source}. "
"Embedding configuration comes from the vector store's server-side registration."
},
)
########################################################
# OpenAI Compatible Endpoints
########################################################
async def build_request_data_from_managed_vector_store(
vector_store: LiteLLM_ManagedVectorStore,
) -> Mapping[str, object]:
"""
Build request params (provider, credential ref, litellm_params) from an
already-resolved managed vector store.
``litellm_embedding_config`` is resolved here, at request-handling time,
instead of at row-creation time: the resolved api_key/api_base/api_version
lives only in the returned per-request mapping and is never persisted back
to the registry cache. Legacy rows that already carry a resolved
(cleartext) config skip the lookup and pass through unchanged.
"""
top_level: Final = MappingProxyType(
{
key: vector_store.get(key)
for key in ("custom_llm_provider", "litellm_credential_name")
if key in vector_store
}
)
litellm_params: Final = vector_store.get("litellm_params") or MappingProxyType({})
embedding_model: Final = litellm_params.get("litellm_embedding_model")
if not embedding_model or litellm_params.get("litellm_embedding_config"):
return MappingProxyType({**top_level, **litellm_params})
from litellm.proxy.proxy_server import prisma_client
resolved_config: Final = await _resolve_embedding_config(
embedding_model=embedding_model, prisma_client=prisma_client
)
if not resolved_config:
return MappingProxyType({**top_level, **litellm_params})
return MappingProxyType({**top_level, **litellm_params, "litellm_embedding_config": resolved_config})
async def _update_request_data_with_litellm_managed_vector_store_registry(
data: dict,
vector_store_id: str,
@ -51,47 +111,14 @@ async def _update_request_data_with_litellm_managed_vector_store_registry(
vector_store_to_run: Final[LiteLLM_ManagedVectorStore | None] = await get_litellm_managed_vector_store(
vector_store_id=vector_store_id
)
if vector_store_to_run is not None:
if user_api_key_dict is not None:
await assert_user_can_access_vector_store(
vector_store=vector_store_to_run,
user_api_key_dict=user_api_key_dict,
)
if "custom_llm_provider" in vector_store_to_run:
data["custom_llm_provider"] = vector_store_to_run.get("custom_llm_provider")
if "litellm_credential_name" in vector_store_to_run:
data["litellm_credential_name"] = vector_store_to_run.get("litellm_credential_name")
if "litellm_params" in vector_store_to_run:
litellm_params = vector_store_to_run.get("litellm_params", {}) or {}
# Resolve ``litellm_embedding_config`` here, at request-handling
# time, instead of at row-creation time. The resolved
# ``api_key`` / ``api_base`` / ``api_version`` lives only in
# this per-request ``data`` dict and is never persisted.
# Legacy rows that already carry a resolved (cleartext)
# ``litellm_embedding_config`` skip the lookup and pass through
# unchanged so the embed call keeps working.
embedding_model: Final = litellm_params.get("litellm_embedding_model")
if embedding_model and not litellm_params.get("litellm_embedding_config"):
from litellm.proxy.proxy_server import prisma_client
resolved_config: Final = await _resolve_embedding_config(
embedding_model=embedding_model, prisma_client=prisma_client
)
if resolved_config:
# Build a fresh dict via spread instead of mutating
# ``litellm_params`` in place — the registry hands back
# a reference to its cached object, so an in-place
# update would persist the resolved cleartext into the
# in-memory cache for the lifetime of the process.
litellm_params = {
**litellm_params,
"litellm_embedding_config": resolved_config,
}
data.update(litellm_params)
return data
if vector_store_to_run is None:
return data
if user_api_key_dict is not None:
await assert_user_can_access_vector_store(
vector_store=vector_store_to_run,
user_api_key_dict=user_api_key_dict,
)
return {**data, **(await build_request_data_from_managed_vector_store(vector_store_to_run))}
@router.post(
@ -130,6 +157,7 @@ async def vector_store_search(
)
data = await _read_request_body(request=request)
reject_caller_embedding_selection_params(payload=data, source="the search request body")
data["vector_store_id"] = vector_store_id
# Check for legacy vector store registry (non-managed vector stores)

4
litellm/proxy/vector_store_endpoints/management_endpoints.py
View file

@ -470,7 +470,7 @@ async def create_vector_store_in_db(
# exposed every env-stored embedding-model credential on the
# ``/vector_store/{new,info,update,list}`` responses. Keep the user's
# raw ``litellm_embedding_model`` reference; resolution now happens in
# ``_update_request_data_with_litellm_managed_vector_store_registry``
# ``build_request_data_from_managed_vector_store``
# at request-handling time so the cleartext config exists only in
# per-request memory and never reaches the database.
if litellm_params:
@ -864,7 +864,7 @@ async def update_vector_store(
# embedding-config auto-resolve previously persisted cleartext
# credentials into the row; resolution now happens at request-
# handling time in
# ``_update_request_data_with_litellm_managed_vector_store_registry``
# ``build_request_data_from_managed_vector_store``
# so this row only ever stores the user-supplied
# ``litellm_embedding_model`` reference.
if "litellm_params" in update_data:

25
litellm/rag/main.py
View file

@ -14,6 +14,7 @@ import contextvars
from collections.abc import Coroutine, Iterator
from contextlib import contextmanager
from functools import partial
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final
import httpx
@ -50,6 +51,21 @@ INGESTION_REGISTRY: Final[dict[str, type[BaseRAGIngestion]]] = {
"vertex_ai": VertexAIRAGIngestion,
}
# Only these retrieval_config keys are forwarded to vector_stores.asearch as
# provider-specific params. The explicit allowlist keeps caller-controlled
# connection overrides (api_base, api_key, ...) away from the search call,
# where they could redirect store credentials to an attacker-chosen host.
_FORWARDABLE_RETRIEVAL_CONFIG_KEYS: Final = frozenset(
{
"aws_region_name",
"vector_bucket_name",
"embedding_model",
"litellm_embedding_model",
"litellm_embedding_config",
"litellm_credential_name",
}
)
def get_ingestion_class(provider: str) -> type[BaseRAGIngestion]:
"""
@ -224,13 +240,20 @@ async def _execute_query_pipeline(
raise ValueError("No query found in messages for RAG query")
# 2. Search vector store
# Forward allowlisted provider retrieval_config extras (region, embedding
# model, bucket, credential refs) to the search call; kwargs win on conflict.
provider_search_params: Final = MappingProxyType(
{k: v for k, v in retrieval_config.items() if k in _FORWARDABLE_RETRIEVAL_CONFIG_KEYS}
)
forwarded_search_params: Final = MappingProxyType({**provider_search_params, **kwargs})
with _suppressed_sub_call_billing():
search_response: Final = await litellm.vector_stores.asearch(
vector_store_id=retrieval_config["vector_store_id"],
query=query_text,
max_num_results=retrieval_config.get("top_k", 10),
custom_llm_provider=retrieval_config.get("custom_llm_provider", "openai"),
**kwargs,
router=router,
**forwarded_search_params,
)
search_provider: Final = retrieval_config.get("custom_llm_provider", "openai")

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

326
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
@ -50,6 +50,7 @@ from litellm.constants import (
DEFAULT_HEALTH_CHECK_INTERVAL,
DEFAULT_HEALTH_CHECK_STALENESS_MULTIPLIER,
DEFAULT_MAX_LRU_CACHE_SIZE,
RUNTIME_UPDATABLE_ROUTER_SETTINGS,
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
)
from litellm.integrations.custom_logger import CustomLogger
@ -143,7 +144,11 @@ from litellm.router_utils.cooldown_handlers import (
from litellm.router_utils.fallback_event_handlers import (
AttemptedFallbackTargets,
_check_non_standard_fallback_format,
get_fallback_model_group,
clear_pre_routing_selection,
fallback_lookup_groups,
get_fallback_model_group_for_lookup_groups,
get_pre_routing_selection,
record_pre_routing_selection,
run_async_fallback,
)
from litellm.router_utils.get_retry_from_policy import (
@ -350,6 +355,13 @@ _PreRoutingStrategyT = TypeVar("_PreRoutingStrategyT")
_ALIAS_PARAMS_NEVER_FORWARDED: Final = frozenset({"model", "api_base", "api_key", "api_version"})
_ALIAS_MARKER_FORWARDED_PARAMS_KWARG: Final = "_alias_marker_forwarded_params"
_RUNTIME_TOGGLEABLE_PRE_CALL_CHECKS: Final[Mapping[str, type[CustomLogger]]] = MappingProxyType(
{
"prompt_caching": PromptCachingDeploymentCheck,
"enforce_model_rate_limits": ModelRateLimitingCheck,
}
)
def _stream_chunks_have_generated_content(chunks: Sequence[ModelResponseStream]) -> bool:
for chunk in chunks:
@ -821,6 +833,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
@ -2067,11 +2080,39 @@ class Router:
if _callback is None:
continue
if self.optional_callbacks is not None and any(
isinstance(callback, type(_callback)) for callback in self.optional_callbacks
):
continue
if self.optional_callbacks is None:
self.optional_callbacks = []
self.optional_callbacks.append(_callback)
litellm.logging_callback_manager.add_litellm_callback(_callback)
def set_optional_pre_call_checks(self, optional_pre_call_checks: OptionalPreCallChecks | None) -> None:
if optional_pre_call_checks is None:
return
requested: Final = frozenset(optional_pre_call_checks)
for name, callback_cls in _RUNTIME_TOGGLEABLE_PRE_CALL_CHECKS.items():
if name not in requested:
self._remove_optional_callbacks_of_type(callback_cls)
self.add_optional_pre_call_checks(optional_pre_call_checks)
def _remove_optional_callbacks_of_type(self, callback_cls: type[CustomLogger]) -> None:
if self.optional_callbacks is None or not any(type(cb) is callback_cls for cb in self.optional_callbacks):
return
self.optional_callbacks = [cb for cb in self.optional_callbacks if type(cb) is not callback_cls]
if any(
router is not self and any(type(cb) is callback_cls for cb in (router.optional_callbacks or []))
for router in tuple(_live_routers)
):
return
for cb in tuple(litellm.callbacks):
if type(cb) is callback_cls:
litellm.logging_callback_manager.remove_callback_from_list_by_object(
litellm.callbacks, cb, require_self=False
)
def print_deployment(self, deployment: dict):
"""
returns a copy of the deployment with the api key masked
@ -2319,7 +2360,7 @@ class Router:
@overload
async def acompletion(
self, model: str, messages: list[AllMessageValues], stream: Literal[True, False] = False, **kwargs
) -> CustomStreamWrapper | ModelResponse:
) -> CustomStreamWrapper | ModelResponse:
...
# fmt: on
@ -4918,6 +4959,19 @@ class Router:
)
response = await response
if self._should_raise_anthropic_refusal_error(
model=model,
original_generic_function=original_generic_function,
response=response,
kwargs=kwargs,
):
from litellm.llms.anthropic.experimental_pass_through.messages.utils import (
safeguard_refusal_error,
)
refusal_details: Final = cast(dict, response["stop_details"]) # cast-ok: gate verified the shape
raise safeguard_refusal_error(model=model, stop_details=refusal_details)
self.success_calls[model_name] += 1
verbose_router_logger.info("ageneric_api_call_with_fallbacks(model=%s)\x1b[32m 200 OK\x1b[0m", model_name)
@ -4964,6 +5018,11 @@ class Router:
# fallback to the original reference for any non-picklable value.
# The original_generic_function is preserved so the per-attempt
# helper knows which underlying API to call on fallback.
# The pre-routing hook stamps its tier selection into this bucket during the primary
# attempt; seeding it before the snapshot gives both the live kwargs and the copy a
# bucket, so the post-call carry-over below always has somewhere to read and write.
kwargs.setdefault("litellm_metadata", {}) # mutable-ok: shared bucket # rebind-ok: stamp must be readable here
fallback_kwargs: Final[dict[str, object]] = kwargs.copy()
if isinstance(fallback_kwargs.get("litellm_metadata"), dict):
fallback_kwargs["litellm_metadata"] = safe_deep_copy(fallback_kwargs["litellm_metadata"])
@ -4973,6 +5032,14 @@ class Router:
response: Final = await self._ageneric_api_call_with_fallbacks(original_function=original_function, **kwargs)
# The snapshot predates the pre-routing hook, so the tier it stamped into the live kwargs
# is carried over write-or-clear: a stale or caller-supplied selection left in the copy
# would key the mid-stream fallback lookup off a tier this attempt never routed to.
clear_pre_routing_selection(fallback_kwargs)
live_pre_routing_selection: Final = get_pre_routing_selection(kwargs)
if live_pre_routing_selection is not None:
record_pre_routing_selection(fallback_kwargs, live_pre_routing_selection)
if kwargs.get("stream") and isinstance(response, BaseResponsesAPIStreamingIterator):
return await self._aresponses_streaming_iterator(
response=response,
@ -5030,6 +5097,10 @@ class Router:
from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import (
aclose_if_supported,
parse_anthropic_error_event,
parse_anthropic_refusal_stop_details,
)
from litellm.llms.anthropic.experimental_pass_through.messages.utils import (
safeguard_refusal_error,
)
source_iterator: Final = response
@ -5068,13 +5139,35 @@ class Router:
continue
if _anthropic_stream_commits_now(chunk, has_generated_content, len(buffered_lifecycle_chunks)):
has_generated_content = True # rebind-ok: real content seen, or the buffer cap was hit
error_event = parse_anthropic_error_event(chunk)
# A transport can split one SSE data line across byte chunks, so pre-content
# detection parses the accumulated buffer plus the current chunk, never the
# chunk alone; the buffer is already capped, which bounds this window too.
parse_window = ( # rebind-ok: freshly computed each iteration, never carried over
b"".join(c for c in (*buffered_lifecycle_chunks, chunk) if isinstance(c, (bytes, bytearray))) # pyright: ignore[reportUnnecessaryIsInstance] # bridge-path chunks are not always bytes at runtime
if not has_generated_content and isinstance(chunk, (bytes, bytearray)) # pyright: ignore[reportUnnecessaryIsInstance] # bridge-path chunks are not always bytes at runtime
else chunk
)
error_event = parse_anthropic_error_event(parse_window)
retriable_pending_error = ( # rebind-ok: freshly computed each iteration, never carried over
not has_generated_content
and error_event is not None
and _is_retriable_anthropic_status(error_event[2])
and not _anthropic_stream_error_is_gateway_verdict(chunk)
)
refusal_stop_details = ( # rebind-ok: freshly computed each iteration, never carried over
parse_anthropic_refusal_stop_details(parse_window)
if not has_generated_content and error_event is None
else None
)
if refusal_stop_details is not None and self._has_content_policy_fallback(model, initial_kwargs):
refusal_error = safeguard_refusal_error(model=model, stop_details=refusal_stop_details)
raise MidStreamFallbackError(
message=refusal_error.message,
model=model,
llm_provider="anthropic",
original_exception=refusal_error,
is_pre_first_chunk=True,
)
if not has_generated_content and not retriable_pending_error and error_event is None:
buffered_lifecycle_chunks = (*buffered_lifecycle_chunks, chunk)
continue
@ -5186,8 +5279,13 @@ class Router:
kwargs=initial_kwargs,
metadata_variable_name="litellm_metadata",
)
# The content-policy dispatch branch matches on the trigger's own type, so a refusal's
# MidStreamFallbackError envelope is unwrapped here or the wrong fallback list is consulted.
fallback_trigger: Final[Exception] = (
e.original_exception if isinstance(e.original_exception, litellm.ContentPolicyViolationError) else e
)
fallback_response = await self.async_function_with_fallbacks_common_utils( # rebind-ok: set on success
e=e,
e=fallback_trigger,
disable_fallbacks=False,
fallbacks=fallbacks,
context_window_fallbacks=context_window_fallbacks,
@ -5243,6 +5341,11 @@ class Router:
# share, leaking primary-deployment metadata into the mid-stream
# fallback request. safe_deep_copy avoids deep-copying the full
# kwargs (which can hold non-deepcopyable logging handles/clients).
# The pre-routing hook stamps its tier selection into this bucket during the primary
# attempt; seeding it before the snapshot gives both the live kwargs and the copy a
# bucket, so the post-call carry-over below always has somewhere to read and write.
kwargs.setdefault("litellm_metadata", {}) # mutable-ok: shared bucket # rebind-ok: stamp must be readable here
fallback_kwargs: Final[dict[str, object]] = kwargs.copy() # mutable-ok: mutated below before re-entry
if isinstance(fallback_kwargs.get("litellm_metadata"), dict):
fallback_kwargs["litellm_metadata"] = safe_deep_copy(fallback_kwargs["litellm_metadata"])
@ -5252,6 +5355,14 @@ class Router:
response: Final = await self._ageneric_api_call_with_fallbacks(original_function=original_function, **kwargs)
# The snapshot predates the pre-routing hook, so the tier it stamped into the live kwargs
# is carried over write-or-clear: a stale or caller-supplied selection left in the copy
# would key the mid-stream fallback lookup off a tier this attempt never routed to.
clear_pre_routing_selection(fallback_kwargs)
live_pre_routing_selection: Final = get_pre_routing_selection(kwargs)
if live_pre_routing_selection is not None:
record_pre_routing_selection(fallback_kwargs, live_pre_routing_selection)
if kwargs.get("stream") and hasattr(response, "__aiter__"):
return await self._aanthropic_messages_streaming_iterator(
response=cast("AsyncIterator[bytes]", response), # cast-ok: stream=True always returns a byte iterator
@ -6299,8 +6410,6 @@ class Router:
"responses",
"generate_content",
"generate_content_stream",
"vector_store_search",
"vector_store_create",
"ocr",
"search",
"video_generation",
@ -6324,6 +6433,8 @@ class Router:
return sync_wrapper
if call_type in (
"vector_store_search",
"vector_store_create",
"vector_store_retrieve",
"vector_store_list",
"vector_store_update",
@ -6335,11 +6446,16 @@ class Router:
client: object | None = None,
**kwargs,
):
if custom_llm_provider and "custom_llm_provider" not in kwargs:
kwargs["custom_llm_provider"] = custom_llm_provider
if kwargs.get("model"):
return self._generic_api_call_with_fallbacks(original_function=original_function, **kwargs)
return original_function(**kwargs)
provider_kwargs: Final = (
MappingProxyType({**kwargs, "custom_llm_provider": custom_llm_provider})
if custom_llm_provider and "custom_llm_provider" not in kwargs
else MappingProxyType(kwargs)
)
if provider_kwargs.get("model"):
return self._generic_api_call_with_fallbacks(original_function=original_function, **provider_kwargs)
if call_type == "vector_store_search":
return original_function(**MappingProxyType({**provider_kwargs, "router": self}))
return original_function(**provider_kwargs)
return vector_store_sync_wrapper
@ -6515,6 +6631,7 @@ class Router:
return await self._init_vector_store_api_endpoints(
original_function=original_function,
custom_llm_provider=custom_llm_provider,
call_type=call_type,
**kwargs,
)
elif call_type in ("afile_delete", "afile_content"):
@ -6555,6 +6672,7 @@ class Router:
self,
original_function: Callable,
custom_llm_provider: str | None = None,
call_type: str | None = None,
**kwargs,
):
"""
@ -6573,6 +6691,13 @@ class Router:
**kwargs,
)
# For search, pass the router so provider transforms can resolve
# router-managed embedding models (e.g. S3 Vectors query embeddings).
# The merge also overrides any client-supplied `router` key.
if call_type == "avector_store_search":
search_kwargs: Final = MappingProxyType({**kwargs, "router": self})
return await original_function(**search_kwargs)
# Otherwise, call the original function directly
return await original_function(**kwargs)
@ -6818,6 +6943,9 @@ class Router:
original_exception: Final = e
fallback_model_group = None
original_model_group: Final[str | None] = kwargs.get("model")
# A pre-routing hook (complexity / auto / adaptive / quality routers) picks a tier
# behind the router name, and fallbacks are configured per tier, not per router.
lookup_groups: Final[tuple[str, ...]] = fallback_lookup_groups(kwargs, model_group)
fallback_failure_exception_str = ""
if disable_fallbacks is True or original_model_group is None:
@ -6862,15 +6990,15 @@ class Router:
]
# Get external fallbacks — handle both standard and non-standard formats
external_fallback_group: list | None = None
if fallbacks is not None and model_group is not None:
if fallbacks is not None and lookup_groups:
if _check_non_standard_fallback_format(fallbacks=fallbacks):
# Non-standard formats (e.g. ["claude-3-haiku"] or
# [{"model": "...", "messages": [...]}]) are passed through directly
external_fallback_group = fallbacks
else:
external_fallback_group, generic_idx = get_fallback_model_group(
external_fallback_group, generic_idx = get_fallback_model_group_for_lookup_groups(
fallbacks=fallbacks,
model_group=cast(str, model_group),
lookup_groups=lookup_groups,
)
if external_fallback_group is None and generic_idx is not None:
external_fallback_group = fallbacks[generic_idx]["*"]
@ -6928,9 +7056,9 @@ class Router:
if isinstance(e, litellm.ContextWindowExceededError):
if context_window_fallbacks is not None:
context_window_fallback_model_group: Final[list[str] | None] = (
self._get_fallback_model_group_from_fallbacks(
self._get_fallback_model_group_for_lookup_groups(
fallbacks=context_window_fallbacks,
model_group=model_group,
lookup_groups=lookup_groups,
)
)
if context_window_fallback_model_group is None:
@ -6961,9 +7089,9 @@ class Router:
elif isinstance(e, litellm.ContentPolicyViolationError):
if content_policy_fallbacks is not None:
content_policy_fallback_model_group: Final[list[str] | None] = (
self._get_fallback_model_group_from_fallbacks(
self._get_fallback_model_group_for_lookup_groups(
fallbacks=content_policy_fallbacks,
model_group=model_group,
lookup_groups=lookup_groups,
)
)
if content_policy_fallback_model_group is None:
@ -6990,14 +7118,14 @@ class Router:
if litellm.expose_router_debug_in_errors:
e.message += f"\n{error_message}"
if fallbacks is not None and model_group is not None:
if fallbacks is not None and lookup_groups:
verbose_router_logger.debug("inside model fallbacks: %s", mask_sensitive_structure(fallbacks))
(
fallback_model_group,
generic_fallback_idx,
) = get_fallback_model_group(
) = get_fallback_model_group_for_lookup_groups(
fallbacks=fallbacks, # if fallbacks = [{"gpt-3.5-turbo": ["claude-3-haiku"]}]
model_group=cast(str, model_group),
lookup_groups=lookup_groups,
)
## if none, check for generic fallback
if fallback_model_group is None and generic_fallback_idx is not None:
@ -7006,12 +7134,12 @@ class Router:
if fallback_model_group is None:
masked_fallbacks: Final = mask_sensitive_structure(fallbacks)
verbose_router_logger.info(
"No fallback model group found for original model_group=%s. Fallbacks=%s",
model_group,
"No fallback model group found for lookup_groups=%s. Fallbacks=%s",
" -> ".join(lookup_groups),
masked_fallbacks,
)
if hasattr(original_exception, "message") and litellm.expose_router_debug_in_errors:
original_exception.message += f"No fallback model group found for original model_group={model_group}. Fallbacks={masked_fallbacks}"
original_exception.message += f"No fallback model group found for lookup_groups={' -> '.join(lookup_groups)}. Fallbacks={masked_fallbacks}"
raise original_exception
input_kwargs.update(
@ -7057,6 +7185,7 @@ class Router:
If it fails after num_retries, fall back to another model group
"""
model_group: Final[str | None] = kwargs.get("model")
clear_pre_routing_selection(kwargs) # pyright: ignore[reportUnknownArgumentType] # **kwargs is untyped at this boundary
if not isinstance(kwargs.get("attempted_targets"), AttemptedFallbackTargets):
_fallback_metadata_key: Final = _get_router_metadata_variable_name(
function_name=getattr(kwargs.get("original_function"), "__name__", None)
@ -7482,6 +7611,24 @@ class Router:
break
return fallback_model_group
def _get_fallback_model_group_for_lookup_groups(
self,
fallbacks: list[dict[str, list[str]]], # mutable-ok: mirrors the sibling resolver's contract
lookup_groups: tuple[str, ...],
) -> list[str] | None: # mutable-ok: mirrors the sibling resolver's contract
"""First lookup group whose exact-key chain resolves (tier first, then requested group)."""
return next(
(
resolved
for resolved in (
self._get_fallback_model_group_from_fallbacks(fallbacks=fallbacks, model_group=group)
for group in lookup_groups
)
if resolved is not None
),
None,
)
def _get_first_default_fallback(self) -> str | None:
"""
Returns the first model from the default_fallbacks list, if it exists.
@ -7897,6 +8044,31 @@ class Router:
return True
return False
def _has_content_policy_fallback(self, model_group: str, kwargs: Mapping[str, Any]) -> bool:
"""
Whether a content-policy fallback would resolve for this request, keyed the same way
async_function_with_fallbacks_common_utils resolves it: the tier a pre-routing hook
selected wins over the requested group. Raising without this returning True would turn
a deliverable response into an error the fallback chain cannot recover from.
"""
content_policy_fallbacks: Final = kwargs.get("content_policy_fallbacks", self.content_policy_fallbacks)
if content_policy_fallbacks is not None:
return (
self._get_fallback_model_group_for_lookup_groups(
fallbacks=content_policy_fallbacks,
lookup_groups=fallback_lookup_groups(kwargs, model_group),
)
is not None
)
if self._has_default_fallbacks():
return True
verbose_router_logger.debug(
"No content-policy fallback available. Returning original response. model=%s, content_policy_fallbacks=%s",
model_group,
content_policy_fallbacks,
)
return False
def _should_raise_content_policy_error(self, model: str, response: ModelResponse, kwargs: dict) -> bool:
"""
Determines if a content policy error should be raised.
@ -7909,27 +8081,26 @@ class Router:
if response.choices[0].finish_reason != "content_filter":
return False
content_policy_fallbacks: Final = kwargs.get("content_policy_fallbacks", self.content_policy_fallbacks)
return self._has_content_policy_fallback(model, kwargs)
### ONLY RAISE ERROR IF CP FALLBACK AVAILABLE ###
if content_policy_fallbacks is not None:
fallback_model_group = None
for item in content_policy_fallbacks: # [{"gpt-3.5-turbo": ["gpt-4"]}]
if list(item.keys())[0] == model:
fallback_model_group = item[model]
break
if fallback_model_group is not None:
return True
elif self._has_default_fallbacks(): # default fallbacks set
return True
verbose_router_logger.debug(
"Content Policy Error occurred. No available fallbacks. Returning original response. model=%s, content_policy_fallbacks=%s",
model,
content_policy_fallbacks,
def _should_raise_anthropic_refusal_error(
self, model: str, original_generic_function: Callable, response: object, kwargs: Mapping[str, Any]
) -> bool:
"""
The /v1/messages twin of _should_raise_content_policy_error: an Anthropic safeguard
refusal (stop_reason "refusal" carrying stop_details) re-enters the fallback chain only
when a content-policy fallback is configured; a plain refusal without stop_details, or
any response with nothing configured, is returned to the client unchanged.
"""
from litellm.llms.anthropic.experimental_pass_through.messages.utils import (
get_safeguard_refusal_stop_details,
)
return False
if getattr(original_generic_function, "__name__", "") != "anthropic_messages":
return False
if get_safeguard_refusal_stop_details(response) is None:
return False
return self._has_content_policy_fallback(model, kwargs)
def _get_healthy_deployments(self, model: str, parent_otel_span: Span | None):
_all_deployments: list = []
@ -8363,6 +8534,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
)
@ -8792,6 +8976,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
@ -9414,8 +9599,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)
@ -9618,6 +9807,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:
@ -11203,27 +11412,6 @@ class Router:
"""
Update the router settings.
"""
# only the following settings are allowed to be configured
_allowed_settings: Final = [
"routing_strategy_args",
"routing_strategy",
"routing_groups",
"allowed_fails",
"cooldown_time",
"num_retries",
"timeout",
"max_retries",
"retry_after",
"fallbacks",
"context_window_fallbacks",
"retry_policy",
"model_group_retry_policy",
"model_group_alias",
"enable_weighted_failover",
"enable_tag_filtering",
"tag_routing_prefix",
]
_int_settings: Final = [
"timeout",
"num_retries",
@ -11236,13 +11424,15 @@ class Router:
rebuild_routing_groups = False
relink_lar1_from_args = False
for var in kwargs:
if var in _allowed_settings:
if var in RUNTIME_UPDATABLE_ROUTER_SETTINGS:
if var in _int_settings:
_casted_value = int(kwargs[var])
setattr(self, var, _casted_value)
elif var == "routing_groups":
self._routing_groups_input = kwargs[var]
rebuild_routing_groups = True
elif var == "optional_pre_call_checks":
self.set_optional_pre_call_checks(kwargs[var])
elif var == "retry_policy":
value = kwargs[var]
if isinstance(value, dict):
@ -12098,6 +12288,7 @@ class Router:
if pre_routing_hook_response is not None:
model = pre_routing_hook_response.model
messages = pre_routing_hook_response.messages
record_pre_routing_selection(request_kwargs, model)
if pre_routing_hook_response.litellm_params:
accepted_tier_params: Final = self._tier_params_the_target_accepts(
model, pre_routing_hook_response.litellm_params, request_kwargs
@ -12213,6 +12404,7 @@ class Router:
if pre_routing_hook_response is not None:
model = pre_routing_hook_response.model
messages = pre_routing_hook_response.messages
record_pre_routing_selection(request_kwargs, model)
if pre_routing_hook_response.litellm_params:
accepted_tier_params: Final = self._tier_params_the_target_accepts(
model, pre_routing_hook_response.litellm_params, request_kwargs

85
litellm/router_utils/fallback_event_handlers.py
View file

@ -214,6 +214,91 @@ def _check_stripped_model_group(model_group: str, fallback_key: str) -> bool:
return False
PRE_ROUTING_SELECTED_MODEL_KEY: Final = "pre_routing_selected_model"
_ROUTER_METADATA_BUCKETS: Final = ("metadata", "litellm_metadata")
def record_pre_routing_selection(request_kwargs: Mapping[str, Any] | None, selected_model: str) -> None:
"""
Remember which model a pre-routing hook picked, so fallback lookup can key off it.
Fallback resolution runs on an outer kwargs dict that ``**kwargs`` already copied, so
writing the model there is invisible by the time routing picks a tier. The metadata
buckets are nested dicts shared by reference across those copies, which is how the
router already carries values back up.
The write goes through the proxy-internal bucket resolver, never into both buckets:
on /v1/messages the top-level ``metadata`` dict is the provider's own request field,
so a blanket write would forward the tier stamp upstream.
"""
from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs
if request_kwargs is None:
return
bucket: Final = request_kwargs.get(get_metadata_variable_name_from_kwargs(request_kwargs))
if isinstance(bucket, dict):
bucket[PRE_ROUTING_SELECTED_MODEL_KEY] = selected_model
def clear_pre_routing_selection(request_kwargs: Mapping[str, object] | None) -> None:
"""
Drop any selection the router did not make itself on this hop.
The buckets carry whatever the caller sent, so an inbound value is the caller
choosing a fallback chain rather than the router choosing a tier. A fallback hop
also inherits the previous hop's selection, which would key its own failure off
the tier that already failed. Clearing at the start of every hop leaves only a
value the pre-routing hook wrote while routing that hop.
"""
if request_kwargs is None:
return
for bucket in (request_kwargs.get(name) for name in _ROUTER_METADATA_BUCKETS):
if isinstance(bucket, dict) and PRE_ROUTING_SELECTED_MODEL_KEY in bucket:
del bucket[PRE_ROUTING_SELECTED_MODEL_KEY]
def get_pre_routing_selection(kwargs: Mapping[str, Any]) -> str | None:
"""The model a pre-routing hook selected for this request, if one did."""
buckets: Final = (kwargs.get(name) for name in _ROUTER_METADATA_BUCKETS)
selections: Final = (bucket.get(PRE_ROUTING_SELECTED_MODEL_KEY) for bucket in buckets if isinstance(bucket, dict))
return next((selected for selected in selections if isinstance(selected, str) and selected), None)
def fallback_lookup_groups(kwargs: Mapping[str, Any], model_group: str | None) -> tuple[str, ...]:
"""
Ordered keys for resolving a fallback chain: the tier a pre-routing hook selected wins,
and the requested group still resolves when no tier-keyed chain exists, so configs keyed
on the router name (the documented contract) keep working behind auto-routers.
"""
ordered: Final = (get_pre_routing_selection(kwargs), model_group)
return tuple(dict.fromkeys(group for group in ordered if group))
def _resolved_a_specific_chain(
fallbacks: list[Any], # mutable-ok: mirrors get_fallback_model_group's contract
result: tuple[list[str] | None, int | None], # mutable-ok: mirrors get_fallback_model_group's contract
) -> bool:
resolved, generic_idx = result
if resolved is None:
return False
return generic_idx is None or resolved is not fallbacks[generic_idx]["*"]
def get_fallback_model_group_for_lookup_groups(
fallbacks: list[Any], # mutable-ok: mirrors get_fallback_model_group's contract
lookup_groups: tuple[str, ...],
) -> tuple[list[str] | None, int | None]: # mutable-ok: mirrors get_fallback_model_group's contract
"""
First lookup group with a specifically-keyed chain wins; the generic "*" chain applies
only after every group missed, so a catch-all cannot shadow a later group's own chain.
"""
results: Final = tuple(get_fallback_model_group(fallbacks=fallbacks, model_group=group) for group in lookup_groups)
specific: Final = next((result for result in results if _resolved_a_specific_chain(fallbacks, result)), None)
if specific is not None:
return specific
return next((result for result in results if result[0] is not None), (None, None))
def get_fallback_model_group(fallbacks: list[Any], model_group: str) -> tuple[list[str] | None, int | None]:
"""
Returns:

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,

13
litellm/types/llms/anthropic_messages/anthropic_response.py
View file

@ -78,6 +78,16 @@ class AnthropicUsage(TypedDict, total=False):
server_tool_use: NotRequired[ReadOnly[ServerToolUsage]]
class AnthropicStopDetails(TypedDict, total=False):
"""
Safeguard verdict accompanying a `stop_reason: "refusal"` response:
https://platform.claude.com/docs/en/build-with-claude/refusals-and-fallback
"""
category: ReadOnly[str | None]
explanation: ReadOnly[str | None]
class AnthropicMessagesResponse(TypedDict, total=False):
"""
Anthropic Messages API Response: https://docs.anthropic.com/en/api/messages
@ -90,7 +100,8 @@ class AnthropicMessagesResponse(TypedDict, total=False):
id: str
model: str | None # This represents the Model type from Anthropic
role: Literal["assistant"] | None
stop_reason: Literal["end_turn", "max_tokens", "stop_sequence", "tool_use"] | None
stop_reason: Literal["end_turn", "max_tokens", "stop_sequence", "tool_use", "refusal"] | None
stop_details: NotRequired[ReadOnly[AnthropicStopDetails | None]]
stop_sequence: str | None
type: Literal["message"] | None
usage: AnthropicUsage | None

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"]

29
litellm/types/router.py
View file

@ -106,6 +106,20 @@ class RetryPolicy(BaseModel):
InternalServerErrorRetries: int | None = None
OptionalPreCallChecks = list[
Literal[
"prompt_caching",
"router_budget_limiting",
"responses_api_deployment_check",
"deployment_affinity",
"session_affinity",
"forward_client_headers_by_model_group",
"enforce_model_rate_limits",
"encrypted_content_affinity",
]
]
class UpdateRouterConfig(BaseModel):
"""
Set of params that you can modify via `router.update_settings()`.
@ -128,6 +142,7 @@ class UpdateRouterConfig(BaseModel):
model_group_alias: dict[str, str | dict] | None = {}
enable_tag_filtering: bool | None = None
tag_routing_prefix: str | None = None
optional_pre_call_checks: OptionalPreCallChecks | None = None
model_config = ConfigDict(protected_namespaces=())
@ -869,20 +884,6 @@ class FallbackAccessCheck(Protocol):
async def __call__(self, *, model: str, request_kwargs: Mapping[str, object], llm_router: "Router") -> bool: ...
OptionalPreCallChecks = list[
Literal[
"prompt_caching",
"router_budget_limiting",
"responses_api_deployment_check",
"deployment_affinity",
"session_affinity",
"forward_client_headers_by_model_group",
"enforce_model_rate_limits",
"encrypted_content_affinity",
]
]
class LiteLLM_RouterFileObject(TypedDict, total=False):
"""
Tracking the litellm params hash, used for mapping the file id to the right model

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

9
litellm/vector_stores/main.py
View file

@ -7,7 +7,7 @@ import builtins
import contextvars
from collections.abc import Coroutine, Mapping
from functools import partial
from typing import Final
from typing import TYPE_CHECKING, Final
import httpx
@ -29,6 +29,9 @@ from litellm.types.vector_stores import (
from litellm.utils import ProviderConfigManager, client
from litellm.vector_stores.utils import VectorStoreRequestUtils
if TYPE_CHECKING:
from litellm.router import Router
####### ENVIRONMENT VARIABLES ###################
# Initialize any necessary instances or variables here
base_llm_http_handler = BaseLLMHTTPHandler()
@ -280,6 +283,7 @@ async def asearch(
timeout: float | httpx.Timeout | None = None,
# LiteLLM specific params,
custom_llm_provider: str | None = None,
router: "Router | None" = None,
**kwargs,
) -> VectorStoreSearchResponse:
"""
@ -308,6 +312,7 @@ async def asearch(
extra_body=extra_body,
timeout=timeout,
custom_llm_provider=custom_llm_provider,
router=router,
**kwargs,
)
@ -347,6 +352,7 @@ def search(
timeout: float | httpx.Timeout | None = None,
# LiteLLM specific params,
custom_llm_provider: str | None = None,
router: "Router | None" = None,
**kwargs,
) -> VectorStoreSearchResponse | Coroutine[object, object, VectorStoreSearchResponse]:
"""
@ -450,6 +456,7 @@ def search(
timeout=timeout or request_timeout,
_is_async=_is_async,
client=kwargs.get("client"),
router=router,
)
return response

1007
model_prices_and_context_window.json
View file

File diff suppressed because it is too large Load diff

4
pyproject.toml
View file

@ -161,7 +161,7 @@ proxy-runtime = [
"mangum>=0.17.0,<1.0",
"azure-ai-contentsafety>=1.0.0,<2.0",
"azure-storage-file-datalake>=12.20.0,<13.0",
"pypdf>=6.12.0,<7.0",
"pypdf>=6.16.1,<7.0",
"llm-sandbox>=0.3.39,<1.0",
"detect-secrets>=1.5.0,<2.0",
]
@ -292,7 +292,7 @@ exclude = [
[tool.uv]
constraint-dependencies = [
"tornado>=6.5.6",
"tornado>=6.5.8",
"aiohttp>=3.14.2,<4.0",
"packaging>=24.0",
"soupsieve>=2.8.4",

75
tests/e2e/junit_properties.py
View file

@ -2,10 +2,16 @@
The e2e suite ships results to Loki/Grafana from a standard pytest JUnit report
(`--junitxml=e2e-report.xml`), not a bespoke log line. JUnit already records
outcome, duration, and node id for every `<testcase>`; the only signals it cannot
derive on its own are the normalized suite package and the coverage-registry cell
ids a test covers. Those ride along as JUnit `<property>` entries via each item's
`user_properties`, attached in `conftest.py::pytest_collection_modifyitems`.
outcome, duration, and node id for every `<testcase>`; the signals it cannot
derive on its own are the normalized suite package, the coverage-registry cell
ids a test covers, and where the test's source lives. Those ride along as JUnit
`<property>` entries via each item's `user_properties`, attached in
`conftest.py::pytest_collection_modifyitems`.
`source` is a property rather than the `file=` / `line=` attributes pytest used
to write, because the `xunit2` family this suite runs on drops those, and
switching families would change the XML for every consumer of it -- the
Buildkite Test Engine upload and the Loki pipeline included.
"""
from __future__ import annotations
@ -14,22 +20,61 @@ from collections.abc import Iterable
import pytest
# Hardcoded because the runner image copies tests/e2e/ to /app/e2e, so nothing
# at runtime names this suite's place in the repo. test_junit_properties.py
# fails from a checkout if it moves.
SUITE_ROOT = "tests/e2e"
def suite_parts(path_part: str) -> tuple[str, ...]:
"""Path components of a suite file relative to tests/e2e, however it ran.
Pytest paths are rootdir-relative, and rootdir moves with the invocation: a
repo-root run gives `tests/e2e/logging/test_x.py`, a suite-cwd run (the
runner image) gives `logging/test_x.py`. Both collapse to the same tuple.
"""
raw = tuple(p for p in path_part.replace("\\", "/").split("/") if p and p != ".")
return raw[2:] if len(raw) >= 3 and raw[0] == "tests" and raw[1] == "e2e" else raw
def package_from_nodeid(nodeid: str) -> str:
"""Top-level suite package under tests/e2e/, or 'root' for top-level files.
Pytest nodeids are relative to the invocation cwd. Repo-root runs look like
`tests/e2e/logging/...`; suite-cwd runs look like `logging/...`. Strip the
`tests/e2e` prefix so package is the suite dir either way.
"""
path_part = nodeid.split("::", 1)[0].replace("\\", "/")
raw = tuple(p for p in path_part.split("/") if p and p != ".")
parts = raw[2:] if len(raw) >= 3 and raw[0] == "tests" and raw[1] == "e2e" else raw
"""Top-level suite package under tests/e2e/, or 'root' for top-level files."""
parts = suite_parts(nodeid.split("::", 1)[0])
if len(parts) <= 1:
return "root"
return parts[0]
def source_from_location(path: str, lineno: int | None) -> str:
"""Repo-relative `path:line` for a test, or '' when nothing is linkable.
`pytest.Item.location` gives a rootdir-relative path and a ZERO-based line.
The path is re-rooted at SUITE_ROOT so consumers need not know how pytest was
started, and the line is emitted ONE-based to match editors, tracebacks and
code hosts. A decorated test anchors at its first decorator, which is where
pytest reports it.
Empty rather than a guess for anything unlinkable: no line, a path reaching
upward, or a path carrying a colon, which is both how an absolute Windows
path arrives and a character `path:line` has no way to represent.
"""
if lineno is None:
return ""
normalized = path.replace("\\", "/")
if normalized.startswith("/") or ":" in normalized or ".." in normalized.split("/"):
return ""
parts = suite_parts(normalized)
if not parts:
return ""
return f"{'/'.join((SUITE_ROOT, *parts))}:{lineno + 1}"
def source_from_item(item: pytest.Item) -> str:
"""Read the repo-relative `path:line` off a pytest Item's reported location."""
path, lineno, _ = item.location
return source_from_location(path, lineno)
def dedupe_covers(marker_args: Iterable[tuple[object, ...]]) -> tuple[str, ...]:
"""Flatten @pytest.mark.covers arg lists into unique, order-preserving cell
ids, dropping anything that is not a non-empty string."""
@ -43,10 +88,12 @@ def covers_from_item(item: pytest.Item) -> tuple[str, ...]:
def result_properties(item: pytest.Item) -> tuple[tuple[str, str], ...]:
"""The custom signals a standard reporter cannot derive: the normalized suite
package and the comma-joined coverage-registry cell ids this test covers."""
package, the comma-joined coverage-registry cell ids this test covers, and the
repo-relative `path:line` its source sits at."""
return (
("package", package_from_nodeid(item.nodeid)),
("covers", ",".join(covers_from_item(item))),
("source", source_from_item(item)),
)

131
tests/e2e/test_junit_properties.py Normal file
View file

@ -0,0 +1,131 @@
"""Harness coverage for the custom JUnit properties.
No proxy and no ``e2e`` marker. Pins the two normalizations that have to agree
about where a suite file lives -- ``package_from_nodeid`` (strip the suite root)
and ``source_from_location`` (re-root at it) -- across both ways the suite is
launched, plus the one-based line offset and the refusal to emit a path that
escapes the suite. The consumers of these properties are the Loki/Grafana
rollups and, for ``source``, the status page's per-test links to GitHub.
"""
from __future__ import annotations
from pathlib import Path
import pytest
from junit_properties import (
SUITE_ROOT,
attach_result_properties,
dedupe_covers,
package_from_nodeid,
result_properties,
source_from_location,
suite_parts,
)
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:
"""The litellm checkout above this file, or None when there isn't one."""
return next((p for p in Path(__file__).resolve().parents if (p / ".git").exists()), None)
class TestSuiteParts:
@pytest.mark.parametrize(
"path",
["logging/test_x.py", "tests/e2e/logging/test_x.py", "./logging/test_x.py", "tests\\e2e\\logging\\test_x.py"],
)
def test_both_invocation_shapes_collapse_to_the_same_components(self, path: str) -> None:
"""A repo-root run and a suite-cwd run report the same file differently;
every downstream signal has to see one spelling."""
assert suite_parts(path) == ("logging", "test_x.py")
def test_top_level_suite_file_keeps_its_single_component(self) -> None:
assert suite_parts("tests/e2e/test_fixture_mode.py") == ("test_fixture_mode.py",)
class TestPackageFromNodeid:
@pytest.mark.parametrize(
("nodeid", "expected"),
[
("logging/test_x.py::TestFoo::test_bar", "logging"),
("tests/e2e/logging/test_x.py::TestFoo::test_bar", "logging"),
("quota_management/spend_tracking/test_x.py::test_bar", "quota_management"),
("test_fixture_mode.py::TestParseFixtureMode::test_known_values_normalize", "root"),
("tests/e2e/test_fixture_mode.py::test_bar", "root"),
],
)
def test_package_is_the_first_dir_under_the_suite_root(self, nodeid: str, expected: str) -> None:
assert package_from_nodeid(nodeid) == expected
class TestSourceFromLocation:
@pytest.mark.parametrize("path", ["a2a/test_a2a_agent_e2e.py", "tests/e2e/a2a/test_a2a_agent_e2e.py"])
def test_path_is_repo_relative_however_pytest_was_started(self, path: str) -> None:
assert source_from_location(path, 40) == "tests/e2e/a2a/test_a2a_agent_e2e.py:41"
def test_line_is_emitted_one_based(self) -> None:
"""pytest.Item.location counts from 0; editors, tracebacks and GitHub's
#L anchor all count from 1, and an off-by-one lands on the decorator."""
assert source_from_location("a2a/test_x.py", 0) == "tests/e2e/a2a/test_x.py:1"
def test_top_level_suite_file_sits_directly_under_the_suite_root(self) -> None:
assert source_from_location("test_fixture_mode.py", 39) == "tests/e2e/test_fixture_mode.py:40"
@pytest.mark.parametrize(
("path", "lineno"),
[
("a2a/test_x.py", None),
("/app/e2e/a2a/test_x.py", 40),
("C:\\app\\e2e\\a2a\\test_x.py", 40),
("../conftest.py", 40),
("", 40),
],
)
def test_nothing_linkable_yields_empty_rather_than_a_guess(self, path: str, lineno: int | None) -> None:
"""A colon is rejected on two counts: it is how a Windows absolute path
arrives, and `path:line` cannot represent one in the path half."""
assert source_from_location(path, lineno) == ""
class TestResultProperties:
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", f"tests/e2e/test_junit_properties.py:{test.__code__.co_firstlineno}"),
)
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 = 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"]
class TestSuiteRoot:
def test_suite_root_names_this_file_s_real_home(self) -> None:
"""SUITE_ROOT is hardcoded because the runner image has no repo to read it
from. Where there IS a checkout, prove the constant still points at us --
otherwise a moved tests/e2e/ ships links that 404."""
root = repo_root()
if root is None:
pytest.skip("no checkout above this file (the runner image copies tests/e2e/ to /app/e2e)")
assert (root / SUITE_ROOT / Path(__file__).name).resolve() == Path(__file__).resolve()
class TestDedupeCovers:
def test_ids_are_unique_order_preserving_and_non_empty_strings(self) -> None:
assert dedupe_covers([("A", "B"), ("B", ""), ("C", 7)]) == ("A", "B", "C")

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

4
tests/proxy_unit_tests/test_proxy_server.py
View file

@ -3076,7 +3076,9 @@ async def test_update_config_success_callback_normalization():
admin_user = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-test"
)
await proxy_server.update_config(config_update, user_api_key_dict=admin_user)
request = MagicMock()
request.json = AsyncMock(return_value={"litellm_settings": {"success_callback": ["SQS", "sQs"]}})
await proxy_server.update_config(config_update, request=request, user_api_key_dict=admin_user)
assert (
"litellm_settings" in upserted

402
tests/router_unit_tests/test_router_anthropic_messages_fallback.py Normal file
View file

@ -0,0 +1,402 @@
"""
Unit tests for safeguard-refusal fallback on the /v1/messages router surface.
An Anthropic safeguard refusal is an HTTP 200 whose body carries
stop_reason "refusal" plus a stop_details object; the router converts it
into a ContentPolicyViolationError so the content-policy fallback chain
runs, but only when a matching fallback is configured. A plain refusal
without stop_details, or any refusal with nothing configured, must reach
the client byte-identical.
The upstream is faked at the HTTP boundary by intercepting the third-party
transport (httpx.AsyncClient.send), so requests run litellm's real
transformation, allowlist, and streaming pipeline end to end.
"""
import json
from typing import Any, AsyncIterator
from unittest.mock import patch
import httpx
import pytest
from litellm import Router
from litellm.router_utils.fallback_event_handlers import (
PRE_ROUTING_SELECTED_MODEL_KEY,
record_pre_routing_selection,
)
REFUSAL_RESPONSE: dict[str, Any] = {
"id": "msg_refusal",
"type": "message",
"role": "assistant",
"model": "claude-fable-5",
"content": [],
"stop_reason": "refusal",
"stop_sequence": None,
"stop_details": {"category": "cyber", "explanation": "flagged"},
"usage": {"input_tokens": 25, "output_tokens": 1},
}
PLAIN_REFUSAL_RESPONSE: dict[str, Any] = {k: v for k, v in REFUSAL_RESPONSE.items() if k != "stop_details"}
OK_RESPONSE: dict[str, Any] = {
"id": "msg_ok",
"type": "message",
"role": "assistant",
"model": "claude-opus-5",
"content": [{"type": "text", "text": "hello"}],
"stop_reason": "end_turn",
"stop_sequence": None,
"usage": {"input_tokens": 25, "output_tokens": 2},
}
def _sse(event: str, data: dict[str, Any]) -> bytes:
return f"event: {event}\ndata: {json.dumps(data)}\n\n".encode()
REFUSAL_STREAM_FRAMES: tuple[bytes, ...] = (
_sse("message_start", {"type": "message_start", "message": {**REFUSAL_RESPONSE, "stop_reason": None}}),
_sse(
"message_delta",
{
"type": "message_delta",
"delta": {"stop_reason": "refusal", "stop_details": {"category": "cyber"}},
"usage": {"output_tokens": 1},
},
),
_sse("message_stop", {"type": "message_stop"}),
)
OK_STREAM_FRAMES: tuple[bytes, ...] = (
_sse("message_start", {"type": "message_start", "message": {**OK_RESPONSE, "stop_reason": None}}),
_sse(
"content_block_delta",
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hello"}},
),
_sse("message_stop", {"type": "message_stop"}),
)
def _split_frames_mid_data_line(frames: tuple[bytes, ...]) -> tuple[bytes, ...]:
"""Split each frame's data line in half, modeling a transport chunk boundary."""
return tuple(part for frame in frames for part in (frame[: len(frame) // 2], frame[len(frame) // 2 :]))
class _FrameStream(httpx.AsyncByteStream):
def __init__(self, frames: tuple[bytes, ...]) -> None:
self._frames = frames
async def __aiter__(self) -> AsyncIterator[bytes]:
for frame in self._frames:
yield frame
async def aclose(self) -> None:
return None
class FakeAnthropicUpstream:
"""Intercepts the third-party transport (httpx.AsyncClient.send): refuses on fable
models, answers on others. The router deliberately does not forward caller-injected
clients, so the transport is the seam that exercises the real litellm pipeline."""
def __init__(
self,
refusal_body: dict[str, Any] = REFUSAL_RESPONSE,
refusal_frames: tuple[bytes, ...] = REFUSAL_STREAM_FRAMES,
) -> None:
self.refusal_body = refusal_body
self.refusal_frames = refusal_frames
self.calls: list[str] = []
self.bodies: list[dict[str, Any]] = []
async def send(self, request: httpx.Request, **kwargs: Any) -> httpx.Response:
body = json.loads(request.content or b"{}")
model = body.get("model", "")
self.calls.append(model)
self.bodies.append(body)
refuses = "fable" in model
if body.get("stream"):
frames = self.refusal_frames if refuses else OK_STREAM_FRAMES
return httpx.Response(
200,
stream=_FrameStream(frames),
headers={"content-type": "text/event-stream"},
request=request,
)
return httpx.Response(200, json=self.refusal_body if refuses else OK_RESPONSE, request=request)
def install(self):
async def _send(_client: httpx.AsyncClient, request: httpx.Request, **kwargs: Any) -> httpx.Response:
return await self.send(request, **kwargs)
return patch("httpx.AsyncClient.send", new=_send)
FABLE_TIER = {
"model_name": "fable-tier",
"litellm_params": {"model": "anthropic/claude-fable-5", "api_key": "sk-test"},
}
OPUS_TARGET = {
"model_name": "opus-target",
"litellm_params": {"model": "anthropic/claude-opus-5", "api_key": "sk-test"},
}
def _router(content_policy_fallbacks: list | None) -> Router:
return Router(model_list=[FABLE_TIER, OPUS_TARGET], content_policy_fallbacks=content_policy_fallbacks)
async def _collect(stream: AsyncIterator[bytes]) -> bytes:
return b"".join([chunk async for chunk in stream])
@pytest.mark.asyncio
async def test_non_streaming_refusal_with_fallback_row_returns_fallback_response():
fake = FakeAnthropicUpstream()
router = _router(content_policy_fallbacks=[{"fable-tier": ["opus-target"]}])
with fake.install():
response = await router.aanthropic_messages(
model="fable-tier", max_tokens=16, messages=[{"role": "user", "content": "hi"}]
)
assert response["stop_reason"] == "end_turn"
assert response["id"] == "msg_ok"
assert len(fake.calls) == 2
assert "claude-opus-5" in fake.calls[1]
@pytest.mark.asyncio
@pytest.mark.parametrize(
"content_policy_fallbacks, upstream_body",
[
(None, REFUSAL_RESPONSE),
([{"unrelated-group": ["opus-target"]}], REFUSAL_RESPONSE),
([{"fable-tier": ["opus-target"]}], PLAIN_REFUSAL_RESPONSE),
],
ids=["nothing-configured", "row-for-other-group", "refusal-without-stop-details"],
)
async def test_non_streaming_refusal_passes_through_untouched(content_policy_fallbacks, upstream_body):
fake = FakeAnthropicUpstream(refusal_body=upstream_body)
router = _router(content_policy_fallbacks=content_policy_fallbacks)
with fake.install():
response = await router.aanthropic_messages(
model="fable-tier", max_tokens=16, messages=[{"role": "user", "content": "hi"}]
)
assert response["stop_reason"] == "refusal"
assert response.get("stop_details") == upstream_body.get("stop_details")
assert len(fake.calls) == 1
@pytest.mark.asyncio
async def test_streaming_refusal_with_fallback_row_streams_fallback_frames():
fake = FakeAnthropicUpstream()
router = _router(content_policy_fallbacks=[{"fable-tier": ["opus-target"]}])
with fake.install():
stream = await router.aanthropic_messages(
model="fable-tier", max_tokens=16, stream=True, messages=[{"role": "user", "content": "hi"}]
)
body = await _collect(stream)
assert b'"refusal"' not in body
assert b"text_delta" in body
assert len(fake.calls) == 2
@pytest.mark.asyncio
async def test_streaming_refusal_split_across_chunks_still_falls_back():
fake = FakeAnthropicUpstream(refusal_frames=_split_frames_mid_data_line(REFUSAL_STREAM_FRAMES))
router = _router(content_policy_fallbacks=[{"fable-tier": ["opus-target"]}])
with fake.install():
stream = await router.aanthropic_messages(
model="fable-tier", max_tokens=16, stream=True, messages=[{"role": "user", "content": "hi"}]
)
body = await _collect(stream)
assert b'"refusal"' not in body
assert b"text_delta" in body
assert len(fake.calls) == 2
@pytest.mark.asyncio
async def test_streaming_refusal_without_fallback_row_passes_frames_through():
fake = FakeAnthropicUpstream()
router = _router(content_policy_fallbacks=None)
with fake.install():
stream = await router.aanthropic_messages(
model="fable-tier", max_tokens=16, stream=True, messages=[{"role": "user", "content": "hi"}]
)
body = await _collect(stream)
assert b'"stop_reason": "refusal"' in body
assert b"stop_details" in body
assert len(fake.calls) == 1
@pytest.mark.asyncio
async def test_streaming_refusal_on_routed_tier_matches_tier_keyed_row_without_inbound_metadata():
"""The pre-routing hook's tier stamp must reach the mid-stream fallback lookup even when the
request carries no metadata bucket at all (the snapshot is taken before the request runs)."""
fake = FakeAnthropicUpstream()
smart_router = {
"model_name": "smart-router",
"litellm_params": {
"model": "auto_router/complexity_router",
"complexity_router_config": {
"tiers": {"SIMPLE": "fable-tier", "MEDIUM": "fable-tier", "COMPLEX": "fable-tier"}
},
"complexity_router_default_model": "fable-tier",
},
"model_info": {"id": "router-1", "db_model": True},
}
router = Router(
model_list=[FABLE_TIER, OPUS_TARGET, smart_router],
content_policy_fallbacks=[{"fable-tier": ["opus-target"]}],
ignore_invalid_deployments=True,
)
with fake.install():
stream = await router.aanthropic_messages(
model="smart-router", max_tokens=16, stream=True, messages=[{"role": "user", "content": "hi"}]
)
body = await _collect(stream)
assert b'"refusal"' not in body
assert b"text_delta" in body
assert len(fake.calls) == 2
@pytest.mark.asyncio
async def test_caller_forged_tier_stamp_cannot_pick_the_streaming_fallback_chain():
fake = FakeAnthropicUpstream()
router = _router(content_policy_fallbacks=[{"forged-tier": ["opus-target"]}])
with fake.install():
stream = await router.aanthropic_messages(
model="fable-tier",
max_tokens=16,
stream=True,
messages=[{"role": "user", "content": "hi"}],
litellm_metadata={PRE_ROUTING_SELECTED_MODEL_KEY: "forged-tier"},
)
body = await _collect(stream)
assert b'"stop_reason": "refusal"' in body
assert len(fake.calls) == 1
@pytest.mark.asyncio
async def test_tier_stamp_never_reaches_provider_bound_metadata():
"""On /v1/messages the top-level metadata dict is Anthropic's own request field, so the
routed-tier stamp must never appear in any upstream body even when the client sends one."""
fake = FakeAnthropicUpstream()
smart_router = {
"model_name": "smart-router",
"litellm_params": {
"model": "auto_router/complexity_router",
"complexity_router_config": {
"tiers": {"SIMPLE": "fable-tier", "MEDIUM": "fable-tier", "COMPLEX": "fable-tier"}
},
"complexity_router_default_model": "fable-tier",
},
"model_info": {"id": "router-1", "db_model": True},
}
router = Router(
model_list=[FABLE_TIER, OPUS_TARGET, smart_router],
content_policy_fallbacks=[{"fable-tier": ["opus-target"]}],
ignore_invalid_deployments=True,
)
with fake.install():
response = await router.aanthropic_messages(
model="smart-router",
max_tokens=16,
messages=[{"role": "user", "content": "hi"}],
metadata={"user_id": "u1"},
)
assert response["stop_reason"] == "end_turn"
assert len(fake.bodies) == 2
for body in fake.bodies:
assert body.get("metadata") == {"user_id": "u1"}
def test_record_pre_routing_selection_writes_only_the_internal_bucket():
"""The Anthropic request's own metadata field must never carry the tier stamp."""
kwargs = {"metadata": {"user_id": "u1"}, "litellm_metadata": {}}
record_pre_routing_selection(kwargs, "tier-x")
assert kwargs["litellm_metadata"] == {PRE_ROUTING_SELECTED_MODEL_KEY: "tier-x"}
assert kwargs["metadata"] == {"user_id": "u1"}
def test_refusal_gate_keys_on_pre_routing_tier_stamp():
router = _router(content_policy_fallbacks=[{"tier-group": ["opus-target"]}])
def anthropic_messages(**kwargs: Any) -> None:
return None
refusal_kwargs = {"litellm_metadata": {PRE_ROUTING_SELECTED_MODEL_KEY: "tier-group"}}
assert (
router._should_raise_anthropic_refusal_error(
model="router-group",
original_generic_function=anthropic_messages,
response=dict(REFUSAL_RESPONSE),
kwargs=refusal_kwargs,
)
is True
)
assert (
router._should_raise_anthropic_refusal_error(
model="router-group",
original_generic_function=anthropic_messages,
response=dict(REFUSAL_RESPONSE),
kwargs={},
)
is False
)
def test_has_content_policy_fallback_default_fallbacks_arm():
router = Router(model_list=[OPUS_TARGET], fallbacks=[{"*": ["opus-target"]}])
assert router._has_content_policy_fallback("any-group", {}) is True
assert router._has_content_policy_fallback("any-group", {"content_policy_fallbacks": [{"other": ["x"]}]}) is False
def test_get_fallback_model_group_for_lookup_groups_orders_tier_before_requested():
router = _router(content_policy_fallbacks=None)
fallbacks = [{"tier1": ["backup-a"]}, {"smart-router": ["backup-b"]}]
assert router._get_fallback_model_group_for_lookup_groups(
fallbacks=fallbacks, lookup_groups=("tier1", "smart-router")
) == ["backup-a"]
assert router._get_fallback_model_group_for_lookup_groups(
fallbacks=fallbacks, lookup_groups=("tier9", "smart-router")
) == ["backup-b"]
assert router._get_fallback_model_group_for_lookup_groups(fallbacks=fallbacks, lookup_groups=()) is None
def test_refusal_gate_ignores_other_generic_call_types():
router = _router(content_policy_fallbacks=[{"fable-tier": ["opus-target"]}])
def aresponses(**kwargs: Any) -> None:
return None
assert (
router._should_raise_anthropic_refusal_error(
model="fable-tier",
original_generic_function=aresponses,
response=dict(REFUSAL_RESPONSE),
kwargs={},
)
is False
)

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.

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