scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

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

Browse source 
# Conflicts:
#	enterprise/litellm_enterprise/proxy/hooks/managed_files.py
This commit is contained in:
mateo-berri 2026-08-06 19:00:51 -07:00
commit 02d60847ee
130 changed files with 5364 additions and 645 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

21
.git-blame-ignore-revs
View file

@ -17,3 +17,24 @@
# style: unify ruff format width on 120 (#31518)
48b5a5a0cc5a694a11219416ee0b6eb6e620e74e
# refactor(imports): move collections.abc names out of typing (#35495)
397e8e4918777e4e60a7f5e88699e0a9a7dabb3d
# refactor(lint): apply every safe ruff autofix and zero 28 strict-rule budgets (#35495)
b604e2b20c6db2099085a2f0e59b7e99e87eed6f
# refactor(logging): drop redundant !s conversion flags from f-strings (#35546)
7b2d3440cba3160277470f7a0180098ae9b87864
# perf: build log messages lazily so filtered-out log records cost nothing (#35703)
c9887a1f94bc1e7e4bdfe64d640f0509a0bc19dd
# feat(lint): enforce Final on locals and freeze function parameters (#35807)
2708620d6a599cc73c1950a942d26ac26a7ed3d4
# chore(lint): remove litellm/types from the ruff lint exclusion (#35926)
4e32a8bf6a1e1af1e04b67c759841ccef44b2235
# chore(lint): strip inert type: ignore comments and zero LIT009/LIT010/LIT011 headroom (#35928)
338e411103ad5d7003e97f34f04fa36bca542dbe

18
.github/workflows/publish-basedpyright-base-counts.yml vendored
View file

@ -43,22 +43,14 @@ jobs:
with:
version: "0.10.9"
- name: Install dependencies
run: |
uv sync --frozen --group proxy-dev --group e2e-dev
# Mirrors test-linting.yml's lint job: basedpyright resolves Prisma's
# generated client only after `prisma generate`, and the published counts
# must match what that job would measure for the same tree.
- name: Generate Prisma client
# The gate provisions its own measurement env (.venv-typecheck: a frozen
# uv sync of its canonical dependency groups plus a generated Prisma
# client), so no install step here can drift from what local runs measure.
- name: Emit basedpyright counts for HEAD
env:
PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache
run: |
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
- name: Emit basedpyright counts for HEAD
run: |
uv run --no-sync python scripts/type_check_gate.py --emit-counts-dir "$RUNNER_TEMP/basedpyright-counts"
python scripts/type_check_gate.py --emit-counts-dir "$RUNNER_TEMP/basedpyright-counts"
counts_file=$(ls "$RUNNER_TEMP"/basedpyright-counts/basedpyright-counts-*.json)
echo "COUNTS_ARTIFACT_NAME=$(basename "$counts_file" .json)" >> "$GITHUB_ENV"

1
.github/workflows/test-linting.yml vendored
View file

@ -115,6 +115,7 @@ jobs:
- name: Check basedpyright budget (delta vs base)
env:
GH_TOKEN: ${{ github.token }}
PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache
run: |
uv run --no-sync python scripts/type_check_gate.py --base "$GATE_BASE_SHA"

1
.gitignore vendored
View file

@ -1,5 +1,6 @@
.python-version
.venv
.venv-typecheck
.venv_policy_test
.env
.claude

10
CLAUDE.md
View file

@ -29,7 +29,7 @@ Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We pref
If you ever make public-facing PR descriptions, comments, issues, commit messages, etc., always follow these guidelines to sound less AI-y:
- don't use emojis
- don't use "—". Instead, reach for ",", ".", conjunction words, ":", ";", etc. in descending order of preference: vary among them, weighted toward the front of the list, and skip "," where it would cause a comma splice or the sentence is getting long. Overusing any one of them, ";" especially, also feels AI-y
- don't use "—". Instead, reach for ",", ".", conjunction words, ":", ";", etc. in descending order of preference: vary among them, weighted toward the front of the list, and skip "," where it would cause a comma splice or the sentence is getting long. Overusing any one of them, ";" especially, also feels AI-y. A word cap does not penalize you for adding more sentences: when writing under tight word budgets, prefer a period split or a conjunction over ";", and keep to at most one ";" per message
- don't use the pattern "It's not X, it's Y", "You're not X, you're Y", etc.
- don't use bulleted or numbered lists unless it would be nonsensical not to. Instead, prefer prose
- don't add a trailing "." at the end of paragraphs (just like this file). That means every paragraph, not just the last one (of the markdown file, PR description, GitHub comment, etc.). Rule of thumb: if you're adding new line(s) before the next sentence, don't add a "."
@ -41,11 +41,11 @@ Python max line length is 120, not 88
When you fix violations gated by `ruff-strict-budget.json`, `type-discipline-budget.json`, or `basedpyright-code-budget.json`, run `make lint-budget-update` and commit the lowered limits so the ceilings ratchet down instead of leaving stale headroom. It measures the working tree, so it must contain exactly the fixes you're committing
`make pre-commit` always saves its complete output to a per-worktree log file and prints that path as its first and last output lines. To inspect a run, read or grep that log instead of re-running the multi-minute checks just to see a different slice, and re-run only after the working tree actually changed
`make pre-commit` saves its complete output to a log file in .git (overwriting previous pre-commit logs) and prints that path as its first and last output lines. To inspect a run, read or grep that log instead of re-running the multi-minute checks just to see a different slice
If you're trying to create a new function that relies on untyped stuff, instead of adding more Any's and pushing `reportAny` / `reportExplicitAny` closer to their basedpyright ceilings, just validate it in the caller with Pydantic (a model or `TypeAdapter` that returns the typed thing or raises will do) and then pass the now typed variable in
If you get an LIT001 or LIT002 fail, refactor the code to follow functional programming best practices rather than introducing mutable data structures. For example, build values in one shot with comprehensions or generators wrapped in `tuple()` / `frozenset()` instead of seeding an empty `list`/`dict`/`set` and mutating it over time. Ideally, `# mutable-ok` is never used; reach for it only as a genuine last resort when an immutable rewrite is truly impossible, and always pair it with a real reason
If you get an LIT001 or LIT002 fail, refactor the code to follow functional programming best practices rather than introducing mutable data structures. For example, build values in one shot with comprehensions or generators wrapped in `tuple()` / `MappingProxyType()` / `frozenset()` instead of seeding an empty `list`/`dict`/`set` and mutating it over time. Ideally, `# mutable-ok` is never used; reach for it only as a genuine last resort when an immutable rewrite is truly impossible, and always pair it with a real reason
Every lint or type suppression must name the exact rule inside brackets and carry a reason comment, e.g. `# pyright: ignore[reportArgumentType] # stubs lack async overload` or `# noqa: TID251 # <reason>`. `# type: ignore` is banned (LIT009): pyrightconfig.json sets `enableTypeIgnoreComments` to false, so it silently does nothing
@ -59,7 +59,7 @@ Do not add `Co-Authored-By: Claude` or any Claude attribution to commit messages
When working on a PR, keep the PR description in sync with new commits being made
Replies/rebuttals to AI PR review bots must be 15-25 word human-readable replies
All GitHub comments must be human-readable and 15-25 words max
Monkeypatching attributes of a class to do testing is an anti-pattern. Prefer dependency-injecting things into classes. That way, at unit test time, you can pass a mocked dependency in
@ -72,7 +72,7 @@ Follow these coding conventions for new/updated code (a three-line fix in a lega
- Composition over inheritance
- Never-nester: early returns over deep nesting
- Don't throw; model failures as values (One function (e.g., raise_public) maps error union to existing public exception contracts via exhaustive match + assert_never)
- No mutation; don't reassign variables, global or local. Instead of mutable lists and dicts, prefer tuples, frozen dataclasses (with slots=True), etc.
- No mutation; don't reassign variables, global or local. Instead of mutable lists and dicts, prefer tuples, frozen dataclasses (with slots=True), `MappingProxyType`, etc.
- Annotate every variable with `: Final` (LIT010). Unpacking and walrus targets cannot carry the annotation, so they are implicitly final. Don't rebind them. Never rebind or mutate function parameters (LIT011); `self`/`cls` attribute stores are the exception. If rebinding or in-place mutation is truly unavoidable, suppress with `# rebind-ok: <reason>` explaining why
- Use dependency injection
- Fully typed; no `Any` or coarse types like `dict[str, Any]` or just `dict`. Every function parameter must be strongly typed

8
Makefile
View file

@ -124,10 +124,10 @@ lint-fetch-base:
git fetch origin litellm_internal_staging
# Mirror test-linting.yml's lint job environment: the proxy-dev group plus a generated
# Prisma client, so basedpyright resolves the same modules CI does (without the generated
# client the DB wrappers typed against it degrade to Unknown, drifting the budget from
# CI's). --inexact tops up the venv instead of pruning the proxy extras gen:api and the
# running proxy need.
# Prisma client, so `basedpyright tests/e2e` resolves the same modules CI does. The
# budget gate itself no longer measures here (scripts/type_check_gate.py provisions its
# own .venv-typecheck). --inexact tops up the venv instead of pruning the proxy extras
# gen:api and the running proxy need.
lint-install:
$(UV) sync --inexact --frozen --group proxy-dev --group e2e-dev
$(UV_RUN) python scripts/prisma_generate_if_needed.py

28
basedpyright-code-budget.json
View file

@ -1,18 +1,18 @@
{
"reportAny": {
"limit": 29204
"limit": 28842
},
"reportArgumentType": {
"limit": 2635
"limit": 2634
},
"reportAssignmentType": {
"limit": 329
},
"reportAttributeAccessIssue": {
"limit": 516
"limit": 514
},
"reportCallIssue": {
"limit": 123
"limit": 117
},
"reportConstantRedefinition": {
"limit": 40
@ -24,7 +24,7 @@
"limit": 19
},
"reportExplicitAny": {
"limit": 9227
"limit": 9103
},
"reportFunctionMemberAccess": {
"limit": 7
@ -54,10 +54,10 @@
"limit": 0
},
"reportMissingParameterType": {
"limit": 5850
"limit": 5843
},
"reportMissingTypeArgument": {
"limit": 15833
"limit": 15816
},
"reportMissingTypeStubs": {
"limit": 40
@ -99,34 +99,34 @@
"limit": 0
},
"reportUnknownArgumentType": {
"limit": 45242
"limit": 45110
},
"reportUnknownLambdaType": {
"limit": 113
},
"reportUnknownMemberType": {
"limit": 40340
"limit": 39838
},
"reportUnknownParameterType": {
"limit": 20293
"limit": 20237
},
"reportUnknownVariableType": {
"limit": 31796
"limit": 31383
},
"reportUnnecessaryCast": {
"limit": 122
},
"reportUnnecessaryComparison": {
"limit": 703
"limit": 701
},
"reportUnnecessaryContains": {
"limit": 5
},
"reportUnnecessaryIsInstance": {
"limit": 865
"limit": 864
},
"reportUntypedBaseClass": {
"limit": 72
"limit": 0
},
"reportUntypedFunctionDecorator": {
"limit": 33

15
enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
View file

@ -498,6 +498,7 @@ class CheckBatchCost:
},
"metadata": {
"user_api_key_user_id": creator_user_id,
"user_api_key_team_id": getattr(job, "team_id", None),
**user_info,
},
},
@ -656,6 +657,20 @@ class CheckBatchCost:
elif response.status in ("failed", "expired", "cancelled"):
try:
from litellm.proxy.openai_files_endpoints.common_utils import (
_is_base64_encoded_unified_file_id,
ensure_batch_response_managed_file_ids,
)
response.id = job.unified_object_id
await ensure_batch_response_managed_file_ids(
response=response,
managed_files_obj=self.proxy_logging_obj.get_proxy_hook("managed_files"),
prisma_client=self.prisma_client,
verbose_proxy_logger=verbose_proxy_logger,
db_batch_object=job,
unified_batch_id=_is_base64_encoded_unified_file_id(job.unified_object_id),
)
update_data = {
"status": response.status,
"file_object": response.model_dump_json(),

49
enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py
View file

@ -1,10 +1,10 @@
"""
Polls LiteLLM_ManagedObjectTable to check if the response is complete.
Cost tracking is handled automatically by litellm.aget_responses().
Cost tracking is handled automatically by the get-responses call.
"""
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING
from typing import TYPE_CHECKING, Dict, Optional, cast
import litellm
from litellm._logging import verbose_proxy_logger
@ -13,11 +13,15 @@ from litellm.constants import (
MAX_OBJECTS_PER_POLL_CYCLE,
STALE_OBJECT_CLEANUP_BATCH_SIZE,
)
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.types.llms.openai import ResponsesAPIResponse
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.router import Router
TERMINAL_RESPONSE_STATUSES = frozenset({"completed", "failed", "cancelled", "incomplete"})
class CheckResponsesCost:
def __init__(
@ -33,6 +37,28 @@ class CheckResponsesCost:
self.prisma_client: PrismaClient = prisma_client
self.llm_router: Router = llm_router
async def _get_response(
self,
response_id: str,
litellm_metadata: Dict[str, str],
) -> ResponsesAPIResponse:
"""Fetch the upstream response, using deployment credentials when available.
LiteLLM-encoded response IDs carry the ``model_id`` of the deployment that
served the original request, so routing through ``llm_router`` applies that
deployment's ``api_base`` / ``api_key`` / ``api_version``, exactly like
``GET /v1/responses/{id}`` does. ``litellm.aget_responses`` on its own only
sees provider env vars, so it fails for every deployment whose credentials
live in the config; the row then never leaves ``queued``.
"""
model_id: Optional[str] = ResponsesAPIRequestUtils.get_model_id_from_response_id(response_id)
if model_id is None or self.llm_router.get_deployment(model_id=model_id) is None:
return await litellm.aget_responses(response_id=response_id, litellm_metadata=litellm_metadata)
router_response = await self.llm_router.aget_responses(
response_id=response_id, litellm_metadata=litellm_metadata
)
return cast(ResponsesAPIResponse, router_response)
async def _expire_stale_rows(
self, cutoff: datetime, batch_size: int
) -> int:
@ -87,8 +113,8 @@ class CheckResponsesCost:
Check if background responses are complete and track their cost.
- Get all status="queued" or "in_progress" and file_purpose="response" jobs
- Query the provider to check if response is complete
- Cost is automatically tracked by litellm.aget_responses()
- Mark completed/failed/cancelled responses as complete in the database
- Cost is automatically tracked by the get-responses call
- Mark responses in a terminal state as complete in the database
"""
try:
await self._cleanup_stale_managed_objects()
@ -134,7 +160,7 @@ class CheckResponsesCost:
litellm_metadata["model"] = model_name
litellm_metadata["model_group"] = model_name # Use same value for model_group
response = await litellm.aget_responses(
response = await self._get_response(
response_id=responses_id_security,
litellm_metadata=litellm_metadata,
)
@ -144,21 +170,14 @@ class CheckResponsesCost:
)
except Exception as e:
verbose_proxy_logger.info(
verbose_proxy_logger.warning(
f"Skipping job {unified_object_id} due to error: {e}"
)
continue
# Check if response is in a terminal state
if response.status == "completed":
if response.status in TERMINAL_RESPONSE_STATUSES:
verbose_proxy_logger.info(
f"Response {unified_object_id} is complete. Cost automatically tracked by aget_responses."
)
completed_jobs.append(job)
elif response.status in ["failed", "cancelled"]:
verbose_proxy_logger.info(
f"Response {unified_object_id} has status {response.status}, marking as complete"
f"Response {unified_object_id} has terminal status {response.status}, marking as complete"
)
completed_jobs.append(job)

38
enterprise/litellm_enterprise/proxy/hooks/managed_files.py
View file

@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any, Dict, Final, List, Literal, Optional, Uni
from uuid import NAMESPACE_URL, uuid5
from fastapi import HTTPException
from pydantic import ValidationError
import litellm
from litellm import Router, verbose_logger
@ -81,6 +82,14 @@ else:
PrismaClient = Any
def _sanitized_parse_error(e: Exception) -> str:
return (
str(e.errors(include_input=False, include_url=False, include_context=False))
if isinstance(e, ValidationError)
else type(e).__name__
)
def _decode_json_blob(blob: object) -> object:
return json.loads(blob) if isinstance(blob, str) else blob
@ -89,12 +98,28 @@ def _parse_managed_batch_row(row: "PrismaManagedObjectRow") -> Optional[LiteLLMB
try:
batch_obj: Final = LiteLLMBatch.model_validate(_decode_json_blob(row.file_object))
except Exception as e:
verbose_logger.warning(f"Failed to parse batch object {row.unified_object_id}: {e}")
verbose_logger.warning(
f"Failed to parse batch object {row.unified_object_id}: {_sanitized_parse_error(e)}"
)
return None
batch_obj.id = row.unified_object_id
return batch_obj
def _parse_managed_file_object(
raw_file_object: object, unified_file_id: str
) -> Optional[OpenAIFileObject]:
if raw_file_object is None:
return None
try:
return OpenAIFileObject.model_validate(raw_file_object)
except Exception as e:
verbose_logger.warning(
f"Failed to parse managed file object {unified_file_id}: {_sanitized_parse_error(e)}"
)
return None
class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
# Class variables or attributes
def __init__(
@ -440,11 +465,14 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
}
)
return [
OpenAIFileObject.model_validate(row.file_object).model_copy(
update={"id": row.unified_file_id}
)
parsed_file_object.model_copy(update={"id": row.unified_file_id})
for row in file_ids
if row.file_object is not None
if (
parsed_file_object := _parse_managed_file_object(
row.file_object, row.unified_file_id
)
)
is not None
]
async def check_managed_file_id_access(

4
enterprise/pyproject.toml
View file

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

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

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

33
litellm/caching/redis_semantic_cache.py
View file

@ -13,6 +13,7 @@ import ast
import asyncio
import json
import os
from collections.abc import Callable, Mapping
from typing import Any, Final, cast
import litellm
@ -47,7 +48,7 @@ class RedisSemanticCache(BaseCache):
similarity_threshold: float | None = None,
embedding_model: str = "text-embedding-ada-002",
index_name: str | None = None,
**kwargs,
**kwargs: object,
):
"""
Initialize the Redis Semantic Cache.
@ -150,11 +151,11 @@ class RedisSemanticCache(BaseCache):
def _init_semantic_cache(
self,
semantic_cache_cls: Any,
semantic_cache_cls: Callable[..., object],
index_name: str,
redis_url: str,
cache_vectorizer: Any,
) -> Any:
cache_vectorizer: object,
) -> object:
def _is_schema_mismatch(exc: ValueError) -> bool:
error_message: Final = str(exc).lower()
return any(phrase in error_message for phrase in ("schema does not match", "index schema"))
@ -206,12 +207,12 @@ class RedisSemanticCache(BaseCache):
def _get_cache_filters(self, key: str) -> dict[str, str]:
return {self.CACHE_KEY_FIELD_NAME: str(key)}
def _get_cache_key_filter_expression(self, key: str) -> Any:
def _get_cache_key_filter_expression(self, key: str) -> object:
from redisvl.query.filter import Tag
return Tag(self.CACHE_KEY_FIELD_NAME) == str(key)
def _cache_hit_matches_key(self, cache_hit: dict[str, Any], key: str) -> bool:
def _cache_hit_matches_key(self, cache_hit: Mapping[str, object], key: str) -> bool:
# Pre-isolation entries with no ``litellm_cache_key`` field cannot be
# safely reassigned to a caller's scope and are treated as misses.
cached_key = cache_hit.get(self.CACHE_KEY_FIELD_NAME)
@ -297,7 +298,7 @@ class RedisSemanticCache(BaseCache):
return
@staticmethod
def _coerce_response_input_value(value: Any) -> Any:
def _coerce_response_input_value(value: object) -> object:
model_dump: Final = getattr(value, "model_dump", None)
if callable(model_dump):
return model_dump()
@ -340,7 +341,7 @@ class RedisSemanticCache(BaseCache):
)
return embedding_response["data"][0]["embedding"]
def _get_cache_logic(self, cached_response: Any) -> Any:
def _get_cache_logic(self, cached_response: Any) -> object:
"""
Process the cached response to prepare it for use.
@ -369,7 +370,7 @@ class RedisSemanticCache(BaseCache):
return cached_response
def set_cache(self, key: str, value: Any, **kwargs) -> None:
def set_cache(self, key: str, value: object, **kwargs) -> None:
"""
Store a value in the semantic cache.
@ -405,7 +406,7 @@ class RedisSemanticCache(BaseCache):
except Exception as e:
print_verbose(f"Error setting {value_str or value} in the Redis semantic cache: {e}")
def get_cache(self, key: str, **kwargs) -> Any:
def get_cache(self, key: str, **kwargs) -> object:
"""
Retrieve a semantically similar cached response.
@ -428,7 +429,7 @@ class RedisSemanticCache(BaseCache):
# Check the cache for semantically similar prompts in this exact
# LiteLLM cache-key scope.
prompt_embedding: Final = self._get_embedding(prompt, metadata=kwargs.get("metadata"))
check_kwargs: Final[dict[str, Any]] = {
check_kwargs: Final[Mapping[str, object]] = {
"prompt": prompt,
"vector": prompt_embedding,
"filter_expression": self._get_cache_key_filter_expression(key),
@ -508,7 +509,7 @@ class RedisSemanticCache(BaseCache):
print_verbose(f"Error generating async embedding: {e}")
raise ValueError(f"Failed to generate embedding: {e}") from e
async def async_set_cache(self, key: str, value: Any, **kwargs) -> None:
async def async_set_cache(self, key: str, value: object, **kwargs) -> None:
"""
Asynchronously store a value in the semantic cache.
@ -548,7 +549,7 @@ class RedisSemanticCache(BaseCache):
except Exception as e:
print_verbose(f"Error in async_set_cache: {e}")
async def async_get_cache(self, key: str, **kwargs) -> Any:
async def async_get_cache(self, key: str, **kwargs) -> object:
"""
Asynchronously retrieve a semantically similar cached response.
@ -573,7 +574,7 @@ class RedisSemanticCache(BaseCache):
# Check the cache for semantically similar prompts in this exact
# LiteLLM cache-key scope.
check_kwargs: Final[dict[str, Any]] = {
check_kwargs: Final[Mapping[str, object]] = {
"prompt": prompt,
"vector": prompt_embedding,
"filter_expression": self._get_cache_key_filter_expression(key),
@ -615,7 +616,7 @@ class RedisSemanticCache(BaseCache):
print_verbose(f"Error in async_get_cache: {e}")
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
async def _index_info(self) -> dict[str, Any]:
async def _index_info(self) -> Mapping[str, object]:
"""
Get information about the Redis index.
@ -625,7 +626,7 @@ class RedisSemanticCache(BaseCache):
aindex: Final = await self.llmcache._get_async_index()
return await aindex.info()
async def async_set_cache_pipeline(self, cache_list: list[tuple[str, Any]], **kwargs) -> None:
async def async_set_cache_pipeline(self, cache_list: list[tuple[str, Any]], **kwargs: object) -> None:
"""
Asynchronously store multiple values in the semantic cache.

3
litellm/integrations/arize/arize_phoenix.py
View file

@ -430,7 +430,8 @@ class ArizePhoenixLogger(OpenTelemetry):
otlp_auth_headers = None
if api_key is not None:
otlp_auth_headers = f"Authorization=Bearer {api_key}"
auth_header_key = "authorization" if protocol == "otlp_grpc" else "Authorization"
otlp_auth_headers = f"{auth_header_key}=Bearer {api_key}"
elif "app.phoenix.arize.com" in endpoint:
raise ValueError("PHOENIX_API_KEY must be set when using Phoenix Cloud (app.phoenix.arize.com).")

23
litellm/integrations/custom_guardrail.py
View file

@ -714,6 +714,29 @@ class CustomGuardrail(CustomLogger):
return result
def supports_scan_only_tool_results(self) -> bool:
"""Whether this guardrail can scan tool-result content.
Guardrails whose own role filtering only ever scans human-authored
messages override this to return False, so configuring them with
``scan_only_tool_results`` is rejected at initialization instead of
silently scanning nothing on every request.
"""
return True
def structured_messages_cover_full_request(self) -> bool:
"""Whether returned ``structured_messages`` span the whole request.
Translation handlers hand guardrails only the in-scope subset of the
conversation and merge a returned ``structured_messages`` list back
into the full request. A guardrail that already rebuilds the complete
conversation itself (like CrowdStrike AIDR with its skip filters
active) overrides this to return True so the handler installs the
returned list as-is instead of merging it a second time, which would
duplicate the out-of-scope messages.
"""
return False
def should_run_guardrail(
self,
data,

96
litellm/integrations/galileo.py
View file

@ -4,8 +4,9 @@ import json
import os
import re
import uuid
from datetime import datetime, timezone
from typing import Any, Final, cast
from collections.abc import Mapping, Sequence
from datetime import datetime, timezone, tzinfo
from typing import Any, Final, TypedDict, cast
import httpx
from pydantic import BaseModel, Field
@ -34,6 +35,17 @@ GALILEO_CLOUD_API_BASE_URL: Final = "https://api.galileo.ai"
GALILEO_MAX_IN_MEMORY_RECORDS: Final = 1000
class GalileoStandardLoggingFields(TypedDict, total=False):
call_type: str
model: str
prompt_tokens: int
completion_tokens: int
total_tokens: int
response_cost: float
startTime: float
endTime: float
class LLMResponse(BaseModel):
latency_ms: int
status_code: int
@ -59,7 +71,7 @@ class LLMResponse(BaseModel):
class GalileoObserve(CustomLogger):
def __init__(self) -> None:
self.in_memory_records: list[dict] = []
self.in_memory_records: list[Mapping[str, object]] = []
self.batch_size = 1
self.api_key = os.getenv("GALILEO_API_KEY")
self.project_id = os.getenv("GALILEO_PROJECT_ID")
@ -176,7 +188,7 @@ class GalileoObserve(CustomLogger):
return False
@staticmethod
def _galileo_input_messages(messages: Any | None, input_text: str) -> list[dict[str, str]]:
def _galileo_input_messages(messages: object, input_text: str) -> list[dict[str, str]]:
if isinstance(messages, dict):
messages = messages.get("messages")
if not messages:
@ -203,11 +215,11 @@ class GalileoObserve(CustomLogger):
return [{"role": "user", "content": input_text}]
@staticmethod
def _local_timezone():
def _local_timezone() -> tzinfo:
return datetime.now().astimezone().tzinfo or timezone.utc
@staticmethod
def _format_created_at(dt: datetime | Any) -> str:
def _format_created_at(dt: object) -> str:
"""Serialize timestamps as UTC ISO-8601 for Galileo."""
if not isinstance(dt, datetime):
return str(dt)
@ -226,7 +238,7 @@ class GalileoObserve(CustomLogger):
return created_at
@staticmethod
def _token_metrics_from_record(record: dict[str, Any]) -> dict[str, Any]:
def _token_metrics_from_record(record: Mapping[str, Any]) -> dict[str, Any]:
num_input_tokens: Final = int(record.get("num_input_tokens") or 0)
num_output_tokens: Final = int(record.get("num_output_tokens") or 0)
num_total_tokens = int(record.get("num_total_tokens") or 0)
@ -244,7 +256,7 @@ class GalileoObserve(CustomLogger):
@staticmethod
def _record_to_v2_span(
record: dict[str, Any],
record: Mapping[str, Any],
*,
trace_id: str,
span_id: str,
@ -275,7 +287,7 @@ class GalileoObserve(CustomLogger):
return span
@staticmethod
def _record_to_v2_trace(record: dict[str, Any]) -> dict[str, Any]:
def _record_to_v2_trace(record: Mapping[str, Any]) -> dict[str, Any]:
trace_id: Final = str(uuid.uuid4())
span_id: Final = str(uuid.uuid4())
created_at: Final = GalileoObserve._normalize_created_at(record.get("created_at", ""))
@ -295,7 +307,7 @@ class GalileoObserve(CustomLogger):
"spans": [GalileoObserve._record_to_v2_span(record, trace_id=trace_id, span_id=span_id)],
}
def _build_traces_payload(self, records: list[dict]) -> dict[str, Any]:
def _build_traces_payload(self, records: Sequence[Mapping[str, Any]]) -> dict[str, Any]:
payload: Final[dict[str, Any]] = {
"traces": [self._record_to_v2_trace(record) for record in records],
"logging_method": "api_direct",
@ -357,7 +369,7 @@ class GalileoObserve(CustomLogger):
@staticmethod
def _log_v2_payload_validation(payload: dict[str, Any]) -> None:
missing_fields: Final[list[str]] = []
traces: Final = payload.get("traces", [])
traces: Final[Sequence[object]] = payload.get("traces", [])
if not traces:
missing_fields.append("traces")
@ -385,7 +397,7 @@ class GalileoObserve(CustomLogger):
)
def _log_flush_payload(self, url: str, payload: dict[str, Any]) -> None:
traces: Final = payload.get("traces", [])
traces: Final[Sequence[object]] = payload.get("traces", [])
verbose_logger.debug(
"Galileo Logger flush URL: %s trace_count=%s",
url,
@ -415,8 +427,8 @@ class GalileoObserve(CustomLogger):
pass
@staticmethod
def _build_prompt(kwargs: dict[str, Any]) -> dict[str, Any]:
optional_params: Final = kwargs.get("optional_params", {}) or {}
def _build_prompt(kwargs: Mapping[str, Any]) -> dict[str, Any]:
optional_params: Final[Mapping[str, object]] = kwargs.get("optional_params", {}) or {}
prompt: Final[dict[str, Any]] = {"messages": kwargs.get("messages")}
if optional_params.get("functions") is not None:
prompt["functions"] = optional_params["functions"]
@ -425,13 +437,13 @@ class GalileoObserve(CustomLogger):
return prompt
@staticmethod
def _serialize_galileo_output(value: Any) -> str:
def _serialize_galileo_output(value: object) -> str:
if value is None:
return ""
if isinstance(value, str):
return value
def _json_default(obj: Any) -> Any:
def _json_default(obj: Any) -> object:
if hasattr(obj, "model_dump"):
return obj.model_dump()
return str(obj)
@ -439,8 +451,8 @@ class GalileoObserve(CustomLogger):
return json.dumps(value, default=_json_default)
@staticmethod
def _prompt_to_input_text(prompt: dict[str, Any]) -> str:
messages: Final = prompt.get("messages")
def _prompt_to_input_text(prompt: Mapping[str, Any]) -> str:
messages: Final[object] = prompt.get("messages")
if messages is not None:
text: Final = GalileoObserve._input_text_from_messages(messages)
if text:
@ -448,7 +460,7 @@ class GalileoObserve(CustomLogger):
return json.dumps(prompt, default=str)
@staticmethod
def _get_chat_content_for_galileo(response_obj: litellm.ModelResponse) -> Any:
def _get_chat_content_for_galileo(response_obj: litellm.ModelResponse) -> object:
if response_obj.choices and len(response_obj.choices) > 0:
message: Final = response_obj["choices"][0]["message"]
if hasattr(message, "json"):
@ -470,23 +482,23 @@ class GalileoObserve(CustomLogger):
@staticmethod
def _get_responses_api_content_for_galileo(
response_obj: ResponsesAPIResponse,
) -> Any:
) -> object:
if hasattr(response_obj, "output") and response_obj.output:
return response_obj.output
return None
@staticmethod
def _langfuse_style_rerank_prompt(kwargs: dict[str, Any]) -> dict[str, Any]:
def _langfuse_style_rerank_prompt(kwargs: Mapping[str, object]) -> dict[str, Any]:
"""Match Langfuse rerank input: prompt = {"messages": kwargs.get("messages")}."""
return {"messages": kwargs.get("messages")}
def _get_galileo_input_output_content(
self,
kwargs: dict[str, Any],
response_obj: Any,
kwargs: Mapping[str, object],
response_obj: object,
level: str = "DEFAULT",
status_message: str | None = None,
) -> tuple[str, str, Any]:
) -> tuple[str, str, object]:
"""
Mirror Langfuse _get_langfuse_input_output_content for Galileo ingest.
@ -582,12 +594,12 @@ class GalileoObserve(CustomLogger):
return self._prompt_to_input_text(prompt), "", kwargs.get("messages") or []
def get_output_str_from_response(self, response_obj: Any, kwargs: dict[str, Any]) -> str:
def get_output_str_from_response(self, response_obj: object, kwargs: Mapping[str, object]) -> str:
_, output_text, _ = self._get_galileo_input_output_content(kwargs=kwargs, response_obj=response_obj)
return output_text
@staticmethod
def _input_text_from_messages(messages: Any) -> str:
def _input_text_from_messages(messages: object) -> str:
"""Return a plain-string summary of the input suitable for the trace-level input field."""
if isinstance(messages, str):
return messages
@ -613,7 +625,13 @@ class GalileoObserve(CustomLogger):
return str(content)
return ""
async def async_log_success_event(self, kwargs: Any, response_obj: Any, start_time: Any, end_time: Any):
async def async_log_success_event(
self,
kwargs: Mapping[str, object],
response_obj: object,
start_time: object,
end_time: object,
) -> None:
verbose_logger.debug("On Async Success")
try:
await self._async_log_success_event_impl(
@ -625,7 +643,13 @@ class GalileoObserve(CustomLogger):
except Exception:
verbose_logger.exception("Galileo Logger: unexpected error in async_log_success_event")
async def _async_log_success_event_impl(self, kwargs: Any, response_obj: Any, start_time: Any, end_time: Any):
async def _async_log_success_event_impl(
self,
kwargs: Mapping[str, Any],
response_obj: object,
start_time: object,
end_time: object,
) -> None:
if not self._is_configured():
verbose_logger.debug(
"Galileo Logger: skipping — GALILEO_PROJECT_ID=%s GALILEO_API_KEY=%s GALILEO_BASE_URL=%s",
@ -635,7 +659,7 @@ class GalileoObserve(CustomLogger):
)
return
slo: Final[dict[str, Any] | None] = kwargs.get("standard_logging_object")
slo: Final[GalileoStandardLoggingFields | None] = kwargs.get("standard_logging_object")
if slo is None:
verbose_logger.debug("Galileo Logger: no standard_logging_object in kwargs, skipping")
return
@ -646,8 +670,8 @@ class GalileoObserve(CustomLogger):
kwargs=kwargs, response_obj=response_obj
)
raw_start: Final = slo.get("startTime")
raw_end: Final = slo.get("endTime")
raw_start: Final[float | None] = slo.get("startTime")
raw_end: Final[float | None] = slo.get("endTime")
if raw_start is None or raw_end is None:
verbose_logger.debug(
"Galileo Logger: standard_logging_object missing startTime/endTime, "
@ -710,7 +734,7 @@ class GalileoObserve(CustomLogger):
if len(self.in_memory_records) >= self.batch_size:
await self.flush_in_memory_records()
async def flush_in_memory_records(self):
async def flush_in_memory_records(self) -> None:
if not self.in_memory_records:
return
@ -774,5 +798,11 @@ class GalileoObserve(CustomLogger):
if not self.use_v2_api and response.status_code in (401, 403):
self.headers = None
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
async def async_log_failure_event(
self,
kwargs: Mapping[str, object],
response_obj: object,
start_time: object,
end_time: object,
) -> None:
verbose_logger.debug("On Async Failure")

7
litellm/litellm_core_utils/get_litellm_params.py
View file

@ -115,8 +115,11 @@ def get_litellm_params(
litellm_request_debug: bool | None = None,
**kwargs,
) -> dict:
_litellm_metadata_dict: Final = litellm_metadata if isinstance(litellm_metadata, dict) else None
resolved_metadata: Final = _litellm_metadata_dict.copy() if not metadata and _litellm_metadata_dict else metadata
# Derive litellm_session_id / litellm_trace_id from metadata when not provided (call chaining)
_meta: Final = metadata or {}
_meta: Final = resolved_metadata or {}
if litellm_session_id is None:
litellm_session_id = _meta.get("session_id") or _meta.get("trace_id")
if litellm_trace_id is None:
@ -139,7 +142,7 @@ def get_litellm_params(
"model_alias_map": model_alias_map,
"completion_call_id": completion_call_id,
"aembedding": aembedding,
"metadata": metadata,
"metadata": resolved_metadata,
"model_info": model_info,
"proxy_server_request": proxy_server_request,
"preset_cache_key": preset_cache_key,

22
litellm/litellm_core_utils/litellm_logging.py
View file

@ -585,8 +585,8 @@ class Logging(LiteLLMLoggingBaseClass):
"""
base_litellm_params: Final[dict[str, Any]] = {}
if "metadata" in kwargs:
base_litellm_params["metadata"] = kwargs["metadata"]
if isinstance(kwargs.get("metadata"), dict):
base_litellm_params["metadata"] = kwargs["metadata"].copy()
if "litellm_metadata" in kwargs and isinstance(kwargs["litellm_metadata"], dict):
base_litellm_params["litellm_metadata"] = kwargs["litellm_metadata"]
if "metadata" not in base_litellm_params:
@ -1399,10 +1399,7 @@ class Logging(LiteLLMLoggingBaseClass):
litellm_params=(self.litellm_params if hasattr(self, "litellm_params") else None)
)
prompt = "" # use for tts cost calc
_input: Final = self.model_call_details.get("input", None)
if _input is not None and isinstance(_input, str):
prompt = _input
prompt = self._prompt_for_cost_calculation()
if cache_hit is None:
cache_hit = self.model_call_details.get("cache_hit", False)
@ -1461,6 +1458,19 @@ class Logging(LiteLLMLoggingBaseClass):
return None
def _prompt_for_cost_calculation(self) -> str:
"""
The raw input string is only priced directly for text-to-speech, which bills per character.
Every other call type gets its billable units from the response usage object, and call types
that carry no usage at all (file content retrieval, and anything else `function_setup` cannot
build messages for) only have the ``"default-message-value"`` placeholder here, so passing the
input along would token-price that placeholder.
"""
if self.call_type not in (CallTypes.speech.value, CallTypes.aspeech.value):
return ""
_input = self.model_call_details.get("input", None)
return _input if isinstance(_input, str) else ""
def _generate_content_result_as_model_response(self, result: object) -> ModelResponse | None:
"""
Native Google :generateContent bodies report token usage under

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

@ -681,6 +681,23 @@ def _get_regional_uplift_multiplier(model_info: ModelInfo, data_residency: str |
return 1.0
def _resolve_reasoning_token_cost(
model_info: ModelInfo,
service_tier: str | None,
completion_base_cost: float,
) -> float:
tier_reasoning_key: Final = _get_service_tier_cost_key("output_cost_per_reasoning_token", service_tier)
if model_info.get(tier_reasoning_key) is not None:
tier_reasoning_cost: Final = _get_cost_per_unit(model_info, tier_reasoning_key, None)
if tier_reasoning_cost is not None:
return tier_reasoning_cost
tier_output_key: Final = _get_service_tier_cost_key("output_cost_per_token", service_tier)
if tier_output_key != "output_cost_per_token" and model_info.get(tier_output_key) is not None:
return completion_base_cost
standard_reasoning_cost: Final = _get_cost_per_unit(model_info, "output_cost_per_reasoning_token", None)
return standard_reasoning_cost if standard_reasoning_cost is not None else completion_base_cost
def generic_cost_per_token(
model: str,
usage: Usage,
@ -817,9 +834,10 @@ def generic_cost_per_token(
## REASONING COST
if not is_text_tokens_total and reasoning_tokens and reasoning_tokens > 0:
_output_cost_per_reasoning_token = _get_cost_per_unit(model_info, "output_cost_per_reasoning_token", None)
_output_cost_per_reasoning_token = (
_output_cost_per_reasoning_token if _output_cost_per_reasoning_token is not None else completion_base_cost
_output_cost_per_reasoning_token = _resolve_reasoning_token_cost(
model_info=model_info,
service_tier=service_tier,
completion_base_cost=completion_base_cost,
)
completion_cost += float(reasoning_tokens) * _output_cost_per_reasoning_token

71
litellm/llms/anthropic/chat/guardrail_translation/handler.py
View file

@ -26,10 +26,13 @@ from litellm.llms.anthropic.experimental_pass_through.adapters.transformation im
)
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
from litellm.llms.base_llm.guardrail_translation.utils import (
anthropic_tool_name,
effective_scan_only_tool_results_for_guardrail,
effective_skip_system_message_for_guardrail,
effective_skip_tool_message_for_guardrail,
openai_messages_without_system,
openai_messages_without_tool,
merge_guardrailed_scoped_messages,
merge_returned_tools_into_request_tools,
scoped_structured_message_indices,
)
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import (
AnthropicPassthroughLoggingHandler,
@ -326,19 +329,25 @@ class AnthropicMessagesHandler(BaseTranslation):
skip_system: Final = effective_skip_system_message_for_guardrail(guardrail_to_apply)
skip_tool: Final = effective_skip_tool_message_for_guardrail(guardrail_to_apply)
scan_only_tool_results: Final = effective_scan_only_tool_results_for_guardrail(guardrail_to_apply)
chat_completion_compatible_request: Final = self._translate_to_openai(data)
structured_messages = cast(
full_structured_messages: Final = cast(
list[AllMessageValues],
chat_completion_compatible_request.get("messages", []),
)
if skip_system:
structured_messages = openai_messages_without_system(structured_messages)
if skip_tool:
structured_messages = openai_messages_without_tool(structured_messages)
scoped_message_indices: Final = scoped_structured_message_indices(
full_structured_messages,
scan_only_tool_results=scan_only_tool_results,
skip_system=skip_system,
skip_tool=skip_tool,
)
structured_messages: Final = [full_structured_messages[index] for index in scoped_message_indices]
tools_to_check: Final[list[ChatCompletionToolParam]] = chat_completion_compatible_request.get("tools", [])
tools_to_check: Final[list[ChatCompletionToolParam]] = (
[] if scan_only_tool_results else chat_completion_compatible_request.get("tools", [])
)
# Step 1: Extract all text content and images
extracted: Final = tuple(
@ -347,6 +356,7 @@ class AnthropicMessagesHandler(BaseTranslation):
msg_idx=msg_idx,
skip_system_message=skip_system,
skip_tool_message=skip_tool,
scan_only_tool_results=scan_only_tool_results,
)
for msg_idx, message in enumerate(messages)
)
@ -388,14 +398,31 @@ class AnthropicMessagesHandler(BaseTranslation):
if converted_tool is not None:
anthropic_tools.append(converted_tool)
# Note: MCP servers are handled separately in the main transformation
data["tools"] = anthropic_tools
data["tools"] = (
merge_returned_tools_into_request_tools(
request_tools=data.get("tools"),
returned_tools=anthropic_tools,
tool_name=anthropic_tool_name,
)
if scan_only_tool_results
else anthropic_tools
)
guardrailed_structured_messages: Final = guardrailed_inputs.get("structured_messages")
if (
guardrailed_structured_messages is not None
and guardrailed_structured_messages is not original_structured_messages
):
self._write_back_structured_messages(data, guardrailed_structured_messages)
self._write_back_structured_messages(
data,
guardrailed_structured_messages
if guardrail_to_apply.structured_messages_cover_full_request()
else merge_guardrailed_scoped_messages(
full_messages=full_structured_messages,
scoped_indices=scoped_message_indices,
guardrailed_scoped=guardrailed_structured_messages,
),
)
else:
# Step 3: Map guardrail responses back to original message structure
await self._apply_guardrail_responses_to_input(
@ -461,6 +488,7 @@ class AnthropicMessagesHandler(BaseTranslation):
msg_idx: int,
skip_system_message: bool = False,
skip_tool_message: bool = False,
scan_only_tool_results: bool = False,
) -> ExtractedInput:
"""
Extract text content and images from a message.
@ -471,6 +499,8 @@ class AnthropicMessagesHandler(BaseTranslation):
content: Final = message.get("content", None)
if isinstance(content, str):
if scan_only_tool_results:
return EMPTY_EXTRACTED_INPUT
return ExtractedInput(scanned=(ScannedText(content, MessageContentTarget(msg_idx)),), images=())
if not isinstance(content, list):
return EMPTY_EXTRACTED_INPUT
@ -481,6 +511,7 @@ class AnthropicMessagesHandler(BaseTranslation):
msg_idx=msg_idx,
content_idx=content_idx,
skip_tool_message=skip_tool_message,
scan_only_tool_results=scan_only_tool_results,
)
for content_idx, content_item in enumerate(content)
if isinstance(content_item, dict)
@ -497,12 +528,16 @@ class AnthropicMessagesHandler(BaseTranslation):
msg_idx: int,
content_idx: int,
skip_tool_message: bool,
scan_only_tool_results: bool = False,
) -> ExtractedInput:
if content_item.get("type") == "tool_result":
if skip_tool_message:
return EMPTY_EXTRACTED_INPUT
return cls._extract_tool_result(content_item=content_item, msg_idx=msg_idx, content_idx=content_idx)
if scan_only_tool_results:
return EMPTY_EXTRACTED_INPUT
text_str: Final = content_item.get("text", None)
return ExtractedInput(
scanned=(
@ -551,22 +586,6 @@ class AnthropicMessagesHandler(BaseTranslation):
data: Final = source.get("data")
return (data,) if data else ()
def _extract_input_tools(
self,
tools: list[dict[str, Any]],
tools_to_check: list[ChatCompletionToolParam],
) -> None:
"""
Extract tools from a message.
"""
## CHECK FOR TOOLS
if tools is not None and isinstance(tools, list):
# TRANSFORM ANTHROPIC TOOLS TO OPENAI TOOLS
openai_tools: Final = self.adapter.translate_anthropic_tools_to_openai(
tools=cast(list[AllAnthropicToolsValues], tools)
)
tools_to_check.extend(openai_tools)
async def _apply_guardrail_responses_to_input(
self,
messages: list[dict[str, Any]],

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

@ -592,7 +592,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
# Anthropic requires additionalProperties=false for object schemas
# See: https://docs.anthropic.com/en/docs/build-with-claude/structured-outputs
if result.get("type") == "object" and "additionalProperties" not in result:
if result.get("type") == "object":
result["additionalProperties"] = False
return result

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

@ -1,7 +1,8 @@
from __future__ import annotations
import json
from typing import Any, Final
from collections.abc import Callable, Iterator, Sequence
from typing import Any, Final, TypeVar
from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage
from litellm.types.llms.openai import AllMessageValues
@ -113,13 +114,131 @@ def effective_skip_tool_message_for_guardrail(guardrail_to_apply: Any) -> bool:
return bool(getattr(litellm, "skip_tool_message_in_guardrail", False))
def _message_role(message: AllMessageValues) -> str:
return str((message or {}).get("role") or "").lower()
def openai_messages_without_system(
messages: list[AllMessageValues],
) -> list[AllMessageValues]:
return [m for m in messages if str((m or {}).get("role") or "").lower() != "system"]
messages: Sequence[AllMessageValues],
) -> tuple[AllMessageValues, ...]:
return tuple(m for m in messages if _message_role(m) != "system")
def openai_messages_without_tool(
messages: list[AllMessageValues],
messages: Sequence[AllMessageValues],
) -> tuple[AllMessageValues, ...]:
return tuple(m for m in messages if _message_role(m) != "tool")
def effective_scan_only_tool_results_for_guardrail(guardrail_to_apply: object) -> bool:
return getattr(guardrail_to_apply, "scan_only_tool_results", None) is True
def role_out_of_guardrail_scope(
role: str,
*,
skip_system_message: bool,
skip_tool_message: bool,
scan_only_tool_results: bool = False,
) -> bool:
if skip_system_message and role == "system":
return True
if skip_tool_message and role == "tool":
return True
return scan_only_tool_results and role not in ("tool", "function")
def scoped_structured_message_indices(
messages: Sequence[AllMessageValues],
*,
scan_only_tool_results: bool,
skip_system: bool,
skip_tool: bool,
) -> tuple[int, ...]:
return tuple(
index
for index, message in enumerate(messages)
if not role_out_of_guardrail_scope(
_message_role(message),
skip_system_message=skip_system,
skip_tool_message=skip_tool,
scan_only_tool_results=scan_only_tool_results,
)
)
ToolT = TypeVar("ToolT")
def openai_tool_name(tool: object) -> str | None:
if not isinstance(tool, dict):
return None
function: Final = tool.get("function")
if isinstance(function, dict):
function_name: Final = function.get("name")
return function_name if isinstance(function_name, str) else None
flat_name: Final = tool.get("name")
return flat_name if isinstance(flat_name, str) else None
def anthropic_tool_name(tool: object) -> str | None:
name: Final = tool.get("name") if isinstance(tool, dict) else None
return name if isinstance(name, str) else None
def merge_returned_tools_into_request_tools(
request_tools: Sequence[ToolT] | None,
returned_tools: Sequence[ToolT],
tool_name: Callable[[ToolT], str | None],
) -> list[ToolT]:
"""Union of the request's tools and guardrail-returned tools, keyed by name.
Under ``scan_only_tool_results`` the guardrail never saw the request's
tools, so a returned list can neither replace them (it would drop every
user-defined function) nor be discarded (it may carry a tool the guardrail
synthesized and told the model to call, like Compresr's retrieve tool).
Keep every request tool and append only returned tools whose names aren't
already taken by a request tool or an earlier returned tool.
"""
originals: Final = tuple(request_tools or ())
taken_names: Final = frozenset(name for tool in originals if (name := tool_name(tool)) is not None)
additions: Final = tuple(
tool
for index, tool in enumerate(returned_tools)
if (name := tool_name(tool)) not in taken_names
and (name is None or all(tool_name(earlier) != name for earlier in returned_tools[:index]))
)
return [*originals, *additions]
def merge_guardrailed_scoped_messages(
full_messages: Sequence[AllMessageValues],
scoped_indices: Sequence[int],
guardrailed_scoped: Sequence[AllMessageValues],
) -> list[AllMessageValues]:
return [m for m in messages if str((m or {}).get("role") or "").lower() != "tool"]
"""Substitute guardrail-returned messages back into the full conversation.
Guardrails only ever see the scoped subset of messages, so a replacement
list they hand back describes that subset, not the whole request. Writing
it over ``data["messages"]`` wholesale would silently drop every
out-of-scope message (system prompt, prior turns). Instead, swap each
returned message into the position its scoped original came from; extra
returned messages land after the last scoped position, and scoped
originals without a counterpart are treated as removed by the guardrail.
When nothing was filtered out this degenerates to the returned list
itself, preserving wholesale-replacement behavior for unscoped guardrails.
"""
replacements: Final = dict(zip(scoped_indices, guardrailed_scoped))
removed: Final = frozenset(scoped_indices[len(guardrailed_scoped) :])
appended: Final = tuple(guardrailed_scoped[len(scoped_indices) :])
last_scoped_index: Final = scoped_indices[-1] if scoped_indices else None
def _merged() -> Iterator[AllMessageValues]:
for index, message in enumerate(full_messages):
if index in removed:
continue
yield replacements.get(index, message)
if index == last_scoped_index:
yield from appended
return list(_merged())

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

@ -23,10 +23,14 @@ from litellm.llms.base_llm.guardrail_translation.base_translation import (
StreamTransformSink,
)
from litellm.llms.base_llm.guardrail_translation.utils import (
effective_scan_only_tool_results_for_guardrail,
effective_skip_system_message_for_guardrail,
effective_skip_tool_message_for_guardrail,
openai_messages_without_system,
openai_messages_without_tool,
merge_guardrailed_scoped_messages,
merge_returned_tools_into_request_tools,
openai_tool_name,
role_out_of_guardrail_scope,
scoped_structured_message_indices,
)
from litellm.main import stream_chunk_builder
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
@ -82,6 +86,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
skip_system: Final = effective_skip_system_message_for_guardrail(guardrail_to_apply)
skip_tool: Final = effective_skip_tool_message_for_guardrail(guardrail_to_apply)
scan_only_tool_results: Final = effective_scan_only_tool_results_for_guardrail(guardrail_to_apply)
texts_to_check: Final[list[str]] = []
images_to_check: Final[list[str]] = []
@ -101,6 +106,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
tool_call_task_mappings=tool_call_task_mappings,
skip_system_message=skip_system,
skip_tool_message=skip_tool,
scan_only_tool_results=scan_only_tool_results,
)
# Step 2: Apply guardrail to all texts and tool calls in batch
@ -110,16 +116,18 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
inputs["images"] = images_to_check
if tool_calls_to_check:
inputs["tool_calls"] = tool_calls_to_check
structured_messages = self.get_structured_messages(data)
structured_messages: Final = self.get_structured_messages(data)
scoped_message_indices: Final = scoped_structured_message_indices(
structured_messages or [],
scan_only_tool_results=scan_only_tool_results,
skip_system=skip_system,
skip_tool=skip_tool,
)
if structured_messages:
if skip_system:
structured_messages = openai_messages_without_system(structured_messages)
if skip_tool:
structured_messages = openai_messages_without_tool(structured_messages)
inputs["structured_messages"] = structured_messages
inputs["structured_messages"] = [structured_messages[index] for index in scoped_message_indices]
# Pass tools (function definitions) to the guardrail
tools: Final = data.get("tools")
if tools:
if tools and not scan_only_tool_results:
inputs["tools"] = tools
# Include model information if available
model: Final = data.get("model")
@ -138,14 +146,30 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
guardrailed_tool_calls: Final = guardrailed_inputs.get("tool_calls", [])
guardrailed_tools: Final = guardrailed_inputs.get("tools")
if guardrailed_tools is not None:
data["tools"] = guardrailed_tools
data["tools"] = (
merge_returned_tools_into_request_tools(
request_tools=tools,
returned_tools=guardrailed_tools,
tool_name=openai_tool_name,
)
if scan_only_tool_results
else guardrailed_tools
)
guardrailed_structured_messages: Final = guardrailed_inputs.get("structured_messages")
if (
guardrailed_structured_messages is not None
and guardrailed_structured_messages is not original_structured_messages
):
data["messages"] = guardrailed_structured_messages
data["messages"] = (
guardrailed_structured_messages
if guardrail_to_apply.structured_messages_cover_full_request()
else merge_guardrailed_scoped_messages(
full_messages=structured_messages or [],
scoped_indices=scoped_message_indices,
guardrailed_scoped=guardrailed_structured_messages,
)
)
else:
# Step 3: Map guardrail responses back to original message structure
if guardrailed_texts and texts_to_check:
@ -194,16 +218,19 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
tool_call_task_mappings: list[tuple[int, int]],
skip_system_message: bool = False,
skip_tool_message: bool = False,
scan_only_tool_results: bool = False,
) -> None:
"""
Extract text content, images, and tool calls from a message.
Override this method to customize text/image/tool call extraction logic.
"""
role: Final = str(message.get("role") or "").lower()
if skip_system_message and role == "system":
return
if skip_tool_message and role == "tool":
if role_out_of_guardrail_scope(
str(message.get("role") or "").lower(),
skip_system_message=skip_system_message,
skip_tool_message=skip_tool_message,
scan_only_tool_results=scan_only_tool_results,
):
return
content: Final = message.get("content", None)

31
litellm/model_prices_and_context_window_backup.json
View file

@ -22176,7 +22176,9 @@
},
"gpt-4.1-2025-04-14": {
"cache_read_input_token_cost": 5e-07,
"cache_read_input_token_cost_priority": 8.75e-07,
"input_cost_per_token": 2e-06,
"input_cost_per_token_priority": 3.5e-06,
"input_cost_per_token_batches": 1e-06,
"litellm_provider": "openai",
"max_input_tokens": 1047576,
@ -22184,6 +22186,7 @@
"max_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 8e-06,
"output_cost_per_token_priority": 1.4e-05,
"output_cost_per_token_batches": 4e-06,
"supported_endpoints": [
"/v1/chat/completions",
@ -22247,7 +22250,9 @@
},
"gpt-4.1-mini-2025-04-14": {
"cache_read_input_token_cost": 1e-07,
"cache_read_input_token_cost_priority": 1.75e-07,
"input_cost_per_token": 4e-07,
"input_cost_per_token_priority": 7e-07,
"input_cost_per_token_batches": 2e-07,
"litellm_provider": "openai",
"max_input_tokens": 1047576,
@ -22255,6 +22260,7 @@
"max_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 1.6e-06,
"output_cost_per_token_priority": 2.8e-06,
"output_cost_per_token_batches": 8e-07,
"supported_endpoints": [
"/v1/chat/completions",
@ -22317,7 +22323,9 @@
},
"gpt-4.1-nano-2025-04-14": {
"cache_read_input_token_cost": 2.5e-08,
"cache_read_input_token_cost_priority": 5e-08,
"input_cost_per_token": 1e-07,
"input_cost_per_token_priority": 2e-07,
"input_cost_per_token_batches": 5e-08,
"litellm_provider": "openai",
"max_input_tokens": 1047576,
@ -22325,6 +22333,7 @@
"max_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 4e-07,
"output_cost_per_token_priority": 8e-07,
"output_cost_per_token_batches": 2e-07,
"supported_endpoints": [
"/v1/chat/completions",
@ -22393,7 +22402,9 @@
},
"gpt-4o-2024-08-06": {
"cache_read_input_token_cost": 1.25e-06,
"cache_read_input_token_cost_priority": 2.125e-06,
"input_cost_per_token": 2.5e-06,
"input_cost_per_token_priority": 4.25e-06,
"input_cost_per_token_batches": 1.25e-06,
"litellm_provider": "openai",
"max_input_tokens": 128000,
@ -22401,6 +22412,7 @@
"max_tokens": 16384,
"mode": "chat",
"output_cost_per_token": 1e-05,
"output_cost_per_token_priority": 1.7e-05,
"output_cost_per_token_batches": 5e-06,
"supports_function_calling": true,
"supports_parallel_function_calling": true,
@ -22413,7 +22425,9 @@
},
"gpt-4o-2024-11-20": {
"cache_read_input_token_cost": 1.25e-06,
"cache_read_input_token_cost_priority": 2.125e-06,
"input_cost_per_token": 2.5e-06,
"input_cost_per_token_priority": 4.25e-06,
"input_cost_per_token_batches": 1.25e-06,
"litellm_provider": "openai",
"max_input_tokens": 128000,
@ -22421,6 +22435,7 @@
"max_tokens": 16384,
"mode": "chat",
"output_cost_per_token": 1e-05,
"output_cost_per_token_priority": 1.7e-05,
"output_cost_per_token_batches": 5e-06,
"supports_function_calling": true,
"supports_parallel_function_calling": true,
@ -22720,7 +22735,9 @@
},
"gpt-4o-mini-2024-07-18": {
"cache_read_input_token_cost": 7.5e-08,
"cache_read_input_token_cost_priority": 1.25e-07,
"input_cost_per_token": 1.5e-07,
"input_cost_per_token_priority": 2.5e-07,
"input_cost_per_token_batches": 7.5e-08,
"litellm_provider": "openai",
"max_input_tokens": 128000,
@ -22728,6 +22745,7 @@
"max_tokens": 16384,
"mode": "chat",
"output_cost_per_token": 6e-07,
"output_cost_per_token_priority": 1e-06,
"output_cost_per_token_batches": 3e-07,
"search_context_cost_per_query": {
"search_context_size_high": 0.03,
@ -25077,6 +25095,7 @@
"cache_read_input_token_cost": 5e-09,
"cache_read_input_token_cost_flex": 2.5e-09,
"input_cost_per_token": 5e-08,
"input_cost_per_token_priority": 2.5e-06,
"input_cost_per_token_flex": 2.5e-08,
"litellm_provider": "openai",
"max_input_tokens": 272000,
@ -29304,13 +29323,19 @@
},
"o3-2025-04-16": {
"cache_read_input_token_cost": 5e-07,
"cache_read_input_token_cost_flex": 2.5e-07,
"cache_read_input_token_cost_priority": 8.75e-07,
"input_cost_per_token": 2e-06,
"input_cost_per_token_flex": 1e-06,
"input_cost_per_token_priority": 3.5e-06,
"litellm_provider": "openai",
"max_input_tokens": 200000,
"max_output_tokens": 100000,
"max_tokens": 100000,
"mode": "chat",
"output_cost_per_token": 8e-06,
"output_cost_per_token_flex": 4e-06,
"output_cost_per_token_priority": 1.4e-05,
"supported_endpoints": [
"/v1/responses",
"/v1/chat/completions",
@ -29525,13 +29550,19 @@
},
"o4-mini-2025-04-16": {
"cache_read_input_token_cost": 2.75e-07,
"cache_read_input_token_cost_flex": 1.375e-07,
"cache_read_input_token_cost_priority": 5e-07,
"input_cost_per_token": 1.1e-06,
"input_cost_per_token_flex": 5.5e-07,
"input_cost_per_token_priority": 2e-06,
"litellm_provider": "openai",
"max_input_tokens": 200000,
"max_output_tokens": 100000,
"max_tokens": 100000,
"mode": "chat",
"output_cost_per_token": 4.4e-06,
"output_cost_per_token_flex": 2.2e-06,
"output_cost_per_token_priority": 8e-06,
"supports_function_calling": true,
"supports_parallel_function_calling": false,
"supports_pdf_input": true,

90
litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py
View file

@ -7,10 +7,13 @@ import contextvars
import json
import os
import re
from collections.abc import Mapping, Sequence
from pathlib import PurePosixPath
from typing import Any, Final
from typing import Any, Final, TypeAlias, TypedDict
from urllib.parse import quote
import httpx
# Tool names emitted from OpenAPI specs must work across all major LLM providers.
# OpenAI/Anthropic/Bedrock all enforce a character class roughly equivalent to
# ^[a-zA-Z0-9_-]+$ on tool names. Many specs (notably GitHub's REST API) use
@ -44,6 +47,41 @@ from litellm.proxy._experimental.mcp_server.tool_registry import (
global_mcp_tool_registry,
)
_OpenAPIParameter: TypeAlias = Mapping[str, Any]
class _OpenAPIJSONSchema(TypedDict, total=False):
properties: Mapping[str, object]
class _OpenAPIMediaType(TypedDict, total=False):
schema: _OpenAPIJSONSchema
class _OpenAPIRequestBody(TypedDict, total=False):
description: str
required: bool
content: Mapping[str, _OpenAPIMediaType]
class _OpenAPIOperation(TypedDict, total=False):
operationId: str
summary: str
description: str
parameters: Sequence[_OpenAPIParameter]
requestBody: _OpenAPIRequestBody
class _OpenAPIPathItem(TypedDict, total=False):
summary: str
description: str
parameters: Sequence[_OpenAPIParameter]
class _OpenAPIComponents(TypedDict, total=False):
parameters: Mapping[str, _OpenAPIParameter]
# Store the base URL and headers globally
BASE_URL: Final = ""
HEADERS: Final[dict[str, str]] = {}
@ -69,7 +107,7 @@ _request_resolved_auth_headers: Final[contextvars.ContextVar[dict[str, str] | No
)
def _sanitize_path_parameter_value(param_value: Any, param_name: str) -> str:
def _sanitize_path_parameter_value(param_value: object, param_name: str) -> str:
"""Ensure path params cannot introduce directory traversal."""
if param_value is None:
return ""
@ -109,7 +147,7 @@ def load_openapi_spec(filepath: str) -> dict[str, Any]:
async def load_openapi_spec_async(filepath: str) -> dict[str, Any]:
if filepath.startswith("http://") or filepath.startswith("https://"):
client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP)
r: Final = await async_safe_get(client, filepath)
r: Final[httpx.Response] = await async_safe_get(client, filepath)
r.raise_for_status()
return r.json()
@ -121,11 +159,11 @@ async def load_openapi_spec_async(filepath: str) -> dict[str, Any]:
return json.load(f)
def get_base_url(spec: dict[str, Any], spec_path: str | None = None) -> str:
def get_base_url(spec: Mapping[str, Any], spec_path: str | None = None) -> str:
"""Extract base URL from OpenAPI spec."""
# OpenAPI 3.x
if "servers" in spec and spec["servers"]:
server_url: Final = spec["servers"][0]["url"]
server_url: Final[str] = spec["servers"][0]["url"]
# If the server URL is relative (starts with /), derive base from spec_path
if server_url.startswith("/") and spec_path:
@ -147,8 +185,8 @@ def get_base_url(spec: dict[str, Any], spec_path: str | None = None) -> str:
return server_url
# OpenAPI 2.x (Swagger)
elif "host" in spec:
scheme: Final = spec.get("schemes", ["https"])[0]
base_path: Final = spec.get("basePath", "")
scheme: Final[str] = spec.get("schemes", ["https"])[0]
base_path: Final[str] = spec.get("basePath", "")
return f"{scheme}://{spec['host']}{base_path}"
# Fallback: derive base URL from spec_path if it's a URL
@ -172,20 +210,24 @@ def get_base_url(spec: dict[str, Any], spec_path: str | None = None) -> str:
return ""
def _resolve_ref(param: dict[str, Any], component_params: dict[str, Any]) -> dict[str, Any] | None:
def _resolve_ref(
param: _OpenAPIParameter, component_params: Mapping[str, _OpenAPIParameter]
) -> _OpenAPIParameter | None:
"""Resolve a single parameter, following a $ref if present.
Returns the resolved param dict, or None if the $ref target is absent from
components (so callers can skip/filter it rather than propagating a stub
with name=None that would corrupt deduplication).
"""
ref: Final = param.get("$ref", "")
ref: Final[str] = param.get("$ref", "")
if not ref.startswith("#/components/parameters/"):
return param
return component_params.get(ref.split("/")[-1])
def _resolve_param_list(raw: list[dict[str, Any]], component_params: dict[str, Any]) -> list[dict[str, Any]]:
def _resolve_param_list(
raw: Sequence[_OpenAPIParameter], component_params: Mapping[str, _OpenAPIParameter]
) -> list[_OpenAPIParameter]:
"""Resolve $refs in a parameter list, dropping any unresolvable entries."""
result: Final = []
for p in raw:
@ -196,9 +238,9 @@ def _resolve_param_list(raw: list[dict[str, Any]], component_params: dict[str, A
def resolve_operation_params(
operation: dict[str, Any],
path_item: dict[str, Any],
components: dict[str, Any],
operation: _OpenAPIOperation,
path_item: _OpenAPIPathItem,
components: _OpenAPIComponents,
) -> dict[str, Any]:
"""Return a copy of *operation* with fully-resolved, merged parameters.
@ -214,7 +256,7 @@ def resolve_operation_params(
merged with the operation-level params; operation-level wins when the
same ``name`` + ``in`` combination appears in both.
"""
component_params: Final = components.get("parameters", {})
component_params: Final[Mapping[str, _OpenAPIParameter]] = components.get("parameters", {})
path_level: Final = _resolve_param_list(path_item.get("parameters", []), component_params)
op_level: Final = _resolve_param_list(operation.get("parameters", []), component_params)
op_keys: Final = {(p["name"], p.get("in")) for p in op_level}
@ -224,7 +266,7 @@ def resolve_operation_params(
return result
def extract_parameters(operation: dict[str, Any]) -> tuple:
def extract_parameters(operation: Mapping[str, Any]) -> tuple[Sequence[str], Sequence[str], Sequence[str]]:
"""Extract parameter names from OpenAPI operation."""
path_params: Final = []
query_params: Final = []
@ -250,7 +292,7 @@ def extract_parameters(operation: dict[str, Any]) -> tuple:
return path_params, query_params, body_params
def build_input_schema(operation: dict[str, Any]) -> dict[str, Any]:
def build_input_schema(operation: Mapping[str, Any]) -> dict[str, Any]:
"""Build MCP input schema from OpenAPI operation."""
properties: Final = {}
required: Final = []
@ -274,12 +316,12 @@ def build_input_schema(operation: dict[str, Any]) -> dict[str, Any]:
# Process requestBody (OpenAPI 3.x)
if "requestBody" in operation:
request_body: Final = operation["requestBody"]
content: Final = request_body.get("content", {})
request_body: Final[_OpenAPIRequestBody] = operation["requestBody"]
content: Final[Mapping[str, _OpenAPIMediaType]] = request_body.get("content", {})
# Try to get JSON schema
if "application/json" in content:
schema: Final = content["application/json"].get("schema", {})
schema: Final[_OpenAPIJSONSchema] = content["application/json"].get("schema", {})
properties["body"] = {
"type": "object",
"description": request_body.get("description", "Request body"),
@ -347,7 +389,7 @@ def _merge_openapi_tool_request_headers(
def create_tool_function(
path: str,
method: str,
operation: dict[str, Any],
operation: Mapping[str, Any],
base_url: str,
headers: dict[str, str] | None = None,
):
@ -373,7 +415,7 @@ def create_tool_function(
path_params, query_params, body_params = extract_parameters(operation)
original_method: Final = method.lower()
async def tool_function(**kwargs: Any) -> str:
async def tool_function(**kwargs: object) -> str:
"""
Dynamically generated tool function.
@ -448,10 +490,10 @@ def create_tool_function(
return tool_function
def register_tools_from_openapi(spec: dict[str, Any], base_url: str):
def register_tools_from_openapi(spec: Mapping[str, Any], base_url: str) -> None:
"""Register MCP tools from OpenAPI specification."""
paths: Final = spec.get("paths", {})
used_names: Final[set] = set()
paths: Final[Mapping[str, Mapping[str, Any]]] = spec.get("paths", {})
used_names: Final = set()
for path, path_item in paths.items():
for method in ["get", "post", "put", "delete", "patch"]:

1
litellm/proxy/_types.py
View file

@ -781,6 +781,7 @@ class LiteLLMRoutes(enum.Enum):
"/model/update",
"/model/delete",
"/user/daily/activity",
"/user/daily/activity/aggregated",
"/user/available_roles", # read-only role metadata; any authenticated user may read
"/user/list", # org admins checked in endpoint; non-admins get 403
"/model/{model_id}/update",

80
litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py
View file

@ -18,8 +18,9 @@ Endpoints:
import json
import re
from collections.abc import Mapping, Sequence
from datetime import datetime, timezone
from typing import Any, Final
from typing import Final, Protocol, TypedDict
from fastapi import APIRouter, Depends, HTTPException
from fastapi.responses import JSONResponse
@ -41,7 +42,30 @@ from litellm.types.proxy.claude_code_endpoints import (
router: Final = APIRouter()
async def _get_prisma_client():
class _PluginRecord(Protocol):
id: str
name: str
version: str | None
description: str | None
manifest_json: str | None
enabled: bool
created_at: datetime | None
updated_at: datetime | None
created_by: str | None
class _MarketplaceEntry(TypedDict, total=False):
name: str
source: object
version: str
description: str
author: object
homepage: object
keywords: object
category: object
async def _get_prisma_client() -> object:
"""Get the prisma client from proxy_server."""
from litellm.proxy.proxy_server import prisma_client
@ -77,12 +101,14 @@ async def get_marketplace():
try:
prisma_client: Final = await _get_prisma_client()
plugins: Final = await ClaudeCodePluginRepository(prisma_client).table.find_many(where={"enabled": True})
plugins: Final[Sequence[_PluginRecord]] = await ClaudeCodePluginRepository(prisma_client).table.find_many(
where={"enabled": True}
)
plugin_list: Final = []
for plugin in plugins:
try:
manifest = json.loads(plugin.manifest_json)
manifest: Mapping[str, object] = json.loads(plugin.manifest_json or "{}")
except json.JSONDecodeError:
verbose_proxy_logger.warning("Plugin %s has invalid manifest JSON, skipping", plugin.name)
continue
@ -92,7 +118,7 @@ async def get_marketplace():
verbose_proxy_logger.warning("Plugin %s has no source field, skipping", plugin.name)
continue
entry: dict[str, Any] = {
entry: _MarketplaceEntry = {
"name": plugin.name,
"source": manifest["source"],
}
@ -137,7 +163,7 @@ async def get_marketplace():
_VALID_GIT_SUBDIR_PATH_RE: Final = re.compile(r"^[a-zA-Z0-9][a-zA-Z0-9._-]*(/[a-zA-Z0-9][a-zA-Z0-9._-]*)*$")
def _validate_plugin_source(source: dict[str, Any]) -> None:
def _validate_plugin_source(source: Mapping[str, str]) -> None:
"""Validate plugin source format, raising HTTPException on invalid input."""
source_type: Final = source.get("source")
if source_type == "github":
@ -179,9 +205,9 @@ def _validate_plugin_source(source: dict[str, Any]) -> None:
)
def _build_plugin_manifest(name: str, spec: PluginSpec) -> dict[str, Any]:
def _build_plugin_manifest(name: str, spec: PluginSpec) -> Mapping[str, object]:
"""Build the stored manifest dict shared by plugin create and update."""
dumped = spec.model_dump(exclude_none=True)
dumped: Final[Mapping[str, object]] = spec.model_dump(exclude_none=True)
return {"name": name, **{key: value for key, value in dumped.items() if value and key != "name"}}
@ -255,14 +281,16 @@ async def register_plugin(
_validate_plugin_source(request.source)
existing = await ClaudeCodePluginRepository(prisma_client).table.find_unique(where={"name": request.name})
existing: Final[_PluginRecord | None] = await ClaudeCodePluginRepository(prisma_client).table.find_unique(
where={"name": request.name}
)
if existing:
raise _name_conflict_error(request.name)
manifest = _build_plugin_manifest(request.name, request)
manifest: Final[Mapping[str, object]] = _build_plugin_manifest(request.name, request)
try:
plugin = await ClaudeCodePluginRepository(prisma_client).table.create(
plugin: Final[_PluginRecord] = await ClaudeCodePluginRepository(prisma_client).table.create(
data={
"name": request.name,
"version": request.version,
@ -326,7 +354,9 @@ async def list_plugins(
prisma_client: Final = await _get_prisma_client()
where: Final = {"enabled": True} if enabled_only else {}
plugins: Final = await ClaudeCodePluginRepository(prisma_client).table.find_many(where=where)
plugins: Final[Sequence[_PluginRecord]] = await ClaudeCodePluginRepository(prisma_client).table.find_many(
where=where
)
plugin_list: Final = []
for p in plugins:
@ -391,7 +421,9 @@ async def get_plugin(
try:
prisma_client: Final = await _get_prisma_client()
plugin: Final = await ClaudeCodePluginRepository(prisma_client).table.find_unique(where={"name": plugin_name})
plugin: Final[_PluginRecord | None] = await ClaudeCodePluginRepository(prisma_client).table.find_unique(
where={"name": plugin_name}
)
if not plugin:
raise HTTPException(
@ -399,7 +431,7 @@ async def get_plugin(
detail={"error": f"Plugin '{plugin_name}' not found"},
)
manifest: Final = json.loads(plugin.manifest_json) if plugin.manifest_json else {}
manifest: Final[Mapping[str, object]] = json.loads(plugin.manifest_json or "{}") if plugin.manifest_json else {}
return {
"id": plugin.id,
@ -477,19 +509,19 @@ async def update_plugin(
from prisma.errors import PrismaError
try:
prisma_client = await _get_prisma_client()
prisma_client: Final = await _get_prisma_client()
_validate_plugin_source(request.source)
existing = await ClaudeCodePluginRepository(prisma_client).table.find_unique(
existing: Final[_PluginRecord | None] = await ClaudeCodePluginRepository(prisma_client).table.find_unique(
where={"name": plugin_name} # mutable-ok: prisma query arguments must be plain dicts
)
if not existing:
raise _error_response(404, f"Plugin '{plugin_name}' not found")
manifest = _build_plugin_manifest(plugin_name, request)
manifest: Final[Mapping[str, object]] = _build_plugin_manifest(plugin_name, request)
plugin = await ClaudeCodePluginRepository(prisma_client).table.update(
plugin: Final[_PluginRecord] = await ClaudeCodePluginRepository(prisma_client).table.update(
where={"name": plugin_name}, # mutable-ok: prisma query arguments must be plain dicts
data={ # mutable-ok: prisma query arguments must be plain dicts
"version": request.version,
@ -540,7 +572,9 @@ async def enable_plugin(
try:
prisma_client: Final = await _get_prisma_client()
plugin: Final = await ClaudeCodePluginRepository(prisma_client).table.find_unique(where={"name": plugin_name})
plugin: Final[_PluginRecord | None] = await ClaudeCodePluginRepository(prisma_client).table.find_unique(
where={"name": plugin_name}
)
if not plugin:
raise HTTPException(
status_code=404,
@ -583,7 +617,9 @@ async def disable_plugin(
try:
prisma_client: Final = await _get_prisma_client()
plugin: Final = await ClaudeCodePluginRepository(prisma_client).table.find_unique(where={"name": plugin_name})
plugin: Final[_PluginRecord | None] = await ClaudeCodePluginRepository(prisma_client).table.find_unique(
where={"name": plugin_name}
)
if not plugin:
raise HTTPException(
status_code=404,
@ -626,7 +662,9 @@ async def delete_plugin(
try:
prisma_client: Final = await _get_prisma_client()
plugin: Final = await ClaudeCodePluginRepository(prisma_client).table.find_unique(where={"name": plugin_name})
plugin: Final[_PluginRecord | None] = await ClaudeCodePluginRepository(prisma_client).table.find_unique(
where={"name": plugin_name}
)
if not plugin:
raise HTTPException(
status_code=404,

2
litellm/proxy/batches_endpoints/endpoints.py
View file

@ -264,6 +264,7 @@ async def create_batch(
detail={"error": "LLM Router not initialized. Ensure models added to proxy."},
)
_create_batch_data.update(disable_fallbacks=True) # pyright: ignore[reportCallIssue] # router flag
response = await llm_router.acreate_batch(**_create_batch_data)
response.input_file_id = input_file_id
response._hidden_params["unified_file_id"] = unified_file_id
@ -961,6 +962,7 @@ async def cancel_batch(
prisma_client=prisma_client,
verbose_proxy_logger=verbose_proxy_logger,
operation="cancel",
user_api_key_dict=user_api_key_dict,
)
### CALL HOOKS ### - modify outgoing data

33
litellm/proxy/common_utils/custom_openapi_spec.py
View file

@ -1,3 +1,4 @@
from collections.abc import Mapping, Sequence
from typing import Any, Final
from litellm._logging import verbose_proxy_logger
@ -26,7 +27,7 @@ class CustomOpenAPISpec:
RESPONSES_API_PATHS = ["/v1/responses", "/responses"]
@staticmethod
def get_pydantic_schema(model_class) -> dict[str, Any] | None:
def get_pydantic_schema(model_class) -> Mapping[str, object] | None:
"""
Get JSON schema from a Pydantic model, handling both v1 and v2 APIs.
@ -53,7 +54,9 @@ class CustomOpenAPISpec:
return None
@staticmethod
def add_schema_to_components(openapi_schema: dict[str, Any], schema_name: str, schema_def: dict[str, Any]) -> None:
def add_schema_to_components(
openapi_schema: dict[str, Any], schema_name: str, schema_def: Mapping[str, object]
) -> None:
"""
Add a schema definition to the OpenAPI components/schemas section.
@ -72,7 +75,7 @@ class CustomOpenAPISpec:
CustomOpenAPISpec._move_defs_to_components(openapi_schema, {schema_name: schema_def})
@staticmethod
def add_request_body_to_paths(openapi_schema: dict[str, Any], paths: list[str], schema_ref: str) -> None:
def add_request_body_to_paths(openapi_schema: dict[str, Any], paths: Sequence[str], schema_ref: str) -> None:
"""
Add request body with expanded form fields for better Swagger UI display.
This keeps the request body but expands it to show individual fields in the UI.
@ -130,7 +133,7 @@ class CustomOpenAPISpec:
openapi_schema["paths"][path]["post"]["parameters"] = filtered_params
@staticmethod
def _move_defs_to_components(openapi_schema: dict[str, Any], defs: dict[str, Any]) -> None:
def _move_defs_to_components(openapi_schema: dict[str, Any], defs: Mapping[str, Mapping[str, Any]]) -> None:
"""
Move $defs from Pydantic v2 schema to OpenAPI components/schemas.
This makes the definitions resolvable in Swagger/OpenAPI viewers.
@ -218,7 +221,7 @@ class CustomOpenAPISpec:
return {"type": "string"}
@staticmethod
def _expand_field_definition(field_def: dict[str, Any]) -> dict[str, Any]:
def _expand_field_definition(field_def: dict[str, object]) -> dict[str, object]:
"""
Expand a Pydantic field definition for inline use in OpenAPI schema.
This creates a full field definition that Swagger UI can render as individual form fields.
@ -234,12 +237,12 @@ class CustomOpenAPISpec:
@staticmethod
def add_request_schema(
openapi_schema: dict[str, Any],
openapi_schema: dict[str, object],
model_class: type,
schema_name: str,
paths: list[str],
paths: Sequence[str],
operation_name: str,
) -> dict[str, Any]:
) -> dict[str, object]:
"""
Generic method to add a request schema to OpenAPI specification.
@ -279,8 +282,8 @@ class CustomOpenAPISpec:
@staticmethod
def add_chat_completion_request_schema(
openapi_schema: dict[str, Any],
) -> dict[str, Any]:
openapi_schema: dict[str, object],
) -> dict[str, object]:
"""
Add ProxyChatCompletionRequest schema to chat completion endpoints for documentation.
This shows the request body in Swagger without runtime validation.
@ -306,7 +309,7 @@ class CustomOpenAPISpec:
return openapi_schema
@staticmethod
def add_embedding_request_schema(openapi_schema: dict[str, Any]) -> dict[str, Any]:
def add_embedding_request_schema(openapi_schema: dict[str, object]) -> dict[str, object]:
"""
Add EmbeddingRequest schema to embedding endpoints for documentation.
This shows the request body in Swagger without runtime validation.
@ -333,8 +336,8 @@ class CustomOpenAPISpec:
@staticmethod
def add_responses_api_request_schema(
openapi_schema: dict[str, Any],
) -> dict[str, Any]:
openapi_schema: dict[str, object],
) -> dict[str, object]:
"""
Add ResponsesAPIRequestParams schema to responses API endpoints for documentation.
This shows the request body in Swagger without runtime validation.
@ -361,8 +364,8 @@ class CustomOpenAPISpec:
@staticmethod
def add_llm_api_request_schema_body(
openapi_schema: dict[str, Any],
) -> dict[str, Any]:
openapi_schema: dict[str, object],
) -> dict[str, object]:
"""
Add LLM API request schema bodies to OpenAPI specification for documentation.

74
litellm/proxy/container_endpoints/ownership.py
View file

@ -1,5 +1,7 @@
import json
from typing import Any, Final
from collections.abc import Mapping, Sequence
from collections.abc import Set as AbstractSet
from typing import TYPE_CHECKING, Any, Final, Protocol
from fastapi import HTTPException
@ -15,6 +17,29 @@ from litellm.proxy.common_utils.resource_ownership import (
from litellm.repositories.table_repositories import ManagedObjectRepository
from litellm.responses.utils import ResponsesAPIRequestUtils
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
class _ManagedObjectRow(Protocol):
model_object_id: str
unified_object_id: str | None
file_purpose: str | None
created_by: str | None
class _ManagedObjectTable(Protocol):
async def find_unique(self, *, where: Mapping[str, str]) -> _ManagedObjectRow | None: ...
async def find_first(self, *, where: Mapping[str, str]) -> _ManagedObjectRow | None: ...
async def find_many(self, *, where: Mapping[str, object]) -> Sequence[_ManagedObjectRow]: ...
async def create(self, *, data: Mapping[str, str]) -> _ManagedObjectRow: ...
async def update(self, *, where: Mapping[str, str], data: Mapping[str, str]) -> _ManagedObjectRow | None: ...
CONTAINER_OBJECT_PURPOSE: Final = "container"
# 60s LRU/TTL cache absorbs every container access check before it reaches
@ -39,7 +64,7 @@ _CONTAINER_STORED_ID_CACHE: Final = InMemoryCache(max_size_in_memory=10000, defa
_ALLOWED_CONTAINER_IDS_CACHE: Final = InMemoryCache(max_size_in_memory=2048, default_ttl=60)
def _allowed_container_ids_cache_key(owner_scopes: list[str]) -> str:
def _allowed_container_ids_cache_key(owner_scopes: Sequence[str]) -> str:
"""JSON-encode the sorted scope list — using a separator like ``|``
would collide for any tenant whose user_id / team_id / org_id /
api_key happens to contain the separator. JSON quoting escapes
@ -86,7 +111,7 @@ async def get_container_forwarding_params(
return params
def _get_response_id(response: Any) -> str | None:
def _get_response_id(response: object) -> str | None:
if response is None:
return None
if isinstance(response, dict):
@ -96,7 +121,7 @@ def _get_response_id(response: Any) -> str | None:
return value if isinstance(value, str) else None
def _dump_response(response: Any) -> dict[str, Any]:
def _dump_response(response: Any) -> dict[str, object]:
if isinstance(response, dict):
return dict(response)
if hasattr(response, "model_dump"):
@ -106,17 +131,17 @@ def _dump_response(response: Any) -> dict[str, Any]:
return {"id": _get_response_id(response)}
async def _get_prisma_client():
async def _get_prisma_client() -> "PrismaClient | None":
from litellm.proxy.proxy_server import prisma_client
return prisma_client
def _custom_llm_provider_from_responses_response(
response: Any,
response: object,
default: str = "openai",
) -> str:
hidden_params: dict[str, Any] = {}
hidden_params: Mapping[str, object] = {}
if isinstance(response, dict):
hidden_params = response.get("_hidden_params") or {}
else:
@ -129,7 +154,7 @@ def _custom_llm_provider_from_responses_response(
async def record_container_owners_from_responses_response(
response: Any,
response: object,
user_api_key_dict: UserAPIKeyAuth,
custom_llm_provider: str | None = None,
) -> None:
@ -160,10 +185,10 @@ async def record_container_owners_from_responses_response(
async def record_container_owner(
response: Any,
response: object,
user_api_key_dict: UserAPIKeyAuth,
custom_llm_provider: str,
) -> Any:
) -> object:
container_id: Final = _get_response_id(response)
if container_id is None:
verbose_proxy_logger.warning("Skipping container ownership tracking because provider response has no id")
@ -195,7 +220,7 @@ async def record_container_owner(
verbose_proxy_logger.warning("Skipping container ownership tracking because prisma_client is None")
return response
table: Final = ManagedObjectRepository(prisma_client).table
table: Final[_ManagedObjectTable] = ManagedObjectRepository(prisma_client).table
existing: Final = await table.find_unique(where={"model_object_id": model_object_id})
if existing is not None:
if getattr(existing, "file_purpose", None) != CONTAINER_OBJECT_PURPOSE:
@ -247,15 +272,16 @@ async def _get_container_owner(original_container_id: str, custom_llm_provider:
if prisma_client is None:
return None
row: Final = await ManagedObjectRepository(prisma_client).table.find_first(
table: Final[_ManagedObjectTable] = ManagedObjectRepository(prisma_client).table
row: Final[_ManagedObjectRow | None] = await table.find_first(
where={
"model_object_id": model_object_id,
"file_purpose": CONTAINER_OBJECT_PURPOSE,
}
)
owner: Final = getattr(row, "created_by", None) if row is not None else None
owner: Final[str | None] = getattr(row, "created_by", None) if row is not None else None
_CONTAINER_OWNER_CACHE.set_cache(model_object_id, owner if owner is not None else _NEGATIVE_OWNER_SENTINEL)
stored_id: Final = getattr(row, "unified_object_id", None) if row is not None else None
stored_id: Final[str | None] = getattr(row, "unified_object_id", None) if row is not None else None
_CONTAINER_STORED_ID_CACHE.set_cache(
model_object_id,
(stored_id if isinstance(stored_id, str) and stored_id else _NEGATIVE_STORED_ID_SENTINEL),
@ -283,13 +309,14 @@ async def _get_stored_container_id(original_container_id: str, custom_llm_provid
if prisma_client is None:
return None
row: Final = await ManagedObjectRepository(prisma_client).table.find_first(
table: Final[_ManagedObjectTable] = ManagedObjectRepository(prisma_client).table
row: Final[_ManagedObjectRow | None] = await table.find_first(
where={
"model_object_id": model_object_id,
"file_purpose": CONTAINER_OBJECT_PURPOSE,
}
)
stored_id: Final = getattr(row, "unified_object_id", None) if row is not None else None
stored_id: Final[str | None] = getattr(row, "unified_object_id", None) if row is not None else None
_CONTAINER_STORED_ID_CACHE.set_cache(
model_object_id,
(stored_id if isinstance(stored_id, str) and stored_id else _NEGATIVE_STORED_ID_SENTINEL),
@ -317,7 +344,7 @@ async def assert_user_can_access_container(
return original_container_id, resolved_provider
def _get_container_list_data(response: Any) -> list[Any] | None:
def _get_container_list_data(response: object) -> Sequence[object] | None:
if response is None:
return None
if isinstance(response, dict):
@ -327,7 +354,7 @@ def _get_container_list_data(response: Any) -> list[Any] | None:
return data if isinstance(data, list) else None
def _set_container_list_data(response: Any, data: list[Any], removed_filtered_items: bool = False) -> Any:
def _set_container_list_data(response: Any, data: list[object], removed_filtered_items: bool = False) -> object:
if isinstance(response, dict):
response["data"] = data
if data:
@ -353,7 +380,7 @@ def _set_container_list_data(response: Any, data: list[Any], removed_filtered_it
async def _get_allowed_container_ids(
user_api_key_dict: UserAPIKeyAuth,
) -> set[str]:
) -> AbstractSet[str]:
owner_scopes: Final = get_resource_owner_scopes(user_api_key_dict)
if not owner_scopes:
return set()
@ -367,7 +394,8 @@ async def _get_allowed_container_ids(
if prisma_client is None:
return set()
rows: Final = await ManagedObjectRepository(prisma_client).table.find_many(
table: Final[_ManagedObjectTable] = ManagedObjectRepository(prisma_client).table
rows: Final[Sequence[_ManagedObjectRow]] = await table.find_many(
where={
"file_purpose": CONTAINER_OBJECT_PURPOSE,
"created_by": {"in": owner_scopes},
@ -382,10 +410,10 @@ async def _get_allowed_container_ids(
async def filter_container_list_response(
response: Any,
response: object,
user_api_key_dict: UserAPIKeyAuth,
custom_llm_provider: str,
) -> Any:
) -> object:
if is_proxy_admin(user_api_key_dict):
return response
@ -394,7 +422,7 @@ async def filter_container_list_response(
return response
allowed_container_ids: Final = await _get_allowed_container_ids(user_api_key_dict)
filtered: Final[list[Any]] = []
filtered: Final[list[object]] = []
for item in data:
container_id = _get_response_id(item)
if container_id is None:

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

@ -26,6 +26,9 @@ from litellm.caching import DualCache
from litellm.exceptions import ModifyResponseException
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.litellm_core_utils.core_helpers import redact_nested_match_and_regex_keys
from litellm.llms.base_llm.guardrail_translation.utils import (
effective_scan_only_tool_results_for_guardrail,
)
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
@ -402,6 +405,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
grounding.append(block)
return grounding
def supports_scan_only_tool_results(self) -> bool:
return self.experimental_use_latest_role_message_only is not True
def _prepare_guardrail_messages_for_role(
self,
messages: list[AllMessageValues] | None,
@ -523,6 +529,11 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
latest_user_index: Final = self._find_latest_message_index(structured_messages, target_role="user")
if latest_user_index is None:
if effective_scan_only_tool_results_for_guardrail(self):
verbose_proxy_logger.warning(
"Bedrock Guardrail: experimental_use_latest_role_message_only scans only the latest "
"user message, so scan_only_tool_results leaves nothing to scan for this request"
)
verbose_proxy_logger.debug("Bedrock Guardrail: no user-role message in request, skipping INPUT scan")
return ApplyGuardrailMessageSelection(None, None, True, skip_scan=True)

94
litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py
View file

@ -9,11 +9,13 @@ import contextlib
import json
import os
import ssl
from collections.abc import AsyncGenerator
from collections.abc import AsyncGenerator, AsyncIterable, Mapping, Sequence
from ssl import SSLContext
from typing import TYPE_CHECKING, Any, Final
from fastapi import HTTPException
from pydantic import BaseModel
from typing_extensions import NotRequired, TypedDict
from websockets.asyncio.client import ClientConnection, connect
from websockets.exceptions import ConnectionClosed
@ -35,8 +37,8 @@ from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.utils import (
CallTypesLiteral,
Choices,
EmbeddingResponse,
ImageResponse,
LLMResponseTypes,
Message,
ModelResponse,
ModelResponseStream,
ResponsesAPIResponse,
@ -50,6 +52,44 @@ class CatoNetworksGuardrailMissingSecrets(Exception):
pass
class _WsSslKwargs(TypedDict, total=False):
ssl: bool | str | SSLContext
class _CatoRequiredAction(TypedDict, total=False):
action_type: str
detection_message: str
class _CatoRedactedMessage(TypedDict):
role: NotRequired[str]
content: str | None
class _CatoRedactedChat(TypedDict, total=False):
all_redacted_messages: Sequence[_CatoRedactedMessage]
class _CatoAnalysisResult(TypedDict, total=False):
policy_drill_down: Mapping[str, object]
class _CatoAnalyzeResponse(TypedDict):
required_action: NotRequired[_CatoRequiredAction | None]
analysis_result: NotRequired[_CatoAnalysisResult]
redacted_chat: NotRequired[_CatoRedactedChat]
class _CatoOutputRedaction(TypedDict):
redacted_output: str
class _CatoStreamMessage(TypedDict, total=False):
verified_chunk: Mapping[str, object]
done: bool
blocking_message: str
class CatoNetworksGuardrail(CustomGuardrail):
@classmethod
def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]:
@ -80,7 +120,7 @@ class CatoNetworksGuardrail(CustomGuardrail):
super().__init__(**kwargs)
@staticmethod
def _build_ws_ssl_kwargs(ssl_verify: bool | str | None, ws_api_base: str) -> dict:
def _build_ws_ssl_kwargs(ssl_verify: bool | str | None, ws_api_base: str) -> _WsSslKwargs:
"""Resolve the ``ssl`` argument for ``websockets.connect``. Mirrors the
``ssl_verify`` handling applied to the HTTP handler so a custom Cato instance
behind TLS honours the same verification settings for streaming."""
@ -156,7 +196,7 @@ class CatoNetworksGuardrail(CustomGuardrail):
return flattened
@staticmethod
def _prompt_inspection_messages(prompt: Any) -> list:
def _prompt_inspection_messages(prompt: object) -> Sequence[Mapping[str, str]]:
"""Synthetic user messages for a legacy completion ``prompt`` (a string
or a list of string prompts)."""
if isinstance(prompt, str):
@ -166,7 +206,7 @@ class CatoNetworksGuardrail(CustomGuardrail):
return []
@staticmethod
def _iter_schema_string_refs(data: dict):
def _iter_schema_string_refs(data: Mapping[str, Any]):
"""Yield ``(container, key)`` for every non-empty schema string the proxy
forwards to the model inside tool/function and structured-output schemas:
each ``tools[].function`` and legacy ``functions[]`` entry plus the
@ -208,7 +248,7 @@ class CatoNetworksGuardrail(CustomGuardrail):
stack.extend(reversed(node))
@classmethod
def _extra_inspection_sources(cls, data: dict) -> list:
def _extra_inspection_sources(cls, data: Mapping[str, Any]) -> Sequence[tuple[str, Sequence[Mapping[str, str]]]]:
"""Text the proxy forwards to the model outside chat ``messages``:
Responses-API ``input`` and ``instructions``, legacy completion
``prompt`` and tool/function/``response_format`` schema strings. Returned
@ -251,7 +291,7 @@ class CatoNetworksGuardrail(CustomGuardrail):
json={"messages": self._inspection_messages(data)},
)
response.raise_for_status()
res: Final = response.json()
res: Final[_CatoAnalyzeResponse] = response.json()
required_action: Final = res.get("required_action")
action_type: Final = required_action and required_action.get("action_type", None)
if action_type is None:
@ -267,7 +307,11 @@ class CatoNetworksGuardrail(CustomGuardrail):
verbose_proxy_logger.error("Cato: %s action", action_type)
return data
def _handle_block_action(self, analysis_result: Any, required_action: Any) -> None:
def _handle_block_action(
self,
analysis_result: _CatoAnalysisResult,
required_action: Any,
) -> None:
detection_message: Final = required_action.get("detection_message", None)
verbose_proxy_logger.info(
"Cato: Violation detected enabled policies: {policies}".format(
@ -348,7 +392,7 @@ class CatoNetworksGuardrail(CustomGuardrail):
hook: str,
key_alias: str | None,
user_email: str | None = None,
) -> dict | None:
) -> _CatoOutputRedaction | None:
call_id: Final = request_data.get("litellm_call_id")
inspection_messages: Final = self._inspection_messages(request_data)
assistant_index: Final = len(inspection_messages)
@ -363,7 +407,7 @@ class CatoNetworksGuardrail(CustomGuardrail):
json={"messages": inspection_messages + [{"role": "assistant", "content": output}]},
)
response.raise_for_status()
res: Final = response.json()
res: Final[_CatoAnalyzeResponse] = response.json()
required_action: Final = res.get("required_action")
action_type: Final = required_action and required_action.get("action_type", None)
if action_type and action_type == "block_action":
@ -378,7 +422,11 @@ class CatoNetworksGuardrail(CustomGuardrail):
return {"redacted_output": redacted_output}
return None
def _handle_block_action_on_output(self, analysis_result: Any, required_action: Any) -> None:
def _handle_block_action_on_output(
self,
analysis_result: _CatoAnalysisResult,
required_action: Any,
) -> None:
detection_message: Final = required_action.get("detection_message", None)
verbose_proxy_logger.info(
"Cato: detected: {detected}, enabled policies: {policies}".format(
@ -422,7 +470,7 @@ class CatoNetworksGuardrail(CustomGuardrail):
)
@staticmethod
def _output_fragments(message: Any) -> list:
def _output_fragments(message: Message) -> Sequence[tuple[tuple[str, int | None], str]]:
"""Assistant text the proxy returns to the caller: ``content`` plus every
``tool_calls[].function.arguments`` string, each tagged with where a
redaction must be written back. ``content`` is only included when present
@ -439,7 +487,7 @@ class CatoNetworksGuardrail(CustomGuardrail):
return fragments
@staticmethod
def _apply_output_fragment(message: Any, target: tuple, redacted: str) -> None:
def _apply_output_fragment(message: Any, target: tuple[str, int | None], redacted: str) -> None:
kind, idx = target
if kind == "content":
message.content = redacted
@ -447,11 +495,11 @@ class CatoNetworksGuardrail(CustomGuardrail):
message.tool_calls[idx].function.arguments = redacted
@staticmethod
def _responses_output_field(item: Any, key: str) -> Any:
def _responses_output_field(item: object, key: str) -> str | Sequence[object] | None:
return item.get(key) if isinstance(item, dict) else getattr(item, key, None)
@classmethod
def _responses_output_fragments(cls, response: ResponsesAPIResponse) -> list:
def _responses_output_fragments(cls, response: ResponsesAPIResponse) -> Sequence[tuple[object, str, str]]:
"""Assistant text the Responses API returns to the caller: every
``output_text`` content block plus every function-call ``arguments``
string, each paired with the ``(container, key)`` a Cato redaction is
@ -474,7 +522,7 @@ class CatoNetworksGuardrail(CustomGuardrail):
return fragments
@staticmethod
def _apply_responses_output_fragment(container: Any, key: str, redacted: str) -> None:
def _apply_responses_output_fragment(container: object, key: str, redacted: str) -> None:
if isinstance(container, dict):
container[key] = redacted
else:
@ -505,8 +553,8 @@ class CatoNetworksGuardrail(CustomGuardrail):
self,
data: dict,
user_api_key_dict: UserAPIKeyAuth,
response: Any | ModelResponse | EmbeddingResponse | ImageResponse,
) -> Any:
response: LLMResponseTypes,
) -> LLMResponseTypes:
user_email: Final = self._resolve_cato_user_email(user_api_key_dict)
if isinstance(response, ModelResponse) and response.choices:
for choice in response.choices:
@ -526,7 +574,7 @@ class CatoNetworksGuardrail(CustomGuardrail):
async def async_post_call_streaming_iterator_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
response,
response: AsyncIterable[object],
request_data: dict,
) -> AsyncGenerator[ModelResponseStream, None]:
from litellm.proxy.proxy_server import StreamingCallbackError
@ -547,7 +595,7 @@ class CatoNetworksGuardrail(CustomGuardrail):
try:
while True:
raw_message = await self._await_cato_message(websocket, sender)
result = json.loads(raw_message)
result: _CatoStreamMessage = json.loads(raw_message)
if verified_chunk := result.get("verified_chunk"):
yield ModelResponseStream.model_validate(verified_chunk)
continue
@ -560,7 +608,7 @@ class CatoNetworksGuardrail(CustomGuardrail):
finally:
await self._cancel_background_task(sender)
async def _await_cato_message(self, websocket: ClientConnection, sender: asyncio.Task) -> Any:
async def _await_cato_message(self, websocket: ClientConnection, sender: asyncio.Task[None]) -> str | bytes:
"""Wait for the next Cato message, surfacing a dead forwarding task instead of blocking."""
from litellm.proxy.proxy_server import StreamingCallbackError
@ -578,7 +626,7 @@ class CatoNetworksGuardrail(CustomGuardrail):
async def forward_the_stream_to_cato(
self,
websocket: ClientConnection,
response_iter: AsyncGenerator[Any, None],
response_iter: AsyncIterable[object],
) -> None:
async for chunk in response_iter:
if isinstance(chunk, BaseModel):

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

@ -362,6 +362,10 @@ class CrowdStrikeAIDRHandler(CustomGuardrail):
tail: Final = guard_output.messages[-num_assistant_messages:] if num_assistant_messages > 0 else []
return [_extract_text_from_message(msg) for msg in tail]
@override
def structured_messages_cover_full_request(self) -> bool:
return effective_skip_system_message_for_guardrail(self) or effective_skip_tool_message_for_guardrail(self)
def _writeback_messages(
self,
structured_messages: list[AllMessageValues],

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

@ -1,7 +1,8 @@
from __future__ import annotations
import os
from typing import TYPE_CHECKING, Any, Final, Literal
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict
from urllib.parse import urlparse
from uuid import uuid4
@ -29,9 +30,31 @@ from litellm.types.proxy.guardrails.guardrail_hooks.hiddenlayer import (
from litellm.types.utils import GenericGuardrailAPIInputs
if TYPE_CHECKING:
from pydantic import BaseModel
from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
class _HiddenlayerEvaluation(TypedDict, total=False):
action: str
threat_level: str
class _HiddenlayerAnalysisEntry(TypedDict, total=False):
name: str
detected: bool
class _HiddenlayerModifiedSide(TypedDict):
messages: Any
class _HiddenlayerResponse(TypedDict, total=False):
evaluation: _HiddenlayerEvaluation
analysis: Sequence[_HiddenlayerAnalysisEntry]
modified_data: Mapping[str, _HiddenlayerModifiedSide]
def is_saas(host: str) -> bool:
"""Checks whether the connection is to the SaaS platform"""
@ -43,7 +66,7 @@ def is_saas(host: str) -> bool:
return False
def _get_jwt(auth_url, api_id, api_key):
def _get_jwt(auth_url, api_id, api_key) -> str:
token_url: Final = f"{auth_url}/oauth2/token?grant_type=client_credentials"
resp: Final = requests.post(token_url, auth=HTTPBasicAuth(api_id, api_key))
@ -139,7 +162,7 @@ class HiddenlayerGuardrail(CustomGuardrail):
if scan_params := inputs.get("structured_messages"):
last_msg: Final = scan_params[-1]
result = await self._call_hiddenlayer(
result: _HiddenlayerResponse = await self._call_hiddenlayer(
project_id,
hl_request_metadata,
{
@ -205,11 +228,11 @@ class HiddenlayerGuardrail(CustomGuardrail):
async def _call_hiddenlayer(
self,
project_id: str | None,
metadata: dict[str, str],
payload: dict[str, Any],
metadata: Mapping[str, str],
payload: Mapping[str, Sequence[Mapping[str, str]]],
input_type: Literal["request", "response"],
) -> dict[str, Any]:
data: Final[dict[str, Any]] = {"metadata": metadata}
) -> _HiddenlayerResponse:
data: Final[dict[str, object]] = {"metadata": metadata}
if input_type == "request":
data["input"] = payload
@ -235,7 +258,7 @@ class HiddenlayerGuardrail(CustomGuardrail):
headers=headers,
)
response.raise_for_status()
result = response.json()
result: _HiddenlayerResponse = response.json()
verbose_proxy_logger.debug("Hiddenlayer reponse: %s", result)
@ -265,7 +288,7 @@ class HiddenlayerGuardrail(CustomGuardrail):
return result
@staticmethod
def get_config_model() -> type[GuardrailConfigModel] | None:
def get_config_model() -> type[GuardrailConfigModel[BaseModel]] | None:
from litellm.types.proxy.guardrails.guardrail_hooks.hiddenlayer import (
HiddenlayerGuardrailConfigModel,
)
@ -343,7 +366,7 @@ class HiddenlayerGuardrailV2(CustomGuardrail):
if "hl-requester-id" not in hl_headers:
hl_headers["hl-requester-id"] = "LiteLLM"
payload: Any
payload: object
if input_type == "request":
payload = {
"messages": inputs.get("structured_messages"),
@ -461,7 +484,7 @@ class HiddenlayerGuardrailV2(CustomGuardrail):
return response
@staticmethod
def get_config_model() -> type[GuardrailConfigModel] | None:
def get_config_model() -> type[GuardrailConfigModel[BaseModel]] | None:
from litellm.types.proxy.guardrails.guardrail_hooks.hiddenlayer import (
HiddenlayerGuardrailConfigModel,
)

58
litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/base.py
View file

@ -2,8 +2,11 @@ import threading
import time
import uuid
from collections import OrderedDict
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final
from typing_extensions import NotRequired, TypedDict
from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.prompt_templates.common_utils import (
convert_content_list_to_str,
@ -15,6 +18,7 @@ from litellm.llms.custom_httpx.http_handler import (
)
if TYPE_CHECKING:
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.llms.openai import AllMessageValues
GRAPH_API_BASE: Final = "https://graph.microsoft.com/v1.0"
@ -25,6 +29,11 @@ GRAPH_SCOPE: Final = "https://graph.microsoft.com/.default"
SCOPE_CACHE_TTL_SECONDS: Final = 3600.0
class GraphTokenResponse(TypedDict):
access_token: str
expires_in: NotRequired[int]
class PurviewGuardrailBase:
"""
Base class for Microsoft Purview guardrails.
@ -41,8 +50,8 @@ class PurviewGuardrailBase:
client_secret: str,
purview_app_name: str = "LiteLLM",
user_id_field: str = "user_id",
**kwargs: Any,
):
**kwargs: object,
) -> None:
# Forward remaining kwargs to the next class in the MRO
# (typically CustomGuardrail).
super().__init__(**kwargs)
@ -59,7 +68,7 @@ class PurviewGuardrailBase:
# Protection scope cache: user_id -> (etag, scope_response, fetched_at)
# Capped at 1000 entries (LRU eviction) to avoid unbounded growth.
self._scope_cache: OrderedDict[str, tuple[str, dict[str, Any], float]] = OrderedDict()
self._scope_cache: OrderedDict[str, tuple[str, Mapping[str, object], float]] = OrderedDict()
self._scope_cache_maxsize = 1000
# Use a threading.Lock (not asyncio.Lock) because this lock is acquired
# from both the proxy's main asyncio event loop and from short-lived
@ -100,7 +109,7 @@ class PurviewGuardrailBase:
headers={"Content-Type": "application/x-www-form-urlencoded"},
)
response.raise_for_status()
token_data: Final = response.json()
token_data: Final[GraphTokenResponse] = response.json()
access_token: Final = token_data["access_token"]
expires_in: Final = int(token_data.get("expires_in", 3599))
# Recompute ``now`` after the await so the expiry reflects when the
@ -117,9 +126,9 @@ class PurviewGuardrailBase:
async def _graph_post(
self,
url: str,
json_body: dict[str, Any],
extra_headers: dict[str, str] | None = None,
) -> tuple[dict[str, Any], dict[str, str]]:
json_body: dict[str, object],
extra_headers: Mapping[str, str] | None = None,
) -> tuple[dict[str, object], dict[str, str]]:
"""POST to Graph API with bearer auth.
Returns:
@ -136,7 +145,7 @@ class PurviewGuardrailBase:
verbose_proxy_logger.debug("Purview Graph POST %s", url)
response: Final = await self.async_handler.post(url=url, headers=headers, json=json_body)
response.raise_for_status()
response_json: Final[dict[str, Any]] = response.json()
response_json: Final[dict[str, object]] = response.json()
response_headers: Final = dict(response.headers)
verbose_proxy_logger.debug("Purview Graph response: %s", response_json)
return response_json, response_headers
@ -145,7 +154,7 @@ class PurviewGuardrailBase:
# Protection scopes
# ------------------------------------------------------------------
async def _compute_protection_scopes(self, user_id: str) -> tuple[str, dict[str, Any]]:
async def _compute_protection_scopes(self, user_id: str) -> tuple[str, Mapping[str, object]]:
"""Call protectionScopes/compute and cache with ETag.
Returns:
@ -161,7 +170,7 @@ class PurviewGuardrailBase:
return cached[0], cached[1]
url: Final = f"{GRAPH_API_BASE}/users/{encoded_user_id}/dataSecurityAndGovernance/protectionScopes/compute"
body: Final[dict[str, Any]] = {
body: Final[dict[str, object]] = {
"activities": "uploadText,downloadText",
"locations": [
{
@ -199,7 +208,7 @@ class PurviewGuardrailBase:
activity: str,
etag: str,
correlation_id: str | None = None,
) -> dict[str, Any]:
) -> dict[str, object]:
"""Call processContent for DLP policy evaluation.
Args:
@ -211,7 +220,7 @@ class PurviewGuardrailBase:
"""
encoded_user_id: Final = self._encode_graph_user_id(user_id)
url: Final = f"{GRAPH_API_BASE}/users/{encoded_user_id}/dataSecurityAndGovernance/processContent"
body: Final[dict[str, Any]] = {
body: Final[dict[str, object]] = {
"contentToProcess": {
"contentEntries": [
{
@ -261,7 +270,7 @@ class PurviewGuardrailBase:
# User ID resolution
# ------------------------------------------------------------------
def _resolve_user_id(self, data: dict[str, Any], user_api_key_dict: Any) -> str | None:
def _resolve_user_id(self, data: Mapping[str, object], user_api_key_dict: "UserAPIKeyAuth") -> str | None:
"""Resolve the Entra user object ID from request data or auth context.
Returns the strongest available identity walking down four sources, in
@ -284,7 +293,10 @@ class PurviewGuardrailBase:
if hasattr(user_api_key_dict, "end_user_id") and user_api_key_dict.end_user_id:
return str(user_api_key_dict.end_user_id)
metadata: Final = data.get("metadata") or data.get("litellm_metadata") or {}
metadata_value: Final[object] = data.get("metadata") or data.get("litellm_metadata") or {}
if not isinstance(metadata_value, Mapping):
return None
metadata: Final[Mapping[str, object]] = metadata_value
uid = metadata.get("user_api_key_user_id")
if uid:
return str(uid)
@ -296,15 +308,15 @@ class PurviewGuardrailBase:
return None
@staticmethod
def _logging_kwargs_metadata(kwargs: dict[str, Any]) -> dict[str, Any]:
def _logging_kwargs_metadata(kwargs: Mapping[str, object]) -> Mapping[str, object]:
"""Metadata dict from ``model_call_details`` / logging kwargs."""
litellm_params: Final = kwargs.get("litellm_params") or {}
litellm_params: Final[object] = kwargs.get("litellm_params") or {}
if not isinstance(litellm_params, dict):
return {}
md: Final = litellm_params.get("metadata")
return md if isinstance(md, dict) else {}
def _resolve_trusted_user_id(self, data: dict[str, Any], user_api_key_dict: Any) -> str | None:
def _resolve_trusted_user_id(self, data: Mapping[str, object], user_api_key_dict: "UserAPIKeyAuth") -> str | None:
"""Resolve user ID from API-key/JWT-bound identity for blocking DLP.
Uses only ``UserAPIKeyAuth.user_id`` (bound on the LiteLLM key or JWT).
@ -325,7 +337,7 @@ class PurviewGuardrailBase:
return None
def _resolve_user_id_from_logging_kwargs(self, kwargs: dict[str, Any]) -> str | None:
def _resolve_user_id_from_logging_kwargs(self, kwargs: Mapping[str, object]) -> str | None:
"""Trusted-identity-only resolver for logging-only hooks.
Uses only the proxy-injected ``user_api_key_user_id`` (populated from
@ -365,7 +377,7 @@ class PurviewGuardrailBase:
# ------------------------------------------------------------------
@staticmethod
def is_token_id_prompt(prompt: Any) -> bool:
def is_token_id_prompt(prompt: str | Sequence[object] | None) -> bool:
"""Return True if ``prompt`` carries OpenAI completions token ids.
Covers every list shape that ``completion_prompt_to_str`` cannot decode
@ -383,7 +395,7 @@ class PurviewGuardrailBase:
return False
@staticmethod
def completion_prompt_to_str(prompt: Any) -> str | None:
def completion_prompt_to_str(prompt: str | Sequence[object] | None) -> str | None:
"""Normalize OpenAI ``/v1/completions`` ``prompt`` for text DLP.
Supports string prompts and list-of-string prompts. List-of-token-id prompts
@ -408,7 +420,7 @@ class PurviewGuardrailBase:
return None
@staticmethod
def _extract_tool_call_args_from_message(message: Any) -> list[str]:
def _extract_tool_call_args_from_message(message: object) -> list[str]:
"""Return plaintext arguments strings from tool_calls and function_call fields.
Covers both the request path (assistant messages in chat histories that
@ -419,7 +431,9 @@ class PurviewGuardrailBase:
args: Final[list[str]] = []
# tool_calls: [{"function": {"arguments": "..."}}]
tool_calls = message.get("tool_calls") if isinstance(message, dict) else getattr(message, "tool_calls", None)
tool_calls: Final[Sequence[object] | None] = (
message.get("tool_calls") if isinstance(message, dict) else getattr(message, "tool_calls", None)
)
if tool_calls:
for tc in tool_calls:
fn = tc.get("function") if isinstance(tc, dict) else getattr(tc, "function", None)

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

@ -22,6 +22,9 @@ from litellm.integrations.custom_guardrail import (
CustomGuardrail,
log_guardrail_information,
)
from litellm.llms.base_llm.guardrail_translation.utils import (
effective_scan_only_tool_results_for_guardrail,
)
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
@ -1561,6 +1564,9 @@ class PanwPrismaAirsHandler(CustomGuardrail):
return scannable
def supports_scan_only_tool_results(self) -> bool:
return False
@staticmethod
def _get_scannable_text_indices(
texts: list[str],
@ -1716,6 +1722,15 @@ class PanwPrismaAirsHandler(CustomGuardrail):
# - latest-user extraction returned None (no user / count mismatch)
if scannable_indices is None:
scannable_indices = self._get_scannable_text_indices(texts, structured_messages)
if (
scannable_indices is not None
and not scannable_indices
and effective_scan_only_tool_results_for_guardrail(self)
):
verbose_proxy_logger.warning(
"PANW Prisma AIRS scans only user, system, and developer messages, "
"so scan_only_tool_results leaves nothing to scan for this request"
)
for i, text in enumerate(texts):
if not text or not text.strip():

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

@ -74,6 +74,9 @@ class PromptSecurityGuardrail(CustomGuardrail):
super().__init__(**kwargs)
def supports_scan_only_tool_results(self) -> bool:
return self.check_tool_results
@log_guardrail_information
async def apply_guardrail(
self,

1
litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/__init__.py
View file

@ -18,6 +18,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"
send_user_api_key_alias=litellm_params.send_user_api_key_alias,
send_user_api_key_user_id=litellm_params.send_user_api_key_user_id,
send_user_api_key_team_id=litellm_params.send_user_api_key_team_id,
timeout=litellm_params.timeout,
guardrail_name=guardrail.get("guardrail_name", ""),
event_hook=litellm_params.mode,
default_on=litellm_params.default_on,

30
litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py
View file

@ -22,9 +22,10 @@ from litellm.types.utils import GenericGuardrailAPIInputs
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.types.guardrails import LitellmParams
from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
GUARDRAIL_TIMEOUT: Final = 5
DEFAULT_GUARDRAIL_TIMEOUT: Final = 5.0
class ZscalerAIGuard(CustomGuardrail):
@ -43,6 +44,7 @@ class ZscalerAIGuard(CustomGuardrail):
send_user_api_key_alias: bool | None = None,
send_user_api_key_user_id: bool | None = None,
send_user_api_key_team_id: bool | None = None,
timeout: float | None = None,
**kwargs,
):
kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks()))
@ -68,6 +70,7 @@ class ZscalerAIGuard(CustomGuardrail):
if send_user_api_key_team_id is not None
else os.getenv("SEND_USER_API_KEY_TEAM_ID", "False").lower() in ("true", "1")
)
self.timeout = self._resolve_timeout(timeout)
verbose_proxy_logger.debug(
"send_user_api_key_alias: %s, \n send_user_api_key_user_id:%s, \n send_user_api_key_team_id:%s",
@ -80,6 +83,29 @@ class ZscalerAIGuard(CustomGuardrail):
verbose_proxy_logger.debug("ZscalerAIGuard Initializing ...")
@staticmethod
def _resolve_timeout(timeout: float | None) -> float:
"""
Resolve the effective per-request timeout, falling back to the default
when it is unset or non-positive.
"""
if timeout is None:
return DEFAULT_GUARDRAIL_TIMEOUT
if timeout <= 0:
verbose_proxy_logger.warning(
"Ignoring non-positive Zscaler AI Guard timeout %s, using %s seconds",
timeout,
DEFAULT_GUARDRAIL_TIMEOUT,
)
return DEFAULT_GUARDRAIL_TIMEOUT
return timeout
def update_in_memory_litellm_params(self, litellm_params: "LitellmParams") -> None:
super().update_in_memory_litellm_params(litellm_params)
self.timeout = self._resolve_timeout(litellm_params.timeout)
@staticmethod
def _resolve_metadata_value(request_data: dict | None, key: str) -> str | None:
"""
@ -267,7 +293,7 @@ class ZscalerAIGuard(CustomGuardrail):
f"{url}",
headers=headers,
json=data,
timeout=GUARDRAIL_TIMEOUT,
timeout=self.timeout,
)
response.raise_for_status()
return response

29
litellm/proxy/guardrails/guardrail_registry.py
View file

@ -14,6 +14,10 @@ from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.llms.base_llm.guardrail_translation.utils import (
effective_scan_only_tool_results_for_guardrail,
effective_skip_tool_message_for_guardrail,
)
from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import (
BedrockGuardrail,
)
@ -487,16 +491,27 @@ class InMemoryGuardrailHandler:
raise ValueError(f"Unsupported guardrail: {guardrail_type}")
if custom_guardrail_callback is not None:
setattr(
custom_guardrail_callback,
for scoping_param in (
"skip_system_message_in_guardrail",
getattr(litellm_params, "skip_system_message_in_guardrail", None),
)
setattr(
custom_guardrail_callback,
"skip_tool_message_in_guardrail",
getattr(litellm_params, "skip_tool_message_in_guardrail", None),
"scan_only_tool_results",
):
setattr(custom_guardrail_callback, scoping_param, getattr(litellm_params, scoping_param, None))
scan_only_tool_results_enabled: Final = effective_scan_only_tool_results_for_guardrail(
custom_guardrail_callback
)
if scan_only_tool_results_enabled and not custom_guardrail_callback.supports_scan_only_tool_results():
raise ValueError(
f"Guardrail {guardrail['guardrail_name']}: scan_only_tool_results is enabled, but this "
"guardrail's role filtering never scans tool results, so no request content would ever "
"be scanned. Remove scan_only_tool_results or the guardrail's role-filtering option."
)
if scan_only_tool_results_enabled and effective_skip_tool_message_for_guardrail(custom_guardrail_callback):
raise ValueError(
f"Guardrail {guardrail['guardrail_name']}: scan_only_tool_results and "
"skip_tool_message_in_guardrail are enabled together, which excludes every message from "
"scanning, so no request content would ever be scanned. Remove one of the two."
)
configured_run_in_parallel: Final = getattr(litellm_params, "run_in_parallel", None)
if configured_run_in_parallel is not None:
custom_guardrail_callback.run_in_parallel = bool(configured_run_in_parallel)

12
litellm/proxy/hooks/proxy_track_cost_callback.py
View file

@ -34,11 +34,17 @@ from litellm.types.utils import (
)
from litellm.utils import get_end_user_id_for_cost_tracking
_PASS_THROUGH_CALL_TYPES: Final[frozenset[str]] = frozenset(
_UNATTRIBUTED_TRACKABLE_CALL_TYPES: Final[frozenset[str]] = frozenset(
{
CallTypes.pass_through.value,
CallTypes.llm_passthrough_route.value,
CallTypes.allm_passthrough_route.value,
# CheckBatchCost's synthetic logging_obj for a completed managed batch only ever
# carries user_api_key_user_id (from LiteLLM_ManagedObjectTable.created_by) and
# user_api_key_team_id (from .team_id) -- both are None for batches created with
# the master key or a team-less key, since the table never stores the raw key
# hash. The batch already incurred real provider cost, so track it regardless.
CallTypes.aretrieve_batch.value,
}
)
@ -440,6 +446,8 @@ def _should_track_cost_callback(
the request with no key/user/team/end-user to attribute spend to. Those
requests still forward real provider traffic that operators expect to see
in request/usage logs, so they are tracked even when unauthenticated.
The same reasoning applies to a completed managed batch's cost event
(see _UNATTRIBUTED_TRACKABLE_CALL_TYPES).
"""
# don't run track cost callback if user opted into disabling spend
@ -448,7 +456,7 @@ def _should_track_cost_callback(
if user_api_key is not None or user_id is not None or team_id is not None or end_user_id is not None:
return True
return call_type in _PASS_THROUGH_CALL_TYPES
return call_type in _UNATTRIBUTED_TRACKABLE_CALL_TYPES
def _get_budget_reservation_from_metadata(metadata: dict) -> dict | None:

33
litellm/proxy/litellm_pre_call_utils.py
View file

@ -151,6 +151,20 @@ LITELLM_METADATA_ROUTES: Final = (
"files",
)
LITELLM_TRACE_CONTROL_METADATA_FIELDS: Final = frozenset(
{
"mask_input",
"mask_output",
"session_id",
"trace_id",
"trace_metadata",
"trace_name",
"trace_release",
"trace_user_id",
"trace_version",
}
)
_UNTRUSTED_ROOT_CONTROL_FIELDS: Final = (
"proxy_server_request",
"standard_logging_object",
@ -458,6 +472,18 @@ def _get_metadata_variable_name(request: Request) -> str:
return "metadata"
def _promoted_trace_control_fields(
requester_metadata: Mapping[str, Any],
litellm_metadata: Mapping[str, Any],
) -> tuple[tuple[str, Any], ...]:
"""Return the caller's trace-control fields that ``litellm_metadata`` does not already set."""
return tuple(
(key, value)
for key, value in requester_metadata.items()
if key in LITELLM_TRACE_CONTROL_METADATA_FIELDS and key not in litellm_metadata
)
def _extract_generic_session_id_from_headers(
normalized: dict[str, str],
) -> str | None:
@ -1670,6 +1696,13 @@ async def add_litellm_data_to_request(
# paths may read from it.
if "metadata" in data and isinstance(data["metadata"], dict):
data[_metadata_variable_name]["requester_metadata"] = copy.deepcopy(data["metadata"])
if _metadata_variable_name == "litellm_metadata":
data[_metadata_variable_name].update(
_promoted_trace_control_fields(
requester_metadata=data[_metadata_variable_name]["requester_metadata"],
litellm_metadata=data[_metadata_variable_name],
)
)
# Merge litellm_metadata into the metadata variable (preserving existing
# values). Runs after the user_api_key_* / _pipeline_managed_guardrails

55
litellm/proxy/management_endpoints/common_daily_activity.py
View file

@ -1,6 +1,6 @@
import asyncio
from collections.abc import Awaitable, Callable, Mapping, Sequence
from datetime import datetime
from datetime import datetime, timedelta, timezone
from types import SimpleNamespace
from typing import TYPE_CHECKING, Final, Protocol
@ -422,26 +422,46 @@ def _adjust_dates_for_timezone(
start_date: str,
end_date: str,
timezone_offset_minutes: int | None,
include_current_utc_day: bool = False,
utc_now: datetime | None = None,
) -> tuple[str, str]:
"""
Pass-through for the local date range; the timezone offset is intentionally ignored here.
Map a caller-local date range onto UTC bucket keys, extending only the live end.
The aggregation table (e.g. LiteLLM_DailyUserSpend) stores spend in whole-UTC-day
buckets keyed on date as YYYY-MM-DD. Any conversion from a local date range to a
UTC date range using only date arithmetic must round to whole UTC days, allowing up
to 24h of slop at each boundary. The previous implementation expanded the SQL range
by an extra full UTC day on whichever side the offset pointed, which pulled in 24h
of unrelated bucket data per boundary and produced approximately 100% over-counting
on single-day queries (e.g. IST May 29 returning UTC May 28 + UTC May 29 in full).
buckets keyed on date as YYYY-MM-DD. Any conversion of an interior local-day
boundary using only date arithmetic must round to whole UTC days, allowing up to
24h of slop at each boundary. A previous implementation expanded the SQL range by
an extra full UTC day on whichever side the offset pointed, which pulled in 24h of
unrelated bucket data per boundary and produced approximately 100% over-counting on
single-day queries (e.g. IST May 29 returning UTC May 28 + UTC May 29 in full).
Sums of single-day queries then exceeded the equivalent multi-day aggregate, which
is mathematically impossible.
is mathematically impossible. Historical dates therefore stay a pass-through: the
local date is the UTC bucket key, trading boundary slop for monotonic, additive
results. Hour-level buckets or pro-rata weighting would fix that properly; both
require data the current schema does not store.
Treating the local date as the UTC date trades a small one-time boundary slop for
correct, monotonic, additive results across single-day and multi-day queries. A
later fix can introduce hour-level buckets or pro-rata weighting on adjacent UTC
days; both require data the current schema does not store.
The end boundary is different when the range reaches the caller's current day. A
caller west of UTC asking for a range ending "today" is asking for data up to now,
but once UTC has rolled past their local midnight, everything they sent since then
sits in the next UTC bucket, which the pass-through excludes: a PT dashboard goes
stale every evening from 5pm until local midnight, showing $0 for anything that
only started accruing that evening. Extending such a range to today's UTC bucket
cannot over-count, because the only part of that bucket outside the caller's range
is the future, and the future is empty. ``timezone_offset_minutes`` follows the
JS ``Date.getTimezoneOffset`` convention: UTC minus local, positive west of UTC.
The extension is strictly opt-in via ``include_current_utc_day`` so a consumer
whose axis or reconciliation expects the range to stop at the requested end date
keeps today's byte-for-byte behaviour; the cost optimization dashboard opts in.
"""
return start_date, end_date
if not include_current_utc_day or timezone_offset_minutes is None:
return start_date, end_date
now: Final = utc_now if utc_now is not None else datetime.now(timezone.utc)
caller_local_today: Final = (now - timedelta(minutes=timezone_offset_minutes)).date().isoformat()
if end_date < caller_local_today:
return start_date, end_date
return start_date, max(end_date, now.date().isoformat())
def _build_where_conditions(
@ -454,10 +474,13 @@ def _build_where_conditions(
api_key: str | list[str] | None,
exclude_entity_ids: list[str] | None = None,
timezone_offset_minutes: int | None = None,
include_current_utc_day: bool = False,
) -> dict[str, "_WhereValue"]:
"""Build prisma where clause for daily activity queries."""
# Adjust dates for timezone if provided
adjusted_start, adjusted_end = _adjust_dates_for_timezone(start_date, end_date, timezone_offset_minutes)
adjusted_start, adjusted_end = _adjust_dates_for_timezone(
start_date, end_date, timezone_offset_minutes, include_current_utc_day
)
where_conditions: Final[dict[str, _WhereValue]] = {
"date": {
@ -903,6 +926,7 @@ async def get_daily_activity(
exclude_entity_ids: list[str] | None = None,
metadata_metrics_func: Callable[[Sequence[DailySpendRecord]], SpendMetrics] | None = None,
timezone_offset_minutes: int | None = None,
include_current_utc_day: bool = False,
resolve_entity_metadata: Callable[[Sequence[DailySpendRecord]], Awaitable[dict[str, dict[str, object]]]]
| None = None,
) -> SpendAnalyticsPaginatedResponse:
@ -936,6 +960,7 @@ async def get_daily_activity(
api_key=api_key,
exclude_entity_ids=exclude_entity_ids,
timezone_offset_minutes=timezone_offset_minutes,
include_current_utc_day=include_current_utc_day,
)
# Get total count for pagination

8
litellm/proxy/management_endpoints/internal_user_endpoints.py
View file

@ -2650,6 +2650,13 @@ async def get_user_daily_activity(
description="Timezone offset in minutes from UTC (e.g., 480 for PST). "
"Matches JavaScript's Date.getTimezoneOffset() convention.",
),
include_current_utc_day: bool = fastapi.Query(
default=False,
description="When the range ends on the caller's current local day, extend it to "
"today's UTC bucket so spend written after the caller's local midnight (in UTC "
"terms) is included. Requires the timezone parameter. Historical ranges are "
"never extended.",
),
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
) -> SpendAnalyticsPaginatedResponse:
"""
@ -2711,6 +2718,7 @@ async def get_user_daily_activity(
page=page,
page_size=page_size,
timezone_offset_minutes=timezone,
include_current_utc_day=include_current_utc_day,
resolve_entity_metadata=lambda records: _resolve_user_email_metadata(prisma_client, records),
)

98
litellm/proxy/management_endpoints/workflow_management_endpoints.py
View file

@ -14,7 +14,8 @@ GET /v1/workflows/runs/{run_id}/messages - Fetch conversation history
"""
import json
from typing import Any, Final, Literal
from collections.abc import Mapping, Sequence
from typing import Final, Literal, Protocol, TypedDict
from fastapi import APIRouter, Depends, HTTPException, Query
@ -43,7 +44,7 @@ router: Final = APIRouter()
_MAX_SEQUENCE_RETRIES: Final = 5
def _json(value: Any) -> str:
def _json(value: object) -> str:
"""Serialize a Python value for prisma-client-py Json fields (must be a string)."""
return json.dumps(value)
@ -62,7 +63,7 @@ def _caller_key(user_api_key_dict: UserAPIKeyAuth) -> str | None:
# Status transitions driven by event_type
_EVENT_STATUS_MAP: Final[dict[str, str]] = {
_EVENT_STATUS_MAP: Final[Mapping[str, str]] = {
"step.started": "running",
"step.failed": "failed",
"hook.waiting": "paused",
@ -77,8 +78,8 @@ _EVENT_STATUS_MAP: Final[dict[str, str]] = {
class WorkflowRunCreateRequest(BaseModel):
workflow_type: str
input: dict[str, Any] | None = None
metadata: dict[str, Any] | None = None
input: Mapping[str, object] | None = None
metadata: Mapping[str, object] | None = None
WorkflowRunStatus = Literal["pending", "running", "paused", "completed", "failed"]
@ -86,14 +87,14 @@ WorkflowRunStatus = Literal["pending", "running", "paused", "completed", "failed
class WorkflowRunUpdateRequest(BaseModel):
status: WorkflowRunStatus | None = None
output: dict[str, Any] | None = None
metadata: dict[str, Any] | None = None
output: Mapping[str, object] | None = None
metadata: Mapping[str, object] | None = None
class WorkflowEventCreateRequest(BaseModel):
event_type: str
step_name: str
data: dict[str, Any] | None = None
data: Mapping[str, object] | None = None
class WorkflowMessageCreateRequest(BaseModel):
@ -102,15 +103,60 @@ class WorkflowMessageCreateRequest(BaseModel):
session_id: str | None = None
class _RunRow(Protocol):
@property
def created_by(self) -> str | None: ...
class _SeqRow(Protocol):
@property
def sequence_number(self) -> int: ...
class _RunCreateData(TypedDict, total=False):
workflow_type: str
created_by: str | None
input: str
metadata: str
class _RunWhere(TypedDict, total=False):
workflow_type: str
status: str | Mapping[str, Sequence[str]]
created_by: str
class _RunUpdateData(TypedDict, total=False):
status: WorkflowRunStatus
output: str
metadata: str
class _EventCreateData(TypedDict, total=False):
run_id: str
event_type: str
step_name: str
sequence_number: int
data: str
class _MessageCreateData(TypedDict, total=False):
run_id: str
role: str
content: str
sequence_number: int
session_id: str
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
async def _get_next_sequence_number(prisma_client: Any, run_id: str, table: str) -> int:
async def _get_next_sequence_number(prisma_client: object, run_id: str, table: str) -> int:
"""Return MAX(sequence_number) + 1 for the given run, for either events or messages."""
if table == "events":
rows = await WorkflowEventRepository(prisma_client).table.find_many(
rows: Sequence[_SeqRow] = await WorkflowEventRepository(prisma_client).table.find_many(
where={"run_id": run_id},
order={"sequence_number": "desc"},
take=1,
@ -125,12 +171,12 @@ async def _get_next_sequence_number(prisma_client: Any, run_id: str, table: str)
async def _require_run(
prisma_client: Any,
prisma_client: object,
run_id: str,
user_api_key_dict: UserAPIKeyAuth | None = None,
) -> Any:
) -> _RunRow:
"""Return the run or raise 404. For non-admin callers, also enforce key ownership."""
run: Final = await WorkflowRunRepository(prisma_client).table.find_unique(where={"run_id": run_id})
run: Final[_RunRow | None] = await WorkflowRunRepository(prisma_client).table.find_unique(where={"run_id": run_id})
if run is None:
raise HTTPException(status_code=404, detail=f"Run '{run_id}' not found")
if user_api_key_dict is not None and not _is_admin(user_api_key_dict):
@ -165,7 +211,7 @@ async def create_workflow_run(
raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)
try:
create_data: Final[dict[str, Any]] = {
create_data: Final[_RunCreateData] = {
"workflow_type": data.workflow_type,
"created_by": _caller_key(user_api_key_dict),
}
@ -173,7 +219,7 @@ async def create_workflow_run(
create_data["input"] = _json(data.input)
if data.metadata is not None:
create_data["metadata"] = _json(data.metadata)
run: Final = await WorkflowRunRepository(prisma_client).table.create(data=create_data)
run: Final[_RunRow] = await WorkflowRunRepository(prisma_client).table.create(data=create_data)
return run
except Exception as e:
verbose_proxy_logger.exception("Error creating workflow run: %s", e)
@ -200,7 +246,7 @@ async def list_workflow_runs(
if prisma_client is None:
raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)
where: Final[dict[str, Any]] = {}
where: Final[_RunWhere] = {}
if workflow_type:
where["workflow_type"] = workflow_type
if status:
@ -214,7 +260,7 @@ async def list_workflow_runs(
where["created_by"] = caller
try:
runs: Final = await WorkflowRunRepository(prisma_client).table.find_many(
runs: Final[Sequence[object]] = await WorkflowRunRepository(prisma_client).table.find_many(
where=where,
order={"created_at": "desc"},
take=limit,
@ -241,7 +287,7 @@ async def get_workflow_run(
raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)
try:
run: Final = await WorkflowRunRepository(prisma_client).table.find_unique(
run: Final[_RunRow | None] = await WorkflowRunRepository(prisma_client).table.find_unique(
where={"run_id": run_id},
include={"events": {"order_by": {"sequence_number": "desc"}, "take": 1}},
)
@ -275,7 +321,7 @@ async def update_workflow_run(
if prisma_client is None:
raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)
update: Final[dict[str, Any]] = {}
update: Final[_RunUpdateData] = {}
if data.status is not None:
update["status"] = data.status
if data.output is not None:
@ -290,7 +336,7 @@ async def update_workflow_run(
await _require_run(prisma_client, run_id, user_api_key_dict)
try:
run: Final = await WorkflowRunRepository(prisma_client).table.update(
run: Final[_RunRow | None] = await WorkflowRunRepository(prisma_client).table.update(
where={"run_id": run_id},
data=update,
)
@ -332,7 +378,7 @@ async def append_workflow_event(
for attempt in range(_MAX_SEQUENCE_RETRIES):
try:
seq = await _get_next_sequence_number(prisma_client, run_id, "events")
event_data: dict[str, Any] = {
event_data: _EventCreateData = {
"run_id": run_id,
"event_type": data.event_type,
"step_name": data.step_name,
@ -342,7 +388,7 @@ async def append_workflow_event(
event_data["data"] = _json(data.data)
async with prisma_client.db.tx() as tx:
event = await tx.litellm_workflowevent.create(data=event_data)
event: object = await tx.litellm_workflowevent.create(data=event_data)
if new_status:
await tx.litellm_workflowrun.update(
where={"run_id": run_id},
@ -389,7 +435,7 @@ async def list_workflow_events(
await _require_run(prisma_client, run_id, _read_scope_caller(user_api_key_dict))
try:
events: Final = await WorkflowEventRepository(prisma_client).table.find_many(
events: Final[Sequence[object]] = await WorkflowEventRepository(prisma_client).table.find_many(
where={"run_id": run_id},
order={"sequence_number": "asc"},
take=limit,
@ -424,7 +470,7 @@ async def append_workflow_message(
for attempt in range(_MAX_SEQUENCE_RETRIES):
try:
seq = await _get_next_sequence_number(prisma_client, run_id, "messages")
msg_data: dict[str, Any] = {
msg_data: _MessageCreateData = {
"run_id": run_id,
"role": data.role,
"content": data.content,
@ -432,7 +478,7 @@ async def append_workflow_message(
}
if data.session_id is not None:
msg_data["session_id"] = data.session_id
msg = await WorkflowMessageRepository(prisma_client).table.create(data=msg_data)
msg: object = await WorkflowMessageRepository(prisma_client).table.create(data=msg_data)
return msg
except Exception as e:
@ -473,7 +519,7 @@ async def list_workflow_messages(
await _require_run(prisma_client, run_id, _read_scope_caller(user_api_key_dict))
try:
messages: Final = await WorkflowMessageRepository(prisma_client).table.find_many(
messages: Final[Sequence[object]] = await WorkflowMessageRepository(prisma_client).table.find_many(
where={"run_id": run_id},
order={"sequence_number": "asc"},
take=limit,

19
litellm/proxy/openai_files_endpoints/common_utils.py
View file

@ -375,7 +375,7 @@ def get_team_provider_credentials(
def _provider_credentials(model_id: str) -> dict | None:
credentials: Final = llm_router.get_deployment_credentials_with_provider(model_id=model_id, team_id=team_id)
if credentials is not None and credentials.get("custom_llm_provider") == custom_llm_provider:
return credentials
return {key: value for key, value in credentials.items() if key != "model"}
return None
# 1. Prefer the team's own BYOK deployment, matched by model_info.team_id.
@ -1178,7 +1178,7 @@ async def update_batch_in_database(
managed_files_obj: The managed_files proxy hook object
prisma_client: Prisma database client
verbose_proxy_logger: Logger instance
db_batch_object: Optional existing database object (for comparison)
db_batch_object: Optional existing database object; fetched by unified_object_id when omitted
operation: Description of operation ("update", "cancel", etc.)
user_api_key_dict: Optional auth context for creating managed file IDs
"""
@ -1191,6 +1191,12 @@ async def update_batch_in_database(
if not prisma_client:
return
effective_db_batch_object: Final = (
db_batch_object
if db_batch_object is not None
else await ManagedObjectRepository(prisma_client).table.find_first(where={"unified_object_id": batch_id})
)
# Always normalize the response's file IDs to unified managed IDs
# (mutates in place) so the caller returns unified IDs to the user
# even when we skip the DB update below for an unchanged status.
@ -1200,16 +1206,17 @@ async def update_batch_in_database(
prisma_client=prisma_client,
verbose_proxy_logger=verbose_proxy_logger,
user_api_key_dict=user_api_key_dict,
db_batch_object=db_batch_object,
db_batch_object=effective_db_batch_object,
unified_batch_id=unified_batch_id,
)
# Only update if status has changed (when db_batch_object is provided)
if db_batch_object and response.status == db_batch_object.status:
if effective_db_batch_object and response.status == effective_db_batch_object.status:
return
if db_batch_object:
if effective_db_batch_object:
verbose_proxy_logger.info(
"Updating batch %s status from %s to %s", batch_id, db_batch_object.status, response.status
"Updating batch %s status from %s to %s", batch_id, effective_db_batch_object.status, response.status
)
else:
verbose_proxy_logger.info("Updating batch %s status to %s after %s", batch_id, response.status, operation)

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

@ -565,6 +565,9 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils):
# real parent span.
_metadata["user_api_key"] = user_api_key_dict.api_key
_metadata["litellm_parent_otel_span"] = user_api_key_dict.parent_otel_span
_metadata.update(
LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=user_api_key_dict)
)
kwargs: Final = {
"litellm_params": {

39
litellm/proxy/policy_engine/attachment_registry.py
View file

@ -6,7 +6,7 @@ This allows the same policy to be attached to multiple scopes.
"""
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Any, Final
from typing import TYPE_CHECKING, Any, Final, TypedDict
from litellm._logging import verbose_proxy_logger
from litellm.repositories.table_repositories import PolicyAttachmentRepository
@ -18,9 +18,18 @@ from litellm.types.proxy.policy_engine import (
)
if TYPE_CHECKING:
from collections.abc import Sequence
from prisma.models import LiteLLM_PolicyAttachmentTable
from litellm.proxy.utils import PrismaClient
class PolicyAttachmentMatch(TypedDict):
policy_name: str
matched_via: str
class AttachmentRegistry:
"""
In-memory registry for storing and managing policy attachments.
@ -40,7 +49,7 @@ class AttachmentRegistry:
```
"""
def __init__(self):
def __init__(self) -> None:
self._attachments: list[PolicyAttachment] = []
self._config_attachments: tuple[PolicyAttachment, ...] = ()
self._initialized: bool = False
@ -98,7 +107,7 @@ class AttachmentRegistry:
"""
return [r["policy_name"] for r in self.get_attached_policies_with_reasons(context)]
def get_attached_policies_with_reasons(self, context: PolicyMatchContext) -> list[dict[str, Any]]:
def get_attached_policies_with_reasons(self, context: PolicyMatchContext) -> list[PolicyAttachmentMatch]:
"""
Get list of policy names and match reasons for the given context.
@ -107,8 +116,8 @@ class AttachmentRegistry:
"""
from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher
results: Final[list[dict[str, Any]]] = []
seen_policies: Final[set] = set()
results: Final[list[PolicyAttachmentMatch]] = []
seen_policies: Final[set[str]] = set()
for attachment in self._attachments:
scope = attachment.to_policy_scope()
@ -280,7 +289,9 @@ class AttachmentRegistry:
PolicyAttachmentDBResponse with the created attachment
"""
try:
created_attachment: Final = await PolicyAttachmentRepository(prisma_client).table.create(
created_attachment: Final[LiteLLM_PolicyAttachmentTable] = await PolicyAttachmentRepository(
prisma_client
).table.create(
data={
"policy_name": attachment_request.policy_name,
"scope": attachment_request.scope,
@ -340,9 +351,9 @@ class AttachmentRegistry:
"""
try:
# Get attachment before deleting
attachment: Final = await PolicyAttachmentRepository(prisma_client).table.find_unique(
where={"attachment_id": attachment_id}
)
attachment: Final[LiteLLM_PolicyAttachmentTable | None] = await PolicyAttachmentRepository(
prisma_client
).table.find_unique(where={"attachment_id": attachment_id})
if attachment is None:
raise Exception(f"Attachment with ID {attachment_id} not found")
@ -375,9 +386,9 @@ class AttachmentRegistry:
PolicyAttachmentDBResponse if found, None otherwise
"""
try:
attachment: Final = await PolicyAttachmentRepository(prisma_client).table.find_unique(
where={"attachment_id": attachment_id}
)
attachment: Final[LiteLLM_PolicyAttachmentTable | None] = await PolicyAttachmentRepository(
prisma_client
).table.find_unique(where={"attachment_id": attachment_id})
if attachment is None:
return None
@ -413,7 +424,9 @@ class AttachmentRegistry:
List of PolicyAttachmentDBResponse objects
"""
try:
attachments: Final = await PolicyAttachmentRepository(prisma_client).table.find_many(
attachments: Final[Sequence[LiteLLM_PolicyAttachmentTable]] = await PolicyAttachmentRepository(
prisma_client
).table.find_many(
order={"created_at": "desc"},
)

36
litellm/proxy/proxy_server.py
View file

@ -1111,6 +1111,10 @@ async def proxy_startup_event(app: FastAPI):
prisma_client=prisma_client,
)
)
ProxyStartupEvent._warn_budget_without_db(
max_budget=litellm.max_budget,
prisma_client=prisma_client,
)
### START BATCH WRITING DB + CHECKING NEW MODELS###
if prisma_client is not None:
@ -3909,6 +3913,10 @@ class ProxyConfig:
# whether an existing request predates the prices it just fetched, and re-serving one
# costs a single fetch where skipping one leaves it priced wrong indefinitely
self.model_cost_map_applied_revision: int = 0
# Keys explicitly set in the YAML config file. Used to give YAML
# precedence over stale DB-cached values for these specific keys
# during periodic config reloads (_update_general_settings).
self._yaml_general_settings_keys: set[str] = set() # mutable-ok: populated once at startup, read-only thereafter # fmt: skip
def is_yaml(self, config_file_path: str) -> bool:
if not os.path.isfile(config_file_path):
@ -4839,6 +4847,11 @@ class ProxyConfig:
_hc_staleness = None
_hc_ignore_transient = False
if general_settings:
# Record which keys were explicitly set in the YAML config file.
# These keys take precedence over DB-cached values during periodic
# reloads (see _update_general_settings).
self._yaml_general_settings_keys = set(general_settings.keys()) # mutable-ok: snapshot of YAML keys at load time # fmt: skip
### LOAD KEY MANAGEMENT SETTINGS FIRST (needed for custom secret manager) ###
key_management_settings: Final = general_settings.get("key_management_settings", None)
if key_management_settings is not None:
@ -6049,7 +6062,15 @@ class ProxyConfig:
## STORE PROMPTS IN SPEND LOGS ##
if "store_prompts_in_spend_logs" in _general_settings:
value = _general_settings["store_prompts_in_spend_logs"]
# If the YAML config explicitly set this key, prefer the YAML value
# over the DB-cached value. This ensures config changes deployed via
# CI/CD take effect without requiring a manual /config/update call.
# When YAML does not set this key, the DB value is used (preserving
# admin UI runtime changes).
if "store_prompts_in_spend_logs" in self._yaml_general_settings_keys:
value = general_settings.get("store_prompts_in_spend_logs")
else:
value = _general_settings["store_prompts_in_spend_logs"]
# Normalize case: handle True/true/TRUE, False/false/FALSE, None/null
if value is None:
general_settings["store_prompts_in_spend_logs"] = None
@ -7808,6 +7829,19 @@ def giveup(e):
class ProxyStartupEvent:
@staticmethod
def _warn_budget_without_db(max_budget: float | None, prisma_client: PrismaClient | None) -> None:
if prisma_client is not None or not max_budget or max_budget <= 0:
return
verbose_proxy_logger.warning(
"A proxy-wide budget (litellm.max_budget=%s) is configured but no database is connected, "
"so the budget will NOT be enforced and requests will never be blocked. Set DATABASE_URL or "
"general_settings.database_url and restart. Redis and fail_closed_budget_enforcement do not "
"cover the proxy-wide budget because there is no global spend counter; Redis alone is not a substitute.",
max_budget,
)
@classmethod
def _initialize_startup_logging(
cls,

66
litellm/proxy/search_endpoints/search_tool_management.py
View file

@ -2,13 +2,15 @@
CRUD ENDPOINTS FOR SEARCH TOOLS
"""
from collections.abc import Awaitable, Callable
from datetime import datetime
from typing import Any, Final
from typing import Any, Final, TypeAlias
from fastapi import APIRouter, Depends, HTTPException
from pydantic import BaseModel
from litellm._logging import verbose_proxy_logger
from litellm.constants import UI_SESSION_TOKEN_TEAM_ID
from litellm.proxy._types import (
LiteLLM_TeamTable,
LitellmUserRoles,
@ -46,9 +48,46 @@ def _convert_datetime_to_str(value: datetime | str | None) -> str | None:
return value
TeamObjectLookup: TypeAlias = Callable[[str, UserAPIKeyAuth], Awaitable[LiteLLM_TeamTable]]
async def _team_object_from_db(team_id: str, user_api_key_dict: UserAPIKeyAuth) -> LiteLLM_TeamTable:
from litellm.proxy.auth.auth_checks import get_team_object
from litellm.proxy.proxy_server import (
prisma_client,
proxy_logging_obj,
user_api_key_cache,
)
return await get_team_object(
team_id=team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=user_api_key_dict.parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
def _allowlist_team_id(user_api_key_dict: UserAPIKeyAuth) -> str | None:
"""
The team whose object_permission allowlist scopes this caller, or None when there is none.
Every Admin UI session key is stamped with UI_SESSION_TOKEN_TEAM_ID, a reserved sentinel that
never has a row in LiteLLM_TeamTable (`/team/new` rejects it as a real team id), so looking it
up would raise 404 instead of resolving a team. It carries no allowlist of its own, so the
caller is scoped by its key-level allowlist alone. Any other team id is looked up for real and
a failed lookup still surfaces.
"""
team_id: Final = user_api_key_dict.team_id
if not team_id or team_id == UI_SESSION_TOKEN_TEAM_ID:
return None
return team_id
async def _filter_visible_search_tools(
search_tools: list[SearchToolInfoResponse],
user_api_key_dict: UserAPIKeyAuth,
lookup_team_object: TeamObjectLookup = _team_object_from_db,
) -> list[SearchToolInfoResponse]:
"""
Drop search tools the caller is not authorized to invoke, applying the same
@ -60,25 +99,12 @@ async def _filter_visible_search_tools(
):
return search_tools
from litellm.proxy.auth.auth_checks import (
can_user_view_search_tool,
get_team_object,
)
from litellm.proxy.proxy_server import (
prisma_client,
proxy_logging_obj,
user_api_key_cache,
)
from litellm.proxy.auth.auth_checks import can_user_view_search_tool
team_object: LiteLLM_TeamTable | None = None
if user_api_key_dict.team_id:
team_object = await get_team_object(
team_id=user_api_key_dict.team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=user_api_key_dict.parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
allowlist_team_id: Final = _allowlist_team_id(user_api_key_dict)
team_object: Final[LiteLLM_TeamTable | None] = (
await lookup_team_object(allowlist_team_id, user_api_key_dict) if allowlist_team_id else None
)
visible: Final[list[SearchToolInfoResponse]] = []
for tool in search_tools:
@ -213,6 +239,8 @@ async def list_search_tools(
visible_search_tools: Final = await _filter_visible_search_tools(search_tool_configs, user_api_key_dict)
return ListSearchToolsResponse(search_tools=visible_search_tools)
except HTTPException:
raise
except Exception as e:
verbose_proxy_logger.exception("Error getting search tools: %s", e)
raise HTTPException(status_code=500, detail=str(e))

4
litellm/router.py
View file

@ -8738,7 +8738,7 @@ class Router:
Example:
credentials = router.get_deployment_credentials_with_provider("gpt-4o-litellm")
# Returns: {"api_key": "sk-...", "custom_llm_provider": "openai", ...}
# Returns: {"api_key": "sk-...", "custom_llm_provider": "openai", "model": "gpt-4o", ...}
"""
# Try to get deployment by model_id first
deployment = self.get_deployment(model_id=model_id)
@ -8797,6 +8797,8 @@ class Router:
# Remove the credential name since we've resolved it
credentials.pop("litellm_credential_name", None)
credentials["model"] = deployment.litellm_params.model
# Add custom_llm_provider
if deployment.litellm_params.custom_llm_provider:
credentials["custom_llm_provider"] = deployment.litellm_params.custom_llm_provider

21
litellm/router_strategy/complexity_router/README.md
View file

@ -171,6 +171,27 @@ If 2+ reasoning markers are detected in the user message, the request is automat
Reasoning markers in the system prompt do **not** trigger the reasoning override. This prevents system prompts like "Think step by step before answering" from forcing all requests to the reasoning tier.
### Harness Reminder Blocks
Agent harnesses inject their own context into the conversation as ordinary message text. That text is plumbing, not something a human asked for, so the router strips complete reminder blocks before classifying and picking a tier. A turn that is nothing but a reminder block strips to empty and is skipped, and the router falls back to the last real ask instead
By default a block is anything between `<system-reminder>` and `</system-reminder>`. `reminder_markers` replaces that with your harness's own delimiters. Many harnesses use a different envelope per agent type, so list every pair you emit:
```yaml
model_list:
- model_name: smart-router
litellm_params:
model: auto_router/complexity_router
complexity_router_config:
reminder_markers:
- open: "<<<BEGIN_CONTEXT>>>"
close: "<<<END_CONTEXT>>>"
- open: "[[SUBAGENT_CONTEXT_BEGIN]]"
close: "[[SUBAGENT_CONTEXT_END]]"
```
Setting `reminder_markers` replaces the built-in `<system-reminder>` pair rather than adding to it, so list that pair too if your harness also emits it. Matching is case-insensitive. Blocks that nest or overlap across pairs are stripped whole. An unclosed delimiter is not a block and is left in place, which keeps prose that merely mentions a delimiter from being eaten
### Code Detection
Technical code keywords are detected case-insensitively and include:

2
litellm/router_strategy/complexity_router/__init__.py
View file

@ -16,6 +16,7 @@ from litellm.router_strategy.complexity_router.config import (
DEFAULT_COMPLEXITY_CONFIG,
ComplexityRouterConfig,
ComplexityTier,
ReminderMarkerPair,
)
__all__ = [
@ -24,5 +25,6 @@ __all__ = [
"ComplexityRouter",
"ComplexityRouterConfig",
"ComplexityTier",
"ReminderMarkerPair",
"classification_system_prompt",
]

71
litellm/router_strategy/complexity_router/complexity_router.py
View file

@ -19,7 +19,7 @@ import asyncio
import random
import re
from collections.abc import Iterator, Mapping, Sequence
from itertools import islice
from itertools import accumulate, islice
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, cast
@ -233,6 +233,7 @@ def _effective_turn_off_message_logging(request_kwargs: Mapping[str, Any] | None
_REMINDER_OPEN: Final = "<system-reminder>"
_REMINDER_CLOSE: Final = "</system-reminder>"
_DEFAULT_REMINDER_MARKERS: Final = ((_REMINDER_OPEN, _REMINDER_CLOSE),)
_TRUNCATION_MARKER: Final = "..."
@ -253,10 +254,8 @@ def _message_text(content: object) -> str:
return content if isinstance(content, str) else ""
def _reminder_block_spans(
lowered: str, open_marker: str = _REMINDER_OPEN, close_marker: str = _REMINDER_CLOSE
) -> Iterator[tuple[int, int]]:
"""Span of each complete reminder block, left to right.
def _reminder_block_spans(lowered: str, open_marker: str, close_marker: str) -> Iterator[tuple[int, int]]:
"""Span of each complete reminder block for one marker pair, left to right.
Literal `str.find`, not a regex: the delimiters are fixed strings, and `<system-reminder>.*?`
retried its lazy quantifier from every opening tag, so repeated unclosed tags were quadratic
@ -272,17 +271,36 @@ def _reminder_block_spans(
yield start, cursor
def _strip_reminder_blocks(text: str, open_marker: str = _REMINDER_OPEN, close_marker: str = _REMINDER_CLOSE) -> str:
"""Remove every complete reminder block from text, keeping everything written around them."""
spans: Final = tuple(_reminder_block_spans(text.lower(), open_marker, close_marker))
def _strip_reminder_blocks(text: str, marker_pairs: tuple[tuple[str, str], ...] = _DEFAULT_REMINDER_MARKERS) -> str:
"""Remove every complete reminder block from text, keeping everything written around them.
Blocks from different pairs can nest or overlap, which the gap construction below would
otherwise mishandle: an inner block's end would resume the kept text partway through the outer
block, leaking the rest of that block into the classified ask. Running the block ends through a
maximum resumes each gap past the furthest block seen so far, which collapses nested and
overlapping spans without a separate merge pass. A single pair's ends already increase, so the
maximum is the identity there and the default path is byte-identical to a plain scan.
Deliberately linear in both the text and the block count. This runs pre-routing on input any
keyholder controls, and both a regex scan and a fold that rebuilds a growing tuple of merged
spans go quadratic on inputs that are cheap to send.
"""
lowered: Final = text.lower()
spans: Final = tuple(
sorted(
span
for open_marker, close_marker in marker_pairs
for span in _reminder_block_spans(lowered, open_marker, close_marker)
)
)
if not spans:
return text.strip()
keep_from: Final = (0, *(end for _, end in spans))
keep_from: Final = (0, *accumulate((end for _, end in spans), max))
keep_to: Final = (*(start for start, _ in spans), len(text))
return " ".join(kept for a, b in zip(keep_from, keep_to) if (kept := text[a:b].strip()))
def _human_text(content: object, open_marker: str = _REMINDER_OPEN, close_marker: str = _REMINDER_CLOSE) -> str:
def _human_text(content: object, marker_pairs: tuple[tuple[str, str], ...] = _DEFAULT_REMINDER_MARKERS) -> str:
"""Message content as the text a human wrote, with complete reminder blocks removed.
Harnesses inject reminders as ordinary text alongside the live ask, so the block is stripped and
@ -291,18 +309,18 @@ def _human_text(content: object, open_marker: str = _REMINDER_OPEN, close_marker
one, and this same string drives escalation keywords and keyword_tier_rules, which choose the
model and therefore the spend. An unclosed tag is not a block and is left intact.
"""
return _strip_reminder_blocks(_message_text(content), open_marker, close_marker)
return _strip_reminder_blocks(_message_text(content), marker_pairs)
def _iter_human_asks_newest_first(
messages: Sequence[Mapping[str, object]], markers: tuple[str, str] = (_REMINDER_OPEN, _REMINDER_CLOSE)
messages: Sequence[Mapping[str, object]],
marker_pairs: tuple[tuple[str, str], ...] = _DEFAULT_REMINDER_MARKERS,
) -> Iterator[str]:
"""Yield user-turn texts that carry a real human ask, newest first, with harness noise removed."""
open_marker, close_marker = markers
return (
text
for msg in reversed(messages)
if msg.get("role") == "user" and (text := _human_text(msg.get("content"), open_marker, close_marker))
if msg.get("role") == "user" and (text := _human_text(msg.get("content"), marker_pairs))
)
@ -341,7 +359,8 @@ def _conversation_is_continuing(messages: Sequence[Mapping[str, object]] | None)
def _newest_turn_ask(
messages: Sequence[Mapping[str, object]], markers: tuple[str, str] = (_REMINDER_OPEN, _REMINDER_CLOSE)
messages: Sequence[Mapping[str, object]],
marker_pairs: tuple[tuple[str, str], ...] = _DEFAULT_REMINDER_MARKERS,
) -> str | None:
"""The human ask on the newest user turn, or None when that turn carries only plumbing.
@ -352,12 +371,12 @@ def _newest_turn_ask(
newest_user_turn: Final = next((msg for msg in reversed(messages) if msg.get("role") == "user"), None)
if newest_user_turn is None:
return None
return _human_text(newest_user_turn.get("content"), *markers) or None
return _human_text(newest_user_turn.get("content"), marker_pairs) or None
def _extract_current_ask_and_system_prompt(
messages: Sequence[Mapping[str, object]],
markers: tuple[str, str] = (_REMINDER_OPEN, _REMINDER_CLOSE),
marker_pairs: tuple[tuple[str, str], ...] = _DEFAULT_REMINDER_MARKERS,
) -> tuple[str | None, str | None]:
"""The last real human ask and the last system prompt; either is None if absent.
@ -365,7 +384,7 @@ def _extract_current_ask_and_system_prompt(
the caller routes to its default model. That is the correct answer rather than a gap to fill:
filling it would hand tier selection to harness-injected text.
"""
current_ask: Final = next(_iter_human_asks_newest_first(messages, markers), None)
current_ask: Final = next(_iter_human_asks_newest_first(messages, marker_pairs), None)
system_prompt: Final = next(
(
text
@ -385,7 +404,7 @@ def _truncate(text: str, limit: int) -> str:
def _iter_context_turns_newest_first(
messages: Sequence[Mapping[str, object]],
include_assistant: bool,
markers: tuple[str, str] = (_REMINDER_OPEN, _REMINDER_CLOSE),
marker_pairs: tuple[tuple[str, str], ...] = _DEFAULT_REMINDER_MARKERS,
) -> Iterator[tuple[str, str]]:
"""Yield (role, text) for turns eligible as classifier context, newest first.
@ -401,7 +420,7 @@ def _iter_context_turns_newest_first(
for msg in reversed(messages)
if isinstance(role := msg.get("role"), str)
and role in roles
and (text := _human_text(msg.get("content"), *markers))
and (text := _human_text(msg.get("content"), marker_pairs))
)
@ -411,7 +430,7 @@ def _extract_prior_turns(
window_size: int,
per_turn_chars: int,
include_assistant: bool,
markers: tuple[str, str] = (_REMINDER_OPEN, _REMINDER_CLOSE),
marker_pairs: tuple[tuple[str, str], ...] = _DEFAULT_REMINDER_MARKERS,
) -> tuple[tuple[str, str], ...]:
"""Up to window_size turns other than current_ask, oldest first, as (role, text).
@ -431,7 +450,7 @@ def _extract_prior_turns(
prior: Final = islice(
(
turn
for turn in _iter_context_turns_newest_first(messages, include_assistant, markers)
for turn in _iter_context_turns_newest_first(messages, include_assistant, marker_pairs)
if turn[1] != current_ask
),
window_size,
@ -556,7 +575,11 @@ class ComplexityRouter(CustomLogger):
if self.config.escalation_keywords is not None
else DEFAULT_ESCALATION_KEYWORDS
)
self._reminder_markers: tuple[str, str] = self.config.reminder_markers or (_REMINDER_OPEN, _REMINDER_CLOSE)
self._reminder_markers: tuple[tuple[str, str], ...] = (
tuple((pair.open, pair.close) for pair in self.config.reminder_markers)
if self.config.reminder_markers
else _DEFAULT_REMINDER_MARKERS
)
# Lazily built on first semantic request and cached for reuse (route
# embeddings are static, only the prompt is embedded per request). The lock
@ -993,7 +1016,7 @@ class ComplexityRouter(CustomLogger):
window_size=self.config.classifier_context_window_size,
per_turn_chars=self.config.classifier_context_per_turn_chars,
include_assistant=include_assistant,
markers=self._reminder_markers,
marker_pairs=self._reminder_markers,
)
if context_enabled
else ()

47
litellm/router_strategy/complexity_router/config.py
View file

@ -59,6 +59,30 @@ class KeywordTierRule(BaseModel):
return self
class ReminderMarkerPair(BaseModel):
"""One open/close delimiter pair a harness wraps injected context in.
Normalizing here rather than at the scan is what makes matching case-insensitive: markers reach
the scan already lowered, so it lowercases only the haystack and never the needles. Stripping
keeps YAML indentation whitespace from becoming part of the delimiter.
"""
open: str = Field(description="Opening delimiter, e.g. '<system-reminder>'")
close: str = Field(description="Closing delimiter, e.g. '</system-reminder>'")
@model_validator(mode="after")
def _normalize(self) -> "ReminderMarkerPair":
open_marker: Final = self.open.strip().lower()
close_marker: Final = self.close.strip().lower()
if not open_marker or not close_marker:
raise ValueError("reminder_markers entries must not be blank")
if open_marker == close_marker:
raise ValueError("reminder_markers open and close must be different strings")
self.open = open_marker
self.close = close_marker
return self
# ─── Default Keyword Lists ───
# Note: Keywords should be full words/phrases to avoid substring false positives.
# The matching logic uses word boundary detection for single-word keywords.
@ -498,12 +522,15 @@ class ComplexityRouterConfig(BaseModel):
description="RoutingPlugin instances that narrow the classified tier's candidate models before selection",
)
reminder_markers: tuple[str, str] | None = Field(
reminder_markers: tuple[ReminderMarkerPair, ...] | None = Field(
default=None,
min_length=1,
description=(
"Override the (open, close) marker pair used to recognize and strip harness-injected "
"reminder blocks before classification. Defaults to Claude Code's convention, "
"('<system-reminder>', '</system-reminder>'), when unset. Matching is case-insensitive."
"Override the delimiter pairs used to recognize and strip harness-injected reminder "
"blocks before classification. A harness that wraps injected context differently per "
"agent type (main, subagent, cron) lists every pair it emits. Replaces, rather than "
"adds to, the built-in default of ('<system-reminder>', '</system-reminder>'), so a "
"harness that also emits that pair lists it too. Matching is case-insensitive."
),
)
@ -601,18 +628,6 @@ class ComplexityRouterConfig(BaseModel):
)
return self
@model_validator(mode="after")
def _normalize_reminder_markers(self) -> "ComplexityRouterConfig":
if self.reminder_markers is None:
return self
open_marker, close_marker = (marker.strip().lower() for marker in self.reminder_markers)
if not open_marker or not close_marker:
raise ValueError("reminder_markers entries must not be blank")
if open_marker == close_marker:
raise ValueError("reminder_markers open and close must be different strings")
self.reminder_markers = (open_marker, close_marker)
return self
def tier_label(self, tier: ComplexityTier) -> str:
"""Operator-facing display name for a tier, falling back to its canonical name."""
return self.tier_labels.get(tier, "").strip() or tier.value

10
litellm/types/guardrails.py
View file

@ -753,6 +753,16 @@ class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch up
),
)
scan_only_tool_results: bool | None = Field(
default=None,
description=(
"When True, unified guardrails only evaluate tool results, the untrusted data an "
"agent feeds back into the model, and skip system, user, and assistant content. "
"Intended for agent harnesses whose own prompt scaffolding is trusted but often "
"trips prompt-attack detectors."
),
)
# Lakera specific params
category_thresholds: LakeraCategoryThresholds | None = Field(
default=None,

9
litellm/types/proxy/guardrails/guardrail_hooks/zscaler_ai_guard.py
View file

@ -79,6 +79,15 @@ class ZscalerAIGuardConfigModel(GuardrailConfigModel):
json_schema_extra={"ui_type": GuardrailParamUITypes.BOOL},
)
timeout: float | None = Field(
default=None,
description=(
"Timeout for each Zscaler AI Guard API call, in seconds. Must be positive. "
"Raise it if scans fail under load with 'Connection timed out'. "
"Defaults to 5 seconds."
),
)
@model_validator(mode="after")
def validate_endpoint_configuration(self) -> "ZscalerAIGuardConfigModel":
"""

8
litellm/types/router.py
View file

@ -200,6 +200,9 @@ class CredentialLiteLLMParams(BaseModel):
aws_bedrock_runtime_endpoint: str | None = None
aws_bedrock_project_id: str | None = None
s3_bucket_name: str | None = None
s3_region_name: str | None = None
s3_encryption_key_id: str | None = None
aws_batch_role_arn: str | None = None
## IBM WATSONX ##
watsonx_region_name: str | None = None
@ -272,11 +275,6 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams):
quality_router_config: dict | None = None
quality_router_default_model: str | None = None
# Batch/File API Params
s3_bucket_name: str | None = None
s3_encryption_key_id: str | None = None
gcs_bucket_name: str | None = None
# Vector Store Params
vector_store_id: str | None = None
milvus_text_field: str | None = None

4
litellm/types/utils.py
View file

@ -258,6 +258,8 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
output_cost_per_video_token: float | None # for gemini omni models with video output
output_vector_size: int | None
output_cost_per_reasoning_token: float | None
output_cost_per_reasoning_token_flex: float | None
output_cost_per_reasoning_token_priority: float | None
output_cost_per_video_per_second: float | None # only for vertex ai models
output_cost_per_audio_per_second: float | None # only for vertex ai models
output_cost_per_second: float | None # for OpenAI Speech models
@ -3308,6 +3310,8 @@ class CustomPricingLiteLLMParams(BaseModel):
output_cost_per_image_token: float | None = None
output_cost_per_video_token: float | None = None
output_cost_per_reasoning_token: float | None = None
output_cost_per_reasoning_token_flex: float | None = None
output_cost_per_reasoning_token_priority: float | None = None
output_cost_per_video_per_second: float | None = None
output_cost_per_audio_per_second: float | None = None
search_context_cost_per_query: dict[str, Any] | None = None

4
litellm/utils.py
View file

@ -5533,6 +5533,10 @@ def _get_model_info_helper(
output_cost_per_audio_token=_model_info.get("output_cost_per_audio_token", None),
output_cost_per_character=_model_info.get("output_cost_per_character", None),
output_cost_per_reasoning_token=_model_info.get("output_cost_per_reasoning_token", None),
output_cost_per_reasoning_token_flex=_model_info.get("output_cost_per_reasoning_token_flex", None),
output_cost_per_reasoning_token_priority=_model_info.get(
"output_cost_per_reasoning_token_priority", None
),
output_cost_per_token_above_128k_tokens=_model_info.get(
"output_cost_per_token_above_128k_tokens", None
),

31
model_prices_and_context_window.json
View file

@ -22251,7 +22251,9 @@
},
"gpt-4.1-2025-04-14": {
"cache_read_input_token_cost": 5e-07,
"cache_read_input_token_cost_priority": 8.75e-07,
"input_cost_per_token": 2e-06,
"input_cost_per_token_priority": 3.5e-06,
"input_cost_per_token_batches": 1e-06,
"litellm_provider": "openai",
"max_input_tokens": 1047576,
@ -22259,6 +22261,7 @@
"max_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 8e-06,
"output_cost_per_token_priority": 1.4e-05,
"output_cost_per_token_batches": 4e-06,
"supported_endpoints": [
"/v1/chat/completions",
@ -22322,7 +22325,9 @@
},
"gpt-4.1-mini-2025-04-14": {
"cache_read_input_token_cost": 1e-07,
"cache_read_input_token_cost_priority": 1.75e-07,
"input_cost_per_token": 4e-07,
"input_cost_per_token_priority": 7e-07,
"input_cost_per_token_batches": 2e-07,
"litellm_provider": "openai",
"max_input_tokens": 1047576,
@ -22330,6 +22335,7 @@
"max_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 1.6e-06,
"output_cost_per_token_priority": 2.8e-06,
"output_cost_per_token_batches": 8e-07,
"supported_endpoints": [
"/v1/chat/completions",
@ -22392,7 +22398,9 @@
},
"gpt-4.1-nano-2025-04-14": {
"cache_read_input_token_cost": 2.5e-08,
"cache_read_input_token_cost_priority": 5e-08,
"input_cost_per_token": 1e-07,
"input_cost_per_token_priority": 2e-07,
"input_cost_per_token_batches": 5e-08,
"litellm_provider": "openai",
"max_input_tokens": 1047576,
@ -22400,6 +22408,7 @@
"max_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 4e-07,
"output_cost_per_token_priority": 8e-07,
"output_cost_per_token_batches": 2e-07,
"supported_endpoints": [
"/v1/chat/completions",
@ -22468,7 +22477,9 @@
},
"gpt-4o-2024-08-06": {
"cache_read_input_token_cost": 1.25e-06,
"cache_read_input_token_cost_priority": 2.125e-06,
"input_cost_per_token": 2.5e-06,
"input_cost_per_token_priority": 4.25e-06,
"input_cost_per_token_batches": 1.25e-06,
"litellm_provider": "openai",
"max_input_tokens": 128000,
@ -22476,6 +22487,7 @@
"max_tokens": 16384,
"mode": "chat",
"output_cost_per_token": 1e-05,
"output_cost_per_token_priority": 1.7e-05,
"output_cost_per_token_batches": 5e-06,
"supports_function_calling": true,
"supports_parallel_function_calling": true,
@ -22488,7 +22500,9 @@
},
"gpt-4o-2024-11-20": {
"cache_read_input_token_cost": 1.25e-06,
"cache_read_input_token_cost_priority": 2.125e-06,
"input_cost_per_token": 2.5e-06,
"input_cost_per_token_priority": 4.25e-06,
"input_cost_per_token_batches": 1.25e-06,
"litellm_provider": "openai",
"max_input_tokens": 128000,
@ -22496,6 +22510,7 @@
"max_tokens": 16384,
"mode": "chat",
"output_cost_per_token": 1e-05,
"output_cost_per_token_priority": 1.7e-05,
"output_cost_per_token_batches": 5e-06,
"supports_function_calling": true,
"supports_parallel_function_calling": true,
@ -22795,7 +22810,9 @@
},
"gpt-4o-mini-2024-07-18": {
"cache_read_input_token_cost": 7.5e-08,
"cache_read_input_token_cost_priority": 1.25e-07,
"input_cost_per_token": 1.5e-07,
"input_cost_per_token_priority": 2.5e-07,
"input_cost_per_token_batches": 7.5e-08,
"litellm_provider": "openai",
"max_input_tokens": 128000,
@ -22803,6 +22820,7 @@
"max_tokens": 16384,
"mode": "chat",
"output_cost_per_token": 6e-07,
"output_cost_per_token_priority": 1e-06,
"output_cost_per_token_batches": 3e-07,
"search_context_cost_per_query": {
"search_context_size_high": 0.03,
@ -25152,6 +25170,7 @@
"cache_read_input_token_cost": 5e-09,
"cache_read_input_token_cost_flex": 2.5e-09,
"input_cost_per_token": 5e-08,
"input_cost_per_token_priority": 2.5e-06,
"input_cost_per_token_flex": 2.5e-08,
"litellm_provider": "openai",
"max_input_tokens": 272000,
@ -29379,13 +29398,19 @@
},
"o3-2025-04-16": {
"cache_read_input_token_cost": 5e-07,
"cache_read_input_token_cost_flex": 2.5e-07,
"cache_read_input_token_cost_priority": 8.75e-07,
"input_cost_per_token": 2e-06,
"input_cost_per_token_flex": 1e-06,
"input_cost_per_token_priority": 3.5e-06,
"litellm_provider": "openai",
"max_input_tokens": 200000,
"max_output_tokens": 100000,
"max_tokens": 100000,
"mode": "chat",
"output_cost_per_token": 8e-06,
"output_cost_per_token_flex": 4e-06,
"output_cost_per_token_priority": 1.4e-05,
"supported_endpoints": [
"/v1/responses",
"/v1/chat/completions",
@ -29600,13 +29625,19 @@
},
"o4-mini-2025-04-16": {
"cache_read_input_token_cost": 2.75e-07,
"cache_read_input_token_cost_flex": 1.375e-07,
"cache_read_input_token_cost_priority": 5e-07,
"input_cost_per_token": 1.1e-06,
"input_cost_per_token_flex": 5.5e-07,
"input_cost_per_token_priority": 2e-06,
"litellm_provider": "openai",
"max_input_tokens": 200000,
"max_output_tokens": 100000,
"max_tokens": 100000,
"mode": "chat",
"output_cost_per_token": 4.4e-06,
"output_cost_per_token_flex": 2.2e-06,
"output_cost_per_token_priority": 8e-06,
"supports_function_calling": true,
"supports_parallel_function_calling": false,
"supports_pdf_input": true,

4
pyproject.toml
View file

@ -66,8 +66,8 @@ proxy = [
"azure-identity>=1.25.2,<2.0",
"azure-storage-blob>=12.28.0,<13.0",
"mcp>=1.28.1,<2.0",
"litellm-proxy-extras==0.4.83",
"litellm-enterprise==0.1.53",
"litellm-proxy-extras==0.4.84",
"litellm-enterprise==0.1.54",
"RestrictedPython>=8.1,<9.0",
"rich>=13.9.4,<14.0",
"InquirerPy>=0.3.4,<1.0",

20
ruff-strict-budget.json
View file

@ -1,30 +1,30 @@
{
"ANN001": {
"limit": 3126
"limit": 3121
},
"ANN002": {
"limit": 71
},
"ANN003": {
"limit": 836
"limit": 834
},
"ANN201": {
"limit": 2037
"limit": 2033
},
"ANN202": {
"limit": 869
"limit": 865
},
"ANN204": {
"limit": 715
"limit": 713
},
"ANN205": {
"limit": 115
"limit": 114
},
"ANN206": {
"limit": 133
},
"ANN401": {
"limit": 1689
"limit": 1630
},
"ASYNC230": {
"limit": 11
@ -42,7 +42,7 @@
"limit": 81
},
"B010": {
"limit": 194
"limit": 190
},
"B018": {
"limit": 2
@ -222,7 +222,7 @@
"limit": 0
},
"RET504": {
"limit": 178
"limit": 177
},
"RUF010": {
"limit": 0
@ -306,7 +306,7 @@
"limit": 0
},
"TID251": {
"limit": 1242
"limit": 1240
},
"TRY002": {
"limit": 528

2
ruff-strict.toml
View file

@ -15,7 +15,7 @@ max-args = 5
"typing.Any".msg = "Use a concrete type. Frozen slots=True dataclass (preferred) / NamedTuple / ReadOnly TypedDict for payloads."
"typing_extensions.Any".msg = "Same as typing.Any."
"typing.List".msg = "tuple[X, ...] for state, Sequence[X] for params."
"typing.Dict".msg = "Frozen dataclass / NamedTuple / ReadOnly TypedDict; create a Mapping alias with concrete value types if truly dynamic."
"typing.Dict".msg = "Frozen dataclass / NamedTuple / ReadOnly TypedDict; if truly dynamic, use MappingProxyType."
"typing.Set".msg = "frozenset[X] or AbstractSet[X]."
"typing.MutableSequence".msg = "Sequence[X]."
"typing.MutableMapping".msg = "See typing.Dict."

20
scripts/check_type_discipline.py
View file

@ -17,13 +17,14 @@ LIT002 Mutable-collection *construction*: a list/dict/set literal or comprehens
a call to a mutable constructor (list/dict/set/deque/defaultdict/Counter/...).
Catches the unannotated seed-then-mutate pattern LIT001 cannot see (`acc = []`).
Build the value in one shot and freeze it: a `tuple`/`frozenset` wrapping a
generator (`tuple(f(x) for x in xs)`), a tuple literal, or a frozen dataclass /
NamedTuple / ReadOnly TypedDict. Generator expressions and `tuple`/`frozenset`
calls are not construction and pass. Annotation-internal lists (`Callable[[int],
str]`) are exempt, as is a value passed directly to a freezing wrapper
(`tuple(...)`, `frozenset(...)`, `MappingProxyType(...)`): it is frozen before
it can escape, though anything mutable nested inside it still counts.
Suppress with `# mutable-ok: <reason>`.
generator (`tuple(f(x) for x in xs)`), a tuple literal, a frozen dataclass /
NamedTuple / ReadOnly TypedDict, or (if it really must be dynamic) a
MappingProxyType wrapping a dict literal or comprehension. Generator expressions
and freezing-wrapper calls (`tuple(...)`, `frozenset(...)`,
`MappingProxyType(...)`) are not construction and pass, as does the value passed
directly to a wrapper: it is frozen before it can escape, though anything
mutable nested inside it still counts. Annotation-internal lists
(`Callable[[int], str]`) are exempt. Suppress with `# mutable-ok: <reason>`.
LIT003 noqa suppression without rule codes or without a reason.
Required shape: `# noqa: TID251 # <reason>`
LIT004 pyright/mypy ignore without bracketed codes or without a reason.
@ -488,8 +489,9 @@ def iter_construction_violations(path: Path, tree: ast.AST, comments: Comments)
path, node.lineno, "LIT002",
f"mutable {kind}: this builds a collection that can be grown or rewritten. "
f"Build it in one shot and freeze it -- a tuple/frozenset wrapping a generator "
f"(`tuple(f(x) for x in xs)`), a tuple literal, or a frozen dataclass / NamedTuple "
f"/ ReadOnly TypedDict (suppress: `# mutable-ok: <reason>`)",
f"(`tuple(f(x) for x in xs)`), a tuple literal, a frozen dataclass / NamedTuple "
f"/ ReadOnly TypedDict, or (if it really must be dynamic) a MappingProxyType "
f"wrapping a dict literal or comprehension (suppress: `# mutable-ok: <reason>`)",
)

36
scripts/prisma_generate_if_needed.py
View file

@ -9,13 +9,21 @@ client (a fresh or reinstalled prisma package) forces a regenerate even when
the stamp matches. The prisma package itself is never imported here: once
generated it re-exports the whole client on import, which costs more than the
generate this script exists to skip.
prisma resolves its generator command (``prisma-client-py``) through a plain
PATH lookup, never through the interpreter that invoked ``prisma generate``,
so the generate runs with this interpreter's own bin directory pinned to the
front of PATH; without that pin the client lands in whichever venv the caller
happened to have on PATH (or the generate fails outright when none is).
"""
import hashlib
import importlib.metadata
import importlib.util
import os
import subprocess
import sys
from collections.abc import Callable, Mapping
from pathlib import Path
REPO_ROOT = Path(__file__).resolve().parent.parent
@ -46,6 +54,25 @@ def client_is_generated() -> bool:
)
def env_with_own_bin_first(base_env: Mapping[str, str]) -> dict[str, str]:
bin_dir = str(Path(sys.executable).parent)
inherited = base_env.get("PATH")
path = os.pathsep.join((bin_dir, inherited)) if inherited else bin_dir
return {**base_env, "PATH": path}
def _run_command(cmd: list[str], cwd: Path, env: dict[str, str]) -> int:
return subprocess.run(cmd, cwd=cwd, env=env).returncode
def run_generate(run: Callable[[list[str], Path, dict[str, str]], int] = _run_command) -> int:
return run(
[sys.executable, "-m", "prisma", "generate", "--schema", str(SCHEMA)],
REPO_ROOT,
env_with_own_bin_first(os.environ),
)
def main() -> int:
version = importlib.metadata.version("prisma")
expected = stamp_value(SCHEMA.read_bytes(), version)
@ -55,12 +82,9 @@ def main() -> int:
f"(prisma {version}); skipping prisma generate"
)
return 0
result = subprocess.run(
[sys.executable, "-m", "prisma", "generate", "--schema", str(SCHEMA)],
cwd=REPO_ROOT,
)
if result.returncode != 0:
return result.returncode
returncode = run_generate()
if returncode != 0:
return returncode
STAMP.write_text(expected)
return 0

149
scripts/type_check_gate.py
View file

@ -12,6 +12,18 @@ a red once two PRs each land near the limit and their sum crosses it: the
bystander's count equals its base, so it is spared, while any PR that actually
grows the rule past its limit still fails.
Installed packages are part of the measurement: a typed dependency that is
present changes what basedpyright can prove (and therefore which diagnostics
fire) versus when it is absent, so counts from two differently provisioned
venvs are not comparable and their comparison produces phantom breaches no
diff hunk explains. The gate therefore provisions its own environment at
``.venv-typecheck`` (a frozen ``uv sync`` of one canonical dependency-group
set, plus a generated Prisma client) and runs every basedpyright pass from it,
so pre-commit, the CI lint job, and the artifact publisher measure one package
set by construction; re-syncs of an up-to-date env are a near-instant no-op.
The group set is folded into the cache and artifact fingerprint, so counts
recorded under a different set are never matched, only recomputed.
The gate runs basedpyright itself, for both the head and the base pass, with
``NODE_OPTIONS`` raised to the heap this repo needs: basedpyright's node
process OOMs at the ~4 GB default, and when callers had to remember the flag,
@ -21,9 +33,9 @@ matters once some rule is over its limit, so when none is the base pass is
skipped outright. When it is needed, it is a second basedpyright pass over a
detached worktree at the merge-base, run under the same environment so import
resolution matches, and its per-rule counts are cached under the repo's git
common dir keyed by merge-base commit,
``pyrightconfig.json``, and ``uv.lock``, so re-runs against the same branch
point pay for it once. A CI workflow publishes every staging commit's counts as
common dir keyed by merge-base commit, ``pyrightconfig.json``, ``uv.lock``,
the Prisma schema, and the dependency-group set, so re-runs against the same
branch point pay for it once. A CI workflow publishes every staging commit's counts as
an artifact (``--emit-counts-dir`` is its entry point), and on a disk-cache miss
the gate first tries to download the merge-base's artifact through the ``gh``
CLI; any fetch failure falls back silently to the local base pass, so the gate
@ -51,7 +63,7 @@ import sys
import tempfile
import zipfile
from collections import Counter
from collections.abc import Callable, Iterator, Mapping
from collections.abc import Callable, Iterator, Mapping, Sequence
from pathlib import Path
from typing import Final, NamedTuple
@ -61,13 +73,23 @@ PYRIGHT_CONFIG = REPO_ROOT / "pyrightconfig.json"
UV_LOCK = REPO_ROOT / "uv.lock"
DEFAULT_BASE = "origin/litellm_internal_staging"
CACHE_FILE_PREFIX = "basedpyright-base-"
CACHE_KEEP_ENTRIES = 8
ARTIFACT_NAME_PREFIX = "basedpyright-counts-"
GH_TIMEOUT_SECONDS = 10
# The one environment every basedpyright pass measures in. The group set is
# the slim one the CI publisher has always installed (not bootstrap's fatter
# --extra proxy env), so the committed budgets stay valid; changing it re-keys
# every cache and artifact fingerprint, so stale counts can never be matched.
TYPECHECK_ENV_DIR = REPO_ROOT / ".venv-typecheck"
TYPECHECK_DEP_GROUPS = ("proxy-dev", "e2e-dev")
PRISMA_GENERATE_SCRIPT = REPO_ROOT / "scripts" / "prisma_generate_if_needed.py"
PRISMA_SCHEMA = REPO_ROOT / "litellm" / "proxy" / "schema.prisma"
# basedpyright's node process needs more than the ~4 GB default heap on this
# repo; appended last so it wins node's last-flag-wins resolution over any
# caller-set value while preserving the caller's other NODE_OPTIONS flags.
NODE_HEAP_OPTION = "--max-old-space-size=12288"
NODE_HEAP_OPTION = "--max-old-space-size=8192"
# Bucket for a basedpyright diagnostic with no `rule`. Counted so it's gated.
UNCODED = "<uncoded>"
@ -129,14 +151,83 @@ def node_options_with_heap(base_env: Mapping[str, str]) -> str:
return f"{base_env.get('NODE_OPTIONS', '')} {NODE_HEAP_OPTION}".strip()
def run_basedpyright(cwd: Path = REPO_ROOT) -> str:
"""One basedpyright pass over `cwd` with the raised node heap exported.
def typecheck_python_version() -> str | None:
"""The interpreter version to build the owned env with, read from
pyrightconfig's `pythonVersion` so the packages installed for basedpyright
to see always come from the same version it type-checks against."""
try:
config = json.loads(PYRIGHT_CONFIG.read_text())
except (OSError, ValueError):
return None
version: Final = config.get("pythonVersion") if isinstance(config, dict) else None
return version if isinstance(version, str) else None
Exit 0 (clean) and 1 (errors found) are both output-bearing runs; anything
else is a crash and fails loudly instead of reading as zero errors."""
exe = shutil.which("basedpyright") or "basedpyright"
def typecheck_env_commands(env_dir: Path = TYPECHECK_ENV_DIR) -> tuple[tuple[str, ...], ...]:
python_pin: Final = typecheck_python_version()
sync: Final = (
"uv",
"sync",
"--frozen",
*(("--python", python_pin) if python_pin else ()),
*(flag for group in TYPECHECK_DEP_GROUPS for flag in ("--group", group)),
)
generate: Final = (str(env_dir / "bin" / "python"), str(PRISMA_GENERATE_SCRIPT))
return (sync, generate)
def _run_provision_step(cmd: tuple[str, ...], env: Mapping[str, str]) -> int:
proc = subprocess.run(
[exe, "--outputjson"],
list(cmd), cwd=REPO_ROOT, env=dict(env), capture_output=True, text=True
)
if proc.returncode != 0:
sys.stderr.write(proc.stdout)
sys.stderr.write(proc.stderr)
return proc.returncode
def ensure_typecheck_env(
env_dir: Path = TYPECHECK_ENV_DIR,
run: Callable[[tuple[str, ...], Mapping[str, str]], int] = _run_provision_step,
) -> Path:
"""Sync the gate-owned venv (and its generated Prisma client) before a
measurement pass. Unconditional on purpose: an up-to-date env makes both
steps near-instant no-ops, and skipping them on a heuristic is how the
measured environment and the fingerprinted one drift apart."""
if not env_dir.exists():
sys.stderr.write(
f"provisioning {env_dir.name} (first run installs packages and "
"generates the Prisma client; re-runs are near-instant no-ops)\n"
)
env: Final = {**os.environ, "UV_PROJECT_ENVIRONMENT": str(env_dir)}
for cmd in typecheck_env_commands(env_dir):
if run(cmd, env) != 0:
raise SystemExit(
f"could not provision the type-check environment at {env_dir}: "
f"`{' '.join(cmd)}` failed"
)
return env_dir
def run_basedpyright(cwd: Path = REPO_ROOT, env_dir: Path = TYPECHECK_ENV_DIR) -> str:
"""One basedpyright pass over `cwd` from the gate-owned venv, with the
raised node heap exported.
`--pythonpath` pins import resolution to the owned env's interpreter; it is
the only pin that works, because basedpyright auto-detects a `.venv` in the
project root and that beats both PATH order and VIRTUAL_ENV, silently
measuring the caller's fatter venv (whose extra typed packages flip
diagnostics) whenever the repo has one. Exit 0 (clean) and 1 (errors
found) are both output-bearing runs; anything else is a crash and fails
loudly instead of reading as zero errors."""
bin_dir: Final = env_dir / "bin"
proc = subprocess.run(
[
str(bin_dir / "basedpyright"),
"--outputjson",
"--pythonpath",
str(bin_dir / "python"),
],
cwd=cwd,
capture_output=True,
text=True,
@ -208,11 +299,16 @@ def over_ceiling(
)
def environment_fingerprints() -> tuple[str, ...]:
return tuple(
hashlib.sha256(path.read_bytes()).hexdigest()
for path in (PYRIGHT_CONFIG, UV_LOCK)
if path.exists()
def environment_fingerprints(
dep_groups: tuple[str, ...] = TYPECHECK_DEP_GROUPS,
) -> tuple[str, ...]:
return (
*(
hashlib.sha256(path.read_bytes()).hexdigest()
for path in (PYRIGHT_CONFIG, UV_LOCK, PRISMA_SCHEMA)
if path.exists()
),
"groups:" + ",".join(dep_groups),
)
@ -272,16 +368,30 @@ def counts_payload(base_point: str, counts: Mapping[str, int]) -> str:
)
def entry_recency(path: Path) -> float:
try:
return path.stat().st_mtime
except OSError:
return 0.0
def evicted_beyond_cap(entries: Sequence[Path], keep: int) -> tuple[Path, ...]:
newest_first: Final = sorted(entries, key=entry_recency, reverse=True)
return tuple(newest_first[keep:])
def store_counts(
directory: Path, path: Path, base_point: str, counts: Mapping[str, int]
) -> None:
directory.mkdir(parents=True, exist_ok=True)
for stale in directory.glob(f"{CACHE_FILE_PREFIX}*.json"):
if stale != path:
stale.unlink(missing_ok=True)
scratch = scratch_path(path)
scratch.write_text(counts_payload(base_point, counts))
scratch.replace(path)
siblings: Final = tuple(
entry for entry in directory.glob(f"{CACHE_FILE_PREFIX}*.json") if entry != path
)
for stale in evicted_beyond_cap(siblings, CACHE_KEEP_ENTRIES - 1):
stale.unlink(missing_ok=True)
def parse_origin_slug(url: str) -> str | None:
@ -560,6 +670,7 @@ def main() -> None:
parser.add_argument("--update", action="store_true")
parser.add_argument("--emit-counts-dir", type=Path)
args = parser.parse_args()
ensure_typecheck_env()
head = count_basedpyright(run_basedpyright())
if args.emit_counts_dir is not None:
cmd_emit_counts(

13
terraform/provider/CHANGELOG.md
View file

@ -7,13 +7,26 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
## [Unreleased]
## [0.4.0] - 2026-08-06
### Fixed
- **organization**: Send `PATCH` instead of `POST` to `/organization/update` and `/organization/member_update`, matching the methods the LiteLLM proxy serves; organization and organization member updates previously failed with a 405
- **team_member**: Include `role` in the update payload so a role change on an existing `litellm_team_member` is applied instead of being silently dropped
### Changed
- The provider source of truth moved to `terraform/provider/` in [BerriAI/litellm](https://github.com/BerriAI/litellm); this repository is now a release mirror. CI in the monorepo statically audits every endpoint the provider calls against the proxy's OpenAPI schema on every change
- **mcp_server**, **vector_store**: `env` and `litellm_params` are now marked sensitive, so they are redacted from plan/apply output, and they are no longer read back from the API into state — the configured value is authoritative. If the proxy returns values that differ from the configuration, that drift is no longer surfaced on refresh
- Dependency updates: `grpc` and `golang.org/x` modules
## [0.3.0] - 2026-07-13
Released from the mirror repository before the source move was complete; this entry backfills it in the monorepo changelog.
### Added
- **model**: Add optional `pricing_base_model` attribute that sets `model_info.base_model` (the cost-map lookup key) independently of routing. Deployments whose routing name differs from the pricing key (for example Azure Data Zone, routed as `azure/gpt-4.1` but priced via `us/gpt-4.1-2025-04-14`) can now be billed correctly without breaking routing. When unset, behavior is unchanged and `base_model` continues to drive both routing and pricing (#47)
## [0.2.2] - 2026-05-13

2
terraform/provider/docs/resources/model.md
View file

@ -118,6 +118,8 @@ The following arguments are supported:
* `base_model` - (Required) string. The actual model identifier from the provider (e.g., "gpt-4", "claude-2").
* `pricing_base_model` - (Optional) string. A pricing key fed to `model_info.base_model` **independently of routing**. When set, `litellm_params.model` still routes via `base_model`, but LiteLLM looks up cost against this key. Useful when the routing/deployment name differs from the cost-map key — e.g. an Azure deployment routed as `azure/gpt-4.1` whose real tier is Data Zone: set `pricing_base_model = "us/gpt-4.1-2025-04-14"` so it is billed at the Data Zone rate. When unset, `base_model` drives pricing as before.
* `litellm_credential_name` - (Optional) string. Name of a LiteLLM credential to use for this model.
* `tier` - (Optional) string. The usage tier for this model. Valid values are `"free"` or `"paid"`. Default: `"free"`.

8
terraform/provider/litellm/resource_model.go
View file

@ -73,6 +73,14 @@ func resourceLiteLLMModel() *schema.Resource {
Type: schema.TypeString,
Required: true,
},
"pricing_base_model": {
// Optional pricing key fed to model_info.base_model, DECOUPLED
// from routing. When set, litellm_params.model still routes via
// base_model, but cost is looked up against this key (e.g.
// "us/gpt-4.1-2025-04-14" for Azure Data Zone pricing).
Type: schema.TypeString,
Optional: true,
},
"tier": {
Type: schema.TypeString,
Optional: true,

21
terraform/provider/litellm/resource_model_crud.go
View file

@ -68,6 +68,14 @@ func createOrUpdateModel(d *schema.ResourceData, m interface{}, isUpdate bool) e
baseModel := d.Get("base_model").(string)
modelName := fmt.Sprintf("%s/%s", customLLMProvider, baseModel)
// Pricing base_model, decoupled from routing. When pricing_base_model is
// set it feeds model_info.base_model (the cost-lookup key) WITHOUT changing
// the routing string above; otherwise base_model drives pricing as before.
pricingBaseModel := baseModel
if v, ok := d.GetOk("pricing_base_model"); ok && v.(string) != "" {
pricingBaseModel = v.(string)
}
// Generate a UUID for new models
modelID := d.Id()
if !isUpdate {
@ -240,7 +248,7 @@ func createOrUpdateModel(d *schema.ResourceData, m interface{}, isUpdate bool) e
ModelInfo: ModelInfo{
ID: modelID,
DBModel: true,
BaseModel: baseModel,
BaseModel: pricingBaseModel,
Tier: d.Get("tier").(string),
Mode: d.Get("mode").(string),
TeamID: d.Get("team_id").(string),
@ -306,7 +314,16 @@ func resourceLiteLLMModelRead(d *schema.ResourceData, m interface{}) error {
d.Set("rpm", GetIntValue(modelResp.LiteLLMParams.RPM, d.Get("rpm").(int)))
d.Set("model_api_base", GetStringValue(modelResp.LiteLLMParams.APIBase, d.Get("model_api_base").(string)))
d.Set("api_version", GetStringValue(modelResp.LiteLLMParams.APIVersion, d.Get("api_version").(string)))
d.Set("base_model", GetStringValue(modelResp.ModelInfo.BaseModel, d.Get("base_model").(string)))
// base_model / pricing_base_model read-back. When pricing_base_model is
// configured, model_info.base_model holds the PRICING key, so recover the
// routing base_model from state (not returned by the API) and read
// pricing_base_model from model_info.
if pbm, ok := d.GetOk("pricing_base_model"); ok && pbm.(string) != "" {
d.Set("base_model", d.Get("base_model").(string))
d.Set("pricing_base_model", GetStringValue(modelResp.ModelInfo.BaseModel, pbm.(string)))
} else {
d.Set("base_model", GetStringValue(modelResp.ModelInfo.BaseModel, d.Get("base_model").(string)))
}
d.Set("tier", GetStringValue(modelResp.ModelInfo.Tier, d.Get("tier").(string)))
d.Set("mode", GetStringValue(modelResp.ModelInfo.Mode, d.Get("mode").(string)))
d.Set("team_id", GetStringValue(modelResp.ModelInfo.TeamID, d.Get("team_id").(string)))

119
tests/guardrails_tests/test_zscaler_ai_guard.py
View file

@ -396,3 +396,122 @@ async def test_apply_guardrail_block_does_not_log_error(mock_api_call):
mock_logger.error.assert_not_called()
assert exc_info.value.status_code == 400
@pytest.mark.asyncio
async def test_send_request_uses_default_timeout_when_unconfigured():
"""
Regression: unconfigured guardrails must keep the historical 5s timeout.
"""
guardrail = ZscalerAIGuard(api_key="test_key", policy_id=1)
assert guardrail.timeout == 5.0
with patch(
"litellm.proxy.guardrails.guardrail_hooks.zscaler_ai_guard.zscaler_ai_guard.get_async_httpx_client"
) as mock_get_client:
mock_client = Mock()
mock_client.post = AsyncMock(return_value=Mock(status_code=200))
mock_get_client.return_value = mock_client
await guardrail._send_request("http://example.com", {}, {})
assert mock_client.post.call_args.kwargs["timeout"] == 5.0
@pytest.mark.asyncio
async def test_send_request_uses_configured_timeout():
"""
Regression for LIT-5222: a configured timeout must reach the HTTP call.
Before the fix _send_request passed a module-level constant, so a slow
upstream failed at 5s with `Timeout passed=5` no matter what was configured.
"""
guardrail = ZscalerAIGuard(api_key="test_key", policy_id=1, timeout=30)
assert guardrail.timeout == 30
with patch(
"litellm.proxy.guardrails.guardrail_hooks.zscaler_ai_guard.zscaler_ai_guard.get_async_httpx_client"
) as mock_get_client:
mock_client = Mock()
mock_client.post = AsyncMock(return_value=Mock(status_code=200))
mock_get_client.return_value = mock_client
await guardrail._send_request("http://example.com", {}, {})
assert mock_client.post.call_args.kwargs["timeout"] == 30
def test_initialize_guardrail_forwards_configured_timeout():
"""
Regression for LIT-5222: the `timeout` key from config.yaml must survive
initialization. It reaches LitellmParams already, but the initializer used
to drop it before it could reach the guardrail instance.
"""
from litellm.proxy.guardrails.guardrail_hooks.zscaler_ai_guard import (
initialize_guardrail,
)
from litellm.types.guardrails import LitellmParams
litellm_params = LitellmParams(
guardrail="zscaler_ai_guard",
mode="pre_call",
api_key="test_key",
api_base="http://example.com",
policy_id=1,
timeout="30",
)
guardrail = initialize_guardrail(
litellm_params, {"guardrail_name": "zscaler-configured-timeout"}
)
assert guardrail.timeout == 30.0
def test_config_model_exposes_timeout_to_dashboard():
"""
The dashboard guardrail form is built from get_config_model(), so the field
has to be declared there for the setting to be reachable outside config.yaml.
"""
config_model = ZscalerAIGuard.get_config_model()
assert config_model is not None
assert "timeout" in config_model.model_fields
@pytest.mark.parametrize("bad_timeout", [0, -1])
def test_non_positive_timeout_falls_back_to_default(bad_timeout):
"""
Regression: httpx rejects a negative timeout and treats 0 as "fail
immediately", so a non-positive value would break every scan instead of
relaxing the limit the operator was trying to raise.
"""
guardrail = ZscalerAIGuard(api_key="test_key", policy_id=1, timeout=bad_timeout)
assert guardrail.timeout == 5.0
def test_update_in_memory_litellm_params_keeps_timeout_resolved():
"""
Regression: the base implementation copies every LitellmParams attribute
onto the guardrail, so an unset timeout would overwrite the resolved value
with None and silently fall back to the shared client's 600s default.
"""
from litellm.types.guardrails import LitellmParams
guardrail = ZscalerAIGuard(api_key="test_key", policy_id=1, timeout=30)
assert guardrail.timeout == 30
guardrail.update_in_memory_litellm_params(
LitellmParams(guardrail="zscaler_ai_guard", mode="pre_call", api_key="test_key")
)
assert guardrail.timeout == 5.0
guardrail.update_in_memory_litellm_params(
LitellmParams(
guardrail="zscaler_ai_guard", mode="pre_call", api_key="test_key", timeout=45
)
)
assert guardrail.timeout == 45.0

72
tests/litellm/llms/anthropic/test_anthropic_schema_filter.py
View file

@ -281,7 +281,10 @@ class TestFilterAnthropicOutputSchema:
"unevaluatedProperties",
):
assert field not in result
assert 'properties whose names match each pattern must satisfy: {"^x": {"type": "string"}}' in result["description"]
assert (
'properties whose names match each pattern must satisfy: {"^x": {"type": "string"}}'
in result["description"]
)
assert 'property names must satisfy: {"pattern": "^[a-z]+$"}' in result["description"]
assert 'dependent required properties: {"first": ["last"]}' in result["description"]
assert 'dependent schemas: {"first": {"required": ["last"]}}' in result["description"]
@ -347,3 +350,70 @@ class TestFilterAnthropicOutputSchema:
"all array items must be unique, minimum number of matching items: 2, "
"maximum number of matching items: 3."
)
def test_coerces_explicit_additional_properties_true(self):
"""An explicit ``additionalProperties: true`` must be coerced to false.
Anthropic rejects anything other than false with:
"output_format.schema: For 'object' type, 'additionalProperties: true' is
not supported".
"""
schema = {
"type": "object",
"additionalProperties": True,
"properties": {"a": {"type": "string"}},
}
result = AnthropicConfig.filter_anthropic_output_schema(schema)
assert result["additionalProperties"] is False
def test_coerces_additional_properties_true_when_nested(self):
"""Nested object schemas are coerced too, at every recursion site."""
schema = {
"type": "object",
"properties": {
"obj": {
"type": "object",
"additionalProperties": True,
"properties": {"a": {"type": "string"}},
},
"rows": {
"type": "array",
"items": {
"type": "object",
"additionalProperties": True,
"properties": {"b": {"type": "string"}},
},
},
},
}
result = AnthropicConfig.filter_anthropic_output_schema(schema)
assert result["properties"]["obj"]["additionalProperties"] is False
assert result["properties"]["rows"]["items"]["additionalProperties"] is False
def test_coerces_additional_properties_sub_schema(self):
"""A sub-schema value (free-form map) is also rejected by Anthropic."""
schema = {
"type": "object",
"additionalProperties": {"type": "string"},
"properties": {"a": {"type": "string"}},
}
result = AnthropicConfig.filter_anthropic_output_schema(schema)
assert result["additionalProperties"] is False
def test_explicit_additional_properties_false_is_preserved(self):
"""The already-correct value must survive untouched."""
schema = {
"type": "object",
"additionalProperties": False,
"properties": {"a": {"type": "string"}},
}
result = AnthropicConfig.filter_anthropic_output_schema(schema)
assert result["additionalProperties"] is False

7
tests/openai_endpoints_tests/test_openai_batches_endpoint.py
View file

@ -414,7 +414,7 @@ async def test_batch_status_sync_from_provider_to_database():
# Verify logger was called with status change message
mock_logger.info.assert_called()
log_message = mock_logger.info.call_args[0][0]
log_message = mock_logger.info.call_args[0][0] % mock_logger.info.call_args[0][1:]
assert "validating" in log_message
assert "completed" in log_message
@ -450,6 +450,9 @@ async def test_batch_cancel_updates_database():
# Mock prisma client
mock_prisma_client = MagicMock()
mock_prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock(
return_value=None
)
mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock()
# Mock managed_files_obj
@ -482,7 +485,7 @@ async def test_batch_cancel_updates_database():
# Verify logger was called
mock_logger.info.assert_called()
log_message = mock_logger.info.call_args[0][0]
log_message = mock_logger.info.call_args[0][0] % mock_logger.info.call_args[0][1:]
assert "cancel" in log_message.lower()
assert "cancelled" in log_message

117
tests/pass_through_unit_tests/test_pass_through_unit_tests.py
View file

@ -30,6 +30,7 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
)
from fastapi import Request
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
_update_metadata_with_tags_in_header,
HttpPassThroughEndpointHelpers,
@ -652,3 +653,119 @@ def test_custom_pricing_used_in_cost_calculation():
print(f"Cache-aware cost: {cache_cost}")
print("✅ Custom pricing parameters are correctly used in cost calculation")
def test_init_kwargs_client_metadata_cannot_spoof_authenticated_identity(
mock_request, mock_user_api_key_dict
):
request = mock_request()
passthrough_payload = PassthroughStandardLoggingPayload(
url="https://test.com",
request_body={},
)
authenticated_key = UserAPIKeyAuth(
api_key="test-key",
user_id="test-user",
team_id="test-team",
end_user_id="test-user",
key_alias="real-key",
team_alias="Real Team",
user_email="real@example.com",
org_id="real-org",
)
result = HttpPassThroughEndpointHelpers._init_kwargs_for_pass_through_endpoint(
request=request,
user_api_key_dict=authenticated_key,
passthrough_logging_payload=passthrough_payload,
litellm_call_id="test-call-id",
logging_obj=LiteLLMLoggingObj(
model="test-model",
messages=[],
stream=False,
call_type="test-call-type",
start_time=datetime.now(),
litellm_call_id="test-call-id",
function_id="test-function-id",
),
_parsed_body={
"litellm_metadata": {
"user_api_key_org_id": "victim-org",
"user_api_key_end_user_id": "victim-end-user",
"user_api_key_user_id": "victim-user",
"user_api_key_team_id": "victim-team",
"user_api_key_team_alias": "Victim Team",
"user_api_key_alias": "victim-key",
"user_api_key_user_email": "victim@example.com",
}
},
)
metadata = result["litellm_params"]["metadata"]
assert metadata["user_api_key_user_id"] == "test-user"
assert metadata["user_api_key_team_id"] == "test-team"
assert metadata["user_api_key_team_alias"] == "Real Team"
assert metadata["user_api_key_alias"] == "real-key"
assert metadata["user_api_key_user_email"] == "real@example.com"
assert metadata["user_api_key_org_id"] == "real-org"
assert metadata["user_api_key_end_user_id"] == "test-user"
def test_init_kwargs_no_authenticated_identity_field_is_client_settable(
mock_request, mock_user_api_key_dict
):
authenticated_key = UserAPIKeyAuth(
api_key="test-key",
user_id="test-user",
team_id="test-team",
end_user_id="test-end-user",
key_alias="real-key",
team_alias="Real Team",
user_email="real@example.com",
org_id="real-org",
organization_alias="Real Org",
project_id="real-project",
project_alias="Real Project",
spend=1.5,
max_budget=10.0,
user_spend=2.5,
user_max_budget=20.0,
team_spend=3.5,
team_max_budget=30.0,
metadata={"real": "auth-metadata"},
)
expected = dict(
LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(
user_api_key_dict=authenticated_key
)
)
assert len(expected) >= 20
spoofed = {key: f"SPOOFED-{key}" for key in expected}
result = HttpPassThroughEndpointHelpers._init_kwargs_for_pass_through_endpoint(
request=mock_request(),
user_api_key_dict=authenticated_key,
passthrough_logging_payload=PassthroughStandardLoggingPayload(
url="https://test.com", request_body={}
),
litellm_call_id="test-call-id",
logging_obj=LiteLLMLoggingObj(
model="test-model",
messages=[],
stream=False,
call_type="test-call-type",
start_time=datetime.now(),
litellm_call_id="test-call-id",
function_id="test-function-id",
),
_parsed_body={"litellm_metadata": dict(spoofed), "metadata": dict(spoofed)},
)
metadata = result["litellm_params"]["metadata"]
survived = {
key: metadata.get(key)
for key in expected
if metadata.get(key) != expected[key]
}
assert survived == {}, f"client-supplied values survived for: {sorted(survived)}"

279
tests/proxy_unit_tests/test_check_batch_cost.py
View file

@ -420,6 +420,134 @@ class TestCheckBatchCost:
), "update() must include batch_processed=True when column is present"
assert update_data["status"] == "complete"
@pytest.mark.asyncio
async def test_completed_batch_with_no_attributable_owner_still_writes_spend_log(
self, check_batch_cost_instance, mock_prisma_client, mock_llm_router
):
"""Regression: a batch created with the master key or a team-less key has
created_by=None and team_id=None on LiteLLM_ManagedObjectTable (the table
never stores the raw key hash). CheckBatchCost's synthetic logging_obj for
such a batch then carries no attributable key/user/team/end-user, and
before the fix _should_track_cost_callback silently skipped the DB write
with no error or warning: batch_processed still became True, but no
LiteLLM_SpendLogs row was ever written.
Unlike the other tests in this file, this one does NOT mock
litellm_logging.Logging or async_success_handler -- it runs the real
logging pipeline through to _ProxyDBLogger, which is the exact gap that
let the original bug ship undetected.
"""
import litellm
from litellm.proxy.hooks.proxy_track_cost_callback import _ProxyDBLogger
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0)
mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock()
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None)
mock_job = MagicMock()
mock_job.id = "job-unattributed-1"
mock_job.unified_object_id = "dW5pZmllZF9iYXRjaF9pZA=="
mock_job.created_by = None
mock_job.team_id = None
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job])
# A real LiteLLMBatch (not a bare MagicMock): this test runs the real
# litellm_logging.Logging pipeline, which type-checks the result via
# isinstance(..., LiteLLMBatch) before it will compute/attach a cost.
from litellm.types.utils import LiteLLMBatch
mock_response = LiteLLMBatch(
id="batch-1",
completion_window="24h",
created_at=1,
endpoint="/v1/chat/completions",
input_file_id="file-input-123",
object="batch",
status="completed",
output_file_id="file-output-123",
)
mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response)
mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"})
mock_deployment = MagicMock()
mock_deployment.litellm_params.custom_llm_provider = "openai"
mock_deployment.litellm_params.model = "gpt-4"
mock_deployment.model_info.model_dump.return_value = {}
mock_llm_router.get_deployment = MagicMock(return_value=mock_deployment)
mock_file_content = MagicMock()
mock_file_content.content = b'{"id":"req-1"}'
decoded_id = "llm_model_id,model-123;llm_batch_id,batch-456;"
db_logger = _ProxyDBLogger()
mock_update_database = AsyncMock()
# Unlike the other tests in this file, this one runs the real
# litellm_logging.Logging pipeline, which calls
# _is_base64_encoded_unified_file_id an extra time (checking result.id
# after it's reset to job.unified_object_id). Key off the argument
# instead of a fixed-length side_effect list so the exact call count
# doesn't matter.
def _fake_is_base64_encoded(file_id):
return decoded_id if file_id == mock_job.unified_object_id else None
with (
patch(
"litellm.proxy.openai_files_endpoints.common_utils._is_base64_encoded_unified_file_id",
side_effect=_fake_is_base64_encoded,
),
patch(
"litellm.proxy.openai_files_endpoints.common_utils.get_model_id_from_unified_batch_id",
return_value="model-123",
),
patch(
"litellm.proxy.openai_files_endpoints.common_utils.get_batch_id_from_unified_batch_id",
return_value="batch-456",
),
patch(
"litellm.files.main.afile_content",
new_callable=AsyncMock,
return_value=mock_file_content,
),
patch(
"litellm.batches.batch_utils._get_file_content_as_dictionary",
return_value=[{"id": "req-1"}],
),
patch(
"litellm.batches.batch_utils.calculate_batch_cost_and_usage",
new_callable=AsyncMock,
return_value=(
0.01,
{"prompt_tokens": 10, "completion_tokens": 5},
["gpt-4"],
),
),
patch(
"litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider",
return_value=("gpt-4", "openai", None, None),
),
patch.object(litellm, "_async_success_callback", [db_logger]),
patch(
"litellm.proxy.proxy_server.proxy_logging_obj",
MagicMock(
db_spend_update_writer=MagicMock(update_database=mock_update_database),
slack_alerting_instance=MagicMock(customer_spend_alert=AsyncMock()),
),
),
patch("litellm.proxy.proxy_server.increment_spend_counters", AsyncMock()),
patch("litellm.proxy.proxy_server.update_cache", AsyncMock()),
):
await check_batch_cost_instance.check_batch_cost()
mock_update_database.assert_awaited_once()
assert mock_update_database.call_args.kwargs["response_cost"] == 0.01
assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, (
"the job must still be marked processed once cost tracking succeeds"
)
@pytest.mark.asyncio
async def test_cost_tracking_failure_leaves_job_unprocessed(
self, check_batch_cost_instance, mock_prisma_client, mock_llm_router
@ -499,7 +627,7 @@ class TestCheckBatchCost:
must be written back with that status and batch_processed=True so it stops being
polled forever.
"""
from unittest.mock import patch
import base64
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
return_value=0
@ -511,7 +639,9 @@ class TestCheckBatchCost:
mock_job = MagicMock()
mock_job.id = "job-terminal-1"
mock_job.unified_object_id = "dW5pZmllZF9iYXRjaF9pZA=="
mock_job.unified_object_id = base64.urlsafe_b64encode(
b"litellm_proxy;model_id:model-123;llm_batch_id:batch-456"
).decode()
mock_job.created_by = "user-1"
assert check_batch_cost_instance._has_batch_processed_column is True
@ -527,23 +657,7 @@ class TestCheckBatchCost:
mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response)
decoded_id = "llm_model_id,model-123;llm_batch_id,batch-456;"
with (
patch(
"litellm.proxy.openai_files_endpoints.common_utils._is_base64_encoded_unified_file_id",
side_effect=[decoded_id, None],
),
patch(
"litellm.proxy.openai_files_endpoints.common_utils.get_model_id_from_unified_batch_id",
return_value="model-123",
),
patch(
"litellm.proxy.openai_files_endpoints.common_utils.get_batch_id_from_unified_batch_id",
return_value="batch-456",
),
):
await check_batch_cost_instance.check_batch_cost()
await check_batch_cost_instance.check_batch_cost()
assert (
mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1
@ -556,6 +670,133 @@ class TestCheckBatchCost:
update_data["batch_processed"] is True
), "terminal-status update() must set batch_processed=True so polling stops"
@pytest.mark.asyncio
@pytest.mark.parametrize("terminal_status", ["failed", "expired", "cancelled"])
async def test_terminal_status_persists_managed_output_file_ids(
self,
check_batch_cost_instance,
mock_prisma_client,
mock_llm_router,
terminal_status,
):
"""A cancelled/failed/expired batch with provider output files must be persisted
with unified managed file IDs, never raw provider IDs. Raw IDs written here leak
to every later GET /batches/{id} and GET /batches because the terminal row is
final (batch_processed=True) and read paths only resolve, never mint.
"""
import base64
import json
from litellm.types.utils import LiteLLMBatch
unified_batch_uid = base64.urlsafe_b64encode(
b"litellm_proxy;model_id:model-123;llm_batch_id:batch-456"
).decode()
raw_output_file_id = "file-terminal-out-abc"
raw_error_file_id = "file-terminal-err-xyz"
raw_input_file_id = "file-terminal-in-123"
unified_input_file_id = base64.urlsafe_b64encode(
b"litellm_proxy:application/octet-stream;unified_id,in-1;target_model_names,gpt-5-batch"
).decode()
unified_output_file_id = base64.urlsafe_b64encode(
f"litellm_proxy:application/octet-stream;unified_id,u-1;llm_output_file_id,{raw_output_file_id}".encode()
).decode()
unified_error_file_id = base64.urlsafe_b64encode(
f"litellm_proxy:application/octet-stream;unified_id,u-2;llm_output_file_id,{raw_error_file_id}".encode()
).decode()
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
return_value=0
)
mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock()
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(
return_value=None
)
input_file_row = MagicMock()
input_file_row.unified_file_id = unified_input_file_id
def find_managed_file(where):
if where["flat_model_file_ids"]["has"] == raw_input_file_id:
return input_file_row
return None
mock_prisma_client.db.litellm_managedfiletable.find_first = AsyncMock(
side_effect=find_managed_file
)
mock_job = MagicMock()
mock_job.id = "job-terminal-mint-1"
mock_job.unified_object_id = unified_batch_uid
mock_job.created_by = "user-1"
mock_job.team_id = "team-1"
check_batch_cost_instance._has_batch_processed_column = True
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
return_value=[mock_job]
)
response = LiteLLMBatch(
id="batch-456",
completion_window="24h",
created_at=1,
endpoint="/v1/chat/completions",
input_file_id=raw_input_file_id,
object="batch",
status=terminal_status,
output_file_id=raw_output_file_id,
error_file_id=raw_error_file_id,
)
mock_llm_router.aretrieve_batch = AsyncMock(return_value=response)
mock_hook = MagicMock()
mock_hook.get_unified_output_file_id.side_effect = [
unified_output_file_id,
unified_error_file_id,
]
mock_hook.store_unified_file_id = AsyncMock()
check_batch_cost_instance.proxy_logging_obj.get_proxy_hook.return_value = (
mock_hook
)
await check_batch_cost_instance.check_batch_cost()
mock_hook.get_unified_output_file_id.assert_any_call(
output_file_id=raw_output_file_id,
model_id="model-123",
model_name="gpt-5-batch",
)
mock_hook.get_unified_output_file_id.assert_any_call(
output_file_id=raw_error_file_id,
model_id="model-123",
model_name="gpt-5-batch",
)
stored = {
next(iter(c.kwargs["model_mappings"].values())): c.kwargs["file_id"]
for c in mock_hook.store_unified_file_id.call_args_list
}
assert stored == {
raw_output_file_id: unified_output_file_id,
raw_error_file_id: unified_error_file_id,
}
for store_call in mock_hook.store_unified_file_id.call_args_list:
assert store_call.kwargs["user_api_key_dict"].user_id == "user-1"
assert store_call.kwargs["user_api_key_dict"].team_id == "team-1"
assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1
update_call = mock_prisma_client.db.litellm_managedobjecttable.update.call_args
assert update_call.kwargs["where"] == {"id": "job-terminal-mint-1"}
update_data = update_call.kwargs["data"]
assert update_data["status"] == terminal_status
assert update_data["batch_processed"] is True
persisted = json.loads(update_data["file_object"])
assert persisted["id"] == unified_batch_uid
assert persisted["input_file_id"] == unified_input_file_id
assert persisted["output_file_id"] == unified_output_file_id
assert persisted["error_file_id"] == unified_error_file_id
assert raw_output_file_id not in update_data["file_object"]
assert raw_error_file_id not in update_data["file_object"]
@pytest.mark.asyncio
async def test_raw_output_file_id_converted_to_managed_id(
self, check_batch_cost_instance, mock_prisma_client, mock_llm_router

275
tests/proxy_unit_tests/test_check_responses_cost.py
View file

@ -449,6 +449,281 @@ class TestCheckResponsesCost:
assert "job-3" in completion_call[1]["where"]["id"]["in"]
assert "job-2" not in completion_call[1]["where"]["id"]["in"]
@pytest.mark.asyncio
async def test_encoded_response_id_is_fetched_through_router(
self, check_responses_cost_instance, mock_prisma_client, mock_llm_router
):
"""
Regression test for https://github.com/BerriAI/litellm/issues/35131
A background response created against a deployment whose credentials only
exist in the config (e.g. Azure api_base/api_key) must be fetched through
the router so the deployment credentials are applied. Calling
litellm.aget_responses directly only sees provider env vars, fails, and
leaves the row in "queued" forever.
"""
from litellm.responses.utils import ResponsesAPIRequestUtils
encoded_response_id = ResponsesAPIRequestUtils._build_responses_api_response_id(
custom_llm_provider="azure",
model_id="deployment-abc",
response_id="resp_upstream_123",
)
mock_job = MagicMock()
mock_job.unified_object_id = encoded_response_id
mock_job.created_by = "test-user"
mock_job.id = "job-router"
mock_job.file_object = {"model": "azure-gpt-5", "id": encoded_response_id}
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
return_value=[mock_job]
)
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
return_value=0
)
mock_llm_router.aget_responses = AsyncMock(
return_value=ResponsesAPIResponse(
id=encoded_response_id,
object="response",
status="completed",
created_at=int(datetime.now().timestamp()),
output=[],
usage=ResponseAPIUsage(
input_tokens=100, output_tokens=50, total_tokens=150
),
)
)
with patch(
"litellm.aget_responses",
new_callable=AsyncMock,
side_effect=AssertionError(
"must not bypass the router for a deployment-scoped response id"
),
) as mock_sdk_aget:
await check_responses_cost_instance.check_responses_cost()
mock_sdk_aget.assert_not_called()
assert (
mock_llm_router.aget_responses.call_args[1]["response_id"]
== encoded_response_id
)
calls = (
mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
)
assert len(calls) == 1
assert calls[0][1]["data"]["status"] == "completed"
assert calls[0][1]["where"]["id"]["in"] == ["job-router"]
@pytest.mark.asyncio
async def test_encrypted_response_id_is_fetched_through_router(
self, check_responses_cost_instance, mock_prisma_client, mock_llm_router, monkeypatch
):
"""
Rows store the *encrypted* response id when responses id security is on.
After decryption the id still carries the deployment model_id, so the
fetch must go through the router (issue #35131).
"""
from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.types.utils import SpecialEnums
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-test-salt-key-for-response-ids")
encoded_response_id = ResponsesAPIRequestUtils._build_responses_api_response_id(
custom_llm_provider="openai",
model_id="deployment-xyz",
response_id="resp_upstream_456",
)
encrypted_response_id = "resp_" + str(
encrypt_value_helper(
value=SpecialEnums.LITELLM_MANAGED_RESPONSE_API_RESPONSE_ID_COMPLETE_STR.value.format(
encoded_response_id, "test-user", "test-team"
)
)
)
mock_job = MagicMock()
mock_job.unified_object_id = encrypted_response_id
mock_job.created_by = "test-user"
mock_job.id = "job-encrypted"
mock_job.file_object = {"model": "gpt-5", "id": encrypted_response_id}
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
return_value=[mock_job]
)
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
return_value=0
)
mock_llm_router.aget_responses = AsyncMock(
return_value=ResponsesAPIResponse(
id=encoded_response_id,
object="response",
status="completed",
created_at=int(datetime.now().timestamp()),
output=[],
usage=None,
)
)
with patch(
"litellm.aget_responses",
new_callable=AsyncMock,
side_effect=AssertionError(
"must not bypass the router for a deployment-scoped response id"
),
) as mock_sdk_aget:
await check_responses_cost_instance.check_responses_cost()
mock_sdk_aget.assert_not_called()
assert (
mock_llm_router.aget_responses.call_args[1]["response_id"]
== encoded_response_id
)
calls = (
mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
)
assert len(calls) == 1
assert calls[0][1]["where"]["id"]["in"] == ["job-encrypted"]
@pytest.mark.asyncio
async def test_response_id_without_model_id_uses_sdk(
self, check_responses_cost_instance, mock_prisma_client, mock_llm_router
):
"""Ids that carry no deployment info can't be routed, so fall back to the SDK."""
mock_job = MagicMock()
mock_job.unified_object_id = "resp_plain_upstream_id"
mock_job.created_by = "test-user"
mock_job.id = "job-plain"
mock_job.file_object = {"model": "gpt-5", "id": "resp_plain_upstream_id"}
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
return_value=[mock_job]
)
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
return_value=0
)
mock_llm_router.aget_responses = AsyncMock(
side_effect=AssertionError("router cannot route an id without a model_id")
)
mock_response = ResponsesAPIResponse(
id="resp_plain_upstream_id",
object="response",
status="completed",
created_at=int(datetime.now().timestamp()),
output=[],
usage=None,
)
with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_sdk_aget:
mock_sdk_aget.return_value = mock_response
await check_responses_cost_instance.check_responses_cost()
mock_sdk_aget.assert_called_once()
mock_llm_router.aget_responses.assert_not_called()
@pytest.mark.asyncio
async def test_missing_deployment_falls_back_to_sdk(
self, check_responses_cost_instance, mock_prisma_client, mock_llm_router
):
"""
An encoded id whose deployment was removed from the router must fall back
to the SDK so provider env credentials can still retrieve it, instead of
failing every poll cycle until stale expiration.
"""
from litellm.responses.utils import ResponsesAPIRequestUtils
encoded_response_id = ResponsesAPIRequestUtils._build_responses_api_response_id(
custom_llm_provider="openai",
model_id="deployment-deleted",
response_id="resp_upstream_789",
)
mock_job = MagicMock()
mock_job.unified_object_id = encoded_response_id
mock_job.created_by = "test-user"
mock_job.id = "job-missing-deployment"
mock_job.file_object = {"model": "gpt-5", "id": encoded_response_id}
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
return_value=[mock_job]
)
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
return_value=0
)
mock_llm_router.get_deployment = MagicMock(return_value=None)
mock_llm_router.aget_responses = AsyncMock(
side_effect=AssertionError("router has no deployment for this model_id")
)
mock_response = ResponsesAPIResponse(
id=encoded_response_id,
object="response",
status="completed",
created_at=int(datetime.now().timestamp()),
output=[],
usage=None,
)
with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_sdk_aget:
mock_sdk_aget.return_value = mock_response
await check_responses_cost_instance.check_responses_cost()
mock_llm_router.get_deployment.assert_called_once_with(model_id="deployment-deleted")
mock_llm_router.aget_responses.assert_not_called()
mock_sdk_aget.assert_called_once()
assert mock_sdk_aget.call_args[1]["response_id"] == encoded_response_id
calls = (
mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
)
assert len(calls) == 1
assert calls[0][1]["data"]["status"] == "completed"
assert calls[0][1]["where"]["id"]["in"] == ["job-missing-deployment"]
@pytest.mark.asyncio
async def test_check_responses_cost_with_incomplete_response(
self, check_responses_cost_instance, mock_prisma_client
):
"""'incomplete' is terminal in the Responses API, so the row must not stay queued."""
mock_job = MagicMock()
mock_job.unified_object_id = "resp_test_incomplete"
mock_job.created_by = "test-user"
mock_job.id = "job-incomplete"
mock_job.file_object = {"model": "gpt-5", "id": "resp_test_incomplete"}
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
return_value=[mock_job]
)
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
return_value=0
)
mock_response = ResponsesAPIResponse(
id="resp_incomplete",
object="response",
status="incomplete",
created_at=int(datetime.now().timestamp()),
output=[],
usage=None,
)
with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_aget:
mock_aget.return_value = mock_response
await check_responses_cost_instance.check_responses_cost()
calls = (
mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
)
assert len(calls) == 1
assert calls[0][1]["data"]["status"] == "completed"
assert calls[0][1]["where"]["id"]["in"] == ["job-incomplete"]
@pytest.mark.asyncio
async def test_check_responses_cost_no_model_in_file_object(
self, check_responses_cost_instance, mock_prisma_client

122
tests/proxy_unit_tests/test_proxy_config_unit_test.py
View file

@ -15,9 +15,7 @@ import os
# this file is to test litellm/proxy
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
sys.path.insert(0, os.path.abspath("../..")) # Adds the parent directory to the system path
import asyncio
import logging
@ -88,25 +86,14 @@ async def test_read_config_file_with_os_environ_vars():
# Read config
proxy_config_instance = ProxyConfig()
current_path = os.path.dirname(os.path.abspath(__file__))
config_path = os.path.join(
current_path, "example_config_yaml", "config_with_env_vars.yaml"
)
config_path = os.path.join(current_path, "example_config_yaml", "config_with_env_vars.yaml")
config = await proxy_config_instance.get_config(config_file_path=config_path)
print(config)
# Add assertions
assert (
config["litellm_settings"]["default_internal_user_params"]["user_role"]
== "admin"
)
assert (
config["litellm_settings"]["s3_callback_params"]["s3_aws_access_key_id"]
== "1234567890"
)
assert (
config["litellm_settings"]["s3_callback_params"]["s3_aws_secret_access_key"]
== "1234567890"
)
assert config["litellm_settings"]["default_internal_user_params"]["user_role"] == "admin"
assert config["litellm_settings"]["s3_callback_params"]["s3_aws_access_key_id"] == "1234567890"
assert config["litellm_settings"]["s3_callback_params"]["s3_aws_secret_access_key"] == "1234567890"
for model in config["model_list"]:
if "azure" in model["litellm_params"]["model"]:
@ -129,17 +116,13 @@ async def test_basic_include_directive():
"""
proxy_config_instance = ProxyConfig()
current_path = os.path.dirname(os.path.abspath(__file__))
config_path = os.path.join(
current_path, "example_config_yaml", "config_with_include.yaml"
)
config_path = os.path.join(current_path, "example_config_yaml", "config_with_include.yaml")
config = await proxy_config_instance.get_config(config_file_path=config_path)
# Verify the included model list was merged
assert len(config["model_list"]) > 0
assert any(
model["model_name"] == "included-model" for model in config["model_list"]
)
assert any(model["model_name"] == "included-model" for model in config["model_list"])
# Verify original config settings remain
assert config["litellm_settings"]["callbacks"] == ["prometheus"]
@ -152,9 +135,7 @@ async def test_missing_include_file():
"""
proxy_config_instance = ProxyConfig()
current_path = os.path.dirname(os.path.abspath(__file__))
config_path = os.path.join(
current_path, "example_config_yaml", "config_with_missing_include.yaml"
)
config_path = os.path.join(current_path, "example_config_yaml", "config_with_missing_include.yaml")
with pytest.raises(FileNotFoundError):
await proxy_config_instance.get_config(config_file_path=config_path)
@ -167,20 +148,14 @@ async def test_multiple_includes():
"""
proxy_config_instance = ProxyConfig()
current_path = os.path.dirname(os.path.abspath(__file__))
config_path = os.path.join(
current_path, "example_config_yaml", "config_with_multiple_includes.yaml"
)
config_path = os.path.join(current_path, "example_config_yaml", "config_with_multiple_includes.yaml")
config = await proxy_config_instance.get_config(config_file_path=config_path)
# Verify models from both included files are present
assert len(config["model_list"]) == 2
assert any(
model["model_name"] == "included-model-1" for model in config["model_list"]
)
assert any(
model["model_name"] == "included-model-2" for model in config["model_list"]
)
assert any(model["model_name"] == "included-model-1" for model in config["model_list"])
assert any(model["model_name"] == "included-model-2" for model in config["model_list"])
# Verify original config settings remain
assert config["litellm_settings"]["callbacks"] == ["prometheus"]
@ -211,8 +186,7 @@ def test_add_callbacks_from_db_config():
# 1 instance of LangfusePromptManagement should exist in litellm.success_callback
num_langfuse_instances = sum(
isinstance(callback, LangfusePromptManagement)
for callback in litellm.success_callback
isinstance(callback, LangfusePromptManagement) for callback in litellm.success_callback
)
assert num_langfuse_instances == 1
assert len(litellm.success_callback) == 2
@ -290,9 +264,7 @@ async def test_json_logs_calls_turn_on_json():
"litellm_settings": {"json_logs": True},
}
with tempfile.NamedTemporaryFile(
mode="w", suffix=".yaml", delete=False
) as temp_file:
with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as temp_file:
yaml.dump(config_content, temp_file)
temp_file_path = temp_file.name
@ -316,3 +288,71 @@ async def test_json_logs_calls_turn_on_json():
# Cleanup
os.unlink(temp_file_path)
litellm.json_logs = False
class TestYamlStorePromptsDbOverride:
"""
Test that YAML store_prompts_in_spend_logs takes precedence over DB-cached value.
When store_model_in_db=true, LiteLLM persists general_settings to the DB.
On periodic reloads, _update_general_settings() must NOT override
YAML-explicit values with stale DB values.
"""
def _make_proxy_config_with_yaml_keys(self, yaml_keys: set) -> "ProxyConfig":
"""Helper: create ProxyConfig with pre-populated _yaml_general_settings_keys."""
proxy_config = ProxyConfig()
proxy_config._yaml_general_settings_keys = yaml_keys
return proxy_config
@pytest.mark.asyncio
async def test_yaml_value_takes_precedence_over_db(self):
"""When YAML sets store_prompts_in_spend_logs=false, DB value (true) should be ignored."""
proxy_config = self._make_proxy_config_with_yaml_keys({"store_prompts_in_spend_logs"})
test_general_settings = {"store_prompts_in_spend_logs": False}
with mock.patch("litellm.proxy.proxy_server.general_settings", test_general_settings):
await proxy_config._update_general_settings(
db_general_settings={"store_prompts_in_spend_logs": True},
)
assert test_general_settings["store_prompts_in_spend_logs"] is False
@pytest.mark.asyncio
async def test_db_value_used_when_yaml_does_not_set_key(self):
"""When YAML does NOT set store_prompts_in_spend_logs, DB value should be used."""
proxy_config = self._make_proxy_config_with_yaml_keys({"master_key", "database_url"})
test_general_settings = {"master_key": "sk-test"}
with mock.patch("litellm.proxy.proxy_server.general_settings", test_general_settings):
await proxy_config._update_general_settings(
db_general_settings={"store_prompts_in_spend_logs": True},
)
assert test_general_settings["store_prompts_in_spend_logs"] is True
@pytest.mark.asyncio
async def test_admin_ui_change_works_when_yaml_omits_key(self):
"""Admin UI change (DB update) should work when YAML doesn't set the key."""
proxy_config = self._make_proxy_config_with_yaml_keys({"master_key"})
test_general_settings = {"master_key": "sk-test"}
with mock.patch("litellm.proxy.proxy_server.general_settings", test_general_settings):
await proxy_config._update_general_settings(
db_general_settings={"store_prompts_in_spend_logs": True},
)
assert test_general_settings["store_prompts_in_spend_logs"] is True
await proxy_config._update_general_settings(
db_general_settings={"store_prompts_in_spend_logs": False},
)
assert test_general_settings["store_prompts_in_spend_logs"] is False
def test_yaml_general_settings_keys_populated_on_load(self):
"""_yaml_general_settings_keys should be empty on init."""
proxy_config = ProxyConfig()
assert proxy_config._yaml_general_settings_keys == set()

2
tests/router_unit_tests/test_router_helper_utils.py
View file

@ -15,6 +15,7 @@ from unittest.mock import patch, MagicMock, AsyncMock
from create_mock_standard_logging_payload import create_standard_logging_payload
from litellm.types.utils import StandardLoggingPayload
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS
@pytest.fixture
@ -1816,6 +1817,7 @@ def test_init_auto_router_deployment_success(mock_auto_router, model_list):
default_model="gpt-5-mini",
embedding_model="text-embedding-3-small",
litellm_router_instance=router,
max_input_chars=DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS,
)
# Verify the auto-router was added to the router's auto_routers dict

9
tests/test_litellm/conftest.py
View file

@ -19,6 +19,7 @@ sys.path.insert(
import asyncio
import litellm
from litellm import utils as litellm_utils_module
from litellm._logging import ALL_LOGGERS
from litellm.litellm_core_utils.prompt_templates import (
image_handling as image_handling_module,
@ -238,6 +239,11 @@ def isolate_litellm_state():
if hasattr(litellm, _attr):
original_state[_attr] = getattr(litellm, _attr)
original_runtime_registered_model_cost = {
model_key: dict(model_value)
for model_key, model_value in litellm_utils_module._runtime_registered_model_cost.items()
}
# Store LiteLLM logger state. Some tests reconfigure handlers/propagation for
# JSON logging and do not restore them, which breaks later caplog-based tests.
logger_state = {}
@ -304,6 +310,9 @@ def isolate_litellm_state():
if hasattr(litellm, attr_name):
setattr(litellm, attr_name, original_value)
litellm_utils_module._runtime_registered_model_cost.clear()
litellm_utils_module._runtime_registered_model_cost.update(original_runtime_registered_model_cost)
# Restore logger configuration mutated by logging-focused tests.
for logger in ALL_LOGGERS:
original_logger_state = logger_state.get(logger.name)

100
tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py
View file

@ -45,9 +45,10 @@ def _build_managed_files_mock(unified_id: str = "file-bWFuYWdlZF9vdXRwdXRfaWQ=")
return mock
def _build_prisma_mock():
def _build_prisma_mock(db_batch_object=None):
mock = MagicMock()
mock.db.litellm_managedfiletable.find_first = AsyncMock(return_value=None)
mock.db.litellm_managedobjecttable.find_first = AsyncMock(return_value=db_batch_object)
mock.db.litellm_managedobjecttable.update = AsyncMock()
return mock
@ -89,6 +90,103 @@ async def test_update_batch_in_database_stores_unified_output_file_id():
assert stored["output_file_id"] != raw_output_file_id
@pytest.mark.asyncio
async def test_cancel_path_registers_output_file_under_batch_owner():
unified_id = "file-bWFuYWdlZF9vdXRwdXRfaWQ="
db_batch_object = SimpleNamespace(
created_by="batch-owner", team_id="batch-team", status="in_progress"
)
response = _build_batch_response(
status="cancelling",
output_file_id="file-raw-output",
hidden_params={"model_id": "my-model", "model_name": "openai/gpt-4o"},
)
mock_managed_files = _build_managed_files_mock(unified_id=unified_id)
mock_prisma = _build_prisma_mock(db_batch_object=db_batch_object)
await update_batch_in_database(
batch_id="batch_managed_ids_test",
unified_batch_id="litellm_proxy;model_id:my-model;llm_batch_id:batch_managed_ids_test",
response=response,
managed_files_obj=mock_managed_files,
prisma_client=mock_prisma,
verbose_proxy_logger=MagicMock(),
operation="cancel",
)
forwarded_auth = mock_managed_files.store_unified_file_id.call_args.kwargs[
"user_api_key_dict"
]
assert forwarded_auth.user_id == "batch-owner"
assert forwarded_auth.team_id == "batch-team"
stored = json.loads(
mock_prisma.db.litellm_managedobjecttable.update.call_args.kwargs["data"][
"file_object"
]
)
assert stored["output_file_id"] == unified_id
@pytest.mark.asyncio
async def test_update_batch_skips_lookup_when_db_batch_object_supplied():
unified_id = "file-bWFuYWdlZF9vdXRwdXRfaWQ="
caller_row = SimpleNamespace(
created_by="caller-owner", team_id="caller-team", status="in_progress"
)
decoy_row = SimpleNamespace(
created_by="decoy-owner", team_id="decoy-team", status="in_progress"
)
response = _build_batch_response(
status="cancelling",
output_file_id="file-raw-output",
hidden_params={"model_id": "my-model", "model_name": "openai/gpt-4o"},
)
mock_managed_files = _build_managed_files_mock(unified_id=unified_id)
mock_prisma = _build_prisma_mock(db_batch_object=decoy_row)
await update_batch_in_database(
batch_id="batch_managed_ids_test",
unified_batch_id="litellm_proxy;model_id:my-model;llm_batch_id:batch_managed_ids_test",
response=response,
managed_files_obj=mock_managed_files,
prisma_client=mock_prisma,
verbose_proxy_logger=MagicMock(),
db_batch_object=caller_row,
operation="retrieve",
)
mock_prisma.db.litellm_managedobjecttable.find_first.assert_not_called()
forwarded_auth = mock_managed_files.store_unified_file_id.call_args.kwargs[
"user_api_key_dict"
]
assert forwarded_auth.user_id == "caller-owner"
assert forwarded_auth.team_id == "caller-team"
@pytest.mark.asyncio
async def test_update_batch_derives_model_id_from_unified_batch_id():
unified_id = "file-bWFuYWdlZF9vdXRwdXRfaWQ="
response = _build_batch_response(output_file_id="file-raw-output", hidden_params={})
mock_managed_files = _build_managed_files_mock(unified_id=unified_id)
mock_prisma = _build_prisma_mock()
await update_batch_in_database(
batch_id="batch_managed_ids_test",
unified_batch_id="litellm_proxy;model_id:model-from-batch-id;llm_batch_id:batch_managed_ids_test",
response=response,
managed_files_obj=mock_managed_files,
prisma_client=mock_prisma,
verbose_proxy_logger=MagicMock(),
user_api_key_dict=UserAPIKeyAuth(user_id="user-abc"),
)
assert (
mock_managed_files.get_unified_output_file_id.call_args.kwargs["model_id"]
== "model-from-batch-id"
)
assert response.output_file_id == unified_id
@pytest.mark.asyncio
async def test_ensure_batch_response_normalizes_error_file_id():
"""Both output_file_id and error_file_id must be normalized to managed IDs."""

42
tests/test_litellm/enterprise/proxy/test_managed_files_hook.py
View file

@ -8,6 +8,7 @@ async_post_call_success_hook when processing completed batch responses.
import asyncio
import base64
import json
import logging
import pytest
from typing import Optional
@ -189,6 +190,47 @@ async def test_get_user_created_file_ids_remaps_stored_raw_provider_id_to_unifie
assert files[0].purpose == raw_provider_object.purpose
@pytest.mark.asyncio
async def test_parse_managed_file_object_warning_omits_rejected_values(caplog):
from litellm_enterprise.proxy.hooks.managed_files import (
_parse_managed_file_object,
)
with caplog.at_level(logging.WARNING):
parsed = _parse_managed_file_object(
{"id": "file-corrupt", "object": "file", "filename": "confidential.jsonl"},
"unified-corrupt",
)
assert parsed is None
assert "unified-corrupt" in caplog.text
assert "bytes" in caplog.text
assert "confidential.jsonl" not in caplog.text
@pytest.mark.asyncio
async def test_get_user_created_file_ids_skips_unparseable_rows():
managed_files = _make_managed_files_instance()
managed_files.prisma_client.db.litellm_managedfiletable.find_many = AsyncMock(
return_value=[
MagicMock(
file_object={"id": "file-corrupt", "object": "file"},
unified_file_id="unified-corrupt",
),
MagicMock(
file_object=_make_file_object().model_dump(),
unified_file_id="unified-valid",
),
]
)
files = await managed_files.get_user_created_file_ids(
_make_user_api_key_dict(), ["file-output-abc"]
)
assert [file.id for file in files] == ["unified-valid"]
@pytest.mark.asyncio
async def test_should_fallback_when_no_router():
"""

40
tests/test_litellm/integrations/arize/test_arize_phoenix.py
View file

@ -37,8 +37,8 @@ class TestArizePhoenixConfig(unittest.TestCase):
# Call the function to get the configuration
config = ArizePhoenixLogger.get_arize_phoenix_config()
# Verify the configuration - now uses standard Authorization Bearer format
self.assertEqual(config.otlp_auth_headers, "Authorization=Bearer test_api_key")
# gRPC metadata keys must be lowercase, so the auth header key is lowercased
self.assertEqual(config.otlp_auth_headers, "authorization=Bearer test_api_key")
self.assertEqual(config.endpoint, "grpc://test.endpoint")
self.assertEqual(config.protocol, "otlp_grpc")
@ -136,7 +136,7 @@ class TestArizePhoenixConfig(unittest.TestCase):
"PHOENIX_COLLECTOR_ENDPOINT": "grpc://localhost:6006",
"PHOENIX_API_KEY": "test_api_key",
},
"Authorization=Bearer test_api_key",
"authorization=Bearer test_api_key",
"grpc://localhost:6006",
"otlp_grpc",
id="explicit grpc endpoint with grpc:// prefix",
@ -215,6 +215,40 @@ def test_get_arize_phoenix_config_expection_on_missing_api_key(monkeypatch, env_
ArizePhoenixLogger.get_arize_phoenix_config()
@pytest.mark.parametrize(
"collector_endpoint, expected_key",
[
pytest.param("grpc://localhost:6006", "authorization", id="grpc prefix"),
pytest.param("http://localhost:4317", "authorization", id="grpc port 4317"),
pytest.param("http://localhost:6006", "Authorization", id="http"),
],
)
def test_get_arize_phoenix_config_auth_header_key_casing(
monkeypatch, collector_endpoint, expected_key
):
"""Regression for #34882: gRPC metadata keys must be lowercase.
HTTP headers are case-insensitive, but the OTLP/gRPC exporter rejects an
uppercase ``Authorization`` metadata key, so span export silently fails.
"""
for key in [
"PHOENIX_API_KEY",
"PHOENIX_COLLECTOR_ENDPOINT",
"PHOENIX_COLLECTOR_HTTP_ENDPOINT",
]:
monkeypatch.delenv(key, raising=False)
monkeypatch.setenv("PHOENIX_API_KEY", "test_api_key")
monkeypatch.setenv("PHOENIX_COLLECTOR_ENDPOINT", collector_endpoint)
config = ArizePhoenixLogger.get_arize_phoenix_config()
assert config.otlp_auth_headers == f"{expected_key}=Bearer test_api_key"
header_key = config.otlp_auth_headers.split("=", 1)[0]
if config.protocol == "otlp_grpc":
assert header_key == header_key.lower()
# ---------------------------------------------------------------------------
# Per-project routing via Resource (not span attributes)
# ---------------------------------------------------------------------------

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

@ -2620,3 +2620,120 @@ def test_fast_service_tier_matches_priority_above_the_context_threshold(_local_m
assert fast == priority
assert fast[0] == pytest.approx(300_000 * 1e-05, rel=1e-9)
assert fast[1] == pytest.approx(1_000 * 4.5e-05, rel=1e-9)
def test_priority_reasoning_tokens_bill_at_the_priority_output_rate(_local_model_cost_map):
"""Regression: gemini-3.5-flash publishes priority output pricing but no priority
reasoning key, so reasoning tokens under priority/fast were billed at the standard
output_cost_per_reasoning_token instead of following the tier's output rate."""
from litellm.types.utils import Usage
usage = Usage(
prompt_tokens=1_000,
completion_tokens=5_000,
completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=4_000),
)
model_info = litellm.get_model_info(model="gemini-3.5-flash", custom_llm_provider="gemini")
standard_output_rate = model_info["output_cost_per_token"]
standard_reasoning_rate = model_info["output_cost_per_reasoning_token"]
priority_output_rate = model_info["output_cost_per_token_priority"]
assert priority_output_rate is not None
assert priority_output_rate != standard_reasoning_rate
standard = generic_cost_per_token(
model="gemini-3.5-flash", usage=usage, custom_llm_provider="gemini", service_tier=None
)
priority = generic_cost_per_token(
model="gemini-3.5-flash", usage=usage, custom_llm_provider="gemini", service_tier="priority"
)
fast = generic_cost_per_token(
model="gemini-3.5-flash", usage=usage, custom_llm_provider="gemini", service_tier="fast"
)
assert standard[1] == pytest.approx(1_000 * standard_output_rate + 4_000 * standard_reasoning_rate, rel=1e-9)
assert priority[1] == pytest.approx(5_000 * priority_output_rate, rel=1e-9)
assert fast == priority
def test_explicit_tier_reasoning_key_wins_over_the_tier_output_rate():
from litellm.types.utils import Usage
model_info = {
"input_cost_per_token": 1e-06,
"output_cost_per_token": 4e-06,
"output_cost_per_reasoning_token": 6e-06,
"input_cost_per_token_priority": 2e-06,
"output_cost_per_token_priority": 8e-06,
"output_cost_per_reasoning_token_priority": 1.2e-05,
}
usage = Usage(
prompt_tokens=100,
completion_tokens=1_000,
completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=600),
)
_, completion_cost = generic_cost_per_token(
model="synthetic-model",
usage=usage,
custom_llm_provider="openai",
service_tier="priority",
model_info=model_info,
)
assert completion_cost == pytest.approx(400 * 8e-06 + 600 * 1.2e-05, rel=1e-9)
def test_null_tier_reasoning_key_falls_back_to_the_tier_output_rate():
"""get_model_info dumps every ModelInfo field, so an unpublished tier reasoning key
arrives as an explicit None and must not shadow the tier output rate."""
from litellm.types.utils import Usage
model_info = {
"input_cost_per_token": 1e-06,
"output_cost_per_token": 4e-06,
"output_cost_per_reasoning_token": 6e-06,
"output_cost_per_reasoning_token_priority": None,
"input_cost_per_token_priority": 2e-06,
"output_cost_per_token_priority": 8e-06,
}
usage = Usage(
prompt_tokens=100,
completion_tokens=1_000,
completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=600),
)
_, completion_cost = generic_cost_per_token(
model="synthetic-model",
usage=usage,
custom_llm_provider="openai",
service_tier="priority",
model_info=model_info,
)
assert completion_cost == pytest.approx(1_000 * 8e-06, rel=1e-9)
def test_tier_request_without_tier_pricing_keeps_the_standard_reasoning_rate():
from litellm.types.utils import Usage
model_info = {
"input_cost_per_token": 1e-06,
"output_cost_per_token": 4e-06,
"output_cost_per_reasoning_token": 6e-06,
}
usage = Usage(
prompt_tokens=100,
completion_tokens=1_000,
completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=600),
)
_, completion_cost = generic_cost_per_token(
model="synthetic-model",
usage=usage,
custom_llm_provider="openai",
service_tier="priority",
model_info=model_info,
)
assert completion_cost == pytest.approx(400 * 4e-06 + 600 * 6e-06, rel=1e-9)

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

@ -162,3 +162,56 @@ class TestGetLitellmParamsDataResidency:
api_base="https://eu.api.openai.com/v1",
)
assert result["data_residency"] is None
class TestMetadataFallsBackToLitellmMetadata:
def test_metadata_falls_back_to_litellm_metadata_when_absent(self):
result = get_litellm_params(litellm_metadata={"trace_id": "trace-1"})
assert result["metadata"] == {"trace_id": "trace-1"}
assert result["litellm_metadata"] == {"trace_id": "trace-1"}
def test_empty_metadata_falls_back_to_litellm_metadata(self):
result = get_litellm_params(metadata={}, litellm_metadata={"trace_id": "trace-1"})
assert result["metadata"] == {"trace_id": "trace-1"}
def test_metadata_wins_when_both_present(self):
result = get_litellm_params(
metadata={"trace_id": "from-metadata"},
litellm_metadata={"trace_id": "from-litellm-metadata"},
)
assert result["metadata"] == {"trace_id": "from-metadata"}
@pytest.mark.parametrize("bad_value", ["not-json-a-string", 12345, ["a"], True])
def test_non_dict_litellm_metadata_is_ignored(self, bad_value):
result = get_litellm_params(litellm_metadata=bad_value)
assert result["metadata"] is None
def test_metadata_stays_none_without_litellm_metadata(self):
result = get_litellm_params(api_key="test-key")
assert result["metadata"] is None
def test_session_and_trace_id_derived_from_litellm_metadata(self):
result = get_litellm_params(
litellm_metadata={"trace_id": "trace-1", "session_id": "session-1"},
)
assert result["litellm_session_id"] == "session-1"
assert result["litellm_trace_id"] == "trace-1"
def test_explicit_session_and_trace_id_are_not_overridden(self):
result = get_litellm_params(
litellm_session_id="explicit-session",
litellm_trace_id="explicit-trace",
litellm_metadata={"trace_id": "trace-1", "session_id": "session-1"},
)
assert result["litellm_session_id"] == "explicit-session"
assert result["litellm_trace_id"] == "explicit-trace"
def test_litellm_metadata_fallback_is_copied_not_aliased(self):
litellm_metadata = {"trace_id": "trace-1"}
result = get_litellm_params(litellm_metadata=litellm_metadata)
assert result["metadata"] == litellm_metadata
assert result["metadata"] is not litellm_metadata
result["metadata"].pop("trace_id")
assert litellm_metadata == {"trace_id": "trace-1"}

77
tests/test_litellm/litellm_core_utils/test_litellm_logging.py
View file

@ -11,6 +11,9 @@ sys.path.insert(
import time
import httpx
from openai._legacy_response import HttpxBinaryResponseContent
import litellm
from litellm.constants import SENTRY_DENYLIST, SENTRY_PII_DENYLIST
from litellm.integrations.custom_logger import CustomLogger
@ -523,6 +526,26 @@ class TestUpdateFromKwargs:
)
assert logging_obj.litellm_params["litellm_call_id"] == "call-empty"
@pytest.mark.parametrize("caller_metadata", [None, "not-a-dict", 42])
def test_non_dict_caller_metadata_does_not_break_the_merge(self, logging_obj, caller_metadata):
logging_obj.update_from_kwargs(
kwargs={"metadata": caller_metadata, "litellm_metadata": {"user_api_key_hash": "hashed"}},
litellm_params={"metadata": {"user_api_key_hash": "hashed", "litellm_api_version": "1.0"}},
)
assert logging_obj.litellm_params["metadata"]["user_api_key_hash"] == "hashed"
def test_does_not_mutate_caller_metadata_dict(self, logging_obj):
caller_metadata: dict = {}
logging_obj.update_from_kwargs(
kwargs={"metadata": caller_metadata, "litellm_metadata": {"user_api_key_hash": "hashed"}},
litellm_params={"metadata": {"user_api_key_hash": "hashed", "litellm_api_version": "1.0"}},
)
assert caller_metadata == {}
assert logging_obj.litellm_params["metadata"]["user_api_key_hash"] == "hashed"
def test_logging_prevent_double_logging(logging_obj):
"""
@ -1771,6 +1794,60 @@ def test_response_cost_calculator_does_not_transform_non_generate_content_dict()
assert not cost
def _file_content_logging_obj(call_type: str) -> LitellmLogging:
logging_obj = LitellmLogging(
model="gemini-3-flash-preview",
messages="default-message-value",
stream=False,
call_type=call_type,
start_time=time.time(),
litellm_call_id=f"file-content-{call_type}",
function_id=f"file-content-{call_type}",
)
logging_obj.model_call_details["custom_llm_provider"] = "vertex_ai"
logging_obj.model_call_details["input"] = "default-message-value"
logging_obj.optional_params = {}
return logging_obj
@pytest.mark.parametrize("call_type", ["afile_content", "file_content"])
def test_file_content_call_is_not_billed(call_type):
"""
Regression for #35130: file content retrieval has no token usage, but ``function_setup``
stores the ``"default-message-value"`` placeholder as the logged input, which the cost
calculator then token-priced, billing every call at exactly 3 * input_cost_per_token.
"""
result = HttpxBinaryResponseContent(httpx.Response(status_code=200, content=b"file contents"))
cost = _file_content_logging_obj(call_type)._response_cost_calculator(result=result)
assert cost == 0.0
@pytest.mark.parametrize("call_type", ["aspeech", "speech"])
def test_speech_call_is_still_priced_from_input_characters(call_type):
"""tts bills per input character, so speech call types must keep passing the input along."""
logging_obj = LitellmLogging(
model="tts-1",
messages="the quick brown fox jumped over the lazy dogs",
stream=False,
call_type=call_type,
start_time=time.time(),
litellm_call_id=f"speech-{call_type}",
function_id=f"speech-{call_type}",
)
logging_obj.model_call_details["custom_llm_provider"] = "openai"
logging_obj.model_call_details["input"] = "the quick brown fox jumped over the lazy dogs"
logging_obj.optional_params = {}
result = HttpxBinaryResponseContent(httpx.Response(status_code=200, content=b"audio bytes"))
cost = logging_obj._response_cost_calculator(result=result)
assert cost is not None
assert cost > 0
def test_sentry_event_scrubber_initialization(monkeypatch):
# Step 1: Create a fake sentry_sdk.scrubber module
mock_event_scrubber_instance = MagicMock()

203
tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py
View file

@ -5,6 +5,7 @@ Tests the handler's ability to process streaming output for Anthropic Messages A
with guardrail transformations, specifically testing edge cases with empty choices.
"""
import json
import os
import sys
from typing import Any, Literal, Optional
@ -760,3 +761,205 @@ class TestAnthropicMessagesToolResultScanning:
assert "skip me POISON" not in guardrail.seen_texts
assert messages[1]["content"][0]["content"] == "skip me POISON"
assert messages[0]["content"] == "keep me [BLOCKED]"
class InputsRecordingGuardrail(MockMaskingGuardrail):
def __init__(self):
super().__init__(guardrail_name="scan-only-capture")
self.captured_inputs: Optional[GenericGuardrailAPIInputs] = None
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict,
input_type: Literal["request", "response"],
logging_obj: Optional[Any] = None,
) -> GenericGuardrailAPIInputs:
self.captured_inputs = inputs
return await super().apply_guardrail(inputs, request_data, input_type, logging_obj)
class StructuredMessagesRewritingGuardrail(CustomGuardrail):
"""Returns a new structured_messages list with a canary redacted, like redaction guardrails do."""
def __init__(self):
super().__init__(guardrail_name="structured-rewrite")
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict,
input_type: Literal["request", "response"],
logging_obj: Optional[Any] = None,
) -> GenericGuardrailAPIInputs:
structured = inputs.get("structured_messages") or []
inputs["structured_messages"] = [
json.loads(json.dumps(message).replace("POISON", "[BLOCKED]")) for message in structured
]
return inputs
class TestAnthropicMessagesScanOnlyToolResults:
def _guardrail(self):
guardrail = InputsRecordingGuardrail()
guardrail.scan_only_tool_results = True
return guardrail
@pytest.mark.asyncio
async def test_structured_write_back_merges_into_the_full_conversation(self):
handler = AnthropicMessagesHandler()
guardrail = StructuredMessagesRewritingGuardrail()
guardrail.scan_only_tool_results = True
data = {
"model": "claude-sonnet-4-5",
"system": "You are a careful agent harness.",
"messages": [
{"role": "user", "content": "fetch the page"},
{
"role": "assistant",
"content": [{"type": "tool_use", "id": "tu1", "name": "Bash", "input": {"cmd": "curl"}}],
},
{
"role": "user",
"content": [{"type": "tool_result", "tool_use_id": "tu1", "content": "fetched POISON page"}],
},
],
}
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
assert data["system"] == "You are a careful agent harness."
assert [m["role"] for m in data["messages"]] == ["user", "assistant", "user"], (
"a redacting guardrail must not strip out-of-scope turns from the request"
)
serialized = json.dumps(data["messages"])
assert "fetch the page" in serialized
assert "tool_use" in serialized
assert "fetched [BLOCKED] page" in serialized
assert "POISON" not in serialized
@pytest.mark.asyncio
async def test_scan_narrows_to_tool_results_and_write_back_stays_aligned(self):
handler = AnthropicMessagesHandler()
guardrail = self._guardrail()
data = {
"model": "claude-sonnet-4-5",
"system": "You are a trusted agent harness with POISON heuristics.",
"tools": [
{
"name": "Bash",
"description": "run a command",
"input_schema": {"type": "object", "properties": {}},
}
],
"messages": [
{"role": "user", "content": "scaffolding POISON prompt"},
{
"role": "assistant",
"content": [{"type": "tool_use", "id": "tu1", "name": "Bash", "input": {"cmd": "curl"}}],
},
{
"role": "user",
"content": [
{"type": "text", "text": "sibling POISON text"},
{"type": "tool_result", "tool_use_id": "tu1", "content": "fetched POISON page"},
],
},
],
}
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
assert guardrail.seen_texts == ["fetched POISON page"], (
"only the tool_result payload may reach the guardrail"
)
assert guardrail.captured_inputs is not None
assert guardrail.captured_inputs.get("tools") is None
assert [m["role"] for m in guardrail.captured_inputs["structured_messages"]] == ["tool"]
assert data["messages"][2]["content"][1]["content"] == "fetched [BLOCKED] page"
assert data["messages"][0]["content"] == "scaffolding POISON prompt", (
"out-of-scope content must come back untouched, not masked or dropped"
)
assert data["messages"][2]["content"][0]["text"] == "sibling POISON text"
@pytest.mark.asyncio
async def test_guardrail_synthesized_tools_are_appended_without_replacing_request_tools(self):
handler = AnthropicMessagesHandler()
guardrail = ToolAppendingGuardrail(guardrail_name="tool-appending")
guardrail.scan_only_tool_results = True
original_tools = [
{
"name": "get_weather",
"description": "Get the weather at a specific location",
"input_schema": {"type": "object", "properties": {"location": {"type": "string"}}},
}
]
data = {
"model": "claude-sonnet-4-5",
"tools": original_tools,
"messages": [
{"role": "user", "content": "what's the weather?"},
{
"role": "assistant",
"content": [{"type": "tool_use", "id": "tu1", "name": "get_weather", "input": {}}],
},
{
"role": "user",
"content": [{"type": "tool_result", "tool_use_id": "tu1", "content": "sunny"}],
},
],
}
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
assert [t["name"] for t in data["tools"]] == ["get_weather", "injected_tool"], (
"a tool the guardrail synthesized must reach the model, converted to Anthropic format, "
"without the request's own tools being replaced or dropped"
)
assert data["tools"][0] == original_tools[0]
@pytest.mark.asyncio
async def test_guardrail_is_not_called_when_the_request_has_no_tool_results(self):
handler = AnthropicMessagesHandler()
guardrail = self._guardrail()
data = {
"model": "claude-sonnet-4-5",
"messages": [{"role": "user", "content": "What is 2 plus 2?"}],
}
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
assert guardrail.captured_inputs is None
assert guardrail.seen_texts == []
@pytest.mark.asyncio
async def test_images_are_scoped_the_same_way_as_texts(self):
handler = AnthropicMessagesHandler()
guardrail = self._guardrail()
data = {
"model": "claude-sonnet-4-5",
"messages": [
{
"role": "user",
"content": [{"type": "image", "source": {"type": "base64", "data": "USER_IMG"}}],
},
{
"role": "user",
"content": [
{
"type": "tool_result",
"tool_use_id": "tu1",
"content": [
{"type": "text", "text": "screenshot POISON"},
{"type": "image", "source": {"type": "base64", "data": "TOOL_IMG"}},
],
}
],
},
],
}
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
assert guardrail.captured_inputs is not None
assert guardrail.captured_inputs.get("images") == ["TOOL_IMG"]

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

@ -1229,3 +1229,338 @@ class TestIncrementalScanRespectsSkipFlags:
assert mock_api.call_count == 1
scanned = [m["content"] for m in mock_api.call_args.kwargs["messages"]]
assert scanned == ["It is sunny in Paris.", "And tomorrow?"]
class StructuredRedactionGuardrail(CustomGuardrail):
"""Captures inputs and returns a new structured_messages list with a canary redacted."""
def __init__(self):
super().__init__(guardrail_name="structured-redaction")
self.captured_inputs: Optional[GenericGuardrailAPIInputs] = None
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict,
input_type: Literal["request", "response"],
logging_obj: Optional[Any] = None,
) -> GenericGuardrailAPIInputs:
self.captured_inputs = inputs
structured = inputs.get("structured_messages") or []
inputs["structured_messages"] = [
{**m, "content": str(m.get("content", "")).replace("POISON", "[BLOCKED]")} for m in structured
]
return inputs
class ToolSynthesizingGuardrail(CustomGuardrail):
"""Appends its own function tool to whatever tools it was given, like a
retrieval/recovery guardrail that injects a tool the model can later call."""
def __init__(self):
super().__init__(guardrail_name="tool-synthesizing")
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict,
input_type: Literal["request", "response"],
logging_obj: Optional[Any] = None,
) -> GenericGuardrailAPIInputs:
tools = list(inputs.get("tools") or [])
tools.append(
{
"type": "function",
"function": {"name": "injected_retrieve", "parameters": {"type": "object", "properties": {}}},
}
)
inputs["tools"] = tools
return inputs
class ToolNameCollidingGuardrail(CustomGuardrail):
"""Returns a tool reusing a request tool's name plus a genuinely new tool."""
def __init__(self):
super().__init__(guardrail_name="tool-name-colliding")
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict,
input_type: Literal["request", "response"],
logging_obj: Optional[Any] = None,
) -> GenericGuardrailAPIInputs:
inputs["tools"] = [
{
"type": "function",
"function": {
"name": "read_file",
"parameters": {"type": "object", "properties": {"hijacked": {"type": "string"}}},
},
},
{
"type": "function",
"function": {"name": "injected_retrieve", "parameters": {"type": "object", "properties": {}}},
},
]
return inputs
class DuplicateToolReturningGuardrail(CustomGuardrail):
"""Returns the same synthesized tool name twice, second copy with a different schema."""
def __init__(self):
super().__init__(guardrail_name="duplicate-tool-returning")
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict,
input_type: Literal["request", "response"],
logging_obj: Optional[Any] = None,
) -> GenericGuardrailAPIInputs:
inputs["tools"] = [
{
"type": "function",
"function": {
"name": "injected_retrieve",
"parameters": {"type": "object", "properties": {"first": {"type": "string"}}},
},
},
{
"type": "function",
"function": {
"name": "injected_retrieve",
"parameters": {"type": "object", "properties": {"second": {"type": "string"}}},
},
},
]
return inputs
class TestScanOnlyToolResults:
def _bedrock_guardrail(self):
from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail
guardrail = BedrockGuardrail(
guardrail_name="bedrock-scan-only-tool-results",
guardrailIdentifier="test-guardrail",
guardrailVersion="DRAFT",
default_on=True,
)
guardrail.scan_only_tool_results = True
return guardrail
@pytest.mark.asyncio
async def test_only_tool_role_content_is_scanned(self):
from unittest.mock import AsyncMock, patch
handler = OpenAIChatCompletionsHandler()
guardrail = self._bedrock_guardrail()
data = {
"messages": [
{"role": "system", "content": "SYSTEM-PROMPT-not-scanned"},
{"role": "user", "content": "USER-PROMPT-not-scanned"},
{
"role": "assistant",
"content": "ASSISTANT-not-scanned",
"tool_calls": [
{
"id": "call_1",
"type": "function",
"function": {"name": "read_file", "arguments": '{"path": "report.html"}'},
}
],
},
{"role": "tool", "tool_call_id": "call_1", "content": "TOOL-RESULT-scanned"},
]
}
with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api:
mock_api.return_value = {"action": "NONE", "output": [], "outputs": []}
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
assert mock_api.call_count == 1
scanned = [m["content"] for m in mock_api.call_args.kwargs["messages"]]
assert scanned == ["TOOL-RESULT-scanned"]
@pytest.mark.asyncio
async def test_legacy_function_role_results_are_scanned(self):
from unittest.mock import AsyncMock, patch
handler = OpenAIChatCompletionsHandler()
guardrail = self._bedrock_guardrail()
data = {
"messages": [
{"role": "user", "content": "USER-PROMPT-not-scanned"},
{"role": "function", "name": "read_file", "content": "FUNCTION-RESULT-scanned"},
{"role": "tool", "tool_call_id": "call_1", "content": "TOOL-RESULT-scanned"},
]
}
with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api:
mock_api.return_value = {"action": "NONE", "output": [], "outputs": []}
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
assert mock_api.call_count == 1
scanned = [m["content"] for m in mock_api.call_args.kwargs["messages"]]
assert scanned == ["FUNCTION-RESULT-scanned", "TOOL-RESULT-scanned"], (
"a tool result sent with the legacy function role must not bypass the scoped scan"
)
@pytest.mark.parametrize("flag_value", [None, "false", 0, object()])
@pytest.mark.asyncio
async def test_scope_narrows_only_when_the_flag_is_actually_true(self, flag_value):
from unittest.mock import AsyncMock, patch
handler = OpenAIChatCompletionsHandler()
guardrail = self._bedrock_guardrail()
guardrail.scan_only_tool_results = flag_value
data = {
"messages": [
{"role": "user", "content": "USER-PROMPT"},
{"role": "tool", "tool_call_id": "call_1", "content": "TOOL-RESULT"},
]
}
with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api:
mock_api.return_value = {"action": "NONE", "output": [], "outputs": []}
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
assert mock_api.call_count == 1
scanned = [m["content"] for m in mock_api.call_args.kwargs["messages"]]
assert scanned == ["USER-PROMPT", "TOOL-RESULT"], (
"anything but an explicit True must leave the whole request in scope"
)
@pytest.mark.parametrize("scan_only_tool_results", [True, False])
@pytest.mark.asyncio
async def test_function_definitions_are_scoped_out_with_the_tool_results_flag(self, scan_only_tool_results):
handler = OpenAIChatCompletionsHandler()
guardrail = StructuredRedactionGuardrail()
guardrail.scan_only_tool_results = scan_only_tool_results
tools = [
{
"type": "function",
"function": {"name": "read_file", "parameters": {"type": "object", "properties": {}}},
}
]
data = {
"messages": [
{"role": "user", "content": "read the report"},
{"role": "tool", "tool_call_id": "call_1", "content": "TOOL-RESULT"},
],
"tools": tools,
}
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
assert guardrail.captured_inputs is not None
expected_tools = None if scan_only_tool_results else tools
assert guardrail.captured_inputs.get("tools") == expected_tools, (
"function definitions must stay out of a tool-results-only scan"
)
@pytest.mark.parametrize("scan_only_tool_results", [True, False])
@pytest.mark.asyncio
async def test_guardrail_synthesized_tools_are_appended_without_replacing_request_tools(
self, scan_only_tool_results
):
handler = OpenAIChatCompletionsHandler()
guardrail = ToolSynthesizingGuardrail()
guardrail.scan_only_tool_results = scan_only_tool_results
original_tools = [
{
"type": "function",
"function": {"name": "read_file", "parameters": {"type": "object", "properties": {}}},
}
]
data = {
"messages": [
{"role": "user", "content": "read the report"},
{"role": "tool", "tool_call_id": "call_1", "content": "TOOL-RESULT"},
],
"tools": original_tools,
}
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
assert [t["function"]["name"] for t in data["tools"]] == ["read_file", "injected_retrieve"], (
"a tool the guardrail synthesized (like a recovery/retrieve tool) must reach the model "
"without the request's own tools being replaced or dropped"
)
assert data["tools"][0] == original_tools[0]
@pytest.mark.asyncio
async def test_returned_tool_name_collisions_keep_the_request_schema(self):
handler = OpenAIChatCompletionsHandler()
guardrail = ToolNameCollidingGuardrail()
guardrail.scan_only_tool_results = True
original_read_file = {
"type": "function",
"function": {"name": "read_file", "parameters": {"type": "object", "properties": {}}},
}
data = {
"messages": [
{"role": "user", "content": "read the report"},
{"role": "tool", "tool_call_id": "call_1", "content": "TOOL-RESULT"},
],
"tools": [original_read_file],
}
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
assert [t["function"]["name"] for t in data["tools"]] == ["read_file", "injected_retrieve"]
assert data["tools"][0] == original_read_file, (
"a returned tool reusing a request tool's name must not replace the request's schema"
)
@pytest.mark.asyncio
async def test_duplicate_returned_tool_names_keep_only_the_first(self):
handler = OpenAIChatCompletionsHandler()
guardrail = DuplicateToolReturningGuardrail()
guardrail.scan_only_tool_results = True
original_read_file = {
"type": "function",
"function": {"name": "read_file", "parameters": {"type": "object", "properties": {}}},
}
data = {
"messages": [
{"role": "user", "content": "read the report"},
{"role": "tool", "tool_call_id": "call_1", "content": "TOOL-RESULT"},
],
"tools": [original_read_file],
}
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
assert [t["function"]["name"] for t in data["tools"]] == ["read_file", "injected_retrieve"], (
"two returned tools sharing a name must not both be forwarded to the provider"
)
assert data["tools"][1]["function"]["parameters"]["properties"] == {"first": {"type": "string"}}
@pytest.mark.asyncio
async def test_structured_write_back_keeps_out_of_scope_messages(self):
handler = OpenAIChatCompletionsHandler()
guardrail = StructuredRedactionGuardrail()
guardrail.scan_only_tool_results = True
data = {
"messages": [
{"role": "system", "content": "SYSTEM-PROMPT"},
{"role": "user", "content": "fetch the page"},
{
"role": "assistant",
"content": "fetching",
"tool_calls": [
{"id": "call_1", "type": "function", "function": {"name": "fetch", "arguments": "{}"}}
],
},
{"role": "tool", "tool_call_id": "call_1", "content": "page says POISON here"},
{"role": "user", "content": "and then?"},
]
}
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
assert [m["role"] for m in data["messages"]] == ["system", "user", "assistant", "tool", "user"], (
"a redacting guardrail must not strip out-of-scope messages from the request"
)
assert data["messages"][0]["content"] == "SYSTEM-PROMPT"
assert data["messages"][3]["content"] == "page says [BLOCKED] here"
assert data["messages"][3]["tool_call_id"] == "call_1"
assert data["messages"][4]["content"] == "and then?"

27
tests/test_litellm/proxy/anthropic_endpoints/test_claude_code_marketplace.py
View file

@ -18,13 +18,14 @@ from litellm.types.proxy.claude_code_endpoints import (
UpdatePluginRequest,
)
from litellm.proxy.anthropic_endpoints.claude_code_endpoints.claude_code_marketplace import (
get_marketplace,
register_plugin,
update_plugin,
)
def _make_mock_prisma():
"""Stateful prisma mock that supports find_unique, create, and update."""
"""Stateful prisma mock that supports find_unique, find_many, create, and update."""
store: dict = {}
mock_client = MagicMock()
@ -34,6 +35,12 @@ def _make_mock_prisma():
async def _find_unique(where):
return store.get(where.get("name"))
async def _find_many(where=None):
records = list(store.values())
if where and "enabled" in where:
return [r for r in records if r.enabled == where["enabled"]]
return records
async def _create(data):
record = MagicMock()
record.id = "test-id"
@ -52,6 +59,7 @@ def _make_mock_prisma():
return record
mock_table.find_unique = AsyncMock(side_effect=_find_unique)
mock_table.find_many = AsyncMock(side_effect=_find_many)
mock_table.create = AsyncMock(side_effect=_create)
mock_table.update = AsyncMock(side_effect=_update)
mock_client.db.litellm_claudecodeplugintable = mock_table
@ -211,6 +219,23 @@ async def test_update_plugin_db_error_maps_to_structured_500():
assert "connection lost" in exc_info.value.detail["error"]
@pytest.mark.asyncio
async def test_get_marketplace_skips_plugin_with_null_manifest():
await register_plugin(
request=RegisterPluginRequest(name="good-plugin", source=_GIT_SUBDIR_SOURCE, version="1.0.0"),
user_api_key_dict=_USER,
)
table = litellm.proxy.proxy_server.prisma_client.db.litellm_claudecodeplugintable
await table.create(data={"name": "null-manifest-plugin", "manifest_json": None, "enabled": True})
response = await get_marketplace()
assert response.status_code == 200
body = json.loads(response.body)
assert [plugin["name"] for plugin in body["plugins"]] == ["good-plugin"]
@pytest.mark.asyncio
async def test_register_plugin_git_subdir_missing_url():
"""git-subdir without url field raises HTTP 400."""

52
tests/test_litellm/proxy/auth/test_route_checks.py
View file

@ -3246,3 +3246,55 @@ def test_internal_user_still_blocked_from_another_users_info():
assert exc_info.value.status_code == 403
assert "key not allowed to access this user's info" in str(exc_info.value.detail)
@pytest.mark.parametrize(
"route",
[
"/user/daily/activity",
"/user/daily/activity/aggregated",
],
)
@pytest.mark.parametrize(
"user_role",
[
LitellmUserRoles.INTERNAL_USER.value,
LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value,
],
)
def test_user_daily_activity_routes_reachable_by_non_admin(route, user_role):
"""Both /user/daily/activity and its /aggregated sibling power the default
"Your Usage" dashboard view, and both handlers self-scope to the caller
(_user_has_admin_view -> require_caller_user_id_for_non_admin -> 403 on a
user_id mismatch). self_managed_routes is the ONLY list that grants either
route to a non-admin, so dropping one from it 401s every internal user's
main Usage page before the handler ever runs.
"""
user_obj = LiteLLM_UserTable(
user_id="test_user",
user_email="test@example.com",
user_role=user_role,
)
valid_token = UserAPIKeyAuth(user_id="test_user", user_role=user_role)
request = MagicMock(spec=Request)
request.query_params = {}
RouteChecks.non_proxy_admin_allowed_routes_check(
user_obj=user_obj,
_user_role=user_role,
route=route,
request=request,
valid_token=valid_token,
request_data={},
)
def test_user_daily_activity_aggregated_not_covered_by_prefix_match():
"""check_route_access is exact-match plus explicit wildcards, so listing the
parent /user/daily/activity does not implicitly cover the /aggregated
sub-path. Pins the reason the sibling needs its own entry.
"""
assert not RouteChecks.check_route_access(
route="/user/daily/activity/aggregated",
allowed_routes=["/user/daily/activity"],
)

15
tests/test_litellm/proxy/batches_endpoints/test_endpoints.py
View file

@ -469,13 +469,14 @@ async def test_create__fallback_body_custom_llm_provider(harness):
@pytest.mark.asyncio
async def test_create__unified_file_id_single_model(harness):
async def test_create__unified_file_id_single_model_disables_cross_model_fallbacks(harness):
set_body(
harness,
{
"input_file_id": "litellm_proxy_unified_id",
"endpoint": "/v1/chat/completions",
"completion_window": "24h",
"disable_fallbacks": False,
},
)
with (
@ -489,6 +490,7 @@ async def test_create__unified_file_id_single_model(harness):
harness.litellm_acreate.assert_not_called()
# model injected from the unified id, input_file_id restored, hidden param set
assert harness.router_kwargs()["model"] == "gpt-4o-mini"
assert harness.router_kwargs()["disable_fallbacks"] is True
assert resp.input_file_id == "litellm_proxy_unified_id"
assert resp._hidden_params["unified_file_id"] == "unified-xyz"
@ -1896,6 +1898,17 @@ async def test_cancel__unified_batch_id_routes_to_router(cancel_harness):
assert cancel_harness.update_batch_in_db.call_args.kwargs["operation"] == "cancel"
@pytest.mark.asyncio
async def test_cancel__db_write_receives_caller_auth(cancel_harness):
"""update_batch_in_database can only mint managed IDs for a cancelled batch's
output files when it has an auth context, so cancel must forward the caller's."""
caller = UserAPIKeyAuth(api_key="sk-test", user_id="user-cancel-1")
with patch.object(endpoints, "_is_base64_encoded_unified_file_id", return_value=UNIFIED_BATCH_ID):
await call_cancel(cancel_harness, "batch-unified-blob", user=caller)
assert cancel_harness.update_batch_in_db.call_args.kwargs["user_api_key_dict"] is caller
@pytest.mark.asyncio
async def test_cancel__unified_missing_model_id_400(cancel_harness):
# unified id with no model_id segment -> get_model_id returns None -> 400.

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

@ -3670,3 +3670,40 @@ async def test_moderation_hook_honors_the_mcp_event_type(mode, call_type, should
"the scan must be logged under the event it actually ran for, so guardrail logs, "
"OTel spans, and Langfuse metadata do not misclassify MCP enforcement as an LLM call"
)
class TestScanOnlyToolResultsWithLatestRoleFilter:
@pytest.mark.asyncio
async def test_warns_and_skips_when_scoped_payload_has_no_user_message(self):
"""scan_only_tool_results hands Bedrock a tool-role-only payload, but
experimental_use_latest_role_message_only scans only the latest user
message: the silent no-op must warn."""
guardrail = BedrockGuardrail(
guardrail_name="bedrock-latest-role-scoped",
guardrailIdentifier="test-guardrail",
guardrailVersion="DRAFT",
default_on=True,
experimental_use_latest_role_message_only=True,
)
guardrail.scan_only_tool_results = True
inputs = {
"texts": ["TOOL-RESULT"],
"structured_messages": [{"role": "tool", "tool_call_id": "call_1", "content": "TOOL-RESULT"}],
}
with (
patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api,
patch(
"litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails.verbose_proxy_logger.warning"
) as mock_warning,
):
result = await guardrail.apply_guardrail(
inputs=inputs,
request_data={"litellm_call_id": "test-call-id"},
input_type="request",
)
mock_api.assert_not_called()
assert result["texts"] == ["TOOL-RESULT"]
warning_text = " ".join(str(arg) for c in mock_warning.call_args_list for arg in c.args)
assert "scan_only_tool_results" in warning_text

28
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py
View file

@ -1696,6 +1696,34 @@ class TestPanwAirsApplyGuardrail:
request_data=request_data, guardrail_name=handler.guardrail_name
)
@pytest.mark.asyncio
async def test_apply_guardrail_warns_when_tool_results_scope_leaves_nothing_scannable(self, handler):
"""scan_only_tool_results hands PANW a tool-role-only payload, but PANW's role
filter only scans user/system/developer rows: the silent no-op must warn."""
handler.scan_only_tool_results = True
inputs: GenericGuardrailAPIInputs = {
"texts": ["TOOL-RESULT"],
"structured_messages": [{"role": "tool", "tool_call_id": "call_1", "content": "TOOL-RESULT"}],
}
request_data = {"litellm_call_id": "test-call-id", "model": "gpt-4"}
with (
patch.object(handler, "_call_panw_api", new_callable=AsyncMock) as mock_api,
patch(
"litellm.proxy.guardrails.guardrail_hooks.panw_prisma_airs.panw_prisma_airs.verbose_proxy_logger.warning"
) as mock_warning,
):
result = await handler.apply_guardrail(
inputs=inputs,
request_data=request_data,
input_type="request",
)
mock_api.assert_not_called()
assert result["texts"] == ["TOOL-RESULT"]
warning_text = " ".join(str(arg) for c in mock_warning.call_args_list for arg in c.args)
assert "scan_only_tool_results" in warning_text
@pytest.mark.asyncio
async def test_apply_guardrail_block(self, handler):
"""Test block action raises HTTPException(400)."""

99
tests/test_litellm/proxy/guardrails/test_guardrail_registry.py
View file

@ -558,3 +558,102 @@ def test_reinitialized_judge_guardrail_uses_lazy_router_provider():
finally:
for cb_list, snapshot in zip(lists, snapshots):
cb_list[:] = snapshot
class TestScanOnlyToolResultsInitRefusal:
"""A guardrail whose role filtering never scans tool results must be rejected at
initialization when configured with scan_only_tool_results, instead of booting a
proxy that silently scans nothing on every request."""
def _initialize(self, name: str, params: dict):
lists = _all_callback_lists()
snapshots = [list(cb_list) for cb_list in lists]
try:
return InMemoryGuardrailHandler().initialize_guardrail(
guardrail={"guardrail_name": name, "litellm_params": params},
)
finally:
for cb_list, snapshot in zip(lists, snapshots):
cb_list[:] = snapshot
def test_panw_prisma_airs_with_scan_only_tool_results_is_rejected(self):
with pytest.raises(ValueError, match="never scans tool results"):
self._initialize(
"panw-scan-only-combo",
{
"guardrail": "panw_prisma_airs",
"mode": "pre_call",
"api_key": "test-key",
"profile_name": "test-profile",
"scan_only_tool_results": True,
},
)
def test_bedrock_latest_role_with_scan_only_tool_results_is_rejected(self):
with pytest.raises(ValueError, match="never scans tool results"):
self._initialize(
"bedrock-latest-role-scan-only-combo",
{
"guardrail": "bedrock",
"mode": "pre_call",
"guardrailIdentifier": "gr-1",
"guardrailVersion": "1",
"experimental_use_latest_role_message_only": True,
"scan_only_tool_results": True,
},
)
def test_bedrock_without_latest_role_accepts_scan_only_tool_results(self):
result = self._initialize(
"bedrock-scan-only-ok",
{
"guardrail": "bedrock",
"mode": "pre_call",
"guardrailIdentifier": "gr-1",
"guardrailVersion": "1",
"scan_only_tool_results": True,
},
)
assert result is not None
def test_prompt_security_default_tool_filtering_rejects_scan_only_tool_results(self, monkeypatch):
monkeypatch.delenv("PROMPT_SECURITY_CHECK_TOOL_RESULTS", raising=False)
with pytest.raises(ValueError, match="never scans tool results"):
self._initialize(
"prompt-security-scan-only-combo",
{
"guardrail": "prompt_security",
"mode": "pre_call",
"api_key": "test-key",
"api_base": "https://ps.example.com",
"scan_only_tool_results": True,
},
)
def test_prompt_security_check_tool_results_accepts_scan_only_tool_results(self, monkeypatch):
monkeypatch.setenv("PROMPT_SECURITY_CHECK_TOOL_RESULTS", "true")
result = self._initialize(
"prompt-security-scan-only-ok",
{
"guardrail": "prompt_security",
"mode": "pre_call",
"api_key": "test-key",
"api_base": "https://ps.example.com",
"scan_only_tool_results": True,
},
)
assert result is not None
def test_skip_tool_message_with_scan_only_tool_results_is_rejected(self):
with pytest.raises(ValueError, match="skip_tool_message_in_guardrail are enabled together"):
self._initialize(
"bedrock-skip-tool-scan-only-combo",
{
"guardrail": "bedrock",
"mode": "pre_call",
"guardrailIdentifier": "gr-1",
"guardrailVersion": "1",
"skip_tool_message_in_guardrail": True,
"scan_only_tool_results": True,
},
)

17
tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py
View file

@ -1186,6 +1186,7 @@ async def test_track_cost_callback_enriches_user_id_for_mcp_style_metadata():
("pass_through_endpoint", True),
("llm_passthrough_route", True),
("allm_passthrough_route", True),
("aretrieve_batch", True),
("acompletion", False),
("call_mcp_tool", False),
(None, False),
@ -1194,7 +1195,14 @@ async def test_track_cost_callback_enriches_user_id_for_mcp_style_metadata():
def test_should_track_cost_callback_pass_through_without_owner(call_type, expected):
"""Regression for LIT-3782: unauthenticated pass-through requests (auth=false)
carry no key/user/team/end-user, yet must still be tracked so they land in
LiteLLM_SpendLogs. Other call types with no owner stay untracked."""
LiteLLM_SpendLogs. Other call types with no owner stay untracked.
aretrieve_batch is included for the same reason: CheckBatchCost's synthetic
logging_obj for a completed managed batch only ever carries
user_api_key_user_id/user_api_key_team_id from LiteLLM_ManagedObjectTable,
both of which are None for a batch created with the master key or a
team-less key (the table never stores the raw key hash). Before this fix,
such a batch's cost silently never reached LiteLLM_SpendLogs."""
assert (
_should_track_cost_callback(
user_api_key=None,
@ -1211,6 +1219,7 @@ def test_should_track_cost_callback_pass_through_without_owner(call_type, expect
"call_type, expect_spend_log",
[
("pass_through_endpoint", True),
("aretrieve_batch", True),
("acompletion", False),
(None, False),
],
@ -1223,7 +1232,11 @@ async def test_track_cost_callback_logs_unauthenticated_pass_through_request(
cost callback with no key/user/team/end-user. Before the fix the spend-log
write was skipped and the request never appeared in request/usage logs. It
must now be written for pass-through call types while other unauthenticated
calls remain skipped."""
calls remain skipped.
aretrieve_batch is included because CheckBatchCost's completed-batch cost
event reaches this same callback with no attributable key/user/team when
the batch was created with the master key or a team-less key."""
logger = _ProxyDBLogger()
kwargs = {

182
tests/test_litellm/proxy/management_endpoints/search_endpoints/test_search_tool_management.py
View file

@ -5,6 +5,7 @@ from datetime import datetime
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from fastapi import HTTPException
from fastapi.testclient import TestClient
sys.path.insert(
@ -23,6 +24,7 @@ import litellm.proxy.proxy_server as ps
# Now we can safely import app
from litellm.proxy.proxy_server import app
from litellm.types.search import SearchToolInfoResponse
client = TestClient(app)
@ -815,3 +817,183 @@ async def test_list_search_tools_admin_with_restricted_key_still_sees_all():
assert response.status_code == 200
names = {t["search_tool_name"] for t in response.json()["search_tools"]}
assert names == {"db-tool-1", "db-tool-2", "db-tool-3"}
def _search_tool_responses(*names: str) -> list[SearchToolInfoResponse]:
return [
SearchToolInfoResponse(
search_tool_id=f"id-{name}",
search_tool_name=name,
litellm_params={"search_provider": "perplexity"},
search_tool_info=None,
created_at=None,
updated_at=None,
is_from_config=False,
)
for name in names
]
def _team_ids_looked_up(lookup: AsyncMock) -> list[str]:
return [awaited.args[0] for awaited in lookup.await_args_list]
@pytest.mark.asyncio
async def test_list_search_tools_dashboard_session_key_does_not_look_up_the_ui_team():
"""
Regression: the Admin UI session key is stamped with the reserved team id
``litellm-dashboard``, which has no row in LiteLLM_TeamTable. Resolving it as a real team
raised 404, which the endpoint reported as a 500, so the Search Tools page was broken for
every non-admin browsing the dashboard.
"""
from litellm.constants import UI_SESSION_TOKEN_TEAM_ID
dashboard_session_user = UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
user_id="internal_user",
team_id=UI_SESSION_TOKEN_TEAM_ID,
)
ui_team_is_not_a_real_team = AsyncMock(
side_effect=HTTPException(
status_code=404,
detail={"error": f"Team doesn't exist in db. Team={UI_SESSION_TOKEN_TEAM_ID}."},
)
)
with (
_mock_search_tool_backend(_scoping_db_tools()),
patch(
"litellm.proxy.auth.auth_checks.get_team_object",
ui_team_is_not_a_real_team,
),
_override_auth(dashboard_session_user),
):
response = TestClient(app).get("/search_tools/list")
assert response.status_code == 200
names = {t["search_tool_name"] for t in response.json()["search_tools"]}
assert names == {"db-tool-1", "db-tool-2", "db-tool-3"}
ui_team_is_not_a_real_team.assert_not_awaited()
@pytest.mark.asyncio
async def test_filter_visible_search_tools_dashboard_session_still_honors_key_allowlist():
"""
Skipping the synthetic team must not widen visibility: a dashboard session whose key
carries a search_tools allowlist stays scoped to it.
"""
from litellm.constants import UI_SESSION_TOKEN_TEAM_ID
from litellm.proxy.search_endpoints.search_tool_management import (
_filter_visible_search_tools,
)
restricted_dashboard_session = UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
user_id="internal_user",
team_id=UI_SESSION_TOKEN_TEAM_ID,
object_permission=LiteLLM_ObjectPermissionTable(
object_permission_id="op-key",
search_tools=["db-tool-3"],
),
)
lookup = AsyncMock()
visible = await _filter_visible_search_tools(
_search_tool_responses("db-tool-1", "db-tool-2", "db-tool-3"),
restricted_dashboard_session,
lookup,
)
assert [t["search_tool_name"] for t in visible] == ["db-tool-3"]
lookup.assert_not_awaited()
@pytest.mark.asyncio
async def test_filter_visible_search_tools_still_applies_a_real_team_allowlist():
"""A caller with a real team is still resolved and scoped by that team's allowlist."""
from litellm.proxy.search_endpoints.search_tool_management import (
_filter_visible_search_tools,
)
team_member = UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
user_id="internal_user",
team_id="team-1",
)
lookup = AsyncMock(
return_value=LiteLLM_TeamTable(
team_id="team-1",
object_permission=LiteLLM_ObjectPermissionTable(
object_permission_id="op-team",
search_tools=["db-tool-2"],
),
)
)
visible = await _filter_visible_search_tools(
_search_tool_responses("db-tool-1", "db-tool-2", "db-tool-3"),
team_member,
lookup,
)
assert [t["search_tool_name"] for t in visible] == ["db-tool-2"]
assert _team_ids_looked_up(lookup) == ["team-1"]
@pytest.mark.asyncio
async def test_filter_visible_search_tools_propagates_a_real_team_lookup_failure():
"""
A caller whose real team cannot be resolved must not fall through to "no team", which
would drop that team's allowlist and show tools the caller may not call.
"""
from litellm.proxy.search_endpoints.search_tool_management import (
_filter_visible_search_tools,
)
team_member = UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
user_id="internal_user",
team_id="deleted-team",
)
lookup = AsyncMock(side_effect=HTTPException(status_code=404, detail={"error": "Team doesn't exist in db."}))
with pytest.raises(HTTPException) as exc_info:
await _filter_visible_search_tools(
_search_tool_responses("db-tool-1", "db-tool-2"),
team_member,
lookup,
)
assert exc_info.value.status_code == 404
assert _team_ids_looked_up(lookup) == ["deleted-team"]
@pytest.mark.asyncio
async def test_list_search_tools_reports_a_missing_real_team_as_404():
"""
The endpoint surfaces a genuine team lookup failure with its own status instead of
masking it as a 500 or quietly returning an unscoped list.
"""
team_member = UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
user_id="internal_user",
team_id="deleted-team",
)
with (
_mock_search_tool_backend(_scoping_db_tools()),
patch(
"litellm.proxy.auth.auth_checks.get_team_object",
AsyncMock(
side_effect=HTTPException(
status_code=404,
detail={"error": "Team doesn't exist in db. Team=deleted-team."},
)
),
),
_override_auth(team_member),
):
response = TestClient(app).get("/search_tools/list")
assert response.status_code == 404
assert "search_tools" not in response.json()

62
tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py
View file

@ -1,6 +1,8 @@
import os
import sys
from datetime import datetime, timezone
from types import SimpleNamespace
from typing import Final
from unittest.mock import AsyncMock, MagicMock
import pytest
@ -870,6 +872,66 @@ class TestAdjustDatesForTimezone:
assert per_day_ends == days
class TestAdjustDatesForTimezoneLiveEnd:
"""
Regression tests for the stale-evening bug: a caller west of UTC whose range
ends on their local "today" was capped at that local date's UTC bucket, so
once UTC rolled past their local midnight (5pm PT), everything sent that
evening sat in the next UTC bucket and the dashboard reported $0 for it
until local midnight. A range that reaches the caller's current day and
opts in via include_current_utc_day must extend to today's UTC bucket; the
only part of that bucket outside the range is the future, which is empty,
so the extension cannot over-count. Callers that do not opt in keep the
pass-through byte for byte.
"""
PT_EVENING_UTC: Final = datetime(2026, 8, 6, 4, 30, tzinfo=timezone.utc)
def test_pt_evening_range_ending_today_extends_to_utc_today(self):
start, end = _adjust_dates_for_timezone(
"2026-07-06", "2026-08-05", 420, include_current_utc_day=True, utc_now=self.PT_EVENING_UTC
)
assert (start, end) == ("2026-07-06", "2026-08-06")
def test_without_opt_in_live_range_keeps_pass_through(self):
start, end = _adjust_dates_for_timezone(
"2026-07-06", "2026-08-05", 420, utc_now=self.PT_EVENING_UTC
)
assert (start, end) == ("2026-07-06", "2026-08-05")
def test_pt_historical_range_is_untouched(self):
start, end = _adjust_dates_for_timezone(
"2026-07-01", "2026-08-04", 420, include_current_utc_day=True, utc_now=self.PT_EVENING_UTC
)
assert (start, end) == ("2026-07-01", "2026-08-04")
def test_east_of_utc_local_today_already_covers_utc_today(self):
ist_evening_utc: Final = datetime(2026, 8, 5, 17, 0, tzinfo=timezone.utc)
start, end = _adjust_dates_for_timezone(
"2026-07-07", "2026-08-06", -330, include_current_utc_day=True, utc_now=ist_evening_utc
)
assert (start, end) == ("2026-07-07", "2026-08-06")
def test_missing_offset_stays_pass_through_even_for_live_range(self):
start, end = _adjust_dates_for_timezone(
"2026-07-06", "2026-08-05", None, include_current_utc_day=True, utc_now=self.PT_EVENING_UTC
)
assert (start, end) == ("2026-07-06", "2026-08-05")
def test_utc_caller_range_ending_today_is_unchanged(self):
utc_noon: Final = datetime(2026, 8, 5, 12, 0, tzinfo=timezone.utc)
start, end = _adjust_dates_for_timezone(
"2026-07-06", "2026-08-05", 0, include_current_utc_day=True, utc_now=utc_noon
)
assert (start, end) == ("2026-07-06", "2026-08-05")
def test_future_end_date_extends_no_further_than_requested(self):
start, end = _adjust_dates_for_timezone(
"2026-07-06", "2026-08-09", 420, include_current_utc_day=True, utc_now=self.PT_EVENING_UTC
)
assert (start, end) == ("2026-07-06", "2026-08-09")
class TestBuildAggregatedSqlQuery:
"""
Asserts the SQL emitted by the aggregated query path stays anchored to the

69
tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py
View file

@ -2294,6 +2294,75 @@ async def test_get_user_daily_activity_aggregated_admin_global_view(monkeypatch)
)
@pytest.mark.asyncio
async def test_get_user_daily_activity_aggregated_non_admin_cannot_view_other_users(
monkeypatch,
):
"""
Same scoping contract as
test_get_user_daily_activity_non_admin_cannot_view_other_users, on the
aggregated route. Non-admins reach this handler now that the route is in
self_managed_routes, so the 403-on-mismatch and default-to-self behaviour
has to hold here too: opening the route must not widen access.
"""
from unittest.mock import AsyncMock, MagicMock, patch
from fastapi import HTTPException
from litellm.proxy.management_endpoints.internal_user_endpoints import (
get_user_daily_activity_aggregated,
)
mock_prisma_client = MagicMock()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
non_admin_key_dict = UserAPIKeyAuth(
user_id="regular-user-123",
user_role=LitellmUserRoles.INTERNAL_USER,
)
# Case 1: Non-admin targets another user's data — 403, helper never reached
with patch(
"litellm.proxy.management_endpoints.internal_user_endpoints.get_daily_activity_aggregated",
new_callable=AsyncMock,
) as mock_get_daily_agg:
with pytest.raises(HTTPException) as exc_info:
await get_user_daily_activity_aggregated(
start_date="2025-01-01",
end_date="2025-01-31",
model=None,
api_key=None,
user_id="other-user-456",
timezone=None,
user_api_key_dict=non_admin_key_dict,
)
assert exc_info.value.status_code == 403
assert "Non-admin users can only view their own spend data" in str(exc_info.value.detail)
mock_get_daily_agg.assert_not_called()
# Case 2: Non-admin omits user_id — scoped to their own user_id, not global
mock_response = MagicMock()
with patch(
"litellm.proxy.management_endpoints.internal_user_endpoints.get_daily_activity_aggregated",
new_callable=AsyncMock,
return_value=mock_response,
) as mock_get_daily_agg:
result = await get_user_daily_activity_aggregated(
start_date="2025-01-01",
end_date="2025-01-31",
model=None,
api_key=None,
user_id=None,
timezone=None,
user_api_key_dict=non_admin_key_dict,
)
assert result is mock_response
mock_get_daily_agg.assert_called_once()
assert mock_get_daily_agg.call_args.kwargs["entity_id"] == "regular-user-123"
@pytest.mark.asyncio
async def test_delete_user_cleans_up_created_by_invitation_links(mocker):
"""

49
tests/test_litellm/proxy/proxy_server/test_lifecycle.py
View file

@ -20,6 +20,7 @@ from __future__ import annotations
import asyncio
import inspect
import json
import logging
import os
from typing import List, Optional, Union
from unittest.mock import AsyncMock, MagicMock, patch
@ -31,6 +32,7 @@ from typing_extensions import TypedDict
import litellm.proxy.proxy_server as ps
from litellm.proxy.proxy_server import (
ProxyStartupEvent,
_initialize_shared_aiohttp_session,
_resolve_pydantic_type,
_resolve_typed_dict_type,
@ -728,3 +730,50 @@ def test_otel_global_provider_published_after_callback_init():
"preset logger will not exist yet and a second generic logger will own "
"the global provider, orphaning gen-ai spans"
)
def test_startup_warns_for_global_budget_without_database(caplog):
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
ProxyStartupEvent._warn_budget_without_db(max_budget=100.0, prisma_client=None)
assert "litellm.max_budget=100.0" in caplog.text
assert "will NOT be enforced" in caplog.text
assert "requests will never be blocked" in caplog.text
def test_startup_does_not_warn_for_global_budget_with_database(caplog):
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
ProxyStartupEvent._warn_budget_without_db(max_budget=100.0, prisma_client=MagicMock())
assert "litellm.max_budget" not in caplog.text
@pytest.mark.parametrize("max_budget", [0, None])
def test_startup_does_not_warn_without_global_budget(caplog, max_budget):
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
ProxyStartupEvent._warn_budget_without_db(max_budget=max_budget, prisma_client=None)
assert "litellm.max_budget" not in caplog.text
def test_proxy_startup_event_warns_for_global_budget_without_database():
"""Pin the lifespan call that prevents silent DB-less budgets.
The call must follow Prisma setup so DB-backed deployments do not false-positive.
Direct ``_warn_budget_without_db`` tests cover the warning behavior itself.
"""
wrapped = getattr(proxy_startup_event, "__wrapped__", proxy_startup_event)
source = inspect.getsource(wrapped)
budget_check_pos = source.find("if prisma_client is not None and litellm.max_budget > 0:")
warn_pos = source.find("_warn_budget_without_db(")
next_startup_section_pos = source.find(
"await ProxyStartupEvent.initialize_scheduled_background_jobs(",
budget_check_pos,
)
assert budget_check_pos != -1, "global budget startup block not found"
assert warn_pos != -1, "DB-less budget warning call not found"
assert next_startup_section_pos != -1, "startup section after budget block not found"
assert budget_check_pos < warn_pos < next_startup_section_pos, (
"DB-less budget warning must run after Prisma setup and the DB-backed budget block"
)

182
tests/test_litellm/proxy/test_litellm_pre_call_utils.py
View file

@ -20,6 +20,7 @@ from litellm.proxy.litellm_pre_call_utils import (
_get_dynamic_logging_metadata,
_get_enforced_params,
_get_metadata_variable_name,
_promoted_trace_control_fields,
_resolve_credential_from_model_config,
_resolve_provider_from_deployment,
_update_model_if_key_alias_exists,
@ -5869,4 +5870,183 @@ async def test_key_level_callback_vars_survive_the_strip():
)
assert updated[TRUSTED_CALLBACK_VARS_FIELD] == {"dd_api_key": "key-dd-key", "dd_site": "us5.datadoghq.com"}
assert updated["dd_site"] == "us5.datadoghq.com"
assert updated["dd_site"] == "us5.datadoghq.com"
class TestPromotedTraceControlFields:
"""LIT-5137: caller metadata trace fields must reach litellm_metadata."""
def _make_request(self, path: str) -> MagicMock:
request = MagicMock(spec=Request)
request.url = MagicMock()
request.url.path = path
request.url.__str__.return_value = f"http://localhost{path}"
request.method = "POST"
request.query_params = {}
request.headers = {"Content-Type": "application/json"}
request.client = MagicMock()
request.client.host = "127.0.0.1"
return request
async def _run(self, path: str, data: dict, headers: dict | None = None) -> dict:
request = self._make_request(path)
if headers is not None:
request.headers = {"Content-Type": "application/json", **headers}
return await add_litellm_data_to_request(
data=data,
request=request,
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
proxy_config=MagicMock(),
general_settings={},
version="test-version",
)
def test_returns_litellm_metadata_for_responses_route(self):
assert _get_metadata_variable_name(self._make_request("/v1/responses")) == "litellm_metadata"
def test_promotes_trace_prefixed_and_allow_listed_fields(self):
requester_metadata = {
"trace_id": "trace-1",
"trace_name": "name-1",
"trace_user_id": "user-1",
"trace_metadata": {"tenant_id": "tenant-1"},
"trace_version": "v1",
"trace_release": "r1",
"session_id": "session-1",
"mask_input": True,
"mask_output": True,
}
promoted = _promoted_trace_control_fields(
requester_metadata=requester_metadata,
litellm_metadata={},
)
assert dict(promoted) == requester_metadata
def test_does_not_promote_unlisted_trace_prefixed_fields(self):
"""trace_public flips a trace to publicly readable, so the allow-list is explicit."""
promoted = _promoted_trace_control_fields(
requester_metadata={"trace_id": "trace-1", "trace_public": True, "trace_tags": ["a"]},
litellm_metadata={},
)
assert dict(promoted) == {"trace_id": "trace-1"}
def test_does_not_promote_non_trace_fields(self):
promoted = _promoted_trace_control_fields(
requester_metadata={
"trace_id": "trace-1",
"tags": ["free-tier"],
"user_api_key": "forged",
"user_api_key_user_id": "forged-user",
"spend_logs_metadata": {"forged": True},
"guardrails": ["disabled"],
"debug_langfuse": True,
"session": "not-session-id",
"existing_trace_id": "victim-trace",
"update_trace_keys": ["input", "output"],
},
litellm_metadata={},
)
assert dict(promoted) == {"trace_id": "trace-1"}
def test_does_not_promote_trace_mutation_controls(self):
"""existing_trace_id + update_trace_keys let a caller overwrite any trace in the project."""
promoted = _promoted_trace_control_fields(
requester_metadata={
"trace_id": "trace-1",
"existing_trace_id": "someone-elses-trace",
"update_trace_keys": ["input", "output"],
},
litellm_metadata={},
)
assert dict(promoted) == {"trace_id": "trace-1"}
def test_existing_litellm_metadata_value_wins(self):
promoted = _promoted_trace_control_fields(
requester_metadata={"trace_id": "from-body", "session_id": "from-body", "trace_name": "from-body"},
litellm_metadata={"trace_id": "from-header", "session_id": "from-header"},
)
assert dict(promoted) == {"trace_name": "from-body"}
def test_empty_requester_metadata_promotes_nothing(self):
assert _promoted_trace_control_fields(requester_metadata={}, litellm_metadata={}) == ()
@pytest.mark.asyncio
async def test_responses_route_end_to_end(self):
caller_metadata = {
"trace_id": "22662678-30c1-41a1-a24b-216d6e5fb83d",
"session_id": "218af06c-28a2-4705-8a0a-5f9970d39326",
"trace_user_id": "user-123",
"trace_metadata": {"tenant_id": "tenant-1"},
"mask_input": True,
}
updated = await self._run(
"/v1/responses",
{"model": "gpt-4.1-mini", "input": "say resp", "metadata": copy.deepcopy(caller_metadata)},
)
litellm_metadata = updated["litellm_metadata"]
for key, value in caller_metadata.items():
assert litellm_metadata[key] == value
assert updated["metadata"] == caller_metadata
@pytest.mark.asyncio
async def test_messages_route_end_to_end(self):
updated = await self._run(
"/v1/messages",
{
"model": "claude-sonnet-4-5",
"max_tokens": 32,
"messages": [{"role": "user", "content": "hi"}],
"metadata": {"trace_id": "msg-trace-1", "session_id": "msg-session-1"},
},
)
assert updated["litellm_metadata"]["trace_id"] == "msg-trace-1"
assert updated["litellm_metadata"]["session_id"] == "msg-session-1"
@pytest.mark.asyncio
async def test_session_id_header_beats_body_metadata(self):
updated = await self._run(
"/v1/responses",
{"model": "gpt-4.1-mini", "input": "say resp", "metadata": {"session_id": "from-body"}},
headers={"x-litellm-session-id": "from-header-12345678"},
)
assert updated["litellm_metadata"]["session_id"] == "from-header-12345678"
@pytest.mark.asyncio
async def test_forged_user_api_key_fields_are_not_promoted(self):
updated = await self._run(
"/v1/responses",
{
"model": "gpt-4.1-mini",
"input": "say resp",
"metadata": {"trace_id": "trace-1", "user_api_key_user_id": "forged", "spend_logs_metadata": {"a": 1}},
},
)
litellm_metadata = updated["litellm_metadata"]
assert litellm_metadata["trace_id"] == "trace-1"
assert litellm_metadata.get("user_api_key_user_id") != "forged"
@pytest.mark.asyncio
async def test_chat_completions_route_is_untouched(self):
updated = await self._run(
"/v1/chat/completions",
{
"model": "gpt-4.1-mini",
"messages": [{"role": "user", "content": "hello"}],
"metadata": {"trace_id": "trace-1", "session_id": "session-1"},
},
)
assert "litellm_metadata" not in updated
assert updated["metadata"]["trace_id"] == "trace-1"
assert updated["metadata"]["session_id"] == "session-1"

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