Merge remote-tracking branch 'origin/main' into HEAD

# Conflicts:
#	litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py
This commit is contained in:
mateo-berri 2026-09-14 18:19:38 -07:00
commit 6574d83eae
220 changed files with 15407 additions and 2198 deletions

View file

@ -147,6 +147,9 @@ commands:
db_name:
type: string
default: circle_test
image:
type: string
default: postgres:14@sha256:6a70deda415ec296f977890e11aba04a0db9f632a362e3fce45e845e3db74f26
steps:
- run:
name: Start PostgreSQL
@ -157,7 +160,7 @@ commands:
-e POSTGRES_PASSWORD=postgres \
-e POSTGRES_DB=<< parameters.db_name >> \
-p 5432:5432 \
postgres:14@sha256:6a70deda415ec296f977890e11aba04a0db9f632a362e3fce45e845e3db74f26
<< parameters.image >>
- wait_for_service:
url: tcp://localhost:5432
timeout: "60"
@ -2912,7 +2915,50 @@ jobs:
exit 1
fi
integration_contracts:
parameters:
suite:
type: string
machine:
image: ubuntu-2204:2024.04.1
resource_class: large
working_directory: ~/project
steps:
- setup_litellm_test_deps
- start_postgres:
image: postgres:16@sha256:e17e86066e5ef83e0952a9347f5c792b7ece00972e2aa787a6986f471b3dd3d5
- start_redis
- run:
name: Run owned integration contracts
command: bash .circleci/scripts/run_integration.sh << parameters.suite >>
no_output_timeout: 15m
- run:
name: Stop owned database and Redis
when: always
command: |
mkdir -p test-results/integration-<< parameters.suite >>
docker logs postgres-db > test-results/integration-<< parameters.suite >>/postgres.log 2>&1 || true
docker logs redis-cache > test-results/integration-<< parameters.suite >>/redis.log 2>&1 || true
docker rm -f postgres-db redis-cache
test -z "$(docker ps -aq --filter name=postgres-db --filter name=redis-cache)"
- store_test_results:
path: test-results
- store_artifacts:
path: test-results
workflows:
integration:
jobs:
- integration_contracts:
name: integration-<< matrix.suite >>
matrix:
parameters:
suite: [management, accounting, providers]
filters:
branches:
only:
- main
- /litellm_.*/
build_and_test:
jobs:
- using_litellm_on_windows:

View file

@ -0,0 +1,142 @@
#!/usr/bin/env bash
set -euo pipefail
suite="${1:?integration suite required}"
results="test-results/integration-${suite}"
mkdir -p "$results"
integration_identity="$(.venv/bin/python -c 'import uuid; print(uuid.uuid4().hex)')"
upstream_pid=""
proxy_pid=""
peer_pid=""
launched_pid=""
guard_created=false
guard_installed=false
guard6_created=false
guard6_installed=false
cleanup() {
original_status=$?
trap - EXIT INT TERM
sudo .venv/bin/python .circleci/scripts/stop_integration_processes.py \
"$integration_identity" "$(id -u)" "$proxy_pid" "$peer_pid" "$upstream_pid" \
> "$results/process-cleanup.txt" 2>&1 || original_status=1
for owned_pid in "$peer_pid" "$proxy_pid" "$upstream_pid"; do
if [ -n "$owned_pid" ]; then
kill -- "-$owned_pid" 2>/dev/null || true
for _ in {1..50}; do
kill -0 -- "-$owned_pid" 2>/dev/null || break
sleep 0.1
done
if kill -0 -- "-$owned_pid" 2>/dev/null; then
kill -KILL -- "-$owned_pid" 2>/dev/null || true
original_status=1
fi
wait "$owned_pid" 2>/dev/null || true
fi
done
if [ "$guard_installed" = true ]; then
sudo iptables -D OUTPUT -m owner --uid-owner "$(id -u)" -j integration_only || original_status=1
fi
if [ "$guard_created" = true ]; then
sudo iptables -F integration_only || original_status=1
sudo iptables -X integration_only || original_status=1
fi
if [ "$guard6_installed" = true ]; then
sudo ip6tables -D OUTPUT -m owner --uid-owner "$(id -u)" -j integration_only || original_status=1
fi
if [ "$guard6_created" = true ]; then
sudo ip6tables -F integration_only || original_status=1
sudo ip6tables -X integration_only || original_status=1
fi
printf '%s\n' "$original_status" > "$results/exit-status.txt"
exit "$original_status"
}
trap cleanup EXIT
trap 'exit 130' INT
trap 'exit 143' TERM
export PATH="$PWD/.venv/bin:$PATH"
export PYTHONPATH="$PWD:$PWD/tests:$PWD/tests/e2e"
export DATABASE_URL="postgresql://postgres:postgres@127.0.0.1:5432/circle_test"
export REDIS_HOST=127.0.0.1 REDIS_PORT=6379
export LITELLM_MASTER_KEY=sk-integration-master LITELLM_SALT_KEY=sk-integration-salt
export LITELLM_MODE=PRODUCTION LITELLM_LOCAL_MODEL_COST_MAP=True
export STORE_MODEL_IN_DB=True AWS_EC2_METADATA_DISABLED=true DO_NOT_TRACK=1
export INTEGRATION_PROXY_URL=http://127.0.0.1:4000
export INTEGRATION_PEER_URL=""
export INTEGRATION_UPSTREAM_URL=http://127.0.0.1:8190
export INTEGRATION_MASTER_KEY="$LITELLM_MASTER_KEY"
export INTEGRATION_SEED="$((16#$(git rev-parse --short=8 HEAD)))"
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma > "$results/prisma-generate.log" 2>&1
sudo iptables -N integration_only
guard_created=true
sudo iptables -A integration_only -o lo -j ACCEPT
sudo iptables -A integration_only -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT
for service in postgres-db redis-cache; do
address="$(docker inspect --format '{{range .NetworkSettings.Networks}}{{.IPAddress}}{{end}}' "$service")"
sudo iptables -A integration_only -d "$address" -j ACCEPT
done
sudo iptables -A integration_only -j REJECT
sudo iptables -I OUTPUT 1 -m owner --uid-owner "$(id -u)" -j integration_only
guard_installed=true
sudo ip6tables -N integration_only
guard6_created=true
sudo ip6tables -A integration_only -o lo -j ACCEPT
sudo ip6tables -A integration_only -j REJECT
sudo ip6tables -I OUTPUT 1 -m owner --uid-owner "$(id -u)" -j integration_only
guard6_installed=true
if curl --noproxy '*' --connect-timeout 2 -s http://198.51.100.1 >/dev/null 2>&1; then
echo "Unexpected outbound network access" >&2
exit 1
fi
sudo iptables -L integration_only -n -v -x > "$results/egress-guard.txt"
awk '$3 == "REJECT" && $1 > 0 { rejected=1 } END { exit !rejected }' "$results/egress-guard.txt"
setsid env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" INTEGRATION_RUN_ID="$integration_identity" \
.venv/bin/python -m integration._support.upstream > "$results/upstream.log" 2>&1 &
upstream_pid=$!
start_proxy() {
local port="$1"
local log_name="$2"
setsid env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" INTEGRATION_RUN_ID="$integration_identity" \
DATABASE_URL="$DATABASE_URL" REDIS_HOST="$REDIS_HOST" REDIS_PORT="$REDIS_PORT" \
LITELLM_MASTER_KEY="$LITELLM_MASTER_KEY" LITELLM_SALT_KEY="$LITELLM_SALT_KEY" \
LITELLM_MODE=PRODUCTION LITELLM_LOCAL_MODEL_COST_MAP=True STORE_MODEL_IN_DB=True \
AWS_EC2_METADATA_DISABLED=true DO_NOT_TRACK=1 \
.venv/bin/python -m integration._support.proxy --config tests/integration/proxy_config.yaml \
--host 127.0.0.1 --port "$port" --num_workers 1 --telemetry False \
--use_prisma_db_push --enforce_prisma_migration_check \
> "$results/$log_name" 2>&1 &
launched_pid=$!
}
start_proxy 4000 proxy.log
proxy_pid="$launched_pid"
.venv/bin/python .circleci/scripts/wait_integration_services.py
if [ "$suite" = management ]; then
export INTEGRATION_PEER_URL=http://127.0.0.1:4001
start_proxy 4001 peer.log
peer_pid="$launched_pid"
.venv/bin/python .circleci/scripts/wait_integration_services.py
fi
if [ "$suite" = providers ]; then
INTEGRATION_RUN_ID="$integration_identity" .venv/bin/python -m pytest --noconftest -o addopts= \
--strict-markers --strict-config -p no:pytest-retry -p no:rerunfailures --timeout=30 \
tests/e2e/test_provider_edge.py::TestReplayMode::test_content_drift_returns_the_miss_status_naming_both_keys \
tests/e2e/test_provider_edge.py::TestReplayMode::test_exhausted_key_returns_the_miss_status \
tests/e2e/test_provider_edge.py::TestReplayLeftover::test_partially_consumed_recording_names_the_leftover \
tests/e2e/test_provider_edge.py::TestStreamingFidelity::test_replay_of_a_stream_makes_no_provider_connection \
--junitxml="$results/replay-controls.xml"
fi
timeout --signal=TERM --kill-after=20s 11m env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" \
INTEGRATION_RUN_ID="$integration_identity" \
DATABASE_URL="$DATABASE_URL" REDIS_HOST="$REDIS_HOST" REDIS_PORT="$REDIS_PORT" \
INTEGRATION_PROXY_URL="$INTEGRATION_PROXY_URL" INTEGRATION_PEER_URL="$INTEGRATION_PEER_URL" \
INTEGRATION_UPSTREAM_URL="$INTEGRATION_UPSTREAM_URL" \
INTEGRATION_MASTER_KEY="$INTEGRATION_MASTER_KEY" LITELLM_MODE=PRODUCTION \
INTEGRATION_SEED="$INTEGRATION_SEED" \
LITELLM_LOCAL_MODEL_COST_MAP=True AWS_EC2_METADATA_DISABLED=true DO_NOT_TRACK=1 \
.venv/bin/python tests/integration/run.py "$suite" --results "$results"

View file

@ -0,0 +1,53 @@
import sys
from typing import Final
import psutil
def is_owned(process: psutil.Process, identity: str, owner_uid: int) -> bool:
try:
return process.uids().real == owner_uid and process.environ().get("INTEGRATION_RUN_ID") == identity
except psutil.NoSuchProcess:
return False
def owned_processes(identity: str, owner_uid: int) -> tuple[psutil.Process, ...]:
return tuple(process for process in psutil.process_iter() if is_owned(process, identity, owner_uid))
def main(identity: str, owner_uid: int, root_pids: tuple[int, ...]) -> int:
assert owner_uid > 0, "The integration process owner must be a non-root UID"
owned: Final = owned_processes(identity, owner_uid)
roots: Final = tuple(process for process in owned if process.pid in root_pids)
for process in roots:
try:
process.terminate()
except psutil.NoSuchProcess:
continue
psutil.wait_procs(roots, timeout=30)
residual: Final = owned_processes(identity, owner_uid)
for process in residual:
try:
process.terminate()
except psutil.NoSuchProcess:
continue
psutil.wait_procs(residual, timeout=10)
remaining: Final = owned_processes(identity, owner_uid)
for process in remaining:
try:
process.kill()
except psutil.NoSuchProcess:
continue
psutil.wait_procs(remaining, timeout=2)
survivors: Final = owned_processes(identity, owner_uid)
print(
f"Owned integration processes: {len(owned)}, roots: {len(roots)}, "
f"residual: {len(residual)}, forced: {len(remaining)}, remaining: {len(survivors)}"
)
for process in remaining:
print(f"Forced cleanup was required for PID {process.pid}")
return 1 if remaining or survivors else 0
if __name__ == "__main__":
raise SystemExit(main(sys.argv[1], int(sys.argv[2]), tuple(int(value) for value in sys.argv[3:] if value)))

View file

@ -0,0 +1,43 @@
import os
import time
from typing import Final
import httpx
from redis import Redis
def main() -> None:
primary: Final = os.environ["INTEGRATION_PROXY_URL"]
peer: Final = os.environ.get("INTEGRATION_PEER_URL")
proxies: Final = (primary, peer) if peer else (primary,)
deadline: Final = time.monotonic() + 90
headers: Final = {"Authorization": f"Bearer {os.environ['INTEGRATION_MASTER_KEY']}"}
with httpx.Client(trust_env=False, timeout=2) as client, Redis(
host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"]), socket_timeout=2
) as cache:
while True:
try:
ready: Final = (
client.get(f"{os.environ['INTEGRATION_UPSTREAM_URL']}/health").status_code == 200
and all(client.get(f"{url}/health/readiness").status_code == 200 for url in proxies)
)
if ready:
for url in proxies:
response: Final = client.get(f"{url}/cache/ping", headers=headers)
response.raise_for_status()
result: Final = response.json()
assert result["status"] == "healthy", result
assert result["cache_type"] == "redis", result
assert result["ping_response"] is True, result
assert result["set_cache_response"] == "success", result
if cache.pubsub_numsub("litellm_proxy.auth_cache_invalidation")[0][1] >= len(proxies):
return
except httpx.TransportError:
pass
if time.monotonic() >= deadline:
raise SystemExit("Integration services or auth-cache subscribers did not become ready")
time.sleep(0.2)
if __name__ == "__main__":
main()

View file

@ -14,6 +14,17 @@ query-filters:
id: py/clear-text-logging-sensitive-data # CWE-312
- exclude:
id: py/polynomial-redos # CWE-730
# Import resolution confuses stdlib types with management_endpoints/types.py.
# The generic cycle query also reports intentional deferred imports.
- exclude:
id: py/cyclic-import
- exclude:
id: py/unsafe-cyclic-import
# Known false positives on live settings and Protocol placeholders.
- exclude:
id: py/unused-global-variable
- exclude:
id: py/ineffectual-statement
paths-ignore:
- tests

View file

@ -1,6 +1,7 @@
from __future__ import annotations
import ast
import json
import operator
import pathlib
import re
@ -498,6 +499,69 @@ def _check_shards() -> int:
return 0
def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozenset[str], tuple[Finding, ...]]:
manifest: Final = repo_root / "tests/integration/contracts.json"
if not manifest.exists():
return frozenset(), ()
entries: Final = json.loads(manifest.read_text())
paths: Final = frozenset(node.split("::", 1)[0] for node in entries["tests"])
circle_path: Final = repo_root / ".circleci/config.yml"
circle: Final = yaml.safe_load(circle_path.read_text()) if circle_path.exists() else {}
steps: Final = circle.get("jobs", {}).get("integration_contracts", {}).get("steps", ())
invoked: Final = any(
".circleci/scripts/run_integration.sh" in scalar.value
for scalar in _scalars(steps, "integration_contracts")
if scalar.key == "command"
)
scheduled: Final = frozenset(
suite
for job in circle.get("workflows", {}).get("integration", {}).get("jobs", ())
if isinstance(job, dict) and "integration_contracts" in job
for suite in job["integration_contracts"]
.get("matrix", {})
.get("parameters", {})
.get("suite", (job["integration_contracts"].get("suite"),))
if isinstance(suite, str)
)
required: Final = frozenset(
group
for group, folders in entries["groups"].items()
if any(any(path.startswith(f"tests/integration/{folder}/") for folder in folders) for path in paths)
)
ungrouped: Final = frozenset(
path
for path in paths
if sum(
any(path.startswith(f"tests/integration/{folder}/") for folder in folders)
for folders in entries["groups"].values()
)
!= 1
)
gha_tokens: Final = _invoked_test_tokens(
scalar
for path in (repo_root / ".github/workflows").glob("*.y*ml")
for scalar in _scalars(yaml.safe_load(path.read_text()), path.name)
)
findings: Final = tuple(
Finding(path, "integration contract is also selected by GitHub Actions")
for path in paths
if any(_token_covers(token, path) for token in gha_tokens)
) + tuple(
Finding(path, "canonical integration test file is missing")
for path in paths
if not (repo_root / path).is_file()
)
group_findings: Final = tuple(
Finding(group, "canonical integration group is not scheduled by CircleCI")
for group in sorted(required - scheduled)
) + tuple(Finding(path, "canonical node must have exactly one integration group") for path in sorted(ungrouped))
if not paths or not invoked or not scheduled:
return frozenset(), findings + (
Finding(str(manifest.relative_to(repo_root)), "dedicated CircleCI runner is missing"),
)
return paths, findings + group_findings
def main() -> int:
if "--shards" in sys.argv[1:]:
return _check_shards()
@ -507,7 +571,8 @@ def main() -> int:
allowlist = _load_allowlist()
scalars = _all_scalars()
test_findings = _uncovered_tests(allowlist, _invoked_test_tokens(scalars))
integration_paths, ownership_findings = _integration_ownership()
test_findings = _uncovered_tests(allowlist, _invoked_test_tokens(scalars) | integration_paths) + ownership_findings
dockerfile_findings = _uncovered_dockerfiles(allowlist, _built_dockerfile_tokens(scalars))
stale_findings = _stale_allowlist_paths(allowlist, test_files=_test_files(), dockerfiles=_dockerfiles())

View file

@ -5,6 +5,7 @@ use serde_json::Value;
#[derive(Deserialize)]
struct Input {
path: String,
model_alias: String,
provider_model: String,
api_base: String,
@ -21,7 +22,8 @@ async fn main() {
Ok(input) => input,
Err(error) => fail(error),
};
let result = litellm_ai_gateway::trace_parity::traced_messages_request(
let result = litellm_ai_gateway::trace_parity::traced_request(
input.path,
input.model_alias,
input.provider_model,
input.api_base,

View file

@ -29,14 +29,15 @@ pub struct TracedGatewayResponse {
pub trace: Vec<litellm_core::observability::FunctionTraceEvent>,
}
pub async fn traced_messages_request(
pub async fn traced_request(
path: String,
model_alias: String,
provider_model: String,
api_base: String,
body: Value,
) -> TracedGatewayResponse {
let trace = litellm_core::observability::FunctionTrace::default();
let result = messages_request(model_alias, provider_model, api_base, body)
let result = request(path, model_alias, provider_model, api_base, body)
.with_subscriber(trace.dispatcher())
.await;
let events = trace.events();
@ -54,7 +55,8 @@ pub async fn traced_messages_request(
}
}
pub async fn messages_request(
pub async fn request(
path: String,
model_alias: String,
provider_model: String,
api_base: String,
@ -75,7 +77,7 @@ pub async fn messages_request(
};
let request = Request::builder()
.method("POST")
.uri("/v1/messages")
.uri(path)
.header(AUTHORIZATION, "Bearer trace-master-key")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(body.to_string()))

View file

@ -538,7 +538,7 @@ context_window_fallbacks: Optional[List] = None
content_policy_fallbacks: Optional[List] = None
allowed_fails: int = 3
allow_dynamic_callback_disabling: bool = True
num_retries_per_request: Optional[int] = None # for the request overall (incl. fallbacks + model retries)
num_retries_per_request: Optional[int] = None # cap on Router retries of one model group; resets per fallback hop
####### SECRET MANAGERS #####################
secret_manager_client: Optional[Any] = (
None # list of instantiated key management clients - e.g. azure kv, infisical, etc.

View file

@ -1,5 +1,6 @@
import ast
import contextvars
import functools
import logging
import os
import sys
@ -225,6 +226,35 @@ class AccessLogRedactionFilter(logging.Filter):
_access_log_filter: Final = AccessLogRedactionFilter()
@functools.lru_cache(maxsize=1)
def _parse_disabled_access_log_paths(raw: str) -> frozenset[str]:
return frozenset(stripped for path in raw.split(",") if (stripped := path.strip()))
def _disabled_access_log_paths() -> frozenset[str]:
"""Read the variable per record so a value loaded later via proxy config
environment_variables or dotenv is honored."""
return _parse_disabled_access_log_paths(os.getenv("LITELLM_DISABLE_ACCESS_LOG_PATHS", ""))
class AccessLogPathFilter(logging.Filter):
"""Drops uvicorn.access records for request paths listed in LITELLM_DISABLE_ACCESS_LOG_PATHS.
uvicorn passes record.args as (client_addr, method, full_path, http_version, status_code).
"""
def filter(self, record: logging.LogRecord) -> bool:
if not isinstance(record.args, tuple) or len(record.args) < 3:
return True
full_path: Final = record.args[2]
if not isinstance(full_path, str):
return True
return full_path.partition("?")[0] not in _disabled_access_log_paths()
_access_log_path_filter: Final = AccessLogPathFilter()
def _get_max_string_length_stdout_log() -> int:
"""Read the limit per record so a value loaded later via proxy config
environment_variables is honored."""
@ -663,6 +693,7 @@ def _redact_third_party_loggers() -> None:
for name in _REDACTED_THIRD_PARTY_LOGGERS:
logging.getLogger(name).addFilter(_secret_filter)
for name in _REDACTED_ACCESS_LOGGERS:
logging.getLogger(name).addFilter(_access_log_path_filter)
logging.getLogger(name).addFilter(_access_log_filter)

View file

@ -679,7 +679,14 @@ class Cache:
cache_key, cached_data, kwargs = self._add_cache_logic(result=result, **kwargs)
self.cache.set_cache(cache_key, cached_data, **kwargs)
except Exception as e:
log_redis_failure(verbose_logger, logging.ERROR, "LiteLLM Cache: exception in add_cache", e)
self._log_add_cache_failure(e)
def _log_add_cache_failure(self, exc: Exception) -> None:
message: Final = "LiteLLM Cache: exception in add_cache"
if isinstance(self.cache, RedisCache):
log_redis_failure(verbose_logger, logging.ERROR, message, exc)
return
verbose_logger.error("%s: %s", message, exc)
async def async_add_cache(self, result, dynamic_cache_object: BaseCache | None = None, **kwargs):
"""
@ -698,7 +705,7 @@ class Cache:
else:
await self.cache.async_set_cache(cache_key, cached_data, **kwargs)
except Exception as e:
log_redis_failure(verbose_logger, logging.ERROR, "LiteLLM Cache: exception in add_cache", e)
self._log_add_cache_failure(e)
def _convert_to_cached_embedding(
self,
@ -877,7 +884,7 @@ class Cache:
else:
await self.cache.async_set_cache_pipeline(cache_list=cache_list, **kwargs)
except Exception as e:
log_redis_failure(verbose_logger, logging.ERROR, "LiteLLM Cache: exception in add_cache", e)
self._log_add_cache_failure(e)
def should_use_cache(self, **kwargs):
"""

View file

@ -15,6 +15,7 @@ import hashlib
import inspect
import json
import logging
import threading
import time
from collections.abc import Awaitable, Callable, Iterator, Sequence
from contextvars import ContextVar
@ -32,6 +33,7 @@ from litellm.constants import (
REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD,
REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT,
REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION,
REDIS_TIMEOUT_LOG_INTERVAL,
)
from litellm.litellm_core_utils.core_helpers import _get_parent_otel_span_from_kwargs
from litellm.litellm_core_utils.coroutine_checker import coroutine_checker
@ -340,7 +342,7 @@ def _explicit_causes(exc: BaseException) -> Iterator[BaseException]:
current = current.__cause__
def _is_redis_timeout_failure(exc: BaseException) -> bool:
def is_redis_timeout_failure(exc: BaseException) -> bool:
"""True when ``exc`` or any exception it was explicitly raised ``from`` is a timeout.
redis-py's blocking pool reports a pool wait timeout as ``ConnectionError`` chained from
@ -414,7 +416,7 @@ def _record_swallowed_redis_failure(breaker: RedisCircuitBreaker, exc: BaseExcep
"""
if not _is_redis_health_failure(exc):
return
breaker.record_failure(is_timeout=_is_redis_timeout_failure(exc))
breaker.record_failure(is_timeout=is_redis_timeout_failure(exc))
_swallowed_redis_failures.set(_swallowed_redis_failures.get() + 1)
@ -422,13 +424,58 @@ class RedisCircuitBreakerOpenError(Exception):
pass
class _RedisTimeoutLogThrottle:
"""Admits one Redis timeout log line per interval and counts the timeouts it suppressed in between."""
def __init__(self, interval: float, clock: Callable[[], float] = time.monotonic) -> None:
self.interval = interval
self._clock = clock
self._lock = threading.Lock()
self._last_logged_at: float | None = None
self._suppressed = 0
def admit(self) -> int | None:
"""Return the number of timeouts suppressed since the last admitted line, or None to suppress this one."""
with self._lock:
now: Final = self._clock()
if self._last_logged_at is not None and now - self._last_logged_at < self.interval:
self._suppressed += 1
return None
suppressed: Final = self._suppressed
self._suppressed = 0
self._last_logged_at = now
return suppressed
_redis_timeout_log_throttle: Final = _RedisTimeoutLogThrottle(REDIS_TIMEOUT_LOG_INTERVAL)
def log_redis_failure(
logger: logging.Logger, level: int, message: str, exc: BaseException, with_traceback: bool = False
) -> None:
if isinstance(exc, RedisCircuitBreakerOpenError):
logger.debug("%s: %s", message, exc)
logger.debug("%s: %s", message, exc, stacklevel=2)
return
logger.log(level, "%s: %s", message, exc, exc_info=exc if with_traceback else None)
exc_info: Final = exc if with_traceback else None
if not is_redis_timeout_failure(exc):
logger.log(level, "%s: %s", message, exc, exc_info=exc_info, stacklevel=2)
return
suppressed: Final = _redis_timeout_log_throttle.admit()
if suppressed is None:
logger.debug("%s: %s", message, exc, stacklevel=2)
return
if suppressed == 0:
logger.log(level, "%s: %s", message, exc, exc_info=exc_info, stacklevel=2)
return
logger.log(
level,
"%s: %s (%d more Redis timeouts since the previous Redis timeout line were logged at DEBUG)",
message,
exc,
suppressed,
exc_info=exc_info,
stacklevel=2,
)
@dataclass(frozen=True, slots=True)
@ -475,7 +522,7 @@ async def _run_under_circuit_breaker(
result: Final = await call()
except Exception as e:
if _is_redis_health_failure(e):
breaker.record_failure(is_timeout=_is_redis_timeout_failure(e))
breaker.record_failure(is_timeout=is_redis_timeout_failure(e))
raise
_exit_circuit_breaker(breaker, admission)
return result
@ -492,7 +539,7 @@ def _run_under_circuit_breaker_sync(
result: Final = call()
except Exception as e:
if _is_redis_health_failure(e):
breaker.record_failure(is_timeout=_is_redis_timeout_failure(e))
breaker.record_failure(is_timeout=is_redis_timeout_failure(e))
raise
_exit_circuit_breaker(breaker, admission)
return result
@ -801,10 +848,8 @@ class RedisCache(BaseCache):
## LOGGING ##
end_time = time.time()
_duration = end_time - start_time
verbose_logger.error(
"LiteLLM Redis Caching: increment_cache() - Got exception from REDIS %s, Writing value=%s",
str(e),
value,
log_redis_failure(
verbose_logger, logging.ERROR, "LiteLLM Redis Caching: increment_cache() - Got exception from REDIS", e
)
raise e
@ -1010,11 +1055,8 @@ class RedisCache(BaseCache):
call_type=f"async_set_cache <- {_get_call_stack_info()}",
)
)
verbose_logger.error(
"LiteLLM Redis Caching: async set() - Got exception from REDIS %s, key=%r, value=%r",
str(e),
key,
value,
log_redis_failure(
verbose_logger, logging.ERROR, "LiteLLM Redis Caching: async set() - Got exception from REDIS", e
)
raise e
@ -1062,10 +1104,8 @@ class RedisCache(BaseCache):
event_metadata={"key": key},
)
)
verbose_logger.error(
"LiteLLM Redis Caching: async set() - Got exception from REDIS %s, Writing value=%s",
str(e),
value,
log_redis_failure(
verbose_logger, logging.ERROR, "LiteLLM Redis Caching: async set() - Got exception from REDIS", e
)
_record_swallowed_redis_failure(self._circuit_breaker, e)
@ -1112,7 +1152,6 @@ class RedisCache(BaseCache):
start_time: Final = time.time()
print_verbose(f"Set Async Redis Cache: key list: {cache_list}\nttl={ttl}, redis_version={self.redis_version}")
cache_value: Final = None
try:
async with _redis_client.pipeline(transaction=False) as pipe:
results: Final = await self._pipeline_helper(pipe, cache_list, ttl)
@ -1149,10 +1188,11 @@ class RedisCache(BaseCache):
)
)
verbose_logger.error(
"LiteLLM Redis Caching: async set_cache_pipeline() - Got exception from REDIS %s, Writing value=%s",
str(e),
cache_value,
log_redis_failure(
verbose_logger,
logging.ERROR,
"LiteLLM Redis Caching: async set_cache_pipeline() - Got exception from REDIS",
e,
)
_record_swallowed_redis_failure(self._circuit_breaker, e)
@ -1191,8 +1231,11 @@ class RedisCache(BaseCache):
end_time=time.time(),
)
)
verbose_logger.error(
"LiteLLM Redis Caching: async_set_cache_pipeline_with_ttls() - Got exception from REDIS %s", str(e)
log_redis_failure(
verbose_logger,
logging.ERROR,
"LiteLLM Redis Caching: async_set_cache_pipeline_with_ttls() - Got exception from REDIS",
e,
)
_record_swallowed_redis_failure(self._circuit_breaker, e)
@ -1235,10 +1278,8 @@ class RedisCache(BaseCache):
)
)
# NON blocking - notify users Redis is throwing an exception
verbose_logger.error(
"LiteLLM Redis Caching: async set() - Got exception from REDIS %s, Writing value=%s",
str(e),
value,
log_redis_failure(
verbose_logger, logging.ERROR, "LiteLLM Redis Caching: async set() - Got exception from REDIS", e
)
raise e
@ -1274,10 +1315,11 @@ class RedisCache(BaseCache):
)
)
# NON blocking - notify users Redis is throwing an exception
verbose_logger.error(
"LiteLLM Redis Caching: async set_cache_sadd() - Got exception from REDIS %s, Writing value=%s",
str(e),
value,
log_redis_failure(
verbose_logger,
logging.ERROR,
"LiteLLM Redis Caching: async set_cache_sadd() - Got exception from REDIS",
e,
)
_record_swallowed_redis_failure(self._circuit_breaker, e)
@ -1359,10 +1401,11 @@ class RedisCache(BaseCache):
parent_otel_span=parent_otel_span,
)
)
verbose_logger.error(
"LiteLLM Redis Caching: async async_increment() - Got exception from REDIS %s, Writing value=%s",
str(e),
value,
log_redis_failure(
verbose_logger,
logging.ERROR,
"LiteLLM Redis Caching: async async_increment() - Got exception from REDIS",
e,
)
raise e
@ -1448,7 +1491,9 @@ class RedisCache(BaseCache):
print_verbose(f"Got Redis Cache: key: {key}, cached_response {cached_response}")
return self._get_cache_logic(cached_response=cached_response)
except Exception as e:
verbose_logger.error("litellm.caching.caching: get() - Got exception from REDIS: %s", e)
log_redis_failure(
verbose_logger, logging.ERROR, "litellm.caching.caching: get() - Got exception from REDIS", e
)
_record_swallowed_redis_failure(self._circuit_breaker, e)
def _run_redis_mget_operation(self, keys: list[str]) -> Sequence[bytes | str | None]:
@ -1526,7 +1571,7 @@ class RedisCache(BaseCache):
end_time=failed_at,
parent_otel_span=parent_otel_span,
)
verbose_logger.error("Error occurred in batch get cache - %s", e)
log_redis_failure(verbose_logger, logging.ERROR, "Error occurred in batch get cache", e)
_record_swallowed_redis_failure(self._circuit_breaker, e)
return key_value_dict
@ -1645,7 +1690,7 @@ class RedisCache(BaseCache):
parent_otel_span=parent_otel_span,
)
)
verbose_logger.error("Error occurred in async batch get cache - %s", e)
log_redis_failure(verbose_logger, logging.ERROR, "Error occurred in async batch get cache", e)
_record_swallowed_redis_failure(self._circuit_breaker, e)
return key_value_dict
@ -1870,9 +1915,11 @@ class RedisCache(BaseCache):
parent_otel_span=_get_parent_otel_span_from_kwargs(kwargs),
)
)
verbose_logger.error(
"LiteLLM Redis Caching: async increment_pipeline() - Got exception from REDIS %s",
str(e),
log_redis_failure(
verbose_logger,
logging.ERROR,
"LiteLLM Redis Caching: async increment_pipeline() - Got exception from REDIS",
e,
)
raise e
@ -1949,7 +1996,7 @@ class RedisCache(BaseCache):
call_type=f"async_rpush <- {_get_call_stack_info()}",
)
)
verbose_logger.error("LiteLLM Redis Cache RPUSH: - Got exception from REDIS : %s", e)
log_redis_failure(verbose_logger, logging.ERROR, "LiteLLM Redis Cache RPUSH: - Got exception from REDIS", e)
raise e
async def _pipeline_rpush_helper(
@ -2017,9 +2064,11 @@ class RedisCache(BaseCache):
call_type=f"async_rpush_pipeline <- {_get_call_stack_info()}",
)
)
verbose_logger.error(
"LiteLLM Redis Caching: async_rpush_pipeline() - Got exception from REDIS %s",
str(e),
log_redis_failure(
verbose_logger,
logging.ERROR,
"LiteLLM Redis Caching: async_rpush_pipeline() - Got exception from REDIS",
e,
)
raise e
@ -2095,7 +2144,7 @@ class RedisCache(BaseCache):
call_type=f"async_lpop <- {_get_call_stack_info()}",
)
)
verbose_logger.error("LiteLLM Redis Cache LPOP: - Got exception from REDIS : %s", e)
log_redis_failure(verbose_logger, logging.ERROR, "LiteLLM Redis Cache LPOP: - Got exception from REDIS", e)
raise e
async def _pipeline_lpop_helper(
@ -2206,8 +2255,10 @@ class RedisCache(BaseCache):
call_type=f"async_lpop_pipeline <- {_get_call_stack_info()}",
)
)
verbose_logger.error(
"LiteLLM Redis Caching: async_lpop_pipeline() - Got exception from REDIS %s",
str(e),
log_redis_failure(
verbose_logger,
logging.ERROR,
"LiteLLM Redis Caching: async_lpop_pipeline() - Got exception from REDIS",
e,
)
raise e

View file

@ -227,6 +227,9 @@ PRE_CALL_EXECUTED_GUARDRAILS_KEY: Final = "_pre_call_executed_guardrails"
# Attribute stamped on log_guardrail_information wrappers so __init_subclass__ does not wrap them again
LOGS_GUARDRAIL_INFORMATION_MARKER: Final = "_litellm_logs_guardrail_information"
# llm_provider stamped on proxy-side rate limit errors when the model resolves to no deployment
PROXY_LLM_PROVIDER_FALLBACK: Final = "litellm_proxy"
# Generic fallback for unknown models
DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET: Final = int(
os.getenv("DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET", 128)
@ -311,6 +314,12 @@ REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS: Final = float(
# RFC 6455 caps the close frame payload at 125 bytes, 2 of which carry the status code
WEBSOCKET_CLOSE_REASON_MAX_BYTES: Final = 123
BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY: Final = "litellm.bedrock_realtime.pending_session_update"
BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY: Final = "litellm.bedrock_realtime.session_committed"
BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY: Final = "litellm.bedrock_realtime.committed_failure"
REALTIME_SESSION_SUCCESS_LOGGED_KEY: Final = "realtime_session_success_logged"
REALTIME_SESSION_FAILURE_LOGGED_KEY: Final = "realtime_session_failure_logged"
# SSL/TLS cipher configuration for faster handshakes
# Strategy: Strongly prefer fast modern ciphers, but allow fallback to commonly supported ones
# This balances performance with broad compatibility
@ -461,6 +470,7 @@ REDIS_CIRCUIT_BREAKER_ENABLED: Final = os.getenv("REDIS_CIRCUIT_BREAKER_ENABLED"
# minimum seconds a timeout-only failure streak must span before it can open the breaker,
# so one event-loop stall timing out many queued calls at once does not trip it
REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION: Final = float(os.getenv("REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION", 5.0))
REDIS_TIMEOUT_LOG_INTERVAL: Final = float(os.getenv("REDIS_TIMEOUT_LOG_INTERVAL", "5.0"))
# Seconds of idle before a Redis cluster connection is validated with a PING and
# reconnected if dead, so a connection silently dropped by a cluster restart
# (e.g. ElastiCache Serverless maintenance) is not reused while broken

View file

@ -97,6 +97,7 @@ from litellm.llms.vertex_ai.cost_calculator import cost_router as google_cost_ro
from litellm.llms.xai.cost_calculator import cost_per_token as xai_cost_per_token
from litellm.responses.utils import ResponseAPILoggingUtils
from litellm.types.agents import LiteLLMSendMessageResponse
from litellm.types.llms.base import CachedTokensDetails
from litellm.types.llms.openai import (
HttpxBinaryResponseContent,
ImageGenerationRequestQuality,
@ -2381,6 +2382,46 @@ def _summable_prompt_token_fields(prompt_tokens_details: BaseModel) -> list[str]
return [attr for attr in field_names if attr != "cache_creation_tokens"]
def _combine_cached_tokens_details(
current: CachedTokensDetails | None, new: CachedTokensDetails
) -> CachedTokensDetails:
def _sum_optional(current_value: int | None, new_value: int | None) -> int | None:
if current_value is None and new_value is None:
return None
return (current_value or 0) + (new_value or 0)
return CachedTokensDetails(
text_tokens=_sum_optional(current.text_tokens if current is not None else None, new.text_tokens),
audio_tokens=_sum_optional(current.audio_tokens if current is not None else None, new.audio_tokens),
image_tokens=_sum_optional(current.image_tokens if current is not None else None, new.image_tokens),
)
def _combine_prompt_tokens_details(
current: PromptTokensDetailsWrapper | None, new: PromptTokensDetailsWrapper
) -> PromptTokensDetailsWrapper:
base: Final = current if current is not None else PromptTokensDetailsWrapper()
base_values: Final = MappingProxyType(
{attr: getattr(base, attr) for attr in type(base).model_fields if hasattr(base, attr)}
)
summed: Final = MappingProxyType(
{
attr: (getattr(base, attr, 0) or 0) + (getattr(new, attr) or 0)
for attr in _summable_prompt_token_fields(new)
if hasattr(new, attr) and isinstance(getattr(new, attr) or 0, (int, float))
}
)
new_cached_tokens_details: Final = getattr(new, "cached_tokens_details", None)
cached_tokens_details: Final = (
_combine_cached_tokens_details(getattr(base, "cached_tokens_details", None), new_cached_tokens_details)
if isinstance(new_cached_tokens_details, CachedTokensDetails)
else getattr(base, "cached_tokens_details", None)
)
return PromptTokensDetailsWrapper(
**MappingProxyType({**base_values, **summed, "cached_tokens_details": cached_tokens_details})
)
class BaseTokenUsageProcessor:
@staticmethod
def combine_usage_objects(usage_objects: list[Usage]) -> Usage:
@ -2389,7 +2430,6 @@ class BaseTokenUsageProcessor:
"""
from litellm.types.utils import (
CompletionTokensDetailsWrapper,
PromptTokensDetailsWrapper,
Usage,
)
@ -2408,27 +2448,10 @@ class BaseTokenUsageProcessor:
and isinstance(current_val, (int, float))
):
setattr(combined, attr, current_val + new_val)
# Handle nested prompt_tokens_details
if hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details:
if not hasattr(combined, "prompt_tokens_details") or not combined.prompt_tokens_details:
combined.prompt_tokens_details = PromptTokensDetailsWrapper()
# Check what keys exist in the model's prompt_tokens_details
# Access model_fields on the class, not the instance, to avoid Pydantic 2.11+ deprecation warnings
for attr in _summable_prompt_token_fields(usage.prompt_tokens_details):
if (
hasattr(usage.prompt_tokens_details, attr)
and not attr.startswith("_")
and not callable(_attribute_value(usage.prompt_tokens_details, attr))
):
current_val = getattr(combined.prompt_tokens_details, attr, 0) or 0
new_val = getattr(usage.prompt_tokens_details, attr, 0) or 0
if new_val is not None and isinstance(new_val, (int, float)):
setattr(
combined.prompt_tokens_details,
attr,
current_val + new_val,
)
combined.prompt_tokens_details = _combine_prompt_tokens_details(
getattr(combined, "prompt_tokens_details", None), usage.prompt_tokens_details
)
# Handle nested completion_tokens_details
if hasattr(usage, "completion_tokens_details") and usage.completion_tokens_details:

View file

@ -379,7 +379,7 @@ class ArizePhoenixPromptManager(CustomPromptManagement):
"""Reload prompts from Arize Phoenix."""
if self.prompt_id:
self._prompt_manager = None # Reset to force reload
self.prompt_manager # This will trigger reload
_ = self.prompt_manager # access triggers lazy reload
def should_run_prompt_management(
self,

View file

@ -406,7 +406,7 @@ class BitBucketPromptManager(CustomPromptManagement):
"""Reload prompts from BitBucket."""
if self.prompt_id:
self._prompt_manager = None # Reset to force reload
self.prompt_manager # This will trigger reload
_ = self.prompt_manager # access triggers lazy reload
def should_run_prompt_management(
self,

View file

@ -884,7 +884,9 @@ class CustomGuardrail(CustomLogger):
"""logging_only: run apply_guardrail on copies of the logged request/response and record the verdict."""
from litellm.llms import get_guardrail_translation_mapping
if not self.uses_apply_guardrail_interface() or self.use_native_lifecycle_hooks:
if not self.uses_apply_guardrail_interface():
return kwargs, result
if not self._event_hook_is_event_type(GuardrailEventHooks.logging_only):
return kwargs, result
try:
translation: Final = get_guardrail_translation_mapping(CallTypes(call_type))()
@ -901,8 +903,18 @@ class CustomGuardrail(CustomLogger):
for key, value in (litellm_params.get("metadata") or {}).items()
if key != "standard_logging_guardrail_information"
}
response: Final = (
kwargs.get("async_complete_streaming_response") or kwargs.get("complete_streaming_response") or result
)
from litellm.types.utils import ModelResponse
output_translation: Final = (
get_guardrail_translation_mapping(CallTypes.acompletion)()
if isinstance(response, ModelResponse)
else translation
)
try:
await self._scan_logged_call(kwargs, result, translation, scratch_metadata)
await self._scan_logged_call(kwargs, response, translation, output_translation, scratch_metadata)
except Exception as e:
verbose_logger.warning("Guardrail %s: logging_only scan raised: %s", self.guardrail_name, e)
recorded: Final = scratch_metadata.get("standard_logging_guardrail_information")
@ -919,8 +931,9 @@ class CustomGuardrail(CustomLogger):
async def _scan_logged_call(
self,
kwargs: dict, # mutable-ok: CustomLogger.async_logging_hook contract
result: object,
response: object | None,
translation: "BaseTranslation",
output_translation: "BaseTranslation",
scratch_metadata: dict, # mutable-ok: apply_guardrail records its verdict into request metadata
) -> None:
optional_params: Final = kwargs.get("optional_params") or {}
@ -934,8 +947,10 @@ class CustomGuardrail(CustomLogger):
"metadata": scratch_metadata,
}
await translation.process_input_messages(data=scratch_request, guardrail_to_apply=self)
await translation.process_output_response(
response=copy.deepcopy(result), guardrail_to_apply=self, request_data=scratch_request
if response is None:
return
await output_translation.process_output_response(
response=copy.deepcopy(response), guardrail_to_apply=self, request_data=scratch_request
)
def supports_scan_only_tool_results(self) -> bool:

View file

@ -12,6 +12,7 @@ from typing import Final
from litellm.integrations.otel.mappers.base import AttributeMap, AttrValue, SpanData
from litellm.integrations.otel.mappers.utils import (
MAX_MESSAGE_ATTRS_PER_SPAN,
MAX_TOOL_DEFINITION_ATTRS_PER_SPAN,
collect,
drop_none,
@ -26,6 +27,8 @@ from litellm.integrations.otel.model.payloads import (
ToolDefinition,
)
_MAX_INDEXED_MESSAGES: Final = MAX_MESSAGE_ATTRS_PER_SPAN // 2
class OpenInferenceMapper:
"""Emits OpenInference attributes for LLM_CALL spans.
@ -84,27 +87,44 @@ class OpenInferenceMapper:
return {}
def _llm_call(self, data: LLMCallSpanData) -> AttributeMap:
outputs: Final = output_messages(data)
indexed_in, indexed_out = self._indexed_split(len(data.messages_in), len(outputs))
return {
**collect(self._LLM_CALL_ATTRS, data),
**collect(self._BLOB_ATTRS, data),
**self._messages("llm.input_messages", "input.value", data.messages_in),
**self._messages("llm.output_messages", "output.value", output_messages(data)),
**self._messages(
"llm.input_messages",
"input.value",
data.messages_in,
self._prompt_positions(len(data.messages_in), indexed_in),
),
**self._messages("llm.output_messages", "output.value", outputs, range(indexed_out)),
**self._tools(data),
}
@staticmethod
def _messages(prefix: str, value_key: str, messages: Sequence[object]) -> AttributeMap:
"""Per-message ``{prefix}.{idx}.message.*`` keys + the ``value_key`` blob."""
def _indexed_split(inputs: int, outputs: int) -> tuple[int, int]:
"""Prompt and response share one allowance; the response is reserved at least half of it."""
indexed_out: Final = min(outputs, max(_MAX_INDEXED_MESSAGES // 2, _MAX_INDEXED_MESSAGES - inputs))
return _MAX_INDEXED_MESSAGES - indexed_out, indexed_out
@staticmethod
def _prompt_positions(total: int, indexed: int) -> tuple[int, ...]:
"""Prompt messages that get per-index attributes: message 0 and the most recent turns."""
if total <= indexed:
return tuple(range(total))
return (0, *range(total - indexed + 1, total))
@staticmethod
def _messages(prefix: str, value_key: str, messages: Sequence[object], positions: Sequence[int]) -> AttributeMap:
"""``{prefix}.{idx}.message.*`` keys for the messages at ``positions`` + the ``value_key`` blob of all."""
parsed: Final = [(m.get("role") if isinstance(m, dict) else None, message_content(m)) for m in messages]
attrs: Final = drop_none(
{
key: value
for idx, (role, content) in enumerate(parsed)
for idx, (role, content) in ((idx, parsed[idx]) for idx in positions)
for key, value in (
(
f"{prefix}.{idx}.message.role",
role if isinstance(role, str) else None,
),
(f"{prefix}.{idx}.message.role", role if isinstance(role, str) else None),
(f"{prefix}.{idx}.message.content", content),
)
}

View file

@ -32,6 +32,14 @@ core telemetry no matter how many vocabularies are configured.
"""
MAX_MESSAGE_ATTRS_PER_SPAN: Final = DEFAULT_SPAN_ATTRIBUTE_LIMIT // 8
"""Span-wide ceiling on per-index chat message attributes, prompt and response together.
An eighth is the largest share that still fits beside the tool ceiling and the core
of every vocabulary at once. The complete conversation still rides the JSON blobs.
"""
def tool_attr_budget(vocabularies: int) -> int:
"""Split the span-wide tool-definition ceiling across active vocabularies."""
return MAX_TOOL_DEFINITION_ATTRS_PER_SPAN // max(vocabularies, 1)

View file

@ -16,6 +16,7 @@ from pydantic import BaseModel
import litellm
from litellm._logging import print_verbose, verbose_logger
from litellm.constants import PROXY_LLM_PROVIDER_FALLBACK
from litellm.exceptions import (
validate_rate_limit_category,
validate_rate_limit_type,
@ -2581,6 +2582,15 @@ class PrometheusLogger(CustomLogger):
)
return None
@staticmethod
def _extract_api_provider_from_exception(exception: Exception) -> str | None:
if not isinstance(exception, litellm.exceptions.RateLimitError):
return None
llm_provider: Final = exception.llm_provider
if not llm_provider or llm_provider == PROXY_LLM_PROVIDER_FALLBACK:
return None
return llm_provider
async def async_post_call_failure_hook(
self,
request_data: dict,
@ -2616,7 +2626,9 @@ class PrometheusLogger(CustomLogger):
_metadata: Final = request_data.get("metadata", {}) or {}
model_id: Final = _metadata.get("model_info", {}).get("id") or request_data.get("model_info", {}).get("id")
rate_limit_category, rate_limit_type = self._extract_rate_limit_labels(original_exception)
api_provider: Final = self._extract_api_provider_from_request_data(request_data)
api_provider: Final = self._extract_api_provider_from_request_data(
request_data
) or self._extract_api_provider_from_exception(original_exception)
enum_values: Final = UserAPIKeyLabelValues(
end_user=user_api_key_dict.end_user_id,
user=user_api_key_dict.user_id,

View file

@ -4,18 +4,10 @@ imported_openAIResponse = True
try:
import io
import logging
import sys
from typing import Any, TypeVar
from typing import Any, Literal, Protocol, TypeVar
from wandb.sdk.data_types import trace_tree
if sys.version_info >= (3, 8):
from typing import Literal, Protocol
else:
from typing import Literal
from typing_extensions import Protocol
logger: Final = logging.getLogger(__name__)
K = TypeVar("K", bound=str)

View file

@ -303,6 +303,16 @@ def get_metadata_variable_name_from_kwargs(
return "litellm_metadata" if "litellm_metadata" in kwargs else "metadata"
def max_retries_per_request_hit(kwargs: Mapping[str, object], num_retries_per_request: int | None) -> bool:
if num_retries_per_request is None:
return False
metadata: Final = kwargs.get(get_metadata_variable_name_from_kwargs(kwargs))
if not isinstance(metadata, Mapping):
return False
attempted_retries: Final = metadata.get("attempted_retries")
return type(attempted_retries) is int and 0 < attempted_retries and num_retries_per_request <= attempted_retries
def get_or_create_metadata_bucket(
request_data: dict,
) -> tuple[Literal["metadata", "litellm_metadata"], dict]:

View file

@ -36,7 +36,7 @@ class CoroutineChecker:
target = callback
if not inspect.isfunction(target) and not inspect.ismethod(target):
try:
call_attr: Final = getattr(target, "__call__", None)
call_attr: Final = getattr(target, "__call__", None) # noqa: B004 # value unwrap so iscoroutinefunction sees through functors
if call_attr is not None:
target = call_attr
except Exception:

View file

@ -34,6 +34,11 @@ rules never mix the two and never use ``extends``. A rule whose
Rules are only consulted after exact and case-insensitive lookups miss, so an
exact cost-map entry always takes precedence over any rule.
Rules flagged with ``fill_missing_for_providers: [..]`` also fill only keys
missing from an exact cost-map entry when the entry's ``litellm_provider`` is
listed, while values already present on the entry win on conflict. Only flagged
capability rules participate in this fill; routing rules never do.
Patterns are matched case-insensitively with ``re.search`` and are not implicitly
anchored: a rule must include ``^`` and ``$`` to bind to the whole model name,
otherwise it matches as a substring. Keeping anchoring in the regex makes the rule
@ -46,17 +51,19 @@ Rules are compiled and classified once, at install time. The match functions are
O(number of rules); callers must only invoke them on a cache miss.
"""
import logging
import re
from collections.abc import Mapping
from dataclasses import dataclass
from typing import Final
from litellm._logging import verbose_logger
verbose_logger: Final = logging.getLogger("LiteLLM")
NAME_FIELD: Final = "name"
PATTERN_FIELD: Final = "pattern"
MODEL_INFO_FIELD: Final = "model_info"
PROVIDER_KEY: Final = "litellm_provider"
LEGACY_EXTENDS_FIELD: Final = "extends"
FILL_MISSING_FOR_PROVIDERS_FIELD: Final = "fill_missing_for_providers"
def _resolve_legacy_extends(rules: list) -> list:
@ -98,11 +105,28 @@ class _RoutingRule:
class _CapabilityRule:
pattern: re.Pattern
model_info: dict
fill_missing_for_providers: frozenset[str]
_CompiledRule = _RoutingRule | _CapabilityRule
def _parse_fill_missing_for_providers(rule: Mapping[str, object], pattern_label: object) -> frozenset[str] | None:
if FILL_MISSING_FOR_PROVIDERS_FIELD not in rule:
return frozenset()
raw_fill_missing_for_providers: Final = rule.get(FILL_MISSING_FOR_PROVIDERS_FIELD)
if not isinstance(raw_fill_missing_for_providers, (list, tuple)) or not all(
isinstance(provider, str) for provider in raw_fill_missing_for_providers
):
verbose_logger.warning(
"LiteLLM: skipping malformed fallback generalization rule %s ('%s' must be a list of provider strings).",
rule.get(NAME_FIELD, pattern_label),
FILL_MISSING_FOR_PROVIDERS_FIELD,
)
return None
return frozenset(raw_fill_missing_for_providers)
def _compile_rule(rule: object) -> tuple[_CompiledRule, ...]:
if not isinstance(rule, dict):
return ()
@ -125,8 +149,17 @@ def _compile_rule(rule: object) -> tuple[_CompiledRule, ...]:
e,
)
return ()
fill_missing_for_providers: Final = _parse_fill_missing_for_providers(rule, pattern)
if fill_missing_for_providers is None:
return ()
if PROVIDER_KEY not in model_info:
return (_CapabilityRule(pattern=compiled, model_info=model_info),)
return (
_CapabilityRule(
pattern=compiled,
model_info=model_info,
fill_missing_for_providers=fill_missing_for_providers,
),
)
provider: Final = model_info[PROVIDER_KEY]
if not isinstance(provider, str):
verbose_logger.warning(
@ -140,7 +173,11 @@ def _compile_rule(rule: object) -> tuple[_CompiledRule, ...]:
return (_RoutingRule(pattern=compiled, provider=provider),)
return (
_RoutingRule(pattern=compiled, provider=provider),
_CapabilityRule(pattern=compiled, model_info=model_info),
_CapabilityRule(
pattern=compiled,
model_info=model_info,
fill_missing_for_providers=fill_missing_for_providers,
),
)
@ -151,6 +188,7 @@ class _FallbackGeneralizations:
self.rules: list = []
self.routing_rules: tuple = ()
self.capability_rules: tuple = ()
self.fill_missing_rules: tuple[_CapabilityRule, ...] = ()
def set_rules(self, rules: list | None) -> None:
installed: Final = rules if isinstance(rules, list) else []
@ -158,6 +196,7 @@ class _FallbackGeneralizations:
self.rules = installed
self.routing_rules = tuple(rule for rule in compiled if isinstance(rule, _RoutingRule))
self.capability_rules = tuple(rule for rule in compiled if isinstance(rule, _CapabilityRule))
self.fill_missing_rules = tuple(rule for rule in self.capability_rules if rule.fill_missing_for_providers)
def match_routing(self, model: str) -> str | None:
if not model:
@ -175,6 +214,21 @@ class _FallbackGeneralizations:
return None
return {key: value for model_info in matched for key, value in model_info.items()}
def match_fill_missing(self, model: str, provider: str) -> Mapping[str, object] | None:
if not model or not provider:
return None
matched = tuple(
rule.model_info
for rule in self.fill_missing_rules
if provider in rule.fill_missing_for_providers and rule.pattern.search(model) is not None
)
if not matched:
return None
fill_missing: Final[Mapping[str, object]] = {
key: value for model_info in matched for key, value in model_info.items() if key != PROVIDER_KEY
}
return fill_missing or None
_registry: Final = _FallbackGeneralizations()
@ -210,3 +264,14 @@ def match_capability_generalizations(model: str) -> dict | None:
capability rule matches. O(number of rules); only call once exact lookups have missed.
"""
return _registry.match_capabilities(model)
def match_fill_missing_generalizations(model: str, provider: str) -> Mapping[str, object] | None:
"""Return flagged capability rules matching ``model`` for ``provider``.
Later rules override earlier ones on key conflicts. Only rules listing
``provider`` in ``fill_missing_for_providers`` contribute. Returns ``None``
when no flagged rule matches. O(number of rules); only call once exact
lookups have matched.
"""
return _registry.match_fill_missing(model, provider)

View file

@ -556,7 +556,7 @@ class Logging(LiteLLMLoggingBaseClass):
# ids leaking into a different, later request on the same thread. Sync
# support is deferred to a follow-up PR with its own safe-restore
# mechanism; async calls (the proxy's only call path) are unaffected.
if supports_correlation_logging:
if supports_correlation_logging and litellm.request_correlation_in_logs:
set_trace_id(self.litellm_trace_id)
set_session_id(self.litellm_session_id)
# set_trace_id()/set_session_id() sanitize (strip control chars, bound
@ -2442,7 +2442,7 @@ class Logging(LiteLLMLoggingBaseClass):
call) would leave the outer request's subsequent log lines stamped with
the nested call's trace_id/session_id instead of its own.
Uses a plain set() of the captured pre-call value rather than
Uses a plain contextvar set() of the captured pre-call value rather than
contextvars.Token-based reset(), since this can end up called from a
different asyncio Task/context than __init__ ran in (e.g. the request
task's own wrapper() finally block, plus async_success_handler
@ -2453,8 +2453,8 @@ class Logging(LiteLLMLoggingBaseClass):
that Task's view of the contextvars, so calling it multiple times
(once per Task involved in this attempt) is required, not just safe.
"""
set_trace_id(self._pre_call_trace_id)
set_session_id(self._pre_call_session_id)
trace_id_var.set(self._pre_call_trace_id)
session_id_var.set(self._pre_call_session_id)
def _restore_correlation_context_if_unclaimed(self) -> None:
"""Guarded variant for __del__-triggered cleanup only.

View file

@ -9,6 +9,8 @@ from types import MappingProxyType
from typing import Any, Final, Literal, TypedDict, cast
from zoneinfo import ZoneInfo, ZoneInfoNotFoundError
from typing_extensions import ReadOnly
import litellm
from litellm._internal_context import current_billing_time
from litellm._logging import verbose_logger
@ -772,6 +774,7 @@ def calculate_cache_writing_cost(
class PromptTokensDetailsResult(TypedDict):
cache_hit_tokens: int
cache_hit_audio_tokens: ReadOnly[int]
cache_creation_tokens: int
cache_creation_token_details: CacheCreationTokenDetails | None
text_tokens: int
@ -802,12 +805,34 @@ def parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult:
)
or None
)
text_tokens: Final = (
cast(int | None, getattr(usage.prompt_tokens_details, "text_tokens", None))
or 0 # default to prompt tokens, if this field is not set
cached_tokens_details: Final = getattr(usage.prompt_tokens_details, "cached_tokens_details", None)
cached_audio_tokens: Final = min(
_get_token_detail_value(cached_tokens_details, "audio_tokens") or 0, cache_hit_tokens
)
cached_text_tokens: Final = min(
_get_token_detail_value(cached_tokens_details, "text_tokens") or 0,
cache_hit_tokens - cached_audio_tokens,
)
cached_image_tokens: Final = min(
_get_token_detail_value(cached_tokens_details, "image_tokens") or 0,
cache_hit_tokens - cached_audio_tokens - cached_text_tokens,
)
text_tokens: Final = max(
(
cast(int | None, getattr(usage.prompt_tokens_details, "text_tokens", None))
or 0 # default to prompt tokens, if this field is not set
)
- cached_text_tokens,
0,
)
audio_tokens: Final = max(
(cast(int | None, getattr(usage.prompt_tokens_details, "audio_tokens", 0)) or 0) - cached_audio_tokens,
0,
)
image_tokens: Final = max(
(cast(int | None, getattr(usage.prompt_tokens_details, "image_tokens", 0)) or 0) - cached_image_tokens,
0,
)
audio_tokens: Final = cast(int | None, getattr(usage.prompt_tokens_details, "audio_tokens", 0)) or 0
image_tokens: Final = cast(int | None, getattr(usage.prompt_tokens_details, "image_tokens", 0)) or 0
video_tokens: Final = _coerce_token_count(getattr(usage.prompt_tokens_details, "video_tokens", 0))
character_count: Final = (
cast(
@ -835,6 +860,7 @@ def parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult:
return PromptTokensDetailsResult(
cache_hit_tokens=cache_hit_tokens,
cache_hit_audio_tokens=cached_audio_tokens,
cache_creation_tokens=cache_creation_tokens,
cache_creation_token_details=cache_creation_token_details,
text_tokens=text_tokens,
@ -918,7 +944,16 @@ def _calculate_input_cost(
prompt_cost = float(prompt_tokens_details["text_tokens"]) * prompt_base_cost
### CACHE READ COST - Now uses tiered pricing
prompt_cost += float(prompt_tokens_details["cache_hit_tokens"]) * cache_read_cost
cache_hit_audio_tokens: Final = prompt_tokens_details["cache_hit_audio_tokens"]
audio_cache_read_rate: Final = _get_cost_per_unit(
model_info,
_get_service_tier_cost_key("cache_read_input_audio_token_cost", service_tier),
None,
)
prompt_cost += float(prompt_tokens_details["cache_hit_tokens"] - cache_hit_audio_tokens) * cache_read_cost
prompt_cost += float(cache_hit_audio_tokens) * (
audio_cache_read_rate if audio_cache_read_rate is not None else cache_read_cost
)
### AUDIO COST
if prompt_tokens_details["audio_tokens"]:
@ -1149,6 +1184,7 @@ def generic_cost_per_token(
### PROCESSING COST
prompt_tokens_details = PromptTokensDetailsResult(
cache_hit_tokens=0,
cache_hit_audio_tokens=0,
cache_creation_tokens=0,
cache_creation_token_details=None,
text_tokens=usage.prompt_tokens,
@ -1319,6 +1355,7 @@ class BilledTokenRates:
input_cost_per_token: float
output_cost_per_token: float
cache_read_input_token_cost: float
cache_read_input_audio_token_cost: float
cache_creation_input_token_cost: float
cache_creation_input_token_cost_above_1hr: float
output_cost_per_reasoning_token: float
@ -1330,6 +1367,7 @@ class BilledTokenRates:
input_cost_per_token=self.input_cost_per_token * multiplier,
output_cost_per_token=self.output_cost_per_token * multiplier,
cache_read_input_token_cost=self.cache_read_input_token_cost * multiplier,
cache_read_input_audio_token_cost=self.cache_read_input_audio_token_cost * multiplier,
cache_creation_input_token_cost=self.cache_creation_input_token_cost * multiplier,
cache_creation_input_token_cost_above_1hr=self.cache_creation_input_token_cost_above_1hr * multiplier,
output_cost_per_reasoning_token=self.output_cost_per_reasoning_token * multiplier,
@ -1353,15 +1391,16 @@ def _reasoning_token_count(usage: Usage) -> int:
return parsed or _coerce_token_count(getattr(usage, "reasoning_tokens", 0))
def _cache_token_counts(usage: Usage) -> tuple[int, int, CacheCreationTokenDetails | None]:
"""(cache read tokens, cache creation tokens, cache creation details): read from prompt_tokens_details
first, then the private top-level counters the Usage constructor mirrors cache tokens onto for
providers/callers that bypass the details."""
def _cache_token_counts(usage: Usage) -> tuple[int, int, int, CacheCreationTokenDetails | None]:
"""(cache read tokens, cached audio tokens, cache creation tokens, cache creation details): read from
prompt_tokens_details first, then the private top-level counters the Usage constructor mirrors cache
tokens onto for providers/callers that bypass the details."""
parsed: Final = parse_prompt_tokens_details(usage) if usage.prompt_tokens_details is not None else None
parsed_read: Final = parsed["cache_hit_tokens"] if parsed is not None else 0
parsed_creation: Final = parsed["cache_creation_tokens"] if parsed is not None else 0
return (
parsed_read or _coerce_token_count(getattr(usage, "_cache_read_input_tokens", 0)),
parsed["cache_hit_audio_tokens"] if parsed is not None else 0,
parsed_creation or _coerce_token_count(getattr(usage, "_cache_creation_input_tokens", 0)),
parsed["cache_creation_token_details"] if parsed is not None else None,
)
@ -1372,11 +1411,13 @@ def _custom_pricing_rates(custom_cost_per_token: CostPerToken) -> BilledTokenRat
cache rates (else the input rate) and reasoning at the output rate, as _cost_per_token_custom_pricing_helper does."""
input_rate: Final = custom_cost_per_token["input_cost_per_token"]
output_rate: Final = custom_cost_per_token["output_cost_per_token"]
cache_read_rate: Final = custom_cost_per_token.get("cache_read_input_token_cost", input_rate)
cache_creation_rate: Final = custom_cost_per_token.get("cache_creation_input_token_cost", input_rate)
return BilledTokenRates(
input_cost_per_token=input_rate,
output_cost_per_token=output_rate,
cache_read_input_token_cost=custom_cost_per_token.get("cache_read_input_token_cost", input_rate),
cache_read_input_token_cost=cache_read_rate,
cache_read_input_audio_token_cost=cache_read_rate,
cache_creation_input_token_cost=cache_creation_rate,
cache_creation_input_token_cost_above_1hr=cache_creation_rate,
output_cost_per_reasoning_token=output_rate,
@ -1413,6 +1454,11 @@ def _cost_map_billed_rates(
completion_base_cost=completion_base_cost,
current_time=billing_time,
)
audio_cache_read_rate: Final = _get_cost_per_unit(
model_info,
_get_service_tier_cost_key("cache_read_input_audio_token_cost", service_tier),
None,
)
multiplier: Final = (
_get_regional_uplift_multiplier(model_info, data_residency)
* get_vertex_regional_endpoint_uplift(model_info, vertex_location)
@ -1422,6 +1468,9 @@ def _cost_map_billed_rates(
input_cost_per_token=prompt_base_cost,
output_cost_per_token=completion_base_cost,
cache_read_input_token_cost=cache_read_cost_rate,
cache_read_input_audio_token_cost=(
audio_cache_read_rate if audio_cache_read_rate is not None else cache_read_cost_rate
),
cache_creation_input_token_cost=cache_creation_cost_rate,
cache_creation_input_token_cost_above_1hr=cache_creation_cost_above_1hr_rate,
output_cost_per_reasoning_token=reasoning_rate,
@ -1494,7 +1543,9 @@ def get_token_type_cost_breakdown(
if rates is None:
return TokenTypeCostBreakdown(0.0, 0.0, 0.0)
cache_read_tokens, cache_creation_tokens, cache_creation_token_details = _cache_token_counts(usage)
cache_read_tokens, cached_audio_tokens, cache_creation_tokens, cache_creation_token_details = _cache_token_counts(
usage
)
cache_creation_cost: Final = (
float(cache_creation_tokens) * rates.cache_creation_input_token_cost
if custom_cost_per_token is not None
@ -1507,7 +1558,10 @@ def get_token_type_cost_breakdown(
)
return TokenTypeCostBreakdown(
reasoning_cost=float(_reasoning_token_count(usage)) * rates.output_cost_per_reasoning_token,
cache_read_cost=float(cache_read_tokens) * rates.cache_read_input_token_cost,
cache_read_cost=(
float(cache_read_tokens - cached_audio_tokens) * rates.cache_read_input_token_cost
+ float(cached_audio_tokens) * rates.cache_read_input_audio_token_cost
),
cache_creation_cost=cache_creation_cost,
rates=rates,
)

View file

@ -1757,7 +1757,7 @@ def convert_to_anthropic_tool_invoke(
anthropic_tool_invoke: Final[list[AnthropicMessagesToolUseParam | dict[str, object]]] = []
for tool in tool_calls:
if not get_attribute_or_key(tool, "type") == "function":
if get_attribute_or_key(tool, "type") != "function":
continue
tool_id = cast(str, get_attribute_or_key(tool, "id"))

View file

@ -10,6 +10,7 @@ from typing_extensions import ReadOnly
import litellm
from litellm._logging import redact_internal_details_from_client_message, verbose_logger
from litellm.constants import REALTIME_SESSION_FAILURE_LOGGED_KEY, REALTIME_SESSION_SUCCESS_LOGGED_KEY
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig
from litellm.types.llms.openai import (
@ -35,9 +36,6 @@ else:
CLIENT_CONNECTION_CLASS = Any
REALTIME_SESSION_SUCCESS_LOGGED_KEY: Final = "realtime_session_success_logged"
@dataclass(frozen=True, slots=True)
class BackendClose:
code: int
@ -1153,6 +1151,7 @@ class RealTimeStreaming:
self._logging_worker.ensure_initialized_and_enqueue(
self.logging_obj.dispatch_failure_handlers(error, traceback.format_exc(), prefer_async_handlers=True)
)
self.logging_obj.model_call_details[REALTIME_SESSION_FAILURE_LOGGED_KEY] = True
@staticmethod
def _detect_beta_header(websocket: ScopedWebSocket) -> bool:

View file

@ -454,7 +454,7 @@ def token_counter(
params: Final = _MessageCountParams(model, custom_tokenizer)
num_tokens = _count_messages(params, new_messages, use_default_image_token_count, default_token_count)
if count_response_tokens is False:
includes_system_message: Final = any([message.get("role", None) == "system" for message in new_messages])
includes_system_message: Final = any(message.get("role", None) == "system" for message in new_messages)
num_tokens += _count_extra(params.count_function, tools, tool_choice, includes_system_message)
else:

View file

@ -0,0 +1,388 @@
from __future__ import annotations
import hashlib
import json
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from itertools import accumulate
from types import MappingProxyType
from typing import Annotated, Final, Literal, Protocol, TypeAlias
import httpx
from pydantic import BaseModel, ConfigDict, Field, JsonValue, StrictInt, TypeAdapter, ValidationError
import litellm
from litellm.llms.anthropic.common_utils import AnthropicModelInfo, is_anthropic_oauth_key
from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler
from litellm.llms.anthropic.experimental_pass_through.messages.transformation import DEFAULT_ANTHROPIC_API_VERSION
from litellm.types.router import LiteLLM_Params
from litellm.types.utils import ModelResponse
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
_HEADERS: Final = TypeAdapter(dict[str, str])
_counter: Final = AnthropicCountTokensHandler()
_NATIVE_HEADERS: Final = frozenset(
(
"host",
"accept",
"accept-encoding",
"connection",
"user-agent",
"content-length",
"content-type",
"x-api-key",
"anthropic-version",
)
)
_DEPLOYMENT_OPTIONS: Final = frozenset(
{
"model",
"api_key",
"api_base",
"custom_llm_provider",
"rpm",
"tpm",
"timeout",
"stream_timeout",
"max_retries",
"num_retries",
"max_parallel_requests",
"input_cost_per_token",
"output_cost_per_token",
"cache_read_input_token_cost",
"cache_creation_input_token_cost",
"cache_creation_input_token_cost_above_1hr",
}
)
class _StrictModel(BaseModel):
model_config = ConfigDict(extra="forbid", frozen=True, strict=True)
class _CacheControl(_StrictModel):
type: Literal["ephemeral"]
ttl: Literal["5m", "1h"] = "5m"
class _Text(_StrictModel):
type: Literal["text"]
text: str = Field(min_length=1, pattern=r"\S")
cache_control: _CacheControl | None = None
class _ToolUse(_StrictModel):
type: Literal["tool_use"]
id: str = Field(min_length=1)
name: str = Field(min_length=1)
input: Mapping[str, JsonValue]
cache_control: _CacheControl | None = None
class _ResultText(_StrictModel):
type: Literal["text"]
text: str
class _ToolResult(_StrictModel):
type: Literal["tool_result"]
tool_use_id: str = Field(min_length=1)
content: str | Annotated[tuple[_ResultText, ...], Field(strict=False)]
is_error: bool | None = None
cache_control: _CacheControl | None = None
_Block: TypeAlias = Annotated[_Text | _ToolUse | _ToolResult, Field(discriminator="type")]
class _Message(_StrictModel):
role: Literal["user", "assistant"]
content: str | Annotated[tuple[_Block, ...], Field(strict=False)]
def blocks(self) -> tuple[_Text | _ToolUse | _ToolResult, ...]:
return (_Text(type="text", text=self.content),) if isinstance(self.content, str) else tuple(self.content)
class _Tool(_StrictModel):
name: str = Field(min_length=1)
description: str | None = None
input_schema: Mapping[str, JsonValue]
type: Literal["custom"] | None = None
class _Request(_StrictModel):
messages: tuple[_Message, ...] = Field(min_length=1, strict=False)
system: str | Annotated[tuple[_ResultText, ...], Field(strict=False)] | None = None
tools: Annotated[tuple[_Tool, ...], Field(strict=False)] | None = None
model: str | None = None
max_tokens: int | None = None
stream: bool | None = None
temperature: float | int | None = None
top_p: float | int | None = None
top_k: int | None = None
stop_sequences: Annotated[tuple[str, ...], Field(strict=False)] | None = None
metadata: Mapping[str, JsonValue] | None = None
@dataclass(frozen=True, slots=True)
class PromptPrefix:
prefix_body: Mapping[str, JsonValue]
fingerprint: str
fingerprints: tuple[str, ...]
ttl_seconds: int
def _digest(value: object) -> str:
return hashlib.sha256(
json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False).encode()
).hexdigest()
def _next_digest(previous: str, boundary: tuple[int, str, Mapping[str, JsonValue]]) -> str:
return _digest((previous, boundary))
def parse_prompt(body: Mapping[str, JsonValue]) -> PromptPrefix | None:
try:
request: Final = _Request.model_validate(body)
blocks: Final = tuple(message.blocks() for message in request.messages)
except ValidationError:
return None
markers: Final = tuple(
(message_index, block_index, block.cache_control)
for message_index, message_blocks in enumerate(blocks)
for block_index, block in enumerate(message_blocks)
if block.cache_control is not None
)
if len(markers) != 1:
return None
message_end, block_end, marker = markers[0]
normalized: Final = _JSON_OBJECT.validate_python(request.model_dump(mode="json", exclude_none=True))
context: Final = MappingProxyType({key: normalized[key] for key in ("system", "tools") if key in normalized})
boundaries: Final = tuple(
(
message_index,
request.messages[message_index].role,
_JSON_OBJECT.validate_python(
block.model_dump(mode="json", exclude=MappingProxyType({"cache_control": True}), exclude_none=True)
),
)
for message_index, message_blocks in enumerate(blocks[: message_end + 1])
for block_index, block in enumerate(message_blocks)
if message_index < message_end or block_index <= block_end
)
hashes: Final = tuple(
accumulate(boundaries, _next_digest, initial=_digest((_JSON_OBJECT.validate_python(context), marker.ttl)))
)[1:]
prefix_messages: Final = tuple(
_Message(
role=request.messages[message_index].role,
content=tuple(
block
for block_index, block in enumerate(message_blocks)
if message_index < message_end or block_index <= block_end
),
)
for message_index, message_blocks in enumerate(blocks[: message_end + 1])
)
return PromptPrefix(
prefix_body=MappingProxyType(
_JSON_OBJECT.validate_python(
_Request(messages=prefix_messages, system=request.system, tools=request.tools).model_dump(
mode="json", exclude_none=True
)
)
),
fingerprint=hashes[-1],
fingerprints=tuple(reversed(hashes[-20:])),
ttl_seconds=3600 if marker.ttl == "1h" else 300,
)
def cache_scope(
caller_key_hash: str,
deployment_id: str,
provider_key: str,
model: str,
anthropic_version: str = DEFAULT_ANTHROPIC_API_VERSION,
) -> str:
return _digest((caller_key_hash, deployment_id, provider_key, model, anthropic_version))
class _TTLUsage(BaseModel):
model_config = ConfigDict(strict=True)
ephemeral_5m_input_tokens: int = Field(default=0, ge=0)
ephemeral_1h_input_tokens: int = Field(default=0, ge=0)
class _CacheUsage(BaseModel):
model_config = ConfigDict(strict=True)
cached_tokens: int = Field(default=0, ge=0)
cache_creation_tokens: int = Field(default=0, ge=0)
cache_creation_token_details: _TTLUsage | None = None
class _Usage(BaseModel):
model_config = ConfigDict(strict=True)
prompt_tokens: int = Field(ge=0)
prompt_tokens_details: _CacheUsage
class _Choice(BaseModel):
finish_reason: str = Field(min_length=1)
class _Response(BaseModel):
model_config = ConfigDict(strict=True)
model: str
usage: _Usage
choices: tuple[_Choice, ...] = Field(min_length=1, strict=False)
class _CountBody(BaseModel):
messages: Sequence[Mapping[str, JsonValue]]
tools: Sequence[Mapping[str, JsonValue]] | None = None
system: str | Sequence[Mapping[str, JsonValue]] | None = None
class _CountResult(BaseModel):
input_tokens: Annotated[StrictInt, Field(ge=0)]
class TokenCounter(Protocol):
async def __call__(self, model: str, api_key: str, body: Mapping[str, JsonValue]) -> int | None: ...
def _count_objects(
values: Sequence[Mapping[str, JsonValue]],
) -> list[dict[str, JsonValue]]: # mutable-ok: the existing provider count API requires JSON lists/dicts
return [dict(value) for value in values] # mutable-ok: serialize read-only inputs at the provider API boundary
async def count_prompt_tokens(model: str, api_key: str, body: Mapping[str, JsonValue]) -> int | None:
native: Final = _CountBody.model_validate(body)
try:
result: Final = _CountResult.model_validate(
await _counter.handle_count_tokens_request(
model=model,
messages=_count_objects(native.messages),
tools=_count_objects(native.tools) if native.tools is not None else None,
system=native.system,
api_key=api_key,
timeout=15.0,
)
)
except Exception: # noqa: BLE001 # provider/count validation failures are unavailable estimates, not zero tokens
return None
return result.input_tokens
@dataclass(frozen=True, slots=True)
class NativePredictionTarget:
model: str
api_key: str
@dataclass(frozen=True, slots=True)
class UnsupportedPredictionTarget:
reason: Literal[
"unsupported_deployment_configuration",
"unsupported_provider_endpoint",
"unsupported_provider",
"unsupported_provider_credentials",
]
def resolve_prediction_target(params: LiteLLM_Params) -> NativePredictionTarget | UnsupportedPredictionTarget:
configured_options: Final = frozenset(params.model_dump(exclude_defaults=True, exclude_none=True))
if configured_options - _DEPLOYMENT_OPTIONS:
return UnsupportedPredictionTarget("unsupported_deployment_configuration")
api_base: Final = AnthropicModelInfo.get_api_base(params.api_base)
if api_base not in ("https://api.anthropic.com", "https://api.anthropic.com/v1/messages"):
return UnsupportedPredictionTarget("unsupported_provider_endpoint")
try:
model, provider, _, _ = litellm.get_llm_provider(
model=params.model, custom_llm_provider=params.custom_llm_provider
)
except Exception: # noqa: BLE001 # the shared provider resolver raises for unknown deployments
return UnsupportedPredictionTarget("unsupported_provider")
if provider != "anthropic":
return UnsupportedPredictionTarget("unsupported_provider")
api_key: Final = AnthropicModelInfo.get_api_key(params.api_key)
if api_key is None or not _supported_provider_key(api_key):
return UnsupportedPredictionTarget("unsupported_provider_credentials")
return NativePredictionTarget(model=model, api_key=api_key)
def _supported_provider_key(api_key: str) -> bool:
return bool(api_key) and not is_anthropic_oauth_key(api_key)
def supported_prediction_headers(headers: Mapping[str, str]) -> bool:
return all(
name.lower() != "anthropic-beta"
and (name.lower() != "anthropic-version" or value == DEFAULT_ANTHROPIC_API_VERSION)
for name, value in headers.items()
)
@dataclass(frozen=True, slots=True)
class ObservedCachePrefix:
prefix: PromptPrefix
scope: str
cached_tokens: int
cache_creation_tokens: int
def parse_observed_cache(
wire: httpx.Request, response_obj: ModelResponse, caller_key_hash: str, deployment_id: str
) -> ObservedCachePrefix | None:
try:
response: Final = _Response.model_validate(response_obj, from_attributes=True)
body: Final = _JSON_OBJECT.validate_json(wire.content)
headers: Final = _HEADERS.validate_python(wire.headers)
except (ValidationError, RuntimeError, httpx.RequestNotRead):
return None
if (
wire.url.scheme != "https"
or wire.url.host != "api.anthropic.com"
or wire.url.path != "/v1/messages"
or wire.url.query
or wire.url.port not in (None, 443)
):
return None
if (
frozenset(headers) - _NATIVE_HEADERS
or not supported_prediction_headers(headers)
or headers.get("anthropic-version") != DEFAULT_ANTHROPIC_API_VERSION
):
return None
provider_key: Final = headers.get("x-api-key", "")
model: Final = body.get("model")
if not _supported_provider_key(provider_key) or not isinstance(model, str) or model != response.model:
return None
prefix: Final = parse_prompt(body)
if prefix is None:
return None
usage: Final = response.usage.prompt_tokens_details
cache_tokens: Final = usage.cached_tokens + usage.cache_creation_tokens
if cache_tokens <= 0 or cache_tokens > response.usage.prompt_tokens:
return None
split: Final = usage.cache_creation_token_details
if usage.cache_creation_tokens and split is None:
return None
if split is not None and (
split.ephemeral_5m_input_tokens + split.ephemeral_1h_input_tokens != usage.cache_creation_tokens
or (prefix.ttl_seconds == 300 and split.ephemeral_1h_input_tokens > 0)
or (prefix.ttl_seconds == 3600 and split.ephemeral_5m_input_tokens > 0)
):
return None
return ObservedCachePrefix(
prefix=prefix,
scope=cache_scope(caller_key_hash, deployment_id, provider_key, model),
cached_tokens=cache_tokens,
cache_creation_tokens=usage.cache_creation_tokens,
)

View file

@ -144,10 +144,13 @@ class AzureFoundryModelInfo(BaseLLMModelInfo):
def get_api_key(api_key: str | None = None) -> str | None:
return api_key or litellm.api_key or get_secret_str("AZURE_AI_API_KEY")
@staticmethod
def get_api_version(api_version: str | None = None) -> str | None:
return api_version or litellm.api_version or get_secret_str("AZURE_API_VERSION")
@property
def api_version(self, api_version: str | None = None) -> str | None:
api_version = api_version or litellm.api_version or get_secret_str("AZURE_API_VERSION")
return api_version
def api_version(self) -> str | None:
return AzureFoundryModelInfo.get_api_version()
def get_token_counter(self) -> BaseTokenCounter | None:
"""

View file

@ -7,13 +7,21 @@ This uses aws_sdk_bedrock_runtime for bidirectional streaming with Nova Sonic.
import asyncio
import contextlib
import json
from collections.abc import AsyncIterator, Mapping
from typing import Final, Protocol
from collections.abc import AsyncIterator, Mapping, MutableMapping
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final, NoReturn, Protocol
from pydantic import JsonValue, TypeAdapter
import litellm
from litellm._logging import _redact_string, verbose_proxy_logger
from litellm.constants import (
BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY,
BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY,
BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY,
REALTIME_SESSION_SUCCESS_LOGGED_KEY,
)
from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
@ -28,6 +36,32 @@ from .transformation import BedrockRealtimeConfig
_CLIENT_MODALITIES_ADAPTER: Final[TypeAdapter["list[str] | None"]] = TypeAdapter(list[str] | None)
_CLIENT_MESSAGE_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
_EMPTY_JSON_OBJECT: Final[Mapping[str, JsonValue]] = MappingProxyType({})
_BEDROCK_STREAM_ERROR_STATUS: Final[Mapping[str, int]] = MappingProxyType(
{
"AccessDeniedException": 403,
"ConflictException": 400,
"InternalServerException": 500,
"ModelErrorException": 424,
"ModelNotReadyException": 429,
"ModelStreamErrorException": 424,
"ModelTimeoutException": 408,
"ResourceNotFoundException": 404,
"ServiceQuotaExceededException": 400,
"ServiceUnavailableException": 503,
"ThrottlingException": 429,
"ValidationException": 400,
}
)
def _as_bedrock_error(error: BaseException) -> BaseException:
status_code: Final = _BEDROCK_STREAM_ERROR_STATUS.get(type(error).__name__)
if status_code is None:
return error
return BedrockError(status_code=status_code, message=f"{type(error).__name__}: {error}")
def _json_dict(value: JsonValue) -> dict[str, JsonValue]:
return value if isinstance(value, dict) else {}
@ -51,6 +85,8 @@ def _should_log_event(openai_message: Mapping[str, object]) -> bool:
class RealtimeClientWebSocket(Protocol):
"""The client-facing websocket surface the realtime bridge talks to."""
scope: MutableMapping[str, object] # mutable-ok: the ASGI scope is the per-connection state store
async def receive_text(self) -> str: ...
async def send_text(self, data: str) -> None: ...
@ -85,6 +121,81 @@ class BedrockBidirectionalStream(Protocol):
async def await_output(self) -> tuple[object, BedrockOutputStream]: ...
@dataclass(frozen=True, slots=True)
class _BridgeOutcome:
logged_events: tuple[OpenAIRealtimeEvents, ...]
provider_failure: BaseException | None
client_disconnected: bool
async def _client_messages(client_ws: RealtimeClientWebSocket, initial_message: str | None) -> AsyncIterator[str]:
if initial_message is not None:
yield initial_message
while True:
try:
yield await client_ws.receive_text()
except Exception as e: # noqa: BLE001 # any receive failure means the client is gone
verbose_proxy_logger.debug("Client to Bedrock forwarding ended: %s", e, exc_info=True)
return
def _pending_session_update(scope: Mapping[str, object]) -> str | None:
"""A fallback attempt on the same websocket replays the session.update the failed attempt never acked."""
if scope.get(BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY) is True:
committed_failure: Final = scope.get(BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY)
raise BedrockError(
status_code=400,
message=(
"Bedrock realtime session already committed to a provider stream; it cannot be replayed"
+ (f". The committed stream failed with: {committed_failure}" if committed_failure else "")
),
)
pending: Final = scope.get(BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY)
return pending if isinstance(pending, str) else None
def _raise_provider_failure(scope: MutableMapping[str, object], failure: BaseException) -> NoReturn:
error: Final = _as_bedrock_error(failure)
verbose_proxy_logger.error("Bedrock Realtime: provider stream failed: %s", _redact_string(str(error)))
if scope.get(BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY) is True:
scope[BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY] = _redact_string(str(error))
raise error from failure
def _parse_client_message(message: str) -> Mapping[str, JsonValue]:
try:
return _json_dict(_CLIENT_MESSAGE_ADAPTER.validate_json(message))
except ValueError:
return _EMPTY_JSON_OBJECT
async def _ack_session_update(
client_ws: RealtimeClientWebSocket,
bedrock_stream: BedrockBidirectionalStream,
transformation_config: BedrockRealtimeConfig,
model: str,
logging_obj: LiteLLMLogging | None,
parsed_client_message: Mapping[str, JsonValue],
) -> bool:
"""Ack the client's session.update once Bedrock accepted the stream; False means the client is gone."""
await bedrock_stream.await_output()
client_ws.scope.pop(BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY, None)
client_ws.scope[BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY] = True # rebind-ok: scope outlives the attempt
if logging_obj is None:
return True
requested_modalities: Final = _CLIENT_MODALITIES_ADAPTER.validate_python(
_json_dict(parsed_client_message.get("session")).get("modalities")
)
try:
await client_ws.send_text(
json.dumps(transformation_config.session_updated_event(model, logging_obj, requested_modalities))
)
except Exception as e: # noqa: BLE001 # any send failure means the client is gone
verbose_proxy_logger.debug("Client to Bedrock forwarding ended: %s", e, exc_info=True)
return False
return True
class BedrockRealtime(BaseAWSLLM):
"""Handler for Bedrock Nova Sonic realtime speech-to-speech API."""
@ -132,6 +243,8 @@ class BedrockRealtime(BaseAWSLLM):
except ImportError:
raise ImportError("Missing aws_sdk_bedrock_runtime. Install with: pip install aws-sdk-bedrock-runtime")
pending_session_update: Final = _pending_session_update(websocket.scope)
# Get AWS region
if aws_region_name is None:
optional_params: Final = {
@ -190,90 +303,106 @@ class BedrockRealtime(BaseAWSLLM):
transformation_config: Final = BedrockRealtimeConfig()
try:
# Initialize the bidirectional stream
bedrock_stream: Final = await open_bidirectional_stream()
bedrock_stream: Final = await open_bidirectional_stream()
verbose_proxy_logger.debug("Bedrock Realtime: Bidirectional stream established")
verbose_proxy_logger.debug("Bedrock Realtime: Bidirectional stream established")
if pending_session_update is None:
await websocket.send_text(json.dumps(transformation_config.session_created_event(model, logging_obj)))
verbose_proxy_logger.debug("Bedrock Realtime: sent session.created to client on connect")
# Track state for transformation
session_state: Final[RealtimeResponseTransformInput] = {
"current_output_item_id": None,
"current_response_id": None,
"current_conversation_id": None,
"current_delta_chunks": None,
"current_item_chunks": None,
"current_delta_type": None,
"session_configuration_request": None,
}
# Track state for transformation
session_state: Final[RealtimeResponseTransformInput] = {
"current_output_item_id": None,
"current_response_id": None,
"current_conversation_id": None,
"current_delta_chunks": None,
"current_item_chunks": None,
"current_delta_type": None,
"session_configuration_request": None,
}
# Create tasks for bidirectional forwarding
client_to_bedrock_task: Final = asyncio.create_task(
self._forward_client_to_bedrock(
websocket,
bedrock_stream,
transformation_config,
model,
session_state,
logging_obj,
outcome: Final = await self._bridge(
websocket,
bedrock_stream,
transformation_config,
model,
session_state,
logging_obj,
initial_message=pending_session_update,
)
logged_events: Final = (
*outcome.logged_events,
*(
leftover_event
for leftover_event in transformation_config.leftover_usage_done_events()
if _should_log_event(leftover_event)
),
)
if logged_events:
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(
logging_obj.dispatch_success_handlers(
list(logged_events), # mutable-ok: realtime spend logging requires a list result
prefer_async_handlers=True,
)
)
logging_obj.model_call_details[REALTIME_SESSION_SUCCESS_LOGGED_KEY] = True
async def forward_bedrock_and_collect_logged_events() -> tuple[OpenAIRealtimeEvents, ...]:
return tuple(
[
event
async for event in self._forward_bedrock_to_client(
bedrock_stream,
websocket,
transformation_config,
model,
logging_obj,
session_state,
)
]
)
bedrock_to_client_task: Final = asyncio.create_task(forward_bedrock_and_collect_logged_events())
# Wait for both tasks to complete
await asyncio.gather(
client_to_bedrock_task,
bedrock_to_client_task,
return_exceptions=True,
if outcome.provider_failure is None:
return
if outcome.client_disconnected:
verbose_proxy_logger.debug(
"Bedrock Realtime: stream failed after the client disconnected: %s", outcome.provider_failure
)
return
_raise_provider_failure(websocket.scope, outcome.provider_failure)
forwarded_logged_events: Final = (
bedrock_to_client_task.result()
if not bedrock_to_client_task.cancelled() and bedrock_to_client_task.exception() is None
else ()
)
logged_events: Final = (
*forwarded_logged_events,
*(
leftover_event
for leftover_event in transformation_config.leftover_usage_done_events()
if _should_log_event(leftover_event)
),
)
if logged_events:
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(
logging_obj.dispatch_success_handlers(
list(logged_events), # mutable-ok: realtime spend logging requires a list result
prefer_async_handlers=True,
)
)
async def _bridge(
self,
websocket: RealtimeClientWebSocket,
bedrock_stream: BedrockBidirectionalStream,
transformation_config: BedrockRealtimeConfig,
model: str,
session_state: RealtimeResponseTransformInput,
logging_obj: LiteLLMLogging,
initial_message: str | None,
) -> _BridgeOutcome:
"""Run both forwarding directions until the client leaves or either side fails."""
logged: Final[list[OpenAIRealtimeEvents]] = [] # mutable-ok: events forwarded before a failure are still spend
except Exception as e:
verbose_proxy_logger.exception("Error in BedrockRealtime.async_realtime: %s", e)
try:
await websocket.close(code=1011, reason=_redact_string(f"Internal error: {e}"))
except Exception:
pass
raise
async def collect_logged_events() -> None:
async for event in self._forward_bedrock_to_client(
bedrock_stream, websocket, transformation_config, model, logging_obj, session_state
):
logged.append(event)
client_task: Final = asyncio.create_task(
self._forward_client_to_bedrock(
websocket, bedrock_stream, transformation_config, model, session_state, logging_obj, initial_message
)
)
bedrock_task: Final = asyncio.create_task(collect_logged_events())
await asyncio.wait((client_task, bedrock_task), return_when=asyncio.FIRST_COMPLETED)
client_disconnected: Final = (
client_task.done() and not client_task.cancelled() and client_task.exception() is None
)
client_task.cancel()
bedrock_task.cancel()
client_outcome, bedrock_outcome = await asyncio.gather(client_task, bedrock_task, return_exceptions=True)
return _BridgeOutcome(
logged_events=tuple(logged),
provider_failure=(
client_outcome
if isinstance(client_outcome, Exception)
else bedrock_outcome
if isinstance(bedrock_outcome, Exception)
else None
),
client_disconnected=client_disconnected,
)
async def _forward_client_to_bedrock(
self,
@ -283,8 +412,12 @@ class BedrockRealtime(BaseAWSLLM):
model: str,
session_state: RealtimeResponseTransformInput,
logging_obj: LiteLLMLogging | None = None,
):
"""Forward messages from client WebSocket to Bedrock stream."""
initial_message: str | None = None,
) -> None:
"""Forward messages from client WebSocket to Bedrock stream.
Returns once the client is gone; provider failures (input stream or readiness) propagate to the caller.
"""
from aws_sdk_bedrock_runtime.models import (
BidirectionalInputPayloadPart,
InvokeModelWithBidirectionalStreamInputChunk,
@ -299,41 +432,28 @@ class BedrockRealtime(BaseAWSLLM):
verbose_proxy_logger.debug("Bedrock Realtime: Sent to Bedrock: %s", bedrock_message[:200])
try:
while True:
# Receive message from client
message = await client_ws.receive_text()
async for message in _client_messages(client_ws, initial_message):
verbose_proxy_logger.debug("Bedrock Realtime: Received from client: %s", message[:200])
parsed_client_message = _parse_client_message(message)
is_session_update = _json_str(parsed_client_message.get("type")) == "session.update"
if is_session_update:
client_ws.scope[BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY] = (
message # rebind-ok: scope outlives the attempt
)
# Transform OpenAI format to Bedrock format
transformed_messages = transformation_config.transform_realtime_request(
message=message,
model=model,
session_configuration_request=session_state.get("session_configuration_request"),
)
# Send transformed messages to Bedrock
for bedrock_message in transformed_messages:
await send_to_bedrock(bedrock_message)
if logging_obj is not None:
client_message_type: str | None = None
requested_modalities: list[str] | None = None
with contextlib.suppress(Exception):
parsed_client_message = _json_dict(_CLIENT_MESSAGE_ADAPTER.validate_json(message))
client_message_type = _json_str(parsed_client_message.get("type"))
if client_message_type == "session.update":
requested_modalities = _CLIENT_MODALITIES_ADAPTER.validate_python(
_json_dict(parsed_client_message.get("session")).get("modalities")
)
if client_message_type == "session.update":
await client_ws.send_text(
json.dumps(
transformation_config.session_updated_event(model, logging_obj, requested_modalities)
)
)
except Exception as e:
verbose_proxy_logger.debug("Client to Bedrock forwarding ended: %s", e, exc_info=True)
if is_session_update and not await _ack_session_update(
client_ws, bedrock_stream, transformation_config, model, logging_obj, parsed_client_message
):
break
finally:
for close_message in transformation_config.session_close_messages():
with contextlib.suppress(Exception):
await send_to_bedrock(close_message)
@ -349,68 +469,71 @@ class BedrockRealtime(BaseAWSLLM):
logging_obj: LiteLLMLogging,
session_state: RealtimeResponseTransformInput,
) -> AsyncIterator[OpenAIRealtimeEvents]:
"""Forward messages from Bedrock to the client, yielding the ones to record for spend logging."""
try:
while True:
# Receive from Bedrock
output = await bedrock_stream.await_output()
result = await output[1].receive()
"""Forward messages from Bedrock to the client, yielding the ones to record for spend logging.
if result is None:
verbose_proxy_logger.debug("Bedrock Realtime: Bedrock stream ended")
break
Provider failures propagate to the caller; the client websocket is only closed on a normal stream end.
"""
payload_bytes = result.value.bytes_ if result.value else None
if payload_bytes:
bedrock_response = payload_bytes.decode("utf-8")
verbose_proxy_logger.debug("Bedrock Realtime: Received from Bedrock: %s", bedrock_response[:200])
# Transform Bedrock format to OpenAI format
realtime_response_transform_input: RealtimeResponseTransformInput = {
"current_output_item_id": session_state.get("current_output_item_id"),
"current_response_id": session_state.get("current_response_id"),
"current_conversation_id": session_state.get("current_conversation_id"),
"current_delta_chunks": session_state.get("current_delta_chunks"),
"current_item_chunks": session_state.get("current_item_chunks"),
"current_delta_type": session_state.get("current_delta_type"),
"session_configuration_request": session_state.get("session_configuration_request"),
}
transformed_response = transformation_config.transform_realtime_response(
message=bedrock_response,
model=model,
logging_obj=logging_obj,
realtime_response_transform_input=realtime_response_transform_input,
)
# Update session state
session_state.update(
{
"current_output_item_id": transformed_response.get("current_output_item_id"),
"current_response_id": transformed_response.get("current_response_id"),
"current_conversation_id": transformed_response.get("current_conversation_id"),
"current_delta_chunks": transformed_response.get("current_delta_chunks"),
"current_item_chunks": transformed_response.get("current_item_chunks"),
"current_delta_type": transformed_response.get("current_delta_type"),
"session_configuration_request": transformed_response.get("session_configuration_request"),
}
)
# Send transformed messages to client
response_value = transformed_response["response"]
openai_messages = response_value if isinstance(response_value, list) else (response_value,)
for openai_message in openai_messages:
message_json = json.dumps(openai_message)
await client_ws.send_text(message_json)
verbose_proxy_logger.debug("Bedrock Realtime: Sent to client: %s", message_json[:200])
if _should_log_event(openai_message):
yield openai_message
except Exception as e:
verbose_proxy_logger.debug("Bedrock to client forwarding ended: %s", e, exc_info=True)
finally:
# Close the client WebSocket
async def send_to_client(message_json: str) -> bool:
try:
await client_ws.close()
except Exception:
pass
await client_ws.send_text(message_json)
except Exception as e: # noqa: BLE001 # any send failure means the client is gone
verbose_proxy_logger.debug("Bedrock to client forwarding ended: %s", e, exc_info=True)
return False
verbose_proxy_logger.debug("Bedrock Realtime: Sent to client: %s", message_json[:200])
return True
output: Final = await bedrock_stream.await_output()
while True:
result = await output[1].receive()
if result is None:
verbose_proxy_logger.debug("Bedrock Realtime: Bedrock stream ended")
with contextlib.suppress(Exception):
await client_ws.close()
return
payload_bytes = result.value.bytes_ if result.value else None
if payload_bytes:
bedrock_response = payload_bytes.decode("utf-8")
verbose_proxy_logger.debug("Bedrock Realtime: Received from Bedrock: %s", bedrock_response[:200])
# Transform Bedrock format to OpenAI format
realtime_response_transform_input: RealtimeResponseTransformInput = {
"current_output_item_id": session_state.get("current_output_item_id"),
"current_response_id": session_state.get("current_response_id"),
"current_conversation_id": session_state.get("current_conversation_id"),
"current_delta_chunks": session_state.get("current_delta_chunks"),
"current_item_chunks": session_state.get("current_item_chunks"),
"current_delta_type": session_state.get("current_delta_type"),
"session_configuration_request": session_state.get("session_configuration_request"),
}
transformed_response = transformation_config.transform_realtime_response(
message=bedrock_response,
model=model,
logging_obj=logging_obj,
realtime_response_transform_input=realtime_response_transform_input,
)
# Update session state
session_state.update(
{
"current_output_item_id": transformed_response.get("current_output_item_id"),
"current_response_id": transformed_response.get("current_response_id"),
"current_conversation_id": transformed_response.get("current_conversation_id"),
"current_delta_chunks": transformed_response.get("current_delta_chunks"),
"current_item_chunks": transformed_response.get("current_item_chunks"),
"current_delta_type": transformed_response.get("current_delta_type"),
"session_configuration_request": transformed_response.get("session_configuration_request"),
}
)
# Send transformed messages to client
response_value = transformed_response["response"]
openai_messages = response_value if isinstance(response_value, list) else (response_value,)
for openai_message in openai_messages:
if not await send_to_client(json.dumps(openai_message)):
return
if _should_log_event(openai_message):
yield openai_message

View file

@ -2592,7 +2592,9 @@ def _complete_custom_openai(
copilot_headers.update(extra_headers)
extra_headers = copilot_headers
if extra_headers is not None:
use_base_llm_http_handler: Final = get_secret_bool("EXPERIMENTAL_OPENAI_BASE_LLM_HTTP_HANDLER")
if extra_headers is not None and not use_base_llm_http_handler:
optional_params["extra_headers"] = extra_headers
if litellm.enable_preview_features and metadata is not None: # [PREVIEW] allow metadata to be passed to OPENAI
@ -2609,8 +2611,6 @@ def _complete_custom_openai(
optional_params[k] = v
## COMPLETION CALL
use_base_llm_http_handler: Final = get_secret_bool("EXPERIMENTAL_OPENAI_BASE_LLM_HTTP_HANDLER")
try:
if use_base_llm_http_handler:
response = base_llm_http_handler.completion(

File diff suppressed because it is too large Load diff

View file

@ -11503,6 +11503,12 @@
"description": "Enable content moderation to check for harmful content (harassment, hate speech, etc.).",
"title": "Content Moderation Check"
},
"contextual_grounding_from_messages": {
"default": false,
"description": "ApplyGuardrail: when True, post-call scans of a request with no grounding_source / query content parts send the system and developer messages as the grounding source and the latest user message as the query, so the guardrail's contextual grounding policy can score the response. Bedrock bills contextual grounding units for these scans and rejects queries, sources and responses over its contextual grounding length limits, so leave this off for guardrails without a contextual grounding policy. Default False: plain messages are never sent as grounding context.",
"title": "Contextual Grounding From Messages",
"type": "boolean"
},
"credentials": {
"anyOf": [
{

View file

@ -887,6 +887,7 @@ class LiteLLMRoutes(enum.Enum):
"/auto_router/validate_complexity_router_config",
# Per-session auto-router read - the endpoint scopes the row to the caller's own key hash
"/auto_router/session",
"/cost/predict-cache",
# Agent registry - reads are role-scoped and writes are proxy-admin-gated
# inside agent_endpoints/endpoints.py
*agent_management_routes,
@ -2635,9 +2636,13 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
None,
description="max file size in MB for /v1/files uploads, for any purpose, if a file is larger than this size it will be rejected before being forwarded to the provider",
)
allowed_file_extensions: tuple[str, ...] | None = Field(
None,
description="the only file extensions (e.g. ['.jsonl', '.pdf', '.txt']) accepted on /v1/files uploads, for any purpose, matched case-insensitively against the uploaded filename. Files with any other extension, or none, are rejected. An empty list rejects every upload. Unset means no allowlist is applied",
)
blocked_file_extensions: tuple[str, ...] | None = Field(
None,
description="file extensions (e.g. ['.exe', '.sh']) rejected on /v1/files uploads, for any purpose, matched case-insensitively against the uploaded filename",
description="file extensions (e.g. ['.exe', '.sh']) rejected on /v1/files uploads, for any purpose, matched case-insensitively against the uploaded filename. Deprecated in favour of allowed_file_extensions; still enforced, after the allowlist, when set",
)
max_response_size_mb: int | None = Field(
None,
@ -4414,6 +4419,9 @@ class TeamInfoResponseObjectTeamTable(LiteLLM_TeamTable):
access_group_mcp_server_ids: list[str] | None = None
access_group_agent_ids: list[str] | None = None
access_group_details: tuple[TeamAccessGroupModelGrant, ...] | None = None
# Parent org's model ceiling, reported only to callers who can manage the team.
# None = no org or not a manager; [] or ["all-proxy-models"] = no ceiling.
organization_models: list[str] | None = None
class TeamInfoResponseObject(TypedDict):
@ -4837,6 +4845,14 @@ class JWTIssuerConfig(BaseModel):
default=None,
description="Issuer-specific claim path to normalize into LiteLLM's end-user id.",
)
virtual_key_claim_field: str | None = Field(
default=None,
description="Issuer-specific claim path used for the virtual key mapping lookup. Falls back to the global field.",
)
unregistered_jwt_client_behavior: UnregisteredJWTClientBehavior | None = Field(
default=None,
description="Issuer-specific policy when the virtual key claim has no mapping. Falls back to the global policy.",
)
model_config = {
"extra": "forbid",
@ -5063,6 +5079,28 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase):
super().__init__(**kwargs)
def get_issuer_config(self, issuer: str | None) -> JWTIssuerConfig | None:
if issuer is None or self.issuers is None:
return None
return next((config for config in self.issuers if config.issuer == issuer), None)
def is_virtual_key_mapping_configured(self) -> bool:
if self.virtual_key_claim_field is not None:
return True
return any(config.virtual_key_claim_field is not None for config in self.issuers or ())
def get_virtual_key_claim_field(self, issuer: str | None) -> str | None:
issuer_config: Final = self.get_issuer_config(issuer)
if issuer_config is not None and issuer_config.virtual_key_claim_field is not None:
return issuer_config.virtual_key_claim_field
return self.virtual_key_claim_field
def get_unregistered_jwt_client_behavior(self, issuer: str | None) -> UnregisteredJWTClientBehavior:
issuer_config: Final = self.get_issuer_config(issuer)
if issuer_config is not None and issuer_config.unregistered_jwt_client_behavior is not None:
return issuer_config.unregistered_jwt_client_behavior
return self.unregistered_jwt_client_behavior
class PrismaCompatibleUpdateDBModel(TypedDict, total=False):
model_name: str

View file

@ -123,6 +123,7 @@ from litellm.repositories.table_repositories import (
from litellm.repositories.team_repository import TeamRepository
from litellm.repositories.user_repository import UserRepository
from litellm.router import Router
from litellm.types.proxy.auth.auth_checks import UserNotFoundError
from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget
from litellm.utils import get_utc_datetime
@ -327,9 +328,23 @@ _safe_json_loads_obj: Final = _typed_json_loads(safe_json_loads)
last_db_access_time: Final = LimitedSizeOrderedDict(max_size=100)
db_cache_expiry: Final = DEFAULT_IN_MEMORY_TTL # refresh every 5s
_TEAM_MEMBERSHIP_INFLIGHT_MAX: Final = 10000
_team_membership_inflight: Final = LimitedSizeOrderedDict(max_size=_TEAM_MEMBERSHIP_INFLIGHT_MAX)
class _TeamMembershipCacheMiss:
__slots__ = ()
_TEAM_MEMBERSHIP_CACHE_MISS: Final = _TeamMembershipCacheMiss()
all_routes: Final = LiteLLMRoutes.openai_routes.value + LiteLLMRoutes.management_routes.value
def _membership_from_shared_load(result: object) -> LiteLLM_TeamMembership | None:
return result if isinstance(result, LiteLLM_TeamMembership) else None
def _log_budget_lookup_failure(entity: str, error: Exception) -> None:
"""
Log a warning when budget lookup fails; cache will not be populated.
@ -880,6 +895,7 @@ async def common_checks(
request_query_params=_safe_get_request_query_params(request=request),
llm_router=llm_router,
request=request,
team_id=valid_token.team_id if valid_token is not None else None,
)
skip_all_budget_checks: Final = skip_budget_checks or (
@ -887,6 +903,22 @@ async def common_checks(
and (route in MODEL_DISCOVERY_ROUTES or not RouteChecks.is_llm_api_route(route=route))
)
membership_user_id: Final = (
valid_token.user_id if valid_token is not None and (bool(_model) or not skip_all_budget_checks) else None
)
team_membership_loaded: Final = team_object is not None and membership_user_id is not None
loaded_team_membership: Final = (
await get_team_membership(
user_id=membership_user_id,
team_id=team_object.team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
if team_object is not None and membership_user_id is not None
else None
)
unpriced_models: Final = (
_unpriced_models_in_request(model=_model, llm_router=llm_router)
if litellm.block_requests_for_models_without_pricing and RouteChecks.is_llm_api_route(route=route)
@ -936,6 +968,8 @@ async def common_checks(
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
team_membership=loaded_team_membership,
team_membership_loaded=team_membership_loaded,
)
# Require trace id for agent keys when agent has require_trace_id_on_calls_by_agent
@ -987,6 +1021,8 @@ async def common_checks(
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
team_membership=loaded_team_membership,
team_membership_loaded=team_membership_loaded,
)
# Run before apply_key_tags_pre_auth injects key metadata.tags into request_body.
@ -1096,6 +1132,8 @@ async def common_checks(
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
team_membership=loaded_team_membership,
team_membership_loaded=team_membership_loaded,
),
_check_end_user_budget(end_user_obj=end_user_object, route=route)
if end_user_object is not None and end_user_object.litellm_budget_table is not None
@ -2141,7 +2179,76 @@ async def get_tag_object(
return tag_objects.get(tag_name)
def _membership_from_cached_payload(
cached: object,
) -> LiteLLM_TeamMembership | None | _TeamMembershipCacheMiss:
if cached is None:
return _TEAM_MEMBERSHIP_CACHE_MISS
if cached == NO_TEAM_MEMBERSHIP_SENTINEL:
return None
cached_membership: Final = CacheCodec.deserialize(cached, model_type=LiteLLM_TeamMembership)
return cached_membership if cached_membership is not None else _TEAM_MEMBERSHIP_CACHE_MISS
@log_db_metrics
async def _fetch_team_membership_from_db(
user_id: str,
team_id: str,
prisma_client: PrismaClient,
user_api_key_cache: UserApiKeyCache,
parent_otel_span: Span | None = None,
proxy_logging_obj: ProxyLogging | None = None,
) -> LiteLLM_TeamMembership | None:
_ = parent_otel_span, proxy_logging_obj
response: Final = await _dictable_table(TeamMembershipRepository(prisma_client)).find_unique(
where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}},
include={"litellm_budget_table": True},
)
membership: Final = None if response is None else LiteLLM_TeamMembership.model_validate(response.dict())
_key: Final = team_membership_reservation_cache_key(user_id=user_id, team_id=team_id)
if membership is None:
await user_api_key_cache.async_set_cache(
key=_key,
value=NO_TEAM_MEMBERSHIP_SENTINEL,
ttl=get_management_object_ttl(user_api_key_cache),
)
else:
await user_api_key_cache.async_set_cache(
key=_key,
value=membership,
model_type=LiteLLM_TeamMembership,
)
return membership
async def _load_team_membership_on_cache_miss(
user_id: str,
team_id: str,
cache_key: str,
prisma_client: PrismaClient,
user_api_key_cache: UserApiKeyCache,
parent_otel_span: Span | None,
proxy_logging_obj: ProxyLogging | None,
) -> LiteLLM_TeamMembership | None:
try:
redis_cached: Final[object] = await user_api_key_cache.async_get_cache(key=cache_key)
redis_membership: Final = _membership_from_cached_payload(redis_cached)
if not isinstance(redis_membership, _TeamMembershipCacheMiss):
return redis_membership
return await _fetch_team_membership_from_db(
user_id=user_id,
team_id=team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
except Exception:
verbose_proxy_logger.exception("Error getting team membership")
return None
async def get_team_membership(
user_id: str,
team_id: str,
@ -2155,54 +2262,42 @@ async def get_team_membership(
Do a isolated check for team membership vs. doing a combined key + team + user + team-membership check, as key might come in frequently for different users/teams. Larger call will slowdown query time. This way we get to cache the constant (key/team/user info) and only update based on the changing value (team membership).
"""
from litellm.proxy._types import LiteLLM_TeamMembership
if prisma_client is None:
raise Exception("No db connected")
if user_id is None or team_id is None:
return None
_key: Final = team_membership_reservation_cache_key(user_id=user_id, team_id=team_id)
# check if in cache
cached: Final[object] = await user_api_key_cache.async_get_cache(key=_key)
if cached == NO_TEAM_MEMBERSHIP_SENTINEL:
return None
cached_membership_obj: Final = CacheCodec.deserialize(cached, model_type=LiteLLM_TeamMembership)
if cached_membership_obj is not None:
return cached_membership_obj
l1_cached: Final[object] = await user_api_key_cache.async_get_cache(key=_key, local_only=True)
l1_membership: Final = _membership_from_cached_payload(l1_cached)
if not isinstance(l1_membership, _TeamMembershipCacheMiss):
return l1_membership
# else, check db
try:
response: Final = await _dictable_table(TeamMembershipRepository(prisma_client)).find_unique(
where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}},
include={"litellm_budget_table": True},
inflight: Final[object] = _team_membership_inflight.get(_key)
if isinstance(inflight, asyncio.Task):
return _membership_from_shared_load(await asyncio.shield(inflight))
if prisma_client is None:
raise Exception("No db connected")
task: Final = asyncio.ensure_future(
_load_team_membership_on_cache_miss(
user_id=user_id,
team_id=team_id,
cache_key=_key,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
)
_team_membership_inflight[_key] = task
if response is None:
await user_api_key_cache.async_set_cache(
key=_key,
value=NO_TEAM_MEMBERSHIP_SENTINEL,
ttl=get_management_object_ttl(user_api_key_cache),
)
return None
def _clear_inflight(_done: object) -> None:
if _team_membership_inflight.get(_key) is task:
_team_membership_inflight.pop(_key, None)
_response: Final = LiteLLM_TeamMembership.model_validate(response.dict())
await user_api_key_cache.async_set_cache(
key=_key,
value=_response,
model_type=LiteLLM_TeamMembership,
)
return _response
except Exception:
verbose_proxy_logger.exception(
"Error getting team membership for user_id: %s, team_id: %s",
user_id,
team_id,
)
return None
task.add_done_callback(_clear_inflight)
return _membership_from_shared_load(await asyncio.shield(task))
def model_in_access_group(model: str, team_models: list[str] | None, llm_router: Router | None) -> bool:
@ -2375,13 +2470,6 @@ async def _backfill_null_user_email(
return updated_row
class UserNotFoundError(ValueError):
"""The user row is provably absent, as opposed to merely unreadable, so a caller that reads a missing row as no user-level limits can key on it without also swallowing a database that would not answer."""
def __init__(self, user_id: str) -> None:
super().__init__(f"User doesn't exist in db. 'user_id'={user_id}. Create user via `/user/new` call.")
@log_db_metrics
async def get_user_object(
user_id: str | None,
@ -2668,6 +2756,12 @@ async def invalidate_team_member_spend_state(
publish_auth_cache_invalidation,
)
inflight: Final[object] = _team_membership_inflight.pop(
team_membership_reservation_cache_key(user_id=user_id, team_id=team_id), None
)
if isinstance(inflight, asyncio.Task) and inflight is not asyncio.current_task():
await asyncio.wait((inflight,))
if new_spend is not None:
from litellm.proxy.proxy_server import SPEND_DB_FLOOR_CACHE_TTL_SECONDS, spend_counter_cache
@ -4122,18 +4216,21 @@ async def _team_member_granted_models(
prisma_client: PrismaClient,
user_api_key_cache: UserApiKeyCache,
proxy_logging_obj: ProxyLogging,
team_membership: LiteLLM_TeamMembership | None = None,
team_membership_loaded: bool = False,
) -> Sequence[str]:
"""The member's own ``allowed_models`` scope; empty when the member is not narrowed below the team."""
if team_object is None or valid_token.user_id is None:
return ()
team_membership: Final = await get_team_membership(
user_id=valid_token.user_id,
team_id=team_object.team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
if not team_membership_loaded:
team_membership = await get_team_membership(
user_id=valid_token.user_id,
team_id=team_object.team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
return () if team_membership is None else _member_allowed_models(team_membership)
@ -4169,6 +4266,8 @@ async def _granted_model_lists(
prisma_client: PrismaClient,
user_api_key_cache: UserApiKeyCache,
proxy_logging_obj: ProxyLogging,
team_membership: LiteLLM_TeamMembership | None = None,
team_membership_loaded: bool = False,
) -> tuple[Sequence[str], ...]:
"""One model allowlist per level that participates in authorizing the request."""
return (
@ -4180,6 +4279,8 @@ async def _granted_model_lists(
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
team_membership=team_membership,
team_membership_loaded=team_membership_loaded,
),
project_object.models if project_object is not None else (),
await _org_granted_models(
@ -4274,6 +4375,8 @@ async def collect_matched_model_access_groups(
prisma_client: PrismaClient | None,
user_api_key_cache: UserApiKeyCache,
proxy_logging_obj: ProxyLogging,
team_membership: LiteLLM_TeamMembership | None = None,
team_membership_loaded: bool = False,
) -> tuple[str, ...]:
"""
The budgeted model access groups that authorized this request, sorted and deduplicated.
@ -4319,6 +4422,8 @@ async def collect_matched_model_access_groups(
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
team_membership=team_membership,
team_membership_loaded=team_membership_loaded,
)
for granted_model in granted_models
)
@ -4334,6 +4439,8 @@ async def stamp_matched_model_access_groups(
prisma_client: PrismaClient | None,
user_api_key_cache: UserApiKeyCache,
proxy_logging_obj: ProxyLogging,
team_membership: LiteLLM_TeamMembership | None = None,
team_membership_loaded: bool = False,
) -> tuple[str, ...]:
"""Record the groups that authorized this request on its auth object, for the post-call spend
writer and the reservation counters, and hand them back for the budget check."""
@ -4350,6 +4457,8 @@ async def stamp_matched_model_access_groups(
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
team_membership=team_membership,
team_membership_loaded=team_membership_loaded,
)
except Exception as e: # noqa: BLE001 # fail-safe: attribution is spend telemetry, it must never break auth
verbose_proxy_logger.debug("model access group attribution failed: %s", e)
@ -4363,7 +4472,7 @@ async def stamp_matched_model_access_groups(
async def can_key_call_model(
model: str | list[str],
llm_model_list: list | None,
llm_model_list: Sequence[object] | None,
valid_token: UserAPIKeyAuth,
llm_router: litellm.Router | None,
) -> Literal[True]:
@ -4410,7 +4519,7 @@ async def can_key_call_model(
async def can_key_call_resolved_model(
model: str,
llm_model_list: list | None,
llm_model_list: Sequence[object] | None,
valid_token: UserAPIKeyAuth,
llm_router: litellm.Router | None,
) -> None:
@ -5152,6 +5261,8 @@ async def _check_team_member_budget(
prisma_client: PrismaClient | None,
user_api_key_cache: UserApiKeyCache,
proxy_logging_obj: ProxyLogging,
team_membership: LiteLLM_TeamMembership | None = None,
team_membership_loaded: bool = False,
):
"""Check if team member is over their max budget within the team."""
if (
@ -5160,23 +5271,25 @@ async def _check_team_member_budget(
and valid_token is not None
and valid_token.user_id is not None
):
team_membership: Final = await get_team_membership(
user_id=valid_token.user_id,
team_id=team_object.team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
if not team_membership_loaded:
team_membership = await get_team_membership(
user_id=valid_token.user_id,
team_id=team_object.team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
loaded_membership = team_membership
# Per-member override wins; otherwise fall back to the team-level
# default configured via team.metadata["team_member_budget_id"].
team_member_budget: float | None = None
if (
team_membership is not None
and team_membership.litellm_budget_table is not None
and team_membership.litellm_budget_table.max_budget is not None
loaded_membership is not None
and loaded_membership.litellm_budget_table is not None
and loaded_membership.litellm_budget_table.max_budget is not None
):
team_member_budget = team_membership.litellm_budget_table.max_budget
team_member_budget = loaded_membership.litellm_budget_table.max_budget
else:
default_budget_id: Final = (team_object.metadata or {}).get("team_member_budget_id")
if isinstance(default_budget_id, str):
@ -5195,7 +5308,7 @@ async def _check_team_member_budget(
team_member_budget = default_budget.max_budget
if team_member_budget is not None:
team_member_spend = (team_membership.spend if team_membership is not None else 0.0) or 0.0
team_member_spend = (loaded_membership.spend if loaded_membership is not None else 0.0) or 0.0
# Read from cross-pod counter (Redis-first) if available
from litellm.proxy.proxy_server import get_current_spend
@ -5224,6 +5337,8 @@ async def _check_team_member_model_access(
prisma_client: Optional["PrismaClient"],
user_api_key_cache: UserApiKeyCache,
proxy_logging_obj: ProxyLogging,
team_membership: LiteLLM_TeamMembership | None = None,
team_membership_loaded: bool = False,
) -> None:
"""
Check if a team member's per-member model scope allows access to the requested model.
@ -5234,22 +5349,24 @@ async def _check_team_member_model_access(
if valid_token.user_id is None or team_object.team_id is None:
return
team_membership: Final = await get_team_membership(
user_id=valid_token.user_id,
team_id=team_object.team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
if not team_membership_loaded:
team_membership = await get_team_membership(
user_id=valid_token.user_id,
team_id=team_object.team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
loaded_membership = team_membership
if (
team_membership is None
or team_membership.litellm_budget_table is None
or not team_membership.litellm_budget_table.allowed_models
loaded_membership is None
or loaded_membership.litellm_budget_table is None
or not loaded_membership.litellm_budget_table.allowed_models
):
return # no per-member restriction — inherit team-level check
member_allowed_models: Final[list[str]] = team_membership.litellm_budget_table.allowed_models
member_allowed_models: Final[list[str]] = loaded_membership.litellm_budget_table.allowed_models
try:
_can_object_call_model(
model=model,

View file

@ -33,7 +33,7 @@ from litellm.proxy.common_utils.http_parsing_utils import extract_nested_form_me
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
LITELLM_PASS_THROUGH_ENDPOINT_MARKER,
)
from litellm.types.router import CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS
from litellm.types.router import CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS, Deployment
from litellm.types.utils import CustomPricingLiteLLMParams
@ -1736,7 +1736,7 @@ def _append_model_candidates(candidates: list[str], value: Any) -> None:
candidates.extend(model for model in model_names if model)
def _dedupe_model_candidates(candidates: list[str]) -> list[str]:
def _dedupe_model_candidates(candidates: Collection[str]) -> list[str]:
deduped: Final[list[str]] = []
for model in candidates:
if model not in deduped:
@ -1845,13 +1845,42 @@ def _resolve_model_id_with_router(model_id: str | None, llm_router: Router | Non
return model_id
def get_cache_prediction_deployments(
*, current_deployment_id: str, candidate_deployment_id: str, llm_router: Router, team_id: str | None
) -> tuple[Deployment, Deployment] | None:
current: Final = llm_router.get_deployment(current_deployment_id)
candidate: Final = llm_router.get_deployment(candidate_deployment_id)
if current is None or candidate is None:
return None
if any(deployment.model_info.team_id not in (None, team_id) for deployment in (current, candidate)):
return None
return current, candidate
def _cache_prediction_model_candidates(
request_data: Mapping[str, object], llm_router: Router | None, team_id: str | None
) -> tuple[str, ...]:
current_id: Final = request_data.get("current_deployment_id")
candidate_id: Final = request_data.get("candidate_deployment_id")
if llm_router is None or not isinstance(current_id, str) or not isinstance(candidate_id, str):
return ()
deployments: Final = get_cache_prediction_deployments(
current_deployment_id=current_id, candidate_deployment_id=candidate_id, llm_router=llm_router, team_id=team_id
)
return tuple(deployment.model_name for deployment in deployments) if deployments is not None else ()
def _extract_model_candidates_from_request(
request_data: dict,
route: str,
request_headers: Mapping[str, object] | None = None,
request_query_params: Mapping[str, object] | None = None,
llm_router: Router | None = None,
team_id: str | None = None,
) -> list[str]:
if route == "/cost/predict-cache":
prediction_models: Final = _cache_prediction_model_candidates(request_data, llm_router, team_id) # pyright: ignore[reportUnknownArgumentType] # the typed reader validates each deployment ID from this legacy payload
return _dedupe_model_candidates(prediction_models)
candidates: Final[list[str]] = []
uses_model_routing_sources: Final = _route_uses_model_routing_sources(route=route)
uses_header_or_query_model_sources: Final = _route_matches_any_marker(
@ -1945,6 +1974,7 @@ def get_model_from_request(
request_query_params: Mapping[str, object] | None = None,
llm_router: Router | None = None,
request: Request | None = None,
team_id: str | None = None,
) -> str | list[str] | None:
"""Resolve the model(s) a request targets, for model-access and budget checks.
@ -1967,6 +1997,7 @@ def get_model_from_request(
request_headers=request_headers,
request_query_params=request_query_params,
llm_router=llm_router,
team_id=team_id,
)
model = _format_model_candidates(candidates)

View file

@ -28,11 +28,11 @@ from litellm.proxy._types import (
)
from litellm.proxy.auth.auth_checks import (
TeamNotFoundError,
UserNotFoundError,
get_team_membership,
get_team_object,
get_user_object,
)
from litellm.types.proxy.auth.auth_checks import UserNotFoundError
if TYPE_CHECKING:
from litellm.proxy._types import Span

View file

@ -136,6 +136,9 @@ class RouteChecks:
# For llm_api_routes, also check registered pass-through endpoints
################################################
if allowed_route == "llm_api_routes":
if route == "/auto_router/session" and RouteChecks._get_request_method(request) == "GET":
return True
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
InitPassThroughEndpointHelpers,
)

View file

@ -191,6 +191,7 @@ def _get_model_from_request_context(
route: str,
request: Request | None,
llm_router: Any | None = None,
team_id: str | None = None,
) -> str | list[str] | None:
return get_model_from_request(
request_data=request_data,
@ -199,6 +200,7 @@ def _get_model_from_request_context(
request_query_params=_safe_get_request_query_params(request=request),
llm_router=llm_router,
request=request,
team_id=team_id,
)
@ -217,7 +219,7 @@ async def _normalize_claude_model(
return
if request is not None and request.scope.get(_CLAUDE_MODEL_NORMALIZED) is True:
return
requested: Final = _get_model_from_request_context(request_data, route, request, llm_router)
requested: Final = _get_model_from_request_context(request_data, route, request, llm_router, valid_token.team_id)
if not isinstance(requested, str) or requested != request_data.get("model"):
return
if not requested.startswith("claude-router-") and not requested.lower().endswith("[1m]"):
@ -996,9 +998,12 @@ async def _resolve_jwt_to_virtual_key(
- Raises HTTPException: REJECT policy hit, missing claim under
REJECT/AUTO_REGISTER, or other policy violations.
"""
virtual_key_claim_field: Final = jwt_handler.litellm_jwtauth.virtual_key_claim_field
raw_issuer: Final = jwt_claims.get(JWTHandler.LITELLM_JWT_ISSUER_CLAIM)
normalized_issuer: Final = raw_issuer if isinstance(raw_issuer, str) else None
virtual_key_claim_field: Final = jwt_handler.litellm_jwtauth.get_virtual_key_claim_field(normalized_issuer)
if virtual_key_claim_field is None:
return None
behavior: Final = jwt_handler.litellm_jwtauth.get_unregistered_jwt_client_behavior(normalized_issuer)
claim_value: Final = get_nested_value(
data=jwt_claims,
@ -1015,7 +1020,6 @@ async def _resolve_jwt_to_virtual_key(
# simply by presenting a JWT that omits the configured field. For
# AUTO_REGISTER there is no stable identity to map without a claim
# value, so we deny rather than create a sentinel-keyed record.
behavior = jwt_handler.litellm_jwtauth.unregistered_jwt_client_behavior
if behavior in (
UnregisteredJWTClientBehavior.REJECT,
UnregisteredJWTClientBehavior.AUTO_REGISTER,
@ -1030,7 +1034,13 @@ async def _resolve_jwt_to_virtual_key(
return None
cache_key: Final = jwt_key_mapping_cache_key(virtual_key_claim_field, str(claim_value))
cached_mapping: Final = await user_api_key_cache.async_get_cache(cache_key)
raw_cached_mapping: Final = await user_api_key_cache.async_get_cache(cache_key)
sentinel_written_by_this_policy: Final = behavior == UnregisteredJWTClientBehavior.AUTO_REGISTER
cached_mapping: Final = (
None
if raw_cached_mapping == _JWT_PROXY_ADMIN_SENTINEL and not sentinel_written_by_this_policy
else raw_cached_mapping
)
if cached_mapping == _JWT_PROXY_ADMIN_SENTINEL:
# Previously resolved to a proxy admin via auth_builder; skip the
@ -1039,7 +1049,6 @@ async def _resolve_jwt_to_virtual_key(
return None
if cached_mapping == "__NO_MAPPING__":
behavior = jwt_handler.litellm_jwtauth.unregistered_jwt_client_behavior
if behavior == UnregisteredJWTClientBehavior.REJECT:
raise HTTPException(
status_code=403,
@ -1102,8 +1111,6 @@ async def _resolve_jwt_to_virtual_key(
)
# No mapping found (DB miss or no DB) — apply no-match policy.
behavior = jwt_handler.litellm_jwtauth.unregistered_jwt_client_behavior
if behavior == UnregisteredJWTClientBehavior.REJECT:
# Cache the miss before raising so repeated rejections are served from
# cache and don't re-query the DB on every request.
@ -1483,7 +1490,7 @@ async def _user_api_key_auth_builder(
# unnecessary DB queries in auth_builder
do_standard_jwt_auth = True
pending_auto_register: _PendingAutoRegister | None = None
if jwt_handler.litellm_jwtauth.virtual_key_claim_field is not None:
if jwt_handler.litellm_jwtauth.is_virtual_key_mapping_configured():
# Decode JWT to get claims without running full auth_builder
jwt_claims: dict | None
if jwt_handler.litellm_jwtauth.oidc_userinfo_enabled and not is_jwt:
@ -1647,6 +1654,7 @@ async def _user_api_key_auth_builder(
route=route,
request=request,
llm_router=llm_router,
team_id=valid_token.team_id,
)
skip_budget_checks = False
if model is not None and llm_router is not None:
@ -1687,6 +1695,7 @@ async def _user_api_key_auth_builder(
route=route,
request=request,
llm_router=llm_router,
team_id=valid_token.team_id,
)
),
)
@ -2086,6 +2095,7 @@ async def _user_api_key_auth_builder(
route=route,
request=request,
llm_router=llm_router,
team_id=valid_token.team_id,
)
skip_budget_checks = False
if model is not None and llm_router is not None:
@ -2204,6 +2214,7 @@ async def _user_api_key_auth_builder(
route=route,
request=request,
llm_router=llm_router,
team_id=valid_token.team_id,
)
current_models = _get_model_names_for_budget_checks(model=current_model)
@ -2234,6 +2245,7 @@ async def _user_api_key_auth_builder(
route=route,
request=request,
llm_router=llm_router,
team_id=valid_token.team_id,
)
current_models = _get_model_names_for_budget_checks(model=current_model)
@ -2729,6 +2741,7 @@ async def _run_centralized_common_checks(
route=route,
request=request,
llm_router=llm_router,
team_id=user_api_key_auth_obj.team_id,
)
# Pin the metadata variable name (litellm_metadata vs metadata) before
@ -2845,12 +2858,14 @@ def _should_skip_budget_checks(
route: str,
request: Request | None,
llm_router: Any | None,
team_id: str | None = None,
) -> bool:
model: Final = _get_model_from_request_context(
request_data=request_data,
route=route,
request=request,
llm_router=llm_router,
team_id=team_id,
)
if model is not None and llm_router is not None:
return _is_model_cost_zero(model=model, llm_router=llm_router)
@ -3296,6 +3311,7 @@ async def _enforce_key_and_fallback_model_access(
route=route,
request=request,
llm_router=llm_router,
team_id=valid_token.team_id,
)
if model is not None:
@ -3403,6 +3419,7 @@ async def _run_post_custom_auth_checks(
route=route,
request=request,
llm_router=llm_router,
team_id=valid_token.team_id,
)
current_models = _get_model_names_for_budget_checks(model=current_model)
@ -3444,6 +3461,7 @@ async def _run_post_custom_auth_checks(
route=route,
request=request,
llm_router=llm_router,
team_id=valid_token.team_id,
)
current_models = _get_model_names_for_budget_checks(model=current_model)

View file

@ -585,7 +585,9 @@ LiteLLM ████████░░░░░░░░░░░░░░
Claude Opus 5 ████████████████████████ $0.38
```
The routed model comes from Claude Code's own transcript, so it only names the tier model when the auto-router deployment sets `return_raw_model_name: true` (the `lite autoroute` wizard does); otherwise it shows the alias you requested. The cost lines come from `GET /auto_router/session?session_id=...`, which any virtual key may call for its own sessions, and are cached for five seconds under a per-user `$TMPDIR/litellm-statusline-<uid>` directory. The baseline is the priciest model in the router's hardest tier, the same counterfactual the auto-router's savings reports use. `lite unconfigure claude` removes the `statusLine` entry only while it still points at that script.
After the first response, the status line uses the latest routed model recorded by `GET /auto_router/session?session_id=...`, so it can show the tier model even when the transcript contains the router alias. If no session record is available, it falls back to Claude Code's transcript. Session records and costs are cached for five seconds under a per-user `$TMPDIR/litellm-statusline-<uid>` directory. The gateway records turns asynchronously, so the display can briefly lag a completed turn. Any virtual key may read its own sessions. The baseline is the priciest model in the router's hardest tier, the same counterfactual the auto-router's savings reports use. `lite unconfigure claude` removes the `statusLine` entry only while it still points at that script
After upgrading the CLI, rerun your original `lite configure claude` command with the same gateway, key and model choice to refresh `~/.litellm/statusline.py`. Keep any explicit `--model` value: omitting it removes the earlier model pin. Package upgrades alone do not refresh this installed copy
`lite codex` registers the same script as a Codex `Stop` hook for the launch, so after each turn Codex prints the same block as a system message. Codex asks once to trust the hook; the answer is remembered for later launches.

View file

@ -7,19 +7,17 @@ status refresh (about every 300ms while typing), so the proxy is asked at most o
TTL per session and every other refresh is served from a small on-disk cache that holds
only the proxy's answer, never the key.
Claude Code pipes a JSON payload on stdin (session_id, transcript_path, model); the routed
model is the `message.model` of the latest foreground assistant line in the transcript,
which is the proxy's response `model` field. That only names the tier model when the
auto-router deployment sets `return_raw_model_name: true`; otherwise it is the alias the
client requested. Codex pipes its Stop event instead (hook_event_name, session_id) and has
no transcript to read, so the routed model comes from the proxy's session record and the
result is printed as a `systemMessage` for the transcript. The proxy key is read from the
agent's own environment (the static token `lite configure claude` writes); nothing here
spawns a credential helper.
Claude Code pipes a JSON payload on stdin (session_id, transcript_path, model). After the
first foreground assistant response, the routed model comes from the proxy's session
record, falling back to the latest foreground assistant `message.model` in the transcript
when no record is available. Codex pipes its Stop event instead (hook_event_name, session_id)
and prints the session record as a `systemMessage` for the transcript. The proxy key is read
from the agent's own environment (the static token `lite configure claude` writes); nothing
here spawns a credential helper.
Cost figures come from GET /auto_router/session on the proxy, which reads the per-session
rollup written by the spend flush. That flush is asynchronous, so a turn's cost lands a
second or two after the turn; the cache TTL absorbs it.
The routed model and cost figures come from GET /auto_router/session on the proxy, which
reads the per-session rollup written by the asynchronous spend flush. The record and cache
can briefly lag a completed turn.
"""
from __future__ import annotations
@ -348,7 +346,8 @@ def status_line(
if not session_id or not credentials.usable:
return render(label, None, config_dir, color_enabled(env))
session: Final = load_session(credentials, session_id, cache_dir, fetch)
return render(label, session, config_dir, color_enabled(env))
routed_label: Final = model_label(session.last_model, config_dir) if session is not None else label
return render(routed_label, session, config_dir, color_enabled(env))
def codex_stop_message(

View file

@ -1571,6 +1571,9 @@ class ProxyBaseLLMRequestProcessing:
) -> dict:
exclude_values: Final = {"", None, "None"}
hidden_params = hidden_params or {}
resolved_call_id: Final = (
call_id or hidden_params.get("litellm_call_id") or (request_data or {}).get("litellm_call_id")
)
timing_values: Final = _timing_values(
hidden_params=hidden_params,
logging_obj=litellm_logging_obj,
@ -1598,7 +1601,7 @@ class ProxyBaseLLMRequestProcessing:
classifier_cost: Final = _classifier_cost_from_request_data(request_data)
headers: Final = {
"x-litellm-call-id": call_id,
"x-litellm-call-id": resolved_call_id,
"x-litellm-model-id": model_id,
"x-litellm-model-name": model_name,
"x-litellm-cache-key": cache_key,
@ -3452,15 +3455,13 @@ class ProxyBaseLLMRequestProcessing:
# a failed request reports no timing, matching /v1/chat/completions
read_timing_from_logging_obj=False,
)
# Extract headers from exception - check both e.headers and e.response.headers
headers = getattr(e, "headers", None) or {}
if not headers:
# Try to get headers from e.response.headers (httpx.Response)
_response: Final = attribute_of(e, "response")
if _response is not None:
_response_headers: Final = getattr(_response, "headers", None)
if _response_headers:
headers = get_response_headers(dict(_response_headers))
_response_headers: Final = getattr(_response, "headers", None) if _response is not None else None
_provider_headers: Final = _response_headers or getattr(e, "litellm_response_headers", None)
if _provider_headers:
headers = get_response_headers(dict(_provider_headers))
headers.update(custom_headers)
# Call response headers hook for failure

View file

@ -0,0 +1,10 @@
import os
from collections.abc import Mapping
def should_hide_default_credentials_hint(general_settings: Mapping[str, object]) -> bool:
return (
os.getenv("LITELLM_HIDE_DEFAULT_CREDENTIALS_HINT", "false").lower() == "true"
or general_settings.get("hide_default_credentials_hint", False) is True
or bool(os.getenv("UI_PASSWORD"))
)

View file

@ -0,0 +1,91 @@
from collections.abc import Mapping
from math import isfinite
from typing import Final
from pydantic import TypeAdapter
import litellm
from litellm.cost_calculator import (
_select_model_name_for_cost_calc, # pyright: ignore[reportPrivateUsage] # shares completion_cost's deployment tariff selection
completion_cost, # pyright: ignore[reportUnknownVariableType] # legacy optional parameters are untyped
)
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.types.management_endpoints.prompt_cache_prediction import CacheTokenBuckets
from litellm.types.utils import CacheCreationTokenDetails, ModelResponse, PromptTokensDetailsWrapper, Usage
_PRICE_ENTRY: Final = TypeAdapter(Mapping[str, object])
def _valid_price(value: object) -> bool:
return isinstance(value, (int, float)) and not isinstance(value, bool) and isfinite(value) and value >= 0
def _has_required_prices(prices: Mapping[str, object], tokens: CacheTokenBuckets) -> bool:
required: Final = (
("input_cost_per_token", True),
("cache_read_input_token_cost", tokens.cache_read_input_tokens > 0),
("cache_creation_input_token_cost", tokens.cache_creation_5m_input_tokens > 0),
("cache_creation_input_token_cost_above_1hr", tokens.cache_creation_1h_input_tokens > 0),
)
if any(needed and not _valid_price(prices.get(key)) for key, needed in required):
return False
return all(
_valid_price(value)
for key, value in prices.items()
if value is not None and any(needed and key.startswith(f"{base}_above_") for base, needed in required)
)
def price_cache_tokens(model: str, deployment_id: str, tokens: CacheTokenBuckets) -> float | None:
try:
selected_model: Final = _select_model_name_for_cost_calc(
model=model,
completion_response=None,
custom_pricing=True,
custom_llm_provider="anthropic",
router_model_id=deployment_id,
)
if selected_model is None:
return None
model_info: Final = litellm.get_model_info(model=selected_model, custom_llm_provider="anthropic")
registry: Final = _PRICE_ENTRY.validate_python(litellm.model_cost) # pyright: ignore[reportUnknownMemberType] # legacy registry is validated at this boundary
price_entry: Final = registry.get(model_info["key"])
if price_entry is None:
return None
prices: Final = _PRICE_ENTRY.validate_python(price_entry)
if not _has_required_prices(prices, tokens):
return None
usage: Final = Usage(
prompt_tokens=tokens.total_tokens,
completion_tokens=0,
total_tokens=tokens.total_tokens,
prompt_tokens_details=PromptTokensDetailsWrapper(
cached_tokens=tokens.cache_read_input_tokens,
cache_creation_tokens=tokens.cache_creation_5m_input_tokens + tokens.cache_creation_1h_input_tokens,
cache_creation_token_details=CacheCreationTokenDetails(
ephemeral_5m_input_tokens=tokens.cache_creation_5m_input_tokens,
ephemeral_1h_input_tokens=tokens.cache_creation_1h_input_tokens,
),
),
)
logging_obj: Final = Logging(
model=model,
messages=[], # mutable-ok: Logging requires a list
stream=False,
call_type="completion",
start_time=None,
litellm_call_id="prompt-cache-prediction",
function_id="prompt-cache-prediction",
)
completion_cost(
completion_response=ModelResponse(model=model, usage=usage),
model=model,
custom_llm_provider="anthropic",
custom_pricing=True,
router_model_id=deployment_id,
litellm_logging_obj=logging_obj,
)
cost: Final = logging_obj.cost_breakdown.get("input_cost") if logging_obj.cost_breakdown is not None else None
return cost if cost is not None and _valid_price(cost) else None
except Exception: # noqa: BLE001 # the shared pricing owners raise plain Exception for unpriceable models
return None

View file

@ -74,10 +74,14 @@ class LatestHealthCheckRow(BaseModel):
_ROWS_ADAPTER: Final = TypeAdapter(tuple[LatestHealthCheckRow, ...])
async def query_latest_health_checks(prisma_client: PrismaClient) -> tuple[LatestHealthCheckRow, ...]:
rows: Final = await prisma_client.db.query_raw(LATEST_HEALTH_CHECKS_SQL)
return _ROWS_ADAPTER.validate_python(rows)
async def fetch_latest_health_checks(prisma_client: PrismaClient) -> tuple[LatestHealthCheckRow, ...]:
try:
rows: Final = await prisma_client.db.query_raw(LATEST_HEALTH_CHECKS_SQL)
return _ROWS_ADAPTER.validate_python(rows)
return await query_latest_health_checks(prisma_client)
except Exception as query_err: # noqa: BLE001 # health decorates other reads; a driver error must not fail them
verbose_proxy_logger.error("Error getting all latest health checks: %s", query_err)
return ()

View file

@ -4,6 +4,7 @@ from typing import Final
from fastapi import APIRouter
from litellm.proxy.common_utils.html_forms.default_credentials_hint import should_hide_default_credentials_hint
from litellm.types.proxy.discovery_endpoints.ui_discovery_endpoints import (
UiDiscoveryEndpoints,
)
@ -23,10 +24,7 @@ async def get_ui_config():
or general_settings.get("auto_redirect_ui_login_to_sso", False) is True
)
admin_ui_disabled: Final = os.getenv("DISABLE_ADMIN_UI", "false").lower() == "true"
hide_default_credentials_hint: Final = bool(
os.getenv("LITELLM_HIDE_DEFAULT_CREDENTIALS_HINT", "false").lower() == "true"
or general_settings.get("hide_default_credentials_hint", False) is True
)
hide_default_credentials_hint: Final = should_hide_default_credentials_hint(general_settings)
sso_configured: Final = has_user_setup_sso()

View file

@ -244,6 +244,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
prompt_attack_threshold: float | None = 0.5,
pii_confidence_threshold: float | None = 0.5,
chunk_budget_chars: int = BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS,
contextual_grounding_from_messages: bool = False,
streaming_buffer_until_moderated: bool | None = None,
streaming_sampling_rate: int | None = None,
streaming_end_of_stream_only: bool | None = None,
@ -265,6 +266,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
self.guardrailVersion = guardrailVersion
self.guardrail_provider = "bedrock"
self.chunk_budget_chars = chunk_budget_chars
self.contextual_grounding_from_messages = contextual_grounding_from_messages
self.experimental_use_latest_role_message_only = bool(kwargs.get("experimental_use_latest_role_message_only"))
# Resource-less, detect-only InvokeGuardrailChecks mode. Present `checks`
@ -459,8 +461,8 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
"""
Flatten a message into text blocks, preserving any contextual-grounding
qualifier carried by the content-block ``type`` (grounding_source / query).
Untagged text keeps ``qualifier=None`` so the payload is unchanged for
callers that do not use grounding.
Untagged text keeps ``qualifier=None``; the OUTPUT scan decides whether to
derive grounding qualifiers from it.
"""
content: Final = message.get("content")
if content is None:
@ -493,6 +495,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
result carrying externally-influenced content can supply fake evidence for the
contextual-grounding check to grade the response against. ``query`` is accepted
from any role (it is the user's question).
With ``contextual_grounding_from_messages`` on, a request with no tagged blocks
falls back to the plain messages: system / developer text is the grounding
source and the latest user message is the query.
"""
grounding: Final[list[QualifiedTextBlock]] = []
for message in messages or []:
@ -504,7 +510,33 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
and role in _GROUNDING_SOURCE_TRUSTED_ROLES
):
grounding.append(block)
return grounding
if grounding or not self.contextual_grounding_from_messages:
return grounding
return self._derive_grounding_blocks_from_plain_messages(messages)
def _derive_grounding_blocks_from_plain_messages(
self, messages: list[AllMessageValues] | None
) -> list[QualifiedTextBlock]:
if not messages:
return []
latest_user_index: Final = self._find_latest_message_index(messages, target_role="user")
if latest_user_index is None:
return []
sources: Final = tuple(
QualifiedTextBlock(text=block.text, qualifier="grounding_source")
for message in messages
if message.get("role") in _GROUNDING_SOURCE_TRUSTED_ROLES
for block in self.get_content_items_for_message(message=message) or []
if block.text
)
queries: Final = tuple(
QualifiedTextBlock(text=block.text, qualifier="query")
for block in self.get_content_items_for_message(message=messages[latest_user_index]) or []
if block.text
)
if not sources or not queries:
return []
return [*sources, *queries]
def supports_scan_only_tool_results(self) -> bool:
return self.experimental_use_latest_role_message_only is not True
@ -3210,6 +3242,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
bedrock_response = await self.make_bedrock_api_request(
source="OUTPUT",
response=synthetic_response,
messages=request_data.get("messages"),
request_data=request_data,
logging_event_type=_log_hook,
)

View file

@ -44,7 +44,7 @@ class JavelinGuardrail(CustomGuardrail):
application: str | None = None,
**kwargs,
):
f"""
"""
Initialize the JavelinGuardrail class.
This calls: {api_base}/{api_version}/guardrail/{guardrail_name}/apply

View file

@ -15,7 +15,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"
# We check the raw guardrail dict because LitellmParams normalizes None → False,
# making it impossible to distinguish "not set" from "explicitly false" via litellm_params.
_raw_default_on: Final = cast(dict[str, Any], guardrail).get("litellm_params", {}).get("default_on")
_default_on: Final = False if _raw_default_on is False else True
_default_on: Final = _raw_default_on is not False
_callback: Final = MCPEndUserPermissionGuardrail(
guardrail_name=guardrail.get("guardrail_name", ""),

View file

@ -1,11 +1,13 @@
from collections.abc import AsyncGenerator, Mapping, Sequence
import time
from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence
from enum import Enum, auto
from typing import TYPE_CHECKING, Any, Final, Literal
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal
import httpx
from fastapi import HTTPException
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
import json
@ -23,6 +25,7 @@ from litellm.litellm_core_utils.core_helpers import (
get_or_create_metadata_bucket,
)
from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
get_async_httpx_client,
httpxSpecialProvider,
)
@ -52,6 +55,7 @@ from litellm.types.utils import (
CallTypes,
CallTypesLiteral,
Choices,
GenericGuardrailAPIInputs,
GuardrailStatus,
ModelResponse,
ModelResponseStream,
@ -118,8 +122,12 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
Supports:
- Pre-call sanitization (sanitizeUserPrompt)
- Post-call sanitization (sanitizeModelResponse)
- logging_only: scans the completed response after it reaches the client and
records the verdict in spend logs without blocking
"""
use_native_lifecycle_hooks: ClassVar[bool] = True
@classmethod
def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]:
return [
@ -128,6 +136,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
GuardrailEventHooks.post_call,
GuardrailEventHooks.pre_mcp_call,
GuardrailEventHooks.during_mcp_call,
GuardrailEventHooks.logging_only,
]
def __init__(
@ -138,6 +147,8 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
credentials: VERTEX_CREDENTIALS_TYPES | None = None,
api_endpoint: str | None = None,
sanitize_error_detail: "bool | None" = True,
async_handler: AsyncHTTPHandler | None = None,
access_token_provider: Callable[[], Awaitable[tuple[str, str]]] | None = None,
**kwargs,
):
# Set supported event hooks if not already provided
@ -154,7 +165,10 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
VertexBase.__init__(self)
# Then set our attributes (this ensures project_id is not overwritten)
self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback)
self.async_handler = async_handler or get_async_httpx_client(
llm_provider=httpxSpecialProvider.GuardrailCallback
)
self.access_token_provider = access_token_provider
self.template_id = template_id
self.project_id = project_id
self.location = location or "us-central1"
@ -278,11 +292,14 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
If file_bytes and file_type are provided, file prompt sanitization is performed.
"""
# Get access token using VertexBase auth
access_token, resolved_project_id = await self._ensure_access_token_async(
credentials=self.credentials,
project_id=self.project_id,
custom_llm_provider="vertex_ai",
)
if self.access_token_provider is not None:
access_token, resolved_project_id = await self.access_token_provider()
else:
access_token, resolved_project_id = await self._ensure_access_token_async(
credentials=self.credentials,
project_id=self.project_id,
custom_llm_provider="vertex_ai",
)
# Use resolved project ID if not explicitly set
if not self.project_id and resolved_project_id:
@ -1096,6 +1113,11 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
add_guardrail_to_applied_guardrails_header,
)
if self.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_call) is not True:
async for chunk in response:
yield chunk
return
all_chunks: Final[Sequence[object]] = tuple([chunk async for chunk in response])
if not all_chunks or self._is_terminal_error_stream(all_chunks):
@ -1213,6 +1235,60 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
for chunk in all_chunks:
yield chunk
@log_guardrail_information
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict,
input_type: Literal["request", "response"],
logging_obj: "LiteLLMLoggingObj | None" = None,
) -> GenericGuardrailAPIInputs:
content: Final = "\n".join(text for text in inputs.get("texts") or () if text)
if not content:
return inputs
source: Final[Literal["user_prompt", "model_response"]] = (
"user_prompt" if input_type == "request" else "model_response"
)
start_time: Final = time.time()
try:
armor_response: Final = await self.make_model_armor_request(
content=content, source=source, request_data=request_data
)
except (ModelArmorAPIError, httpx.HTTPError) as e:
error_end_time: Final = time.time()
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response=str(e),
request_data=request_data,
guardrail_status="guardrail_failed_to_respond",
guardrail_provider="model_armor",
start_time=start_time,
end_time=error_end_time,
duration=error_end_time - start_time,
)
return inputs
flagged: Final = self._should_block_content(armor_response, allow_sanitization=False)
end_time: Final = time.time()
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response=self._build_logging_response(armor_response),
request_data=request_data,
guardrail_status="guardrail_flagged" if flagged else "success",
guardrail_provider="model_armor",
start_time=start_time,
end_time=end_time,
duration=end_time - start_time,
)
if flagged and not self._event_hook_is_event_type(GuardrailEventHooks.logging_only):
raise HTTPException(
status_code=400,
detail=self._build_block_error_detail(
"Response blocked by Model Armor" if input_type == "response" else "Content blocked by Model Armor",
armor_response,
),
)
return inputs
@staticmethod
def get_config_model() -> type["GuardrailConfigModel"] | None:
"""

View file

@ -30,6 +30,7 @@ if TYPE_CHECKING:
_SANITIZE_FILE_FAIL_OPEN_TIMEOUT_SECONDS: Final = 30.0
_SANITIZE_FILE_QUEUED_STATUSES: Final = frozenset({"created", "in progress"})
_PROTECT_ROLES: Final = frozenset({"system", "user", "assistant"})
@ -553,16 +554,18 @@ class PromptSecurityGuardrail(CustomGuardrail):
"metadata": result.get("metadata", {}),
"violations": result.get("metadata", {}).get("violations", []),
}
elif status == "in progress":
verbose_proxy_logger.debug(
"Prompt Security Guardrail: File sanitization in progress (attempt %d/%d)",
attempt + 1,
self.max_poll_attempts,
)
continue
else:
if status not in _SANITIZE_FILE_QUEUED_STATUSES:
raise HTTPException(status_code=500, detail=f"Unexpected sanitization status: {status}")
verbose_proxy_logger.debug(
"Prompt Security Guardrail: File sanitization status=%s for jobId=%s (attempt %d/%d)",
status,
job_id,
attempt + 1,
self.max_poll_attempts,
)
raise HTTPException(status_code=408, detail="File sanitization timeout")
def _raise_if_file_blocked(self, sanitization_result: _SanitizeResult, resource_name: str) -> None:

View file

@ -23,6 +23,7 @@ def initialize_bedrock(litellm_params: LitellmParams, guardrail: Guardrail):
prompt_attack_threshold=litellm_params.prompt_attack_threshold,
pii_confidence_threshold=litellm_params.pii_confidence_threshold,
chunk_budget_chars=litellm_params.chunk_budget_chars,
contextual_grounding_from_messages=litellm_params.contextual_grounding_from_messages,
default_on=litellm_params.default_on,
disable_exception_on_block=litellm_params.disable_exception_on_block,
mask_request_content=litellm_params.mask_request_content,

View file

@ -45,7 +45,10 @@ from litellm.proxy.auth.auth_utils import (
from litellm.proxy.auth.model_checks import get_key_models
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
from litellm.proxy.db.health_check_latest import LatestHealthCheckRow
from litellm.proxy.db.health_check_latest import (
LatestHealthCheckRow,
query_latest_health_checks,
)
from litellm.proxy.db.proxy_worker_heartbeat import count_live_proxy_workers
from litellm.proxy.health_check import (
ADMIN_ONLY_HEALTH_DISPLAY_PARAMS,
@ -876,7 +879,7 @@ async def _save_background_health_checks_to_db(
)
# Step 3: Get latest health checks for all models in one query to compare status
latest_checks: Final = await prisma_client.get_all_latest_health_checks()
latest_checks: Final = await query_latest_health_checks(prisma_client)
latest_checks_map: Final = {}
for check in latest_checks:
# Use model_id as primary key, fallback to model_name

View file

@ -9,6 +9,7 @@ from .max_budget_per_session_limiter import _PROXY_MaxBudgetPerSessionHandler
from .max_iterations_limiter import _PROXY_MaxIterationsHandler
from .parallel_request_limiter import _PROXY_MaxParallelRequestsHandler
from .parallel_request_limiter_v3 import _PROXY_MaxParallelRequestsHandler_v3
from .prompt_cache_prediction import PromptCacheObserver
from .responses_id_security import ResponsesIDSecurity
from .sensitive_data_routing import _PROXY_SensitiveDataRoutingHandler
@ -25,6 +26,7 @@ PROXY_HOOKS: Final = {
"max_iterations_limiter": _PROXY_MaxIterationsHandler,
"max_budget_per_session_limiter": _PROXY_MaxBudgetPerSessionHandler,
"sensitive_data_routing": _PROXY_SensitiveDataRoutingHandler,
"prompt_cache_prediction": PromptCacheObserver,
}
## FEATURE FLAG HOOKS ##

View file

@ -9,10 +9,12 @@ import binascii
import logging
import os
import uuid
from collections.abc import Awaitable, Callable, Mapping, Sequence, Set
from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence, Set
from contextlib import asynccontextmanager
from contextvars import ContextVar
from dataclasses import dataclass, field
from datetime import datetime
from types import MappingProxyType
from typing import (
TYPE_CHECKING,
Any,
@ -23,6 +25,7 @@ from typing import (
TypedDict,
)
from pydantic import TypeAdapter
from typing_extensions import NotRequired, ReadOnly
from litellm import DualCache
@ -84,6 +87,9 @@ else:
InternalUsageCache = Any
_REQUEST_RATE_LIMIT_DATA: Final = TypeAdapter(Mapping[str, object])
BATCH_RATE_LIMITER_SCRIPT: Final = """
local results = {}
local now = tonumber(ARGV[1])
@ -2673,12 +2679,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
Returns list of descriptors for API key, user, team, team member, end user,
model-specific, agent, and agent-session limits.
"""
from litellm.proxy.auth.auth_utils import (
get_team_model_rpm_limit,
get_team_model_tpm_limit,
)
descriptors: Final = []
descriptors: Final[list[RateLimitDescriptor]] = [] # mutable-ok: existing descriptor helpers append in place
# API Key rate limits
if user_api_key_dict.api_key and (
@ -2803,34 +2804,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
descriptors=descriptors,
)
if (
get_team_model_rpm_limit(user_api_key_dict) is not None
or get_team_model_tpm_limit(user_api_key_dict) is not None
):
_tpm_limit_for_team_model: Final = get_team_model_tpm_limit(user_api_key_dict) or {}
_rpm_limit_for_team_model: Final = get_team_model_rpm_limit(user_api_key_dict) or {}
should_check_rate_limit = False
if requested_model in _tpm_limit_for_team_model or requested_model in _rpm_limit_for_team_model:
should_check_rate_limit = True
if should_check_rate_limit:
model_specific_tpm_limit = None
model_specific_rpm_limit = None
if requested_model in _tpm_limit_for_team_model:
model_specific_tpm_limit = _tpm_limit_for_team_model[requested_model]
if requested_model in _rpm_limit_for_team_model:
model_specific_rpm_limit = _rpm_limit_for_team_model[requested_model]
descriptors.append(
RateLimitDescriptor(
key="model_per_team",
value=f"{user_api_key_dict.team_id}:{requested_model}",
rate_limit={
"requests_per_unit": model_specific_rpm_limit,
"tokens_per_unit": model_specific_tpm_limit,
"window_size": self.window_size,
},
)
)
self._add_team_model_rate_limit_descriptor_from_metadata(
user_api_key_dict=user_api_key_dict,
requested_model=requested_model if isinstance(requested_model, str) else None,
descriptors=descriptors,
)
# Agent-level and session-level rate limits
resolved_agent_id: Final = self._get_resolved_agent_id(user_api_key_dict, data)
@ -3416,6 +3394,108 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
requested_model,
)
async def _build_request_rate_limit_descriptors(
self,
user_api_key_dict: UserAPIKeyAuth,
data: Mapping[str, object],
call_type: str | None,
) -> list[RateLimitDescriptor]: # mutable-ok: the shared generation reservation helpers require a list
metadata: Final = _REQUEST_RATE_LIMIT_DATA.validate_python(
user_api_key_dict.metadata or MappingProxyType({}) # pyright: ignore[reportUnknownMemberType] # validates the legacy auth metadata boundary
)
rpm_value: Final = metadata.get("rpm_limit_type")
tpm_value: Final = metadata.get("tpm_limit_type")
rpm_limit_type: Final = rpm_value if isinstance(rpm_value, str) else None
tpm_limit_type: Final = tpm_value if isinstance(tpm_value, str) else None
model_value: Final = data.get("model")
requested_model: Final = model_value if isinstance(model_value, str) else None
model_has_failures: Final = (
await self._check_model_has_recent_failures(
model=requested_model,
parent_otel_span=user_api_key_dict.parent_otel_span,
)
if requested_model and self._is_dynamic_rate_limiting_enabled(rpm_limit_type, tpm_limit_type)
else False
)
descriptors: Final = self._create_rate_limit_descriptors( # pyright: ignore[reportUnknownMemberType] # legacy helper reads a dictionary with validated keys
user_api_key_dict=user_api_key_dict,
data=dict(data), # mutable-ok: legacy descriptor helpers accept a request dictionary
rpm_limit_type=rpm_limit_type,
tpm_limit_type=tpm_limit_type,
model_has_failures=model_has_failures,
call_type=call_type,
)
self._add_project_model_rate_limit_descriptor_from_metadata(
user_api_key_dict=user_api_key_dict,
requested_model=requested_model,
descriptors=descriptors,
)
self.add_project_io_token_rate_limit_descriptors_from_metadata(
user_api_key_dict=user_api_key_dict,
requested_model=requested_model,
descriptors=descriptors,
)
return [ # mutable-ok: the shared generation reservation helpers require a list
*descriptors,
*self.create_organization_rate_limit_descriptor(user_api_key_dict, requested_model),
]
async def _release_request_capacity_when_admitted(
self,
admission: asyncio.Task[RateLimitResponse],
acquisition: ParallelSlotAcquisition,
user_api_key_dict: UserAPIKeyAuth,
) -> None:
response: Final = await admission
if response["overall_code"] == "OK":
await self._release_parallel_request_slots(acquisition, user_api_key_dict.parent_otel_span)
@asynccontextmanager
async def request_capacity(
self,
user_api_key_dict: UserAPIKeyAuth,
model: str,
*,
request_data: Mapping[str, object] | None = None,
) -> AsyncGenerator[None, None]:
"""Charge one non-generation provider request to RPM and hold its concurrency slot."""
data: Final = MappingProxyType({**(request_data or MappingProxyType({})), "model": model})
descriptors: Final = await self._build_request_rate_limit_descriptors(user_api_key_dict, data, None)
acquisition: Final = ParallelSlotAcquisition(
slot_id=uuid.uuid4().hex,
counter_keys=[ # mutable-ok: the shared slot-release contract requires a list
self.create_rate_limit_keys(d["key"], d["value"], "max_parallel_requests")
for d in descriptors
if d["rate_limit"] is not None and d["rate_limit"].get("max_parallel_requests") is not None
],
)
admission: Final = asyncio.create_task(
self.should_rate_limit(
descriptors=descriptors,
parent_otel_span=user_api_key_dict.parent_otel_span,
skip_tpm_check=True,
parallel_slot_id=acquisition["slot_id"],
)
)
try:
response: Final = await asyncio.shield(admission)
if response["overall_code"] == "OVER_LIMIT":
self._handle_rate_limit_error(response, descriptors, model)
yield
finally:
cleanup: Final = asyncio.create_task(
self._release_request_capacity_when_admitted(admission, acquisition, user_api_key_dict)
)
cancellation: asyncio.CancelledError | None = None # rebind-ok: retain cancellation until cleanup finishes
while not cleanup.done():
try:
await asyncio.shield(cleanup)
except asyncio.CancelledError as exc:
cancellation = exc # rebind-ok: retain the latest cancellation without interrupting slot release
cleanup.result()
if cancellation is not None:
raise cancellation
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
@ -3444,59 +3524,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
call_type=call_type,
)
# Get rate limit types from metadata
metadata: Final = user_api_key_dict.metadata or {}
rpm_limit_type: Final = metadata.get("rpm_limit_type")
tpm_limit_type: Final = metadata.get("tpm_limit_type")
# For dynamic mode, check if the model has recent failures
model_has_failures = False
requested_model: Final = data.get("model", None)
if (
self._is_dynamic_rate_limiting_enabled(
rpm_limit_type=rpm_limit_type,
tpm_limit_type=tpm_limit_type,
)
and requested_model
):
model_has_failures = await self._check_model_has_recent_failures(
model=requested_model,
parent_otel_span=user_api_key_dict.parent_otel_span,
)
# Create rate limit descriptors
descriptors: Final = self._create_rate_limit_descriptors(
request_data: Final = _REQUEST_RATE_LIMIT_DATA.validate_python(data)
model_value: Final = request_data.get("model")
requested_model: Final = model_value if isinstance(model_value, str) else None
descriptors: Final = await self._build_request_rate_limit_descriptors(
user_api_key_dict=user_api_key_dict,
data=data,
rpm_limit_type=rpm_limit_type,
tpm_limit_type=tpm_limit_type,
model_has_failures=model_has_failures,
data=request_data,
call_type=call_type,
)
# Add team model rate limits from team_metadata
self._add_team_model_rate_limit_descriptor_from_metadata(
user_api_key_dict=user_api_key_dict,
requested_model=requested_model,
descriptors=descriptors,
)
# Project Level Rate Limits
self._add_project_model_rate_limit_descriptor_from_metadata(
user_api_key_dict=user_api_key_dict,
requested_model=requested_model,
descriptors=descriptors,
)
self.add_project_io_token_rate_limit_descriptors_from_metadata(
user_api_key_dict=user_api_key_dict,
requested_model=requested_model,
descriptors=descriptors,
)
# Org Level Rate Limits
descriptors.extend(self.create_organization_rate_limit_descriptor(user_api_key_dict, requested_model))
# Only check rate limits if we have descriptors with actual limits
if descriptors:
# First pass: RPM and max_parallel_requests sliding-window check.

View file

@ -0,0 +1,142 @@
from __future__ import annotations
import asyncio
import time
from collections.abc import Callable, Mapping
from datetime import datetime
from typing import TYPE_CHECKING, Final, Literal
import httpx
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError
from litellm.caching.dual_cache import DualCache
from litellm.integrations.custom_logger import CustomLogger
from litellm.llms.anthropic.prompt_cache_prediction import PromptPrefix, parse_observed_cache
from litellm.types.utils import ModelResponse
if TYPE_CHECKING:
from litellm.proxy.utils import InternalUsageCache
_RETENTION_SECONDS: Final = 86_400
class CacheObservation(BaseModel):
model_config = ConfigDict(extra="forbid", frozen=True, strict=True)
fingerprint: str = Field(pattern=r"^[0-9a-f]{64}$")
cached_tokens: int = Field(gt=0)
observed_at: float = Field(ge=0, allow_inf_nan=False)
expires_at: float = Field(ge=0, allow_inf_nan=False)
_CACHE_ENTRY: Final[TypeAdapter[CacheObservation | str | None]] = TypeAdapter(CacheObservation | str | None)
def _cache_key(scope: str, fingerprint: str) -> str:
return f"prompt-cache-observation:{scope}:{fingerprint}"
async def lookup(
cache: DualCache, scope: str, prefix: PromptPrefix, now: float | None = None
) -> CacheObservation | None:
checked_at: Final = time.time() if now is None else now
exact: Final = await _read_exact(cache, scope, prefix.fingerprint)
if exact is not None and exact.expires_at > checked_at:
return exact
older: Final = await asyncio.gather(
*(_read_exact(cache, scope, fingerprint) for fingerprint in prefix.fingerprints[1:])
)
observations: Final = tuple(observation for observation in (exact, *older) if observation is not None)
return next(
(observation for observation in observations if observation.expires_at > checked_at),
next(iter(observations), None),
)
async def _read_exact(cache: DualCache, scope: str, fingerprint: str) -> CacheObservation | None:
try:
value: Final = _CACHE_ENTRY.validate_python(await cache.async_get_cache(_cache_key(scope, fingerprint), ttl=1)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # validate the legacy cache's untyped result at the I/O boundary
if value is None:
return None
observation: Final = CacheObservation.model_validate_json(value) if isinstance(value, str) else value
except ValidationError:
return None
return observation if observation.fingerprint == fingerprint else None
class _Metadata(BaseModel):
model_config = ConfigDict(strict=True)
user_api_key_hash: str = Field(min_length=1)
class _Logged(BaseModel):
model_config = ConfigDict(strict=True)
status: Literal["success"]
model_id: str = Field(min_length=1)
metadata: _Metadata
class _Event(BaseModel):
model_config = ConfigDict(strict=True, arbitrary_types_allowed=True)
call_type: Literal["anthropic_messages"]
custom_llm_provider: Literal["anthropic"]
cache_hit: bool | None = None
httpx_response: httpx.Response
first_api_call_start_time: datetime
standard_logging_object: _Logged
stream: bool = False
prompt_cache_response_complete: bool = False
class PromptCacheObserver(CustomLogger):
def __init__(self, internal_usage_cache: InternalUsageCache, clock: Callable[[], float] = time.time) -> None:
super().__init__() # pyright: ignore[reportUnknownMemberType] # base callback constructor accepts untyped kwargs
self.cache = internal_usage_cache.dual_cache
self.clock = clock
async def async_log_success_event(
self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime
) -> None:
if not isinstance(response_obj, ModelResponse):
return
try:
event: Final = _Event.model_validate(kwargs)
wire: Final = event.httpx_response.request
except (ValidationError, RuntimeError, httpx.RequestNotRead):
return
if (
event.cache_hit
or event.httpx_response.status_code != 200
or (event.stream and not event.prompt_cache_response_complete)
):
return
observed: Final = parse_observed_cache(
wire,
response_obj,
event.standard_logging_object.metadata.user_api_key_hash,
event.standard_logging_object.model_id,
)
if observed is None:
return
prefix: Final = observed.prefix
scope: Final = observed.scope
cache_tokens: Final = observed.cached_tokens
now: Final = self.clock()
started: Final = event.first_api_call_start_time.timestamp()
if started > now:
return
if observed.cache_creation_tokens == 0:
previous: Final = await _read_exact(self.cache, scope, prefix.fingerprint)
if previous is None or previous.fingerprint != prefix.fingerprint or previous.cached_tokens != cache_tokens:
return
observation: Final = CacheObservation(
fingerprint=prefix.fingerprint,
cached_tokens=cache_tokens,
observed_at=now,
expires_at=started + prefix.ttl_seconds,
)
key: Final = _cache_key(scope, prefix.fingerprint)
payload: Final = observation.model_dump_json()
await self.cache.async_set_cache(key, payload, ttl=_RETENTION_SECONDS) # pyright: ignore[reportUnknownMemberType] # legacy cache accepts a serialized validated observation
if self.cache.redis_cache is not None:
await self.cache.async_set_cache(key, payload, local_only=True, ttl=1) # pyright: ignore[reportUnknownMemberType] # keep the local copy short-lived while Redis retains stale evidence

View file

@ -6,11 +6,10 @@ from typing import Final
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.constants import PROXY_LLM_PROVIDER_FALLBACK
from litellm.types.router import ModelGroupInfo
from litellm.types.utils import PriorityReservationDict
PROXY_LLM_PROVIDER_FALLBACK: Final = "litellm_proxy"
def resolve_llm_provider_for_rate_limit(
model: str | None,

View file

@ -28,6 +28,7 @@ from litellm.proxy._types import (
UserAPIKeyAuth,
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.management_endpoints.prompt_cache_prediction import router as prompt_cache_prediction_router
from litellm.types.utils import (
CostBreakdown,
CostPerToken,
@ -39,6 +40,7 @@ from litellm.types.utils import (
)
router: Final = APIRouter()
router.include_router(prompt_cache_prediction_router)
@dataclass(frozen=True, slots=True)

View file

@ -0,0 +1,278 @@
import time
from collections.abc import Mapping
from types import MappingProxyType
from typing import Annotated, Final
from fastapi import APIRouter, Depends, HTTPException, Request
from pydantic import BaseModel, JsonValue, TypeAdapter
import litellm
from litellm._internal_context import current_billing_time, pinned_billing_time
from litellm.caching.caching import DualCache
from litellm.integrations.custom_logger import CustomLogger
from litellm.llms.anthropic.prompt_cache_prediction import (
PromptPrefix,
TokenCounter,
UnsupportedPredictionTarget,
cache_scope,
count_prompt_tokens,
parse_prompt,
resolve_prediction_target,
supported_prediction_headers,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.auth_checks import can_key_call_resolved_model
from litellm.proxy.auth.auth_utils import get_cache_prediction_deployments
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.http_parsing_utils import (
_read_request_body, # pyright: ignore[reportPrivateUsage, reportUnknownVariableType] # canonical parsed-body owner; validate its legacy result at the endpoint boundary
)
from litellm.proxy.common_utils.prompt_cache_pricing import price_cache_tokens
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
_PROXY_MaxParallelRequestsHandler_v3, # pyright: ignore[reportPrivateUsage] # use the configured proxy limiter's shared capacity owner
)
from litellm.proxy.hooks.prompt_cache_prediction import lookup
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
from litellm.types.management_endpoints.prompt_cache_prediction import (
CacheCostScenario,
CacheEvidence,
CachePredictionArm,
CachePredictionRequest,
CachePredictionResponse,
CacheTokenBuckets,
)
from litellm.types.router import Deployment
from litellm.utils import get_prompt_cache_min_tokens
router: Final = APIRouter()
_REQUEST_DATA: Final = TypeAdapter(Mapping[str, object])
class _CallerSettings(BaseModel):
config: Mapping[str, object] | None = None
def has_request_transforms() -> bool:
from litellm.proxy.hooks import PROXY_HOOKS
builtins: Final = frozenset(PROXY_HOOKS.values())
hooks: Final = ("async_pre_call_hook", "async_pre_request_hook", "async_pre_call_deployment_hook")
callbacks: Final = litellm.logging_callback_manager.get_custom_loggers_for_type(callback_type=CustomLogger)
return any(
type(callback) not in builtins
and any(getattr(type(callback), hook) is not getattr(CustomLogger, hook) for hook in hooks)
for callback in callbacks
)
def _buckets(prefix_tokens: int, suffix_tokens: int, read_tokens: int, ttl_seconds: int) -> CacheTokenBuckets:
return CacheTokenBuckets(
uncached_input_tokens=suffix_tokens,
cache_read_input_tokens=read_tokens,
cache_creation_5m_input_tokens=prefix_tokens - read_tokens if ttl_seconds == 300 else 0,
cache_creation_1h_input_tokens=prefix_tokens - read_tokens if ttl_seconds == 3600 else 0,
)
def _scenario(model: str, deployment_id: str, tokens: CacheTokenBuckets) -> CacheCostScenario | None:
cost: Final = price_cache_tokens(model=model, deployment_id=deployment_id, tokens=tokens)
return CacheCostScenario(tokens=tokens, input_cost=cost) if cost is not None else None
def _capacity_counter(
limiter: _PROXY_MaxParallelRequestsHandler_v3,
caller: UserAPIKeyAuth,
model_name: str,
request_data: Mapping[str, object],
) -> TokenCounter:
async def count(model: str, api_key: str, body: Mapping[str, JsonValue]) -> int | None:
async with limiter.request_capacity(caller, model_name, request_data=request_data):
return await count_prompt_tokens(model, api_key, body)
return count
def _capacity_request_data(
http_request: Request, caller: UserAPIKeyAuth, request_data: Mapping[str, object]
) -> Mapping[str, object]:
# The parsed-body cache retains only original top-level keys. Replay the
# shared idempotent tag merges on limiter-only data when auth added metadata.
data: Final = dict(request_data) # mutable-ok: the existing tag merge owners accept a dictionary out-param
LiteLLMProxyRequestSetup.apply_client_tag_policy_pre_auth(http_request, data, caller) # pyright: ignore[reportUnknownMemberType] # legacy tag owner takes the validated capacity dictionary
LiteLLMProxyRequestSetup.apply_key_tags_pre_auth(data, caller) # pyright: ignore[reportUnknownMemberType] # legacy tag owner merges trusted key tags into capacity metadata
return MappingProxyType(data)
async def predict_arm(
deployment: Deployment,
body: Mapping[str, JsonValue],
prefix: PromptPrefix,
caller_key_hash: str,
cache: DualCache,
token_counter: TokenCounter,
) -> CachePredictionArm:
deployment_id: Final = deployment.model_info.id or ""
params: Final = deployment.litellm_params
unknown: Final = CachePredictionArm(deployment_id=deployment_id, model=params.model)
if deployment.model_info.blocked:
return unknown.model_copy(update=MappingProxyType({"reason": "unsupported_deployment_configuration"}))
target: Final = resolve_prediction_target(params)
if isinstance(target, UnsupportedPredictionTarget):
return unknown.model_copy(update=MappingProxyType({"reason": target.reason}))
model: Final = target.model
api_key: Final = target.api_key
total_count: Final = await token_counter(model, api_key, body)
prefix_count: Final = await token_counter(model, api_key, prefix.prefix_body)
if total_count is None or prefix_count is None or total_count < prefix_count:
return unknown.model_copy(update=MappingProxyType({"reason": "token_count_unavailable"}))
scope: Final = cache_scope(caller_key_hash, deployment_id, api_key, model)
observation: Final = await lookup(cache, scope, prefix)
exact: Final = observation is not None and observation.fingerprint == prefix.fingerprint
cacheable: Final = observation.cached_tokens if exact and observation is not None else prefix_count
if cacheable > total_count or (observation is not None and observation.cached_tokens > cacheable):
return unknown.model_copy(update=MappingProxyType({"reason": "inconsistent_prefix_token_count"}))
suffix: Final = total_count - cacheable
evidence: Final = (
CacheEvidence(observed_at=observation.observed_at, expires_at=observation.expires_at)
if observation is not None
else None
)
if cacheable < get_prompt_cache_min_tokens(params.model):
disabled: Final = _scenario(model, deployment_id, CacheTokenBuckets(uncached_input_tokens=total_count))
if disabled is None:
return unknown.model_copy(update=MappingProxyType({"reason": "pricing_unavailable"}))
return CachePredictionArm(
deployment_id=deployment_id,
model=model,
cache_state="disabled",
reason="below_cache_minimum",
estimate=disabled,
cold=disabled,
warm=disabled,
token_count_source="anthropic_count_tokens",
)
fresh: Final = observation is not None and observation.expires_at > time.time()
read: Final = observation.cached_tokens if fresh and observation is not None else 0
with pinned_billing_time(current_billing_time()):
cold: Final = _scenario(model, deployment_id, _buckets(cacheable, suffix, 0, prefix.ttl_seconds))
warm: Final = _scenario(model, deployment_id, _buckets(cacheable, suffix, cacheable, prefix.ttl_seconds))
estimate: Final = _scenario(model, deployment_id, _buckets(cacheable, suffix, read, prefix.ttl_seconds))
if cold is None or warm is None or estimate is None:
return unknown.model_copy(update=MappingProxyType({"reason": "pricing_unavailable"}))
return CachePredictionArm(
deployment_id=deployment_id,
model=model,
cache_state="warm" if fresh and exact else "partial" if fresh else "stale" if observation else "unknown",
reason=None if fresh else "observation_expired" if observation else "no_compatible_observation",
estimate=estimate,
cold=cold,
warm=warm,
evidence=evidence,
token_count_source="anthropic_count_tokens",
)
@router.post(
"/cost/predict-cache",
tags=["Cost Tracking"], # mutable-ok: FastAPI requires a list for OpenAPI tags
response_model=CachePredictionResponse,
)
async def predict_cache_cost(
request: CachePredictionRequest,
http_request: Request,
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
) -> CachePredictionResponse:
"""Compare the next native Anthropic request on two configured deployment IDs.
Estimates use provider token counting and recent successful cache telemetry for this key.
Unknown cache state uses the cold scenario when prices/counts are available. Cache observations
do not guarantee retention. v0 supports one message-content breakpoint, text and client tools;
system/tool-only breakpoints, thinking, images, nondefault Anthropic versions, beta headers and
request transforms are unknown.
Each provider count consumes one RPM unit and holds concurrency capacity; a comparison uses
up to four counts. The legacy rate limiter returns unknown without contacting the provider.
This endpoint does not generate tokens, prewarm caches, choose a model or alter routing.
"""
from litellm.proxy.proxy_server import llm_router, proxy_logging_obj
if llm_router is None:
raise HTTPException(status_code=503, detail="Model router is unavailable")
deployments: Final = get_cache_prediction_deployments(
current_deployment_id=request.current_deployment_id,
candidate_deployment_id=request.candidate_deployment_id,
llm_router=llm_router,
team_id=user_api_key_dict.team_id,
)
if deployments is None:
raise HTTPException(status_code=404, detail="Deployment not found")
current, candidate = deployments
for deployment in (current, candidate):
await can_key_call_resolved_model(
model=deployment.model_name,
llm_model_list=llm_router.get_model_list(),
valid_token=user_api_key_dict,
llm_router=llm_router,
)
prefix: Final = parse_prompt(request.request)
caller: Final = user_api_key_dict.api_key
caller_settings: Final = _CallerSettings.model_validate(user_api_key_dict, from_attributes=True)
unsupported_transform: Final = bool(caller_settings.config) or has_request_transforms()
unsupported_headers: Final = not supported_prediction_headers(http_request.headers)
limiter: Final = proxy_logging_obj.get_proxy_hook("parallel_request_limiter")
if (
prefix is None
or not caller
or unsupported_transform
or unsupported_headers
or not isinstance(limiter, _PROXY_MaxParallelRequestsHandler_v3)
):
reason: Final = (
"unsupported_provider_headers"
if unsupported_headers
else "unsupported_request_transform"
if unsupported_transform
else "unsupported_prompt_shape"
if prefix is None
else "caller_identity_unavailable"
if not caller
else "limiter_unavailable"
)
return CachePredictionResponse(
stay=CachePredictionArm(deployment_id=request.current_deployment_id, reason=reason),
switch=CachePredictionArm(deployment_id=request.candidate_deployment_id, reason=reason),
switch_delta=None,
cache_rebuild_penalty=None,
)
request_data: Final = _capacity_request_data(
http_request, user_api_key_dict, _REQUEST_DATA.validate_python(await _read_request_body(http_request))
)
stay: Final = await predict_arm(
current,
request.request,
prefix,
caller,
proxy_logging_obj.internal_usage_cache.dual_cache,
_capacity_counter(limiter, user_api_key_dict, current.model_name, request_data),
)
switch: Final = (
stay
if current.model_info.id == candidate.model_info.id
else await predict_arm(
candidate,
request.request,
prefix,
caller,
proxy_logging_obj.internal_usage_cache.dual_cache,
_capacity_counter(limiter, user_api_key_dict, candidate.model_name, request_data),
)
)
return CachePredictionResponse(
stay=stay,
switch=switch,
switch_delta=(switch.estimate.input_cost - stay.estimate.input_cost)
if switch.estimate is not None and stay.estimate is not None
else None,
cache_rebuild_penalty=(switch.estimate.input_cost - switch.warm.input_cost)
if switch.estimate is not None and switch.warm is not None
else None,
)

View file

@ -156,6 +156,7 @@ from litellm.repositories.verification_token_repository import (
VerificationTokenRepository,
)
from litellm.router import Router
from litellm.types.proxy.auth.auth_checks import UserNotFoundError
from litellm.types.proxy.management_endpoints.common_daily_activity import (
SpendAnalyticsPaginatedResponse,
)
@ -179,6 +180,7 @@ from litellm.types.proxy.management_endpoints.team_endpoints import (
if TYPE_CHECKING:
from prisma import Prisma
from prisma import models as prisma_models
from prisma import types as prisma_types
router: Final = APIRouter()
@ -429,27 +431,26 @@ async def _refresh_cached_team(
)
async def _can_manage_team(
team_obj: LiteLLM_TeamTable,
user_api_key_dict: UserAPIKeyAuth,
) -> bool:
"""True for a proxy admin, an admin of this team, or an org admin for the team's organization."""
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN:
return True
if _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team_obj):
return True
return await _is_user_org_admin_for_team(user_api_key_dict=user_api_key_dict, team_obj=team_obj)
async def _verify_team_access(
team_obj: LiteLLM_TeamTable,
user_api_key_dict: UserAPIKeyAuth,
) -> None:
"""
Verify the caller is authorized to manage the given team.
Access is granted if:
- Caller is a proxy admin, OR
- Caller is an org admin for the team's organization, OR
- Caller is a team admin of this team
Raises HTTPException(403) otherwise.
"""
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN:
return
if _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team_obj):
return
if await _is_user_org_admin_for_team(user_api_key_dict=user_api_key_dict, team_obj=team_obj):
"""Raise HTTPException(403) unless the caller can manage the given team."""
if await _can_manage_team(team_obj=team_obj, user_api_key_dict=user_api_key_dict):
return
raise HTTPException(
@ -4368,6 +4369,20 @@ async def _hydrate_member_user_details(
return tuple(hydrate(m) for m in members)
class _OrganizationModelsRow(BaseModel):
models: list[str] = [] # mutable-ok: pydantic field default
class _TeamRowWithOrganization(BaseModel):
litellm_organization_table: _OrganizationModelsRow | None = None
def _parent_organization_models(team_row: BaseModel) -> list[str] | None:
"""Return the parent org's model allow-list, or None when the team has no org."""
organization: Final = _TeamRowWithOrganization.model_validate(team_row.model_dump()).litellm_organization_table
return organization.models if organization is not None else None
async def _resolve_team_access_group_resources(
_team_info: TeamInfoResponseObjectTeamTable,
) -> TeamInfoResponseObjectTeamTable:
@ -4439,7 +4454,11 @@ async def team_info(
try:
team_info: BaseModel | None = await _team_db(prisma_client).find_unique(
where={"team_id": team_id},
include={"litellm_model_table": True, "object_permission": True},
include={
"litellm_model_table": True,
"object_permission": True,
"litellm_organization_table": True,
},
)
if team_info is None:
raise Exception
@ -4448,9 +4467,12 @@ async def team_info(
status_code=status.HTTP_404_NOT_FOUND,
detail={"message": f"Team not found, passed team id: {team_id}."},
)
await validate_membership(
user_api_key_dict=user_api_key_dict,
team_table=LiteLLM_TeamTable.model_validate(team_info.model_dump()),
team_table: Final = LiteLLM_TeamTable.model_validate(team_info.model_dump())
await validate_membership(user_api_key_dict=user_api_key_dict, team_table=team_table)
organization_models: Final[list[str] | None] = (
_parent_organization_models(team_info)
if await _can_manage_team(team_obj=team_table, user_api_key_dict=user_api_key_dict)
else None
)
## GET ALL KEYS ##
@ -4510,7 +4532,10 @@ async def team_info(
members=resolved_team_info.members_with_roles,
)
hydrated_team_info: Final = resolved_team_info.model_copy(
update={"members_with_roles": hydrated_members} # mutable-ok: pydantic update payload
update={ # mutable-ok: pydantic update payload
"members_with_roles": hydrated_members,
"organization_models": organization_models,
}
)
response_object: Final = TeamInfoResponseObject(
@ -4857,6 +4882,26 @@ async def _get_org_admin_org_ids(
return org_ids if org_ids else None
async def _get_user_team_ids_from_db(
user_id: str,
prisma_client: PrismaClient,
user_api_key_cache: UserApiKeyCache,
proxy_logging_obj: ProxyLogging,
) -> tuple[str, ...]:
try:
user: Final = await get_user_object(
user_id=user_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
user_id_upsert=False,
proxy_logging_obj=proxy_logging_obj,
check_db_only=True,
)
except UserNotFoundError:
return ()
return tuple(user.teams or ()) if user is not None else ()
async def _build_team_list_where_conditions(
prisma_client: PrismaClient,
team_id: str | None,
@ -4867,12 +4912,16 @@ async def _build_team_list_where_conditions(
search: str | None = None,
search_team_id_match: TeamIdSearchMatch = "exact",
org_admin_org_ids: list[str] | None = None,
own_team_ids: tuple[str, ...] = (),
user_api_key_cache: UserApiKeyCache | None = None,
proxy_logging_obj: ProxyLogging | None = None,
) -> dict[str, object] | None:
"""
Build where conditions for team list query.
An org admin listing their own teams sees the union of the teams in the
orgs they administer and `own_team_ids`, the teams they are a member of.
Returns None when the query is guaranteed to yield no results (e.g. user
has no team memberships), allowing the caller to skip the DB round-trip.
"""
@ -4895,6 +4944,11 @@ async def _build_team_list_where_conditions(
if organization_id:
where_conditions["organization_id"] = organization_id
elif org_admin_org_ids is not None and own_team_ids:
org_or_membership_scope: Final[prisma_types.LiteLLM_TeamTableWhereInput] = {
"OR": [{"organization_id": {"in": org_admin_org_ids}}, {"team_id": {"in": list(own_team_ids)}}]
}
where_conditions["AND"] = [org_or_membership_scope]
elif org_admin_org_ids is not None:
# Org admin: always scope to their orgs, even when filtering by user_id.
where_conditions["organization_id"] = {"in": org_admin_org_ids}
@ -5026,66 +5080,72 @@ async def _enforce_list_team_v2_access(
prisma_client: PrismaClient,
user_api_key_cache: UserApiKeyCache,
proxy_logging_obj: ProxyLogging,
) -> tuple[str | None, list[str] | None]:
) -> tuple[str | None, list[str] | None, tuple[str, ...]]:
"""Enforce access control for list_team_v2.
- Proxy admins and admin viewers can query any teams.
- Org admins can query teams within their organizations.
- Org admins can query teams within their organizations, plus the teams
they are a member of when listing their own teams.
- Regular users can only query their own teams.
Returns the (possibly overridden) user_id and org_admin_org_ids.
Returns the (possibly overridden) user_id, org_admin_org_ids and, for an
org admin's own query, the caller's own team ids.
"""
is_proxy_admin: Final = _user_has_admin_view(user_api_key_dict)
org_admin_org_ids: list[str] | None = None
caller_user_id: Final = user_api_key_dict.user_id
if is_proxy_admin:
return user_id, org_admin_org_ids
return user_id, None, ()
# Always check org admin status so that even own-queries see
# the full set of organisation teams, not just direct memberships.
if user_api_key_dict.user_id:
org_admin_org_ids = await _get_org_admin_org_ids(
user_id=user_api_key_dict.user_id,
org_admin_org_ids: Final = (
await _get_org_admin_org_ids(
user_id=caller_user_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
if caller_user_id
else None
)
if org_admin_org_ids is not None:
if caller_user_id and org_admin_org_ids is not None:
# Org admin: validate org_id filter if provided
if organization_id and organization_id not in org_admin_org_ids:
raise HTTPException(
status_code=403,
detail={"error": "You can only view teams within your organizations."},
)
# When the caller is an org admin querying their own teams (or no
# specific user), null out user_id so that
# _build_team_list_where_conditions scopes only by organization_id
# — org admins should see all teams in their orgs, not just teams
# they are a direct member of. Keep user_id when the org admin
# explicitly queries a *different* user's teams.
if user_id is None or user_id == user_api_key_dict.user_id:
user_id = None
is_own_query: Final = user_id is None or user_id == caller_user_id
own_team_ids: Final = (
await _get_user_team_ids_from_db(
user_id=caller_user_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
if is_own_query
else ()
)
verbose_proxy_logger.debug(
"list_team_v2: org admin access for user=%s, org_ids=%s, user_id_filter=%s",
user_api_key_dict.user_id,
_sanitize_for_log(caller_user_id),
org_admin_org_ids,
user_id,
_sanitize_for_log(None if is_own_query else user_id),
)
else:
# Not an org admin — fall back to standard route check
if not allowed_route_check_inside_route(user_api_key_dict=user_api_key_dict, requested_user_id=user_id):
raise HTTPException(
status_code=401,
detail={
"error": f"Only admin users can query all teams/other teams. Your user role={user_api_key_dict.user_role}"
},
)
# Regular user — auto-inject caller's user_id
if user_id is None:
user_id = user_api_key_dict.user_id
return None if is_own_query else user_id, org_admin_org_ids, own_team_ids
return user_id, org_admin_org_ids
# Not an org admin — fall back to standard route check
if not allowed_route_check_inside_route(user_api_key_dict=user_api_key_dict, requested_user_id=user_id):
raise HTTPException(
status_code=401,
detail={
"error": f"Only admin users can query all teams/other teams. Your user role={user_api_key_dict.user_role}"
},
)
# Regular user — auto-inject caller's user_id
return user_id if user_id is not None else caller_user_id, None, ()
@router.get(
@ -5163,7 +5223,7 @@ async def list_team_v2(
)
# --- Access control ---
user_id, org_admin_org_ids = await _enforce_list_team_v2_access(
user_id, org_admin_org_ids, own_team_ids = await _enforce_list_team_v2_access(
user_api_key_dict=user_api_key_dict,
user_id=user_id,
organization_id=organization_id,
@ -5195,6 +5255,7 @@ async def list_team_v2(
search=search,
search_team_id_match=search_team_id_match,
org_admin_org_ids=org_admin_org_ids,
own_team_ids=own_team_ids,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
@ -5291,17 +5352,16 @@ async def _authorize_and_filter_teams(
- Proxy admins: all teams (or filtered by user_id if provided).
- Org admins: teams from their orgs (scoped to user_id if provided).
- Own query (user_id matches caller): teams the user is a member of.
- Own query (user_id matches caller): teams the user is a member of, across all orgs.
- Others: 401.
"""
is_proxy_admin: Final = _user_has_admin_view(user_api_key_dict)
is_own_query: Final = (
user_id is not None and user_api_key_dict.user_id is not None and user_api_key_dict.user_id == user_id
)
allowed_org_ids: list[str] | None = None
if not is_proxy_admin:
is_own_query: Final = (
user_id is not None and user_api_key_dict.user_id is not None and user_api_key_dict.user_id == user_id
)
# Check if user is an org admin (even for own queries, so they see org teams)
if user_api_key_dict.user_id is not None:
caller_user: Final = await get_user_object(
@ -5328,33 +5388,30 @@ async def _authorize_and_filter_teams(
},
)
if allowed_org_ids is not None:
# Org admin: query DB for teams in their orgs
if allowed_org_ids is not None and not is_own_query:
org_teams: Final = await _raw_team_db(TeamRepository(prisma_client)).find_many(
where={"organization_id": {"in": allowed_org_ids}},
include={"litellm_model_table": True},
)
if not user_id:
return list(org_teams)
# Filter org teams to only those where the target user is a member
return [
team
for team in org_teams
if team.members_with_roles and any(m.get("user_id") == user_id for m in team.members_with_roles)
]
elif user_id:
# Regular user: fetch all and filter by membership (Prisma can't filter JSON arrays)
response: Final = await _raw_team_db(TeamRepository(prisma_client)).find_many(
include={"litellm_model_table": True}
)
return [
team
for team in response
if team.members_with_roles and any(m.get("user_id") == user_id for m in team.members_with_roles)
]
else:
response: Final = await _raw_team_db(TeamRepository(prisma_client)).find_many(include={"litellm_model_table": True})
if not user_id:
# Proxy admin: all teams
return list(await _raw_team_db(TeamRepository(prisma_client)).find_many(include={"litellm_model_table": True}))
return list(response)
# Prisma can't filter JSON arrays, so membership is filtered in Python
return [
team
for team in response
if team.members_with_roles and any(m.get("user_id") == user_id for m in team.members_with_roles)
]
@router.get("/team/list", tags=["team management"], dependencies=[Depends(user_api_key_auth)])

View file

@ -101,6 +101,7 @@ from litellm.proxy.common_utils.admin_ui_utils import (
admin_ui_disabled,
show_missing_vars_in_env,
)
from litellm.proxy.common_utils.html_forms.default_credentials_hint import should_hide_default_credentials_hint
from litellm.proxy.common_utils.html_forms.jwt_display_template import (
jwt_display_template,
)
@ -1110,10 +1111,7 @@ async def google_login(
from fastapi.responses import HTMLResponse
hide_default_credentials_hint: Final = (
os.getenv("LITELLM_HIDE_DEFAULT_CREDENTIALS_HINT", "false").lower() == "true"
or general_settings.get("hide_default_credentials_hint", False) is True
)
hide_default_credentials_hint: Final = should_hide_default_credentials_hint(general_settings)
form_response: Final = HTMLResponse(
content=build_ui_login_form(
show_deprecation_banner=True,

View file

@ -115,15 +115,15 @@ async def run_team_metadata_validation(
"error": f"custom_team_metadata_validate is an Enterprise feature. {CommonProxyErrors.not_premium_user.value}"
},
)
if not (
inspect.iscoroutinefunction(validator) or inspect.iscoroutinefunction(getattr(validator, "__call__", None))
):
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail={ # mutable-ok: HTTPException.detail has no immutable form
"error": "custom_team_metadata_validate must be an async function"
},
)
if not inspect.iscoroutinefunction(validator):
validator_call: Final = getattr(validator, "__call__", None) # noqa: B004 # value unwrap for the functor check
if not inspect.iscoroutinefunction(validator_call):
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail={ # mutable-ok: HTTPException.detail has no immutable form
"error": "custom_team_metadata_validate must be an async function"
},
)
try:
raw_result: Final = await asyncio.wait_for(validator(payload), timeout=timeout_seconds)

View file

@ -78,6 +78,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import (
)
from litellm.proxy.openai_files_endpoints.general_upload_validation import (
MB,
check_allowed_extension,
check_blocked_extension,
check_unsafe_filename,
check_upload_file_size,
@ -473,6 +474,11 @@ async def create_file(
if general_size_failure is not None:
raise_upload_validation_failure(general_size_failure)
allowed_extensions: Final = coerce_optional_str_list_setting(general_settings.get("allowed_file_extensions"))
allowed_extension_failure: Final = check_allowed_extension(file.filename, allowed_extensions)
if allowed_extension_failure is not None:
raise_upload_validation_failure(allowed_extension_failure)
blocked_extensions: Final = coerce_optional_str_list_setting(general_settings.get("blocked_file_extensions"))
blocked_extension_failure: Final = check_blocked_extension(file.filename, blocked_extensions)
if blocked_extension_failure is not None:

View file

@ -2,8 +2,8 @@
Upload validation applied to every purpose at POST /v1/files.
batch_file_validation.py checks the JSONL shape of purpose="batch" uploads; this
module applies the same fast-fail-before-forwarding shape (size cap, blocked
extensions, path-traversal filenames) regardless of purpose.
module applies the same fast-fail-before-forwarding shape (size cap, allowed and
blocked extensions, path-traversal filenames) regardless of purpose.
"""
from dataclasses import dataclass
@ -31,10 +31,9 @@ def coerce_optional_int_setting(raw: object) -> int | None:
raise TypeError(f"expected an integer, got {raw!r}")
def coerce_optional_str_list_setting(raw: object) -> tuple[str, ...]:
"""A general_settings value declared as an optional list of strings, e.g. blocked_file_extensions."""
def coerce_optional_str_list_setting(raw: object) -> tuple[str, ...] | None:
if raw is None:
return ()
return None
if not isinstance(raw, list) or not all(isinstance(item, str) for item in raw):
raise TypeError(f"expected a list of strings, got {raw!r}")
return tuple(raw)
@ -46,6 +45,11 @@ class UploadedFileTooLarge:
limit_mb: int
@dataclass(frozen=True, slots=True)
class UploadedFileExtensionNotAllowed:
extension: str
@dataclass(frozen=True, slots=True)
class UploadedFileBlockedExtension:
extension: str
@ -56,7 +60,9 @@ class UploadedFileUnsafeFilename:
filename: str
UploadValidationFailure = UploadedFileTooLarge | UploadedFileBlockedExtension | UploadedFileUnsafeFilename
UploadValidationFailure = (
UploadedFileTooLarge | UploadedFileExtensionNotAllowed | UploadedFileBlockedExtension | UploadedFileUnsafeFilename
)
def _file_size_bytes(file_source: bytes | BinaryIO) -> int:
@ -81,19 +87,35 @@ def check_upload_file_size(
return None
def _normalized_extension(filename: str | None) -> str:
if not filename:
return ""
try:
return Path(safe_filename(filename)).suffix.lower()
except ValueError:
return ""
def check_allowed_extension(
filename: str | None,
allowed_extensions: tuple[str, ...] | None,
) -> UploadedFileExtensionNotAllowed | None:
if allowed_extensions is None:
return None
extension: Final = _normalized_extension(filename)
normalized_allowed: Final = frozenset(item.lower() for item in allowed_extensions)
if extension and extension in normalized_allowed:
return None
return UploadedFileExtensionNotAllowed(extension=extension)
def check_blocked_extension(
filename: str | None,
blocked_extensions: tuple[str, ...],
blocked_extensions: tuple[str, ...] | None,
) -> UploadedFileBlockedExtension | None:
if not blocked_extensions or not filename:
if not blocked_extensions:
return None
try:
extension: Final = Path(safe_filename(filename)).suffix.lower()
except ValueError:
return None
# The uploaded name's extension is normalized above; blocked_extensions comes
# straight from config.yaml or the DB and is normalized here too, so a
# differently-cased entry (".EXE") still catches a lowercase upload.
extension: Final = _normalized_extension(filename)
normalized_blocked: Final = frozenset(item.lower() for item in blocked_extensions)
if extension and extension in normalized_blocked:
return UploadedFileBlockedExtension(extension=extension)
@ -128,6 +150,17 @@ def raise_upload_validation_failure(failure: UploadValidationFailure) -> NoRetur
param="file",
code=413,
)
case UploadedFileExtensionNotAllowed(extension=extension):
raise ProxyException(
message=(
(f"File extension '{extension}'" if extension else "A file without an extension")
+ " is not in this proxy's allowed_file_extensions setting. "
"The file was not forwarded to the provider."
),
type="invalid_request_error",
param="file",
code=400,
)
case UploadedFileBlockedExtension(extension=extension):
raise ProxyException(
message=(

View file

@ -270,6 +270,24 @@ class PassThroughStreamingHandler:
- Vertex AI
- OpenAI
"""
from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import (
_is_message_stop_chunk, # pyright: ignore[reportPrivateUsage] # both native stream paths share terminal-event detection
_is_provider_error_chunk, # pyright: ignore[reportPrivateUsage] # provider errors must not become cache evidence
)
# Transport reads can split event names and JSON payloads. Recognize terminal
# events only after the shared SSE framer has reassembled the collected bytes.
complete_frames, incomplete_tail = split_complete_sse_frames(
b"".join(raw_bytes) if endpoint_type == EndpointType.ANTHROPIC else b""
)
litellm_logging_obj.model_call_details[ # rebind-ok: stamp evidence on the per-request state read by callbacks
"prompt_cache_response_complete"
] = (
endpoint_type == EndpointType.ANTHROPIC
and not incomplete_tail.strip()
and _is_message_stop_chunk(complete_frames)
and not _is_provider_error_chunk(complete_frames)
)
try:
(
standard_logging_response_object,

View file

@ -250,7 +250,7 @@ import litellm._redis
from litellm import Router
from litellm._logging import _redact_string, verbose_proxy_logger, verbose_router_logger
from litellm.caching.caching import DualCache, RedisCache
from litellm.caching.redis_cache import RedisCircuitBreakerOpenError
from litellm.caching.redis_cache import RedisCircuitBreakerOpenError, is_redis_timeout_failure
from litellm.caching.redis_cluster_cache import RedisClusterCache
from litellm.constants import (
_REALTIME_BODY_CACHE_SIZE,
@ -274,6 +274,8 @@ from litellm.constants import (
PROXY_BUDGET_RESCHEDULER_MAX_TIME,
PROXY_BUDGET_RESCHEDULER_MIN_TIME,
PROXY_CONFIG_RELOAD_INTERVAL_SECONDS,
REALTIME_SESSION_FAILURE_LOGGED_KEY,
REALTIME_SESSION_SUCCESS_LOGGED_KEY,
ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG,
USER_SPEND_ALERTS_JOB_ID,
WEEKLY_SPEND_REPORT_JOB_ID,
@ -365,6 +367,7 @@ from litellm.proxy.common_utils.healthy_model_filter import (
get_hidden_unhealthy_model_names,
is_healthy_only_listing_default,
)
from litellm.proxy.common_utils.html_forms.default_credentials_hint import should_hide_default_credentials_hint
from litellm.proxy.common_utils.html_forms.ui_login import build_ui_login_form
from litellm.proxy.common_utils.http_parsing_utils import (
_read_request_body,
@ -3450,8 +3453,10 @@ async def _invalidate_spend_counter(counter_key: str):
async def _apply_spend_counter_increments(pending: Sequence[PendingSpendIncrement]) -> None:
try:
await increment_spend_counters_pipeline(pending=pending)
except RedisCircuitBreakerOpenError:
return
except Exception as e:
if isinstance(e, RedisCircuitBreakerOpenError) or is_redis_timeout_failure(e):
return
raise
async def increment_spend_counters_pipeline(pending: Sequence[PendingSpendIncrement]) -> None:
@ -7065,6 +7070,9 @@ class ProxyConfig:
if "max_file_size_mb" not in self._yaml_general_settings_keys:
general_settings["max_file_size_mb"] = _general_settings.get("max_file_size_mb")
if "allowed_file_extensions" not in self._yaml_general_settings_keys:
general_settings["allowed_file_extensions"] = _general_settings.get("allowed_file_extensions")
if "blocked_file_extensions" not in self._yaml_general_settings_keys:
general_settings["blocked_file_extensions"] = _general_settings.get("blocked_file_extensions")
@ -11893,6 +11901,13 @@ async def _release_realtime_budget_reservation(user_api_key_dict: UserAPIKeyAuth
)
async def _release_realtime_max_parallel_slot(user_api_key_dict: UserAPIKeyAuth) -> None:
release_like_http_disconnect: Final = (
proxy_logging_obj._arelease_max_parallel_requests_on_disconnect # pyright: ignore[reportPrivateUsage] # shared
)
await release_like_http_disconnect(user_api_key_dict)
async def _reject_realtime_session(
websocket: WebSocket,
user_api_key_dict: UserAPIKeyAuth,
@ -11912,6 +11927,7 @@ async def _reject_realtime_session(
await websocket.close(code=code, reason=reason)
finally:
await _release_realtime_budget_reservation(user_api_key_dict)
await _release_realtime_max_parallel_slot(user_api_key_dict)
@app.websocket("/openai/v1/realtime")
@ -12015,6 +12031,9 @@ async def realtime_websocket_endpoint(
websocket, user_api_key_dict, code=1011, reason="Pre-call error", error_message=str(e)
)
return
except BaseException:
await _release_realtime_max_parallel_slot(user_api_key_dict)
raise
# Phase 2: route to upstream LLM.
try:
@ -12044,12 +12063,10 @@ async def realtime_websocket_endpoint(
except Exception: # noqa: BLE001 # the lower layer may have closed the socket already; closing twice is not an error
verbose_proxy_logger.debug("Could not close realtime client websocket; it is already gone")
finally:
from litellm.litellm_core_utils.realtime_streaming import (
REALTIME_SESSION_SUCCESS_LOGGED_KEY,
)
if not litellm_logging_obj.model_call_details.get(REALTIME_SESSION_SUCCESS_LOGGED_KEY):
await _release_realtime_budget_reservation(user_api_key_dict)
if not litellm_logging_obj.model_call_details.get(REALTIME_SESSION_FAILURE_LOGGED_KEY):
await _release_realtime_max_parallel_slot(user_api_key_dict)
######################################################################
@ -15033,6 +15050,13 @@ def _get_proxy_model_info(model: dict) -> dict:
return _translate_model_name_for_response(model)
def _model_info_json_response(data: Sequence[Mapping[str, object]] | Mapping[str, object]) -> Response:
return Response(
content=orjson.dumps({"data": data}, default=jsonable_encoder, option=orjson.OPT_NON_STR_KEYS),
media_type="application/json",
)
@router.get(
"/model/info",
tags=["model management"],
@ -15080,7 +15104,7 @@ async def model_info_v1(
`model_info.direct_access` when the proxy database is connected.
Returns:
Returns a dictionary containing information about each model.
A JSON response whose `data` list holds one entry per model.
Example Response:
```json
@ -15128,7 +15152,7 @@ async def model_info_v1(
deployment_dict=_deployment_info_dict,
excluded_keys={"litellm_credential_name"},
)
return {"data": _deployment_info_dict}
return _model_info_json_response(_deployment_info_dict)
if llm_model_list is None:
raise HTTPException(
@ -15179,7 +15203,7 @@ async def model_info_v1(
llm_router=llm_router,
user_api_key_dict=user_api_key_dict,
)
return {"data": single_model_list}
return _model_info_json_response(single_model_list)
# Return router deployments (same source as /v2/model/info), not wildcard-
# expanded model names from get_complete_model_list(). Team-scoped rows
@ -15247,7 +15271,7 @@ async def model_info_v1(
visible_models: Final = [model for model in all_models if model.get("model_name") not in hidden_names]
verbose_proxy_logger.debug("all_models: %s", visible_models)
return {"data": visible_models}
return _model_info_json_response(visible_models)
@router.get(
@ -15816,10 +15840,7 @@ async def fallback_login(request: Request):
from fastapi.responses import HTMLResponse
hide_default_credentials_hint: Final = (
os.getenv("LITELLM_HIDE_DEFAULT_CREDENTIALS_HINT", "false").lower() == "true"
or general_settings.get("hide_default_credentials_hint", False) is True
)
hide_default_credentials_hint: Final = should_hide_default_credentials_hint(general_settings)
return HTMLResponse(
content=build_ui_login_form(
show_deprecation_banner=False,
@ -17036,6 +17057,7 @@ _GENERAL_SETTINGS_CONFIG_LIST_FIELD_TYPES: Final[Mapping[str, str]] = MappingPro
"max_request_size_mb": "Integer",
"max_batch_file_size_mb": "Integer",
"max_file_size_mb": "Integer",
"allowed_file_extensions": "List",
"blocked_file_extensions": "List",
"max_response_size_mb": "Integer",
"proxy_config_reload_interval_seconds": "Integer",

View file

@ -159,7 +159,7 @@ def _get_spend_logs_metadata(
requester_ip_address=None,
additional_usage_values=None,
applied_guardrails=None,
status=None or "success",
status="success",
error_information=None,
proxy_server_request=None,
batch_models=None,

View file

@ -4,6 +4,7 @@ import copy
import hashlib
import inspect
import json
import math
import os
import smtplib
import ssl
@ -6404,7 +6405,7 @@ class PrismaClient:
return None
try:
value: Final = float(response_time_ms)
return value if value == value and value not in (float("inf"), float("-inf")) else None
return value if math.isfinite(value) else None
except (ValueError, TypeError):
verbose_proxy_logger.warning("Invalid response_time_ms value: %s", response_time_ms)
return None

View file

@ -43,6 +43,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.responses.litellm_completion_transformation.session_handler import (
ResponsesSessionHandler,
)
from litellm.types.llms.base import CachedTokensDetails
from litellm.types.llms.openai import (
AllMessageValues,
ChatCompletionAssistantMessage,
@ -2816,27 +2817,24 @@ class LiteLLMCompletionResponsesConfig:
# Translate prompt_tokens_details to input_tokens_details
if hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details is not None:
prompt_details: Final = usage.prompt_tokens_details
input_details_dict: Final[dict[str, int]] = {}
if hasattr(prompt_details, "cached_tokens") and prompt_details.cached_tokens is not None:
input_details_dict["cached_tokens"] = prompt_details.cached_tokens
else:
input_details_dict["cached_tokens"] = 0
if hasattr(prompt_details, "text_tokens") and prompt_details.text_tokens is not None:
input_details_dict["text_tokens"] = prompt_details.text_tokens
if hasattr(prompt_details, "audio_tokens") and prompt_details.audio_tokens is not None:
input_details_dict["audio_tokens"] = prompt_details.audio_tokens
cache_write_tokens = getattr(prompt_details, "cache_write_tokens", None) or getattr(
cached_tokens_details: Final = getattr(prompt_details, "cached_tokens_details", None)
cache_write_tokens: Final = getattr(prompt_details, "cache_write_tokens", None) or getattr(
prompt_details, "cache_creation_tokens", None
)
if cache_write_tokens is not None:
input_details_dict["cache_write_tokens"] = cache_write_tokens
if input_details_dict:
response_usage.input_tokens_details = InputTokensDetails(**input_details_dict)
cache_write_extra: Final[Mapping[str, int]] = (
MappingProxyType({"cache_write_tokens": cache_write_tokens})
if cache_write_tokens is not None
else MappingProxyType({})
)
response_usage.input_tokens_details = InputTokensDetails(
cached_tokens=prompt_details.cached_tokens if prompt_details.cached_tokens is not None else 0,
text_tokens=prompt_details.text_tokens,
audio_tokens=prompt_details.audio_tokens,
cached_tokens_details=(
cached_tokens_details if isinstance(cached_tokens_details, CachedTokensDetails) else None
),
**cache_write_extra,
)
# Translate completion_tokens_details to output_tokens_details
if hasattr(usage, "completion_tokens_details") and usage.completion_tokens_details is not None:

View file

@ -1179,6 +1179,9 @@ class ResponseAPILoggingUtils:
audio_tokens=getattr(response_api_usage.input_tokens_details, "audio_tokens", None),
text_tokens=getattr(response_api_usage.input_tokens_details, "text_tokens", None),
image_tokens=getattr(response_api_usage.input_tokens_details, "image_tokens", None),
cached_tokens_details=getattr(
response_api_usage.input_tokens_details, "cached_tokens_details", None
),
cache_write_tokens=getattr(response_api_usage.input_tokens_details, "cache_write_tokens", None),
)
completion_tokens_details: CompletionTokensDetailsWrapper | None = None

View file

@ -96,7 +96,6 @@ from litellm.litellm_core_utils.request_timeout_resolver import (
from litellm.litellm_core_utils.secret_redaction import redact_string
from litellm.litellm_core_utils.sensitive_data_masker import (
SensitiveDataMasker,
mask_credentials_in_payload,
mask_sensitive_structure,
)
from litellm.litellm_core_utils.token_counter import offload_token_count
@ -623,20 +622,6 @@ def _replay_live_router_model_cost() -> None:
set_live_deployment_replay(_replay_live_router_model_cost)
# Kwargs that carry no signal about the failed attempt, so log_retry drops them from a
# breadcrumb entirely: the request payload, the proxy's snapshot of the inbound request (its body
# aliases the live request metadata, earlier breadcrumbs included, so copying it would nest every
# breadcrumb inside the next one), and the router-internal walk state. Credentials are handled
# separately by mask_credentials_in_payload, which scrubs credential-named values from whatever
# kwargs remain rather than trying to enumerate every credential-bearing key here.
RETRY_BREADCRUMB_EXCLUDED_KWARGS: Final = frozenset(
(
"messages",
"original_function",
"attempted_targets",
"proxy_server_request",
)
)
RETRY_BREADCRUMB_LIMIT: Final = 4
@ -1553,6 +1538,18 @@ class Router:
return False
return sum(len(self.model_name_to_deployment_indices.get(member) or ()) for member in group.models) > 1
def team_model_has_alternatives(self, deployment_id: str) -> bool:
deployment: Final = self.get_deployment(model_id=deployment_id)
if deployment is None:
return False
team_id: Final = deployment.model_info.team_id
public_model_name: Final = deployment.model_info.team_public_model_name
if team_id is None or public_model_name is None:
return False
sibling_indices: Final = self.team_model_to_deployment_indices.get((team_id, public_model_name)) or ()
routable_siblings: Final = self._filter_blocked_deployments([self.model_list[idx] for idx in sibling_indices])
return len(routable_siblings) > 1
_OVERRIDABLE_ROUTING_STRATEGIES: frozenset[str] = frozenset({"simple-shuffle", *_DEFAULT_SELECTOR_ATTR_BY_STRATEGY})
def _get_request_routing_strategy_override(self, request_kwargs: dict | None) -> str | None:
@ -4108,16 +4105,16 @@ class Router:
models: Final = [m.strip() for m in model.split(",")]
async def _async_completion_no_exceptions(
model: str, messages: list[dict[str, str]], stream: bool, **kwargs: Any
model_name: str, messages: list[dict[str, str]], stream: bool, **kwargs: Any
) -> ModelResponse | CustomStreamWrapper | Exception:
"""
Wrapper around self.acompletion that catches exceptions and returns them as a result
"""
try:
result = await self.acompletion(model=model, messages=messages, stream=stream, **kwargs)
result = await self.acompletion(model=model_name, messages=messages, stream=stream, **kwargs)
return result
except asyncio.CancelledError:
verbose_router_logger.debug("Received 'task.cancel'. Cancelling call w/ model=%s.", model)
verbose_router_logger.debug("Received 'task.cancel'. Cancelling call w/ model=%s.", model_name)
raise
except Exception as e:
return e
@ -4144,9 +4141,9 @@ class Router:
except KeyError:
pass
for model in models:
for model_name in models:
task = asyncio.create_task(
_async_completion_no_exceptions(model=model, messages=messages, stream=stream, **kwargs)
_async_completion_no_exceptions(model_name=model_name, messages=messages, stream=stream, **kwargs)
)
pending_tasks.append(task)
@ -8374,31 +8371,30 @@ class Router:
def log_retry(self, kwargs: dict, e: Exception) -> dict:
"""
When a retry or fallback happens, log the details of the just failed model call - similar to Sentry breadcrumbing
When a retry or fallback happens, record which model group, deployment and attempt just failed and why
"""
from litellm.types.router import RetryAttemptRecord
_metadata_var: Final = "litellm_metadata" if "litellm_metadata" in kwargs else "metadata"
request_metadata: Final[Mapping[str, object]] = kwargs[_metadata_var]
attempt_kwargs: Final = MappingProxyType(
{k: v for k, v in kwargs.items() if k != _metadata_var and k not in RETRY_BREADCRUMB_EXCLUDED_KWARGS}
)
attempt_metadata: Final = MappingProxyType(
{k: v for k, v in request_metadata.items() if k != "previous_models"}
)
previous_model: Final = MappingProxyType(
{
"exception_type": type(e).__name__,
"exception_string": str(e),
**attempt_kwargs,
_metadata_var: attempt_metadata,
}
)
model_group: Final = kwargs.get("model")
model_info: Final = request_metadata.get("model_info")
deployment_id: Final = model_info.get("id") if isinstance(model_info, Mapping) else None
attempted_retries: Final = request_metadata.get("attempted_retries")
attempt_record: Final[RetryAttemptRecord] = {
"model_group": model_group if isinstance(model_group, str) else None,
"deployment_id": deployment_id if isinstance(deployment_id, str) else None,
"exception_type": type(e).__name__,
"exception_string": str(e),
"attempted_retries": attempted_retries if type(attempted_retries) is int else None,
}
earlier_breadcrumbs: Final = request_metadata.get("previous_models")
kept_breadcrumbs: Final[tuple[object, ...]] = (
tuple(earlier_breadcrumbs)[-(RETRY_BREADCRUMB_LIMIT - 1) :]
if isinstance(earlier_breadcrumbs, (list, tuple))
else ()
)
breadcrumbs: Final = (*kept_breadcrumbs, mask_credentials_in_payload(previous_model))
breadcrumbs: Final = (*kept_breadcrumbs, attempt_record)
kwargs[_metadata_var]["previous_models"] = breadcrumbs # rebind-ok: the logging object already holds this dict
return kwargs
@ -13878,6 +13874,7 @@ class Router:
cooldown_time=_cooldown_time,
enable_pre_call_checks=self.enable_pre_call_checks,
cooldown_list=_cooldown_list,
model_ids=model_ids,
)
if strategy == "simple-shuffle":
@ -13910,6 +13907,7 @@ class Router:
cooldown_time=_cooldown_time,
enable_pre_call_checks=self.enable_pre_call_checks,
cooldown_list=_cooldown_list,
model_ids=model_ids,
)
self._override_selector_pre_call_check(strategy, strategy_selector, deployment)
verbose_router_logger.info(
@ -13987,6 +13985,11 @@ class Router:
model=model,
llm_provider="",
)
pass_through_model_ids: Final = tuple(
deployment["model_info"]["id"]
for deployment in pass_through_deployments
if "id" in deployment.get("model_info", {})
)
# 4. Apply health-check and cooldown filtering
parent_otel_span: Final[Span | None] = _get_parent_otel_span_from_kwargs(request_kwargs)
@ -14024,6 +14027,7 @@ class Router:
cooldown_time=_cooldown_time,
enable_pre_call_checks=self.enable_pre_call_checks,
cooldown_list=_cooldown_list,
model_ids=pass_through_model_ids,
)
# 6. Apply load balancing strategy
@ -14057,6 +14061,7 @@ class Router:
cooldown_time=_cooldown_time,
enable_pre_call_checks=self.enable_pre_call_checks,
cooldown_list=_cooldown_list,
model_ids=model_ids,
)
self._override_selector_pre_call_check(strategy, strategy_selector, deployment)

View file

@ -343,8 +343,9 @@ def _should_cooldown_deployment(
model_group: Final = litellm_router_instance.get_model_group(id=deployment)
is_single_deployment_model_group = False
if model_group is not None and len(model_group) == 1:
is_single_deployment_model_group = not litellm_router_instance.routing_group_has_alternatives(
requested_model_group
is_single_deployment_model_group = not (
litellm_router_instance.routing_group_has_alternatives(requested_model_group)
or litellm_router_instance.team_model_has_alternatives(deployment)
)
## CHECK DEPLOYMENT-LEVEL POLICY FIRST (overrides router-level)

View file

@ -93,4 +93,5 @@ async def async_raise_no_deployment_exception(
cooldown_time=_cooldown_time,
enable_pre_call_checks=litellm_router_instance.enable_pre_call_checks,
cooldown_list=cooldown_list_ids,
model_ids=model_ids,
)

View file

@ -99,23 +99,13 @@ def setup(
def check_limits(kwargs: Mapping[str, object]) -> None:
import litellm
from litellm.litellm_core_utils.core_helpers import max_retries_per_request_hit
current_cost: Final = litellm._current_cost # pyright: ignore[reportPrivateUsage] # shared SDK budget counter has no public accessor
if litellm.max_budget and current_cost > litellm.max_budget:
raise litellm.BudgetExceededError(current_cost=current_cost, max_budget=litellm.max_budget)
metadata: Final = kwargs.get("metadata")
if isinstance(metadata, Mapping):
typed_metadata: Final = cast( # cast-ok: runtime Mapping check establishes read-only metadata
Mapping[str, object], metadata
)
previous: Final = typed_metadata.get("previous_models")
if (
isinstance(previous, list)
and litellm.num_retries_per_request is not None
and len(cast(list[object], previous)) # cast-ok: runtime list check establishes the retry history
>= litellm.num_retries_per_request
):
raise RuntimeError("Max retries per request hit!")
if max_retries_per_request_hit(kwargs, litellm.num_retries_per_request):
raise RuntimeError("Max retries per request hit!")
def finalize(

View file

@ -552,6 +552,16 @@ class BedrockGuardrailConfigModel(BaseModel):
"still rejects is bisected automatically, so this value only trades round trips against "
"batch size and cannot fail a request on its own.",
)
contextual_grounding_from_messages: bool = Field(
default=False,
description="ApplyGuardrail: when True, post-call scans of a request with no grounding_source / "
"query content parts send the system and developer messages as the grounding source and "
"the latest user message as the query, so the guardrail's contextual grounding policy can "
"score the response. Bedrock bills contextual grounding units for these scans and rejects "
"queries, sources and responses over its contextual grounding length limits, so leave this "
"off for guardrails without a contextual grounding policy. Default False: plain messages "
"are never sent as grounding context.",
)
class BedrockGuardrailStreamingParams(BaseModel):

View file

@ -75,3 +75,9 @@ class HiddenParams(OpenAIObject):
data: Final = super().model_dump(**kwargs)
data["_response_ms"] = self._response_ms
return data
class CachedTokensDetails(BaseModel):
text_tokens: int | None = None
audio_tokens: int | None = None
image_tokens: int | None = None

View file

@ -91,6 +91,8 @@ from litellm.types.responses.main import (
OutputImageGenerationCall,
)
from .base import CachedTokensDetails
FileContent = IO[bytes] | bytes | PathLike
FileTypes = (
@ -1288,6 +1290,7 @@ class OutputTokensDetails(BaseLiteLLMOpenAIResponseObject):
class InputTokensDetails(BaseLiteLLMOpenAIResponseObject):
audio_tokens: int | None = None
cached_tokens: int = 0
cached_tokens_details: CachedTokensDetails | None = None
text_tokens: int | None = None
model_config = {"extra": "allow"}
@ -2254,10 +2257,17 @@ class OpenAIRealtimeInputAudioTranscriptionCompleted(TypedDict):
usage: NotRequired[ReadOnly[Mapping[str, object]]]
class OpenAIRealtimeCachedTokensDetails(TypedDict, total=False):
text_tokens: ReadOnly[int]
audio_tokens: ReadOnly[int]
image_tokens: ReadOnly[int]
class OpenAIRealtimeUsageTokenDetails(TypedDict):
audio_tokens: ReadOnly[int]
text_tokens: ReadOnly[int]
cached_tokens: NotRequired[ReadOnly[int]]
cached_tokens_details: NotRequired[ReadOnly[OpenAIRealtimeCachedTokensDetails]]
class OpenAIRealtimeResponseUsage(TypedDict):

View file

@ -0,0 +1,67 @@
from collections.abc import Mapping
from typing import Annotated, Literal, TypeAlias
from pydantic import BaseModel, ConfigDict, Field, JsonValue, StrictInt
TokenCount: TypeAlias = Annotated[StrictInt, Field(ge=0)]
class CacheTokenBuckets(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
uncached_input_tokens: TokenCount = 0
cache_read_input_tokens: TokenCount = 0
cache_creation_5m_input_tokens: TokenCount = 0
cache_creation_1h_input_tokens: TokenCount = 0
@property
def total_tokens(self) -> int:
return (
self.uncached_input_tokens
+ self.cache_read_input_tokens
+ self.cache_creation_5m_input_tokens
+ self.cache_creation_1h_input_tokens
)
class CacheEvidence(BaseModel):
model_config = ConfigDict(frozen=True)
observed_at: float
expires_at: float
source: Literal["provider_usage"] = "provider_usage"
confidence: Literal["observed"] = "observed"
class CacheCostScenario(BaseModel):
tokens: CacheTokenBuckets
input_cost: float
class CachePredictionArm(BaseModel):
deployment_id: str
model: str | None = None
cache_state: Literal["warm", "partial", "stale", "unknown", "disabled"] = "unknown"
reason: str | None = None
estimate: CacheCostScenario | None = None
cold: CacheCostScenario | None = None
warm: CacheCostScenario | None = None
evidence: CacheEvidence | None = None
token_count_source: Literal["anthropic_count_tokens"] | None = None
class CachePredictionRequest(BaseModel):
model_config = ConfigDict(extra="forbid")
current_deployment_id: str = Field(min_length=1, max_length=256)
candidate_deployment_id: str = Field(min_length=1, max_length=256)
request: Mapping[str, JsonValue]
class CachePredictionResponse(BaseModel):
stay: CachePredictionArm
switch: CachePredictionArm
switch_delta: float | None
cache_rebuild_penalty: float | None
pricing_basis: Literal["input_before_discounts_and_margins"] = "input_before_discounts_and_margins"
cache_guarantee: Literal[False] = False

View file

@ -0,0 +1,8 @@
"""Failure values raised by `litellm/proxy/auth/auth_checks.py`. Kept free of `litellm` imports so any proxy module can import them without joining the `litellm.proxy` import cycle."""
class UserNotFoundError(ValueError):
"""The user row is provably absent, as opposed to merely unreadable, so a caller that reads a missing row as no user-level limits can key on it without also swallowing a database that would not answer."""
def __init__(self, user_id: str) -> None:
super().__init__(f"User doesn't exist in db. 'user_id'={user_id}. Create user via `/user/new` call.")

View file

@ -645,6 +645,7 @@ class RouterErrors(enum.Enum):
user_defined_ratelimit_error = "Deployment over user-defined ratelimit."
no_deployments_available = "No deployments available for selected model"
all_deployments_in_cooldown = "All deployments for selected model are in cooldown"
no_deployments_with_tag_routing = "Not allowed to access model due to tags configuration"
no_deployments_with_provider_budget_routing = "No deployments available - crossed budget"
no_healthy_deployments = "There are no healthy deployments for this model"
@ -868,6 +869,11 @@ class RouterRateLimitErrorBasic(ValueError):
super().__init__(_message)
class RouterErrorTypes(str, enum.Enum):
rate_limit_error = "rate_limit_error"
all_deployments_in_cooldown = "all_deployments_in_cooldown"
class RouterRateLimitError(ValueError):
def __init__(
self,
@ -875,12 +881,25 @@ class RouterRateLimitError(ValueError):
cooldown_time: float,
enable_pre_call_checks: bool,
cooldown_list: list,
model_ids: Sequence[str] = (),
) -> None:
self.model = model
self.cooldown_time = cooldown_time
self.enable_pre_call_checks = enable_pre_call_checks
self.cooldown_list = cooldown_list
_message = f"{RouterErrors.no_deployments_available.value}, Try again in {cooldown_time} seconds. Passed model={model}. pre-call-checks={enable_pre_call_checks}, cooldown_list={cooldown_list}"
self.all_deployments_in_cooldown = bool(model_ids) and frozenset(model_ids) <= frozenset(cooldown_list)
self.type = (
RouterErrorTypes.all_deployments_in_cooldown.value
if self.all_deployments_in_cooldown
else RouterErrorTypes.rate_limit_error.value
)
_reason: Final = (
f" {RouterErrors.all_deployments_in_cooldown.value}." if self.all_deployments_in_cooldown else ""
)
_message: Final = (
f"{RouterErrors.no_deployments_available.value}, Try again in {cooldown_time} seconds.{_reason} "
f"Passed model={model}. pre-call-checks={enable_pre_call_checks}, cooldown_list={cooldown_list}"
)
super().__init__(_message)
@ -889,6 +908,14 @@ class RouterModelGroupAliasItem(TypedDict):
hidden: bool # if 'True', don't return on `.get_model_list`
class RetryAttemptRecord(TypedDict):
model_group: ReadOnly[str | None]
deployment_id: ReadOnly[str | None]
exception_type: ReadOnly[str]
exception_string: ReadOnly[str]
attempted_retries: ReadOnly[int | None]
VALID_LITELLM_ENVIRONMENTS = [
"development",
"staging",

View file

@ -48,6 +48,7 @@ from litellm._logging import verbose_logger
from litellm._uuid import uuid
from litellm.types.llms.base import (
BaseLiteLLMOpenAIResponseObject,
CachedTokensDetails,
LiteLLMPydanticObjectBase,
)
from litellm.types.mcp import MCPServerCostInfo
@ -252,6 +253,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
cache_creation_input_token_cost_priority: float | None # OpenAI priority service tier pricing
cache_creation_input_token_cost_ultrafast: ReadOnly[float | None] # OpenAI ultrafast service tier pricing
cache_read_input_token_cost: float | None
cache_read_input_audio_token_cost: ReadOnly[float | None]
cache_read_input_token_cost_flex: float | None # OpenAI flex service tier pricing
cache_read_input_token_cost_priority: float | None # OpenAI priority service tier pricing
cache_read_input_token_cost_ultrafast: ReadOnly[float | None] # OpenAI ultrafast service tier pricing
@ -1710,6 +1712,9 @@ class PromptTokensDetailsWrapper(
cache_creation_token_details: CacheCreationTokenDetails | None = None
"""Details of cache creation tokens sent to the model. Used for tracking 5m/1h cache creation tokens for Anthropic prompt caching."""
cached_tokens_details: CachedTokensDetails | None = None
"""Details of cached (cache-hit) tokens sent to the model. OpenAI realtime naming; carries the per-modality cache-read split."""
def __setattr__(self, name: str, value: object) -> None:
super().__setattr__(name, value)
if name == "cache_write_tokens":
@ -1756,6 +1761,8 @@ class PromptTokensDetailsWrapper(
del self.cache_creation_tokens
if self.cache_creation_token_details is None:
del self.cache_creation_token_details
if self.cached_tokens_details is None:
del self.cached_tokens_details
class ServerToolUse(BaseModel):
@ -3754,6 +3761,16 @@ all_litellm_params = (
"model_file_id_mapping",
"litellm_logging_obj",
"litellm_call_id",
"completion_call_id",
"model_alias_map",
"custom_prompt_dict",
"stream_response",
"cost_per_query",
"ssl_verify",
"data_residency",
"async_call",
"aembedding",
"allm_passthrough_route",
"_litellm_strip_stream_usage",
"use_client",
"id",

View file

@ -84,6 +84,7 @@ from litellm.constants import (
from litellm.litellm_core_utils.core_helpers import normalize_drop_params
from litellm.litellm_core_utils.fallback_generalizations import (
match_capability_generalizations,
match_fill_missing_generalizations,
)
from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload
@ -254,6 +255,7 @@ from litellm.types.utils import (
)
_CALL_TYPE_ENUM_MAP: Final[dict] = {ct.value: ct for ct in CallTypes}
_BACKFILL_MODES: Final = frozenset({"chat", "responses"})
# +-----------------------------------------------+
# | |
@ -1260,15 +1262,6 @@ async def _client_async_logging_helper(
async_coroutine=logging_obj.async_success_handler(result=result, start_time=start_time, end_time=end_time)
)
################################################
# Sync Logging Worker
################################################
logging_obj.handle_sync_success_callbacks_for_async_calls(
result=result,
start_time=start_time,
end_time=end_time,
)
def _get_wrapper_num_retries(kwargs: dict[str, Any], exception: Exception) -> tuple[int | None, dict[str, Any]]:
"""
@ -1500,6 +1493,8 @@ def post_call_processing(
def client(original_function):
from litellm.litellm_core_utils.core_helpers import max_retries_per_request_hit
Rules: Final = litellm_utils.Rules
rules_obj: Final = Rules()
@ -1510,12 +1505,8 @@ def client(original_function):
call_type = original_function.__name__
if _is_async_request(kwargs):
# [OPTIONAL] CHECK MAX RETRIES / REQUEST
if litellm.num_retries_per_request is not None:
# check if previous_models passed in as ['litellm_params']['metadata]['previous_models']
previous_models = (kwargs.get("metadata") or {}).get("previous_models", None)
if previous_models is not None:
if litellm.num_retries_per_request <= len(previous_models):
raise Exception("Max retries per request hit!")
if max_retries_per_request_hit(kwargs, litellm.num_retries_per_request):
raise Exception("Max retries per request hit!")
# MODEL CALL
result = original_function(*args, **kwargs)
@ -1574,12 +1565,8 @@ def client(original_function):
)
# [OPTIONAL] CHECK MAX RETRIES / REQUEST
if litellm.num_retries_per_request is not None:
# check if previous_models passed in as ['litellm_params']['metadata]['previous_models']
previous_models = (kwargs.get("metadata") or {}).get("previous_models", None)
if previous_models is not None:
if litellm.num_retries_per_request <= len(previous_models):
raise Exception("Max retries per request hit!")
if max_retries_per_request_hit(kwargs, litellm.num_retries_per_request):
raise Exception("Max retries per request hit!")
# [OPTIONAL] CHECK CACHE
print_verbose(
@ -5819,6 +5806,14 @@ def _get_model_info_helper(
):
_model_info = None
if _model_info is not None and key is not None and _model_info.get("mode", "chat") in _BACKFILL_MODES:
fill_missing: Final = match_fill_missing_generalizations(key, _model_info.get("litellm_provider", ""))
if fill_missing is not None:
_model_info = {
**{k: v for k, v in fill_missing.items() if k not in _model_info},
**_model_info,
}
if _model_info is None:
generalization: Final = _get_model_info_from_generalization(
model=model,
@ -5882,6 +5877,7 @@ def _get_model_info_helper(
"cache_creation_input_token_cost_ultrafast", None
),
cache_read_input_token_cost=_model_info.get("cache_read_input_token_cost", None),
cache_read_input_audio_token_cost=_model_info.get("cache_read_input_audio_token_cost", None),
prompt_cache_min_tokens=_model_info.get("prompt_cache_min_tokens", None),
cache_read_input_token_cost_above_200k_tokens=_model_info.get(
"cache_read_input_token_cost_above_200k_tokens", None
@ -6264,7 +6260,7 @@ def function_to_dict(input_function) -> dict:
"enum": param_enum,
}
parameters[param_name] = dict([(k, v) for k, v in param_dict.items() if isinstance(v, str)])
parameters[param_name] = {k: v for k, v in param_dict.items() if isinstance(v, str)}
# Check if the parameter has no default value (i.e., it's required)
if param.default == param.empty:

File diff suppressed because it is too large Load diff

View file

@ -5,5 +5,5 @@ reason = "diskcache has no fixed release published; remove this entry once one e
[[IgnoredVulns]]
id = "GHSA-h7x2-h6g9-p789"
ignoreUntil = 2026-09-14
reason = "mlflow has no fixed release published; remove this entry once one exists"
ignoreUntil = 2026-10-14
reason = "mlflow has no fixed release published (3.16.0, 2026-09-04, and master still store gateway secret api_base unvalidated); remove this entry once one exists"

View file

@ -4,11 +4,12 @@ lint.ignore = ["F405", "E402", "F403"]
# That gives editors and `ruff check --fix` the diagnostic, which the gate script cannot.
lint.extend-select = [
"T20", "PGH004", "RUF008", "RUF009", "RUF100",
"B033", "FURB136", "FURB168", "FURB188", "I001", "PERF402", "PIE790", "PIE800", "PLC0208",
"PLR0402", "PLR1711", "PLR1730", "PLR2044", "PLW0133", "PYI030", "PYI041", "PYI064", "RET501",
"RUF010", "RUF022", "RUF023", "RUF051", "S113", "SIM114", "SIM118", "TC005", "UP006", "UP007",
"UP008",
"UP012", "UP018", "UP024", "UP032", "UP034", "UP035", "UP037", "UP045",
"B004", "B018", "B021", "B033", "FURB136", "FURB168", "FURB188", "I001", "PERF402", "PIE790",
"PIE800", "PLC0208", "PLR0124", "PLR0402", "PLR0206", "PLR1704", "PLR1711", "PLR1730", "PLR2044",
"PLW0133", "PYI030", "PYI041", "PYI064", "RET501", "RUF010", "RUF022", "RUF023", "RUF051", "S113",
"SIM114", "SIM118", "SIM201", "SIM211", "SIM222", "TC005", "UP006", "UP007", "UP008",
"UP012", "UP018", "UP024", "UP032", "UP034", "UP035", "UP036", "UP037", "UP045",
"C404", "C419",
]
# RUF100 (unused-noqa) only knows the rules enabled in THIS config, so it would strip
# `# noqa` directives that protect rules enforced elsewhere. List those codes as external

View file

@ -0,0 +1,19 @@
# Integration contracts
These tests exercise a running gateway, PostgreSQL and Redis with an owned local upstream. CircleCI owns this suite. Tests are grouped by behavior, with no automatic test retries or fallback to paid provider calls
Use `tests/integration/run.py management`, `accounting` or `providers` to run a selected group. Set `INTEGRATION_PROXY_URL`, `INTEGRATION_UPSTREAM_URL`, `INTEGRATION_MASTER_KEY` and `DATABASE_URL` to an isolated test deployment. The runner selects the new domain directories explicitly; the legacy OCI and sandbox selections remain separate
Management also requires `INTEGRATION_PEER_URL`, `REDIS_HOST` and `REDIS_PORT`. CircleCI starts two directly addressed proxy processes sharing only that job's stores. The test-only CLI wrapper supplies enterprise route entitlement, following the existing behavior suite's convention. It does not qualify license validation; run it with one worker and no reload
The generated lifecycle models use 20 examples, eight steps, generation and shrinking, with isolated resources per example. HTTP operation caps include generation and shrinking and exempt cleanup. Local qualification defaults to seed 4106601; CircleCI derives its exploration seed from the checked-out revision. Use `--seed` to reproduce a run. Actual installed Hypothesis version, settings and seed are written beside the execution manifest
Reuse the existing canned provider handlers through `_support/upstream.py`. It rejects internal request fields and exposes actual received requests for independent assertions. Register every created resource for cleanup immediately, keep expected values independent of production calculations, and assert readback plus the runtime effect of a change
The CircleCI workflow starts its own database and Redis, restricts test-phase egress to its owned services and writes JUnit plus an executed-node manifest. Missing setup, skipped tests, failed cleanup or a selected test without a passed call fail qualification. Existing GitHub Actions jobs do not own these tests
Define integration contract IDs and their canonical test nodes in `contracts.json`. Every node must declare the same IDs with `covers`. The runner checks exact collected and passed selections against that mapping. These IDs belong to this CircleCI suite and must not be added to the separate E2E coverage registry. A manifest declaration alone does not mean a test passed
Provider sentinels currently use the controlled server, not live recordings. The provider shard also runs the existing strict replay controls for changed requests, exhausted interactions, leftover interactions and no provider connection. Future recorded scenarios must use that replay-only implementation; missing recordings cannot fall back to a real provider. The observation endpoint is destructive and the current selection runs serially against one owned upstream
Fixtures must contain synthetic data only. Keep private incident records and source documents out of code, fixtures, logs and PR descriptions

View file

View file

View file

@ -0,0 +1,170 @@
from __future__ import annotations
import os
import time
import uuid
from hashlib import sha256
from collections.abc import Callable, Iterator, Mapping
from contextlib import ExitStack, contextmanager
from dataclasses import dataclass
from typing import Final, TypeVar
import httpx
from pydantic import JsonValue, TypeAdapter
from integration._support.database import read_rows
JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
T = TypeVar("T")
def object_value(value: JsonValue) -> dict[str, JsonValue]:
return JSON_OBJECT.validate_python(value)
def string_value(value: JsonValue) -> str:
assert isinstance(value, str), f"Expected a string, received {type(value).__name__}"
return value
def eventually(read: Callable[[], T], satisfied: Callable[[T], bool], seconds: float = 10) -> T:
deadline: Final = time.monotonic() + seconds
while True:
observed: Final = read()
if satisfied(observed):
return observed
assert time.monotonic() < deadline, f"State did not converge: {observed!r}"
time.sleep(0.1)
@dataclass(frozen=True, slots=True)
class Gateway:
client: httpx.Client
key: str
upstream_url: str
def request(
self,
method: str,
path: str,
body: Mapping[str, JsonValue] | None = None,
*,
key: str | None = None,
params: Mapping[str, str] | None = None,
) -> httpx.Response:
return self.client.request(
method,
path,
json=body,
params=params,
headers={"Authorization": f"Bearer {self.key if key is None else key}"},
)
def post(self, path: str, body: Mapping[str, JsonValue], *, key: str | None = None) -> dict[str, JsonValue]:
response: Final = self.request("POST", path, body, key=key)
assert response.status_code == 200, f"POST {path}: {response.status_code} {response.text}"
return JSON_OBJECT.validate_json(response.content)
def get(self, path: str, params: Mapping[str, str] | None = None) -> dict[str, JsonValue]:
response: Final = self.request("GET", path, params=params)
assert response.status_code == 200, f"GET {path}: {response.status_code} {response.text}"
return JSON_OBJECT.validate_json(response.content)
def chat(self, model: str, *, key: str | None = None, text: str = "integration control") -> dict[str, JsonValue]:
return self.post(
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": text}]},
key=key,
)
@contextmanager
def scenario(self) -> Iterator[Scenario]:
with ExitStack() as cleanups:
yield Scenario(self, cleanups)
@dataclass(frozen=True, slots=True)
class Scenario:
gateway: Gateway
cleanups: ExitStack
def key(self, **fields: JsonValue) -> str:
created: Final = self.gateway.post("/key/generate", fields)
token: Final = string_value(created["key"])
self.cleanups.callback(self.delete_key, token)
return token
def team(self, **fields: JsonValue) -> str:
created: Final = self.gateway.post("/team/new", {"team_alias": f"integration-{uuid.uuid4().hex}", **fields})
identity: Final = string_value(created["team_id"])
self.cleanups.callback(self.delete_team, identity)
return identity
def delete_team(self, identity: str) -> None:
self.gateway.post("/team/delete", {"team_ids": [identity]})
assert read_rows('SELECT team_id FROM "LiteLLM_TeamTable" WHERE team_id = %s', (identity,)) == []
def project(self, team_id: str, **fields: JsonValue) -> str:
created: Final = self.gateway.post(
"/project/new", {"team_id": team_id, "project_alias": f"integration-{uuid.uuid4().hex}", **fields}
)
identity: Final = string_value(created["project_id"])
self.cleanups.callback(self.delete_project, identity)
return identity
def delete_project(self, identity: str) -> None:
response: Final = self.gateway.request("DELETE", "/project/delete", {"project_ids": [identity]})
assert response.status_code == 200, response.text
assert read_rows('SELECT project_id FROM "LiteLLM_ProjectTable" WHERE project_id = %s', (identity,)) == []
def user(self, **fields: JsonValue) -> str:
created: Final = self.gateway.post(
"/user/new", {"user_id": f"integration-{uuid.uuid4().hex}", "auto_create_key": False, **fields}
)
identity: Final = string_value(created["user_id"])
self.cleanups.callback(self.delete_user, identity)
return identity
def delete_user(self, identity: str) -> None:
response: Final = self.gateway.request("POST", "/user/delete", {"user_ids": [identity]})
assert response.status_code == 200 and response.json() == 1, response.text
assert read_rows('SELECT user_id FROM "LiteLLM_UserTable" WHERE user_id = %s', (identity,)) == []
def delete_key(self, token: str) -> None:
self.gateway.post("/key/delete", {"keys": [token]})
response: Final = self.gateway.request("GET", "/key/info", params={"key": sha256(token.encode()).hexdigest()})
assert response.status_code == 404, f"Deleted key remains readable: {response.status_code}"
def delete_model(self, identity: str) -> None:
self.gateway.post("/model/delete", {"id": identity})
entries: Final = self.gateway.get("/model/info")["data"]
assert isinstance(entries, list)
assert all(object_value(object_value(entry)["model_info"])["id"] != identity for entry in entries)
assert read_rows('SELECT model_id FROM "LiteLLM_ProxyModelTable" WHERE model_id = %s', (identity,)) == []
def model(self, **parameters: JsonValue) -> str:
name: Final = f"integration-{uuid.uuid4().hex}"
created: Final = self.gateway.post(
"/model/new",
{
"model_name": name,
"litellm_params": {
"model": "openai/gpt-4o-mini",
"api_key": "integration-provider-key",
"api_base": f"{self.gateway.upstream_url}/v1",
**parameters,
},
"model_info": {},
},
)
identity: Final = string_value(object_value(created["model_info"])["id"])
self.cleanups.callback(self.delete_model, identity)
return name
@contextmanager
def gateway_from_environment() -> Iterator[Gateway]:
url: Final = os.environ["INTEGRATION_PROXY_URL"]
upstream: Final = os.environ["INTEGRATION_UPSTREAM_URL"]
with httpx.Client(base_url=url, timeout=15, trust_env=False) as client:
yield Gateway(client, os.environ["INTEGRATION_MASTER_KEY"], upstream)

View file

@ -0,0 +1,14 @@
import os
from typing import Final
import psycopg
from psycopg.rows import dict_row
from pydantic import JsonValue, TypeAdapter
ROWS: Final = TypeAdapter(list[dict[str, JsonValue]])
def read_rows(query: str, parameters: tuple[str, ...]) -> list[dict[str, JsonValue]]:
with psycopg.connect(os.environ["DATABASE_URL"], row_factory=dict_row) as connection:
connection.execute("SET TRANSACTION READ ONLY")
return ROWS.validate_python(connection.execute(query, parameters).fetchall())

View file

@ -0,0 +1,52 @@
from typing import Final
from dataclasses import dataclass
from collections.abc import Iterator, Sequence
from contextlib import contextmanager
import httpx
from hypothesis import Phase, settings
from integration._support.client import Gateway
LIFECYCLE_SETTINGS: Final = settings(
max_examples=20,
stateful_step_count=8,
deadline=None,
database=None,
phases=(Phase.generate, Phase.shrink),
print_blob=True,
)
@dataclass(slots=True)
class RequestBudget:
limit: int
requests: int = 0
cleaning: bool = False
def observe(self, _request: httpx.Request) -> None:
if self.cleaning:
return
self.requests += 1
assert self.requests <= self.limit, f"Generated HTTP operation budget exceeded: {self.limit}"
@contextmanager
def cleanup(self) -> Iterator[None]:
self.cleaning = True
try:
yield
finally:
self.cleaning = False
@contextmanager
def bounded_http_requests(gateways: Sequence[Gateway], limit: int) -> Iterator[RequestBudget]:
budget: Final = RequestBudget(limit)
for gateway in gateways:
gateway.client.event_hooks["request"].append(budget.observe)
try:
yield budget
finally:
for gateway in gateways:
gateway.client.event_hooks["request"].remove(budget.observe)
print(f"Generated HTTP operations: {budget.requests}/{budget.limit}; cleanup excluded")

View file

@ -0,0 +1,31 @@
import json
from pathlib import Path
from typing import Final
from pydantic import TypeAdapter
MAPPING: Final = TypeAdapter(dict[str, tuple[str, ...]])
OWNED_DIRECTORIES: Final = frozenset(
{
"management",
"authorization",
"database",
"pricing",
"spend",
"routing",
"providers",
"streaming",
"configuration",
"mcp",
"observability",
"compatibility",
}
)
def contracts() -> dict[str, tuple[str, ...]]:
document: Final = json.loads((Path(__file__).resolve().parents[1] / "contracts.json").read_bytes())
result: Final = MAPPING.validate_python(document["tests"])
if not result or any(not values or any(not value.strip() for value in values) for values in result.values()):
raise ValueError("Integration manifest must contain nodes with contract IDs")
return result

View file

@ -0,0 +1,16 @@
"""Run the normal single-process CLI with the existing behavior-suite test entitlement."""
from unittest.mock import patch
from litellm import run_server
def main() -> None:
with patch( # test-quality-ok: route entitlement only; license validation is outside these HTTP/DB contracts
"litellm.proxy.auth.litellm_license.LicenseCheck.is_premium", return_value=True
):
run_server()
if __name__ == "__main__":
main()

View file

@ -0,0 +1,122 @@
from __future__ import annotations
import argparse
from dataclasses import dataclass, field
from collections import deque
from queue import SimpleQueue
from typing import Final
import uvicorn
from pydantic import JsonValue, TypeAdapter
from starlette.applications import Starlette
from starlette.requests import Request
from starlette.responses import JSONResponse, Response
from starlette.routing import Route
from _fake_openai_endpoint_server import chat_completions, completions, embeddings, health, moderations
JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
INTERNAL_FIELDS: Final = frozenset(
{
"litellm_params",
"litellm_logging_obj",
"litellm_call_id",
"litellm_metadata",
"proxy_server_request",
"rpm",
"tpm",
"timeout",
"stream_chunk_size",
}
)
@dataclass(frozen=True, slots=True)
class Observation:
path: str
authorization: str
body: dict[str, JsonValue]
@dataclass(frozen=True, slots=True)
class Provider:
observations: SimpleQueue[Observation] = field(default_factory=SimpleQueue)
scripts: dict[str, deque[int]] = field(default_factory=dict)
async def chat(self, request: Request) -> Response:
body: Final = JSON_OBJECT.validate_json(await request.body())
self.observations.put(Observation(request.url.path, request.headers.get("authorization", ""), body))
leaked: Final = tuple(sorted(INTERNAL_FIELDS.intersection(body)))
if leaked:
return JSONResponse({"error": {"message": f"Unexpected provider fields: {leaked}"}}, status_code=400)
messages: Final = body.get("messages")
if not isinstance(body.get("model"), str) or not isinstance(messages, list) or not messages:
return JSONResponse({"error": {"message": "model and nonempty messages are required"}}, status_code=400)
if any(
not isinstance(message, dict)
or message.get("role") not in {"system", "developer", "user", "assistant", "tool"}
or "content" not in message
for message in messages
):
return JSONResponse({"error": {"message": "Invalid selected message contract"}}, status_code=400)
script: Final = self.scripts.get(str(body["model"]))
if script is not None:
if not script:
return JSONResponse({"error": {"message": "Script exhausted", "type": "api_error"}}, status_code=500)
status: Final = script.popleft()
if status != 200:
return JSONResponse(
{"error": {"message": "Controlled provider failure", "type": "api_error", "code": str(status)}},
status_code=status,
)
return await chat_completions(request)
async def script(self, request: Request) -> Response:
name: Final = request.path_params["model"]
if request.method in {"DELETE", "GET"} and name not in self.scripts:
return JSONResponse({"error": "Script not found"}, status_code=404)
if request.method == "GET":
return JSONResponse({"remaining": list(self.scripts[name])})
if request.method == "DELETE":
remaining: Final = self.scripts.pop(name)
return JSONResponse({"remaining": list(remaining)})
body: Final = JSON_OBJECT.validate_json(await request.body())
statuses: Final = body.get("statuses")
if not isinstance(statuses, list) or not statuses or any(type(value) is not int for value in statuses):
return JSONResponse({"error": "A nonempty list of HTTP status codes is required"}, status_code=400)
self.scripts[name] = deque(int(str(value)) for value in statuses)
return JSONResponse({"configured": len(statuses)})
async def observed(self, _request: Request) -> Response:
values: Final = tuple(self.observations.get() for _ in range(self.observations.qsize()))
return JSONResponse(
{
"requests": [
{"path": value.path, "authorization": value.authorization, "body": value.body} for value in values
]
}
)
def app(self) -> Starlette:
return Starlette(
routes=[
Route("/health", health),
Route("/__observations", self.observed),
Route("/__scripts/{model}", self.script, methods=["POST", "DELETE", "GET"]),
Route("/v1/chat/completions", self.chat, methods=["POST"]),
Route("/v1/completions", completions, methods=["POST"]),
Route("/v1/embeddings", embeddings, methods=["POST"]),
Route("/v1/moderations", moderations, methods=["POST"]),
]
)
def main() -> None:
parser: Final = argparse.ArgumentParser()
parser.add_argument("--port", type=int, default=8190)
arguments: Final = parser.parse_args()
uvicorn.run(Provider().app(), host="127.0.0.1", port=arguments.port, access_log=False)
if __name__ == "__main__":
main()

View file

@ -0,0 +1,197 @@
from contextlib import ExitStack
from hashlib import sha256
from typing import Final
import os
import psycopg
import pytest
from hypothesis import strategies as st
from hypothesis.stateful import RuleBasedStateMachine, invariant, rule, run_state_machine_as_test
from integration._support.client import Gateway, eventually, object_value
from integration._support.database import read_rows
from integration._support.generation import LIFECYCLE_SETTINGS, bounded_http_requests
def assert_serving(gateway: Gateway, model: str, key: str, status: int, error_type: str = "auth_error") -> None:
response: Final = eventually(
lambda: gateway.request(
"POST", "/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": "warmed policy control"}]}, key=key,
),
lambda value: value.status_code == status,
seconds=3,
)
if status == 200:
assert response.json()["usage"]["total_tokens"] == 40
assert response.json()["choices"][0]["message"]["content"] == (
"Hello! This is a mock response from the fake OpenAI endpoint."
)
else:
assert response.json()["error"]["type"] == error_type
@pytest.mark.covers("mgmt.key.update.two_workers_enforce_warmed_policy")
def test_generated_policy_changes_reach_both_warmed_workers(gateway: Gateway, peer: Gateway) -> None:
class Policies(RuleBasedStateMachine):
def __init__(self) -> None:
super().__init__()
self.resources = ExitStack()
try:
scenario = self.resources.enter_context(gateway.scenario())
self.models = (scenario.model(), scenario.model())
self.allowed = 0
self.blocked = False
self.key = scenario.key(models=[self.models[0]], blocked=False)
self.control = scenario.key(models=list(self.models))
for worker in (gateway, peer):
assert_serving(worker, self.models[0], self.key, 200)
assert_serving(worker, self.models[1], self.control, 200)
except BaseException:
with budget.cleanup():
self.resources.close()
raise
@rule(index=st.integers(min_value=0, max_value=1))
def model_grant(self, index: int) -> None:
gateway.post("/key/update", {"key": self.key, "models": [self.models[index]]})
self.allowed = index
@rule(blocked=st.booleans())
def block(self, blocked: bool) -> None:
gateway.post("/key/update", {"key": self.key, "blocked": blocked})
self.blocked = blocked
@invariant()
def both_workers_enforce_policy(self) -> None:
rows: Final = read_rows(
'SELECT models, blocked FROM "LiteLLM_VerificationToken" WHERE token = %s',
(sha256(self.key.encode()).hexdigest(),),
)
assert rows == [{"models": [self.models[self.allowed]], "blocked": self.blocked}]
for worker in (gateway, peer):
for index, model in enumerate(self.models):
status: Final = 401 if self.blocked else 200 if index == self.allowed else 403
kind: Final = "auth_error" if self.blocked else "key_model_access_denied"
assert_serving(worker, model, self.key, status, kind)
assert_serving(worker, self.models[1], self.control, 200)
def teardown(self) -> None:
with budget.cleanup():
self.resources.close()
with bounded_http_requests((gateway, peer), limit=3000) as budget:
run_state_machine_as_test(Policies, settings=LIFECYCLE_SETTINGS)
@pytest.mark.covers("mgmt.user.scim.deactivation_includes_nullable_blocked_keys")
def test_scim_deactivation_blocks_null_and_false_keys_but_preserves_other_owners(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
model: Final = scenario.model()
user: Final = scenario.user(user_role="internal_user")
other: Final = scenario.user(user_role="internal_user")
null_key: Final = scenario.key(user_id=user, models=[model])
false_key: Final = scenario.key(user_id=user, models=[model], blocked=False)
manual: Final = scenario.key(user_id=user, models=[model], blocked=True)
control: Final = scenario.key(user_id=other, models=[model])
team: Final = scenario.team(models=[model])
service: Final = gateway.post("/key/service-account/generate", {"team_id": team, "models": [model]})
service_key: Final = service["key"]
assert isinstance(service_key, str)
scenario.cleanups.callback(scenario.delete_key, service_key)
with psycopg.connect(os.environ["DATABASE_URL"]) as connection:
connection.execute(
'UPDATE "LiteLLM_VerificationToken" SET blocked = NULL WHERE token = %s',
(sha256(null_key.encode()).hexdigest(),),
)
assert read_rows(
'SELECT user_id, blocked FROM "LiteLLM_VerificationToken" WHERE token = %s',
(sha256(null_key.encode()).hexdigest(),),
) == [{"user_id": user, "blocked": None}]
assert read_rows(
'SELECT user_id FROM "LiteLLM_VerificationToken" WHERE token = %s',
(sha256(service_key.encode()).hexdigest(),),
) == [{"user_id": None}]
for token in (null_key, false_key, control, service_key):
assert_serving(gateway, model, token, 200)
for active in (False, True):
response: Final = gateway.request(
"PATCH", f"/scim/v2/Users/{user}",
{"schemas": ["urn:ietf:params:scim:api:messages:2.0:PatchOp"],
"Operations": [{"op": "replace", "path": "active", "value": active}]},
)
assert response.status_code == 200, response.text
for token in (null_key, false_key):
rows: Final = read_rows(
'SELECT blocked, metadata FROM "LiteLLM_VerificationToken" WHERE token = %s',
(sha256(token.encode()).hexdigest(),),
)
assert rows[0]["blocked"] is not active
assert object_value(rows[0]["metadata"]).get("scim_blocked") is (None if active else True)
assert_serving(gateway, model, token, 200 if active else 401)
assert_serving(gateway, model, manual, 401)
for token in (control, service_key):
assert_serving(gateway, model, token, 200)
@pytest.mark.covers("mgmt.team.member_update.demoted_role_cannot_write")
def test_warmed_team_role_demotion_prevents_later_management_writes(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
model: Final = scenario.model()
user: Final = scenario.user(user_role="internal_user")
team: Final = scenario.team(models=[model], members_with_roles=[{"user_id": user, "role": "admin"}])
control_team: Final = scenario.team(models=[model])
caller: Final = scenario.key(
user_id=user, team_id=team, models=[model], allowed_routes=["/team/update", "/v1/chat/completions"]
)
gateway.chat(model, key=caller)
changed: Final = gateway.request("POST", "/team/update", {"team_id": team, "team_alias": "permitted"}, key=caller)
assert changed.status_code == 200, changed.text
unrelated_before: Final = read_rows(
'SELECT team_alias FROM "LiteLLM_TeamTable" WHERE team_id = %s', (control_team,)
)
unrelated: Final = gateway.request(
"POST", "/team/update", {"team_id": control_team, "team_alias": "must-not-persist"}, key=caller
)
assert unrelated.status_code == 403, unrelated.text
assert read_rows(
'SELECT team_alias FROM "LiteLLM_TeamTable" WHERE team_id = %s', (control_team,)
) == unrelated_before
gateway.post("/team/member_update", {"team_id": team, "user_id": user, "role": "user"})
for target in (team, control_team):
before: Final = read_rows('SELECT team_alias FROM "LiteLLM_TeamTable" WHERE team_id = %s', (target,))
denied: Final = gateway.request(
"POST", "/team/update", {"team_id": target, "team_alias": "must-not-persist"}, key=caller
)
assert denied.status_code == 403, denied.text
assert read_rows('SELECT team_alias FROM "LiteLLM_TeamTable" WHERE team_id = %s', (target,)) == before
roster: Final = read_rows('SELECT members_with_roles FROM "LiteLLM_TeamTable" WHERE team_id = %s', (team,))
members: Final = roster[0]["members_with_roles"]
assert isinstance(members, list)
assert next(object_value(member)["role"] for member in members if object_value(member)["user_id"] == user) == "user"
assert_serving(gateway, model, caller, 200)
@pytest.mark.covers("mgmt.key.update.expiry_changes_reach_warmed_workers")
def test_expiry_and_explicit_clear_reach_both_warmed_workers(gateway: Gateway, peer: Gateway) -> None:
with gateway.scenario() as scenario:
model: Final = scenario.model()
key: Final = scenario.key(models=[model], duration="1h")
control: Final = scenario.key(models=[model])
for worker in (gateway, peer):
assert_serving(worker, model, key, 200)
gateway.post("/key/update", {"key": key, "duration": "0s"})
assert read_rows(
"SELECT expires <= timezone('UTC', now()) AS expired FROM \"LiteLLM_VerificationToken\" WHERE token = %s",
(sha256(key.encode()).hexdigest(),),
) == [{"expired": True}]
for worker in (gateway, peer):
assert_serving(worker, model, key, 401, "expired_key")
assert_serving(worker, model, control, 200)
gateway.post("/key/update", {"key": key, "duration": None})
assert read_rows(
'SELECT expires IS NULL AS cleared FROM "LiteLLM_VerificationToken" WHERE token = %s',
(sha256(key.encode()).hexdigest(),),
) == [{"cleared": True}]
for worker in (gateway, peer):
assert_serving(worker, model, key, 200)

View file

@ -0,0 +1,123 @@
import uuid
from typing import Final
import httpx
import pytest
from integration._support.client import Gateway, object_value, string_value
from integration._support.database import read_rows
def model_identity(gateway: Gateway, alias: str) -> str:
entries: Final = gateway.get("/model/info")["data"]
assert isinstance(entries, list)
entry: Final = next(object_value(value) for value in entries if object_value(value)["model_name"] == alias)
return string_value(object_value(entry["model_info"])["id"])
@pytest.mark.covers("mgmt.model.block.changes_serving_and_preserves_control")
def test_model_block_changes_actual_route_and_leaves_other_route_working(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
model: Final = scenario.model()
other: Final = scenario.model()
identity: Final = model_identity(gateway, model)
gateway.chat(model)
gateway.chat(other)
gateway.post("/model/block", {"model_id": identity})
assert read_rows(
'SELECT blocked FROM "LiteLLM_ProxyModelTable" WHERE model_id = %s', (identity,)
) == [{"blocked": True}]
response: Final = gateway.request(
"POST", "/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": "blocked deployment"}]},
)
assert response.status_code == 403, response.text
assert response.json()["error"]["type"] == "permission_error"
assert response.json()["error"]["message"] == "litellm.PermissionDeniedError: Model is blocked"
assert object_value(gateway.chat(other)["usage"])["total_tokens"] == 40
gateway.post("/model/unblock", {"model_id": identity})
assert read_rows(
'SELECT blocked FROM "LiteLLM_ProxyModelTable" WHERE model_id = %s', (identity,)
) == [{"blocked": False}]
assert object_value(gateway.chat(model)["usage"])["total_tokens"] == 40
@pytest.mark.covers("mgmt.router_settings.update.changes_observed_attempt_count")
def test_saved_retry_setting_controls_real_attempts_and_restores(gateway: Gateway) -> None:
with httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, gateway.scenario() as scenario:
original: Final = object_value(gateway.get("/router/settings")["current_values"])["num_retries"]
provider_model: Final = f"retry-{uuid.uuid4().hex}"
model: Final = scenario.model(model=f"openai/{provider_model}", input_cost_per_token=0, output_cost_per_token=0)
def remove_script() -> None:
response: Final = upstream.delete(f"/__scripts/{provider_model}")
assert response.status_code in (200, 404), response.text
assert upstream.get(f"/__scripts/{provider_model}").status_code == 404
scenario.cleanups.callback(remove_script)
try:
for generation, retries in enumerate((0, 1, original)):
gateway.post("/config/update", {"router_settings": {"num_retries": retries}})
assert object_value(gateway.get("/router/settings")["current_values"])["num_retries"] == retries
configured: Final = upstream.post(f"/__scripts/{provider_model}", json={"statuses": [500, 200]})
assert configured.status_code == 200, configured.text
upstream.get("/__observations").raise_for_status()
response: Final = gateway.request(
"POST", "/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": f"{provider_model} attempt {generation}"}]},
)
observed: Final = upstream.get("/__observations")
observed.raise_for_status()
requests: Final = observed.json()["requests"]
assert len(requests) == (1 if retries == 0 else 2), (response.status_code, response.text, requests)
assert all(value["body"]["model"] == provider_model for value in requests)
assert response.status_code == (500 if retries == 0 else 200), response.text
if retries != 0:
assert response.json()["usage"]["total_tokens"] == 40
remaining: Final = upstream.delete(f"/__scripts/{provider_model}")
assert remaining.status_code == 200, remaining.text
assert remaining.json()["remaining"] == ([200] if retries == 0 else [])
finally:
gateway.post("/config/update", {"router_settings": {"num_retries": original}})
assert object_value(gateway.get("/router/settings")["current_values"])["num_retries"] == original
@pytest.mark.covers("mgmt.credential.update.saved_value_reaches_wire")
def test_credential_value_update_and_model_reload_reach_provider(gateway: Gateway) -> None:
with gateway.scenario() as scenario, httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream:
name: Final = f"credential-{uuid.uuid4().hex}"
gateway.post("/credentials", {
"credential_name": name, "credential_values": {"api_key": "synthetic-credential-first"}, "credential_info": {}
})
def remove_credential() -> None:
response: Final = gateway.request("DELETE", f"/credentials/{name}")
assert response.status_code == 200, response.text
assert read_rows(
'SELECT credential_name FROM "LiteLLM_CredentialsTable" WHERE credential_name = %s', (name,)
) == []
scenario.cleanups.callback(remove_credential)
model: Final = scenario.model(api_key=None, litellm_credential_name=name)
identity: Final = model_identity(gateway, model)
for value in ("synthetic-credential-first", "synthetic-credential-second"):
patched: Final = gateway.request("PATCH", f"/credentials/{name}", {
"credential_name": name, "credential_values": {"api_key": value}, "credential_info": {}
})
assert patched.status_code == 200, patched.text
rows: Final = read_rows(
'SELECT credential_values FROM "LiteLLM_CredentialsTable" WHERE credential_name = %s', (name,)
)
assert len(rows) == 1
stored: Final = object_value(rows[0]["credential_values"])
assert isinstance(stored["api_key"], str) and stored["api_key"] != value
for reload in (False, True):
if reload:
response: Final = gateway.request("PATCH", f"/model/{identity}/update", {"model_info": {"description": value}})
assert response.status_code == 200, response.text
upstream.get("/__observations").raise_for_status()
assert object_value(gateway.chat(model, text=f"{name} {value} reload={reload}")["usage"])["total_tokens"] == 40
observed: Final = upstream.get("/__observations")
observed.raise_for_status()
assert len(observed.json()["requests"]) == 1, (value, reload, observed.text)
assert observed.json()["requests"][0]["authorization"] == f"Bearer {value}"

View file

@ -0,0 +1,103 @@
from __future__ import annotations
import json
import os
from importlib.metadata import version
from collections.abc import Generator, Iterator
from pathlib import Path
from typing import Final
import pytest
import httpx
from redis import Redis
from integration._support.client import Gateway, eventually, gateway_from_environment
from integration._support.manifest import OWNED_DIRECTORIES, contracts
from integration._support.generation import LIFECYCLE_SETTINGS
COLLECTED: Final = pytest.StashKey[tuple[str, ...]]()
REPORTS: Final = pytest.StashKey[list[pytest.TestReport]]()
def pytest_configure(config: pytest.Config) -> None:
config.addinivalue_line("markers", "integration: owned real-service integration contracts")
config.addinivalue_line("markers", "covers(*ids): independently asserted behavior contracts")
config.stash[REPORTS] = []
def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item]) -> None:
manifest: Final = contracts()
root: Final = Path(__file__).parent
owned: Final = tuple(
item
for item in items
if item.path.is_relative_to(root) and item.path.relative_to(root).parts[0] in OWNED_DIRECTORIES
)
if owned and os.environ.get("GITHUB_ACTIONS") == "true":
raise pytest.UsageError("Integration contracts are owned by CircleCI")
for item in owned:
if item.nodeid not in manifest:
raise pytest.UsageError(f"Integration node missing from manifest: {item.nodeid}")
item.add_marker(pytest.mark.integration)
declared: Final = tuple(value for mark in item.iter_markers("covers") for value in mark.args)
if set(declared) != set(manifest[item.nodeid]):
raise pytest.UsageError(f"Contract mapping differs for {item.nodeid}")
config.stash[COLLECTED] = tuple(item.nodeid for item in owned)
@pytest.hookimpl(wrapper=True)
def pytest_runtest_makereport(
item: pytest.Item, call: pytest.CallInfo[None]
) -> Generator[None, pytest.TestReport, pytest.TestReport]:
report: Final = yield
item.config.stash[REPORTS].append(report)
return report
def pytest_sessionfinish(session: pytest.Session, exitstatus: int) -> None:
destination: Final = os.environ.get("INTEGRATION_RESULTS_DIR")
if destination is None:
return
collected: Final = session.config.stash.get(COLLECTED, ())
reports: Final = tuple(report for report in session.config.stash[REPORTS] if report.nodeid in collected)
passed: Final = tuple(report.nodeid for report in reports if report.when == "call" and report.passed)
complete: Final = (
exitstatus == 0
and bool(collected)
and sorted(collected) == sorted(passed)
and all(report.passed for report in reports)
)
output: Final = Path(destination)
output.mkdir(parents=True, exist_ok=True)
(output / "execution.json").write_text(
json.dumps({
"collected": collected, "passed": passed, "complete": complete, "exitstatus": exitstatus,
"hypothesis_version": version("hypothesis"),
"hypothesis_seed": session.config.getoption("hypothesis_seed"),
"generation": {
"max_examples": LIFECYCLE_SETTINGS.max_examples,
"stateful_step_count": LIFECYCLE_SETTINGS.stateful_step_count,
"database": str(LIFECYCLE_SETTINGS.database),
"phases": [phase.name for phase in LIFECYCLE_SETTINGS.phases],
},
}, indent=2)
+ "\n"
)
if not complete and exitstatus == 0:
session.exitstatus = pytest.ExitCode.TESTS_FAILED
@pytest.fixture
def gateway() -> Iterator[Gateway]:
with gateway_from_environment() as value:
yield value
@pytest.fixture
def peer(gateway: Gateway) -> Iterator[Gateway]:
url: Final = os.environ["INTEGRATION_PEER_URL"]
assert url.rstrip("/") != str(gateway.client.base_url).rstrip("/")
with Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) as cache:
eventually(lambda: cache.pubsub_numsub("litellm_proxy.auth_cache_invalidation")[0][1], lambda count: count >= 2)
with httpx.Client(base_url=url, timeout=15, trust_env=False) as client:
yield Gateway(client, gateway.key, gateway.upstream_url)

View file

@ -0,0 +1,80 @@
{
"groups": {
"management": [
"management",
"authorization",
"configuration"
],
"accounting": [
"pricing",
"spend"
],
"database": [
"database"
],
"providers": [
"providers",
"routing",
"streaming"
],
"extensions": [
"mcp",
"observability",
"compatibility"
]
},
"tests": {
"tests/integration/management/test_key_updates.py::test_update_preserves_independent_fields_and_serving": [
"mgmt.key.update.preserves_independent_fields"
],
"tests/integration/pricing/test_configured_prices.py::test_custom_price_is_reported_and_charged": [
"quota_management.spend_tracking.custom_price.matches_input_rates"
],
"tests/integration/providers/test_request_boundary.py::test_internal_request_state_does_not_reach_provider": [
"other.provider_wire.internal_parameters_filtered"
],
"tests/integration/pricing/test_configured_prices.py::test_default_prices_survive_nullable_sibling_and_reload": [
"quota_management.spend_tracking.default_prices.survive_nullable_sibling_reload"
],
"tests/integration/providers/test_request_boundary.py::test_upstream_rejects_corruption_and_accepts_supported_metadata": [
"other.provider_wire.validator_rejects_corruption"
],
"tests/integration/pricing/test_configured_prices.py::test_loaded_router_preserves_cached_defaults_during_real_requests": [
"quota_management.spend_tracking.default_prices.loaded_router_preserves_cached_defaults"
],
"tests/integration/management/test_partial_update_sequences.py::test_generated_partial_updates_preserve_persisted_and_effective_state": [
"mgmt.key.update.generated_sequences_preserve_state"
],
"tests/integration/management/test_partial_update_sequences.py::test_zero_false_and_empty_values_are_not_treated_as_omission": [
"mgmt.key.update.false_zero_and_empty_values_affect_serving"
],
"tests/integration/management/test_partial_update_sequences.py::test_project_omission_clear_and_invalid_update_have_distinct_effects": [
"mgmt.key.update.project_clear_preserves_scope",
"mgmt.key.update.invalid_batch_is_atomic"
],
"tests/integration/authorization/test_warmed_policy.py::test_generated_policy_changes_reach_both_warmed_workers": [
"mgmt.key.update.two_workers_enforce_warmed_policy"
],
"tests/integration/authorization/test_warmed_policy.py::test_scim_deactivation_blocks_null_and_false_keys_but_preserves_other_owners": [
"mgmt.user.scim.deactivation_includes_nullable_blocked_keys"
],
"tests/integration/authorization/test_warmed_policy.py::test_warmed_team_role_demotion_prevents_later_management_writes": [
"mgmt.team.member_update.demoted_role_cannot_write"
],
"tests/integration/configuration/test_effective_settings.py::test_model_block_changes_actual_route_and_leaves_other_route_working": [
"mgmt.model.block.changes_serving_and_preserves_control"
],
"tests/integration/configuration/test_effective_settings.py::test_saved_retry_setting_controls_real_attempts_and_restores": [
"mgmt.router_settings.update.changes_observed_attempt_count"
],
"tests/integration/configuration/test_effective_settings.py::test_credential_value_update_and_model_reload_reach_provider": [
"mgmt.credential.update.saved_value_reaches_wire"
],
"tests/integration/management/test_partial_update_sequences.py::test_denied_key_update_preserves_saved_grants_and_serving": [
"mgmt.key.update.denied_request_preserves_effective_state"
],
"tests/integration/authorization/test_warmed_policy.py::test_expiry_and_explicit_clear_reach_both_warmed_workers": [
"mgmt.key.update.expiry_changes_reach_warmed_workers"
]
}
}

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