scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-19 00:01:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 87

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

Browse source
This commit is contained in:
moe-berri 2026-09-14 17:16:55 -07:00
commit 49d0e9da59
179 changed files with 12489 additions and 1994 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

48
.circleci/config.yml
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:

142
.circleci/scripts/run_integration.sh Normal file
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"

53
.circleci/scripts/stop_integration_processes.py Normal file
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)))

43
.circleci/scripts/wait_integration_services.py Normal file
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()

11
.github/codeql/codeql-config.yml vendored
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

67
.github/scripts/assert_ci_coverage.py vendored
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())

4
litellm-rust/crates/ai-gateway/src/bin/trace_parity_gateway.rs
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,

10
litellm-rust/crates/ai-gateway/src/trace_parity.rs
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()))

2
litellm/__init__.py
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.

13
litellm/caching/caching.py
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):
"""

155
litellm/caching/redis_cache.py
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

10
litellm/constants.py
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

2
litellm/integrations/arize/arize_phoenix_prompt_manager.py
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,

2
litellm/integrations/bitbucket/bitbucket_prompt_manager.py
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,

2
litellm/integrations/prometheus.py
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,
@ -41,7 +42,6 @@ from litellm.proxy._types import (
LiteLLM_UserTable,
UserAPIKeyAuth,
)
from litellm.proxy.hooks.rate_limiter_utils import PROXY_LLM_PROVIDER_FALLBACK
from litellm.repositories.base_repository import BaseRepository
from litellm.repositories.organization_repository import OrganizationRepository
from litellm.repositories.team_repository import TeamRepository

10
litellm/integrations/weights_biases.py
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)

10
litellm/litellm_core_utils/core_helpers.py
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]:

2
litellm/litellm_core_utils/coroutine_checker.py
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:

73
litellm/litellm_core_utils/fallback_generalizations.py
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)

2
litellm/litellm_core_utils/prompt_templates/factory.py
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"))

5
litellm/litellm_core_utils/realtime_streaming.py
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:

2
litellm/litellm_core_utils/token_counter.py
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:

388
litellm/llms/anthropic/prompt_cache_prediction.py Normal file
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,
)

9
litellm/llms/azure_ai/common_utils.py
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:
"""

449
litellm/llms/bedrock/realtime/handler.py
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

6
litellm/main.py
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(

1059
litellm/model_prices_and_context_window_backup.json
View file

File diff suppressed because it is too large Load diff

4
litellm/proxy/_types.py
View file

@ -893,6 +893,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,
@ -4427,6 +4428,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):

273
litellm/proxy/auth/auth_checks.py
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,

35
litellm/proxy/auth/auth_utils.py
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)

2
litellm/proxy/auth/resolvers/grants.py
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

3
litellm/proxy/auth/route_checks.py
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,
)

15
litellm/proxy/auth/user_api_key_auth.py
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]"):
@ -1652,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:
@ -1692,6 +1695,7 @@ async def _user_api_key_auth_builder(
route=route,
request=request,
llm_router=llm_router,
team_id=valid_token.team_id,
)
),
)
@ -2091,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:
@ -2209,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)
@ -2239,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)
@ -2734,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
@ -2850,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)
@ -3301,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:
@ -3408,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)
@ -3449,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)

4
litellm/proxy/client/cli/README.md
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.

25
litellm/proxy/client/cli/commands/statusline_script.py
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(

5
litellm/proxy/common_request_processing.py
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,

10
litellm/proxy/common_utils/html_forms/default_credentials_hint.py Normal file
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"))
)

91
litellm/proxy/common_utils/prompt_cache_pricing.py Normal file
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

8
litellm/proxy/db/health_check_latest.py
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 ()

6
litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py
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()

2
litellm/proxy/guardrails/guardrail_hooks/javelin/javelin.py
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

2
litellm/proxy/guardrails/guardrail_hooks/mcp_end_user_permission/__init__.py
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", ""),

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

@ -27,6 +27,7 @@ if TYPE_CHECKING:
_SANITIZE_FILE_FAIL_OPEN_TIMEOUT_SECONDS: Final = 30.0
_SANITIZE_FILE_QUEUED_STATUSES: Final = frozenset({"created", "in progress"})
class PromptSecurityGuardrailMissingSecrets(Exception):
@ -512,16 +513,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:

7
litellm/proxy/health_endpoints/_health_endpoints.py
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

2
litellm/proxy/hooks/__init__.py
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 ##

204
litellm/proxy/hooks/parallel_request_limiter_v3.py
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])
@ -2681,12 +2687,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 (
@ -2811,34 +2812,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)
@ -3425,6 +3403,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,
@ -3453,59 +3533,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.

142
litellm/proxy/hooks/prompt_cache_prediction.py Normal file
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

3
litellm/proxy/hooks/rate_limiter_utils.py
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,

2
litellm/proxy/management_endpoints/cost_tracking_settings.py
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)

278
litellm/proxy/management_endpoints/prompt_cache_prediction.py Normal file
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,
)

207
litellm/proxy/management_endpoints/team_endpoints.py
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)])

6
litellm/proxy/management_endpoints/ui_sso.py
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,

18
litellm/proxy/management_helpers/team_metadata_validation.py
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)

18
litellm/proxy/pass_through_endpoints/streaming_handler.py
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,

48
litellm/proxy/proxy_server.py
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:
@ -11909,6 +11914,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,
@ -11928,6 +11940,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")
@ -12031,6 +12044,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:
@ -12060,12 +12076,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)
######################################################################
@ -15049,6 +15063,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"],
@ -15096,7 +15117,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
@ -15144,7 +15165,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(
@ -15195,7 +15216,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
@ -15263,7 +15284,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(
@ -15832,10 +15853,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,

2
litellm/proxy/spend_tracking/spend_tracking_utils.py
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,

3
litellm/proxy/utils.py
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

68
litellm/router.py
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

5
litellm/router_utils/cooldown_handlers.py
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)

16
litellm/rust_bridge/lifecycle.py
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(

67
litellm/types/management_endpoints/prompt_cache_prediction.py Normal file
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

8
litellm/types/proxy/auth/auth_checks.py Normal file
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.")

8
litellm/types/router.py
View file

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

10
litellm/types/utils.py
View file

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

39
litellm/utils.py
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,
@ -6265,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:

1059
model_prices_and_context_window.json
View file

File diff suppressed because it is too large Load diff

11
ruff.toml
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

19
tests/integration/README.md Normal file
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

0
tests/integration/__init__.py Normal file
View file

0
tests/integration/_support/__init__.py Normal file
View file

170
tests/integration/_support/client.py Normal 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)

14
tests/integration/_support/database.py Normal file
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())

52
tests/integration/_support/generation.py Normal file
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")

31
tests/integration/_support/manifest.py Normal file
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

16
tests/integration/_support/proxy.py Normal file
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()

122
tests/integration/_support/upstream.py Normal file
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()

197
tests/integration/authorization/test_warmed_policy.py Normal file
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)

123
tests/integration/configuration/test_effective_settings.py Normal file
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}"

103
tests/integration/conftest.py Normal file
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)

80
tests/integration/contracts.json Normal file
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"
]
}
}

40
tests/integration/management/test_key_updates.py Normal file
View file

@ -0,0 +1,40 @@
from typing import Final
from hashlib import sha256
import pytest
from integration._support.client import Gateway, object_value
from integration._support.database import read_rows
@pytest.mark.covers("mgmt.key.update.preserves_independent_fields")
def test_update_preserves_independent_fields_and_serving(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
model: Final = scenario.model()
key: Final = scenario.key(models=[model], key_alias="before", metadata={"retained": "value"})
gateway.chat(model, key=key)
gateway.post("/key/update", {"key": key, "key_alias": "after"})
info: Final = object_value(gateway.get("/key/info", {"key": key})["info"])
assert info["key_alias"] == "after"
assert info["models"] == [model]
assert object_value(info["metadata"])["retained"] == "value"
response: Final = gateway.chat(model, key=key)
assert object_value(response["usage"])["total_tokens"] == 40
replacement: Final = scenario.model()
gateway.post("/key/update", {"key": key, "models": [replacement]})
saved: Final = read_rows(
'SELECT key_alias, models, metadata FROM "LiteLLM_VerificationToken" WHERE token = %s',
(sha256(key.encode()).hexdigest(),),
)
assert len(saved) == 1
assert saved[0]["models"] == [replacement]
assert saved[0]["key_alias"] == "after"
denied: Final = gateway.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": "old grant"}]},
key=key,
)
assert denied.status_code == 403, denied.text
assert object_value(object_value(denied.json())["error"])["type"] == "key_model_access_denied"
assert object_value(gateway.chat(replacement, key=key)["usage"])["total_tokens"] == 40

200
tests/integration/management/test_partial_update_sequences.py Normal file
View file

@ -0,0 +1,200 @@
from contextlib import ExitStack
from hashlib import sha256
from typing import Final
import pytest
from hypothesis import strategies as st
from hypothesis.stateful import RuleBasedStateMachine, invariant, rule, run_state_machine_as_test
from pydantic import JsonValue
from integration._support.client import Gateway, object_value
from integration._support.database import read_rows
from integration._support.generation import LIFECYCLE_SETTINGS, bounded_http_requests
@pytest.mark.covers("mgmt.key.update.generated_sequences_preserve_state")
def test_generated_partial_updates_preserve_persisted_and_effective_state(gateway: Gateway) -> None:
class KeyUpdates(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.key = scenario.key(models=[self.models[0]], key_alias="initial", metadata={"revision": "initial"})
self.expected: dict[str, JsonValue] = {
"models": [self.models[0]], "key_alias": "initial", "metadata": {"revision": "initial"}
}
gateway.chat(self.models[0], key=self.key)
except BaseException:
with budget.cleanup():
self.resources.close()
raise
@rule(alias=st.sampled_from(("first", "second", "", "unicode-λ")))
def alias(self, alias: str) -> None:
gateway.post("/key/update", {"key": self.key, "key_alias": alias})
self.expected["key_alias"] = alias
@rule(index=st.integers(min_value=0, max_value=1), both=st.booleans())
def grant(self, index: int, both: bool) -> None:
models: Final = list(self.models) if both else [self.models[index]]
gateway.post("/key/update", {"key": self.key, "models": models})
self.expected["models"] = models
@rule(value=st.sampled_from(("", "a", "different", "λ")))
def metadata(self, value: str) -> None:
gateway.post("/key/update", {"key": self.key, "metadata": {"revision": value}})
self.expected["metadata"] = {"revision": value}
@invariant()
def persisted_state_and_serving_match(self) -> None:
rows: Final = read_rows(
'SELECT models, key_alias, metadata FROM "LiteLLM_VerificationToken" WHERE token = %s',
(sha256(self.key.encode()).hexdigest(),),
)
assert rows == [self.expected]
info: Final = object_value(gateway.get("/key/info", {"key": self.key})["info"])
assert {field: info[field] for field in self.expected} == self.expected
for model in self.models:
response: Final = gateway.request(
"POST", "/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": "generated update control"}]},
key=self.key,
)
if model in self.expected["models"]:
assert response.status_code == 200, response.text
assert response.json()["usage"]["total_tokens"] == 40
else:
assert response.status_code == 403, response.text
assert response.json()["error"]["type"] == "key_model_access_denied"
def teardown(self) -> None:
with budget.cleanup():
self.resources.close()
with bounded_http_requests((gateway,), limit=2000) as budget:
run_state_machine_as_test(KeyUpdates, settings=LIFECYCLE_SETTINGS)
@pytest.mark.covers("mgmt.key.update.false_zero_and_empty_values_affect_serving")
def test_zero_false_and_empty_values_are_not_treated_as_omission(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
models: Final = (scenario.model(), scenario.model())
key: Final = scenario.key(models=[models[0]], max_budget=0, metadata={"ordinary": "value"})
denied: Final = gateway.request(
"POST", "/v1/chat/completions",
{"model": models[0], "messages": [{"role": "user", "content": "zero budget"}]}, key=key,
)
assert denied.status_code == 429, denied.text
assert denied.json()["error"]["type"] == "budget_exceeded"
gateway.post("/key/update", {"key": key, "max_budget": 1, "models": [], "metadata": {}})
info: Final = object_value(gateway.get("/key/info", {"key": key})["info"])
assert (info["models"], info["metadata"], info["max_budget"]) == ([], {}, 1)
for model in models:
assert object_value(gateway.chat(model, key=key)["usage"])["total_tokens"] == 40
gateway.post("/key/update", {"key": key, "blocked": True})
blocked: Final = gateway.request(
"POST", "/v1/chat/completions",
{"model": models[0], "messages": [{"role": "user", "content": "blocked control"}]}, key=key,
)
assert blocked.status_code == 401, blocked.text
assert blocked.json()["error"]["type"] == "auth_error"
gateway.post("/key/update", {"key": key, "blocked": False})
assert object_value(gateway.chat(models[0], key=key)["usage"])["total_tokens"] == 40
rows: Final = read_rows(
'SELECT blocked, models, metadata, max_budget FROM "LiteLLM_VerificationToken" WHERE token = %s',
(sha256(key.encode()).hexdigest(),),
)
assert rows == [{"blocked": False, "models": [], "metadata": {}, "max_budget": 1.0}]
gateway.post("/key/update", {"key": key, "max_budget": 0})
assert read_rows(
'SELECT max_budget FROM "LiteLLM_VerificationToken" WHERE token = %s',
(sha256(key.encode()).hexdigest(),),
) == [{"max_budget": 0.0}]
zero_after_update: Final = gateway.request(
"POST", "/v1/chat/completions",
{"model": models[0], "messages": [{"role": "user", "content": "updated zero budget"}]}, key=key,
)
assert zero_after_update.status_code == 429, zero_after_update.text
assert zero_after_update.json()["error"]["type"] == "budget_exceeded"
gateway.post("/key/update", {"key": key, "max_budget": None})
assert read_rows(
'SELECT max_budget FROM "LiteLLM_VerificationToken" WHERE token = %s',
(sha256(key.encode()).hexdigest(),),
) == [{"max_budget": None}]
assert object_value(gateway.chat(models[0], key=key)["usage"])["total_tokens"] == 40
@pytest.mark.covers("mgmt.key.update.project_clear_preserves_scope", "mgmt.key.update.invalid_batch_is_atomic")
def test_project_omission_clear_and_invalid_update_have_distinct_effects(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
model: Final = scenario.model()
outside: Final = scenario.model()
team: Final = scenario.team(models=[model])
project: Final = scenario.project(team, models=[model])
other: Final = scenario.project(team, models=[model])
key: Final = scenario.key(team_id=team, project_id=project, models=[model], key_alias="before", max_budget=5)
gateway.chat(model, key=key)
gateway.post("/key/update", {"key": key, "key_alias": "after"})
digest: Final = sha256(key.encode()).hexdigest()
def saved() -> list[dict[str, object]]:
return read_rows(
'SELECT key_alias, project_id, team_id, models, max_budget FROM "LiteLLM_VerificationToken" '
'WHERE token = %s', (digest,),
)
before: Final = saved()
assert before == [{"key_alias": "after", "project_id": project, "team_id": team, "models": [model], "max_budget": 5}]
for invalid in (other, ""):
rejected: Final = gateway.request(
"POST", "/key/update", {"key": key, "project_id": invalid, "key_alias": "must-not-persist"}
)
assert rejected.status_code == 400, rejected.text
assert saved() == before
gateway.chat(model, key=key)
gateway.post("/project/update", {"project_id": project, "blocked": True})
denied: Final = gateway.request(
"POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": "blocked project"}]},
key=key,
)
assert denied.status_code == 401, denied.text
assert denied.json()["error"]["type"] == "auth_error"
for _ in range(2):
gateway.post("/key/update", {"key": key, "project_id": None})
assert saved() == [{**before[0], "project_id": None}]
assert object_value(gateway.chat(model, key=key)["usage"])["total_tokens"] == 40
outside_request: Final = gateway.request(
"POST", "/v1/chat/completions",
{"model": outside, "messages": [{"role": "user", "content": "detached scope control"}]}, key=key,
)
assert outside_request.status_code == 403, outside_request.text
assert outside_request.json()["error"]["type"] == "key_model_access_denied"
@pytest.mark.covers("mgmt.key.update.denied_request_preserves_effective_state")
def test_denied_key_update_preserves_saved_grants_and_serving(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
model: Final = scenario.model()
outside: Final = scenario.model()
owner: Final = scenario.user(user_role="internal_user")
other: Final = scenario.user(user_role="internal_user")
key: Final = scenario.key(user_id=owner, models=[model], key_alias="unchanged", max_budget=2)
caller: Final = scenario.key(user_id=other, models=[model], allowed_routes=["/key/update", "/v1/chat/completions"])
gateway.chat(model, key=key)
denied: Final = gateway.request(
"POST", "/key/update", {"key": key, "key_alias": "wrong", "models": [outside], "max_budget": 0}, key=caller
)
assert denied.status_code == 403, denied.text
assert read_rows(
'SELECT user_id, models, key_alias, max_budget FROM "LiteLLM_VerificationToken" WHERE token = %s',
(sha256(key.encode()).hexdigest(),),
) == [{"user_id": owner, "models": [model], "key_alias": "unchanged", "max_budget": 2.0}]
assert object_value(gateway.chat(model, key=key)["usage"])["total_tokens"] == 40
rejected: Final = gateway.request(
"POST", "/v1/chat/completions",
{"model": outside, "messages": [{"role": "user", "content": "unchanged scope"}]}, key=key,
)
assert rejected.status_code == 403, rejected.text
assert rejected.json()["error"]["type"] == "key_model_access_denied"

148
tests/integration/pricing/test_configured_prices.py Normal file
View file

@ -0,0 +1,148 @@
from collections.abc import Iterator, Mapping
from typing import Final
from pathlib import Path
import uuid
import pytest
import yaml
from integration._support.client import Gateway, eventually, object_value, string_value
from integration._support.database import read_rows
@pytest.mark.covers("quota_management.spend_tracking.custom_price.matches_input_rates")
def test_custom_price_is_reported_and_charged(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002)
response: Final = gateway.request(
"POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": "price control"}]}
)
assert response.status_code == 200, response.text
assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(20 * 0.001 + 20 * 0.002)
entries: Final = gateway.get("/model/info")["data"]
assert isinstance(entries, list)
matching: Final = tuple(object_value(entry) for entry in entries if object_value(entry)["model_name"] == model)
assert len(matching) == 1
params: Final = object_value(matching[0]["litellm_params"])
assert params["input_cost_per_token"] == 0.001
assert params["output_cost_per_token"] == 0.002
@pytest.mark.covers("quota_management.spend_tracking.default_prices.survive_nullable_sibling_reload")
def test_default_prices_survive_nullable_sibling_and_reload(gateway: Gateway) -> None:
for registration_order in (("custom", "omitted", "nullable"), ("nullable", "omitted", "custom")):
with gateway.scenario() as scenario:
configured: Final = {
"custom": {"input_cost_per_token": 0.001, "output_cost_per_token": 0.002},
"omitted": {},
"nullable": {"input_cost_per_token": None, "output_cost_per_token": None},
}
rates: Final = {
"custom": (0.001, 0.002),
"omitted": (0.00000015, 0.0000006),
"nullable": (0.00000015, 0.0000006),
}
models: Final = {kind: scenario.model(**configured[kind]) for kind in registration_order}
def observe_requests(
registration_order: tuple[str, ...],
models: Mapping[str, str],
rates: Mapping[str, tuple[float, float]],
) -> Iterator[tuple[str, float]]:
for generation in range(2):
entries: Final = gateway.get("/model/info")["data"]
assert isinstance(entries, list)
kinds: Final = tuple(reversed(registration_order)) if generation else registration_order
for index, kind in enumerate(kinds):
model: Final = models[kind]
target: Final = next(
object_value(entry) for entry in entries if object_value(entry)["model_name"] == model
)
info: Final = object_value(target["model_info"])
assert info["input_cost_per_token"] == rates[kind][0]
assert info["output_cost_per_token"] == rates[kind][1]
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{
"model": model,
"messages": [
{"role": "user", "content": f"price {generation * len(registration_order) + index}"}
],
},
)
assert response.status_code == 200, response.text
expected: Final = 20 * rates[kind][0] + 20 * rates[kind][1]
assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(expected, rel=1e-6)
request_id: Final = string_value(object_value(response.json())["id"])
yield request_id, expected
target: Final = next(
object_value(entry)
for entry in entries
if object_value(entry)["model_name"] == models["nullable"]
)
identity: Final = string_value(object_value(target["model_info"])["id"])
updated: Final = gateway.request(
"PATCH", f"/model/{identity}/update", {"model_info": {"description": "reload price contract"}}
)
assert updated.status_code == 200, updated.text
observations: Final = tuple(observe_requests(registration_order, models, rates))
for request_id, expected in observations:
rows: Final = eventually(
lambda request_id=request_id: read_rows(
'SELECT request_id, spend, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" '
"WHERE request_id = %s",
(request_id,),
),
lambda values: len(values) == 1,
seconds=70,
)
assert rows[0]["prompt_tokens"] == 20
assert rows[0]["completion_tokens"] == 20
assert float(rows[0]["spend"]) == pytest.approx(expected, rel=1e-6)
@pytest.mark.covers("quota_management.spend_tracking.default_prices.loaded_router_preserves_cached_defaults")
def test_loaded_router_preserves_cached_defaults_during_real_requests(gateway: Gateway, tmp_path: Path) -> None:
from litellm import Router
aliases: Final = (f"pricing-{uuid.uuid4().hex}", f"pricing-{uuid.uuid4().hex}")
path: Final = tmp_path / "models.yaml"
path.write_text(
yaml.safe_dump(
{
"model_list": [
{
"model_name": alias,
"litellm_params": {
"model": "openai/gpt-4o-mini",
"api_key": "integration-provider-key",
"api_base": f"{gateway.upstream_url}/v1",
},
"model_info": {"id": alias, **pricing},
}
for alias, pricing in zip(
aliases, ({}, {"input_cost_per_token": None, "output_cost_per_token": None}), strict=True
)
]
}
)
)
for reverse in (False, True):
configured: Final = yaml.safe_load(path.read_text())["model_list"]
router: Final = Router(model_list=list(reversed(configured)) if reverse else configured, num_retries=0)
try:
for alias in (*aliases, *reversed(aliases)):
result: Final = router.completion(
model=alias, messages=[{"role": "user", "content": "router price control"}]
)
assert result.usage.prompt_tokens == 20
assert result.usage.completion_tokens == 20
deployment: Final = router.get_deployment(model_id=alias)
assert deployment is not None
info: Final = router.get_router_model_info(deployment=deployment, received_model_name=alias)
assert info["input_cost_per_token"] == 0.00000015
assert info["output_cost_per_token"] == 0.0000006
finally:
router.reset()

57
tests/integration/providers/test_request_boundary.py Normal file
View file

@ -0,0 +1,57 @@
from typing import Final
import httpx
import pytest
from integration._support.client import Gateway, JSON_OBJECT, object_value
@pytest.mark.covers("other.provider_wire.internal_parameters_filtered")
def test_internal_request_state_does_not_reach_provider(gateway: Gateway) -> None:
with gateway.scenario() as scenario, httpx.Client(base_url=gateway.upstream_url, trust_env=False) as upstream:
upstream.get("/__observations").raise_for_status()
model: Final = scenario.model()
key: Final = scenario.key(models=[model], tpm_limit=10000, rpm_limit=100)
result: Final = gateway.post(
"/v1/chat/completions",
{
"model": model,
"messages": [{"role": "user", "content": "wire contract"}],
"temperature": 0.4,
"max_tokens": 20,
"timeout": 12,
},
key=key,
)
assert object_value(result["usage"])["total_tokens"] == 40
observations: Final = JSON_OBJECT.validate_json(upstream.get("/__observations").content)["requests"]
assert isinstance(observations, list)
assert len(observations) == 1
observed: Final = object_value(observations[0])
body: Final = object_value(observed["body"])
assert body["model"] == "gpt-4o-mini"
assert body["messages"] == [{"role": "user", "content": "wire contract"}]
assert body["temperature"] == 0.4
assert body["max_tokens"] == 20
assert observed["authorization"] == "Bearer integration-provider-key"
assert "litellm_metadata" not in body
assert "litellm_params" not in body
assert "timeout" not in body
assert "tpm" not in body
@pytest.mark.covers("other.provider_wire.validator_rejects_corruption")
def test_upstream_rejects_corruption_and_accepts_supported_metadata(gateway: Gateway) -> None:
with httpx.Client(base_url=gateway.upstream_url, trust_env=False) as upstream:
missing: Final = upstream.post("/v1/chat/completions", json={"model": "gpt-4o-mini"})
assert missing.status_code == 400
assert (
object_value(object_value(missing.json())["error"])["message"] == "model and nonempty messages are required"
)
body: Final = {"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "strict control"}]}
leaked: Final = upstream.post("/v1/chat/completions", json={**body, "litellm_metadata": {"hidden": "value"}})
assert leaked.status_code == 400
assert "litellm_metadata" in str(object_value(object_value(leaked.json())["error"])["message"])
valid: Final = upstream.post("/v1/chat/completions", json={**body, "metadata": {"purpose": "synthetic"}})
assert valid.status_code == 200, valid.text
assert object_value(object_value(valid.json())["usage"])["total_tokens"] == 40

17
tests/integration/proxy_config.yaml Normal file
View file

@ -0,0 +1,17 @@
model_list: []
general_settings:
master_key: os.environ/LITELLM_MASTER_KEY
database_url: os.environ/DATABASE_URL
store_model_in_db: true
disable_spend_logs: false
proxy_batch_write_at: 1
litellm_settings:
enable_redis_auth_cache: true
cache: true
cache_params:
type: redis
host: os.environ/REDIS_HOST
port: os.environ/REDIS_PORT
router_settings:
num_retries: 0
disable_cooldowns: true

71
tests/integration/run.py Normal file
View file

@ -0,0 +1,71 @@
from __future__ import annotations
import argparse
import json
import os
import subprocess
import sys
from pathlib import Path
from types import MappingProxyType
from typing import Final
GROUPS: Final = MappingProxyType(json.loads(Path(__file__).with_name("contracts.json").read_text())["groups"])
def main() -> int:
parser: Final = argparse.ArgumentParser()
parser.add_argument("group", choices=tuple(GROUPS))
parser.add_argument("--results", type=Path, default=Path("test-results/integration"))
parser.add_argument("--seed", type=int, default=int(os.environ.get("INTEGRATION_SEED", "4106601")))
options: Final = parser.parse_args()
root: Final = Path(__file__).resolve().parents[2]
selected: Final = tuple(
str(path.relative_to(root))
for folder in GROUPS[options.group]
for path in sorted((root / "tests/integration" / folder).glob("test_*.py"))
)
if not selected:
parser.error(f"No integration contracts selected for {options.group}")
output: Final = options.results.resolve()
output.mkdir(parents=True, exist_ok=True)
manifest: Final = json.loads((root / "tests/integration/contracts.json").read_text())["tests"]
expected: Final = sorted(node for node in manifest if node.split("::", 1)[0] in selected)
if not expected or set(selected) != {node.split("::", 1)[0] for node in expected}:
parser.error("Every selected file must have canonical manifest nodes")
environment: Final = {
**os.environ,
"PYTHONPATH": os.pathsep.join((str(root), str(root / "tests"), str(root / "tests/e2e"))),
"INTEGRATION_RESULTS_DIR": str(output),
"LITELLM_LOCAL_MODEL_COST_MAP": "True",
}
result: Final = subprocess.call(
[
sys.executable,
"-m",
"pytest",
*selected,
"-vv",
"--strict-markers",
"-p",
"no:pytest-retry",
"-p",
"no:rerunfailures",
"--timeout=90",
"--durations=15",
f"--hypothesis-seed={options.seed}",
f"--junitxml={output / 'junit.xml'}",
],
cwd=root,
env=environment,
)
if result != 0:
return result
evidence: Final = json.loads((output / "execution.json").read_text())
if not evidence["complete"] or sorted(evidence["passed"]) != expected or sorted(evidence["collected"]) != expected:
print("Executed integration nodes differ from the canonical manifest", file=sys.stderr)
return 1
return 0
if __name__ == "__main__":
raise SystemExit(main())

1
tests/proxy_behavior/auth/test_auth_object_prefetch.py
View file

@ -66,6 +66,7 @@ async def test_join_binds_the_membership_to_the_requested_team(prisma):
cache = _frozen_cache()
refs = AuthObjectRefs(user_id=user_id, team_id=team_a, membership_user_id=user_id, organization_id=org_id)
await prefetch_auth_objects(refs=refs, user_api_key_cache=cache, prisma_client=prisma)
assert cache.in_memory_cache.get_cache(f"org_id:{org_id}") is not None
dead_db = _dead_db()
membership = await get_team_membership(

6
tests/proxy_unit_tests/test_proxy_server.py
View file

@ -2143,7 +2143,7 @@ async def test_model_info_alias_without_prisma(hidden):
user_api_key_dict=UserAPIKeyAuth(models=[]),
)
models = resp["data"]
models = json.loads(resp.body)["data"]
alias_found = any(
m["model_name"] == model_alias
@ -2207,7 +2207,7 @@ async def test_proxy_model_group_alias_checks(prisma_client, hidden): # noqa: F
resp = await model_info_v1(
user_api_key_dict=UserAPIKeyAuth(models=[]),
)
models = resp["data"]
models = json.loads(resp.body)["data"]
is_model_alias_in_list = False
for item in models:
if model_alias == item["model_name"]:
@ -2284,7 +2284,7 @@ async def test_proxy_model_group_info_rerank(prisma_client): # noqa: F811 # py
resp = await model_info_v1(
user_api_key_dict=UserAPIKeyAuth(models=[]),
)
models = resp["data"]
models = json.loads(resp.body)["data"]
assert models[0]["model_info"]["mode"] == "rerank"
resp = await model_group_info(
user_api_key_dict=UserAPIKeyAuth(models=[]),

29
tests/router_unit_tests/test_router_helper_utils.py
View file

@ -1,3 +1,4 @@
import json
import os
import traceback
from dotenv import load_dotenv
@ -628,17 +629,29 @@ def test_deployment_callback_respects_cooldown_time(model_list):
assert mock_set.call_args.kwargs["time_to_cooldown"] == 0
def test_log_retry(model_list):
"""Test if the '_log_retry' function is working correctly"""
import time
@pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"])
def test_log_retry(model_list, metadata_key):
"""log_retry appends one flat record per failed attempt and copies neither the request kwargs nor
the request metadata into it"""
router = Router(model_list=model_list)
new_kwargs = router.log_retry(
kwargs={"metadata": {}},
e=Exception(),
kwargs={
"model": "gpt-3.5-turbo",
"api_key": "sk-must-not-be-recorded",
"messages": [{"role": "user", "content": "hi"}],
metadata_key: {"model_info": {"id": "deployment-1"}, "attempted_retries": 2, "user_api_key": "sk-proxy"},
},
e=litellm.RateLimitError(message="slow down", llm_provider="openai", model="gpt-3.5-turbo"),
)
assert "metadata" in new_kwargs
assert "previous_models" in new_kwargs["metadata"]
assert json.loads(json.dumps(new_kwargs[metadata_key]["previous_models"])) == [
{
"model_group": "gpt-3.5-turbo",
"deployment_id": "deployment-1",
"exception_type": "RateLimitError",
"exception_string": "litellm.RateLimitError: slow down",
"attempted_retries": 2,
}
]
def test_update_usage(model_list):

2
tests/rust-python-harness/AGENTS.md
View file

@ -63,7 +63,7 @@ tests/rust-python-harness/
- Examples: `run e2e_parity --surface sdk --function ocr`, `run unit_tests_parity --function ocr --pytest-arg=-x`, or `run all --function ocr`
- `cli/catalog.py` discovers strategies, validates their Python definitions, and orders them; `cli/__init__.py` builds the Click command tree; `cli/commands.py` runs selected cases
- `e2e_parity/` compares SDK objects, exceptions, callbacks, and streams, or gateway HTTP responses
- `trace_parity/` compares mapped operations, call counts, and required execution ordering; before running it rebuilds the native bridge with the `trace-parity` feature whenever `litellm-rust` sources are newer than the installed extension (`shared/native_build.py`)
- `trace_parity/` prints every collected Python call under `litellm/` and every Rust span without comparing them; mappings only filter the separate unit-test mapping strategy. Before running it rebuilds the native bridge with the `trace-parity` feature whenever `litellm-rust` sources are newer than the installed extension (`shared/native_build.py`)
- E2E and trace strategies load their registered module cases and run surface-specific execution from their folders
- `unit_tests_mapping/contracts.py` owns typed harness-side mapping contracts, per-function contracts live below `cases/`, and `mappings.py` exports the registry; live test discovery derives unmapped Python and Rust-only tests without an exhaustive manifest
- `unit_tests_mapping/runner.py` validates confirmed mappings against the live Python and Rust inventories and attaches the derived status report

15
tests/rust-python-harness/cli/__init__.py
View file

@ -58,16 +58,29 @@ def _strategy_command(strategy: Strategy) -> click.Command:
help=runner_argument.help,
)
)
for runner_option in strategy.definition.runner_options:
name: Final = runner_option.option.removeprefix("--").replace("-", "_")
params.append(
click.Option(
(runner_option.option, name),
type=click.Choice(runner_option.choices),
help=runner_option.help,
)
)
def run_strategy(
sdk_functions: tuple[str, ...],
surface: str | None = None,
runner_args: tuple[str, ...] = (),
**runner_options: str | None,
) -> int:
selected_functions: Final = cast(frozenset[SdkFunction], frozenset(sdk_functions))
selected_surface: Final = cast(Surface | None, surface)
cases: Final = select_cases((strategy,), selected_functions, selected_surface)
return run_command((strategy,), cases, runner_args)
option_args: Final = tuple(
f"--{name.replace('_', '-')}={value}" for name, value in runner_options.items() if value is not None
)
return run_command((strategy,), cases, (*runner_args, *option_args))
return click.Command(
strategy.id,

19
tests/rust-python-harness/cli/catalog.py
View file

@ -21,9 +21,7 @@ def _load_strategy_module(name: str, folder: Path, prefix: str | None) -> Module
if prefix is not None:
return importlib.import_module(f"{prefix}.{name}")
module_name: Final = _synthetic_module_name(folder)
spec: Final = importlib.util.spec_from_file_location(
module_name, folder / "__init__.py"
)
spec: Final = importlib.util.spec_from_file_location(module_name, folder / "__init__.py")
if spec is None or spec.loader is None:
raise ValueError(f"{folder}: cannot load strategy package")
module: Final = importlib.util.module_from_spec(spec)
@ -59,9 +57,7 @@ def _load_strategy(name: str, folder: Path, prefix: str | None) -> Strategy:
if duplicates:
raise ValueError(f"{folder}: duplicate strategy cases: {duplicates}")
expected: Final = frozenset(
(surface, function)
for surface in (definition.surfaces or (None,))
for function in SDK_FUNCTIONS
(surface, function) for surface in (definition.surfaces or (None,)) for function in SDK_FUNCTIONS
)
actual: Final = frozenset(keys)
if actual != expected:
@ -73,8 +69,7 @@ def _load_strategy(name: str, folder: Path, prefix: str | None) -> Strategy:
incompatible: Final = tuple(
(case.surface, case.sdk_function)
for case in definition.cases
if case.spec.disposition is CaseDisposition.RUNNABLE
and not isinstance(case.spec, definition.runnable_spec)
if case.spec.disposition is CaseDisposition.RUNNABLE and not isinstance(case.spec, definition.runnable_spec)
)
if incompatible:
raise ValueError(f"{folder}: runnable cases do not match {definition.runnable_spec.__name__}: {incompatible}")
@ -102,14 +97,10 @@ def _load_strategy(name: str, folder: Path, prefix: str | None) -> Strategy:
def load_catalog(root: Path | None = None) -> tuple[Strategy, ...]:
resolved: Final = STRATEGIES_ROOT if root is None else root
prefix: Final = _STRATEGIES_PACKAGE.__name__ if resolved == STRATEGIES_ROOT else None
folders: Final = tuple(
info.name for info in pkgutil.iter_modules([str(resolved)]) if info.ispkg
)
folders: Final = tuple(info.name for info in pkgutil.iter_modules([str(resolved)]) if info.ispkg)
if not folders:
raise ValueError(f"No strategy packages found below {resolved}")
strategies: Final = tuple(
_load_strategy(name, resolved / name, prefix) for name in sorted(folders)
)
strategies: Final = tuple(_load_strategy(name, resolved / name, prefix) for name in sorted(folders))
ids: Final = [strategy.id for strategy in strategies]
if len(set(ids)) != len(ids):
raise ValueError(f"Duplicate strategy id in {resolved}")

6
tests/rust-python-harness/cli/commands.py
View file

@ -21,8 +21,7 @@ def select_cases(
case
for strategy in strategies
for case in strategy.cases
if (not sdk_functions or case.sdk_function in sdk_functions)
and (surface is None or case.surface == surface)
if (not sdk_functions or case.sdk_function in sdk_functions) and (surface is None or case.surface == surface)
)
@ -32,8 +31,7 @@ def run_command(
runner_args: Sequence[str] = (),
) -> int:
grouped: Final = {
strategy.id: tuple(case for case in cases if case.strategy_id == strategy.id)
for strategy in strategies
strategy.id: tuple(case for case in cases if case.strategy_id == strategy.id) for strategy in strategies
}
visible: Final = tuple(strategy for strategy in strategies if grouped[strategy.id])
runners: Final = tuple(replace(strategy, cases=grouped[strategy.id]) for strategy in visible)

21
tests/rust-python-harness/cli/test_cli.py
View file

@ -242,7 +242,7 @@ def _assert_unavailable_cell(strategy: Strategy, case: HarnessCase, section_titl
def test_every_unavailable_case_finishes_and_explains_itself() -> None:
section_titles: Final = {
"e2e_parity": "End-to-end parity outcomes",
"trace_parity": "trace comparisons",
"trace_parity": "traces",
"unit_tests_mapping": "Python/Rust unit-test mappings",
"unit_tests_parity": "Python backend parity outcomes",
"unit_tests_rust": "Native Rust unit-test outcomes",
@ -359,6 +359,25 @@ def test_strategy_command_forwards_repeated_filters_and_runner_arguments(
]
def test_trace_command_forwards_engine_and_scenario(monkeypatch: pytest.MonkeyPatch) -> None:
cli: Final = importlib.import_module("tests.rust-python-harness.cli")
captured: list[tuple[str, ...]] = []
def capture_run(
strategies: Sequence[Strategy],
cases: Sequence[HarnessCase],
runner_args: Sequence[str] = (),
) -> int:
del strategies, cases
captured.append(tuple(runner_args))
return 0
monkeypatch.setattr(cli, "run_command", capture_run)
assert main(["run", "trace_parity", "--scenario", "async-mistral", "--engine", "python"]) == 0
assert captured == [("async-mistral", "--engine=python")]
def test_omitted_surface_selects_every_strategy_surface(monkeypatch: pytest.MonkeyPatch) -> None:
cli: Final = importlib.import_module("tests.rust-python-harness.cli")
selected: list[str] = []

4
tests/rust-python-harness/conftest.py
View file

@ -19,9 +19,7 @@ def subprocess_test_environment(monkeypatch: pytest.MonkeyPatch) -> None:
def cargo_project(tmp_path: Path) -> Callable[[str, str], Path]:
def create(package: str, source: str) -> Path:
manifest: Final = tmp_path / "Cargo.toml"
manifest.write_text(
f'[package]\nname = "{package}"\nversion = "0.1.0"\nedition = "2021"\n[workspace]\n'
)
manifest.write_text(f'[package]\nname = "{package}"\nversion = "0.1.0"\nedition = "2021"\n[workspace]\n')
(tmp_path / "src").mkdir()
(tmp_path / "src/lib.rs").write_text(source)
return manifest

21
tests/rust-python-harness/shared/native_build.py
View file

@ -16,6 +16,11 @@ _RUST_ROOT: Final = "litellm-rust"
_LOCKFILE: Final = "Cargo.lock"
_SOURCE_SUFFIXES: Final = frozenset({".rs", ".toml"})
_FAILURE_OUTPUT_LINES: Final = 15
_TRACE_CHECK: Final = (
"from litellm.rust_bridge import get_native_bridge; "
"bridge = get_native_bridge(); "
"raise SystemExit(0 if bridge is not None and getattr(bridge, '_trace', None) is not None else 1)"
)
def needs_rebuild(native_mtime: float | None, newest_source_mtime: float | None) -> bool:
@ -73,6 +78,17 @@ def _rebuild(repo_root: Path) -> tuple[bool, str]:
return completed.returncode == 0, "\n".join(lines[-_FAILURE_OUTPUT_LINES:])
def _installed_bridge_has_trace(repo_root: Path) -> bool:
completed: Final = subprocess.run(
(sys.executable, "-c", _TRACE_CHECK),
cwd=repo_root,
stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL,
check=False,
)
return completed.returncode == 0
def trace_bridge_error() -> str | None:
bridge: Final = get_native_bridge()
if bridge is None:
@ -85,7 +101,10 @@ def trace_bridge_error() -> str | None:
def ensure_trace_bridge(repo_root: Path) -> str | None:
native_path: Final = _native_module_path()
native_mtime: Final = native_path.stat().st_mtime if native_path is not None and native_path.exists() else None
if needs_rebuild(native_mtime, _newest_source_mtime(repo_root)):
rebuild_required: Final = needs_rebuild(
native_mtime, _newest_source_mtime(repo_root)
) or not _installed_bridge_has_trace(repo_root)
if rebuild_required:
print(f"Rebuilding native Rust bridge ({BRIDGE_FEATURE} feature)...", flush=True)
succeeded: Final
output: Final

1
tests/rust-python-harness/shared/parity/fixtures/recording.py
View file

@ -191,6 +191,7 @@ class _RecordingHandler(LocalHttpHandler):
self.end_headers()
self.wfile.write(body)
def _recording_provider(spec: UpstreamEndpoint) -> AbstractContextManager[_RecordingProvider]:
return serve_in_thread(_RecordingProvider(spec))

2
tests/rust-python-harness/shared/parity/fixtures/store.py
View file

@ -15,6 +15,8 @@ from .cassette import deserialize_cassette, serialize_cassette
from .recording import RecordedInteraction
FIXTURE_SCHEMA_VERSION: Final = 1
class FixtureInput(Protocol):
def canonical_input(self) -> dict[str, object]: ...

3
tests/rust-python-harness/shared/parity/fixtures/test_pipeline.py
View file

@ -45,8 +45,8 @@ class _Upstream(LocalHttpServer):
super().__init__(("127.0.0.1", 0), _UpstreamHandler)
self.response_status: Final = status
class _UpstreamHandler(LocalHttpHandler):
class _UpstreamHandler(LocalHttpHandler):
def do_POST(self) -> None:
length: Final = int(self.headers.get("content-length") or "0")
self.rfile.read(length)
@ -59,6 +59,7 @@ class _UpstreamHandler(LocalHttpHandler):
self.end_headers()
self.wfile.write(body)
def _upstream(status: int = 200) -> AbstractContextManager[_Upstream]:
return serve_in_thread(_Upstream(status))

1
tests/rust-python-harness/shared/parity/fixtures/test_recording.py
View file

@ -238,6 +238,7 @@ class _ControlledUpstreamHandler(LocalHttpHandler):
self.end_headers()
self.wfile.write(body)
def _controlled_upstream(
stream_chunks: tuple[bytes, ...] = _SSE_CHUNKS,
) -> AbstractContextManager[_ControlledUpstream]:

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