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 88

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

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

# Conflicts:
#	litellm/integrations/custom_guardrail.py
This commit is contained in:
mateo 2026-09-15 00:26:16 +00:00
commit 0cedfaaf16
66 changed files with 7560 additions and 522 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())

25
litellm/integrations/custom_guardrail.py
View file

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

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

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(

1047
litellm/model_prices_and_context_window_backup.json
View file

File diff suppressed because it is too large Load diff

1
litellm/proxy/_types.py
View file

@ -887,6 +887,7 @@ class LiteLLMRoutes(enum.Enum):
"/auto_router/validate_complexity_router_config",
# Per-session auto-router read - the endpoint scopes the row to the caller's own key hash
"/auto_router/session",
"/cost/predict-cache",
# Agent registry - reads are role-scoped and writes are proxy-admin-gated
# inside agent_endpoints/endpoints.py
*agent_management_routes,

5
litellm/proxy/auth/auth_checks.py
View file

@ -895,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 (
@ -4471,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]:
@ -4518,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:

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)

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)

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,

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=[],
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 ()

92
litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py
View file

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

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])
@ -2661,12 +2667,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 (
@ -2791,34 +2792,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)
@ -3396,6 +3374,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),
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 [
*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=[
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,
@ -3424,59 +3504,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

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

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,

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

10
litellm/types/utils.py
View file

@ -3755,6 +3755,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",

1047
model_prices_and_context_window.json
View file

File diff suppressed because it is too large Load diff

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

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

@ -2625,7 +2625,7 @@ class TestLoggingOnlyApplyGuardrail:
assert [e["guardrail_status"] for e in entries] == ["success"]
@pytest.mark.asyncio
async def test_native_lifecycle_hook_guardrail_is_left_alone(self):
async def test_native_lifecycle_hook_guardrail_scans_in_logging_only(self):
class _NativeHooks(_ApplyOnlyObserver):
use_native_lifecycle_hooks = True
@ -2634,9 +2634,9 @@ class TestLoggingOnlyApplyGuardrail:
out_kwargs, out_response = await guardrail.async_logging_hook(kwargs, response, CallTypes.acompletion.value)
assert guardrail.calls == []
assert out_kwargs is kwargs
assert guardrail.calls == [("request", ["hello there"]), ("response", ["general kenobi"])]
assert out_response is response
assert out_kwargs["standard_logging_object"]["guardrail_information"]
@pytest.mark.asyncio
async def test_aresponses_scans_logged_messages_when_input_is_cleared(self):
@ -2890,6 +2890,61 @@ class TestCustomGuardrailPostCallSuccessDeploymentHook:
assert len(_guardrail_entries(request_data)) == 1
class _NativeLifecycleLoggingGuardrail(CustomGuardrail):
"""Native lifecycle guardrail that also implements apply_guardrail, like the azure guards."""
use_native_lifecycle_hooks: ClassVar[bool] = True
def __init__(self):
from litellm.types.guardrails import GuardrailEventHooks
super().__init__(
guardrail_name="native-logging-guardrail",
event_hook=GuardrailEventHooks.logging_only,
)
self.calls: list[tuple[Literal["request", "response"], list[str]]] = []
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict[str, object],
input_type: Literal["request", "response"],
logging_obj: "LiteLLMLoggingObj | None" = None,
) -> GenericGuardrailAPIInputs:
self.calls.append((input_type, list(inputs.get("texts") or [])))
return inputs
@pytest.mark.asyncio
async def test_native_lifecycle_guardrail_logging_only_scans_assembled_response():
"""A use_native_lifecycle_hooks guardrail accepts mode logging_only and its
async_logging_hook scans kwargs["async_complete_streaming_response"], not the raw result."""
from litellm.types.utils import Choices, Message, ModelResponse
guardrail = _NativeLifecycleLoggingGuardrail()
assembled = ModelResponse(
choices=[Choices(message=Message(role="assistant", content="assembled stream text"))]
)
sentinel_result = object()
kwargs = {
"model": "gpt-5.4-mini",
"messages": [{"role": "user", "content": "hi"}],
"litellm_call_id": "call-1",
"litellm_params": {"metadata": {}},
"optional_params": {},
"standard_logging_object": {"guardrail_information": None},
"async_complete_streaming_response": assembled,
}
out_kwargs, out_result = await guardrail.async_logging_hook(
kwargs=kwargs, result=sentinel_result, call_type=CallTypes.acompletion.value
)
assert out_result is sentinel_result
assert ("response", ["assembled stream text"]) in guardrail.calls
assert out_kwargs["standard_logging_object"]["guardrail_information"]
class TestPreCallHookResponseIsNotLoggedVerbatim:
"""Regression for LIT-6935: a pre_call hook returning the request payload leaked the prompt
into ``guardrail_response`` and from there onto OTEL guardrail spans."""

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

@ -4560,7 +4560,7 @@ GEMINI_35_FLASH_LITE_TIER_RATES_BY_SURFACE = [
("gemini", "priority", 5.4e-07, 4.5e-06, 5e-08),
("vertex_ai", None, 3e-07, 2.5e-06, 3e-08),
("vertex_ai", "flex", 1.5e-07, 1.25e-06, 1.5e-08),
("vertex_ai", "priority", 5.4e-07, 4.5e-06, 5e-08),
("vertex_ai", "priority", 5.4e-07, 4.5e-06, 5.4e-08),
]

209
tests/test_litellm/llms/anthropic/test_anthropic_prompt_cache_prediction.py Normal file
View file

@ -0,0 +1,209 @@
import json
from collections.abc import Mapping
from datetime import datetime
from types import SimpleNamespace
from typing import Final
import httpx
import pytest
import respx
from pydantic import JsonValue
import litellm
from litellm.caching.dual_cache import DualCache
from litellm.caching.llm_caching_handler import LLMClientCache
from litellm.llms.anthropic.count_tokens import handler as count_handler
from litellm.llms.anthropic.experimental_pass_through.messages.transformation import DEFAULT_ANTHROPIC_API_VERSION
from litellm.llms.anthropic.prompt_cache_prediction import (
NativePredictionTarget,
cache_scope,
count_prompt_tokens,
parse_observed_cache,
parse_prompt,
resolve_prediction_target,
supported_prediction_headers,
)
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.models.credentials import CredentialItem
from litellm.proxy import proxy_server
from litellm.proxy.hooks.prompt_cache_prediction import PromptCacheObserver, lookup
from litellm.proxy.management_endpoints.prompt_cache_prediction import predict_arm
from litellm.proxy.utils import InternalUsageCache
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
from litellm.types.utils import CacheCreationTokenDetails, ModelResponse, PromptTokensDetailsWrapper, Usage
_MODEL: Final = "claude-sonnet-5"
_KEY: Final = "test-provider-key"
_CALLER: Final = "test-caller-hash"
_DEPLOYMENT: Final = "test-native-deployment"
def _body() -> dict[str, JsonValue]:
return {
"model": _MODEL,
"system": "Keep this context",
"tools": [{"name": "lookup", "input_schema": {"type": "object"}}],
"messages": [{"role": "user", "content": [
{"type": "text", "text": "A cacheable prefix", "cache_control": {"type": "ephemeral"}}
]}],
}
@pytest.mark.parametrize("version", [None, "2099-01-01", DEFAULT_ANTHROPIC_API_VERSION])
@pytest.mark.asyncio
async def test_observer_records_only_version_supported_by_token_counter(version: str | None) -> None:
cache: Final = DualCache()
observer: Final = PromptCacheObserver(InternalUsageCache(dual_cache=cache), clock=lambda: 1010.0)
body: Final = _body()
prefix: Final = parse_prompt(body)
assert prefix is not None
headers: Final = {"x-api-key": _KEY, **({"anthropic-version": version} if version is not None else {})}
wire: Final = httpx.Request("POST", "https://api.anthropic.com/v1/messages", headers=headers, json=body)
response: Final = ModelResponse(
model=_MODEL,
usage=Usage(
prompt_tokens=311,
completion_tokens=2,
total_tokens=313,
prompt_tokens_details=PromptTokensDetailsWrapper(
cached_tokens=100,
cache_creation_tokens=200,
cache_creation_token_details=CacheCreationTokenDetails(
ephemeral_5m_input_tokens=200, ephemeral_1h_input_tokens=0
),
),
),
)
await observer.async_log_success_event(
{
"call_type": "anthropic_messages",
"custom_llm_provider": "anthropic",
"httpx_response": httpx.Response(200, request=wire),
"first_api_call_start_time": datetime.fromtimestamp(1000.0),
"standard_logging_object": {
"status": "success", "model_id": _DEPLOYMENT,
"metadata": {"user_api_key_hash": _CALLER},
},
},
response,
datetime.fromtimestamp(1010.0),
datetime.fromtimestamp(1010.0),
)
default_scope: Final = cache_scope(_CALLER, _DEPLOYMENT, _KEY, _MODEL)
found: Final = await lookup(cache, default_scope, prefix, now=1010.0)
assert (found is not None) == (version == DEFAULT_ANTHROPIC_API_VERSION)
if version != DEFAULT_ANTHROPIC_API_VERSION:
other_scope: Final = cache_scope(_CALLER, _DEPLOYMENT, _KEY, _MODEL, version or "")
assert await lookup(cache, other_scope, prefix, now=1010.0) is None
@pytest.mark.parametrize("headers, supported", [
({}, True),
({"Anthropic-Version": DEFAULT_ANTHROPIC_API_VERSION}, True),
({"anthropic-version": "2099-01-01"}, False),
({"Anthropic-Beta": ""}, False),
({"anthropic-beta": "future-feature"}, False),
])
def test_prediction_header_eligibility(headers: Mapping[str, str], supported: bool) -> None:
assert supported_prediction_headers(headers) is supported
@pytest.mark.asyncio
async def test_provider_count_uses_same_version_and_preserves_native_input(monkeypatch: pytest.MonkeyPatch) -> None:
body: Final = _body()
requests: Final[list[httpx.Request]] = []
def provider(request: httpx.Request) -> httpx.Response:
requests.append(request)
return httpx.Response(200, json={"input_tokens": 311})
client: Final = AsyncHTTPHandler()
await client.client.aclose()
client.client = httpx.AsyncClient(transport=httpx.MockTransport(provider))
monkeypatch.setattr(count_handler, "get_async_httpx_client", lambda **kwargs: client)
try:
assert await count_prompt_tokens(_MODEL, _KEY, body) == 311
finally:
await client.client.aclose()
assert len(requests) == 1
assert requests[0].headers["anthropic-version"] == DEFAULT_ANTHROPIC_API_VERSION
assert requests[0].url == "https://api.anthropic.com/v1/messages/count_tokens"
assert json.loads(requests[0].content) == body
@pytest.mark.parametrize("source", ["static", "database"])
@pytest.mark.asyncio
async def test_environment_credential_matches_native_count_and_observed_scope(
source: str, monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setenv("LIT7658_PROVIDER_KEY", _KEY)
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache())
params: Final = {
"model": f"anthropic/{_MODEL}", "api_key": "os.environ/LIT7658_PROVIDER_KEY",
"api_base": "https://api.anthropic.com",
}
router: Final = litellm.Router(model_list=[{
"model_name": "test-native", "litellm_params": dict(params), "model_info": {"id": _DEPLOYMENT},
}] if source == "static" else [], num_retries=0)
if source == "database":
monkeypatch.setattr(proxy_server, "llm_router", router)
assert proxy_server.ProxyConfig()._add_deployment([SimpleNamespace(
model_id=_DEPLOYMENT, model_name="test-native", model_info={}, litellm_params=dict(params),
)]) == 1
deployment: Final = router.get_deployment(_DEPLOYMENT)
assert deployment is not None
target: Final = resolve_prediction_target(deployment.litellm_params)
assert isinstance(target, NativePredictionTarget)
body: Final = _body()
with respx.mock() as upstream:
native: Final = upstream.post("https://api.anthropic.com/v1/messages").respond(200, json={
"id": "msg_test", "type": "message", "role": "assistant", "model": _MODEL,
"content": [{"type": "text", "text": "Hello"}], "stop_reason": "end_turn", "stop_sequence": None,
"usage": {"input_tokens": 11, "output_tokens": 1, "cache_read_input_tokens": 300},
})
counter: Final = upstream.post("https://api.anthropic.com/v1/messages/count_tokens").respond(
200, json={"input_tokens": 311},
)
await router.aanthropic_messages(
model="test-native", max_tokens=1, **{key: value for key, value in body.items() if key != "model"},
)
assert await count_prompt_tokens(target.model, target.api_key, body) == 311
assert native.call_count == counter.call_count == 1
assert native.calls.last.request.headers["x-api-key"] == counter.calls.last.request.headers["x-api-key"] == _KEY
observed: Final = parse_observed_cache(native.calls.last.request, ModelResponse(
model=_MODEL, usage=Usage(
prompt_tokens=311, completion_tokens=1, total_tokens=312,
prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=300),
),
), _CALLER, _DEPLOYMENT)
assert observed is not None
assert observed.scope == cache_scope(_CALLER, _DEPLOYMENT, target.api_key, target.model)
@pytest.mark.parametrize("inline_key", [None, _KEY])
@pytest.mark.asyncio
async def test_named_credential_is_explicitly_unsupported_before_count(
inline_key: str | None, monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(litellm, "credential_list", [CredentialItem(
credential_name="test-named", credential_info={}, credential_values={"api_key": "test-named-provider-key"},
)])
deployment: Final = Deployment(
model_name="test-native",
litellm_params=LiteLLM_Params(
model=f"anthropic/{_MODEL}", api_key=inline_key, litellm_credential_name="test-named",
),
model_info=ModelInfo(id=_DEPLOYMENT),
)
body: Final = _body()
prefix: Final = parse_prompt(body)
assert prefix is not None
async def count(model: str, api_key: str, body: Mapping[str, JsonValue]) -> int | None:
pytest.fail("Unsupported named credentials must not reach provider counting")
arm: Final = await predict_arm(deployment, body, prefix, _CALLER, DualCache(), count)
assert arm.cache_state == "unknown"
assert arm.reason == "unsupported_deployment_configuration"
assert arm.estimate is None and arm.cold is None and arm.warm is None

226
tests/test_litellm/proxy/auth/test_auth_utils.py
View file

@ -433,6 +433,232 @@ def test_get_model_from_request_no_request_extracts_model():
)
def _cache_prediction_router():
from litellm.router import Router
return Router(model_list=[
{
"model_name": group,
"litellm_params": {"model": "anthropic/claude-sonnet-5", "api_key": "test-provider-key"},
"model_info": {"id": deployment_id, "team_id": team_id},
}
for group, deployment_id, team_id in (
("current-group", "current-id", None), ("candidate-group", "candidate-id", None),
("own-group", "own-id", "prediction-team"), ("foreign-group", "foreign-id", "foreign-team"),
)
])
@pytest.mark.parametrize("candidate,team_id,expected", [
("candidate-id", None, ["current-group", "candidate-group"]),
("current-id", None, "current-group"),
("missing-id", None, None),
("candidate-group", None, None),
("own-id", None, None),
("own-id", "prediction-team", ["current-group", "own-group"]),
("foreign-id", "prediction-team", None),
])
def test_cache_prediction_auth_resolves_only_exact_deployment_ids(candidate, team_id, expected):
assert get_model_from_request(
request_data={
"current_deployment_id": "current-id", "candidate_deployment_id": candidate,
"request": {"model": "caller-controlled-provider-model"},
},
route="/cost/predict-cache",
llm_router=_cache_prediction_router(),
team_id=team_id,
) == expected
def _cache_prediction_auth_app(
monkeypatch, allowed_routes, user_models, metadata=None, *, team_id=None, key_models=None, team_models=None
):
import importlib
from unittest.mock import AsyncMock
from fastapi import FastAPI
import litellm.proxy.proxy_server as proxy_server
from litellm.caching.dual_cache import DualCache
from litellm.proxy._types import LiteLLM_TeamTableCachedObj, LiteLLM_UserTable, LitellmUserRoles, ProxyException
from litellm.proxy.auth import auth_checks
from litellm.proxy.hooks.parallel_request_limiter_v3 import _PROXY_MaxParallelRequestsHandler_v3
from litellm.proxy.management_endpoints import prompt_cache_prediction as endpoint
from litellm.proxy.utils import InternalUsageCache, ProxyLogging
auth = importlib.import_module("litellm.proxy.auth.user_api_key_auth")
router = _cache_prediction_router()
allowed_models = ["current-group", "candidate-group", "own-group"]
token = UserAPIKeyAuth(
api_key="test-proxy-key-hash", user_id="prediction-user", user_role=LitellmUserRoles.INTERNAL_USER,
models=allowed_models if key_models is None else key_models, team_id=team_id,
team_models=allowed_models if team_models is None else team_models,
allowed_routes=allowed_routes, metadata=metadata or {},
)
user = LiteLLM_UserTable(
user_id=token.user_id, user_role=LitellmUserRoles.INTERNAL_USER.value, models=user_models,
)
async def authenticate(request, request_data, **_headers):
await auth._enforce_key_and_fallback_model_access(
valid_token=token, request_data=request_data, route=request.url.path, request=request,
llm_model_list=router.get_model_list(), llm_router=router,
)
return token
monkeypatch.setattr(auth, "_user_api_key_auth_builder", authenticate)
monkeypatch.setattr(auth, "get_user_object", AsyncMock(return_value=user))
team = LiteLLM_TeamTableCachedObj(team_id=team_id, models=token.team_models) if team_id else None
monkeypatch.setattr(auth, "get_team_object", AsyncMock(return_value=team))
monkeypatch.setattr(auth_checks, "get_team_object", AsyncMock(return_value=team))
monkeypatch.setattr(auth_checks, "get_team_membership", AsyncMock(return_value=None))
monkeypatch.setattr(auth, "get_global_proxy_spend", AsyncMock(return_value=0))
monkeypatch.setattr(proxy_server, "master_key", "test-master-key")
monkeypatch.setattr(proxy_server, "user_custom_auth", None)
monkeypatch.setattr(proxy_server, "general_settings", {})
monkeypatch.setattr(proxy_server, "llm_router", router)
monkeypatch.setattr(proxy_server, "llm_model_list", router.get_model_list())
monkeypatch.setattr(proxy_server, "prisma_client", None)
monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache())
logging = ProxyLogging(user_api_key_cache=DualCache())
logging.proxy_hook_mapping["parallel_request_limiter"] = _PROXY_MaxParallelRequestsHandler_v3(
InternalUsageCache(dual_cache=DualCache())
)
monkeypatch.setattr(proxy_server, "proxy_logging_obj", logging)
counts = AsyncMock(return_value=6_000)
monkeypatch.setattr(endpoint, "count_prompt_tokens", counts)
app = FastAPI()
app.include_router(endpoint.router)
app.add_exception_handler(ProxyException, proxy_server.openai_exception_handler)
return app, counts
def _cache_prediction_payload(candidate="candidate-id", current="current-id"):
return {
"current_deployment_id": current, "candidate_deployment_id": candidate,
"request": {"messages": [{"role": "user", "content": [{
"type": "text", "text": "Stable cached context",
"cache_control": {"type": "ephemeral"},
}]}]},
}
@pytest.mark.asyncio
@pytest.mark.parametrize("allowed_routes,user_models,candidate,status_code", [
(["/chat/completions"], ["current-group", "candidate-group"], "candidate-id", 403),
(["/cost/predict-cache"], ["current-group"], "candidate-id", 403),
(["/cost/*"], ["current-group", "candidate-group"], "candidate-id", 200),
(["/cost/predict-cache"], ["current-group"], "current-id", 200),
(["/cost/predict-cache"], ["current-group"], "missing-id", 404),
])
async def test_cache_prediction_authorizes_route_and_personal_models_before_provider_counts(
monkeypatch, allowed_routes, user_models, candidate, status_code
):
import httpx
app, counts = _cache_prediction_auth_app(monkeypatch, allowed_routes, user_models)
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client:
response = await client.post("/cost/predict-cache", json=_cache_prediction_payload(candidate))
assert response.status_code == status_code, response.text
if status_code == 200:
assert counts.await_count == (2 if candidate == "current-id" else 4)
else:
assert counts.await_count == 0
@pytest.mark.asyncio
@pytest.mark.parametrize("arm", ["current_deployment_id", "candidate_deployment_id"])
@pytest.mark.parametrize("team_id,key_models,user_models,team_models", [
(None, ["*"], ["*"], None),
(None, ["current-group", "candidate-group"], ["*"], None),
(None, ["*"], ["current-group", "candidate-group"], None),
("prediction-team", ["*"], ["*"], ["current-group", "candidate-group"]),
])
async def test_cache_prediction_hides_foreign_and_missing_ids_before_model_authorization(
monkeypatch, arm, team_id, key_models, user_models, team_models
):
import httpx
app, counts = _cache_prediction_auth_app(
monkeypatch, ["/cost/predict-cache"], user_models,
team_id=team_id, key_models=key_models, team_models=team_models,
)
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client:
missing = await client.post("/cost/predict-cache", json={**_cache_prediction_payload(), arm: "missing-id"})
foreign = await client.post("/cost/predict-cache", json={**_cache_prediction_payload(), arm: "foreign-id"})
assert missing.status_code == foreign.status_code == 404, foreign.text
assert missing.json() == foreign.json() == {"detail": "Deployment not found"}
assert counts.await_count == 0
@pytest.mark.asyncio
@pytest.mark.parametrize("arm", ["current_deployment_id", "candidate_deployment_id"])
@pytest.mark.parametrize("key_models,team_models,status_code", [
(["*"], ["*"], 200),
(["current-group", "candidate-group"], ["*"], 403),
(["*"], ["current-group", "candidate-group"], 403),
])
async def test_cache_prediction_checks_each_visible_team_deployment_model(
monkeypatch, arm, key_models, team_models, status_code
):
import httpx
app, counts = _cache_prediction_auth_app(
monkeypatch, ["/cost/predict-cache"], ["*"],
team_id="prediction-team", key_models=key_models, team_models=team_models,
)
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client:
response = await client.post("/cost/predict-cache", json={**_cache_prediction_payload(), arm: "own-id"})
assert response.status_code == status_code, response.text
assert counts.await_count == (4 if status_code == 200 else 0)
@pytest.mark.asyncio
@pytest.mark.parametrize("arm", ["current_deployment_id", "candidate_deployment_id"])
async def test_cache_prediction_checks_each_visible_personal_deployment_model(monkeypatch, arm):
import httpx
app, counts = _cache_prediction_auth_app(monkeypatch, ["/cost/predict-cache"], ["current-group"])
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client:
response = await client.post(
"/cost/predict-cache", json={**_cache_prediction_payload(candidate="current-id"), arm: "candidate-id"}
)
assert response.status_code == 403, response.text
assert response.json()["error"]["type"] == "user_model_access_denied"
assert counts.await_count == 0
@pytest.mark.asyncio
@pytest.mark.parametrize("header_tag,key_tags,limit,status_code,provider_calls", [
("limited", [], 1, 429, 1),
(None, ["limited"], 1, 429, 1),
("limited", ["limited"], 4, 200, 4),
("unlimited", [], 1, 200, 4),
])
async def test_cache_prediction_preserves_authenticated_header_and_key_tag_rpm(
monkeypatch, header_tag, key_tags, limit, status_code, provider_calls
):
import httpx
app, counts = _cache_prediction_auth_app(
monkeypatch, ["/cost/predict-cache"], ["current-group", "candidate-group"],
metadata={"tag_rpm_limit": {"limited": limit}, "tags": key_tags},
)
headers = {"x-litellm-tags": header_tag} if header_tag else {}
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client:
response = await client.post("/cost/predict-cache", json=_cache_prediction_payload(), headers=headers)
assert response.status_code == status_code, response.text
assert counts.await_count == provider_calls
if limit == 4:
exhausted = await client.post("/cost/predict-cache", json=_cache_prediction_payload(), headers=headers)
assert exhausted.status_code == 429, exhausted.text
assert counts.await_count == 4
assert all("metadata" not in call.args[2] for call in counts.await_args_list)
def test_get_model_from_request_supports_google_model_names_with_slashes():
assert (
get_model_from_request(

105
tests/test_litellm/proxy/common_utils/test_prompt_cache_pricing.py Normal file
View file

@ -0,0 +1,105 @@
from typing import Final
import pytest
import litellm
from litellm.proxy.common_utils.prompt_cache_pricing import price_cache_tokens
from litellm.types.management_endpoints.prompt_cache_prediction import CacheTokenBuckets
@pytest.mark.parametrize(
("model", "expected"),
[("anthropic/claude-sonnet-4-5", 1.26), ("anthropic/claude-sonnet-4-6", 0.63)],
)
def test_prices_all_cache_buckets_at_total_context_tier(model: str, expected: float) -> None:
tokens: Final = CacheTokenBuckets(
uncached_input_tokens=100_000,
cache_read_input_tokens=50_000,
cache_creation_5m_input_tokens=20_000,
cache_creation_1h_input_tokens=40_000,
)
assert price_cache_tokens(model, "unconfigured-deployment", tokens) == pytest.approx(expected)
@pytest.mark.parametrize(("total", "expected"), [(200_000, 0.387), (200_001, 0.774006)])
def test_long_context_tier_starts_above_threshold(total: int, expected: float) -> None:
tokens: Final = CacheTokenBuckets(
uncached_input_tokens=total - 100_000,
cache_creation_1h_input_tokens=10_000,
cache_read_input_tokens=90_000,
)
actual: Final = price_cache_tokens("anthropic/claude-sonnet-4-5", "unconfigured-deployment", tokens)
assert actual == pytest.approx(expected)
def test_deployment_tariff_wins_without_proxy_discounts_or_margins(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(litellm, "model_cost", litellm.model_cost.copy())
litellm.Router(
model_list=[
{
"model_name": "cache-pricing-test",
"litellm_params": {
"model": "anthropic/claude-sonnet-4-6",
"api_key": "test-only",
"input_cost_per_token": 0.00001,
"output_cost_per_token": 0.00002,
"cache_read_input_token_cost": 0.000001,
"cache_creation_input_token_cost": 0.0000125,
"cache_creation_input_token_cost_above_1hr": 0.00002,
},
"model_info": {"id": "cache-pricing-test-a"},
}
]
)
monkeypatch.setattr(litellm, "cost_discount_config", {"anthropic": 0.5})
monkeypatch.setattr(litellm, "cost_margin_config", {"global": {"percentage": 0.3, "fixed_amount": 1.0}})
tokens: Final = CacheTokenBuckets(
uncached_input_tokens=3_000,
cache_read_input_tokens=4_000,
cache_creation_5m_input_tokens=1_000,
cache_creation_1h_input_tokens=2_000,
)
assert price_cache_tokens("anthropic/claude-sonnet-4-6", "cache-pricing-test-a", tokens) == pytest.approx(0.0865)
@pytest.mark.parametrize("rate", [None, -1.0, float("nan"), float("inf"), "0.00001", True])
def test_unknown_for_absent_or_invalid_active_cache_rate(monkeypatch: pytest.MonkeyPatch, rate: object) -> None:
monkeypatch.setitem(
litellm.model_cost,
"cache-pricing-invalid",
{
"litellm_provider": "anthropic",
"mode": "chat",
"input_cost_per_token": 0.00001,
"output_cost_per_token": 0.00002,
"cache_creation_input_token_cost_above_1hr": rate,
},
)
tokens: Final = CacheTokenBuckets(cache_creation_1h_input_tokens=4_000)
assert price_cache_tokens("anthropic/claude-sonnet-4-6", "cache-pricing-invalid", tokens) is None
def test_missing_input_price_is_unknown_even_when_get_model_info_defaults_to_zero(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setitem(litellm.model_cost, "cache-pricing-missing", {"litellm_provider": "anthropic", "mode": "chat"})
tokens: Final = CacheTokenBuckets(uncached_input_tokens=4_000)
assert price_cache_tokens("cache-pricing-missing", "unconfigured-deployment", tokens) is None
def test_explicit_free_pricing_is_not_unknown(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setitem(
litellm.model_cost,
"cache-pricing-free",
{
"litellm_provider": "anthropic",
"mode": "chat",
"input_cost_per_token": 0.0,
"output_cost_per_token": 0.0,
"cache_read_input_token_cost": 0.0,
"cache_creation_input_token_cost": 0.0,
"cache_creation_input_token_cost_above_1hr": 0.0,
},
)
tokens: Final = CacheTokenBuckets(uncached_input_tokens=100, cache_read_input_tokens=5_000)
assert price_cache_tokens("anthropic/claude-sonnet-4-6", "cache-pricing-free", tokens) == 0.0

10
tests/test_litellm/proxy/db/test_health_check_latest.py
View file

@ -8,6 +8,7 @@ from litellm.proxy.db.health_check_latest import (
LATEST_HEALTH_CHECKS_SQL,
fetch_latest_health_checks,
fetch_latest_health_checks_for_models,
query_latest_health_checks,
)
@ -83,6 +84,15 @@ async def test_fetch_all_degrades_to_no_rows_when_the_query_fails():
assert await fetch_latest_health_checks(prisma) == ()
@pytest.mark.asyncio
async def test_query_all_raises_when_the_query_fails_instead_of_reading_as_an_empty_table():
"""The background save decides what to write from this read; a failure has to be told apart from no rows."""
prisma = _prisma([])
prisma.db.query_raw.side_effect = RuntimeError("db down")
with pytest.raises(RuntimeError, match="db down"):
await query_latest_health_checks(prisma)
@pytest.mark.asyncio
async def test_fetch_all_degrades_to_no_rows_for_a_malformed_row():
assert await fetch_latest_health_checks(_prisma([{"unexpected": "shape"}])) == ()

391
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py
View file

@ -2,6 +2,8 @@ import asyncio
import base64
import io
import json
from collections.abc import Iterator, Sequence
from typing import cast
from unittest.mock import AsyncMock, MagicMock, Mock, patch
import pytest
@ -14,7 +16,7 @@ import litellm
import litellm.types.utils
from litellm._logging import verbose_proxy_logger
from litellm.caching import DualCache
from litellm.llms.custom_httpx.http_handler import MaskedHTTPStatusError
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, MaskedHTTPStatusError
from litellm.proxy.guardrails.anthropic_sse import anthropic_sse_error_frames
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails.guardrail_hooks.model_armor import ModelArmorGuardrail
@ -4929,3 +4931,390 @@ def test_every_responses_delta_event_is_in_the_scanned_set():
}
assert not missing
assert "response.mcp_call_arguments.delta" in _RESPONSES_DELTA_EVENT_TYPES
def _clean_armor_response() -> dict[str, object]:
return {
"sanitizationResult": {
"filterMatchState": "NO_MATCH_FOUND",
"filterResults": {},
}
}
def _flagged_armor_response() -> dict[str, object]:
return {
"sanitizationResult": {
"filterMatchState": "MATCH_FOUND",
"filterResults": {"rai": {"raiFilterResult": {"matchState": "MATCH_FOUND"}}},
}
}
class _FakeArmorHandler(AsyncHTTPHandler):
def __init__(self, responses: Sequence[dict[str, object] | Exception]):
self.responses: Iterator[dict[str, object] | Exception] = iter(responses)
self.calls: list[dict[str, object]] = []
self.raise_on_call: Exception | None = None
async def post(
self,
url: str,
json: dict[str, object] | None = None,
headers: dict[str, str] | None = None,
**kwargs: object,
) -> httpx.Response:
if self.raise_on_call is not None:
raise self.raise_on_call
if json is not None:
self.calls.append(json)
response: dict[str, object] | Exception = next(self.responses)
if isinstance(response, Exception):
raise response
return httpx.Response(200, json=response, request=httpx.Request("POST", url))
async def _async_token_provider() -> tuple[str, str]:
return ("test-token", "test-project")
def _logging_only_guardrail(
responses: Sequence[dict[str, object] | Exception] = (_clean_armor_response(), _clean_armor_response()),
) -> ModelArmorGuardrail:
handler = _FakeArmorHandler(responses)
guardrail = ModelArmorGuardrail(
template_id="test-template",
project_id="test-project",
location="us-central1",
guardrail_name="model-armor-logging",
event_hook=GuardrailEventHooks.logging_only,
async_handler=handler,
access_token_provider=_async_token_provider,
)
return guardrail
def _logged_kwargs() -> dict[str, object]:
return {
"model": "gpt-4o",
"messages": [{"role": "user", "content": "hi"}],
"litellm_call_id": "call-1",
"litellm_params": {"metadata": {}},
"optional_params": {},
"standard_logging_object": {"guardrail_information": None},
}
def _chat_response(text: str) -> litellm.ModelResponse:
return litellm.ModelResponse(
choices=[
litellm.types.utils.Choices(
message=litellm.types.utils.Message(role="assistant", content=text)
)
]
)
def _stream_chunk(text: str) -> litellm.ModelResponseStream:
return litellm.ModelResponseStream(
choices=[
litellm.types.utils.StreamingChoices(
delta=litellm.types.utils.Delta(content=text)
)
]
)
def _metadata_entries(kwargs: dict[str, object]) -> list[dict[str, object]]:
standard_logging_object = cast(dict[str, object], kwargs["standard_logging_object"])
entries = standard_logging_object.get("guardrail_information") or []
return cast(list[dict[str, object]], entries)
def test_logging_only_mode_is_accepted_and_keeps_native_hooks():
guardrail = _logging_only_guardrail()
assert guardrail.event_hook == GuardrailEventHooks.logging_only
assert guardrail.use_native_lifecycle_hooks is True
assert GuardrailEventHooks.logging_only in ModelArmorGuardrail.get_supported_event_hooks()
post_call_guardrail = ModelArmorGuardrail(
template_id="test-template",
project_id="test-project",
location="us-central1",
guardrail_name="model-armor-post",
event_hook=GuardrailEventHooks.post_call,
)
assert post_call_guardrail._deployment_hook_target() is post_call_guardrail
@pytest.mark.asyncio
async def test_logging_only_stream_yields_chunks_without_waiting_for_scan():
"""A logging_only guardrail must pass stream chunks straight through; the scan happens
afterwards on the assembled response via async_logging_hook."""
guardrail = _logging_only_guardrail(
[_clean_armor_response(), _clean_armor_response()]
)
handler = cast(_FakeArmorHandler, guardrail.async_handler)
handler.raise_on_call = AssertionError("logging_only must not scan the stream")
produced = 0
async def gen():
nonlocal produced
for i in range(3):
produced += 1
yield _stream_chunk(f"chunk-{i} ")
hook_iter = guardrail.async_post_call_streaming_iterator_hook(
user_api_key_dict=UserAPIKeyAuth(),
response=gen(),
request_data={"metadata": {}, "guardrails": ["model-armor-logging"]},
)
first = await hook_iter.__anext__()
assert produced == 1
chunks = [first]
async for chunk in hook_iter:
chunks.append(chunk)
assert len(chunks) == 3
assert handler.calls == []
handler.raise_on_call = None
response = _chat_response("all clear")
kwargs = _logged_kwargs()
out_kwargs, out_result = await guardrail.async_logging_hook(
kwargs=kwargs, result=response, call_type="acompletion"
)
assert out_result is response
entries = _metadata_entries(out_kwargs)
assert len(entries) >= 1
entry = entries[-1]
assert entry["guardrail_status"] == "success"
assert entry["guardrail_mode"] == "logging_only"
assert entry["guardrail_provider"] == "model_armor"
@pytest.mark.asyncio
async def test_logging_only_records_flagged_verdict_without_altering_response():
guardrail = _logging_only_guardrail(
[_flagged_armor_response(), _flagged_armor_response()]
)
response = _chat_response("flagged output")
kwargs = _logged_kwargs()
out_kwargs, out_result = await guardrail.async_logging_hook(
kwargs=kwargs, result=response, call_type="acompletion"
)
assert out_result is response
entries = _metadata_entries(out_kwargs)
assert entries[-1]["guardrail_status"] == "guardrail_flagged"
assert entries[-1]["guardrail_mode"] == "logging_only"
@pytest.mark.asyncio
async def test_logging_only_records_model_armor_api_error():
guardrail = _logging_only_guardrail(
[
ModelArmorAPIError("Model Armor API error (upstream 500)"),
ModelArmorAPIError("Model Armor API error (upstream 500)"),
]
)
response = _chat_response("some output")
kwargs = _logged_kwargs()
out_kwargs, out_result = await guardrail.async_logging_hook(
kwargs=kwargs, result=response, call_type="acompletion"
)
assert out_result is response
entries = _metadata_entries(out_kwargs)
assert entries[-1]["guardrail_status"] == "guardrail_failed_to_respond"
@pytest.mark.asyncio
async def test_logging_only_scans_assembled_responses_api_stream():
"""The terminal ResponseCompletedEvent is an envelope; the scan must run on the
assembled ResponsesAPIResponse kept in kwargs."""
from openai.types.responses import ResponseOutputMessage, ResponseOutputText
from litellm.types.llms.openai import (
ResponseCompletedEvent,
ResponsesAPIResponse,
ResponsesAPIStreamEvents,
)
assembled = ResponsesAPIResponse(
id="resp-1",
created_at=1700000000,
output=[
ResponseOutputMessage(
id="msg-1",
type="message",
role="assistant",
status="completed",
content=[
ResponseOutputText(
annotations=[], text="assembled output text", type="output_text"
)
],
)
],
)
event = ResponseCompletedEvent(
type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, response=assembled
)
guardrail = _logging_only_guardrail()
kwargs = _logged_kwargs()
del kwargs["messages"]
kwargs["input"] = "hello"
kwargs["async_complete_streaming_response"] = assembled
out_kwargs, _ = await guardrail.async_logging_hook(
kwargs=kwargs, result=event, call_type="aresponses"
)
handler = cast(_FakeArmorHandler, guardrail.async_handler)
response_scans = [call for call in handler.calls if "modelResponseData" in call]
assert response_scans, "expected a model_response scan of the assembled response"
assert "assembled output text" in response_scans[0]["modelResponseData"]["text"]
assert _metadata_entries(out_kwargs)
@pytest.mark.asyncio
async def test_logging_only_scans_anthropic_messages_model_response():
"""/v1/messages logs a ModelResponse; the output scan must extract the assistant text."""
guardrail = _logging_only_guardrail()
kwargs = _logged_kwargs()
kwargs["messages"] = [{"role": "user", "content": [{"type": "text", "text": "hi"}]}]
response = _chat_response("anthropic assembled text")
out_kwargs, out_result = await guardrail.async_logging_hook(
kwargs=kwargs, result=response, call_type="anthropic_messages"
)
assert out_result is response
handler = cast(_FakeArmorHandler, guardrail.async_handler)
response_scans = [call for call in handler.calls if "modelResponseData" in call]
assert response_scans
assert "anthropic assembled text" in response_scans[0]["modelResponseData"]["text"]
assert _metadata_entries(out_kwargs)
@pytest.mark.asyncio
async def test_logging_only_skips_output_scan_when_no_assembled_response():
guardrail = _logging_only_guardrail()
kwargs = _logged_kwargs()
await guardrail.async_logging_hook(kwargs=kwargs, result=None, call_type="acompletion")
handler = cast(_FakeArmorHandler, guardrail.async_handler)
assert all("modelResponseData" not in call for call in handler.calls)
@pytest.mark.asyncio
async def test_native_post_call_mode_ignores_logging_hook():
handler = _FakeArmorHandler([_clean_armor_response()])
guardrail = ModelArmorGuardrail(
template_id="test-template",
project_id="test-project",
location="us-central1",
guardrail_name="model-armor-post",
event_hook=GuardrailEventHooks.post_call,
async_handler=handler,
access_token_provider=_async_token_provider,
)
response = _chat_response("some output")
kwargs = _logged_kwargs()
out_kwargs, out_result = await guardrail.async_logging_hook(
kwargs=kwargs, result=response, call_type="acompletion"
)
assert out_kwargs is kwargs
assert out_result is response
assert handler.calls == []
@pytest.mark.asyncio
async def test_apply_guardrail_records_flagged_without_raising():
guardrail = _logging_only_guardrail([_flagged_armor_response()])
request_data = {"metadata": {}}
inputs = {"texts": ["forbidden output"]}
result = await guardrail.apply_guardrail(
inputs=inputs,
request_data=request_data,
input_type="response",
)
assert result == inputs
entries = request_data["metadata"]["standard_logging_guardrail_information"]
assert entries[-1]["guardrail_status"] == "guardrail_flagged"
@pytest.mark.asyncio
async def test_logging_only_records_transport_error():
guardrail = _logging_only_guardrail([httpx.ConnectError("boom"), httpx.ConnectError("boom")])
response = _chat_response("some output")
kwargs = _logged_kwargs()
out_kwargs, out_result = await guardrail.async_logging_hook(
kwargs=kwargs, result=response, call_type="acompletion"
)
assert out_result is response
entries = _metadata_entries(out_kwargs)
failed = [e for e in entries if e["guardrail_status"] == "guardrail_failed_to_respond"]
assert failed
assert all(e["guardrail_provider"] == "model_armor" for e in failed)
@pytest.mark.asyncio
async def test_logging_only_flagged_prompt_still_scans_response():
"""A flagged input scan must not abort the output scan; both verdicts are recorded."""
guardrail = _logging_only_guardrail(
[_flagged_armor_response(), _flagged_armor_response()]
)
response = _chat_response("flagged output")
kwargs = _logged_kwargs()
out_kwargs, _ = await guardrail.async_logging_hook(
kwargs=kwargs, result=response, call_type="acompletion"
)
handler = cast(_FakeArmorHandler, guardrail.async_handler)
sources = ["user_prompt" if "userPromptData" in call else "model_response" for call in handler.calls]
assert sources == ["user_prompt", "model_response"]
entries = _metadata_entries(out_kwargs)
flagged = [e for e in entries if e["guardrail_status"] == "guardrail_flagged"]
assert len(flagged) == 2
@pytest.mark.asyncio
async def test_apply_guardrail_raises_on_flagged_when_not_logging_only():
"""The /guardrails/apply_guardrail endpoint calls apply_guardrail directly; a
non-logging_only instance must signal the block so flagged text is not returned as clean."""
handler = _FakeArmorHandler([_flagged_armor_response()])
guardrail = ModelArmorGuardrail(
template_id="test-template",
project_id="test-project",
location="us-central1",
guardrail_name="model-armor-pre",
event_hook=GuardrailEventHooks.pre_call,
async_handler=handler,
access_token_provider=_async_token_provider,
)
request_data = {"metadata": {}}
with pytest.raises(HTTPException) as exc_info:
await guardrail.apply_guardrail(
inputs={"texts": ["forbidden prompt"]},
request_data=request_data,
input_type="request",
)
assert exc_info.value.status_code == 400
entries = request_data["metadata"]["standard_logging_guardrail_information"]
flagged = [e for e in entries if e["guardrail_status"] == "guardrail_flagged"]
assert len(flagged) == 1

92
tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py
View file

@ -497,6 +497,98 @@ async def test_file_sanitization_modify_can_rewrite_when_blocking_disabled(monke
assert base64.b64decode(result["file"]["data"]) == b"name,email\nAlice,[REDACTED]\n"
@pytest.mark.asyncio
async def test_file_sanitization_keeps_polling_through_queued_statuses(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key")
monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security")
guardrail = PromptSecurityGuardrail(guardrail_name="test-guard", event_hook="pre_call", default_on=True)
guardrail.poll_interval = 0
upload_response = Response(
json={"jobId": "queued-job"},
status_code=200,
request=Request(method="POST", url="https://test.prompt.security/api/sanitizeFile"),
)
poll_request = Request(method="GET", url="https://test.prompt.security/api/sanitizeFile")
poll_responses = [
Response(json={"status": "created"}, status_code=200, request=poll_request),
Response(json={"status": "in progress"}, status_code=200, request=poll_request),
Response(
json={"status": "done", "content": "clean", "metadata": {"action": "allow", "violations": []}},
status_code=200,
request=poll_request,
),
]
with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=upload_response)):
with patch.object(guardrail.async_handler, "get", AsyncMock(side_effect=poll_responses)) as poll_mock:
result = await guardrail.sanitize_file_content(b"image-content", "image.png")
assert poll_mock.await_count == 3
assert result["action"] == "allow"
assert result["content"] == "clean"
@pytest.mark.asyncio
async def test_file_sanitization_never_finishing_job_times_out(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key")
monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security")
guardrail = PromptSecurityGuardrail(
guardrail_name="test-guard", event_hook="pre_call", default_on=True, file_sanitization_fail_open=False
)
guardrail.poll_interval = 0
guardrail.max_poll_attempts = 3
upload_response = Response(
json={"jobId": "stuck-job"},
status_code=200,
request=Request(method="POST", url="https://test.prompt.security/api/sanitizeFile"),
)
poll_response = Response(
json={"status": "created"},
status_code=200,
request=Request(method="GET", url="https://test.prompt.security/api/sanitizeFile"),
)
with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=upload_response)):
with patch.object(guardrail.async_handler, "get", AsyncMock(return_value=poll_response)) as poll_mock:
with pytest.raises(HTTPException) as exc_info:
await guardrail.sanitize_file_content(b"file-content", "document.pdf")
assert poll_mock.await_count == 3
assert exc_info.value.status_code == 408
assert exc_info.value.detail == "File sanitization timeout"
@pytest.mark.asyncio
@pytest.mark.parametrize("poll_body", [{"status": "failed"}, {}])
async def test_file_sanitization_terminal_failure_does_not_fail_open(monkeypatch: pytest.MonkeyPatch, poll_body):
monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key")
monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security")
guardrail = PromptSecurityGuardrail(guardrail_name="test-guard", event_hook="pre_call", default_on=True)
guardrail.poll_interval = 0
upload_response = Response(
json={"jobId": "failed-job"},
status_code=200,
request=Request(method="POST", url="https://test.prompt.security/api/sanitizeFile"),
)
poll_response = Response(
json=poll_body,
status_code=200,
request=Request(method="GET", url="https://test.prompt.security/api/sanitizeFile"),
)
with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=upload_response)):
with patch.object(guardrail.async_handler, "get", AsyncMock(return_value=poll_response)) as poll_mock:
with pytest.raises(HTTPException) as exc_info:
await guardrail.sanitize_file_content(b"file-content", "document.pdf")
assert poll_mock.await_count == 1
assert exc_info.value.status_code == 500
assert exc_info.value.detail == f"Unexpected sanitization status: {poll_body.get('status')}"
@pytest.mark.asyncio
@pytest.mark.parametrize(
"timeout",

245
tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py
View file

@ -6284,3 +6284,248 @@ async def test_an_open_circuit_breaker_reads_the_sliding_window_locally_without_
assert isinstance(values, list)
assert [record.getMessage() for record in caplog.records if record.levelno >= logging.WARNING] == []
assert any("circuit breaker is open" in record.getMessage() for record in caplog.records)
@pytest.mark.parametrize(
"limits, request_data, counter_scope",
[
({"rpm_limit": 1}, {}, "api_key"),
({"user_id": "u", "user_rpm_limit": 1}, {}, "user"),
({"team_id": "t", "team_rpm_limit": 1}, {}, "team"),
(
{"team_id": "t", "user_id": "u", "team_member_rpm_limit": 1},
{},
"team_member",
),
({"end_user_id": "e", "end_user_rpm_limit": 1}, {}, "end_user"),
(
{"metadata": {"model_rpm_limit": {"test-model": 1}}},
{},
"model_per_key",
),
(
{"metadata": {"tag_rpm_limit": {"test-tag": 1}}},
{"metadata": {"tags": ["test-tag"]}},
"tag_per_key",
),
(
{
"team_id": "t",
"metadata": {"model_rpm_limit": {"test-model": 100}},
"team_metadata": {"model_rpm_limit": {"test-model": 1}},
},
{},
"model_per_team",
),
(
{"project_id": "p", "project_metadata": {"model_rpm_limit": {"test-model": 1}}},
{},
"model_per_project",
),
({"org_id": "o", "organization_rpm_limit": 1}, {}, "organization"),
(
{"org_id": "o", "organization_metadata": {"model_rpm_limit": {"test-model": 1}}},
{},
"model_per_organization",
),
],
)
@pytest.mark.parametrize("request_kind", ["count", "generation"])
@pytest.mark.asyncio
async def test_request_capacity_enforces_shared_rpm_scopes(
limits, request_data, counter_scope, request_kind
):
cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache))
auth = UserAPIKeyAuth(api_key=hash_token("sk-count-rpm"), **limits)
async def request():
if request_kind == "generation":
await handler.async_pre_call_hook(
user_api_key_dict=auth,
cache=cache,
data={**request_data, "model": "test-model"},
call_type="acompletion",
)
return
async with handler.request_capacity(auth, "test-model", request_data=request_data):
pass
await request()
with pytest.raises(HTTPException) as exc:
await request()
assert exc.value.status_code == 429
assert counter_scope in str(exc.value.detail)
@pytest.mark.asyncio
async def test_request_capacity_keeps_dynamic_rpm_policy(monkeypatch):
import litellm.proxy.proxy_server as proxy_server
router = Router(model_list=[{
"model_name": "test-model",
"litellm_params": {"model": "openai/gpt-test", "api_key": "test-key"},
"model_info": {"id": "test-deployment"},
}])
monkeypatch.setattr(proxy_server, "llm_router", router)
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache()))
auth = UserAPIKeyAuth(
api_key=hash_token("sk-count-dynamic"),
rpm_limit=1,
metadata={"rpm_limit_type": "dynamic"},
)
for _ in range(2):
async with handler.request_capacity(auth, "test-model"):
pass
router.cache.set_cache("test-deployment:fails", 100, ttl=60, local_only=True)
async with handler.request_capacity(auth, "test-model"):
pass
with pytest.raises(HTTPException) as exc:
async with handler.request_capacity(auth, "test-model"):
pytest.fail("dynamic RPM must enforce after deployment failures")
assert exc.value.status_code == 429
@pytest.mark.asyncio
async def test_request_capacity_skips_tokens_and_preserves_parent_stash():
cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache))
auth = UserAPIKeyAuth(
api_key=hash_token("sk-count-tpm"),
rpm_limit=5,
tpm_limit=1,
max_parallel_requests=1,
project_id="p",
project_metadata={
"model_tpm_limit": {"test-model": 1},
"model_itpm_limit": {"test-model": 1},
"model_otpm_limit": {"test-model": 1},
},
)
token_scopes = (
("api_key", auth.api_key),
("model_per_project", "p:test-model"),
("model_per_project_itpm", "p:test-model"),
("model_per_project_otpm", "p:test-model"),
)
for scope, value in token_scopes:
token_key = handler.create_rate_limit_keys(scope, value, "tokens")
await cache.async_set_cache(token_key, 100, ttl=60)
await cache.async_set_cache(f"{{{scope}:{value}}}:window", int(time.time()), ttl=60)
parent = get_or_create_request_stash()
parent.reserved_tokens = 123
parent.parallel_slot = ParallelSlotAcquisition(slot_id="parent", counter_keys=["parent-gauge"])
for _ in range(2):
async with handler.request_capacity(auth, "test-model"):
assert get_request_stash() is parent
assert parent.parallel_slot["slot_id"] == "parent"
assert parent.reserved_tokens == 123
for scope, value in token_scopes:
assert await cache.async_get_cache(handler.create_rate_limit_keys(scope, value, "tokens")) == 100
@pytest.mark.parametrize("exit_mode", ["success", "failure", "cancel"])
@pytest.mark.asyncio
async def test_request_capacity_releases_exact_parallel_slot(exit_mode):
cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache))
auth = UserAPIKeyAuth(api_key=hash_token("sk-count-parallel"), max_parallel_requests=1)
entered = asyncio.Event()
finish = asyncio.Event()
async def provider():
async with handler.request_capacity(auth, "test-model"):
entered.set()
await finish.wait()
if exit_mode == "failure":
raise RuntimeError("provider failed")
task = asyncio.create_task(provider())
await asyncio.wait_for(entered.wait(), timeout=2)
try:
for _ in range(2):
with pytest.raises(HTTPException) as exc:
async with handler.request_capacity(auth, "test-model"):
pytest.fail("rejected request freed the occupied slot")
assert exc.value.status_code == 429
finally:
if exit_mode == "cancel":
task.cancel()
else:
finish.set()
if exit_mode == "success":
await task
else:
with pytest.raises(asyncio.CancelledError if exit_mode == "cancel" else RuntimeError):
await task
async with handler.request_capacity(auth, "test-model"):
pass
class _DelayedCapacityUsageCache:
def __init__(self):
self.delegate = InternalUsageCache(DualCache())
self.dual_cache = self.delegate.dual_cache
self.acquired = asyncio.Event()
self.finish_admission = asyncio.Event()
self.releasing = asyncio.Event()
self.finish_release = asyncio.Event()
async def async_get_cache(self, *args, **kwargs):
return await self.delegate.async_get_cache(*args, **kwargs)
async def async_batch_get_cache(self, *args, **kwargs):
return await self.delegate.async_batch_get_cache(*args, **kwargs)
async def async_set_cache(self, key, value, **kwargs):
await self.delegate.async_set_cache(key=key, value=value, **kwargs)
if not key.endswith(":max_parallel_requests"):
return
if value:
self.acquired.set()
await self.finish_admission.wait()
else:
self.releasing.set()
await self.finish_release.wait()
@pytest.mark.asyncio
async def test_request_capacity_finishes_admission_and_release_despite_repeated_cancel():
cache = _DelayedCapacityUsageCache()
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=cache)
auth = UserAPIKeyAuth(api_key=hash_token("sk-count-cancel-admission"), max_parallel_requests=1)
async def provider():
async with handler.request_capacity(auth, "test-model"):
pytest.fail("cancelled admission entered provider body")
task = asyncio.create_task(provider())
await asyncio.wait_for(cache.acquired.wait(), timeout=2)
task.cancel()
await asyncio.sleep(0)
cache.finish_admission.set()
await asyncio.wait_for(cache.releasing.wait(), timeout=2)
task.cancel()
await asyncio.sleep(0)
task.cancel()
await asyncio.sleep(0)
assert not task.done()
cache.finish_release.set()
with pytest.raises(asyncio.CancelledError):
await asyncio.wait_for(task, timeout=2)
async with handler.request_capacity(auth, "test-model"):
pass
@pytest.mark.asyncio
async def test_request_capacity_rejection_keeps_existing_redis_mirror():
cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache))
auth = UserAPIKeyAuth(api_key=hash_token("sk-count-mirror"), max_parallel_requests=1)
counter_key = handler.create_rate_limit_keys("api_key", auth.api_key, "max_parallel_requests")
await cache.async_set_cache(counter_key, 1, ttl=60, local_only=True)
for _ in range(2):
with pytest.raises(HTTPException) as exc:
async with handler.request_capacity(auth, "test-model"):
pytest.fail("rejection released another request's mirrored slot")
assert exc.value.status_code == 429
assert await cache.async_get_cache(counter_key, local_only=True) == 1

300
tests/test_litellm/proxy/hooks/test_prompt_cache_observer.py Normal file
View file

@ -0,0 +1,300 @@
import asyncio
import json
import time
from datetime import datetime
import httpx
import pytest
import litellm
from litellm.caching.dual_cache import DualCache
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.llms.anthropic.prompt_cache_prediction import cache_scope, parse_prompt
from litellm.proxy.hooks.prompt_cache_prediction import (
PromptCacheObserver,
lookup,
)
from litellm.proxy.utils import InternalUsageCache
from litellm.types.utils import ModelResponse
MODEL = "claude-sonnet-5"
CALLER = "a" * 64
DEPLOYMENT = "native-deployment"
KEY = "test-provider-key"
def body(ttl="5m", texts=("private cache prefix",)):
return {
"model": MODEL,
"max_tokens": 2,
"system": "private system instructions",
"tools": [{"name": "lookup", "input_schema": {"type": "object"}}],
"messages": [{"role": "user", "content": [
{"type": "text", "text": text, **(
{"cache_control": {"type": "ephemeral", "ttl": ttl}}
if index == len(texts) - 1 else {}
)}
for index, text in enumerate(texts)
]}],
}
def usage(ttl="5m", read=100, write=200):
return {
"input_tokens": 11,
"output_tokens": 2,
"cache_read_input_tokens": read,
"cache_creation_input_tokens": write,
"cache_creation": {
"ephemeral_5m_input_tokens": write if ttl == "5m" else 0,
"ephemeral_1h_input_tokens": write if ttl == "1h" else 0,
},
}
def event(request_body, started=1000.0, headers=None, **overrides):
request = httpx.Request(
"POST", "https://api.anthropic.com/v1/messages", json=request_body,
headers={"x-api-key": KEY, "anthropic-version": "2023-06-01", **(headers or {})},
)
return {
"call_type": "anthropic_messages",
"custom_llm_provider": "anthropic",
"cache_hit": False,
"httpx_response": httpx.Response(200, request=request),
"first_api_call_start_time": datetime.fromtimestamp(started),
"standard_logging_object": {
"status": "success", "model_id": DEPLOYMENT,
"metadata": {"user_api_key_hash": CALLER},
},
**overrides,
}
async def observe(cache, request_body=None, native_usage=None, now=1010.0, **overrides):
observer = PromptCacheObserver(InternalUsageCache(dual_cache=cache), clock=lambda: now)
response = ModelResponse(
model=MODEL,
usage=AnthropicConfig().calculate_usage(native_usage or usage(), reasoning_content=None),
)
await observer.async_log_success_event(
event(request_body or body(), **overrides), response,
datetime.fromtimestamp(now), datetime.fromtimestamp(now),
)
def scope(**overrides):
return cache_scope(**{
"caller_key_hash": CALLER, "deployment_id": DEPLOYMENT,
"provider_key": KEY, "model": MODEL, **overrides,
})
@pytest.mark.parametrize("ttl,expires", [("5m", 1300), ("1h", 4600)])
@pytest.mark.asyncio
async def test_observed_cache_count_and_request_start_expiry_survive_as_stale(ttl, expires):
cache = DualCache()
request_body = body(ttl=ttl)
await observe(cache, request_body, usage(ttl=ttl))
prefix = parse_prompt(request_body)
observed = await lookup(cache, scope(), prefix, now=1200)
assert observed.cached_tokens == 300
assert observed.observed_at == 1010
assert observed.expires_at == expires
assert await lookup(cache, scope(), prefix, now=expires) == observed
saved = json.dumps(cache.in_memory_cache.cache_dict)
assert "private cache prefix" not in saved
assert "private system instructions" not in saved
assert KEY not in saved
assert CALLER not in saved
@pytest.mark.parametrize("changed", [
{"caller_key_hash": "b" * 64}, {"deployment_id": "other"},
{"provider_key": "rotated"}, {"model": "claude-opus-5"},
{"anthropic_version": "different"},
])
@pytest.mark.asyncio
async def test_cache_evidence_is_isolated_by_every_scope_dimension(changed):
cache = DualCache()
await observe(cache)
assert await lookup(cache, scope(**changed), parse_prompt(body()), now=1010) is None
@pytest.mark.asyncio
async def test_append_only_prefix_finds_prior_evidence_but_edit_or_context_change_does_not():
cache = DualCache()
await observe(cache)
extended = parse_prompt(body(texts=("private cache prefix", "new turn")))
prior = await lookup(cache, scope(), extended, now=1010)
assert prior.cached_tokens == 300
assert prior.fingerprint != extended.fingerprint
for changed in (
body(texts=("edited prefix", "new turn")),
{**body(), "system": "different system"},
{**body(), "tools": [{"name": "other", "input_schema": {"type": "object"}}]},
body(ttl="1h"),
):
assert await lookup(cache, scope(), parse_prompt(changed), now=1010) is None
outside_lookback = parse_prompt(body(texts=("private cache prefix", *[str(i) for i in range(20)])))
assert await lookup(cache, scope(), outside_lookback, now=1010) is None
@pytest.mark.parametrize("change", [
{"thinking": {"type": "enabled", "budget_tokens": 1024}},
{"tool_choice": {"type": "auto"}},
{"cache_control": {"type": "ephemeral"}},
{"tools": [{"type": "web_search_20250305", "name": "web_search"}]},
{"system": [{"type": "text", "text": "system", "cache_control": {"type": "ephemeral"}}]},
{"messages": [{"role": "user", "content": [{"type": "image", "source": {}}]}]},
{"messages": [{"role": "user", "content": "no breakpoint"}]},
])
def test_unsupported_or_ambiguous_shapes_have_no_cache_identity(change):
assert parse_prompt({**body(), **change}) is None
duplicate = body()
duplicate["messages"][0]["content"].append(duplicate["messages"][0]["content"][0])
assert parse_prompt(duplicate) is None
@pytest.mark.parametrize("overrides", [
{"cache_hit": True}, {"call_type": "completion"},
{"custom_llm_provider": "bedrock"}, {"stream": True},
{"headers": {"anthropic-beta": "unverified-feature"}},
{"headers": {"x-custom-header": "unverified"}},
{"standard_logging_object": {"status": "success", "model_id": DEPLOYMENT, "metadata": {}}},
])
@pytest.mark.asyncio
async def test_unverified_source_never_creates_observations(overrides):
cache = DualCache()
await observe(cache, **overrides)
assert await lookup(cache, scope(), parse_prompt(body()), now=1010) is None
@pytest.mark.parametrize("native_usage", [
usage(write=0),
{**usage(), "cache_creation": None},
{**usage(), "cache_creation": {"ephemeral_5m_input_tokens": 199, "ephemeral_1h_input_tokens": 0}},
usage(ttl="1h"),
{**usage(), "cache_creation_input_tokens": -200},
])
@pytest.mark.asyncio
async def test_missing_or_contradictory_telemetry_cannot_create_observations(native_usage):
cache = DualCache()
await observe(cache, native_usage=native_usage)
assert await lookup(cache, scope(), parse_prompt(body()), now=1010) is None
@pytest.mark.asyncio
async def test_pure_read_refresh_requires_prior_matching_evidence():
cache = DualCache()
await observe(cache, native_usage=usage(read=300, write=0))
assert await lookup(cache, scope(), parse_prompt(body()), now=1010) is None
await observe(cache)
await observe(cache, native_usage=usage(read=300, write=0), started=1100, now=1110)
assert (await lookup(cache, scope(), parse_prompt(body()), now=1110)).expires_at == 1400
class RecordingObserver(PromptCacheObserver):
def __init__(self, cache):
super().__init__(InternalUsageCache(dual_cache=cache))
self.finished = asyncio.Event()
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
await super().async_log_success_event(kwargs, response_obj, start_time, end_time)
self.finished.set()
def native_response():
return {
"id": "msg_prediction", "type": "message", "role": "assistant", "model": MODEL,
"content": [{"type": "text", "text": "ok"}], "stop_reason": "end_turn",
"stop_sequence": None, "usage": usage(ttl="1h"),
}
def stream_response(completed, provider_error=False):
response = native_response()
events = [
{"type": "message_start", "message": {**response, "content": [], "stop_reason": None}},
{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "ok"}},
{"type": "content_block_stop", "index": 0},
{"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 2}},
]
if completed:
events.append({"type": "message_stop"})
if provider_error:
events.append({"type": "error", "error": {"type": "overloaded_error", "message": "temporary failure"}})
return "".join(f"event: {item['type']}\ndata: {json.dumps(item)}\n\n" for item in events)
class TransportChunks(httpx.AsyncByteStream):
def __init__(self, payload, chunk_size, fragment_error_only=False):
self.payload = payload.encode()
self.chunk_size = chunk_size or len(self.payload)
self.prefix_length = self.payload.index(b"event: error") if fragment_error_only else 0
async def __aiter__(self):
if self.prefix_length:
yield self.payload[:self.prefix_length]
for offset in range(self.prefix_length, len(self.payload), self.chunk_size):
yield self.payload[offset:offset + self.chunk_size]
@pytest.mark.parametrize("stream,completed,provider_error,transport", [
(False, True, False, "whole"),
(True, True, False, "whole"),
(True, False, False, "whole"),
(True, True, True, "whole"),
(True, True, False, "fragmented"),
(True, False, False, "fragmented"),
(True, True, True, "fragmented"),
(True, True, True, "fragmented_error"),
(True, True, False, "unterminated"),
])
@pytest.mark.asyncio
async def test_native_production_callback_records_only_completed_wire_requests(stream, completed, provider_error, transport):
cache = DualCache()
observer = RecordingObserver(cache)
litellm.logging_callback_manager.add_litellm_callback(observer)
def provider(request):
if stream:
payload = stream_response(completed, provider_error)
if transport == "unterminated":
payload = payload.removesuffix("\n\n")
return httpx.Response(
200, request=request, headers={"content-type": "text/event-stream"},
stream=TransportChunks(
payload, 1 if transport.startswith("fragmented") else None,
fragment_error_only=transport == "fragmented_error",
),
)
return httpx.Response(200, request=request, json=native_response())
client = AsyncHTTPHandler()
await client.client.aclose()
client.client = httpx.AsyncClient(transport=httpx.MockTransport(provider))
try:
request_body = body(ttl="1h")
before = time.time()
result = await litellm.anthropic_messages(
**{**request_body, "model": f"anthropic/{MODEL}"},
api_key=KEY, client=client, stream=stream, model_info={"id": DEPLOYMENT},
litellm_metadata={"user_api_key_hash": CALLER, "model_info": {"id": DEPLOYMENT}},
)
if stream:
async for _ in result:
pass
await asyncio.wait_for(observer.finished.wait(), timeout=5)
found = await lookup(cache, scope(), parse_prompt(request_body))
if completed and not provider_error and transport != "unterminated":
assert found is not None
assert found.cached_tokens == 300
assert before + 3600 <= found.expires_at <= time.time() + 3600
else:
assert found is None
finally:
litellm.logging_callback_manager.remove_callback_from_all_lists(observer)
await client.client.aclose()

698
tests/test_litellm/proxy/management_endpoints/test_prompt_cache_prediction.py Normal file
View file

@ -0,0 +1,698 @@
import asyncio
import time
from collections.abc import Iterator, Mapping
from dataclasses import dataclass
from typing import Final, Literal
import httpx
import pytest
from fastapi import FastAPI, Request
from pydantic import JsonValue
import litellm
from litellm.caching.caching import DualCache
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
from litellm.proxy.common_utils.http_parsing_utils import _read_request_body, _safe_set_request_parsed_body
from litellm.proxy.hooks.parallel_request_limiter_v3 import _PROXY_MaxParallelRequestsHandler_v3
from litellm.llms.anthropic.prompt_cache_prediction import PromptPrefix, cache_scope, parse_prompt
from litellm.proxy.hooks.prompt_cache_prediction import (
CacheObservation,
_cache_key,
)
from litellm.proxy.management_endpoints import prompt_cache_prediction as endpoint
from litellm.proxy.utils import InternalUsageCache
from litellm.types.management_endpoints.prompt_cache_prediction import CachePredictionResponse
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
_PROVIDER_KEY: Final = "cache-prediction-test-provider-key"
_CALLER: Final = "cache-prediction-test-caller-hash"
@pytest.fixture(autouse=True)
def anthropic_endpoint_environment(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.delenv("ANTHROPIC_API_BASE", raising=False)
monkeypatch.delenv("ANTHROPIC_BASE_URL", raising=False)
def _body(ttl: str = "5m", *, extended: bool = False) -> dict[str, JsonValue]:
blocks: Final[list[JsonValue]] = [
{"type": "text", "text": "Stable context"},
*([{"type": "text", "text": "Appended context"}] if extended else []),
]
return {
"max_tokens": 10,
"system": "Follow the project conventions",
"messages": [
{
"role": "user",
"content": [
*blocks[:-1],
{**blocks[-1], "cache_control": {"type": "ephemeral", "ttl": ttl}},
{"type": "text", "text": "Follow-up question"},
],
}
],
}
def _prefix(body: Mapping[str, JsonValue]) -> PromptPrefix:
prefix: Final = parse_prompt(body)
assert prefix is not None
return prefix
def _deployment(
deployment_id: str = "sonnet",
model: str = "claude-sonnet-5",
*,
team_id: str | None = None,
api_base: str | None = None,
) -> Deployment:
return Deployment(
model_name=deployment_id,
litellm_params=LiteLLM_Params(model=f"anthropic/{model}", api_key=_PROVIDER_KEY, api_base=api_base),
model_info=ModelInfo(id=deployment_id, team_id=team_id),
)
@dataclass(frozen=True)
class Counts:
total: int | None = 6_000
prefix: int | None = 5_000
async def __call__(self, model: str, api_key: str, body: Mapping[str, JsonValue]) -> int | None:
assert api_key == _PROVIDER_KEY
assert model.startswith("claude-")
return self.total if "max_tokens" in body else self.prefix
async def _observe(
cache: DualCache,
body: Mapping[str, JsonValue],
*,
deployment_id: str = "sonnet",
model: str = "claude-sonnet-5",
cached_tokens: int = 5_000,
expired: bool = False,
caller: str = _CALLER,
) -> None:
prefix: Final = _prefix(body)
now: Final = time.time()
observation: Final = CacheObservation(
fingerprint=prefix.fingerprint,
cached_tokens=cached_tokens,
observed_at=now - 400 if expired else now - 10,
expires_at=now - 100 if expired else now + 290,
)
scope: Final = cache_scope(caller, deployment_id, _PROVIDER_KEY, model)
await cache.async_set_cache(_cache_key(scope, prefix.fingerprint), observation.model_dump_json(), ttl=3_600)
@pytest.mark.asyncio
@pytest.mark.parametrize(("ttl", "cold_cost"), [("5m", 0.0145), ("1h", 0.022)])
async def test_unobserved_cache_prices_cold_and_warm_bounds(ttl: str, cold_cost: float) -> None:
body: Final = _body(ttl)
arm: Final = await endpoint.predict_arm(_deployment(), body, _prefix(body), _CALLER, DualCache(), Counts())
assert arm.cache_state == "unknown"
assert arm.reason == "no_compatible_observation"
assert arm.evidence is None
assert arm.estimate is not None and arm.cold is not None and arm.warm is not None
assert arm.estimate.input_cost == pytest.approx(cold_cost)
assert arm.cold.input_cost == pytest.approx(cold_cost)
assert arm.warm.input_cost == pytest.approx(0.003)
assert arm.cold.tokens.uncached_input_tokens == 1_000
assert arm.cold.tokens.cache_read_input_tokens == 0
assert arm.cold.tokens.cache_creation_5m_input_tokens == (5_000 if ttl == "5m" else 0)
assert arm.cold.tokens.cache_creation_1h_input_tokens == (5_000 if ttl == "1h" else 0)
assert arm.warm.tokens.cache_read_input_tokens == 5_000
@pytest.mark.asyncio
@pytest.mark.parametrize(
("cached_tokens", "warm_cost", "cold_cost"), [(5_400, 0.00228, 0.0147), (4_600, 0.00372, 0.0143)]
)
@pytest.mark.parametrize("expired", [False, True])
async def test_exact_prefix_conserves_total_with_observed_count_in_all_scenarios(
cached_tokens: int, warm_cost: float, cold_cost: float, expired: bool
) -> None:
cache: Final = DualCache()
body: Final = _body()
await _observe(cache, body, cached_tokens=cached_tokens, expired=expired)
arm: Final = await endpoint.predict_arm(_deployment(), body, _prefix(body), _CALLER, cache, Counts())
assert arm.cache_state == ("stale" if expired else "warm")
assert arm.evidence is not None
assert arm.estimate is not None and arm.warm is not None and arm.cold is not None
assert arm.warm.tokens.cache_read_input_tokens == cached_tokens
assert arm.warm.tokens.cache_creation_5m_input_tokens == 0
assert arm.cold.tokens.cache_creation_5m_input_tokens == cached_tokens
assert arm.cold.tokens.cache_read_input_tokens == 0
for scenario in (arm.estimate, arm.cold, arm.warm):
assert scenario.tokens.total_tokens == 6_000
assert scenario.tokens.uncached_input_tokens == 6_000 - cached_tokens
assert arm.warm.input_cost == pytest.approx(warm_cost)
assert arm.cold.input_cost == pytest.approx(cold_cost)
assert arm.estimate.input_cost == pytest.approx(cold_cost if expired else warm_cost)
@pytest.mark.asyncio
async def test_observed_prefix_larger_than_full_request_returns_unknown() -> None:
cache: Final = DualCache()
body: Final = _body()
await _observe(cache, body, cached_tokens=6_001)
arm: Final = await endpoint.predict_arm(_deployment(), body, _prefix(body), _CALLER, cache, Counts())
assert arm.cache_state == "unknown"
assert arm.reason == "inconsistent_prefix_token_count"
assert arm.estimate is None and arm.cold is None and arm.warm is None
@pytest.mark.asyncio
@pytest.mark.parametrize(("ttl", "expected"), [("5m", 0.0053), ("1h", 0.0068)])
async def test_append_only_prefix_reads_old_tokens_and_writes_extension(ttl: str, expected: float) -> None:
cache: Final = DualCache()
await _observe(cache, _body(ttl), cached_tokens=4_000)
body: Final = _body(ttl, extended=True)
arm: Final = await endpoint.predict_arm(_deployment(), body, _prefix(body), _CALLER, cache, Counts())
assert arm.cache_state == "partial"
assert arm.estimate is not None
assert arm.estimate.tokens.cache_read_input_tokens == 4_000
assert arm.estimate.tokens.cache_creation_5m_input_tokens == (1_000 if ttl == "5m" else 0)
assert arm.estimate.tokens.cache_creation_1h_input_tokens == (1_000 if ttl == "1h" else 0)
assert arm.estimate.input_cost == pytest.approx(expected)
@pytest.mark.asyncio
async def test_expired_observation_estimates_a_cold_rebuild() -> None:
cache: Final = DualCache()
body: Final = _body()
await _observe(cache, body, expired=True)
arm: Final = await endpoint.predict_arm(_deployment(), body, _prefix(body), _CALLER, cache, Counts())
assert arm.cache_state == "stale"
assert arm.reason == "observation_expired"
assert arm.evidence is not None and arm.evidence.expires_at < time.time()
assert arm.estimate is not None and arm.cold is not None
assert arm.estimate.tokens.cache_read_input_tokens == 0
assert arm.estimate.tokens.cache_creation_5m_input_tokens == 5_000
assert arm.estimate.input_cost == arm.cold.input_cost
@pytest.mark.asyncio
async def test_below_model_minimum_prices_all_input_as_uncached() -> None:
body: Final = _body()
arm: Final = await endpoint.predict_arm(
_deployment(), body, _prefix(body), _CALLER, DualCache(), Counts(total=1_500, prefix=1_000)
)
assert arm.cache_state == "disabled"
assert arm.reason == "below_cache_minimum"
assert arm.estimate is not None
assert arm.estimate.tokens.uncached_input_tokens == 1_500
assert arm.estimate.tokens.cache_read_input_tokens == 0
assert arm.estimate.tokens.cache_creation_5m_input_tokens == 0
assert arm.estimate.input_cost == pytest.approx(0.003)
@pytest.mark.asyncio
@pytest.mark.parametrize("counts", [Counts(total=None), Counts(prefix=None), Counts(total=4_000)])
async def test_unavailable_or_inconsistent_token_counts_return_null_estimates(counts: Counts) -> None:
body: Final = _body()
arm: Final = await endpoint.predict_arm(_deployment(), body, _prefix(body), _CALLER, DualCache(), counts)
assert arm.cache_state == "unknown"
assert arm.reason == "token_count_unavailable"
assert arm.estimate is None and arm.cold is None and arm.warm is None
@pytest.mark.asyncio
@pytest.mark.parametrize("counts", [Counts(), Counts(total=1_500, prefix=1_000)])
async def test_missing_prices_return_unknown_and_null_estimates(
monkeypatch: pytest.MonkeyPatch, counts: Counts
) -> None:
monkeypatch.setitem(
litellm.model_cost,
"claude-cache-unpriced-5",
{"litellm_provider": "anthropic", "mode": "chat"},
)
body: Final = _body()
arm: Final = await endpoint.predict_arm(
_deployment("cache-prediction-unpriced", "claude-cache-unpriced-5"),
body,
_prefix(body),
_CALLER,
DualCache(),
counts,
)
assert arm.cache_state == "unknown"
assert arm.reason == "pricing_unavailable"
assert arm.estimate is None and arm.cold is None and arm.warm is None
@pytest.mark.asyncio
async def test_custom_api_base_from_environment_returns_unknown_before_counting(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setenv("ANTHROPIC_API_BASE", "https://custom.invalid")
body: Final = _body()
arm: Final = await endpoint.predict_arm(
_deployment(), body, _prefix(body), _CALLER, DualCache(), _unexpected_count
)
assert arm.cache_state == "unknown"
assert arm.reason == "unsupported_provider_endpoint"
assert arm.estimate is None and arm.cold is None and arm.warm is None
@pytest.mark.asyncio
async def test_explicit_official_api_base_overrides_custom_environment(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("ANTHROPIC_API_BASE", "https://custom.invalid")
body: Final = _body()
arm: Final = await endpoint.predict_arm(
_deployment(api_base="https://api.anthropic.com"), body, _prefix(body), _CALLER, DualCache(), Counts()
)
assert arm.cache_state == "unknown"
assert arm.reason == "no_compatible_observation"
assert arm.estimate is not None
assert arm.estimate.input_cost == pytest.approx(0.0145)
@dataclass(frozen=True)
class _ProxyLogging:
internal_usage_cache: InternalUsageCache
parallel_limiter: CustomLogger | None
def get_proxy_hook(self, hook: str) -> CustomLogger | None:
return self.parallel_limiter if hook == "parallel_request_limiter" else None
def _app(
monkeypatch: pytest.MonkeyPatch,
cache: DualCache,
*,
caller: UserAPIKeyAuth | None = None,
current_team: str | None = None,
candidate_team: str | None = None,
counts: endpoint.TokenCounter = Counts(),
limiter: CustomLogger | Literal["default"] | None = "default",
) -> FastAPI:
import litellm.proxy.proxy_server as proxy_server
model_list: Final = [
_deployment("opus", "claude-opus-5", team_id=current_team).model_dump(exclude_unset=True),
_deployment("sonnet", team_id=candidate_team).model_dump(exclude_unset=True),
]
router: Final = litellm.Router(model_list=model_list)
monkeypatch.setattr(proxy_server, "llm_router", router)
monkeypatch.setattr(proxy_server, "llm_model_list", model_list)
monkeypatch.setattr(endpoint, "count_prompt_tokens", counts)
app: Final = FastAPI()
app.include_router(endpoint.router)
app.add_exception_handler(ProxyException, proxy_server.openai_exception_handler)
if caller is not None:
usage_cache: Final = InternalUsageCache(cache)
configured_limiter: Final = (
_PROXY_MaxParallelRequestsHandler_v3(usage_cache) if isinstance(limiter, str) else limiter
)
monkeypatch.setattr(proxy_server, "proxy_logging_obj", _ProxyLogging(usage_cache, configured_limiter))
app.dependency_overrides[endpoint.user_api_key_auth] = lambda: caller
return app
async def _post(
app: FastAPI,
body: Mapping[str, JsonValue],
*,
current_deployment_id: str = "opus",
candidate_deployment_id: str = "sonnet",
) -> httpx.Response:
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client:
return await client.post(
"/cost/predict-cache",
json={
"current_deployment_id": current_deployment_id,
"candidate_deployment_id": candidate_deployment_id,
"request": body,
},
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("warm_deployment", "warm_model", "expected_delta", "expected_penalty"),
[("sonnet", "claude-sonnet-5", -0.03325, 0.0), ("opus", "claude-opus-5", 0.007, 0.0115)],
)
async def test_switch_delta_accounts_for_each_deployment_cache(
monkeypatch: pytest.MonkeyPatch,
warm_deployment: str,
warm_model: str,
expected_delta: float,
expected_penalty: float,
) -> None:
cache: Final = DualCache()
body: Final = _body()
await _observe(cache, body, deployment_id=warm_deployment, model=warm_model)
app: Final = _app(monkeypatch, cache, caller=UserAPIKeyAuth(api_key=_CALLER))
response: Final = await _post(app, body)
assert response.status_code == 200, response.text
result: Final = CachePredictionResponse.model_validate(response.json())
assert result.switch_delta == pytest.approx(expected_delta)
assert result.cache_rebuild_penalty == pytest.approx(expected_penalty)
assert result.cache_guarantee is False
assert result.pricing_basis == "input_before_discounts_and_margins"
if warm_deployment == "sonnet":
assert result.switch.cache_state == "warm"
assert result.stay.cache_state == "unknown"
else:
assert result.stay.cache_state == "warm"
assert result.switch.cache_state == "unknown"
@pytest.mark.asyncio
async def test_missing_caller_identity_cannot_reuse_observations(monkeypatch: pytest.MonkeyPatch) -> None:
cache: Final = DualCache()
body: Final = _body()
await _observe(cache, body)
response: Final = await _post(
_app(monkeypatch, cache, caller=UserAPIKeyAuth(api_key=None), counts=_unexpected_count), body
)
assert response.status_code == 200, response.text
result: Final = CachePredictionResponse.model_validate(response.json())
assert result.stay.reason == result.switch.reason == "caller_identity_unavailable"
assert result.stay.estimate is None and result.switch.estimate is None
assert result.switch_delta is None and result.cache_rebuild_penalty is None
@pytest.mark.asyncio
async def test_unauthenticated_request_is_rejected(monkeypatch: pytest.MonkeyPatch) -> None:
import litellm.proxy.proxy_server as proxy_server
monkeypatch.setattr(proxy_server, "master_key", "cache-prediction-test-master-key")
response: Final = await _post(_app(monkeypatch, DualCache()), _body())
assert response.status_code == 401, response.text
@pytest.mark.asyncio
@pytest.mark.parametrize("arm", ["current", "candidate"])
@pytest.mark.parametrize("caller_team", [None, "own-team"])
@pytest.mark.parametrize("restricted", [False, True])
async def test_foreign_and_missing_deployments_have_identical_authenticated_responses(
monkeypatch: pytest.MonkeyPatch, arm: str, caller_team: str | None, restricted: bool
) -> None:
allowed: Final = ("sonnet",) if arm == "current" else ("opus",)
app: Final = _app(
monkeypatch,
DualCache(),
caller=UserAPIKeyAuth(api_key=_CALLER, team_id=caller_team, models=list(allowed) if restricted else []),
current_team="foreign-team" if arm == "current" else None,
candidate_team="foreign-team" if arm == "candidate" else None,
counts=_unexpected_count,
)
foreign: Final = await _post(app, _body())
missing: Final = await _post(
app,
_body(),
current_deployment_id="missing-deployment" if arm == "current" else "opus",
candidate_deployment_id="missing-deployment" if arm == "candidate" else "sonnet",
)
assert foreign.status_code == missing.status_code == 404
assert foreign.json() == missing.json() == {"detail": "Deployment not found"}
@pytest.mark.asyncio
@pytest.mark.parametrize("deployment_team", [None, "own-team"])
async def test_visible_public_and_own_team_deployments_remain_available(
monkeypatch: pytest.MonkeyPatch, deployment_team: str | None
) -> None:
app: Final = _app(
monkeypatch,
DualCache(),
caller=UserAPIKeyAuth(api_key=_CALLER, team_id="own-team"),
current_team=deployment_team,
candidate_team=deployment_team,
)
response: Final = await _post(app, _body())
assert response.status_code == 200, response.text
result: Final = CachePredictionResponse.model_validate(response.json())
assert result.stay.estimate is not None and result.switch.estimate is not None
@pytest.mark.asyncio
@pytest.mark.parametrize("arm", ["current", "candidate"])
async def test_visible_deployment_outside_key_model_permissions_is_forbidden(
monkeypatch: pytest.MonkeyPatch, arm: str
) -> None:
allowed: Final = "sonnet" if arm == "current" else "opus"
denied: Final = "opus" if arm == "current" else "sonnet"
app: Final = _app(monkeypatch, DualCache(), caller=UserAPIKeyAuth(api_key=_CALLER, models=[allowed]))
response: Final = await _post(app, _body())
assert response.status_code == 403, response.text
assert denied in response.text
@pytest.mark.asyncio
async def test_other_callers_warm_cache_is_not_prediction_evidence(monkeypatch: pytest.MonkeyPatch) -> None:
cache: Final = DualCache()
body: Final = _body()
await _observe(cache, body, caller="other-caller")
response: Final = await _post(_app(monkeypatch, cache, caller=UserAPIKeyAuth(api_key=_CALLER)), body)
assert response.status_code == 200, response.text
result: Final = CachePredictionResponse.model_validate(response.json())
assert result.switch.cache_state == "unknown"
assert result.switch.reason == "no_compatible_observation"
assert result.switch.evidence is None
assert result.switch.estimate is not None
assert result.switch.estimate.tokens.cache_read_input_tokens == 0
@pytest.mark.asyncio
async def test_count_failure_nulls_switch_comparison(monkeypatch: pytest.MonkeyPatch) -> None:
app: Final = _app(
monkeypatch, DualCache(), caller=UserAPIKeyAuth(api_key=_CALLER), counts=Counts(total=None)
)
response: Final = await _post(app, _body())
assert response.status_code == 200, response.text
result: Final = CachePredictionResponse.model_validate(response.json())
assert result.stay.reason == result.switch.reason == "token_count_unavailable"
assert result.stay.estimate is None and result.switch.estimate is None
assert result.switch_delta is None and result.cache_rebuild_penalty is None
@pytest.mark.asyncio
@pytest.mark.parametrize("limiter", [None, CustomLogger()])
async def test_missing_or_unsupported_limiter_returns_unknown_before_counting(
monkeypatch: pytest.MonkeyPatch, limiter: CustomLogger | None
) -> None:
app: Final = _app(
monkeypatch, DualCache(), caller=UserAPIKeyAuth(api_key=_CALLER), counts=_unexpected_count, limiter=limiter
)
response: Final = await _post(app, _body())
assert response.status_code == 200, response.text
result: Final = CachePredictionResponse.model_validate(response.json())
assert result.stay.reason == result.switch.reason == "limiter_unavailable"
assert result.stay.estimate is None and result.switch.estimate is None
assert result.switch_delta is None and result.cache_rebuild_penalty is None
@pytest.mark.asyncio
async def test_occupied_parallel_capacity_rejects_before_provider_count(monkeypatch: pytest.MonkeyPatch) -> None:
cache: Final = DualCache()
limiter: Final = _PROXY_MaxParallelRequestsHandler_v3(InternalUsageCache(cache))
caller: Final = UserAPIKeyAuth(api_key=_CALLER, max_parallel_requests=1)
app: Final = _app(monkeypatch, cache, caller=caller, counts=_unexpected_count, limiter=limiter)
async with limiter.request_capacity(caller, "opus"):
response: Final = await _post(app, _body())
assert response.status_code == 429, response.text
assert "max_parallel_requests" in response.text
recovered: Final = await _post(_app(monkeypatch, cache, caller=caller, limiter=limiter), _body())
assert recovered.status_code == 200, recovered.text
@pytest.mark.asyncio
async def test_each_count_consumes_the_deployment_group_rpm_limit(monkeypatch: pytest.MonkeyPatch) -> None:
calls: Final = asyncio.Queue[str]()
async def count(model: str, api_key: str, body: Mapping[str, JsonValue]) -> int | None:
calls.put_nowait(model)
return await Counts()(model, api_key, body)
caller: Final = UserAPIKeyAuth(api_key=_CALLER, metadata={"model_rpm_limit": {"sonnet": 1}})
app: Final = _app(monkeypatch, DualCache(), caller=caller, counts=count)
response: Final = await _post(app, _body())
assert response.status_code == 429, response.text
assert calls.qsize() == 3
assert tuple(calls.get_nowait() for _ in range(3)) == (
"claude-opus-5", "claude-opus-5", "claude-sonnet-5"
)
@pytest.mark.asyncio
@pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"])
async def test_each_count_preserves_auth_cached_request_tag_limits(
monkeypatch: pytest.MonkeyPatch, metadata_key: str
) -> None:
calls: Final = asyncio.Queue[str]()
caller: Final = UserAPIKeyAuth(api_key=_CALLER, metadata={"tag_rpm_limit": {"cache-cost": 1}})
async def count(model: str, api_key: str, body: Mapping[str, JsonValue]) -> int | None:
calls.put_nowait(model)
return await Counts()(model, api_key, body)
async def authenticated_request(request: Request) -> UserAPIKeyAuth:
data: Final = await _read_request_body(request)
_safe_set_request_parsed_body(request, {**data, metadata_key: {"tags": ["cache-cost"]}})
return caller
app: Final = _app(monkeypatch, DualCache(), caller=caller, counts=count)
app.dependency_overrides[endpoint.user_api_key_auth] = authenticated_request
response: Final = await _post(app, _body())
assert response.status_code == 429, response.text
assert "tag_per_key" in response.text
assert calls.qsize() == 1
assert calls.get_nowait() == "claude-opus-5"
@pytest.mark.asyncio
async def test_provider_counter_failure_releases_parallel_capacity(monkeypatch: pytest.MonkeyPatch) -> None:
cache: Final = DualCache()
limiter: Final = _PROXY_MaxParallelRequestsHandler_v3(InternalUsageCache(cache))
caller: Final = UserAPIKeyAuth(api_key=_CALLER, max_parallel_requests=1)
async def fail_count(model: str, api_key: str, body: Mapping[str, JsonValue]) -> int | None:
raise RuntimeError("provider counter failed")
app: Final = _app(monkeypatch, cache, caller=caller, counts=fail_count, limiter=limiter)
with pytest.raises(RuntimeError, match="provider counter failed"):
await _post(app, _body())
recovered: Final = await _post(_app(monkeypatch, cache, caller=caller, limiter=limiter), _body())
assert recovered.status_code == 200, recovered.text
assert recovered.json()["switch"]["estimate"]["input_cost"] == pytest.approx(0.0145)
@pytest.mark.asyncio
async def test_cancelled_provider_counter_releases_parallel_capacity(monkeypatch: pytest.MonkeyPatch) -> None:
cache: Final = DualCache()
limiter: Final = _PROXY_MaxParallelRequestsHandler_v3(InternalUsageCache(cache))
caller: Final = UserAPIKeyAuth(api_key=_CALLER, max_parallel_requests=1)
started: Final = asyncio.Event()
release: Final = asyncio.Event()
async def wait_count(model: str, api_key: str, body: Mapping[str, JsonValue]) -> int | None:
started.set()
await release.wait()
return await Counts()(model, api_key, body)
app: Final = _app(monkeypatch, cache, caller=caller, counts=wait_count, limiter=limiter)
pending: Final = asyncio.create_task(_post(app, _body()))
try:
await asyncio.wait_for(started.wait(), timeout=5)
pending.cancel()
with pytest.raises(asyncio.CancelledError):
await pending
release.set()
recovered: Final = await asyncio.wait_for(_post(app, _body()), timeout=5)
assert recovered.status_code == 200, recovered.text
assert recovered.json()["switch"]["estimate"]["input_cost"] == pytest.approx(0.0145)
finally:
pending.cancel()
release.set()
await asyncio.gather(pending, return_exceptions=True)
async def _unexpected_count(model: str, api_key: str, body: Mapping[str, JsonValue]) -> int | None:
pytest.fail("Unsupported prediction must return before contacting the token counter")
class RequestMutator(CustomLogger):
async def async_pre_call_hook(
self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict[str, object], call_type: str
) -> dict[str, object]:
return {**data, "system": "Injected policy"}
@pytest.fixture
def request_mutator() -> Iterator[RequestMutator]:
callback: Final = RequestMutator()
litellm.logging_callback_manager.add_litellm_callback(callback)
try:
yield callback
finally:
litellm.logging_callback_manager.remove_callback_from_all_lists(callback)
@pytest.mark.asyncio
async def test_request_transform_callback_returns_unknown_before_token_counting(
monkeypatch: pytest.MonkeyPatch, request_mutator: RequestMutator
) -> None:
app: Final = _app(
monkeypatch, DualCache(), caller=UserAPIKeyAuth(api_key=_CALLER), counts=_unexpected_count
)
response: Final = await _post(app, _body())
assert response.status_code == 200, response.text
result: Final = CachePredictionResponse.model_validate(response.json())
assert result.stay.cache_state == result.switch.cache_state == "unknown"
assert result.stay.reason == result.switch.reason == "unsupported_request_transform"
assert result.stay.estimate is None and result.switch.estimate is None
assert result.switch_delta is None and result.cache_rebuild_penalty is None
@pytest.mark.asyncio
async def test_key_config_returns_unknown_before_token_counting(monkeypatch: pytest.MonkeyPatch) -> None:
app: Final = _app(
monkeypatch,
DualCache(),
caller=UserAPIKeyAuth(api_key=_CALLER, config={"model_list": []}),
counts=_unexpected_count,
)
response: Final = await _post(app, _body())
assert response.status_code == 200, response.text
result: Final = CachePredictionResponse.model_validate(response.json())
assert result.stay.cache_state == result.switch.cache_state == "unknown"
assert result.stay.reason == result.switch.reason == "unsupported_request_transform"
assert result.stay.estimate is None and result.switch.estimate is None
assert result.switch_delta is None and result.cache_rebuild_penalty is None
@pytest.mark.parametrize("headers", [
{"anthropic-version": "2099-01-01"},
{"anthropic-beta": "future-feature"},
])
@pytest.mark.asyncio
async def test_unsupported_provider_headers_cannot_reuse_default_version_evidence(
monkeypatch: pytest.MonkeyPatch, headers: dict[str, str]
) -> None:
cache: Final = DualCache()
await _observe(cache, _body(), deployment_id="sonnet")
app: Final = _app(
monkeypatch, cache, caller=UserAPIKeyAuth(api_key=_CALLER), counts=_unexpected_count
)
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client:
response: Final = await client.post(
"/cost/predict-cache",
headers=headers,
json={"current_deployment_id": "opus", "candidate_deployment_id": "sonnet", "request": _body()},
)
assert response.status_code == 200, response.text
result: Final = CachePredictionResponse.model_validate(response.json())
assert result.stay.cache_state == result.switch.cache_state == "unknown"
assert result.stay.reason == result.switch.reason == "unsupported_provider_headers"
assert result.stay.estimate is None and result.switch.estimate is None
assert result.switch_delta is None and result.cache_rebuild_penalty is None

45
tests/test_litellm/proxy/test_common_request_processing.py
View file

@ -1402,6 +1402,51 @@ class TestProxyBaseLLMRequestProcessing:
assert "x-litellm-key-spend" in headers_7
assert float(headers_7["x-litellm-key-spend"]) == 0.001 # Should use original spend on error
@pytest.mark.parametrize(
("hidden_params", "request_data", "expected_call_id"),
[
(
{"litellm_call_id": "call-from-hidden-params"},
{"litellm_call_id": "call-from-request"},
"call-from-hidden-params",
),
({}, {"litellm_call_id": "call-from-request"}, "call-from-request"),
({"model_id": "m-1"}, {"litellm_call_id": "call-from-request"}, "call-from-request"),
],
)
def test_get_custom_headers_call_id_falls_back_to_hidden_params_then_request_data(
self, hidden_params, request_data, expected_call_id
):
mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth)
mock_user_api_key_dict.tpm_limit = None
mock_user_api_key_dict.rpm_limit = None
mock_user_api_key_dict.max_budget = None
mock_user_api_key_dict.spend = 0.0
headers = ProxyBaseLLMRequestProcessing.get_custom_headers(
user_api_key_dict=mock_user_api_key_dict,
hidden_params=hidden_params,
request_data=request_data,
)
assert headers["x-litellm-call-id"] == expected_call_id
def test_get_custom_headers_explicit_call_id_wins_over_fallbacks(self):
mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth)
mock_user_api_key_dict.tpm_limit = None
mock_user_api_key_dict.rpm_limit = None
mock_user_api_key_dict.max_budget = None
mock_user_api_key_dict.spend = 0.0
headers = ProxyBaseLLMRequestProcessing.get_custom_headers(
user_api_key_dict=mock_user_api_key_dict,
call_id="explicit-call-id",
hidden_params={"litellm_call_id": "call-from-hidden-params"},
request_data={"litellm_call_id": "call-from-request"},
)
assert headers["x-litellm-call-id"] == "explicit-call-id"
@pytest.mark.asyncio
async def test_queue_time_seconds_is_set_in_metadata(self, monkeypatch):
"""

27
tests/test_litellm/proxy/test_health_check_functions.py
View file

@ -374,7 +374,7 @@ async def test_save_background_health_checks_to_db():
"""Test the main background health check save function"""
mock_prisma = MagicMock()
mock_prisma.save_health_check_result = AsyncMock()
mock_prisma.get_all_latest_health_checks = AsyncMock(return_value=[])
mock_prisma.db.query_raw = AsyncMock(return_value=[])
model_list = [
{
@ -398,9 +398,9 @@ async def test_save_background_health_checks_to_db():
"background_health_check",
)
# Should call get_all_latest_health_checks and save_health_check_result, and report completion
# Should read the latest rows and save_health_check_result, and report completion
assert persisted is True
mock_prisma.get_all_latest_health_checks.assert_called_once()
mock_prisma.db.query_raw.assert_awaited_once()
mock_prisma.save_health_check_result.assert_called_once()
call_kwargs = mock_prisma.save_health_check_result.call_args[1]
@ -493,7 +493,7 @@ def _one_model_setup():
@pytest.mark.asyncio
async def test_save_background_health_checks_to_db_returns_false_when_a_write_fails():
mock_prisma = MagicMock()
mock_prisma.get_all_latest_health_checks = AsyncMock(return_value=[])
mock_prisma.db.query_raw = AsyncMock(return_value=[])
mock_prisma.save_health_check_result = AsyncMock(return_value=None)
model_list, healthy_endpoints, unhealthy_endpoints = _one_model_setup()
@ -504,6 +504,23 @@ async def test_save_background_health_checks_to_db_returns_false_when_a_write_fa
assert (persisted, mock_prisma.save_health_check_result.await_count) == (False, 1)
@pytest.mark.asyncio
async def test_save_background_health_checks_to_db_writes_nothing_when_the_latest_row_read_fails(mock_prisma):
"""
A failed dedup read must not read as an empty table. Treated that way, every model was written on every
cycle by every pod while the read kept failing, which is what filled the table in production.
"""
mock_prisma.db.query_raw = AsyncMock(side_effect=RuntimeError("db down"))
mock_prisma.save_health_check_result = AsyncMock(return_value={"id": "row"})
model_list, healthy_endpoints, unhealthy_endpoints = _one_model_setup()
persisted = await _save_background_health_checks_to_db(
mock_prisma, model_list, healthy_endpoints, unhealthy_endpoints, 1234567890.0, "background_health_check"
)
assert (persisted, mock_prisma.save_health_check_result.await_count) == (False, 0)
@pytest.mark.asyncio
async def test_save_background_health_checks_to_db_no_prisma():
"""Test graceful handling when no prisma client"""
@ -515,7 +532,7 @@ async def test_save_background_health_checks_to_db_no_prisma():
async def test_save_background_health_checks_to_db_exception_handling():
"""Test exception handling in background health check save"""
mock_prisma = MagicMock()
mock_prisma.get_all_latest_health_checks = AsyncMock(side_effect=Exception("DB Error"))
mock_prisma.db.query_raw = AsyncMock(side_effect=Exception("DB Error"))
model_list = [
{

39
tests/test_litellm/proxy/test_proxy_server.py
View file

@ -12795,6 +12795,45 @@ async def test_moderations_reraises_proxy_exception_unwrapped():
mock_logging.post_call_failure_hook.assert_awaited_once()
@pytest.mark.asyncio
async def test_moderations_response_carries_litellm_call_id_header():
from fastapi import Response
from litellm.types.utils import ModerationCreateResponse
call_id = "moderation-call-id-123"
moderation_response = ModerationCreateResponse(id="modr-1", model="omni-moderation-latest", results=[])
moderation_response._hidden_params = {"litellm_call_id": call_id, "model_id": "mod-deployment-1"}
async def fake_llm_call():
return moderation_response
async def passthrough_add_litellm_data(data, **kwargs):
return {**data, "litellm_call_id": call_id}
request = MagicMock()
request.body = AsyncMock(return_value=b'{"input": "hi"}')
fastapi_response = Response()
user_api_key_dict = UserAPIKeyAuth(api_key="sk-test", spend=0.0)
with (
patch.object(proxy_server_module, "add_litellm_data_to_request", new=passthrough_add_litellm_data), # test-quality-ok: the route reads this module global, no injection point
patch.object(proxy_server_module, "route_request", new=AsyncMock(return_value=fake_llm_call())), # test-quality-ok: fakes the provider call so the response headers assembled by the real route are observable
patch.object(proxy_server_module, "proxy_logging_obj") as mock_logging, # test-quality-ok: module global, no injection point
):
mock_logging.pre_call_hook = AsyncMock(side_effect=lambda user_api_key_dict, data, call_type: data)
mock_logging.update_request_status = AsyncMock()
result = await proxy_server_module.moderations(
request=request,
fastapi_response=fastapi_response,
user_api_key_dict=user_api_key_dict,
)
assert result is moderation_response
assert fastapi_response.headers["x-litellm-call-id"] == call_id
assert fastapi_response.headers["x-litellm-model-id"] == "mod-deployment-1"
@pytest.mark.asyncio
async def test_init_agents_in_db_rebuilds_registry_under_agent_reconcile_lock(monkeypatch):
from litellm.proxy.agent_endpoints.agent_registry import (

54
tests/test_litellm/test_assert_ci_coverage.py
View file

@ -9,8 +9,12 @@ the question neither covers: whether the job that globs a file then deselects it
"""
import importlib.util
import json
import sys
from pathlib import Path
from typing import Final
import yaml
_REPO_ROOT = Path(__file__).resolve().parents[2]
_MODULE_PATH = _REPO_ROOT / ".github" / "scripts" / "assert_ci_coverage.py"
@ -20,6 +24,56 @@ sys.modules[_spec.name] = coverage # @dataclass(slots=True) rebuilds via sys.mo
_spec.loader.exec_module(coverage)
def test_integration_manifest_requires_exclusive_scheduled_circleci_owner(tmp_path: Path) -> None:
test_path: Final = "tests/integration/management/test_contract.py"
test_file: Final = tmp_path / test_path
test_file.parent.mkdir(parents=True)
test_file.write_text("def test_contract(): pass\n")
(tmp_path / "tests/integration/contracts.json").write_text(
json.dumps({"groups": {"management": ["management"]}, "tests": {f"{test_path}::test_contract": ["mgmt.test"]}})
)
paths, findings = coverage._integration_ownership(tmp_path)
assert not paths
assert [finding.detail for finding in findings] == ["dedicated CircleCI runner is missing"]
circle: Final = tmp_path / ".circleci/config.yml"
circle.parent.mkdir()
circle.write_text(
yaml.safe_dump(
{
"jobs": {
"integration_contracts": {
"steps": [{"run": {"command": "bash .circleci/scripts/run_integration.sh management"}}]
}
},
"workflows": {"integration": {"jobs": [{"integration_contracts": {"suite": "management"}}]}},
}
)
)
paths, findings = coverage._integration_ownership(tmp_path)
assert paths == frozenset({test_path})
assert findings == ()
configured: Final = yaml.safe_load(circle.read_text())
configured["workflows"]["integration"]["jobs"] = [
{"integration_contracts": {"matrix": {"parameters": {"suite": ["providers"]}}}}
]
circle.write_text(yaml.safe_dump(configured))
_, findings = coverage._integration_ownership(tmp_path)
assert [(finding.subject, finding.detail) for finding in findings] == [
("management", "canonical integration group is not scheduled by CircleCI")
]
configured["workflows"]["integration"]["jobs"][0]["integration_contracts"]["matrix"]["parameters"]["suite"] = [
"management"
]
circle.write_text(yaml.safe_dump(configured))
workflow: Final = tmp_path / ".github/workflows/test.yml"
workflow.parent.mkdir(parents=True)
workflow.write_text(yaml.safe_dump({"jobs": {"tests": {"steps": [{"run": "pytest tests/integration"}]}}}))
_, findings = coverage._integration_ownership(tmp_path)
assert [(finding.subject, finding.detail) for finding in findings] == [
(test_path, "integration contract is also selected by GitHub Actions")
]
def test_an_ancestor_directory_covers_a_file_but_does_not_name_it():
# The whole point of the split: `tests/x` answers "does it run?" but not
# "which shard owns it?" — accepting it for the latter is how a new child

55
tests/test_litellm/test_main.py
View file

@ -3795,3 +3795,58 @@ def test_azure_ai_speech_on_a_foundry_host_uses_the_azure_openai_deployment_rout
assert route.called
assert response.content == b"mp3-bytes"
FORWARDED_CLIENT_HEADERS: Final = {"x-forwarded-for": "10.0.0.1", "x-amzn-trace-id": "Root=1-lit7694"}
def _chat_completion_json() -> Mapping[str, object]:
return {
"id": "chatcmpl-lit7694",
"object": "chat.completion",
"created": 1,
"model": "gpt-5.4",
"choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
}
def _chat_completion_sse() -> bytes:
chunk: Final = {
"id": "chatcmpl-lit7694",
"object": "chat.completion.chunk",
"created": 1,
"model": "gpt-5.4",
"choices": [{"index": 0, "delta": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
}
return f"data: {json.dumps(chunk)}\n\ndata: [DONE]\n\n".encode()
@pytest.mark.parametrize("stream", [False, True])
def test_bridged_responses_with_openai_http_handler_keeps_forwarded_headers_out_of_the_body(
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, stream: bool
):
monkeypatch.setenv("EXPERIMENTAL_OPENAI_BASE_LLM_HTTP_HANDLER", "true")
route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock(
return_value=httpx.Response(200, content=_chat_completion_sse(), headers={"content-type": "text/event-stream"})
if stream
else httpx.Response(200, json=_chat_completion_json())
)
response: Final = litellm.responses(
model="openai/gpt-5.4",
input="Reply with the single word ok",
stream=stream,
use_chat_completions_api=True,
headers=dict(FORWARDED_CLIENT_HEADERS),
api_key="sk-test",
)
if stream:
list(response)
assert route.called
request: Final = route.calls.last.request
body: Final = json.loads(request.content)
assert "extra_headers" not in body
assert body["model"] == "gpt-5.4"
assert {k: request.headers[k] for k in FORWARDED_CLIENT_HEADERS} == FORWARDED_CLIENT_HEADERS

21
tests/test_litellm/test_utils.py
View file

@ -29,6 +29,7 @@ from litellm._logging import (
verbose_logger,
)
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
from litellm.litellm_core_utils.thread_pool_executor import executor as logging_executor
from litellm.proxy.utils import is_valid_api_key
from litellm.types.utils import (
@ -5389,6 +5390,26 @@ def test_websearch_interception_control_fields_never_reach_the_provider():
assert set(WEBSEARCH_INTERNAL_CONTROL_FIELDS) <= set(all_litellm_params)
def test_get_litellm_params_keys_never_reach_the_provider():
"""Bridges (chat <-> Responses, agentic loop follow-ups) forward litellm_params as
`completion()` kwargs. Any key the param builder does not recognize is swept into
extra_body, and OpenAI rejects the call with `Unknown parameter: 'model_alias_map'`.
"""
litellm_param_keys = frozenset(get_litellm_params()) - {"drop_params"}
kwargs = {
"a_real_provider_specific_param": 1,
"model_alias_map": {"alias": "gpt-5.4"},
**{key: "configured-value" for key in litellm_param_keys - {"model_alias_map"}},
}
non_default = get_non_default_completion_params(kwargs)
assert non_default == {"a_real_provider_specific_param": 1}, (
"litellm params leaked into the provider params: "
f"{sorted(set(non_default) - {'a_real_provider_specific_param'})}"
)
def test_bedrock_batch_params_never_reach_the_provider():
"""A Bedrock managed-batch deployment carries aws_batch_role_arn / s3_* /
bedrock_tags in its litellm_params, and the same deployment also serves chat.

164
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -3427,6 +3427,35 @@ export interface paths {
patch?: never;
trace?: never;
};
"/cost/predict-cache": {
parameters: {
query?: never;
header?: never;
path?: never;
cookie?: never;
};
get?: never;
put?: never;
/**
* Predict Cache Cost
* @description 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.
*/
post: operations["predict_cache_cost_cost_predict_cache_post"];
delete?: never;
options?: never;
head?: never;
patch?: never;
trace?: never;
};
"/credentials": {
parameters: {
query?: never;
@ -24702,6 +24731,31 @@ export interface components {
/** Failed Requests */
failed_requests: number;
};
/** CacheCostScenario */
CacheCostScenario: {
/** Input Cost */
input_cost: number;
tokens: components["schemas"]["CacheTokenBuckets"];
};
/** CacheEvidence */
CacheEvidence: {
/**
* Confidence
* @default observed
* @constant
*/
confidence: "observed";
/** Expires At */
expires_at: number;
/** Observed At */
observed_at: number;
/**
* Source
* @default provider_usage
* @constant
*/
source: "provider_usage";
};
/** CachePingResponse */
CachePingResponse: {
/** Cache Type */
@ -24719,6 +24773,59 @@ export interface components {
/** Status */
status: string;
};
/** CachePredictionArm */
CachePredictionArm: {
/**
* Cache State
* @default unknown
* @enum {string}
*/
cache_state: "warm" | "partial" | "stale" | "unknown" | "disabled";
cold?: components["schemas"]["CacheCostScenario"] | null;
/** Deployment Id */
deployment_id: string;
estimate?: components["schemas"]["CacheCostScenario"] | null;
evidence?: components["schemas"]["CacheEvidence"] | null;
/** Model */
model?: string | null;
/** Reason */
reason?: string | null;
/** Token Count Source */
token_count_source?: "anthropic_count_tokens" | null;
warm?: components["schemas"]["CacheCostScenario"] | null;
};
/** CachePredictionRequest */
CachePredictionRequest: {
/** Candidate Deployment Id */
candidate_deployment_id: string;
/** Current Deployment Id */
current_deployment_id: string;
/** Request */
request: {
[key: string]: components["schemas"]["JsonValue"];
};
};
/** CachePredictionResponse */
CachePredictionResponse: {
/**
* Cache Guarantee
* @default false
* @constant
*/
cache_guarantee: false;
/** Cache Rebuild Penalty */
cache_rebuild_penalty: number | null;
/**
* Pricing Basis
* @default input_before_discounts_and_margins
* @constant
*/
pricing_basis: "input_before_discounts_and_margins";
stay: components["schemas"]["CachePredictionArm"];
switch: components["schemas"]["CachePredictionArm"];
/** Switch Delta */
switch_delta: number | null;
};
/** CacheSettingsField */
CacheSettingsField: {
/** Field Default */
@ -24800,6 +24907,29 @@ export interface components {
*/
status: string;
};
/** CacheTokenBuckets */
CacheTokenBuckets: {
/**
* Cache Creation 1H Input Tokens
* @default 0
*/
cache_creation_1h_input_tokens: number;
/**
* Cache Creation 5M Input Tokens
* @default 0
*/
cache_creation_5m_input_tokens: number;
/**
* Cache Read Input Tokens
* @default 0
*/
cache_read_input_tokens: number;
/**
* Uncached Input Tokens
* @default 0
*/
uncached_input_tokens: number;
};
/**
* CallTypes
* @enum {string}
@ -28421,6 +28551,7 @@ export interface components {
/** Updated By */
updated_by?: string | null;
};
JsonValue: unknown;
/** KeyHealthResponse */
KeyHealthResponse: {
/**
@ -45322,6 +45453,39 @@ export interface operations {
};
};
};
predict_cache_cost_cost_predict_cache_post: {
parameters: {
query?: never;
header?: never;
path?: never;
cookie?: never;
};
requestBody: {
content: {
"application/json": components["schemas"]["CachePredictionRequest"];
};
};
responses: {
/** @description Successful Response */
200: {
headers: {
[name: string]: unknown;
};
content: {
"application/json": components["schemas"]["CachePredictionResponse"];
};
};
/** @description Validation Error */
422: {
headers: {
[name: string]: unknown;
};
content: {
"application/json": components["schemas"]["HTTPValidationError"];
};
};
};
};
get_credentials_credentials_get: {
parameters: {
query?: never;