mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-16 23:41:43 +00:00
chore: merge origin/main into litellm_lit6314_guardrail_metadata_transfer
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
commit
9acdebf563
267 changed files with 18196 additions and 2472 deletions
|
|
@ -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,69 @@ jobs:
|
|||
exit 1
|
||||
fi
|
||||
|
||||
provider_replay_harness:
|
||||
docker:
|
||||
- *python312_image
|
||||
working_directory: ~/project
|
||||
resource_class: medium
|
||||
steps:
|
||||
- setup_litellm_test_deps
|
||||
- run:
|
||||
name: Test provider replay harness
|
||||
command: |
|
||||
mkdir -p test-results/provider-replay-harness
|
||||
uv run --no-sync pytest -q --noconftest -o addopts= -o pythonpath=tests/e2e -p no:rerunfailures \
|
||||
--junitxml=test-results/provider-replay-harness/junit.xml \
|
||||
tests/e2e/test_provider_edge.py tests/e2e/test_fixture_bundle.py \
|
||||
tests/e2e/test_fixture_canonical.py tests/e2e/test_fixture_mode.py \
|
||||
tests/code_coverage_tests/test_provider_replay_harness.py
|
||||
- store_test_results:
|
||||
path: test-results/provider-replay-harness
|
||||
|
||||
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:
|
||||
|
|
@ -2921,6 +2986,7 @@ workflows:
|
|||
only:
|
||||
- main
|
||||
- /litellm_.*/
|
||||
- provider_replay_harness
|
||||
- base_sdk_install:
|
||||
filters: *main_branches
|
||||
- local_testing_part1:
|
||||
|
|
|
|||
142
.circleci/scripts/run_integration.sh
Normal file
142
.circleci/scripts/run_integration.sh
Normal 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
53
.circleci/scripts/stop_integration_processes.py
Normal 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
43
.circleci/scripts/wait_integration_services.py
Normal 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
11
.github/codeql/codeql-config.yml
vendored
|
|
@ -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
|
||||
|
|
|
|||
15
.github/e2e-stack/assert_tests_ran.py
vendored
15
.github/e2e-stack/assert_tests_ran.py
vendored
|
|
@ -3,6 +3,9 @@ import xml.etree.ElementTree as ET
|
|||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "tests/e2e"))
|
||||
from coverage_registry.management_cases import MANAGEMENT_CASES
|
||||
|
||||
|
||||
def main() -> int:
|
||||
selected: Final = tuple(sys.argv[2:])
|
||||
|
|
@ -16,6 +19,17 @@ def main() -> int:
|
|||
case.get("file") for case in cases if all(case.find(tag) is None for tag in ("skipped", "failure", "error"))
|
||||
)
|
||||
missing: Final = tuple(path for path in selected if path not in passed)
|
||||
required_nodes: Final = frozenset(case.node for case in MANAGEMENT_CASES if case.node.split("::", 1)[0] in selected)
|
||||
passed_nodes: Final = frozenset(
|
||||
prop.get("value")
|
||||
for case in cases
|
||||
if all(case.find(tag) is None for tag in ("skipped", "failure", "error"))
|
||||
for prop in case.findall("./properties/property")
|
||||
if prop.get("name") == "management_node"
|
||||
)
|
||||
missing_nodes: Final = required_nodes - passed_nodes
|
||||
for node in sorted(missing_nodes):
|
||||
_ = sys.stdout.write(f"::error::required management case did not pass: {node}\n")
|
||||
for path in selected:
|
||||
collected: Final = sum(case.get("file") == path for case in cases)
|
||||
skipped: Final = sum(case.get("file") == path and case.find("skipped") is not None for case in cases)
|
||||
|
|
@ -27,6 +41,7 @@ def main() -> int:
|
|||
if (
|
||||
selected
|
||||
and not missing
|
||||
and not missing_nodes
|
||||
and not any(case.find(tag) is not None for case in cases for tag in ("failure", "error"))
|
||||
):
|
||||
return 0
|
||||
|
|
|
|||
5
.github/e2e-stack/oidc-profile.sh
vendored
Executable file
5
.github/e2e-stack/oidc-profile.sh
vendored
Executable file
|
|
@ -0,0 +1,5 @@
|
|||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
|
||||
cd "${REPO_ROOT}"
|
||||
exec uv run --no-sync python tests/e2e/idp.py "$@"
|
||||
2
.github/e2e-stack/select_tests.py
vendored
2
.github/e2e-stack/select_tests.py
vendored
|
|
@ -12,6 +12,8 @@ UNSUPPORTED: Final = re.compile(
|
|||
HARNESS: Final = re.compile(
|
||||
r"^tests/e2e/[A-Za-z0-9_.-]+\.(py|ini)$"
|
||||
r"|^tests/e2e/idp_realm\.json$"
|
||||
r"|^tests/e2e/management/(management_client|jwt_actors|conftest)\.py$"
|
||||
r"|^tests/e2e/coverage_registry/management_cases\.py$"
|
||||
r"|^tests/e2e/gateway/"
|
||||
r"|^\.github/e2e-stack/"
|
||||
r"|^\.github/workflows/test-e2e-changed\.yml$"
|
||||
|
|
|
|||
67
.github/scripts/assert_ci_coverage.py
vendored
67
.github/scripts/assert_ci_coverage.py
vendored
|
|
@ -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())
|
||||
|
||||
|
|
|
|||
2
.github/workflows/_test-unit-base.yml
vendored
2
.github/workflows/_test-unit-base.yml
vendored
|
|
@ -57,6 +57,7 @@ permissions:
|
|||
|
||||
env:
|
||||
UV_PYTHON: "3.12"
|
||||
LITELLM_LOCAL_MODEL_COST_MAP: "True"
|
||||
|
||||
jobs:
|
||||
run:
|
||||
|
|
@ -113,6 +114,7 @@ jobs:
|
|||
if: steps.changes.outputs.decision != 'skip'
|
||||
timeout-minutes: 8
|
||||
run: |
|
||||
diff -u model_prices_and_context_window.json litellm/model_prices_and_context_window_backup.json
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml
|
||||
uv run --no-sync python -c 'import os, sys; print(sys.version); assert f"{sys.version_info.major}.{sys.version_info.minor}" == os.environ["UV_PYTHON"]'
|
||||
|
||||
|
|
|
|||
5
.github/workflows/test-code-quality.yml
vendored
5
.github/workflows/test-code-quality.yml
vendored
|
|
@ -178,7 +178,7 @@ jobs:
|
|||
version: "0.10.9"
|
||||
|
||||
- name: Install dependencies
|
||||
run: uv sync --frozen --extra proxy --python 3.10
|
||||
run: uv sync --frozen --extra proxy --extra cli --python 3.10
|
||||
|
||||
- run: uv run --no-sync python --version
|
||||
|
||||
|
|
@ -187,3 +187,6 @@ jobs:
|
|||
|
||||
- name: Check litellm CLI
|
||||
run: uv run --no-sync litellm --version
|
||||
|
||||
- name: Check lite CLI
|
||||
run: uv run --no-sync lite version
|
||||
|
|
|
|||
2
.github/workflows/test-e2e-changed.yml
vendored
2
.github/workflows/test-e2e-changed.yml
vendored
|
|
@ -183,7 +183,7 @@ jobs:
|
|||
log="${RUNNER_TEMP}/e2e-pass-${pass}.log"
|
||||
echo "::group::pass ${pass} of 3"
|
||||
set +e
|
||||
uv run --no-sync pytest "${test_files[@]}" --rootdir=. -v -p no:cacheprovider \
|
||||
uv run --no-sync pytest "${test_files[@]}" --rootdir=. -v --reruns 0 -p no:cacheprovider \
|
||||
-o junit_family=xunit1 --junitxml="${report}" > "${log}" 2>&1
|
||||
status=$?
|
||||
uv run --no-sync python .github/e2e-stack/assert_tests_ran.py "${report}" "${test_files[@]}"
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ use serde_json::Value;
|
|||
|
||||
#[derive(Deserialize)]
|
||||
struct Input {
|
||||
path: String,
|
||||
model_alias: String,
|
||||
provider_model: String,
|
||||
api_base: String,
|
||||
|
|
@ -21,7 +22,8 @@ async fn main() {
|
|||
Ok(input) => input,
|
||||
Err(error) => fail(error),
|
||||
};
|
||||
let result = litellm_ai_gateway::trace_parity::traced_messages_request(
|
||||
let result = litellm_ai_gateway::trace_parity::traced_request(
|
||||
input.path,
|
||||
input.model_alias,
|
||||
input.provider_model,
|
||||
input.api_base,
|
||||
|
|
|
|||
|
|
@ -29,14 +29,15 @@ pub struct TracedGatewayResponse {
|
|||
pub trace: Vec<litellm_core::observability::FunctionTraceEvent>,
|
||||
}
|
||||
|
||||
pub async fn traced_messages_request(
|
||||
pub async fn traced_request(
|
||||
path: String,
|
||||
model_alias: String,
|
||||
provider_model: String,
|
||||
api_base: String,
|
||||
body: Value,
|
||||
) -> TracedGatewayResponse {
|
||||
let trace = litellm_core::observability::FunctionTrace::default();
|
||||
let result = messages_request(model_alias, provider_model, api_base, body)
|
||||
let result = request(path, model_alias, provider_model, api_base, body)
|
||||
.with_subscriber(trace.dispatcher())
|
||||
.await;
|
||||
let events = trace.events();
|
||||
|
|
@ -54,7 +55,8 @@ pub async fn traced_messages_request(
|
|||
}
|
||||
}
|
||||
|
||||
pub async fn messages_request(
|
||||
pub async fn request(
|
||||
path: String,
|
||||
model_alias: String,
|
||||
provider_model: String,
|
||||
api_base: String,
|
||||
|
|
@ -75,7 +77,7 @@ pub async fn messages_request(
|
|||
};
|
||||
let request = Request::builder()
|
||||
.method("POST")
|
||||
.uri("/v1/messages")
|
||||
.uri(path)
|
||||
.header(AUTHORIZATION, "Bearer trace-master-key")
|
||||
.header(CONTENT_TYPE, "application/json")
|
||||
.body(Body::from(body.to_string()))
|
||||
|
|
|
|||
|
|
@ -538,7 +538,7 @@ context_window_fallbacks: Optional[List] = None
|
|||
content_policy_fallbacks: Optional[List] = None
|
||||
allowed_fails: int = 3
|
||||
allow_dynamic_callback_disabling: bool = True
|
||||
num_retries_per_request: Optional[int] = None # for the request overall (incl. fallbacks + model retries)
|
||||
num_retries_per_request: Optional[int] = None # cap on Router retries of one model group; resets per fallback hop
|
||||
####### SECRET MANAGERS #####################
|
||||
secret_manager_client: Optional[Any] = (
|
||||
None # list of instantiated key management clients - e.g. azure kv, infisical, etc.
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import ast
|
||||
import contextvars
|
||||
import functools
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
|
|
@ -225,6 +226,35 @@ class AccessLogRedactionFilter(logging.Filter):
|
|||
_access_log_filter: Final = AccessLogRedactionFilter()
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=1)
|
||||
def _parse_disabled_access_log_paths(raw: str) -> frozenset[str]:
|
||||
return frozenset(stripped for path in raw.split(",") if (stripped := path.strip()))
|
||||
|
||||
|
||||
def _disabled_access_log_paths() -> frozenset[str]:
|
||||
"""Read the variable per record so a value loaded later via proxy config
|
||||
environment_variables or dotenv is honored."""
|
||||
return _parse_disabled_access_log_paths(os.getenv("LITELLM_DISABLE_ACCESS_LOG_PATHS", ""))
|
||||
|
||||
|
||||
class AccessLogPathFilter(logging.Filter):
|
||||
"""Drops uvicorn.access records for request paths listed in LITELLM_DISABLE_ACCESS_LOG_PATHS.
|
||||
|
||||
uvicorn passes record.args as (client_addr, method, full_path, http_version, status_code).
|
||||
"""
|
||||
|
||||
def filter(self, record: logging.LogRecord) -> bool:
|
||||
if not isinstance(record.args, tuple) or len(record.args) < 3:
|
||||
return True
|
||||
full_path: Final = record.args[2]
|
||||
if not isinstance(full_path, str):
|
||||
return True
|
||||
return full_path.partition("?")[0] not in _disabled_access_log_paths()
|
||||
|
||||
|
||||
_access_log_path_filter: Final = AccessLogPathFilter()
|
||||
|
||||
|
||||
def _get_max_string_length_stdout_log() -> int:
|
||||
"""Read the limit per record so a value loaded later via proxy config
|
||||
environment_variables is honored."""
|
||||
|
|
@ -663,6 +693,7 @@ def _redact_third_party_loggers() -> None:
|
|||
for name in _REDACTED_THIRD_PARTY_LOGGERS:
|
||||
logging.getLogger(name).addFilter(_secret_filter)
|
||||
for name in _REDACTED_ACCESS_LOGGERS:
|
||||
logging.getLogger(name).addFilter(_access_log_path_filter)
|
||||
logging.getLogger(name).addFilter(_access_log_filter)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ import litellm
|
|||
from litellm._logging import verbose_logger
|
||||
from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import parse_prompt_tokens_details
|
||||
from litellm.llms.vertex_ai.batches.transformation import vertex_prompt_tokens_details
|
||||
from litellm.types.llms.openai import Batch
|
||||
from litellm.types.utils import ModelInfo, Usage
|
||||
from litellm.utils import token_counter
|
||||
|
|
@ -356,6 +357,7 @@ def calculate_vertex_ai_batch_cost_and_usage(
|
|||
prompt_tokens=_prompt,
|
||||
completion_tokens=_completion,
|
||||
total_tokens=_total,
|
||||
prompt_tokens_details=vertex_prompt_tokens_details(usage_metadata),
|
||||
)
|
||||
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -679,7 +679,14 @@ class Cache:
|
|||
cache_key, cached_data, kwargs = self._add_cache_logic(result=result, **kwargs)
|
||||
self.cache.set_cache(cache_key, cached_data, **kwargs)
|
||||
except Exception as e:
|
||||
log_redis_failure(verbose_logger, logging.ERROR, "LiteLLM Cache: exception in add_cache", e)
|
||||
self._log_add_cache_failure(e)
|
||||
|
||||
def _log_add_cache_failure(self, exc: Exception) -> None:
|
||||
message: Final = "LiteLLM Cache: exception in add_cache"
|
||||
if isinstance(self.cache, RedisCache):
|
||||
log_redis_failure(verbose_logger, logging.ERROR, message, exc)
|
||||
return
|
||||
verbose_logger.error("%s: %s", message, exc)
|
||||
|
||||
async def async_add_cache(self, result, dynamic_cache_object: BaseCache | None = None, **kwargs):
|
||||
"""
|
||||
|
|
@ -698,7 +705,7 @@ class Cache:
|
|||
else:
|
||||
await self.cache.async_set_cache(cache_key, cached_data, **kwargs)
|
||||
except Exception as e:
|
||||
log_redis_failure(verbose_logger, logging.ERROR, "LiteLLM Cache: exception in add_cache", e)
|
||||
self._log_add_cache_failure(e)
|
||||
|
||||
def _convert_to_cached_embedding(
|
||||
self,
|
||||
|
|
@ -877,7 +884,7 @@ class Cache:
|
|||
else:
|
||||
await self.cache.async_set_cache_pipeline(cache_list=cache_list, **kwargs)
|
||||
except Exception as e:
|
||||
log_redis_failure(verbose_logger, logging.ERROR, "LiteLLM Cache: exception in add_cache", e)
|
||||
self._log_add_cache_failure(e)
|
||||
|
||||
def should_use_cache(self, **kwargs):
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ import hashlib
|
|||
import inspect
|
||||
import json
|
||||
import logging
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Awaitable, Callable, Iterator, Sequence
|
||||
from contextvars import ContextVar
|
||||
|
|
@ -32,6 +33,7 @@ from litellm.constants import (
|
|||
REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD,
|
||||
REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT,
|
||||
REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION,
|
||||
REDIS_TIMEOUT_LOG_INTERVAL,
|
||||
)
|
||||
from litellm.litellm_core_utils.core_helpers import _get_parent_otel_span_from_kwargs
|
||||
from litellm.litellm_core_utils.coroutine_checker import coroutine_checker
|
||||
|
|
@ -340,7 +342,7 @@ def _explicit_causes(exc: BaseException) -> Iterator[BaseException]:
|
|||
current = current.__cause__
|
||||
|
||||
|
||||
def _is_redis_timeout_failure(exc: BaseException) -> bool:
|
||||
def is_redis_timeout_failure(exc: BaseException) -> bool:
|
||||
"""True when ``exc`` or any exception it was explicitly raised ``from`` is a timeout.
|
||||
|
||||
redis-py's blocking pool reports a pool wait timeout as ``ConnectionError`` chained from
|
||||
|
|
@ -414,7 +416,7 @@ def _record_swallowed_redis_failure(breaker: RedisCircuitBreaker, exc: BaseExcep
|
|||
"""
|
||||
if not _is_redis_health_failure(exc):
|
||||
return
|
||||
breaker.record_failure(is_timeout=_is_redis_timeout_failure(exc))
|
||||
breaker.record_failure(is_timeout=is_redis_timeout_failure(exc))
|
||||
_swallowed_redis_failures.set(_swallowed_redis_failures.get() + 1)
|
||||
|
||||
|
||||
|
|
@ -422,13 +424,58 @@ class RedisCircuitBreakerOpenError(Exception):
|
|||
pass
|
||||
|
||||
|
||||
class _RedisTimeoutLogThrottle:
|
||||
"""Admits one Redis timeout log line per interval and counts the timeouts it suppressed in between."""
|
||||
|
||||
def __init__(self, interval: float, clock: Callable[[], float] = time.monotonic) -> None:
|
||||
self.interval = interval
|
||||
self._clock = clock
|
||||
self._lock = threading.Lock()
|
||||
self._last_logged_at: float | None = None
|
||||
self._suppressed = 0
|
||||
|
||||
def admit(self) -> int | None:
|
||||
"""Return the number of timeouts suppressed since the last admitted line, or None to suppress this one."""
|
||||
with self._lock:
|
||||
now: Final = self._clock()
|
||||
if self._last_logged_at is not None and now - self._last_logged_at < self.interval:
|
||||
self._suppressed += 1
|
||||
return None
|
||||
suppressed: Final = self._suppressed
|
||||
self._suppressed = 0
|
||||
self._last_logged_at = now
|
||||
return suppressed
|
||||
|
||||
|
||||
_redis_timeout_log_throttle: Final = _RedisTimeoutLogThrottle(REDIS_TIMEOUT_LOG_INTERVAL)
|
||||
|
||||
|
||||
def log_redis_failure(
|
||||
logger: logging.Logger, level: int, message: str, exc: BaseException, with_traceback: bool = False
|
||||
) -> None:
|
||||
if isinstance(exc, RedisCircuitBreakerOpenError):
|
||||
logger.debug("%s: %s", message, exc)
|
||||
logger.debug("%s: %s", message, exc, stacklevel=2)
|
||||
return
|
||||
logger.log(level, "%s: %s", message, exc, exc_info=exc if with_traceback else None)
|
||||
exc_info: Final = exc if with_traceback else None
|
||||
if not is_redis_timeout_failure(exc):
|
||||
logger.log(level, "%s: %s", message, exc, exc_info=exc_info, stacklevel=2)
|
||||
return
|
||||
suppressed: Final = _redis_timeout_log_throttle.admit()
|
||||
if suppressed is None:
|
||||
logger.debug("%s: %s", message, exc, stacklevel=2)
|
||||
return
|
||||
if suppressed == 0:
|
||||
logger.log(level, "%s: %s", message, exc, exc_info=exc_info, stacklevel=2)
|
||||
return
|
||||
logger.log(
|
||||
level,
|
||||
"%s: %s (%d more Redis timeouts since the previous Redis timeout line were logged at DEBUG)",
|
||||
message,
|
||||
exc,
|
||||
suppressed,
|
||||
exc_info=exc_info,
|
||||
stacklevel=2,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -475,7 +522,7 @@ async def _run_under_circuit_breaker(
|
|||
result: Final = await call()
|
||||
except Exception as e:
|
||||
if _is_redis_health_failure(e):
|
||||
breaker.record_failure(is_timeout=_is_redis_timeout_failure(e))
|
||||
breaker.record_failure(is_timeout=is_redis_timeout_failure(e))
|
||||
raise
|
||||
_exit_circuit_breaker(breaker, admission)
|
||||
return result
|
||||
|
|
@ -492,7 +539,7 @@ def _run_under_circuit_breaker_sync(
|
|||
result: Final = call()
|
||||
except Exception as e:
|
||||
if _is_redis_health_failure(e):
|
||||
breaker.record_failure(is_timeout=_is_redis_timeout_failure(e))
|
||||
breaker.record_failure(is_timeout=is_redis_timeout_failure(e))
|
||||
raise
|
||||
_exit_circuit_breaker(breaker, admission)
|
||||
return result
|
||||
|
|
@ -801,10 +848,8 @@ class RedisCache(BaseCache):
|
|||
## LOGGING ##
|
||||
end_time = time.time()
|
||||
_duration = end_time - start_time
|
||||
verbose_logger.error(
|
||||
"LiteLLM Redis Caching: increment_cache() - Got exception from REDIS %s, Writing value=%s",
|
||||
str(e),
|
||||
value,
|
||||
log_redis_failure(
|
||||
verbose_logger, logging.ERROR, "LiteLLM Redis Caching: increment_cache() - Got exception from REDIS", e
|
||||
)
|
||||
raise e
|
||||
|
||||
|
|
@ -1010,11 +1055,8 @@ class RedisCache(BaseCache):
|
|||
call_type=f"async_set_cache <- {_get_call_stack_info()}",
|
||||
)
|
||||
)
|
||||
verbose_logger.error(
|
||||
"LiteLLM Redis Caching: async set() - Got exception from REDIS %s, key=%r, value=%r",
|
||||
str(e),
|
||||
key,
|
||||
value,
|
||||
log_redis_failure(
|
||||
verbose_logger, logging.ERROR, "LiteLLM Redis Caching: async set() - Got exception from REDIS", e
|
||||
)
|
||||
raise e
|
||||
|
||||
|
|
@ -1062,10 +1104,8 @@ class RedisCache(BaseCache):
|
|||
event_metadata={"key": key},
|
||||
)
|
||||
)
|
||||
verbose_logger.error(
|
||||
"LiteLLM Redis Caching: async set() - Got exception from REDIS %s, Writing value=%s",
|
||||
str(e),
|
||||
value,
|
||||
log_redis_failure(
|
||||
verbose_logger, logging.ERROR, "LiteLLM Redis Caching: async set() - Got exception from REDIS", e
|
||||
)
|
||||
_record_swallowed_redis_failure(self._circuit_breaker, e)
|
||||
|
||||
|
|
@ -1112,7 +1152,6 @@ class RedisCache(BaseCache):
|
|||
start_time: Final = time.time()
|
||||
|
||||
print_verbose(f"Set Async Redis Cache: key list: {cache_list}\nttl={ttl}, redis_version={self.redis_version}")
|
||||
cache_value: Final = None
|
||||
try:
|
||||
async with _redis_client.pipeline(transaction=False) as pipe:
|
||||
results: Final = await self._pipeline_helper(pipe, cache_list, ttl)
|
||||
|
|
@ -1149,10 +1188,11 @@ class RedisCache(BaseCache):
|
|||
)
|
||||
)
|
||||
|
||||
verbose_logger.error(
|
||||
"LiteLLM Redis Caching: async set_cache_pipeline() - Got exception from REDIS %s, Writing value=%s",
|
||||
str(e),
|
||||
cache_value,
|
||||
log_redis_failure(
|
||||
verbose_logger,
|
||||
logging.ERROR,
|
||||
"LiteLLM Redis Caching: async set_cache_pipeline() - Got exception from REDIS",
|
||||
e,
|
||||
)
|
||||
_record_swallowed_redis_failure(self._circuit_breaker, e)
|
||||
|
||||
|
|
@ -1191,8 +1231,11 @@ class RedisCache(BaseCache):
|
|||
end_time=time.time(),
|
||||
)
|
||||
)
|
||||
verbose_logger.error(
|
||||
"LiteLLM Redis Caching: async_set_cache_pipeline_with_ttls() - Got exception from REDIS %s", str(e)
|
||||
log_redis_failure(
|
||||
verbose_logger,
|
||||
logging.ERROR,
|
||||
"LiteLLM Redis Caching: async_set_cache_pipeline_with_ttls() - Got exception from REDIS",
|
||||
e,
|
||||
)
|
||||
_record_swallowed_redis_failure(self._circuit_breaker, e)
|
||||
|
||||
|
|
@ -1235,10 +1278,8 @@ class RedisCache(BaseCache):
|
|||
)
|
||||
)
|
||||
# NON blocking - notify users Redis is throwing an exception
|
||||
verbose_logger.error(
|
||||
"LiteLLM Redis Caching: async set() - Got exception from REDIS %s, Writing value=%s",
|
||||
str(e),
|
||||
value,
|
||||
log_redis_failure(
|
||||
verbose_logger, logging.ERROR, "LiteLLM Redis Caching: async set() - Got exception from REDIS", e
|
||||
)
|
||||
raise e
|
||||
|
||||
|
|
@ -1274,10 +1315,11 @@ class RedisCache(BaseCache):
|
|||
)
|
||||
)
|
||||
# NON blocking - notify users Redis is throwing an exception
|
||||
verbose_logger.error(
|
||||
"LiteLLM Redis Caching: async set_cache_sadd() - Got exception from REDIS %s, Writing value=%s",
|
||||
str(e),
|
||||
value,
|
||||
log_redis_failure(
|
||||
verbose_logger,
|
||||
logging.ERROR,
|
||||
"LiteLLM Redis Caching: async set_cache_sadd() - Got exception from REDIS",
|
||||
e,
|
||||
)
|
||||
_record_swallowed_redis_failure(self._circuit_breaker, e)
|
||||
|
||||
|
|
@ -1359,10 +1401,11 @@ class RedisCache(BaseCache):
|
|||
parent_otel_span=parent_otel_span,
|
||||
)
|
||||
)
|
||||
verbose_logger.error(
|
||||
"LiteLLM Redis Caching: async async_increment() - Got exception from REDIS %s, Writing value=%s",
|
||||
str(e),
|
||||
value,
|
||||
log_redis_failure(
|
||||
verbose_logger,
|
||||
logging.ERROR,
|
||||
"LiteLLM Redis Caching: async async_increment() - Got exception from REDIS",
|
||||
e,
|
||||
)
|
||||
raise e
|
||||
|
||||
|
|
@ -1448,7 +1491,9 @@ class RedisCache(BaseCache):
|
|||
print_verbose(f"Got Redis Cache: key: {key}, cached_response {cached_response}")
|
||||
return self._get_cache_logic(cached_response=cached_response)
|
||||
except Exception as e:
|
||||
verbose_logger.error("litellm.caching.caching: get() - Got exception from REDIS: %s", e)
|
||||
log_redis_failure(
|
||||
verbose_logger, logging.ERROR, "litellm.caching.caching: get() - Got exception from REDIS", e
|
||||
)
|
||||
_record_swallowed_redis_failure(self._circuit_breaker, e)
|
||||
|
||||
def _run_redis_mget_operation(self, keys: list[str]) -> Sequence[bytes | str | None]:
|
||||
|
|
@ -1526,7 +1571,7 @@ class RedisCache(BaseCache):
|
|||
end_time=failed_at,
|
||||
parent_otel_span=parent_otel_span,
|
||||
)
|
||||
verbose_logger.error("Error occurred in batch get cache - %s", e)
|
||||
log_redis_failure(verbose_logger, logging.ERROR, "Error occurred in batch get cache", e)
|
||||
_record_swallowed_redis_failure(self._circuit_breaker, e)
|
||||
return key_value_dict
|
||||
|
||||
|
|
@ -1645,7 +1690,7 @@ class RedisCache(BaseCache):
|
|||
parent_otel_span=parent_otel_span,
|
||||
)
|
||||
)
|
||||
verbose_logger.error("Error occurred in async batch get cache - %s", e)
|
||||
log_redis_failure(verbose_logger, logging.ERROR, "Error occurred in async batch get cache", e)
|
||||
_record_swallowed_redis_failure(self._circuit_breaker, e)
|
||||
return key_value_dict
|
||||
|
||||
|
|
@ -1870,9 +1915,11 @@ class RedisCache(BaseCache):
|
|||
parent_otel_span=_get_parent_otel_span_from_kwargs(kwargs),
|
||||
)
|
||||
)
|
||||
verbose_logger.error(
|
||||
"LiteLLM Redis Caching: async increment_pipeline() - Got exception from REDIS %s",
|
||||
str(e),
|
||||
log_redis_failure(
|
||||
verbose_logger,
|
||||
logging.ERROR,
|
||||
"LiteLLM Redis Caching: async increment_pipeline() - Got exception from REDIS",
|
||||
e,
|
||||
)
|
||||
raise e
|
||||
|
||||
|
|
@ -1949,7 +1996,7 @@ class RedisCache(BaseCache):
|
|||
call_type=f"async_rpush <- {_get_call_stack_info()}",
|
||||
)
|
||||
)
|
||||
verbose_logger.error("LiteLLM Redis Cache RPUSH: - Got exception from REDIS : %s", e)
|
||||
log_redis_failure(verbose_logger, logging.ERROR, "LiteLLM Redis Cache RPUSH: - Got exception from REDIS", e)
|
||||
raise e
|
||||
|
||||
async def _pipeline_rpush_helper(
|
||||
|
|
@ -2017,9 +2064,11 @@ class RedisCache(BaseCache):
|
|||
call_type=f"async_rpush_pipeline <- {_get_call_stack_info()}",
|
||||
)
|
||||
)
|
||||
verbose_logger.error(
|
||||
"LiteLLM Redis Caching: async_rpush_pipeline() - Got exception from REDIS %s",
|
||||
str(e),
|
||||
log_redis_failure(
|
||||
verbose_logger,
|
||||
logging.ERROR,
|
||||
"LiteLLM Redis Caching: async_rpush_pipeline() - Got exception from REDIS",
|
||||
e,
|
||||
)
|
||||
raise e
|
||||
|
||||
|
|
@ -2095,7 +2144,7 @@ class RedisCache(BaseCache):
|
|||
call_type=f"async_lpop <- {_get_call_stack_info()}",
|
||||
)
|
||||
)
|
||||
verbose_logger.error("LiteLLM Redis Cache LPOP: - Got exception from REDIS : %s", e)
|
||||
log_redis_failure(verbose_logger, logging.ERROR, "LiteLLM Redis Cache LPOP: - Got exception from REDIS", e)
|
||||
raise e
|
||||
|
||||
async def _pipeline_lpop_helper(
|
||||
|
|
@ -2206,8 +2255,10 @@ class RedisCache(BaseCache):
|
|||
call_type=f"async_lpop_pipeline <- {_get_call_stack_info()}",
|
||||
)
|
||||
)
|
||||
verbose_logger.error(
|
||||
"LiteLLM Redis Caching: async_lpop_pipeline() - Got exception from REDIS %s",
|
||||
str(e),
|
||||
log_redis_failure(
|
||||
verbose_logger,
|
||||
logging.ERROR,
|
||||
"LiteLLM Redis Caching: async_lpop_pipeline() - Got exception from REDIS",
|
||||
e,
|
||||
)
|
||||
raise e
|
||||
|
|
|
|||
|
|
@ -205,21 +205,35 @@ def _extract_anthropic_tool_exchange_spans(
|
|||
return spans, None
|
||||
|
||||
|
||||
def _message_has_cache_control(message: Mapping[str, object]) -> bool:
|
||||
if message.get("cache_control") is not None:
|
||||
return True
|
||||
content: Final = message.get("content")
|
||||
if isinstance(content, list):
|
||||
return any(isinstance(part, Mapping) and part.get("cache_control") is not None for part in content)
|
||||
return False
|
||||
|
||||
|
||||
def get_protected_indices(messages: Sequence[Mapping[str, object]]) -> tuple[int, ...]:
|
||||
"""
|
||||
Return indices of messages that must never be compressed:
|
||||
- All system messages
|
||||
- The last user message
|
||||
- The last assistant message
|
||||
- Any message carrying an Anthropic cache_control breakpoint
|
||||
|
||||
The last user message is what the model is being asked to act on right now,
|
||||
so compressing it replaces the live instruction with a marker. Compression
|
||||
guardrails share this policy; see the Headroom guardrail.
|
||||
guardrails share this policy; see the Headroom guardrail. A cache_control
|
||||
breakpoint pins the provider's prompt-cache prefix to that row's exact
|
||||
bytes, so rewriting a marked row anywhere in history turns the next
|
||||
request's cache read into a cache write.
|
||||
"""
|
||||
system_indices: Final = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "system")
|
||||
last_user: Final = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "user")[-1:]
|
||||
last_assistant = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "assistant")[-1:]
|
||||
return system_indices + last_user + last_assistant
|
||||
assistant_indices: Final = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "assistant")
|
||||
cache_control_indices: Final = tuple(index for index, msg in enumerate(messages) if _message_has_cache_control(msg))
|
||||
return tuple(dict.fromkeys(system_indices + last_user + assistant_indices[-1:] + cache_control_indices))
|
||||
|
||||
|
||||
def _combine_scores(
|
||||
|
|
@ -421,7 +435,7 @@ def compress(
|
|||
combined_scores = bm25_scores
|
||||
|
||||
# Protected messages are never compressed
|
||||
protected_indices: Final = get_protected_indices(normalized_messages)
|
||||
protected_indices: Final = get_protected_indices(original_messages)
|
||||
kept_indices: set[int] = set(protected_indices)
|
||||
|
||||
tool_exchange_spans: list[set[int]] = []
|
||||
|
|
|
|||
|
|
@ -227,6 +227,9 @@ PRE_CALL_EXECUTED_GUARDRAILS_KEY: Final = "_pre_call_executed_guardrails"
|
|||
# Attribute stamped on log_guardrail_information wrappers so __init_subclass__ does not wrap them again
|
||||
LOGS_GUARDRAIL_INFORMATION_MARKER: Final = "_litellm_logs_guardrail_information"
|
||||
|
||||
# llm_provider stamped on proxy-side rate limit errors when the model resolves to no deployment
|
||||
PROXY_LLM_PROVIDER_FALLBACK: Final = "litellm_proxy"
|
||||
|
||||
# Generic fallback for unknown models
|
||||
DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET: Final = int(
|
||||
os.getenv("DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET", 128)
|
||||
|
|
@ -311,6 +314,12 @@ REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS: Final = float(
|
|||
# RFC 6455 caps the close frame payload at 125 bytes, 2 of which carry the status code
|
||||
WEBSOCKET_CLOSE_REASON_MAX_BYTES: Final = 123
|
||||
|
||||
BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY: Final = "litellm.bedrock_realtime.pending_session_update"
|
||||
BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY: Final = "litellm.bedrock_realtime.session_committed"
|
||||
BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY: Final = "litellm.bedrock_realtime.committed_failure"
|
||||
REALTIME_SESSION_SUCCESS_LOGGED_KEY: Final = "realtime_session_success_logged"
|
||||
REALTIME_SESSION_FAILURE_LOGGED_KEY: Final = "realtime_session_failure_logged"
|
||||
|
||||
# SSL/TLS cipher configuration for faster handshakes
|
||||
# Strategy: Strongly prefer fast modern ciphers, but allow fallback to commonly supported ones
|
||||
# This balances performance with broad compatibility
|
||||
|
|
@ -461,6 +470,7 @@ REDIS_CIRCUIT_BREAKER_ENABLED: Final = os.getenv("REDIS_CIRCUIT_BREAKER_ENABLED"
|
|||
# minimum seconds a timeout-only failure streak must span before it can open the breaker,
|
||||
# so one event-loop stall timing out many queued calls at once does not trip it
|
||||
REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION: Final = float(os.getenv("REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION", 5.0))
|
||||
REDIS_TIMEOUT_LOG_INTERVAL: Final = float(os.getenv("REDIS_TIMEOUT_LOG_INTERVAL", "5.0"))
|
||||
# Seconds of idle before a Redis cluster connection is validated with a PING and
|
||||
# reconnected if dead, so a connection silently dropped by a cluster restart
|
||||
# (e.g. ElastiCache Serverless maintenance) is not reused while broken
|
||||
|
|
|
|||
|
|
@ -2278,6 +2278,19 @@ def default_video_cost_calculator(
|
|||
return 0.0
|
||||
|
||||
|
||||
def _batch_rate(
|
||||
model_info: ModelInfo,
|
||||
key: Literal[
|
||||
"input_cost_per_audio_token_batches",
|
||||
"input_cost_per_image_token_batches",
|
||||
"input_cost_per_video_token_batches",
|
||||
],
|
||||
fallback: float,
|
||||
) -> float:
|
||||
rate: Final = model_info.get(key)
|
||||
return fallback if rate is None else rate
|
||||
|
||||
|
||||
def batch_cost_calculator(
|
||||
usage: Usage,
|
||||
model: str,
|
||||
|
|
@ -2337,7 +2350,29 @@ def batch_cost_calculator(
|
|||
total_prompt_cost = 0.0
|
||||
total_completion_cost = 0.0
|
||||
if input_cost_per_token_batches is not None:
|
||||
total_prompt_cost = usage.prompt_tokens * input_cost_per_token_batches
|
||||
batch_details: Final = parse_prompt_tokens_details(usage)
|
||||
audio_tokens, image_tokens, video_tokens = (
|
||||
batch_details["audio_tokens"],
|
||||
batch_details["image_tokens"],
|
||||
batch_details["video_tokens"],
|
||||
)
|
||||
modality_rates: Final = (
|
||||
_batch_rate(model_info, "input_cost_per_audio_token_batches", input_cost_per_token_batches),
|
||||
_batch_rate(model_info, "input_cost_per_image_token_batches", input_cost_per_token_batches),
|
||||
_batch_rate(model_info, "input_cost_per_video_token_batches", input_cost_per_token_batches),
|
||||
)
|
||||
total_prompt_cost = sum(
|
||||
tokens * rate
|
||||
for tokens, rate in zip(
|
||||
(
|
||||
max((usage.prompt_tokens or 0) - audio_tokens - image_tokens - video_tokens, 0),
|
||||
audio_tokens,
|
||||
image_tokens,
|
||||
video_tokens,
|
||||
),
|
||||
(input_cost_per_token_batches, *modality_rates),
|
||||
)
|
||||
)
|
||||
elif input_cost_per_token:
|
||||
details: Final = parse_prompt_tokens_details(usage)
|
||||
cache_read_tokens: Final = details["cache_hit_tokens"]
|
||||
|
|
|
|||
|
|
@ -379,7 +379,7 @@ class ArizePhoenixPromptManager(CustomPromptManagement):
|
|||
"""Reload prompts from Arize Phoenix."""
|
||||
if self.prompt_id:
|
||||
self._prompt_manager = None # Reset to force reload
|
||||
self.prompt_manager # This will trigger reload
|
||||
_ = self.prompt_manager # access triggers lazy reload
|
||||
|
||||
def should_run_prompt_management(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -406,7 +406,7 @@ class BitBucketPromptManager(CustomPromptManagement):
|
|||
"""Reload prompts from BitBucket."""
|
||||
if self.prompt_id:
|
||||
self._prompt_manager = None # Reset to force reload
|
||||
self.prompt_manager # This will trigger reload
|
||||
_ = self.prompt_manager # access triggers lazy reload
|
||||
|
||||
def should_run_prompt_management(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -884,7 +884,9 @@ class CustomGuardrail(CustomLogger):
|
|||
"""logging_only: run apply_guardrail on copies of the logged request/response and record the verdict."""
|
||||
from litellm.llms import get_guardrail_translation_mapping
|
||||
|
||||
if not self.uses_apply_guardrail_interface() or self.use_native_lifecycle_hooks:
|
||||
if not self.uses_apply_guardrail_interface():
|
||||
return kwargs, result
|
||||
if not self._event_hook_is_event_type(GuardrailEventHooks.logging_only):
|
||||
return kwargs, result
|
||||
try:
|
||||
translation: Final = get_guardrail_translation_mapping(CallTypes(call_type))()
|
||||
|
|
@ -901,8 +903,18 @@ class CustomGuardrail(CustomLogger):
|
|||
for key, value in (litellm_params.get("metadata") or {}).items()
|
||||
if key != "standard_logging_guardrail_information"
|
||||
}
|
||||
response: Final = (
|
||||
kwargs.get("async_complete_streaming_response") or kwargs.get("complete_streaming_response") or result
|
||||
)
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
output_translation: Final = (
|
||||
get_guardrail_translation_mapping(CallTypes.acompletion)()
|
||||
if isinstance(response, ModelResponse)
|
||||
else translation
|
||||
)
|
||||
try:
|
||||
await self._scan_logged_call(kwargs, result, translation, scratch_metadata)
|
||||
await self._scan_logged_call(kwargs, response, translation, output_translation, scratch_metadata)
|
||||
except Exception as e:
|
||||
verbose_logger.warning("Guardrail %s: logging_only scan raised: %s", self.guardrail_name, e)
|
||||
recorded: Final = scratch_metadata.get("standard_logging_guardrail_information")
|
||||
|
|
@ -919,8 +931,9 @@ class CustomGuardrail(CustomLogger):
|
|||
async def _scan_logged_call(
|
||||
self,
|
||||
kwargs: dict, # mutable-ok: CustomLogger.async_logging_hook contract
|
||||
result: object,
|
||||
response: object | None,
|
||||
translation: "BaseTranslation",
|
||||
output_translation: "BaseTranslation",
|
||||
scratch_metadata: dict, # mutable-ok: apply_guardrail records its verdict into request metadata
|
||||
) -> None:
|
||||
optional_params: Final = kwargs.get("optional_params") or {}
|
||||
|
|
@ -934,8 +947,10 @@ class CustomGuardrail(CustomLogger):
|
|||
"metadata": scratch_metadata,
|
||||
}
|
||||
await translation.process_input_messages(data=scratch_request, guardrail_to_apply=self)
|
||||
await translation.process_output_response(
|
||||
response=copy.deepcopy(result), guardrail_to_apply=self, request_data=scratch_request
|
||||
if response is None:
|
||||
return
|
||||
await output_translation.process_output_response(
|
||||
response=copy.deepcopy(response), guardrail_to_apply=self, request_data=scratch_request
|
||||
)
|
||||
|
||||
def supports_scan_only_tool_results(self) -> bool:
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ from typing import Final
|
|||
|
||||
from litellm.integrations.otel.mappers.base import AttributeMap, AttrValue, SpanData
|
||||
from litellm.integrations.otel.mappers.utils import (
|
||||
MAX_MESSAGE_ATTRS_PER_SPAN,
|
||||
MAX_TOOL_DEFINITION_ATTRS_PER_SPAN,
|
||||
collect,
|
||||
drop_none,
|
||||
|
|
@ -26,6 +27,8 @@ from litellm.integrations.otel.model.payloads import (
|
|||
ToolDefinition,
|
||||
)
|
||||
|
||||
_MAX_INDEXED_MESSAGES: Final = MAX_MESSAGE_ATTRS_PER_SPAN // 2
|
||||
|
||||
|
||||
class OpenInferenceMapper:
|
||||
"""Emits OpenInference attributes for LLM_CALL spans.
|
||||
|
|
@ -84,27 +87,44 @@ class OpenInferenceMapper:
|
|||
return {}
|
||||
|
||||
def _llm_call(self, data: LLMCallSpanData) -> AttributeMap:
|
||||
outputs: Final = output_messages(data)
|
||||
indexed_in, indexed_out = self._indexed_split(len(data.messages_in), len(outputs))
|
||||
return {
|
||||
**collect(self._LLM_CALL_ATTRS, data),
|
||||
**collect(self._BLOB_ATTRS, data),
|
||||
**self._messages("llm.input_messages", "input.value", data.messages_in),
|
||||
**self._messages("llm.output_messages", "output.value", output_messages(data)),
|
||||
**self._messages(
|
||||
"llm.input_messages",
|
||||
"input.value",
|
||||
data.messages_in,
|
||||
self._prompt_positions(len(data.messages_in), indexed_in),
|
||||
),
|
||||
**self._messages("llm.output_messages", "output.value", outputs, range(indexed_out)),
|
||||
**self._tools(data),
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _messages(prefix: str, value_key: str, messages: Sequence[object]) -> AttributeMap:
|
||||
"""Per-message ``{prefix}.{idx}.message.*`` keys + the ``value_key`` blob."""
|
||||
def _indexed_split(inputs: int, outputs: int) -> tuple[int, int]:
|
||||
"""Prompt and response share one allowance; the response is reserved at least half of it."""
|
||||
indexed_out: Final = min(outputs, max(_MAX_INDEXED_MESSAGES // 2, _MAX_INDEXED_MESSAGES - inputs))
|
||||
return _MAX_INDEXED_MESSAGES - indexed_out, indexed_out
|
||||
|
||||
@staticmethod
|
||||
def _prompt_positions(total: int, indexed: int) -> tuple[int, ...]:
|
||||
"""Prompt messages that get per-index attributes: message 0 and the most recent turns."""
|
||||
if total <= indexed:
|
||||
return tuple(range(total))
|
||||
return (0, *range(total - indexed + 1, total))
|
||||
|
||||
@staticmethod
|
||||
def _messages(prefix: str, value_key: str, messages: Sequence[object], positions: Sequence[int]) -> AttributeMap:
|
||||
"""``{prefix}.{idx}.message.*`` keys for the messages at ``positions`` + the ``value_key`` blob of all."""
|
||||
parsed: Final = [(m.get("role") if isinstance(m, dict) else None, message_content(m)) for m in messages]
|
||||
attrs: Final = drop_none(
|
||||
{
|
||||
key: value
|
||||
for idx, (role, content) in enumerate(parsed)
|
||||
for idx, (role, content) in ((idx, parsed[idx]) for idx in positions)
|
||||
for key, value in (
|
||||
(
|
||||
f"{prefix}.{idx}.message.role",
|
||||
role if isinstance(role, str) else None,
|
||||
),
|
||||
(f"{prefix}.{idx}.message.role", role if isinstance(role, str) else None),
|
||||
(f"{prefix}.{idx}.message.content", content),
|
||||
)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -32,6 +32,14 @@ core telemetry no matter how many vocabularies are configured.
|
|||
"""
|
||||
|
||||
|
||||
MAX_MESSAGE_ATTRS_PER_SPAN: Final = DEFAULT_SPAN_ATTRIBUTE_LIMIT // 8
|
||||
"""Span-wide ceiling on per-index chat message attributes, prompt and response together.
|
||||
|
||||
An eighth is the largest share that still fits beside the tool ceiling and the core
|
||||
of every vocabulary at once. The complete conversation still rides the JSON blobs.
|
||||
"""
|
||||
|
||||
|
||||
def tool_attr_budget(vocabularies: int) -> int:
|
||||
"""Split the span-wide tool-definition ceiling across active vocabularies."""
|
||||
return MAX_TOOL_DEFINITION_ATTRS_PER_SPAN // max(vocabularies, 1)
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ from pydantic import BaseModel
|
|||
|
||||
import litellm
|
||||
from litellm._logging import print_verbose, verbose_logger
|
||||
from litellm.constants import PROXY_LLM_PROVIDER_FALLBACK
|
||||
from litellm.exceptions import (
|
||||
validate_rate_limit_category,
|
||||
validate_rate_limit_type,
|
||||
|
|
@ -2581,6 +2582,15 @@ class PrometheusLogger(CustomLogger):
|
|||
)
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _extract_api_provider_from_exception(exception: Exception) -> str | None:
|
||||
if not isinstance(exception, litellm.exceptions.RateLimitError):
|
||||
return None
|
||||
llm_provider: Final = exception.llm_provider
|
||||
if not llm_provider or llm_provider == PROXY_LLM_PROVIDER_FALLBACK:
|
||||
return None
|
||||
return llm_provider
|
||||
|
||||
async def async_post_call_failure_hook(
|
||||
self,
|
||||
request_data: dict,
|
||||
|
|
@ -2616,7 +2626,9 @@ class PrometheusLogger(CustomLogger):
|
|||
_metadata: Final = request_data.get("metadata", {}) or {}
|
||||
model_id: Final = _metadata.get("model_info", {}).get("id") or request_data.get("model_info", {}).get("id")
|
||||
rate_limit_category, rate_limit_type = self._extract_rate_limit_labels(original_exception)
|
||||
api_provider: Final = self._extract_api_provider_from_request_data(request_data)
|
||||
api_provider: Final = self._extract_api_provider_from_request_data(
|
||||
request_data
|
||||
) or self._extract_api_provider_from_exception(original_exception)
|
||||
enum_values: Final = UserAPIKeyLabelValues(
|
||||
end_user=user_api_key_dict.end_user_id,
|
||||
user=user_api_key_dict.user_id,
|
||||
|
|
|
|||
|
|
@ -4,18 +4,10 @@ imported_openAIResponse = True
|
|||
try:
|
||||
import io
|
||||
import logging
|
||||
import sys
|
||||
from typing import Any, TypeVar
|
||||
from typing import Any, Literal, Protocol, TypeVar
|
||||
|
||||
from wandb.sdk.data_types import trace_tree
|
||||
|
||||
if sys.version_info >= (3, 8):
|
||||
from typing import Literal, Protocol
|
||||
else:
|
||||
from typing import Literal
|
||||
|
||||
from typing_extensions import Protocol
|
||||
|
||||
logger: Final = logging.getLogger(__name__)
|
||||
|
||||
K = TypeVar("K", bound=str)
|
||||
|
|
|
|||
|
|
@ -303,6 +303,16 @@ def get_metadata_variable_name_from_kwargs(
|
|||
return "litellm_metadata" if "litellm_metadata" in kwargs else "metadata"
|
||||
|
||||
|
||||
def max_retries_per_request_hit(kwargs: Mapping[str, object], num_retries_per_request: int | None) -> bool:
|
||||
if num_retries_per_request is None:
|
||||
return False
|
||||
metadata: Final = kwargs.get(get_metadata_variable_name_from_kwargs(kwargs))
|
||||
if not isinstance(metadata, Mapping):
|
||||
return False
|
||||
attempted_retries: Final = metadata.get("attempted_retries")
|
||||
return type(attempted_retries) is int and 0 < attempted_retries and num_retries_per_request <= attempted_retries
|
||||
|
||||
|
||||
def get_or_create_metadata_bucket(
|
||||
request_data: dict,
|
||||
) -> tuple[Literal["metadata", "litellm_metadata"], dict]:
|
||||
|
|
|
|||
|
|
@ -36,7 +36,7 @@ class CoroutineChecker:
|
|||
target = callback
|
||||
if not inspect.isfunction(target) and not inspect.ismethod(target):
|
||||
try:
|
||||
call_attr: Final = getattr(target, "__call__", None)
|
||||
call_attr: Final = getattr(target, "__call__", None) # noqa: B004 # value unwrap so iscoroutinefunction sees through functors
|
||||
if call_attr is not None:
|
||||
target = call_attr
|
||||
except Exception:
|
||||
|
|
|
|||
|
|
@ -34,6 +34,11 @@ rules never mix the two and never use ``extends``. A rule whose
|
|||
Rules are only consulted after exact and case-insensitive lookups miss, so an
|
||||
exact cost-map entry always takes precedence over any rule.
|
||||
|
||||
Rules flagged with ``fill_missing_for_providers: [..]`` also fill only keys
|
||||
missing from an exact cost-map entry when the entry's ``litellm_provider`` is
|
||||
listed, while values already present on the entry win on conflict. Only flagged
|
||||
capability rules participate in this fill; routing rules never do.
|
||||
|
||||
Patterns are matched case-insensitively with ``re.search`` and are not implicitly
|
||||
anchored: a rule must include ``^`` and ``$`` to bind to the whole model name,
|
||||
otherwise it matches as a substring. Keeping anchoring in the regex makes the rule
|
||||
|
|
@ -46,17 +51,19 @@ Rules are compiled and classified once, at install time. The match functions are
|
|||
O(number of rules); callers must only invoke them on a cache miss.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import re
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
|
||||
verbose_logger: Final = logging.getLogger("LiteLLM")
|
||||
NAME_FIELD: Final = "name"
|
||||
PATTERN_FIELD: Final = "pattern"
|
||||
MODEL_INFO_FIELD: Final = "model_info"
|
||||
PROVIDER_KEY: Final = "litellm_provider"
|
||||
LEGACY_EXTENDS_FIELD: Final = "extends"
|
||||
FILL_MISSING_FOR_PROVIDERS_FIELD: Final = "fill_missing_for_providers"
|
||||
|
||||
|
||||
def _resolve_legacy_extends(rules: list) -> list:
|
||||
|
|
@ -98,11 +105,28 @@ class _RoutingRule:
|
|||
class _CapabilityRule:
|
||||
pattern: re.Pattern
|
||||
model_info: dict
|
||||
fill_missing_for_providers: frozenset[str]
|
||||
|
||||
|
||||
_CompiledRule = _RoutingRule | _CapabilityRule
|
||||
|
||||
|
||||
def _parse_fill_missing_for_providers(rule: Mapping[str, object], pattern_label: object) -> frozenset[str] | None:
|
||||
if FILL_MISSING_FOR_PROVIDERS_FIELD not in rule:
|
||||
return frozenset()
|
||||
raw_fill_missing_for_providers: Final = rule.get(FILL_MISSING_FOR_PROVIDERS_FIELD)
|
||||
if not isinstance(raw_fill_missing_for_providers, (list, tuple)) or not all(
|
||||
isinstance(provider, str) for provider in raw_fill_missing_for_providers
|
||||
):
|
||||
verbose_logger.warning(
|
||||
"LiteLLM: skipping malformed fallback generalization rule %s ('%s' must be a list of provider strings).",
|
||||
rule.get(NAME_FIELD, pattern_label),
|
||||
FILL_MISSING_FOR_PROVIDERS_FIELD,
|
||||
)
|
||||
return None
|
||||
return frozenset(raw_fill_missing_for_providers)
|
||||
|
||||
|
||||
def _compile_rule(rule: object) -> tuple[_CompiledRule, ...]:
|
||||
if not isinstance(rule, dict):
|
||||
return ()
|
||||
|
|
@ -125,8 +149,17 @@ def _compile_rule(rule: object) -> tuple[_CompiledRule, ...]:
|
|||
e,
|
||||
)
|
||||
return ()
|
||||
fill_missing_for_providers: Final = _parse_fill_missing_for_providers(rule, pattern)
|
||||
if fill_missing_for_providers is None:
|
||||
return ()
|
||||
if PROVIDER_KEY not in model_info:
|
||||
return (_CapabilityRule(pattern=compiled, model_info=model_info),)
|
||||
return (
|
||||
_CapabilityRule(
|
||||
pattern=compiled,
|
||||
model_info=model_info,
|
||||
fill_missing_for_providers=fill_missing_for_providers,
|
||||
),
|
||||
)
|
||||
provider: Final = model_info[PROVIDER_KEY]
|
||||
if not isinstance(provider, str):
|
||||
verbose_logger.warning(
|
||||
|
|
@ -140,7 +173,11 @@ def _compile_rule(rule: object) -> tuple[_CompiledRule, ...]:
|
|||
return (_RoutingRule(pattern=compiled, provider=provider),)
|
||||
return (
|
||||
_RoutingRule(pattern=compiled, provider=provider),
|
||||
_CapabilityRule(pattern=compiled, model_info=model_info),
|
||||
_CapabilityRule(
|
||||
pattern=compiled,
|
||||
model_info=model_info,
|
||||
fill_missing_for_providers=fill_missing_for_providers,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -151,6 +188,7 @@ class _FallbackGeneralizations:
|
|||
self.rules: list = []
|
||||
self.routing_rules: tuple = ()
|
||||
self.capability_rules: tuple = ()
|
||||
self.fill_missing_rules: tuple[_CapabilityRule, ...] = ()
|
||||
|
||||
def set_rules(self, rules: list | None) -> None:
|
||||
installed: Final = rules if isinstance(rules, list) else []
|
||||
|
|
@ -158,6 +196,7 @@ class _FallbackGeneralizations:
|
|||
self.rules = installed
|
||||
self.routing_rules = tuple(rule for rule in compiled if isinstance(rule, _RoutingRule))
|
||||
self.capability_rules = tuple(rule for rule in compiled if isinstance(rule, _CapabilityRule))
|
||||
self.fill_missing_rules = tuple(rule for rule in self.capability_rules if rule.fill_missing_for_providers)
|
||||
|
||||
def match_routing(self, model: str) -> str | None:
|
||||
if not model:
|
||||
|
|
@ -175,6 +214,21 @@ class _FallbackGeneralizations:
|
|||
return None
|
||||
return {key: value for model_info in matched for key, value in model_info.items()}
|
||||
|
||||
def match_fill_missing(self, model: str, provider: str) -> Mapping[str, object] | None:
|
||||
if not model or not provider:
|
||||
return None
|
||||
matched = tuple(
|
||||
rule.model_info
|
||||
for rule in self.fill_missing_rules
|
||||
if provider in rule.fill_missing_for_providers and rule.pattern.search(model) is not None
|
||||
)
|
||||
if not matched:
|
||||
return None
|
||||
fill_missing: Final[Mapping[str, object]] = {
|
||||
key: value for model_info in matched for key, value in model_info.items() if key != PROVIDER_KEY
|
||||
}
|
||||
return fill_missing or None
|
||||
|
||||
|
||||
_registry: Final = _FallbackGeneralizations()
|
||||
|
||||
|
|
@ -210,3 +264,14 @@ def match_capability_generalizations(model: str) -> dict | None:
|
|||
capability rule matches. O(number of rules); only call once exact lookups have missed.
|
||||
"""
|
||||
return _registry.match_capabilities(model)
|
||||
|
||||
|
||||
def match_fill_missing_generalizations(model: str, provider: str) -> Mapping[str, object] | None:
|
||||
"""Return flagged capability rules matching ``model`` for ``provider``.
|
||||
|
||||
Later rules override earlier ones on key conflicts. Only rules listing
|
||||
``provider`` in ``fill_missing_for_providers`` contribute. Returns ``None``
|
||||
when no flagged rule matches. O(number of rules); only call once exact
|
||||
lookups have matched.
|
||||
"""
|
||||
return _registry.match_fill_missing(model, provider)
|
||||
|
|
|
|||
|
|
@ -956,12 +956,16 @@ def _calculate_input_cost(
|
|||
)
|
||||
|
||||
### AUDIO COST
|
||||
if prompt_tokens_details["audio_tokens"]:
|
||||
if prompt_tokens_details["audio_tokens"] and not (
|
||||
prompt_tokens_details["audio_length_seconds"] and model_info.get("input_cost_per_audio_per_second") is not None
|
||||
):
|
||||
audio_cost_key: Final = _get_service_tier_cost_key("input_cost_per_audio_token", service_tier)
|
||||
prompt_cost += calculate_cost_component(model_info, audio_cost_key, prompt_tokens_details["audio_tokens"])
|
||||
|
||||
### IMAGE TOKEN COST
|
||||
if prompt_tokens_details["image_tokens"]:
|
||||
if prompt_tokens_details["image_tokens"] and not (
|
||||
prompt_tokens_details["image_count"] and model_info.get("input_cost_per_image") is not None
|
||||
):
|
||||
# For image token costs:
|
||||
# First check if input_cost_per_image_token is available. If not, default to generic input_cost_per_token.
|
||||
image_token_cost_key = "input_cost_per_image_token"
|
||||
|
|
@ -970,7 +974,9 @@ def _calculate_input_cost(
|
|||
prompt_cost += calculate_cost_component(model_info, image_token_cost_key, prompt_tokens_details["image_tokens"])
|
||||
|
||||
### VIDEO TOKEN COST
|
||||
if prompt_tokens_details["video_tokens"]:
|
||||
if prompt_tokens_details["video_tokens"] and not (
|
||||
prompt_tokens_details["video_length_seconds"] and model_info.get("input_cost_per_video_per_second") is not None
|
||||
):
|
||||
video_token_cost_key = "input_cost_per_video_token"
|
||||
if model_info.get(video_token_cost_key) is None:
|
||||
video_token_cost_key = "input_cost_per_token"
|
||||
|
|
|
|||
|
|
@ -1757,7 +1757,7 @@ def convert_to_anthropic_tool_invoke(
|
|||
anthropic_tool_invoke: Final[list[AnthropicMessagesToolUseParam | dict[str, object]]] = []
|
||||
|
||||
for tool in tool_calls:
|
||||
if not get_attribute_or_key(tool, "type") == "function":
|
||||
if get_attribute_or_key(tool, "type") != "function":
|
||||
continue
|
||||
|
||||
tool_id = cast(str, get_attribute_or_key(tool, "id"))
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ from typing_extensions import ReadOnly
|
|||
|
||||
import litellm
|
||||
from litellm._logging import redact_internal_details_from_client_message, verbose_logger
|
||||
from litellm.constants import REALTIME_SESSION_FAILURE_LOGGED_KEY, REALTIME_SESSION_SUCCESS_LOGGED_KEY
|
||||
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
||||
from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig
|
||||
from litellm.types.llms.openai import (
|
||||
|
|
@ -35,9 +36,6 @@ else:
|
|||
CLIENT_CONNECTION_CLASS = Any
|
||||
|
||||
|
||||
REALTIME_SESSION_SUCCESS_LOGGED_KEY: Final = "realtime_session_success_logged"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class BackendClose:
|
||||
code: int
|
||||
|
|
@ -1153,6 +1151,7 @@ class RealTimeStreaming:
|
|||
self._logging_worker.ensure_initialized_and_enqueue(
|
||||
self.logging_obj.dispatch_failure_handlers(error, traceback.format_exc(), prefer_async_handlers=True)
|
||||
)
|
||||
self.logging_obj.model_call_details[REALTIME_SESSION_FAILURE_LOGGED_KEY] = True
|
||||
|
||||
@staticmethod
|
||||
def _detect_beta_header(websocket: ScopedWebSocket) -> bool:
|
||||
|
|
|
|||
|
|
@ -454,7 +454,7 @@ def token_counter(
|
|||
params: Final = _MessageCountParams(model, custom_tokenizer)
|
||||
num_tokens = _count_messages(params, new_messages, use_default_image_token_count, default_token_count)
|
||||
if count_response_tokens is False:
|
||||
includes_system_message: Final = any([message.get("role", None) == "system" for message in new_messages])
|
||||
includes_system_message: Final = any(message.get("role", None) == "system" for message in new_messages)
|
||||
num_tokens += _count_extra(params.count_function, tools, tool_choice, includes_system_message)
|
||||
|
||||
else:
|
||||
|
|
|
|||
|
|
@ -44,6 +44,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import (
|
|||
scoped_structured_message_indices,
|
||||
stream_item_field,
|
||||
stream_item_fingerprint,
|
||||
unappliable_request_rewrite,
|
||||
)
|
||||
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import (
|
||||
AnthropicPassthroughLoggingHandler,
|
||||
|
|
@ -570,6 +571,8 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
preserve_system_messages=has_midturn_system_message,
|
||||
)
|
||||
else:
|
||||
if guardrailed_texts and len(guardrailed_texts) != len(scanned):
|
||||
raise unappliable_request_rewrite(guardrail_to_apply.guardrail_name)
|
||||
# Step 3: Map guardrail responses back to original message structure
|
||||
await self._apply_guardrail_responses_to_input(
|
||||
messages=messages,
|
||||
|
|
|
|||
388
litellm/llms/anthropic/prompt_cache_prediction.py
Normal file
388
litellm/llms/anthropic/prompt_cache_prediction.py
Normal file
|
|
@ -0,0 +1,388 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from itertools import accumulate
|
||||
from types import MappingProxyType
|
||||
from typing import Annotated, Final, Literal, Protocol, TypeAlias
|
||||
|
||||
import httpx
|
||||
from pydantic import BaseModel, ConfigDict, Field, JsonValue, StrictInt, TypeAdapter, ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm.llms.anthropic.common_utils import AnthropicModelInfo, is_anthropic_oauth_key
|
||||
from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler
|
||||
from litellm.llms.anthropic.experimental_pass_through.messages.transformation import DEFAULT_ANTHROPIC_API_VERSION
|
||||
from litellm.types.router import LiteLLM_Params
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
||||
_HEADERS: Final = TypeAdapter(dict[str, str])
|
||||
_counter: Final = AnthropicCountTokensHandler()
|
||||
|
||||
|
||||
_NATIVE_HEADERS: Final = frozenset(
|
||||
(
|
||||
"host",
|
||||
"accept",
|
||||
"accept-encoding",
|
||||
"connection",
|
||||
"user-agent",
|
||||
"content-length",
|
||||
"content-type",
|
||||
"x-api-key",
|
||||
"anthropic-version",
|
||||
)
|
||||
)
|
||||
|
||||
_DEPLOYMENT_OPTIONS: Final = frozenset(
|
||||
{
|
||||
"model",
|
||||
"api_key",
|
||||
"api_base",
|
||||
"custom_llm_provider",
|
||||
"rpm",
|
||||
"tpm",
|
||||
"timeout",
|
||||
"stream_timeout",
|
||||
"max_retries",
|
||||
"num_retries",
|
||||
"max_parallel_requests",
|
||||
"input_cost_per_token",
|
||||
"output_cost_per_token",
|
||||
"cache_read_input_token_cost",
|
||||
"cache_creation_input_token_cost",
|
||||
"cache_creation_input_token_cost_above_1hr",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
class _StrictModel(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid", frozen=True, strict=True)
|
||||
|
||||
|
||||
class _CacheControl(_StrictModel):
|
||||
type: Literal["ephemeral"]
|
||||
ttl: Literal["5m", "1h"] = "5m"
|
||||
|
||||
|
||||
class _Text(_StrictModel):
|
||||
type: Literal["text"]
|
||||
text: str = Field(min_length=1, pattern=r"\S")
|
||||
cache_control: _CacheControl | None = None
|
||||
|
||||
|
||||
class _ToolUse(_StrictModel):
|
||||
type: Literal["tool_use"]
|
||||
id: str = Field(min_length=1)
|
||||
name: str = Field(min_length=1)
|
||||
input: Mapping[str, JsonValue]
|
||||
cache_control: _CacheControl | None = None
|
||||
|
||||
|
||||
class _ResultText(_StrictModel):
|
||||
type: Literal["text"]
|
||||
text: str
|
||||
|
||||
|
||||
class _ToolResult(_StrictModel):
|
||||
type: Literal["tool_result"]
|
||||
tool_use_id: str = Field(min_length=1)
|
||||
content: str | Annotated[tuple[_ResultText, ...], Field(strict=False)]
|
||||
is_error: bool | None = None
|
||||
cache_control: _CacheControl | None = None
|
||||
|
||||
|
||||
_Block: TypeAlias = Annotated[_Text | _ToolUse | _ToolResult, Field(discriminator="type")]
|
||||
|
||||
|
||||
class _Message(_StrictModel):
|
||||
role: Literal["user", "assistant"]
|
||||
content: str | Annotated[tuple[_Block, ...], Field(strict=False)]
|
||||
|
||||
def blocks(self) -> tuple[_Text | _ToolUse | _ToolResult, ...]:
|
||||
return (_Text(type="text", text=self.content),) if isinstance(self.content, str) else tuple(self.content)
|
||||
|
||||
|
||||
class _Tool(_StrictModel):
|
||||
name: str = Field(min_length=1)
|
||||
description: str | None = None
|
||||
input_schema: Mapping[str, JsonValue]
|
||||
type: Literal["custom"] | None = None
|
||||
|
||||
|
||||
class _Request(_StrictModel):
|
||||
messages: tuple[_Message, ...] = Field(min_length=1, strict=False)
|
||||
system: str | Annotated[tuple[_ResultText, ...], Field(strict=False)] | None = None
|
||||
tools: Annotated[tuple[_Tool, ...], Field(strict=False)] | None = None
|
||||
model: str | None = None
|
||||
max_tokens: int | None = None
|
||||
stream: bool | None = None
|
||||
temperature: float | int | None = None
|
||||
top_p: float | int | None = None
|
||||
top_k: int | None = None
|
||||
stop_sequences: Annotated[tuple[str, ...], Field(strict=False)] | None = None
|
||||
metadata: Mapping[str, JsonValue] | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PromptPrefix:
|
||||
prefix_body: Mapping[str, JsonValue]
|
||||
fingerprint: str
|
||||
fingerprints: tuple[str, ...]
|
||||
ttl_seconds: int
|
||||
|
||||
|
||||
def _digest(value: object) -> str:
|
||||
return hashlib.sha256(
|
||||
json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False).encode()
|
||||
).hexdigest()
|
||||
|
||||
|
||||
def _next_digest(previous: str, boundary: tuple[int, str, Mapping[str, JsonValue]]) -> str:
|
||||
return _digest((previous, boundary))
|
||||
|
||||
|
||||
def parse_prompt(body: Mapping[str, JsonValue]) -> PromptPrefix | None:
|
||||
try:
|
||||
request: Final = _Request.model_validate(body)
|
||||
blocks: Final = tuple(message.blocks() for message in request.messages)
|
||||
except ValidationError:
|
||||
return None
|
||||
markers: Final = tuple(
|
||||
(message_index, block_index, block.cache_control)
|
||||
for message_index, message_blocks in enumerate(blocks)
|
||||
for block_index, block in enumerate(message_blocks)
|
||||
if block.cache_control is not None
|
||||
)
|
||||
if len(markers) != 1:
|
||||
return None
|
||||
message_end, block_end, marker = markers[0]
|
||||
normalized: Final = _JSON_OBJECT.validate_python(request.model_dump(mode="json", exclude_none=True))
|
||||
context: Final = MappingProxyType({key: normalized[key] for key in ("system", "tools") if key in normalized})
|
||||
boundaries: Final = tuple(
|
||||
(
|
||||
message_index,
|
||||
request.messages[message_index].role,
|
||||
_JSON_OBJECT.validate_python(
|
||||
block.model_dump(mode="json", exclude=MappingProxyType({"cache_control": True}), exclude_none=True)
|
||||
),
|
||||
)
|
||||
for message_index, message_blocks in enumerate(blocks[: message_end + 1])
|
||||
for block_index, block in enumerate(message_blocks)
|
||||
if message_index < message_end or block_index <= block_end
|
||||
)
|
||||
hashes: Final = tuple(
|
||||
accumulate(boundaries, _next_digest, initial=_digest((_JSON_OBJECT.validate_python(context), marker.ttl)))
|
||||
)[1:]
|
||||
prefix_messages: Final = tuple(
|
||||
_Message(
|
||||
role=request.messages[message_index].role,
|
||||
content=tuple(
|
||||
block
|
||||
for block_index, block in enumerate(message_blocks)
|
||||
if message_index < message_end or block_index <= block_end
|
||||
),
|
||||
)
|
||||
for message_index, message_blocks in enumerate(blocks[: message_end + 1])
|
||||
)
|
||||
return PromptPrefix(
|
||||
prefix_body=MappingProxyType(
|
||||
_JSON_OBJECT.validate_python(
|
||||
_Request(messages=prefix_messages, system=request.system, tools=request.tools).model_dump(
|
||||
mode="json", exclude_none=True
|
||||
)
|
||||
)
|
||||
),
|
||||
fingerprint=hashes[-1],
|
||||
fingerprints=tuple(reversed(hashes[-20:])),
|
||||
ttl_seconds=3600 if marker.ttl == "1h" else 300,
|
||||
)
|
||||
|
||||
|
||||
def cache_scope(
|
||||
caller_key_hash: str,
|
||||
deployment_id: str,
|
||||
provider_key: str,
|
||||
model: str,
|
||||
anthropic_version: str = DEFAULT_ANTHROPIC_API_VERSION,
|
||||
) -> str:
|
||||
return _digest((caller_key_hash, deployment_id, provider_key, model, anthropic_version))
|
||||
|
||||
|
||||
class _TTLUsage(BaseModel):
|
||||
model_config = ConfigDict(strict=True)
|
||||
ephemeral_5m_input_tokens: int = Field(default=0, ge=0)
|
||||
ephemeral_1h_input_tokens: int = Field(default=0, ge=0)
|
||||
|
||||
|
||||
class _CacheUsage(BaseModel):
|
||||
model_config = ConfigDict(strict=True)
|
||||
cached_tokens: int = Field(default=0, ge=0)
|
||||
cache_creation_tokens: int = Field(default=0, ge=0)
|
||||
cache_creation_token_details: _TTLUsage | None = None
|
||||
|
||||
|
||||
class _Usage(BaseModel):
|
||||
model_config = ConfigDict(strict=True)
|
||||
prompt_tokens: int = Field(ge=0)
|
||||
prompt_tokens_details: _CacheUsage
|
||||
|
||||
|
||||
class _Choice(BaseModel):
|
||||
finish_reason: str = Field(min_length=1)
|
||||
|
||||
|
||||
class _Response(BaseModel):
|
||||
model_config = ConfigDict(strict=True)
|
||||
model: str
|
||||
usage: _Usage
|
||||
choices: tuple[_Choice, ...] = Field(min_length=1, strict=False)
|
||||
|
||||
|
||||
class _CountBody(BaseModel):
|
||||
messages: Sequence[Mapping[str, JsonValue]]
|
||||
tools: Sequence[Mapping[str, JsonValue]] | None = None
|
||||
system: str | Sequence[Mapping[str, JsonValue]] | None = None
|
||||
|
||||
|
||||
class _CountResult(BaseModel):
|
||||
input_tokens: Annotated[StrictInt, Field(ge=0)]
|
||||
|
||||
|
||||
class TokenCounter(Protocol):
|
||||
async def __call__(self, model: str, api_key: str, body: Mapping[str, JsonValue]) -> int | None: ...
|
||||
|
||||
|
||||
def _count_objects(
|
||||
values: Sequence[Mapping[str, JsonValue]],
|
||||
) -> list[dict[str, JsonValue]]: # mutable-ok: the existing provider count API requires JSON lists/dicts
|
||||
return [dict(value) for value in values] # mutable-ok: serialize read-only inputs at the provider API boundary
|
||||
|
||||
|
||||
async def count_prompt_tokens(model: str, api_key: str, body: Mapping[str, JsonValue]) -> int | None:
|
||||
native: Final = _CountBody.model_validate(body)
|
||||
try:
|
||||
result: Final = _CountResult.model_validate(
|
||||
await _counter.handle_count_tokens_request(
|
||||
model=model,
|
||||
messages=_count_objects(native.messages),
|
||||
tools=_count_objects(native.tools) if native.tools is not None else None,
|
||||
system=native.system,
|
||||
api_key=api_key,
|
||||
timeout=15.0,
|
||||
)
|
||||
)
|
||||
except Exception: # noqa: BLE001 # provider/count validation failures are unavailable estimates, not zero tokens
|
||||
return None
|
||||
return result.input_tokens
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class NativePredictionTarget:
|
||||
model: str
|
||||
api_key: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class UnsupportedPredictionTarget:
|
||||
reason: Literal[
|
||||
"unsupported_deployment_configuration",
|
||||
"unsupported_provider_endpoint",
|
||||
"unsupported_provider",
|
||||
"unsupported_provider_credentials",
|
||||
]
|
||||
|
||||
|
||||
def resolve_prediction_target(params: LiteLLM_Params) -> NativePredictionTarget | UnsupportedPredictionTarget:
|
||||
configured_options: Final = frozenset(params.model_dump(exclude_defaults=True, exclude_none=True))
|
||||
if configured_options - _DEPLOYMENT_OPTIONS:
|
||||
return UnsupportedPredictionTarget("unsupported_deployment_configuration")
|
||||
api_base: Final = AnthropicModelInfo.get_api_base(params.api_base)
|
||||
if api_base not in ("https://api.anthropic.com", "https://api.anthropic.com/v1/messages"):
|
||||
return UnsupportedPredictionTarget("unsupported_provider_endpoint")
|
||||
try:
|
||||
model, provider, _, _ = litellm.get_llm_provider(
|
||||
model=params.model, custom_llm_provider=params.custom_llm_provider
|
||||
)
|
||||
except Exception: # noqa: BLE001 # the shared provider resolver raises for unknown deployments
|
||||
return UnsupportedPredictionTarget("unsupported_provider")
|
||||
if provider != "anthropic":
|
||||
return UnsupportedPredictionTarget("unsupported_provider")
|
||||
api_key: Final = AnthropicModelInfo.get_api_key(params.api_key)
|
||||
if api_key is None or not _supported_provider_key(api_key):
|
||||
return UnsupportedPredictionTarget("unsupported_provider_credentials")
|
||||
return NativePredictionTarget(model=model, api_key=api_key)
|
||||
|
||||
|
||||
def _supported_provider_key(api_key: str) -> bool:
|
||||
return bool(api_key) and not is_anthropic_oauth_key(api_key)
|
||||
|
||||
|
||||
def supported_prediction_headers(headers: Mapping[str, str]) -> bool:
|
||||
return all(
|
||||
name.lower() != "anthropic-beta"
|
||||
and (name.lower() != "anthropic-version" or value == DEFAULT_ANTHROPIC_API_VERSION)
|
||||
for name, value in headers.items()
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ObservedCachePrefix:
|
||||
prefix: PromptPrefix
|
||||
scope: str
|
||||
cached_tokens: int
|
||||
cache_creation_tokens: int
|
||||
|
||||
|
||||
def parse_observed_cache(
|
||||
wire: httpx.Request, response_obj: ModelResponse, caller_key_hash: str, deployment_id: str
|
||||
) -> ObservedCachePrefix | None:
|
||||
try:
|
||||
response: Final = _Response.model_validate(response_obj, from_attributes=True)
|
||||
body: Final = _JSON_OBJECT.validate_json(wire.content)
|
||||
headers: Final = _HEADERS.validate_python(wire.headers)
|
||||
except (ValidationError, RuntimeError, httpx.RequestNotRead):
|
||||
return None
|
||||
if (
|
||||
wire.url.scheme != "https"
|
||||
or wire.url.host != "api.anthropic.com"
|
||||
or wire.url.path != "/v1/messages"
|
||||
or wire.url.query
|
||||
or wire.url.port not in (None, 443)
|
||||
):
|
||||
return None
|
||||
if (
|
||||
frozenset(headers) - _NATIVE_HEADERS
|
||||
or not supported_prediction_headers(headers)
|
||||
or headers.get("anthropic-version") != DEFAULT_ANTHROPIC_API_VERSION
|
||||
):
|
||||
return None
|
||||
provider_key: Final = headers.get("x-api-key", "")
|
||||
model: Final = body.get("model")
|
||||
if not _supported_provider_key(provider_key) or not isinstance(model, str) or model != response.model:
|
||||
return None
|
||||
prefix: Final = parse_prompt(body)
|
||||
if prefix is None:
|
||||
return None
|
||||
usage: Final = response.usage.prompt_tokens_details
|
||||
cache_tokens: Final = usage.cached_tokens + usage.cache_creation_tokens
|
||||
if cache_tokens <= 0 or cache_tokens > response.usage.prompt_tokens:
|
||||
return None
|
||||
split: Final = usage.cache_creation_token_details
|
||||
if usage.cache_creation_tokens and split is None:
|
||||
return None
|
||||
if split is not None and (
|
||||
split.ephemeral_5m_input_tokens + split.ephemeral_1h_input_tokens != usage.cache_creation_tokens
|
||||
or (prefix.ttl_seconds == 300 and split.ephemeral_1h_input_tokens > 0)
|
||||
or (prefix.ttl_seconds == 3600 and split.ephemeral_5m_input_tokens > 0)
|
||||
):
|
||||
return None
|
||||
return ObservedCachePrefix(
|
||||
prefix=prefix,
|
||||
scope=cache_scope(caller_key_hash, deployment_id, provider_key, model),
|
||||
cached_tokens=cache_tokens,
|
||||
cache_creation_tokens=usage.cache_creation_tokens,
|
||||
)
|
||||
|
|
@ -144,10 +144,13 @@ class AzureFoundryModelInfo(BaseLLMModelInfo):
|
|||
def get_api_key(api_key: str | None = None) -> str | None:
|
||||
return api_key or litellm.api_key or get_secret_str("AZURE_AI_API_KEY")
|
||||
|
||||
@staticmethod
|
||||
def get_api_version(api_version: str | None = None) -> str | None:
|
||||
return api_version or litellm.api_version or get_secret_str("AZURE_API_VERSION")
|
||||
|
||||
@property
|
||||
def api_version(self, api_version: str | None = None) -> str | None:
|
||||
api_version = api_version or litellm.api_version or get_secret_str("AZURE_API_VERSION")
|
||||
return api_version
|
||||
def api_version(self) -> str | None:
|
||||
return AzureFoundryModelInfo.get_api_version()
|
||||
|
||||
def get_token_counter(self) -> BaseTokenCounter | None:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -1,8 +1,8 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import Callable, Iterator, Sequence
|
||||
from typing import Final, TypeVar
|
||||
from collections.abc import Callable, Iterator, Mapping, Sequence
|
||||
from typing import Final, TypeVar, cast # noqa: TID251 # a rebuilt chat row has no typed constructor across roles
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
|
@ -364,3 +364,67 @@ def merge_guardrailed_scoped_messages(
|
|||
yield from appended
|
||||
|
||||
return list(_merged())
|
||||
|
||||
|
||||
def _content_part_text(part: object) -> str | None:
|
||||
if not isinstance(part, Mapping):
|
||||
return None
|
||||
text: Final = part.get("text")
|
||||
return text if isinstance(text, str) else None
|
||||
|
||||
|
||||
def message_slot_texts(message: Mapping[str, object]) -> tuple[str, ...]:
|
||||
content: Final = message.get("content")
|
||||
if isinstance(content, str):
|
||||
return (content,)
|
||||
if isinstance(content, list):
|
||||
return tuple(text for part in content if (text := _content_part_text(part)) is not None)
|
||||
return ()
|
||||
|
||||
|
||||
def message_text_slot_count(message: AllMessageValues) -> int:
|
||||
return len(message_slot_texts(message))
|
||||
|
||||
|
||||
def _part_with_text(part: object, text: str) -> object:
|
||||
if not isinstance(part, Mapping):
|
||||
return part
|
||||
return {**part, "text": text} # mutable-ok: content parts stay JSON-plain dicts
|
||||
|
||||
|
||||
def _content_with_slot_texts(content: Sequence[object], texts: Sequence[str]) -> Sequence[object]:
|
||||
remaining_texts: Final = iter(texts)
|
||||
return [ # mutable-ok: message content stays a JSON list
|
||||
_part_with_text(part, next(remaining_texts)) if _content_part_text(part) is not None else part
|
||||
for part in content
|
||||
]
|
||||
|
||||
|
||||
def message_with_slot_texts(message: AllMessageValues, texts: Sequence[str]) -> AllMessageValues | None:
|
||||
"""Swap one rewritten text into each text slot of a chat row, in order.
|
||||
|
||||
A slot is a string ``content`` or one list part carrying a string ``text``;
|
||||
images and other parts ride along untouched. Returns None unless the counts
|
||||
line up exactly, so a rewrite never lands on the wrong slot.
|
||||
"""
|
||||
if message_text_slot_count(message) != len(texts):
|
||||
return None
|
||||
content: Final = message.get("content")
|
||||
if not isinstance(content, (str, list)):
|
||||
return message
|
||||
rewritten_content: Final = texts[0] if isinstance(content, str) else _content_with_slot_texts(content, texts)
|
||||
rewritten: Final = {**message, "content": rewritten_content} # mutable-ok: chat rows stay JSON-plain dicts
|
||||
return cast("AllMessageValues", rewritten) # cast-ok: the same row with only its text slots swapped
|
||||
|
||||
|
||||
class UnappliableRequestRewrite(Exception):
|
||||
def __init__(self, guardrail_name: str) -> None:
|
||||
super().__init__(
|
||||
f"Guardrail '{guardrail_name}' rewrote the request in a way this endpoint cannot apply, "
|
||||
"so the request was rejected rather than sent unrewritten"
|
||||
)
|
||||
self.guardrail_name: Final = guardrail_name
|
||||
|
||||
|
||||
def unappliable_request_rewrite(guardrail_name: str | None) -> UnappliableRequestRewrite:
|
||||
return UnappliableRequestRewrite(guardrail_name or "unknown")
|
||||
|
|
|
|||
|
|
@ -7,13 +7,21 @@ This uses aws_sdk_bedrock_runtime for bidirectional streaming with Nova Sonic.
|
|||
import asyncio
|
||||
import contextlib
|
||||
import json
|
||||
from collections.abc import AsyncIterator, Mapping
|
||||
from typing import Final, Protocol
|
||||
from collections.abc import AsyncIterator, Mapping, MutableMapping
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
from typing import Final, NoReturn, Protocol
|
||||
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm._logging import _redact_string, verbose_proxy_logger
|
||||
from litellm.constants import (
|
||||
BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY,
|
||||
BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY,
|
||||
BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY,
|
||||
REALTIME_SESSION_SUCCESS_LOGGED_KEY,
|
||||
)
|
||||
from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
|
||||
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
||||
|
|
@ -28,6 +36,32 @@ from .transformation import BedrockRealtimeConfig
|
|||
_CLIENT_MODALITIES_ADAPTER: Final[TypeAdapter["list[str] | None"]] = TypeAdapter(list[str] | None)
|
||||
_CLIENT_MESSAGE_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
|
||||
|
||||
_EMPTY_JSON_OBJECT: Final[Mapping[str, JsonValue]] = MappingProxyType({})
|
||||
|
||||
_BEDROCK_STREAM_ERROR_STATUS: Final[Mapping[str, int]] = MappingProxyType(
|
||||
{
|
||||
"AccessDeniedException": 403,
|
||||
"ConflictException": 400,
|
||||
"InternalServerException": 500,
|
||||
"ModelErrorException": 424,
|
||||
"ModelNotReadyException": 429,
|
||||
"ModelStreamErrorException": 424,
|
||||
"ModelTimeoutException": 408,
|
||||
"ResourceNotFoundException": 404,
|
||||
"ServiceQuotaExceededException": 400,
|
||||
"ServiceUnavailableException": 503,
|
||||
"ThrottlingException": 429,
|
||||
"ValidationException": 400,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _as_bedrock_error(error: BaseException) -> BaseException:
|
||||
status_code: Final = _BEDROCK_STREAM_ERROR_STATUS.get(type(error).__name__)
|
||||
if status_code is None:
|
||||
return error
|
||||
return BedrockError(status_code=status_code, message=f"{type(error).__name__}: {error}")
|
||||
|
||||
|
||||
def _json_dict(value: JsonValue) -> dict[str, JsonValue]:
|
||||
return value if isinstance(value, dict) else {}
|
||||
|
|
@ -51,6 +85,8 @@ def _should_log_event(openai_message: Mapping[str, object]) -> bool:
|
|||
class RealtimeClientWebSocket(Protocol):
|
||||
"""The client-facing websocket surface the realtime bridge talks to."""
|
||||
|
||||
scope: MutableMapping[str, object] # mutable-ok: the ASGI scope is the per-connection state store
|
||||
|
||||
async def receive_text(self) -> str: ...
|
||||
|
||||
async def send_text(self, data: str) -> None: ...
|
||||
|
|
@ -85,6 +121,81 @@ class BedrockBidirectionalStream(Protocol):
|
|||
async def await_output(self) -> tuple[object, BedrockOutputStream]: ...
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _BridgeOutcome:
|
||||
logged_events: tuple[OpenAIRealtimeEvents, ...]
|
||||
provider_failure: BaseException | None
|
||||
client_disconnected: bool
|
||||
|
||||
|
||||
async def _client_messages(client_ws: RealtimeClientWebSocket, initial_message: str | None) -> AsyncIterator[str]:
|
||||
if initial_message is not None:
|
||||
yield initial_message
|
||||
while True:
|
||||
try:
|
||||
yield await client_ws.receive_text()
|
||||
except Exception as e: # noqa: BLE001 # any receive failure means the client is gone
|
||||
verbose_proxy_logger.debug("Client to Bedrock forwarding ended: %s", e, exc_info=True)
|
||||
return
|
||||
|
||||
|
||||
def _pending_session_update(scope: Mapping[str, object]) -> str | None:
|
||||
"""A fallback attempt on the same websocket replays the session.update the failed attempt never acked."""
|
||||
if scope.get(BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY) is True:
|
||||
committed_failure: Final = scope.get(BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY)
|
||||
raise BedrockError(
|
||||
status_code=400,
|
||||
message=(
|
||||
"Bedrock realtime session already committed to a provider stream; it cannot be replayed"
|
||||
+ (f". The committed stream failed with: {committed_failure}" if committed_failure else "")
|
||||
),
|
||||
)
|
||||
pending: Final = scope.get(BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY)
|
||||
return pending if isinstance(pending, str) else None
|
||||
|
||||
|
||||
def _raise_provider_failure(scope: MutableMapping[str, object], failure: BaseException) -> NoReturn:
|
||||
error: Final = _as_bedrock_error(failure)
|
||||
verbose_proxy_logger.error("Bedrock Realtime: provider stream failed: %s", _redact_string(str(error)))
|
||||
if scope.get(BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY) is True:
|
||||
scope[BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY] = _redact_string(str(error))
|
||||
raise error from failure
|
||||
|
||||
|
||||
def _parse_client_message(message: str) -> Mapping[str, JsonValue]:
|
||||
try:
|
||||
return _json_dict(_CLIENT_MESSAGE_ADAPTER.validate_json(message))
|
||||
except ValueError:
|
||||
return _EMPTY_JSON_OBJECT
|
||||
|
||||
|
||||
async def _ack_session_update(
|
||||
client_ws: RealtimeClientWebSocket,
|
||||
bedrock_stream: BedrockBidirectionalStream,
|
||||
transformation_config: BedrockRealtimeConfig,
|
||||
model: str,
|
||||
logging_obj: LiteLLMLogging | None,
|
||||
parsed_client_message: Mapping[str, JsonValue],
|
||||
) -> bool:
|
||||
"""Ack the client's session.update once Bedrock accepted the stream; False means the client is gone."""
|
||||
await bedrock_stream.await_output()
|
||||
client_ws.scope.pop(BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY, None)
|
||||
client_ws.scope[BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY] = True # rebind-ok: scope outlives the attempt
|
||||
if logging_obj is None:
|
||||
return True
|
||||
requested_modalities: Final = _CLIENT_MODALITIES_ADAPTER.validate_python(
|
||||
_json_dict(parsed_client_message.get("session")).get("modalities")
|
||||
)
|
||||
try:
|
||||
await client_ws.send_text(
|
||||
json.dumps(transformation_config.session_updated_event(model, logging_obj, requested_modalities))
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # any send failure means the client is gone
|
||||
verbose_proxy_logger.debug("Client to Bedrock forwarding ended: %s", e, exc_info=True)
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
class BedrockRealtime(BaseAWSLLM):
|
||||
"""Handler for Bedrock Nova Sonic realtime speech-to-speech API."""
|
||||
|
||||
|
|
@ -132,6 +243,8 @@ class BedrockRealtime(BaseAWSLLM):
|
|||
except ImportError:
|
||||
raise ImportError("Missing aws_sdk_bedrock_runtime. Install with: pip install aws-sdk-bedrock-runtime")
|
||||
|
||||
pending_session_update: Final = _pending_session_update(websocket.scope)
|
||||
|
||||
# Get AWS region
|
||||
if aws_region_name is None:
|
||||
optional_params: Final = {
|
||||
|
|
@ -190,90 +303,106 @@ class BedrockRealtime(BaseAWSLLM):
|
|||
|
||||
transformation_config: Final = BedrockRealtimeConfig()
|
||||
|
||||
try:
|
||||
# Initialize the bidirectional stream
|
||||
bedrock_stream: Final = await open_bidirectional_stream()
|
||||
bedrock_stream: Final = await open_bidirectional_stream()
|
||||
|
||||
verbose_proxy_logger.debug("Bedrock Realtime: Bidirectional stream established")
|
||||
verbose_proxy_logger.debug("Bedrock Realtime: Bidirectional stream established")
|
||||
|
||||
if pending_session_update is None:
|
||||
await websocket.send_text(json.dumps(transformation_config.session_created_event(model, logging_obj)))
|
||||
verbose_proxy_logger.debug("Bedrock Realtime: sent session.created to client on connect")
|
||||
|
||||
# Track state for transformation
|
||||
session_state: Final[RealtimeResponseTransformInput] = {
|
||||
"current_output_item_id": None,
|
||||
"current_response_id": None,
|
||||
"current_conversation_id": None,
|
||||
"current_delta_chunks": None,
|
||||
"current_item_chunks": None,
|
||||
"current_delta_type": None,
|
||||
"session_configuration_request": None,
|
||||
}
|
||||
# Track state for transformation
|
||||
session_state: Final[RealtimeResponseTransformInput] = {
|
||||
"current_output_item_id": None,
|
||||
"current_response_id": None,
|
||||
"current_conversation_id": None,
|
||||
"current_delta_chunks": None,
|
||||
"current_item_chunks": None,
|
||||
"current_delta_type": None,
|
||||
"session_configuration_request": None,
|
||||
}
|
||||
|
||||
# Create tasks for bidirectional forwarding
|
||||
client_to_bedrock_task: Final = asyncio.create_task(
|
||||
self._forward_client_to_bedrock(
|
||||
websocket,
|
||||
bedrock_stream,
|
||||
transformation_config,
|
||||
model,
|
||||
session_state,
|
||||
logging_obj,
|
||||
outcome: Final = await self._bridge(
|
||||
websocket,
|
||||
bedrock_stream,
|
||||
transformation_config,
|
||||
model,
|
||||
session_state,
|
||||
logging_obj,
|
||||
initial_message=pending_session_update,
|
||||
)
|
||||
|
||||
logged_events: Final = (
|
||||
*outcome.logged_events,
|
||||
*(
|
||||
leftover_event
|
||||
for leftover_event in transformation_config.leftover_usage_done_events()
|
||||
if _should_log_event(leftover_event)
|
||||
),
|
||||
)
|
||||
if logged_events:
|
||||
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(
|
||||
logging_obj.dispatch_success_handlers(
|
||||
list(logged_events), # mutable-ok: realtime spend logging requires a list result
|
||||
prefer_async_handlers=True,
|
||||
)
|
||||
)
|
||||
logging_obj.model_call_details[REALTIME_SESSION_SUCCESS_LOGGED_KEY] = True
|
||||
|
||||
async def forward_bedrock_and_collect_logged_events() -> tuple[OpenAIRealtimeEvents, ...]:
|
||||
return tuple(
|
||||
[
|
||||
event
|
||||
async for event in self._forward_bedrock_to_client(
|
||||
bedrock_stream,
|
||||
websocket,
|
||||
transformation_config,
|
||||
model,
|
||||
logging_obj,
|
||||
session_state,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
bedrock_to_client_task: Final = asyncio.create_task(forward_bedrock_and_collect_logged_events())
|
||||
|
||||
# Wait for both tasks to complete
|
||||
await asyncio.gather(
|
||||
client_to_bedrock_task,
|
||||
bedrock_to_client_task,
|
||||
return_exceptions=True,
|
||||
if outcome.provider_failure is None:
|
||||
return
|
||||
if outcome.client_disconnected:
|
||||
verbose_proxy_logger.debug(
|
||||
"Bedrock Realtime: stream failed after the client disconnected: %s", outcome.provider_failure
|
||||
)
|
||||
return
|
||||
_raise_provider_failure(websocket.scope, outcome.provider_failure)
|
||||
|
||||
forwarded_logged_events: Final = (
|
||||
bedrock_to_client_task.result()
|
||||
if not bedrock_to_client_task.cancelled() and bedrock_to_client_task.exception() is None
|
||||
else ()
|
||||
)
|
||||
logged_events: Final = (
|
||||
*forwarded_logged_events,
|
||||
*(
|
||||
leftover_event
|
||||
for leftover_event in transformation_config.leftover_usage_done_events()
|
||||
if _should_log_event(leftover_event)
|
||||
),
|
||||
)
|
||||
if logged_events:
|
||||
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(
|
||||
logging_obj.dispatch_success_handlers(
|
||||
list(logged_events), # mutable-ok: realtime spend logging requires a list result
|
||||
prefer_async_handlers=True,
|
||||
)
|
||||
)
|
||||
async def _bridge(
|
||||
self,
|
||||
websocket: RealtimeClientWebSocket,
|
||||
bedrock_stream: BedrockBidirectionalStream,
|
||||
transformation_config: BedrockRealtimeConfig,
|
||||
model: str,
|
||||
session_state: RealtimeResponseTransformInput,
|
||||
logging_obj: LiteLLMLogging,
|
||||
initial_message: str | None,
|
||||
) -> _BridgeOutcome:
|
||||
"""Run both forwarding directions until the client leaves or either side fails."""
|
||||
logged: Final[list[OpenAIRealtimeEvents]] = [] # mutable-ok: events forwarded before a failure are still spend
|
||||
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception("Error in BedrockRealtime.async_realtime: %s", e)
|
||||
try:
|
||||
await websocket.close(code=1011, reason=_redact_string(f"Internal error: {e}"))
|
||||
except Exception:
|
||||
pass
|
||||
raise
|
||||
async def collect_logged_events() -> None:
|
||||
async for event in self._forward_bedrock_to_client(
|
||||
bedrock_stream, websocket, transformation_config, model, logging_obj, session_state
|
||||
):
|
||||
logged.append(event)
|
||||
|
||||
client_task: Final = asyncio.create_task(
|
||||
self._forward_client_to_bedrock(
|
||||
websocket, bedrock_stream, transformation_config, model, session_state, logging_obj, initial_message
|
||||
)
|
||||
)
|
||||
bedrock_task: Final = asyncio.create_task(collect_logged_events())
|
||||
|
||||
await asyncio.wait((client_task, bedrock_task), return_when=asyncio.FIRST_COMPLETED)
|
||||
client_disconnected: Final = (
|
||||
client_task.done() and not client_task.cancelled() and client_task.exception() is None
|
||||
)
|
||||
client_task.cancel()
|
||||
bedrock_task.cancel()
|
||||
client_outcome, bedrock_outcome = await asyncio.gather(client_task, bedrock_task, return_exceptions=True)
|
||||
|
||||
return _BridgeOutcome(
|
||||
logged_events=tuple(logged),
|
||||
provider_failure=(
|
||||
client_outcome
|
||||
if isinstance(client_outcome, Exception)
|
||||
else bedrock_outcome
|
||||
if isinstance(bedrock_outcome, Exception)
|
||||
else None
|
||||
),
|
||||
client_disconnected=client_disconnected,
|
||||
)
|
||||
|
||||
async def _forward_client_to_bedrock(
|
||||
self,
|
||||
|
|
@ -283,8 +412,12 @@ class BedrockRealtime(BaseAWSLLM):
|
|||
model: str,
|
||||
session_state: RealtimeResponseTransformInput,
|
||||
logging_obj: LiteLLMLogging | None = None,
|
||||
):
|
||||
"""Forward messages from client WebSocket to Bedrock stream."""
|
||||
initial_message: str | None = None,
|
||||
) -> None:
|
||||
"""Forward messages from client WebSocket to Bedrock stream.
|
||||
|
||||
Returns once the client is gone; provider failures (input stream or readiness) propagate to the caller.
|
||||
"""
|
||||
from aws_sdk_bedrock_runtime.models import (
|
||||
BidirectionalInputPayloadPart,
|
||||
InvokeModelWithBidirectionalStreamInputChunk,
|
||||
|
|
@ -299,41 +432,28 @@ class BedrockRealtime(BaseAWSLLM):
|
|||
verbose_proxy_logger.debug("Bedrock Realtime: Sent to Bedrock: %s", bedrock_message[:200])
|
||||
|
||||
try:
|
||||
while True:
|
||||
# Receive message from client
|
||||
message = await client_ws.receive_text()
|
||||
async for message in _client_messages(client_ws, initial_message):
|
||||
verbose_proxy_logger.debug("Bedrock Realtime: Received from client: %s", message[:200])
|
||||
parsed_client_message = _parse_client_message(message)
|
||||
is_session_update = _json_str(parsed_client_message.get("type")) == "session.update"
|
||||
if is_session_update:
|
||||
client_ws.scope[BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY] = (
|
||||
message # rebind-ok: scope outlives the attempt
|
||||
)
|
||||
|
||||
# Transform OpenAI format to Bedrock format
|
||||
transformed_messages = transformation_config.transform_realtime_request(
|
||||
message=message,
|
||||
model=model,
|
||||
session_configuration_request=session_state.get("session_configuration_request"),
|
||||
)
|
||||
|
||||
# Send transformed messages to Bedrock
|
||||
for bedrock_message in transformed_messages:
|
||||
await send_to_bedrock(bedrock_message)
|
||||
|
||||
if logging_obj is not None:
|
||||
client_message_type: str | None = None
|
||||
requested_modalities: list[str] | None = None
|
||||
with contextlib.suppress(Exception):
|
||||
parsed_client_message = _json_dict(_CLIENT_MESSAGE_ADAPTER.validate_json(message))
|
||||
client_message_type = _json_str(parsed_client_message.get("type"))
|
||||
if client_message_type == "session.update":
|
||||
requested_modalities = _CLIENT_MODALITIES_ADAPTER.validate_python(
|
||||
_json_dict(parsed_client_message.get("session")).get("modalities")
|
||||
)
|
||||
if client_message_type == "session.update":
|
||||
await client_ws.send_text(
|
||||
json.dumps(
|
||||
transformation_config.session_updated_event(model, logging_obj, requested_modalities)
|
||||
)
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.debug("Client to Bedrock forwarding ended: %s", e, exc_info=True)
|
||||
if is_session_update and not await _ack_session_update(
|
||||
client_ws, bedrock_stream, transformation_config, model, logging_obj, parsed_client_message
|
||||
):
|
||||
break
|
||||
finally:
|
||||
for close_message in transformation_config.session_close_messages():
|
||||
with contextlib.suppress(Exception):
|
||||
await send_to_bedrock(close_message)
|
||||
|
|
@ -349,68 +469,71 @@ class BedrockRealtime(BaseAWSLLM):
|
|||
logging_obj: LiteLLMLogging,
|
||||
session_state: RealtimeResponseTransformInput,
|
||||
) -> AsyncIterator[OpenAIRealtimeEvents]:
|
||||
"""Forward messages from Bedrock to the client, yielding the ones to record for spend logging."""
|
||||
try:
|
||||
while True:
|
||||
# Receive from Bedrock
|
||||
output = await bedrock_stream.await_output()
|
||||
result = await output[1].receive()
|
||||
"""Forward messages from Bedrock to the client, yielding the ones to record for spend logging.
|
||||
|
||||
if result is None:
|
||||
verbose_proxy_logger.debug("Bedrock Realtime: Bedrock stream ended")
|
||||
break
|
||||
Provider failures propagate to the caller; the client websocket is only closed on a normal stream end.
|
||||
"""
|
||||
|
||||
payload_bytes = result.value.bytes_ if result.value else None
|
||||
if payload_bytes:
|
||||
bedrock_response = payload_bytes.decode("utf-8")
|
||||
verbose_proxy_logger.debug("Bedrock Realtime: Received from Bedrock: %s", bedrock_response[:200])
|
||||
|
||||
# Transform Bedrock format to OpenAI format
|
||||
realtime_response_transform_input: RealtimeResponseTransformInput = {
|
||||
"current_output_item_id": session_state.get("current_output_item_id"),
|
||||
"current_response_id": session_state.get("current_response_id"),
|
||||
"current_conversation_id": session_state.get("current_conversation_id"),
|
||||
"current_delta_chunks": session_state.get("current_delta_chunks"),
|
||||
"current_item_chunks": session_state.get("current_item_chunks"),
|
||||
"current_delta_type": session_state.get("current_delta_type"),
|
||||
"session_configuration_request": session_state.get("session_configuration_request"),
|
||||
}
|
||||
|
||||
transformed_response = transformation_config.transform_realtime_response(
|
||||
message=bedrock_response,
|
||||
model=model,
|
||||
logging_obj=logging_obj,
|
||||
realtime_response_transform_input=realtime_response_transform_input,
|
||||
)
|
||||
|
||||
# Update session state
|
||||
session_state.update(
|
||||
{
|
||||
"current_output_item_id": transformed_response.get("current_output_item_id"),
|
||||
"current_response_id": transformed_response.get("current_response_id"),
|
||||
"current_conversation_id": transformed_response.get("current_conversation_id"),
|
||||
"current_delta_chunks": transformed_response.get("current_delta_chunks"),
|
||||
"current_item_chunks": transformed_response.get("current_item_chunks"),
|
||||
"current_delta_type": transformed_response.get("current_delta_type"),
|
||||
"session_configuration_request": transformed_response.get("session_configuration_request"),
|
||||
}
|
||||
)
|
||||
|
||||
# Send transformed messages to client
|
||||
response_value = transformed_response["response"]
|
||||
openai_messages = response_value if isinstance(response_value, list) else (response_value,)
|
||||
for openai_message in openai_messages:
|
||||
message_json = json.dumps(openai_message)
|
||||
await client_ws.send_text(message_json)
|
||||
verbose_proxy_logger.debug("Bedrock Realtime: Sent to client: %s", message_json[:200])
|
||||
if _should_log_event(openai_message):
|
||||
yield openai_message
|
||||
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.debug("Bedrock to client forwarding ended: %s", e, exc_info=True)
|
||||
finally:
|
||||
# Close the client WebSocket
|
||||
async def send_to_client(message_json: str) -> bool:
|
||||
try:
|
||||
await client_ws.close()
|
||||
except Exception:
|
||||
pass
|
||||
await client_ws.send_text(message_json)
|
||||
except Exception as e: # noqa: BLE001 # any send failure means the client is gone
|
||||
verbose_proxy_logger.debug("Bedrock to client forwarding ended: %s", e, exc_info=True)
|
||||
return False
|
||||
verbose_proxy_logger.debug("Bedrock Realtime: Sent to client: %s", message_json[:200])
|
||||
return True
|
||||
|
||||
output: Final = await bedrock_stream.await_output()
|
||||
while True:
|
||||
result = await output[1].receive()
|
||||
|
||||
if result is None:
|
||||
verbose_proxy_logger.debug("Bedrock Realtime: Bedrock stream ended")
|
||||
with contextlib.suppress(Exception):
|
||||
await client_ws.close()
|
||||
return
|
||||
|
||||
payload_bytes = result.value.bytes_ if result.value else None
|
||||
if payload_bytes:
|
||||
bedrock_response = payload_bytes.decode("utf-8")
|
||||
verbose_proxy_logger.debug("Bedrock Realtime: Received from Bedrock: %s", bedrock_response[:200])
|
||||
|
||||
# Transform Bedrock format to OpenAI format
|
||||
realtime_response_transform_input: RealtimeResponseTransformInput = {
|
||||
"current_output_item_id": session_state.get("current_output_item_id"),
|
||||
"current_response_id": session_state.get("current_response_id"),
|
||||
"current_conversation_id": session_state.get("current_conversation_id"),
|
||||
"current_delta_chunks": session_state.get("current_delta_chunks"),
|
||||
"current_item_chunks": session_state.get("current_item_chunks"),
|
||||
"current_delta_type": session_state.get("current_delta_type"),
|
||||
"session_configuration_request": session_state.get("session_configuration_request"),
|
||||
}
|
||||
|
||||
transformed_response = transformation_config.transform_realtime_response(
|
||||
message=bedrock_response,
|
||||
model=model,
|
||||
logging_obj=logging_obj,
|
||||
realtime_response_transform_input=realtime_response_transform_input,
|
||||
)
|
||||
|
||||
# Update session state
|
||||
session_state.update(
|
||||
{
|
||||
"current_output_item_id": transformed_response.get("current_output_item_id"),
|
||||
"current_response_id": transformed_response.get("current_response_id"),
|
||||
"current_conversation_id": transformed_response.get("current_conversation_id"),
|
||||
"current_delta_chunks": transformed_response.get("current_delta_chunks"),
|
||||
"current_item_chunks": transformed_response.get("current_item_chunks"),
|
||||
"current_delta_type": transformed_response.get("current_delta_type"),
|
||||
"session_configuration_request": transformed_response.get("session_configuration_request"),
|
||||
}
|
||||
)
|
||||
|
||||
# Send transformed messages to client
|
||||
response_value = transformed_response["response"]
|
||||
openai_messages = response_value if isinstance(response_value, list) else (response_value,)
|
||||
for openai_message in openai_messages:
|
||||
if not await send_to_client(json.dumps(openai_message)):
|
||||
return
|
||||
if _should_log_event(openai_message):
|
||||
yield openai_message
|
||||
|
|
|
|||
|
|
@ -42,6 +42,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import (
|
|||
stream_item_field,
|
||||
stream_item_fingerprint,
|
||||
stream_item_items,
|
||||
unappliable_request_rewrite,
|
||||
)
|
||||
from litellm.main import stream_chunk_builder
|
||||
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
|
||||
|
|
@ -196,6 +197,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
else:
|
||||
# Step 3: Map guardrail responses back to original message structure
|
||||
if guardrailed_texts and texts_to_check:
|
||||
if len(guardrailed_texts) != len(text_task_mappings):
|
||||
raise unappliable_request_rewrite(guardrail_to_apply.guardrail_name)
|
||||
await self._apply_guardrail_responses_to_input_texts(
|
||||
messages=messages,
|
||||
responses=guardrailed_texts,
|
||||
|
|
|
|||
|
|
@ -56,6 +56,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import (
|
|||
stream_item_field,
|
||||
stream_item_fingerprint,
|
||||
stream_item_items,
|
||||
unappliable_request_rewrite,
|
||||
)
|
||||
from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge_guardrailed_tools
|
||||
from litellm.responses.litellm_completion_transformation.transformation import (
|
||||
|
|
@ -495,13 +496,13 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
data["instructions"] = written_back.instructions # rebind-ok: data is an out-param
|
||||
elif isinstance(input_data, str):
|
||||
guardrailed_texts: Final = guardrailed_inputs.get("texts") or ()
|
||||
if len(guardrailed_texts) > 1:
|
||||
raise unappliable_request_rewrite(guardrail_to_apply.guardrail_name)
|
||||
data["input"] = guardrailed_texts[0] if guardrailed_texts else input_data # rebind-ok: data is an out-param
|
||||
else:
|
||||
rewritten_texts: Final = guardrailed_inputs.get("texts") or ()
|
||||
if len(rewritten_texts) != len(extracted.task_mappings):
|
||||
from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite
|
||||
|
||||
raise UnappliableRequestRewrite(guardrail_to_apply.guardrail_name or "unknown")
|
||||
raise unappliable_request_rewrite(guardrail_to_apply.guardrail_name)
|
||||
await self._apply_guardrail_responses_to_input(
|
||||
messages=input_data,
|
||||
responses=rewritten_texts,
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
from collections.abc import Mapping
|
||||
from typing import Any, Final
|
||||
from urllib.parse import unquote
|
||||
|
||||
|
|
@ -8,7 +9,36 @@ from litellm.llms.vertex_ai.common_utils import (
|
|||
)
|
||||
from litellm.types.llms.openai import BatchJobStatus, CreateBatchRequest
|
||||
from litellm.types.llms.vertex_ai import *
|
||||
from litellm.types.utils import LiteLLMBatch
|
||||
from litellm.types.utils import LiteLLMBatch, PromptTokensDetailsWrapper
|
||||
|
||||
|
||||
def vertex_prompt_tokens_details(
|
||||
usage_metadata: Mapping[str, object],
|
||||
) -> PromptTokensDetailsWrapper | None:
|
||||
raw_details: Final = usage_metadata.get("promptTokensDetails")
|
||||
if not isinstance(raw_details, list):
|
||||
return None
|
||||
|
||||
def _normalize(detail: object) -> tuple[str, int] | None:
|
||||
if not isinstance(detail, Mapping):
|
||||
return None
|
||||
modality: Final = detail.get("modality")
|
||||
token_count: Final = detail.get("tokenCount")
|
||||
if not isinstance(modality, str) or not isinstance(token_count, int):
|
||||
return None
|
||||
return modality.upper(), token_count
|
||||
|
||||
parsed_details: Final = tuple(_normalize(detail) for detail in raw_details)
|
||||
normalized: Final = tuple(detail for detail in parsed_details if detail is not None)
|
||||
if len(normalized) != len(parsed_details):
|
||||
return None
|
||||
|
||||
return PromptTokensDetailsWrapper(
|
||||
text_tokens=sum(token_count for modality, token_count in normalized if modality in ("TEXT", "DOCUMENT")),
|
||||
audio_tokens=sum(token_count for modality, token_count in normalized if modality == "AUDIO"),
|
||||
image_tokens=sum(token_count for modality, token_count in normalized if modality == "IMAGE"),
|
||||
video_tokens=sum(token_count for modality, token_count in normalized if modality == "VIDEO"),
|
||||
)
|
||||
|
||||
|
||||
class VertexAIBatchTransformation:
|
||||
|
|
|
|||
|
|
@ -298,8 +298,6 @@ def transform_openai_input_gemini_embed_content(
|
|||
|
||||
|
||||
_IMAGE_MIME_TYPES: Final = frozenset({"image/png", "image/jpeg"})
|
||||
_VIDEO_TOKENS_PER_SECOND: Final = 258.0
|
||||
_AUDIO_TOKENS_PER_SECOND: Final = 32.0
|
||||
_usage_metadata_adapter: Final = TypeAdapter(UsageMetadata)
|
||||
|
||||
|
||||
|
|
@ -339,11 +337,12 @@ def _is_image_element(
|
|||
return False
|
||||
|
||||
|
||||
def _count_input_images(
|
||||
def _is_image_only_input(
|
||||
input: GeminiEmbeddingInput,
|
||||
resolved_files: Mapping[str, Mapping[str, str]],
|
||||
) -> int:
|
||||
return sum(1 for element in _flatten_input(input) if _is_image_element(element, resolved_files))
|
||||
) -> bool:
|
||||
elements: Final = _flatten_input(input)
|
||||
return bool(elements) and all(_is_image_element(element, resolved_files) for element in elements)
|
||||
|
||||
|
||||
def _tokens_for_modality(details: Sequence[PromptTokensDetails], modality: str) -> int:
|
||||
|
|
@ -372,30 +371,29 @@ def _usage_from_embed_content_response(
|
|||
total_tokens: Final = usage_metadata.get("totalTokenCount") or prompt_tokens
|
||||
|
||||
details: Final[Sequence[PromptTokensDetails]] = usage_metadata.get("promptTokensDetails") or ()
|
||||
if not details:
|
||||
return Usage(
|
||||
prompt_tokens=prompt_tokens,
|
||||
total_tokens=total_tokens,
|
||||
prompt_tokens_details=PromptTokensDetailsWrapper(
|
||||
text_tokens=0,
|
||||
image_tokens=prompt_tokens if _is_image_only_input(input, resolved_files) else 0,
|
||||
),
|
||||
)
|
||||
|
||||
text_tokens: Final = _tokens_for_modality(details, "TEXT")
|
||||
audio_tokens: Final = _tokens_for_modality(details, "AUDIO")
|
||||
image_tokens: Final = _tokens_for_modality(details, "IMAGE")
|
||||
video_tokens: Final = _tokens_for_modality(details, "VIDEO")
|
||||
image_count: Final = _count_input_images(input, resolved_files)
|
||||
|
||||
video_length_seconds: Final = video_tokens / _VIDEO_TOKENS_PER_SECOND if video_tokens > 0 else 0.0
|
||||
audio_length_seconds: Final = audio_tokens / _AUDIO_TOKENS_PER_SECOND if audio_tokens > 0 else 0.0
|
||||
|
||||
# generic_cost_per_token rewrites text_tokens to the full prompt minus
|
||||
# other modalities when both text_tokens and image_count are zero. For
|
||||
# video, that misallocates video tokens to text; a 1-token floor sidesteps
|
||||
# the rewrite and keeps billing on input_cost_per_video_per_second.
|
||||
needs_video_text_floor: Final = video_length_seconds > 0 and text_tokens == 0 and image_count == 0
|
||||
resolved_text_tokens: Final = 1 if needs_video_text_floor else text_tokens
|
||||
|
||||
return Usage(
|
||||
prompt_tokens=prompt_tokens,
|
||||
total_tokens=total_tokens,
|
||||
prompt_tokens_details=PromptTokensDetailsWrapper(
|
||||
text_tokens=resolved_text_tokens,
|
||||
text_tokens=text_tokens,
|
||||
audio_tokens=audio_tokens,
|
||||
image_count=image_count,
|
||||
video_length_seconds=video_length_seconds,
|
||||
audio_length_seconds=audio_length_seconds,
|
||||
image_tokens=image_tokens,
|
||||
video_tokens=video_tokens,
|
||||
),
|
||||
)
|
||||
|
||||
|
|
@ -415,8 +413,7 @@ def process_embed_content_response(
|
|||
model_response: EmbeddingResponse to populate
|
||||
model: Model name
|
||||
response_json: Raw JSON response from embedContent endpoint
|
||||
resolved_files: Mapping of file references (files/abc) to {mime_type, uri},
|
||||
used to bill resolved image references at the per-image rate
|
||||
resolved_files: Mapping of file references to resolved metadata
|
||||
|
||||
Returns:
|
||||
EmbeddingResponse with single embedding
|
||||
|
|
|
|||
|
|
@ -2592,7 +2592,9 @@ def _complete_custom_openai(
|
|||
copilot_headers.update(extra_headers)
|
||||
extra_headers = copilot_headers
|
||||
|
||||
if extra_headers is not None:
|
||||
use_base_llm_http_handler: Final = get_secret_bool("EXPERIMENTAL_OPENAI_BASE_LLM_HTTP_HANDLER")
|
||||
|
||||
if extra_headers is not None and not use_base_llm_http_handler:
|
||||
optional_params["extra_headers"] = extra_headers
|
||||
|
||||
if litellm.enable_preview_features and metadata is not None: # [PREVIEW] allow metadata to be passed to OPENAI
|
||||
|
|
@ -2609,8 +2611,6 @@ def _complete_custom_openai(
|
|||
optional_params[k] = v
|
||||
|
||||
## COMPLETION CALL
|
||||
use_base_llm_http_handler: Final = get_secret_bool("EXPERIMENTAL_OPENAI_BASE_LLM_HTTP_HANDLER")
|
||||
|
||||
try:
|
||||
if use_base_llm_http_handler:
|
||||
response = base_llm_http_handler.completion(
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -11503,6 +11503,12 @@
|
|||
"description": "Enable content moderation to check for harmful content (harassment, hate speech, etc.).",
|
||||
"title": "Content Moderation Check"
|
||||
},
|
||||
"contextual_grounding_from_messages": {
|
||||
"default": false,
|
||||
"description": "ApplyGuardrail: when True, post-call scans of a request with no grounding_source / query content parts send the system and developer messages as the grounding source and the latest user message as the query, so the guardrail's contextual grounding policy can score the response. Bedrock bills contextual grounding units for these scans and rejects queries, sources and responses over its contextual grounding length limits, so leave this off for guardrails without a contextual grounding policy. Default False: plain messages are never sent as grounding context.",
|
||||
"title": "Contextual Grounding From Messages",
|
||||
"type": "boolean"
|
||||
},
|
||||
"credentials": {
|
||||
"anyOf": [
|
||||
{
|
||||
|
|
|
|||
|
|
@ -887,6 +887,7 @@ class LiteLLMRoutes(enum.Enum):
|
|||
"/auto_router/validate_complexity_router_config",
|
||||
# Per-session auto-router read - the endpoint scopes the row to the caller's own key hash
|
||||
"/auto_router/session",
|
||||
"/cost/predict-cache",
|
||||
# Agent registry - reads are role-scoped and writes are proxy-admin-gated
|
||||
# inside agent_endpoints/endpoints.py
|
||||
*agent_management_routes,
|
||||
|
|
@ -2635,9 +2636,13 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
|
|||
None,
|
||||
description="max file size in MB for /v1/files uploads, for any purpose, if a file is larger than this size it will be rejected before being forwarded to the provider",
|
||||
)
|
||||
allowed_file_extensions: tuple[str, ...] | None = Field(
|
||||
None,
|
||||
description="the only file extensions (e.g. ['.jsonl', '.pdf', '.txt']) accepted on /v1/files uploads, for any purpose, matched case-insensitively against the uploaded filename. Files with any other extension, or none, are rejected. An empty list rejects every upload. Unset means no allowlist is applied",
|
||||
)
|
||||
blocked_file_extensions: tuple[str, ...] | None = Field(
|
||||
None,
|
||||
description="file extensions (e.g. ['.exe', '.sh']) rejected on /v1/files uploads, for any purpose, matched case-insensitively against the uploaded filename",
|
||||
description="file extensions (e.g. ['.exe', '.sh']) rejected on /v1/files uploads, for any purpose, matched case-insensitively against the uploaded filename. Deprecated in favour of allowed_file_extensions; still enforced, after the allowlist, when set",
|
||||
)
|
||||
max_response_size_mb: int | None = Field(
|
||||
None,
|
||||
|
|
@ -4414,6 +4419,9 @@ class TeamInfoResponseObjectTeamTable(LiteLLM_TeamTable):
|
|||
access_group_mcp_server_ids: list[str] | None = None
|
||||
access_group_agent_ids: list[str] | None = None
|
||||
access_group_details: tuple[TeamAccessGroupModelGrant, ...] | None = None
|
||||
# Parent org's model ceiling, reported only to callers who can manage the team.
|
||||
# None = no org or not a manager; [] or ["all-proxy-models"] = no ceiling.
|
||||
organization_models: list[str] | None = None
|
||||
|
||||
|
||||
class TeamInfoResponseObject(TypedDict):
|
||||
|
|
|
|||
|
|
@ -123,6 +123,7 @@ from litellm.repositories.table_repositories import (
|
|||
from litellm.repositories.team_repository import TeamRepository
|
||||
from litellm.repositories.user_repository import UserRepository
|
||||
from litellm.router import Router
|
||||
from litellm.types.proxy.auth.auth_checks import UserNotFoundError
|
||||
from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget
|
||||
from litellm.utils import get_utc_datetime
|
||||
|
||||
|
|
@ -327,9 +328,23 @@ _safe_json_loads_obj: Final = _typed_json_loads(safe_json_loads)
|
|||
last_db_access_time: Final = LimitedSizeOrderedDict(max_size=100)
|
||||
db_cache_expiry: Final = DEFAULT_IN_MEMORY_TTL # refresh every 5s
|
||||
|
||||
_TEAM_MEMBERSHIP_INFLIGHT_MAX: Final = 10000
|
||||
_team_membership_inflight: Final = LimitedSizeOrderedDict(max_size=_TEAM_MEMBERSHIP_INFLIGHT_MAX)
|
||||
|
||||
|
||||
class _TeamMembershipCacheMiss:
|
||||
__slots__ = ()
|
||||
|
||||
|
||||
_TEAM_MEMBERSHIP_CACHE_MISS: Final = _TeamMembershipCacheMiss()
|
||||
|
||||
all_routes: Final = LiteLLMRoutes.openai_routes.value + LiteLLMRoutes.management_routes.value
|
||||
|
||||
|
||||
def _membership_from_shared_load(result: object) -> LiteLLM_TeamMembership | None:
|
||||
return result if isinstance(result, LiteLLM_TeamMembership) else None
|
||||
|
||||
|
||||
def _log_budget_lookup_failure(entity: str, error: Exception) -> None:
|
||||
"""
|
||||
Log a warning when budget lookup fails; cache will not be populated.
|
||||
|
|
@ -880,6 +895,7 @@ async def common_checks(
|
|||
request_query_params=_safe_get_request_query_params(request=request),
|
||||
llm_router=llm_router,
|
||||
request=request,
|
||||
team_id=valid_token.team_id if valid_token is not None else None,
|
||||
)
|
||||
|
||||
skip_all_budget_checks: Final = skip_budget_checks or (
|
||||
|
|
@ -887,6 +903,22 @@ async def common_checks(
|
|||
and (route in MODEL_DISCOVERY_ROUTES or not RouteChecks.is_llm_api_route(route=route))
|
||||
)
|
||||
|
||||
membership_user_id: Final = (
|
||||
valid_token.user_id if valid_token is not None and (bool(_model) or not skip_all_budget_checks) else None
|
||||
)
|
||||
team_membership_loaded: Final = team_object is not None and membership_user_id is not None
|
||||
loaded_team_membership: Final = (
|
||||
await get_team_membership(
|
||||
user_id=membership_user_id,
|
||||
team_id=team_object.team_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
if team_object is not None and membership_user_id is not None
|
||||
else None
|
||||
)
|
||||
|
||||
unpriced_models: Final = (
|
||||
_unpriced_models_in_request(model=_model, llm_router=llm_router)
|
||||
if litellm.block_requests_for_models_without_pricing and RouteChecks.is_llm_api_route(route=route)
|
||||
|
|
@ -936,6 +968,8 @@ async def common_checks(
|
|||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
team_membership=loaded_team_membership,
|
||||
team_membership_loaded=team_membership_loaded,
|
||||
)
|
||||
|
||||
# Require trace id for agent keys when agent has require_trace_id_on_calls_by_agent
|
||||
|
|
@ -987,6 +1021,8 @@ async def common_checks(
|
|||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
team_membership=loaded_team_membership,
|
||||
team_membership_loaded=team_membership_loaded,
|
||||
)
|
||||
|
||||
# Run before apply_key_tags_pre_auth injects key metadata.tags into request_body.
|
||||
|
|
@ -1096,6 +1132,8 @@ async def common_checks(
|
|||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
team_membership=loaded_team_membership,
|
||||
team_membership_loaded=team_membership_loaded,
|
||||
),
|
||||
_check_end_user_budget(end_user_obj=end_user_object, route=route)
|
||||
if end_user_object is not None and end_user_object.litellm_budget_table is not None
|
||||
|
|
@ -2141,7 +2179,76 @@ async def get_tag_object(
|
|||
return tag_objects.get(tag_name)
|
||||
|
||||
|
||||
def _membership_from_cached_payload(
|
||||
cached: object,
|
||||
) -> LiteLLM_TeamMembership | None | _TeamMembershipCacheMiss:
|
||||
if cached is None:
|
||||
return _TEAM_MEMBERSHIP_CACHE_MISS
|
||||
if cached == NO_TEAM_MEMBERSHIP_SENTINEL:
|
||||
return None
|
||||
cached_membership: Final = CacheCodec.deserialize(cached, model_type=LiteLLM_TeamMembership)
|
||||
return cached_membership if cached_membership is not None else _TEAM_MEMBERSHIP_CACHE_MISS
|
||||
|
||||
|
||||
@log_db_metrics
|
||||
async def _fetch_team_membership_from_db(
|
||||
user_id: str,
|
||||
team_id: str,
|
||||
prisma_client: PrismaClient,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
parent_otel_span: Span | None = None,
|
||||
proxy_logging_obj: ProxyLogging | None = None,
|
||||
) -> LiteLLM_TeamMembership | None:
|
||||
_ = parent_otel_span, proxy_logging_obj
|
||||
response: Final = await _dictable_table(TeamMembershipRepository(prisma_client)).find_unique(
|
||||
where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}},
|
||||
include={"litellm_budget_table": True},
|
||||
)
|
||||
membership: Final = None if response is None else LiteLLM_TeamMembership.model_validate(response.dict())
|
||||
_key: Final = team_membership_reservation_cache_key(user_id=user_id, team_id=team_id)
|
||||
if membership is None:
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=_key,
|
||||
value=NO_TEAM_MEMBERSHIP_SENTINEL,
|
||||
ttl=get_management_object_ttl(user_api_key_cache),
|
||||
)
|
||||
else:
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=_key,
|
||||
value=membership,
|
||||
model_type=LiteLLM_TeamMembership,
|
||||
)
|
||||
return membership
|
||||
|
||||
|
||||
async def _load_team_membership_on_cache_miss(
|
||||
user_id: str,
|
||||
team_id: str,
|
||||
cache_key: str,
|
||||
prisma_client: PrismaClient,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
parent_otel_span: Span | None,
|
||||
proxy_logging_obj: ProxyLogging | None,
|
||||
) -> LiteLLM_TeamMembership | None:
|
||||
try:
|
||||
redis_cached: Final[object] = await user_api_key_cache.async_get_cache(key=cache_key)
|
||||
redis_membership: Final = _membership_from_cached_payload(redis_cached)
|
||||
if not isinstance(redis_membership, _TeamMembershipCacheMiss):
|
||||
return redis_membership
|
||||
|
||||
return await _fetch_team_membership_from_db(
|
||||
user_id=user_id,
|
||||
team_id=team_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
except Exception:
|
||||
verbose_proxy_logger.exception("Error getting team membership")
|
||||
return None
|
||||
|
||||
|
||||
async def get_team_membership(
|
||||
user_id: str,
|
||||
team_id: str,
|
||||
|
|
@ -2155,54 +2262,42 @@ async def get_team_membership(
|
|||
|
||||
Do a isolated check for team membership vs. doing a combined key + team + user + team-membership check, as key might come in frequently for different users/teams. Larger call will slowdown query time. This way we get to cache the constant (key/team/user info) and only update based on the changing value (team membership).
|
||||
"""
|
||||
from litellm.proxy._types import LiteLLM_TeamMembership
|
||||
|
||||
if prisma_client is None:
|
||||
raise Exception("No db connected")
|
||||
|
||||
if user_id is None or team_id is None:
|
||||
return None
|
||||
|
||||
_key: Final = team_membership_reservation_cache_key(user_id=user_id, team_id=team_id)
|
||||
|
||||
# check if in cache
|
||||
cached: Final[object] = await user_api_key_cache.async_get_cache(key=_key)
|
||||
if cached == NO_TEAM_MEMBERSHIP_SENTINEL:
|
||||
return None
|
||||
cached_membership_obj: Final = CacheCodec.deserialize(cached, model_type=LiteLLM_TeamMembership)
|
||||
if cached_membership_obj is not None:
|
||||
return cached_membership_obj
|
||||
l1_cached: Final[object] = await user_api_key_cache.async_get_cache(key=_key, local_only=True)
|
||||
l1_membership: Final = _membership_from_cached_payload(l1_cached)
|
||||
if not isinstance(l1_membership, _TeamMembershipCacheMiss):
|
||||
return l1_membership
|
||||
|
||||
# else, check db
|
||||
try:
|
||||
response: Final = await _dictable_table(TeamMembershipRepository(prisma_client)).find_unique(
|
||||
where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}},
|
||||
include={"litellm_budget_table": True},
|
||||
inflight: Final[object] = _team_membership_inflight.get(_key)
|
||||
if isinstance(inflight, asyncio.Task):
|
||||
return _membership_from_shared_load(await asyncio.shield(inflight))
|
||||
|
||||
if prisma_client is None:
|
||||
raise Exception("No db connected")
|
||||
|
||||
task: Final = asyncio.ensure_future(
|
||||
_load_team_membership_on_cache_miss(
|
||||
user_id=user_id,
|
||||
team_id=team_id,
|
||||
cache_key=_key,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
)
|
||||
_team_membership_inflight[_key] = task
|
||||
|
||||
if response is None:
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=_key,
|
||||
value=NO_TEAM_MEMBERSHIP_SENTINEL,
|
||||
ttl=get_management_object_ttl(user_api_key_cache),
|
||||
)
|
||||
return None
|
||||
def _clear_inflight(_done: object) -> None:
|
||||
if _team_membership_inflight.get(_key) is task:
|
||||
_team_membership_inflight.pop(_key, None)
|
||||
|
||||
_response: Final = LiteLLM_TeamMembership.model_validate(response.dict())
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=_key,
|
||||
value=_response,
|
||||
model_type=LiteLLM_TeamMembership,
|
||||
)
|
||||
|
||||
return _response
|
||||
except Exception:
|
||||
verbose_proxy_logger.exception(
|
||||
"Error getting team membership for user_id: %s, team_id: %s",
|
||||
user_id,
|
||||
team_id,
|
||||
)
|
||||
return None
|
||||
task.add_done_callback(_clear_inflight)
|
||||
return _membership_from_shared_load(await asyncio.shield(task))
|
||||
|
||||
|
||||
def model_in_access_group(model: str, team_models: list[str] | None, llm_router: Router | None) -> bool:
|
||||
|
|
@ -2375,13 +2470,6 @@ async def _backfill_null_user_email(
|
|||
return updated_row
|
||||
|
||||
|
||||
class UserNotFoundError(ValueError):
|
||||
"""The user row is provably absent, as opposed to merely unreadable, so a caller that reads a missing row as no user-level limits can key on it without also swallowing a database that would not answer."""
|
||||
|
||||
def __init__(self, user_id: str) -> None:
|
||||
super().__init__(f"User doesn't exist in db. 'user_id'={user_id}. Create user via `/user/new` call.")
|
||||
|
||||
|
||||
@log_db_metrics
|
||||
async def get_user_object(
|
||||
user_id: str | None,
|
||||
|
|
@ -2668,6 +2756,12 @@ async def invalidate_team_member_spend_state(
|
|||
publish_auth_cache_invalidation,
|
||||
)
|
||||
|
||||
inflight: Final[object] = _team_membership_inflight.pop(
|
||||
team_membership_reservation_cache_key(user_id=user_id, team_id=team_id), None
|
||||
)
|
||||
if isinstance(inflight, asyncio.Task) and inflight is not asyncio.current_task():
|
||||
await asyncio.wait((inflight,))
|
||||
|
||||
if new_spend is not None:
|
||||
from litellm.proxy.proxy_server import SPEND_DB_FLOOR_CACHE_TTL_SECONDS, spend_counter_cache
|
||||
|
||||
|
|
@ -4122,18 +4216,21 @@ async def _team_member_granted_models(
|
|||
prisma_client: PrismaClient,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
team_membership: LiteLLM_TeamMembership | None = None,
|
||||
team_membership_loaded: bool = False,
|
||||
) -> Sequence[str]:
|
||||
"""The member's own ``allowed_models`` scope; empty when the member is not narrowed below the team."""
|
||||
if team_object is None or valid_token.user_id is None:
|
||||
return ()
|
||||
|
||||
team_membership: Final = await get_team_membership(
|
||||
user_id=valid_token.user_id,
|
||||
team_id=team_object.team_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
if not team_membership_loaded:
|
||||
team_membership = await get_team_membership(
|
||||
user_id=valid_token.user_id,
|
||||
team_id=team_object.team_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
return () if team_membership is None else _member_allowed_models(team_membership)
|
||||
|
||||
|
||||
|
|
@ -4169,6 +4266,8 @@ async def _granted_model_lists(
|
|||
prisma_client: PrismaClient,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
team_membership: LiteLLM_TeamMembership | None = None,
|
||||
team_membership_loaded: bool = False,
|
||||
) -> tuple[Sequence[str], ...]:
|
||||
"""One model allowlist per level that participates in authorizing the request."""
|
||||
return (
|
||||
|
|
@ -4180,6 +4279,8 @@ async def _granted_model_lists(
|
|||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
team_membership=team_membership,
|
||||
team_membership_loaded=team_membership_loaded,
|
||||
),
|
||||
project_object.models if project_object is not None else (),
|
||||
await _org_granted_models(
|
||||
|
|
@ -4274,6 +4375,8 @@ async def collect_matched_model_access_groups(
|
|||
prisma_client: PrismaClient | None,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
team_membership: LiteLLM_TeamMembership | None = None,
|
||||
team_membership_loaded: bool = False,
|
||||
) -> tuple[str, ...]:
|
||||
"""
|
||||
The budgeted model access groups that authorized this request, sorted and deduplicated.
|
||||
|
|
@ -4319,6 +4422,8 @@ async def collect_matched_model_access_groups(
|
|||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
team_membership=team_membership,
|
||||
team_membership_loaded=team_membership_loaded,
|
||||
)
|
||||
for granted_model in granted_models
|
||||
)
|
||||
|
|
@ -4334,6 +4439,8 @@ async def stamp_matched_model_access_groups(
|
|||
prisma_client: PrismaClient | None,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
team_membership: LiteLLM_TeamMembership | None = None,
|
||||
team_membership_loaded: bool = False,
|
||||
) -> tuple[str, ...]:
|
||||
"""Record the groups that authorized this request on its auth object, for the post-call spend
|
||||
writer and the reservation counters, and hand them back for the budget check."""
|
||||
|
|
@ -4350,6 +4457,8 @@ async def stamp_matched_model_access_groups(
|
|||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
team_membership=team_membership,
|
||||
team_membership_loaded=team_membership_loaded,
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # fail-safe: attribution is spend telemetry, it must never break auth
|
||||
verbose_proxy_logger.debug("model access group attribution failed: %s", e)
|
||||
|
|
@ -4363,7 +4472,7 @@ async def stamp_matched_model_access_groups(
|
|||
|
||||
async def can_key_call_model(
|
||||
model: str | list[str],
|
||||
llm_model_list: list | None,
|
||||
llm_model_list: Sequence[object] | None,
|
||||
valid_token: UserAPIKeyAuth,
|
||||
llm_router: litellm.Router | None,
|
||||
) -> Literal[True]:
|
||||
|
|
@ -4410,7 +4519,7 @@ async def can_key_call_model(
|
|||
|
||||
async def can_key_call_resolved_model(
|
||||
model: str,
|
||||
llm_model_list: list | None,
|
||||
llm_model_list: Sequence[object] | None,
|
||||
valid_token: UserAPIKeyAuth,
|
||||
llm_router: litellm.Router | None,
|
||||
) -> None:
|
||||
|
|
@ -5152,6 +5261,8 @@ async def _check_team_member_budget(
|
|||
prisma_client: PrismaClient | None,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
team_membership: LiteLLM_TeamMembership | None = None,
|
||||
team_membership_loaded: bool = False,
|
||||
):
|
||||
"""Check if team member is over their max budget within the team."""
|
||||
if (
|
||||
|
|
@ -5160,23 +5271,25 @@ async def _check_team_member_budget(
|
|||
and valid_token is not None
|
||||
and valid_token.user_id is not None
|
||||
):
|
||||
team_membership: Final = await get_team_membership(
|
||||
user_id=valid_token.user_id,
|
||||
team_id=team_object.team_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
if not team_membership_loaded:
|
||||
team_membership = await get_team_membership(
|
||||
user_id=valid_token.user_id,
|
||||
team_id=team_object.team_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
loaded_membership = team_membership
|
||||
|
||||
# Per-member override wins; otherwise fall back to the team-level
|
||||
# default configured via team.metadata["team_member_budget_id"].
|
||||
team_member_budget: float | None = None
|
||||
if (
|
||||
team_membership is not None
|
||||
and team_membership.litellm_budget_table is not None
|
||||
and team_membership.litellm_budget_table.max_budget is not None
|
||||
loaded_membership is not None
|
||||
and loaded_membership.litellm_budget_table is not None
|
||||
and loaded_membership.litellm_budget_table.max_budget is not None
|
||||
):
|
||||
team_member_budget = team_membership.litellm_budget_table.max_budget
|
||||
team_member_budget = loaded_membership.litellm_budget_table.max_budget
|
||||
else:
|
||||
default_budget_id: Final = (team_object.metadata or {}).get("team_member_budget_id")
|
||||
if isinstance(default_budget_id, str):
|
||||
|
|
@ -5195,7 +5308,7 @@ async def _check_team_member_budget(
|
|||
team_member_budget = default_budget.max_budget
|
||||
|
||||
if team_member_budget is not None:
|
||||
team_member_spend = (team_membership.spend if team_membership is not None else 0.0) or 0.0
|
||||
team_member_spend = (loaded_membership.spend if loaded_membership is not None else 0.0) or 0.0
|
||||
|
||||
# Read from cross-pod counter (Redis-first) if available
|
||||
from litellm.proxy.proxy_server import get_current_spend
|
||||
|
|
@ -5224,6 +5337,8 @@ async def _check_team_member_model_access(
|
|||
prisma_client: Optional["PrismaClient"],
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
team_membership: LiteLLM_TeamMembership | None = None,
|
||||
team_membership_loaded: bool = False,
|
||||
) -> None:
|
||||
"""
|
||||
Check if a team member's per-member model scope allows access to the requested model.
|
||||
|
|
@ -5234,22 +5349,24 @@ async def _check_team_member_model_access(
|
|||
if valid_token.user_id is None or team_object.team_id is None:
|
||||
return
|
||||
|
||||
team_membership: Final = await get_team_membership(
|
||||
user_id=valid_token.user_id,
|
||||
team_id=team_object.team_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
if not team_membership_loaded:
|
||||
team_membership = await get_team_membership(
|
||||
user_id=valid_token.user_id,
|
||||
team_id=team_object.team_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
loaded_membership = team_membership
|
||||
|
||||
if (
|
||||
team_membership is None
|
||||
or team_membership.litellm_budget_table is None
|
||||
or not team_membership.litellm_budget_table.allowed_models
|
||||
loaded_membership is None
|
||||
or loaded_membership.litellm_budget_table is None
|
||||
or not loaded_membership.litellm_budget_table.allowed_models
|
||||
):
|
||||
return # no per-member restriction — inherit team-level check
|
||||
|
||||
member_allowed_models: Final[list[str]] = team_membership.litellm_budget_table.allowed_models
|
||||
member_allowed_models: Final[list[str]] = loaded_membership.litellm_budget_table.allowed_models
|
||||
try:
|
||||
_can_object_call_model(
|
||||
model=model,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -28,11 +28,11 @@ from litellm.proxy._types import (
|
|||
)
|
||||
from litellm.proxy.auth.auth_checks import (
|
||||
TeamNotFoundError,
|
||||
UserNotFoundError,
|
||||
get_team_membership,
|
||||
get_team_object,
|
||||
get_user_object,
|
||||
)
|
||||
from litellm.types.proxy.auth.auth_checks import UserNotFoundError
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy._types import Span
|
||||
|
|
|
|||
|
|
@ -136,6 +136,9 @@ class RouteChecks:
|
|||
# For llm_api_routes, also check registered pass-through endpoints
|
||||
################################################
|
||||
if allowed_route == "llm_api_routes":
|
||||
if route == "/auto_router/session" and RouteChecks._get_request_method(request) == "GET":
|
||||
return True
|
||||
|
||||
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
|
||||
InitPassThroughEndpointHelpers,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -585,7 +585,9 @@ LiteLLM ████████░░░░░░░░░░░░░░
|
|||
Claude Opus 5 ████████████████████████ $0.38
|
||||
```
|
||||
|
||||
The routed model comes from Claude Code's own transcript, so it only names the tier model when the auto-router deployment sets `return_raw_model_name: true` (the `lite autoroute` wizard does); otherwise it shows the alias you requested. The cost lines come from `GET /auto_router/session?session_id=...`, which any virtual key may call for its own sessions, and are cached for five seconds under a per-user `$TMPDIR/litellm-statusline-<uid>` directory. The baseline is the priciest model in the router's hardest tier, the same counterfactual the auto-router's savings reports use. `lite unconfigure claude` removes the `statusLine` entry only while it still points at that script.
|
||||
After the first response, the status line uses the latest routed model recorded by `GET /auto_router/session?session_id=...`, so it can show the tier model even when the transcript contains the router alias. If no session record is available, it falls back to Claude Code's transcript. Session records and costs are cached for five seconds under a per-user `$TMPDIR/litellm-statusline-<uid>` directory. The gateway records turns asynchronously, so the display can briefly lag a completed turn. Any virtual key may read its own sessions. The baseline is the priciest model in the router's hardest tier, the same counterfactual the auto-router's savings reports use. `lite unconfigure claude` removes the `statusLine` entry only while it still points at that script
|
||||
|
||||
After upgrading the CLI, rerun your original `lite configure claude` command with the same gateway, key and model choice to refresh `~/.litellm/statusline.py`. Keep any explicit `--model` value: omitting it removes the earlier model pin. Package upgrades alone do not refresh this installed copy
|
||||
|
||||
`lite codex` registers the same script as a Codex `Stop` hook for the launch, so after each turn Codex prints the same block as a system message. Codex asks once to trust the hook; the answer is remembered for later launches.
|
||||
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@ import os
|
|||
import tempfile
|
||||
from collections.abc import Callable, Mapping
|
||||
from dataclasses import dataclass
|
||||
from enum import StrEnum
|
||||
from enum import Enum
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import Annotated, Final
|
||||
|
|
@ -25,7 +25,7 @@ LITELLM_PROXY_API_KEY_ENV: Final = "LITELLM_PROXY_API_KEY"
|
|||
_REJECTED_STATUSES: Final = frozenset((401, 403))
|
||||
|
||||
|
||||
class ListingFailure(StrEnum):
|
||||
class ListingFailure(str, Enum):
|
||||
"""Why a proxy could not be listed, decided once where the HTTP outcome is classified.
|
||||
|
||||
`unreachable` means no response at all; the other kinds prove the proxy answered, so callers
|
||||
|
|
|
|||
|
|
@ -7,19 +7,17 @@ status refresh (about every 300ms while typing), so the proxy is asked at most o
|
|||
TTL per session and every other refresh is served from a small on-disk cache that holds
|
||||
only the proxy's answer, never the key.
|
||||
|
||||
Claude Code pipes a JSON payload on stdin (session_id, transcript_path, model); the routed
|
||||
model is the `message.model` of the latest foreground assistant line in the transcript,
|
||||
which is the proxy's response `model` field. That only names the tier model when the
|
||||
auto-router deployment sets `return_raw_model_name: true`; otherwise it is the alias the
|
||||
client requested. Codex pipes its Stop event instead (hook_event_name, session_id) and has
|
||||
no transcript to read, so the routed model comes from the proxy's session record and the
|
||||
result is printed as a `systemMessage` for the transcript. The proxy key is read from the
|
||||
agent's own environment (the static token `lite configure claude` writes); nothing here
|
||||
spawns a credential helper.
|
||||
Claude Code pipes a JSON payload on stdin (session_id, transcript_path, model). After the
|
||||
first foreground assistant response, the routed model comes from the proxy's session
|
||||
record, falling back to the latest foreground assistant `message.model` in the transcript
|
||||
when no record is available. Codex pipes its Stop event instead (hook_event_name, session_id)
|
||||
and prints the session record as a `systemMessage` for the transcript. The proxy key is read
|
||||
from the agent's own environment (the static token `lite configure claude` writes); nothing
|
||||
here spawns a credential helper.
|
||||
|
||||
Cost figures come from GET /auto_router/session on the proxy, which reads the per-session
|
||||
rollup written by the spend flush. That flush is asynchronous, so a turn's cost lands a
|
||||
second or two after the turn; the cache TTL absorbs it.
|
||||
The routed model and cost figures come from GET /auto_router/session on the proxy, which
|
||||
reads the per-session rollup written by the asynchronous spend flush. The record and cache
|
||||
can briefly lag a completed turn.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
|
@ -348,7 +346,8 @@ def status_line(
|
|||
if not session_id or not credentials.usable:
|
||||
return render(label, None, config_dir, color_enabled(env))
|
||||
session: Final = load_session(credentials, session_id, cache_dir, fetch)
|
||||
return render(label, session, config_dir, color_enabled(env))
|
||||
routed_label: Final = model_label(session.last_model, config_dir) if session is not None else label
|
||||
return render(routed_label, session, config_dir, color_enabled(env))
|
||||
|
||||
|
||||
def codex_stop_message(
|
||||
|
|
|
|||
|
|
@ -1571,6 +1571,9 @@ class ProxyBaseLLMRequestProcessing:
|
|||
) -> dict:
|
||||
exclude_values: Final = {"", None, "None"}
|
||||
hidden_params = hidden_params or {}
|
||||
resolved_call_id: Final = (
|
||||
call_id or hidden_params.get("litellm_call_id") or (request_data or {}).get("litellm_call_id")
|
||||
)
|
||||
timing_values: Final = _timing_values(
|
||||
hidden_params=hidden_params,
|
||||
logging_obj=litellm_logging_obj,
|
||||
|
|
@ -1598,7 +1601,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
classifier_cost: Final = _classifier_cost_from_request_data(request_data)
|
||||
|
||||
headers: Final = {
|
||||
"x-litellm-call-id": call_id,
|
||||
"x-litellm-call-id": resolved_call_id,
|
||||
"x-litellm-model-id": model_id,
|
||||
"x-litellm-model-name": model_name,
|
||||
"x-litellm-cache-key": cache_key,
|
||||
|
|
@ -3452,15 +3455,13 @@ class ProxyBaseLLMRequestProcessing:
|
|||
# a failed request reports no timing, matching /v1/chat/completions
|
||||
read_timing_from_logging_obj=False,
|
||||
)
|
||||
# Extract headers from exception - check both e.headers and e.response.headers
|
||||
headers = getattr(e, "headers", None) or {}
|
||||
if not headers:
|
||||
# Try to get headers from e.response.headers (httpx.Response)
|
||||
_response: Final = attribute_of(e, "response")
|
||||
if _response is not None:
|
||||
_response_headers: Final = getattr(_response, "headers", None)
|
||||
if _response_headers:
|
||||
headers = get_response_headers(dict(_response_headers))
|
||||
_response_headers: Final = getattr(_response, "headers", None) if _response is not None else None
|
||||
_provider_headers: Final = _response_headers or getattr(e, "litellm_response_headers", None)
|
||||
if _provider_headers:
|
||||
headers = get_response_headers(dict(_provider_headers))
|
||||
headers.update(custom_headers)
|
||||
|
||||
# Call response headers hook for failure
|
||||
|
|
|
|||
|
|
@ -0,0 +1,10 @@
|
|||
import os
|
||||
from collections.abc import Mapping
|
||||
|
||||
|
||||
def should_hide_default_credentials_hint(general_settings: Mapping[str, object]) -> bool:
|
||||
return (
|
||||
os.getenv("LITELLM_HIDE_DEFAULT_CREDENTIALS_HINT", "false").lower() == "true"
|
||||
or general_settings.get("hide_default_credentials_hint", False) is True
|
||||
or bool(os.getenv("UI_PASSWORD"))
|
||||
)
|
||||
91
litellm/proxy/common_utils/prompt_cache_pricing.py
Normal file
91
litellm/proxy/common_utils/prompt_cache_pricing.py
Normal file
|
|
@ -0,0 +1,91 @@
|
|||
from collections.abc import Mapping
|
||||
from math import isfinite
|
||||
from typing import Final
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm.cost_calculator import (
|
||||
_select_model_name_for_cost_calc, # pyright: ignore[reportPrivateUsage] # shares completion_cost's deployment tariff selection
|
||||
completion_cost, # pyright: ignore[reportUnknownVariableType] # legacy optional parameters are untyped
|
||||
)
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.types.management_endpoints.prompt_cache_prediction import CacheTokenBuckets
|
||||
from litellm.types.utils import CacheCreationTokenDetails, ModelResponse, PromptTokensDetailsWrapper, Usage
|
||||
|
||||
_PRICE_ENTRY: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def _valid_price(value: object) -> bool:
|
||||
return isinstance(value, (int, float)) and not isinstance(value, bool) and isfinite(value) and value >= 0
|
||||
|
||||
|
||||
def _has_required_prices(prices: Mapping[str, object], tokens: CacheTokenBuckets) -> bool:
|
||||
required: Final = (
|
||||
("input_cost_per_token", True),
|
||||
("cache_read_input_token_cost", tokens.cache_read_input_tokens > 0),
|
||||
("cache_creation_input_token_cost", tokens.cache_creation_5m_input_tokens > 0),
|
||||
("cache_creation_input_token_cost_above_1hr", tokens.cache_creation_1h_input_tokens > 0),
|
||||
)
|
||||
if any(needed and not _valid_price(prices.get(key)) for key, needed in required):
|
||||
return False
|
||||
return all(
|
||||
_valid_price(value)
|
||||
for key, value in prices.items()
|
||||
if value is not None and any(needed and key.startswith(f"{base}_above_") for base, needed in required)
|
||||
)
|
||||
|
||||
|
||||
def price_cache_tokens(model: str, deployment_id: str, tokens: CacheTokenBuckets) -> float | None:
|
||||
try:
|
||||
selected_model: Final = _select_model_name_for_cost_calc(
|
||||
model=model,
|
||||
completion_response=None,
|
||||
custom_pricing=True,
|
||||
custom_llm_provider="anthropic",
|
||||
router_model_id=deployment_id,
|
||||
)
|
||||
if selected_model is None:
|
||||
return None
|
||||
model_info: Final = litellm.get_model_info(model=selected_model, custom_llm_provider="anthropic")
|
||||
registry: Final = _PRICE_ENTRY.validate_python(litellm.model_cost) # pyright: ignore[reportUnknownMemberType] # legacy registry is validated at this boundary
|
||||
price_entry: Final = registry.get(model_info["key"])
|
||||
if price_entry is None:
|
||||
return None
|
||||
prices: Final = _PRICE_ENTRY.validate_python(price_entry)
|
||||
if not _has_required_prices(prices, tokens):
|
||||
return None
|
||||
usage: Final = Usage(
|
||||
prompt_tokens=tokens.total_tokens,
|
||||
completion_tokens=0,
|
||||
total_tokens=tokens.total_tokens,
|
||||
prompt_tokens_details=PromptTokensDetailsWrapper(
|
||||
cached_tokens=tokens.cache_read_input_tokens,
|
||||
cache_creation_tokens=tokens.cache_creation_5m_input_tokens + tokens.cache_creation_1h_input_tokens,
|
||||
cache_creation_token_details=CacheCreationTokenDetails(
|
||||
ephemeral_5m_input_tokens=tokens.cache_creation_5m_input_tokens,
|
||||
ephemeral_1h_input_tokens=tokens.cache_creation_1h_input_tokens,
|
||||
),
|
||||
),
|
||||
)
|
||||
logging_obj: Final = Logging(
|
||||
model=model,
|
||||
messages=[], # mutable-ok: Logging requires a list
|
||||
stream=False,
|
||||
call_type="completion",
|
||||
start_time=None,
|
||||
litellm_call_id="prompt-cache-prediction",
|
||||
function_id="prompt-cache-prediction",
|
||||
)
|
||||
completion_cost(
|
||||
completion_response=ModelResponse(model=model, usage=usage),
|
||||
model=model,
|
||||
custom_llm_provider="anthropic",
|
||||
custom_pricing=True,
|
||||
router_model_id=deployment_id,
|
||||
litellm_logging_obj=logging_obj,
|
||||
)
|
||||
cost: Final = logging_obj.cost_breakdown.get("input_cost") if logging_obj.cost_breakdown is not None else None
|
||||
return cost if cost is not None and _valid_price(cost) else None
|
||||
except Exception: # noqa: BLE001 # the shared pricing owners raise plain Exception for unpriceable models
|
||||
return None
|
||||
|
|
@ -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 ()
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ from typing import Final
|
|||
|
||||
from fastapi import APIRouter
|
||||
|
||||
from litellm.proxy.common_utils.html_forms.default_credentials_hint import should_hide_default_credentials_hint
|
||||
from litellm.types.proxy.discovery_endpoints.ui_discovery_endpoints import (
|
||||
UiDiscoveryEndpoints,
|
||||
)
|
||||
|
|
@ -23,10 +24,7 @@ async def get_ui_config():
|
|||
or general_settings.get("auto_redirect_ui_login_to_sso", False) is True
|
||||
)
|
||||
admin_ui_disabled: Final = os.getenv("DISABLE_ADMIN_UI", "false").lower() == "true"
|
||||
hide_default_credentials_hint: Final = bool(
|
||||
os.getenv("LITELLM_HIDE_DEFAULT_CREDENTIALS_HINT", "false").lower() == "true"
|
||||
or general_settings.get("hide_default_credentials_hint", False) is True
|
||||
)
|
||||
hide_default_credentials_hint: Final = should_hide_default_credentials_hint(general_settings)
|
||||
|
||||
sso_configured: Final = has_user_setup_sso()
|
||||
|
||||
|
|
|
|||
|
|
@ -244,6 +244,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
prompt_attack_threshold: float | None = 0.5,
|
||||
pii_confidence_threshold: float | None = 0.5,
|
||||
chunk_budget_chars: int = BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS,
|
||||
contextual_grounding_from_messages: bool = False,
|
||||
streaming_buffer_until_moderated: bool | None = None,
|
||||
streaming_sampling_rate: int | None = None,
|
||||
streaming_end_of_stream_only: bool | None = None,
|
||||
|
|
@ -265,6 +266,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
self.guardrailVersion = guardrailVersion
|
||||
self.guardrail_provider = "bedrock"
|
||||
self.chunk_budget_chars = chunk_budget_chars
|
||||
self.contextual_grounding_from_messages = contextual_grounding_from_messages
|
||||
self.experimental_use_latest_role_message_only = bool(kwargs.get("experimental_use_latest_role_message_only"))
|
||||
|
||||
# Resource-less, detect-only InvokeGuardrailChecks mode. Present `checks`
|
||||
|
|
@ -459,8 +461,8 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
"""
|
||||
Flatten a message into text blocks, preserving any contextual-grounding
|
||||
qualifier carried by the content-block ``type`` (grounding_source / query).
|
||||
Untagged text keeps ``qualifier=None`` so the payload is unchanged for
|
||||
callers that do not use grounding.
|
||||
Untagged text keeps ``qualifier=None``; the OUTPUT scan decides whether to
|
||||
derive grounding qualifiers from it.
|
||||
"""
|
||||
content: Final = message.get("content")
|
||||
if content is None:
|
||||
|
|
@ -493,6 +495,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
result carrying externally-influenced content can supply fake evidence for the
|
||||
contextual-grounding check to grade the response against. ``query`` is accepted
|
||||
from any role (it is the user's question).
|
||||
|
||||
With ``contextual_grounding_from_messages`` on, a request with no tagged blocks
|
||||
falls back to the plain messages: system / developer text is the grounding
|
||||
source and the latest user message is the query.
|
||||
"""
|
||||
grounding: Final[list[QualifiedTextBlock]] = []
|
||||
for message in messages or []:
|
||||
|
|
@ -504,7 +510,33 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
and role in _GROUNDING_SOURCE_TRUSTED_ROLES
|
||||
):
|
||||
grounding.append(block)
|
||||
return grounding
|
||||
if grounding or not self.contextual_grounding_from_messages:
|
||||
return grounding
|
||||
return self._derive_grounding_blocks_from_plain_messages(messages)
|
||||
|
||||
def _derive_grounding_blocks_from_plain_messages(
|
||||
self, messages: list[AllMessageValues] | None
|
||||
) -> list[QualifiedTextBlock]:
|
||||
if not messages:
|
||||
return []
|
||||
latest_user_index: Final = self._find_latest_message_index(messages, target_role="user")
|
||||
if latest_user_index is None:
|
||||
return []
|
||||
sources: Final = tuple(
|
||||
QualifiedTextBlock(text=block.text, qualifier="grounding_source")
|
||||
for message in messages
|
||||
if message.get("role") in _GROUNDING_SOURCE_TRUSTED_ROLES
|
||||
for block in self.get_content_items_for_message(message=message) or []
|
||||
if block.text
|
||||
)
|
||||
queries: Final = tuple(
|
||||
QualifiedTextBlock(text=block.text, qualifier="query")
|
||||
for block in self.get_content_items_for_message(message=messages[latest_user_index]) or []
|
||||
if block.text
|
||||
)
|
||||
if not sources or not queries:
|
||||
return []
|
||||
return [*sources, *queries]
|
||||
|
||||
def supports_scan_only_tool_results(self) -> bool:
|
||||
return self.experimental_use_latest_role_message_only is not True
|
||||
|
|
@ -3210,6 +3242,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
bedrock_response = await self.make_bedrock_api_request(
|
||||
source="OUTPUT",
|
||||
response=synthetic_response,
|
||||
messages=request_data.get("messages"),
|
||||
request_data=request_data,
|
||||
logging_event_type=_log_hook,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@
|
|||
|
||||
import fnmatch
|
||||
import os
|
||||
from collections.abc import Mapping
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Optional
|
||||
|
||||
import httpx
|
||||
|
|
@ -24,7 +24,7 @@ from litellm.llms.custom_httpx.http_handler import (
|
|||
httpxSpecialProvider,
|
||||
)
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.llms.openai import ChatCompletionToolParam
|
||||
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
|
||||
GenericGuardrailAPIMetadata,
|
||||
GenericGuardrailAPIRequest,
|
||||
|
|
@ -150,6 +150,26 @@ def _extract_inbound_headers(
|
|||
return None
|
||||
|
||||
|
||||
def _structured_rows_to_write_back(
|
||||
original_rows: Sequence[AllMessageValues] | None,
|
||||
shown_rows: Sequence[AllMessageValues] | None,
|
||||
returned_rows: Sequence[AllMessageValues],
|
||||
) -> tuple[AllMessageValues, ...] | None:
|
||||
"""The request model drops row keys its message types do not declare, so a
|
||||
row the server echoes back verbatim is restored to the original row object.
|
||||
A server that echoes every row back unchanged has not rewritten anything
|
||||
per row, so its answer is read from texts, as it was before rows could be
|
||||
returned at all."""
|
||||
if original_rows is None or shown_rows is None or len(returned_rows) != len(original_rows):
|
||||
return tuple(returned_rows)
|
||||
if all(returned == shown for shown, returned in zip(shown_rows, returned_rows)):
|
||||
return None
|
||||
return tuple(
|
||||
original if returned == shown else returned
|
||||
for original, shown, returned in zip(original_rows, shown_rows, returned_rows)
|
||||
)
|
||||
|
||||
|
||||
class GenericGuardrailAPI(CustomGuardrail):
|
||||
"""
|
||||
Generic Guardrail API integration for LiteLLM.
|
||||
|
|
@ -322,6 +342,8 @@ class GenericGuardrailAPI(CustomGuardrail):
|
|||
texts: list,
|
||||
images: list[str] | None,
|
||||
tools: list[ChatCompletionToolParam] | None,
|
||||
structured_messages: Sequence[AllMessageValues] | None,
|
||||
shown_messages: Sequence[AllMessageValues] | None,
|
||||
guardrail_response: GenericGuardrailAPIResponse,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
# Action is NONE or no modifications needed
|
||||
|
|
@ -336,6 +358,13 @@ class GenericGuardrailAPI(CustomGuardrail):
|
|||
return_inputs["tools"] = guardrail_response.tools
|
||||
elif tools:
|
||||
return_inputs["tools"] = tools
|
||||
rows_to_write_back: Final = (
|
||||
_structured_rows_to_write_back(structured_messages, shown_messages, guardrail_response.structured_messages)
|
||||
if guardrail_response.structured_messages
|
||||
else None
|
||||
)
|
||||
if rows_to_write_back is not None:
|
||||
return_inputs["structured_messages"] = list(rows_to_write_back) # mutable-ok: guardrail inputs take a list
|
||||
if guardrail_response.stream_holdback_chars is not None:
|
||||
return_inputs["stream_holdback_chars"] = guardrail_response.stream_holdback_chars
|
||||
return return_inputs
|
||||
|
|
@ -473,6 +502,8 @@ class GenericGuardrailAPI(CustomGuardrail):
|
|||
texts=texts,
|
||||
images=images,
|
||||
tools=tools,
|
||||
structured_messages=structured_messages,
|
||||
shown_messages=guardrail_request.structured_messages,
|
||||
guardrail_response=guardrail_response,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -44,7 +44,7 @@ class JavelinGuardrail(CustomGuardrail):
|
|||
application: str | None = None,
|
||||
**kwargs,
|
||||
):
|
||||
f"""
|
||||
"""
|
||||
Initialize the JavelinGuardrail class.
|
||||
|
||||
This calls: {api_base}/{api_version}/guardrail/{guardrail_name}/apply
|
||||
|
|
|
|||
|
|
@ -15,7 +15,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"
|
|||
# We check the raw guardrail dict because LitellmParams normalizes None → False,
|
||||
# making it impossible to distinguish "not set" from "explicitly false" via litellm_params.
|
||||
_raw_default_on: Final = cast(dict[str, Any], guardrail).get("litellm_params", {}).get("default_on")
|
||||
_default_on: Final = False if _raw_default_on is False else True
|
||||
_default_on: Final = _raw_default_on is not False
|
||||
|
||||
_callback: Final = MCPEndUserPermissionGuardrail(
|
||||
guardrail_name=guardrail.get("guardrail_name", ""),
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@ import asyncio
|
|||
import base64
|
||||
import os
|
||||
from collections.abc import Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, Literal, Optional
|
||||
|
||||
import httpx
|
||||
|
|
@ -14,11 +15,13 @@ from litellm.integrations.custom_guardrail import (
|
|||
CustomGuardrail,
|
||||
log_guardrail_information,
|
||||
)
|
||||
from litellm.llms.base_llm.guardrail_translation.utils import message_slot_texts, message_with_slot_texts
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
get_async_httpx_client,
|
||||
httpxSpecialProvider,
|
||||
)
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -27,12 +30,37 @@ if TYPE_CHECKING:
|
|||
|
||||
|
||||
_SANITIZE_FILE_FAIL_OPEN_TIMEOUT_SECONDS: Final = 30.0
|
||||
_SANITIZE_FILE_QUEUED_STATUSES: Final = frozenset({"created", "in progress"})
|
||||
_PROTECT_ROLES: Final = frozenset({"system", "user", "assistant"})
|
||||
|
||||
|
||||
class PromptSecurityGuardrailMissingSecrets(Exception):
|
||||
pass
|
||||
|
||||
|
||||
def _inputs_with_structured_messages(
|
||||
inputs: GenericGuardrailAPIInputs, rewritten_messages: Sequence[AllMessageValues] | None
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
if rewritten_messages is None:
|
||||
return inputs
|
||||
patched: Final[GenericGuardrailAPIInputs] = {
|
||||
**inputs,
|
||||
"structured_messages": list(rewritten_messages), # mutable-ok: the TypedDict field is declared as a list
|
||||
}
|
||||
return patched
|
||||
|
||||
|
||||
def _inputs_with_modifications(
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
modified_texts: list[str],
|
||||
rewritten_messages: Sequence[AllMessageValues] | None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
if not modified_texts:
|
||||
return _inputs_with_structured_messages(inputs, rewritten_messages)
|
||||
with_texts: Final[GenericGuardrailAPIInputs] = {**inputs, "texts": modified_texts}
|
||||
return _inputs_with_structured_messages(with_texts, rewritten_messages)
|
||||
|
||||
|
||||
class _ProtectVerdict(TypedDict, total=False):
|
||||
"""One side (``prompt`` or ``response``) of an ``/api/protect`` verdict."""
|
||||
|
||||
|
|
@ -275,14 +303,39 @@ class PromptSecurityGuardrail(CustomGuardrail):
|
|||
detail="Blocked by Prompt Security, Violations: " + ", ".join(violations),
|
||||
)
|
||||
elif action == "modify":
|
||||
# Extract modified texts from modified_messages
|
||||
modified_messages: Final = result.get("modified_messages", [])
|
||||
modified_texts: Final = self._extract_texts_from_messages(modified_messages)
|
||||
if modified_texts:
|
||||
inputs["texts"] = modified_texts
|
||||
return _inputs_with_modifications(
|
||||
inputs,
|
||||
self._extract_texts_from_messages(modified_messages),
|
||||
self._structured_messages_with_modifications(structured_messages, modified_messages),
|
||||
)
|
||||
|
||||
return inputs
|
||||
|
||||
def _is_sent_to_protect(self, message: Mapping[str, object]) -> bool:
|
||||
return self.check_tool_results or message.get("role") in _PROTECT_ROLES
|
||||
|
||||
def _structured_messages_with_modifications(
|
||||
self,
|
||||
structured_messages: Sequence[AllMessageValues],
|
||||
modified_messages: Sequence[Mapping[str, object]],
|
||||
) -> tuple[AllMessageValues, ...] | None:
|
||||
sent_indices: Final = tuple(
|
||||
index for index, message in enumerate(structured_messages) if self._is_sent_to_protect(message)
|
||||
)
|
||||
if not sent_indices or len(sent_indices) != len(modified_messages):
|
||||
return None
|
||||
rewritten: Final = tuple(
|
||||
message_with_slot_texts(structured_messages[index], self._extract_texts_from_messages((modified,)))
|
||||
for index, modified in zip(sent_indices, modified_messages)
|
||||
)
|
||||
replacements: Final = MappingProxyType(
|
||||
{index: message for index, message in zip(sent_indices, rewritten) if message is not None}
|
||||
)
|
||||
if len(replacements) != len(sent_indices):
|
||||
return None
|
||||
return tuple(replacements.get(index, message) for index, message in enumerate(structured_messages))
|
||||
|
||||
async def _apply_guardrail_on_response(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
|
|
@ -346,19 +399,7 @@ class PromptSecurityGuardrail(CustomGuardrail):
|
|||
return inputs
|
||||
|
||||
def _extract_texts_from_messages(self, messages: Sequence[Mapping[str, object]]) -> list[str]:
|
||||
"""Extract text content from messages."""
|
||||
texts: Final = []
|
||||
for message in messages:
|
||||
content = message.get("content")
|
||||
if isinstance(content, str):
|
||||
texts.append(content)
|
||||
elif isinstance(content, list):
|
||||
for item in content:
|
||||
if isinstance(item, dict) and item.get("type") == "text":
|
||||
text = item.get("text")
|
||||
if text:
|
||||
texts.append(text)
|
||||
return texts
|
||||
return [text for message in messages for text in message_slot_texts(message)]
|
||||
|
||||
async def _process_standalone_images(self, images: list[str], user_api_key_alias: str | None) -> None:
|
||||
"""Process standalone images from inputs (data URLs)."""
|
||||
|
|
@ -512,16 +553,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:
|
||||
|
|
@ -678,14 +721,13 @@ class PromptSecurityGuardrail(CustomGuardrail):
|
|||
|
||||
This allows checking tool results for indirect prompt injection when enabled.
|
||||
"""
|
||||
supported_roles: Final = ["system", "user", "assistant"]
|
||||
filtered_messages: Final = []
|
||||
transformed_count = 0
|
||||
filtered_count = 0
|
||||
|
||||
for message in messages:
|
||||
role = message.get("role", "")
|
||||
if role in supported_roles:
|
||||
if role in _PROTECT_ROLES:
|
||||
filtered_messages.append(message)
|
||||
else:
|
||||
if self.check_tool_results:
|
||||
|
|
|
|||
|
|
@ -23,6 +23,7 @@ def initialize_bedrock(litellm_params: LitellmParams, guardrail: Guardrail):
|
|||
prompt_attack_threshold=litellm_params.prompt_attack_threshold,
|
||||
pii_confidence_threshold=litellm_params.pii_confidence_threshold,
|
||||
chunk_budget_chars=litellm_params.chunk_budget_chars,
|
||||
contextual_grounding_from_messages=litellm_params.contextual_grounding_from_messages,
|
||||
default_on=litellm_params.default_on,
|
||||
disable_exception_on_block=litellm_params.disable_exception_on_block,
|
||||
mask_request_content=litellm_params.mask_request_content,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 ##
|
||||
|
|
|
|||
|
|
@ -9,10 +9,12 @@ import binascii
|
|||
import logging
|
||||
import os
|
||||
import uuid
|
||||
from collections.abc import Awaitable, Callable, Mapping, Sequence, Set
|
||||
from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence, Set
|
||||
from contextlib import asynccontextmanager
|
||||
from contextvars import ContextVar
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from types import MappingProxyType
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
|
|
@ -23,6 +25,7 @@ from typing import (
|
|||
TypedDict,
|
||||
)
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
from typing_extensions import NotRequired, ReadOnly
|
||||
|
||||
from litellm import DualCache
|
||||
|
|
@ -84,6 +87,9 @@ else:
|
|||
InternalUsageCache = Any
|
||||
|
||||
|
||||
_REQUEST_RATE_LIMIT_DATA: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
BATCH_RATE_LIMITER_SCRIPT: Final = """
|
||||
local results = {}
|
||||
local now = tonumber(ARGV[1])
|
||||
|
|
@ -2673,12 +2679,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
Returns list of descriptors for API key, user, team, team member, end user,
|
||||
model-specific, agent, and agent-session limits.
|
||||
"""
|
||||
from litellm.proxy.auth.auth_utils import (
|
||||
get_team_model_rpm_limit,
|
||||
get_team_model_tpm_limit,
|
||||
)
|
||||
|
||||
descriptors: Final = []
|
||||
descriptors: Final[list[RateLimitDescriptor]] = [] # mutable-ok: existing descriptor helpers append in place
|
||||
|
||||
# API Key rate limits
|
||||
if user_api_key_dict.api_key and (
|
||||
|
|
@ -2803,34 +2804,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
descriptors=descriptors,
|
||||
)
|
||||
|
||||
if (
|
||||
get_team_model_rpm_limit(user_api_key_dict) is not None
|
||||
or get_team_model_tpm_limit(user_api_key_dict) is not None
|
||||
):
|
||||
_tpm_limit_for_team_model: Final = get_team_model_tpm_limit(user_api_key_dict) or {}
|
||||
_rpm_limit_for_team_model: Final = get_team_model_rpm_limit(user_api_key_dict) or {}
|
||||
should_check_rate_limit = False
|
||||
if requested_model in _tpm_limit_for_team_model or requested_model in _rpm_limit_for_team_model:
|
||||
should_check_rate_limit = True
|
||||
|
||||
if should_check_rate_limit:
|
||||
model_specific_tpm_limit = None
|
||||
model_specific_rpm_limit = None
|
||||
if requested_model in _tpm_limit_for_team_model:
|
||||
model_specific_tpm_limit = _tpm_limit_for_team_model[requested_model]
|
||||
if requested_model in _rpm_limit_for_team_model:
|
||||
model_specific_rpm_limit = _rpm_limit_for_team_model[requested_model]
|
||||
descriptors.append(
|
||||
RateLimitDescriptor(
|
||||
key="model_per_team",
|
||||
value=f"{user_api_key_dict.team_id}:{requested_model}",
|
||||
rate_limit={
|
||||
"requests_per_unit": model_specific_rpm_limit,
|
||||
"tokens_per_unit": model_specific_tpm_limit,
|
||||
"window_size": self.window_size,
|
||||
},
|
||||
)
|
||||
)
|
||||
self._add_team_model_rate_limit_descriptor_from_metadata(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
requested_model=requested_model if isinstance(requested_model, str) else None,
|
||||
descriptors=descriptors,
|
||||
)
|
||||
|
||||
# Agent-level and session-level rate limits
|
||||
resolved_agent_id: Final = self._get_resolved_agent_id(user_api_key_dict, data)
|
||||
|
|
@ -3416,6 +3394,108 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
requested_model,
|
||||
)
|
||||
|
||||
async def _build_request_rate_limit_descriptors(
|
||||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
data: Mapping[str, object],
|
||||
call_type: str | None,
|
||||
) -> list[RateLimitDescriptor]: # mutable-ok: the shared generation reservation helpers require a list
|
||||
metadata: Final = _REQUEST_RATE_LIMIT_DATA.validate_python(
|
||||
user_api_key_dict.metadata or MappingProxyType({}) # pyright: ignore[reportUnknownMemberType] # validates the legacy auth metadata boundary
|
||||
)
|
||||
rpm_value: Final = metadata.get("rpm_limit_type")
|
||||
tpm_value: Final = metadata.get("tpm_limit_type")
|
||||
rpm_limit_type: Final = rpm_value if isinstance(rpm_value, str) else None
|
||||
tpm_limit_type: Final = tpm_value if isinstance(tpm_value, str) else None
|
||||
model_value: Final = data.get("model")
|
||||
requested_model: Final = model_value if isinstance(model_value, str) else None
|
||||
model_has_failures: Final = (
|
||||
await self._check_model_has_recent_failures(
|
||||
model=requested_model,
|
||||
parent_otel_span=user_api_key_dict.parent_otel_span,
|
||||
)
|
||||
if requested_model and self._is_dynamic_rate_limiting_enabled(rpm_limit_type, tpm_limit_type)
|
||||
else False
|
||||
)
|
||||
descriptors: Final = self._create_rate_limit_descriptors( # pyright: ignore[reportUnknownMemberType] # legacy helper reads a dictionary with validated keys
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
data=dict(data), # mutable-ok: legacy descriptor helpers accept a request dictionary
|
||||
rpm_limit_type=rpm_limit_type,
|
||||
tpm_limit_type=tpm_limit_type,
|
||||
model_has_failures=model_has_failures,
|
||||
call_type=call_type,
|
||||
)
|
||||
self._add_project_model_rate_limit_descriptor_from_metadata(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
requested_model=requested_model,
|
||||
descriptors=descriptors,
|
||||
)
|
||||
self.add_project_io_token_rate_limit_descriptors_from_metadata(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
requested_model=requested_model,
|
||||
descriptors=descriptors,
|
||||
)
|
||||
return [ # mutable-ok: the shared generation reservation helpers require a list
|
||||
*descriptors,
|
||||
*self.create_organization_rate_limit_descriptor(user_api_key_dict, requested_model),
|
||||
]
|
||||
|
||||
async def _release_request_capacity_when_admitted(
|
||||
self,
|
||||
admission: asyncio.Task[RateLimitResponse],
|
||||
acquisition: ParallelSlotAcquisition,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
) -> None:
|
||||
response: Final = await admission
|
||||
if response["overall_code"] == "OK":
|
||||
await self._release_parallel_request_slots(acquisition, user_api_key_dict.parent_otel_span)
|
||||
|
||||
@asynccontextmanager
|
||||
async def request_capacity(
|
||||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
model: str,
|
||||
*,
|
||||
request_data: Mapping[str, object] | None = None,
|
||||
) -> AsyncGenerator[None, None]:
|
||||
"""Charge one non-generation provider request to RPM and hold its concurrency slot."""
|
||||
data: Final = MappingProxyType({**(request_data or MappingProxyType({})), "model": model})
|
||||
descriptors: Final = await self._build_request_rate_limit_descriptors(user_api_key_dict, data, None)
|
||||
acquisition: Final = ParallelSlotAcquisition(
|
||||
slot_id=uuid.uuid4().hex,
|
||||
counter_keys=[ # mutable-ok: the shared slot-release contract requires a list
|
||||
self.create_rate_limit_keys(d["key"], d["value"], "max_parallel_requests")
|
||||
for d in descriptors
|
||||
if d["rate_limit"] is not None and d["rate_limit"].get("max_parallel_requests") is not None
|
||||
],
|
||||
)
|
||||
admission: Final = asyncio.create_task(
|
||||
self.should_rate_limit(
|
||||
descriptors=descriptors,
|
||||
parent_otel_span=user_api_key_dict.parent_otel_span,
|
||||
skip_tpm_check=True,
|
||||
parallel_slot_id=acquisition["slot_id"],
|
||||
)
|
||||
)
|
||||
try:
|
||||
response: Final = await asyncio.shield(admission)
|
||||
if response["overall_code"] == "OVER_LIMIT":
|
||||
self._handle_rate_limit_error(response, descriptors, model)
|
||||
yield
|
||||
finally:
|
||||
cleanup: Final = asyncio.create_task(
|
||||
self._release_request_capacity_when_admitted(admission, acquisition, user_api_key_dict)
|
||||
)
|
||||
cancellation: asyncio.CancelledError | None = None # rebind-ok: retain cancellation until cleanup finishes
|
||||
while not cleanup.done():
|
||||
try:
|
||||
await asyncio.shield(cleanup)
|
||||
except asyncio.CancelledError as exc:
|
||||
cancellation = exc # rebind-ok: retain the latest cancellation without interrupting slot release
|
||||
cleanup.result()
|
||||
if cancellation is not None:
|
||||
raise cancellation
|
||||
|
||||
async def async_pre_call_hook(
|
||||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
|
|
@ -3444,59 +3524,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
call_type=call_type,
|
||||
)
|
||||
|
||||
# Get rate limit types from metadata
|
||||
metadata: Final = user_api_key_dict.metadata or {}
|
||||
rpm_limit_type: Final = metadata.get("rpm_limit_type")
|
||||
tpm_limit_type: Final = metadata.get("tpm_limit_type")
|
||||
|
||||
# For dynamic mode, check if the model has recent failures
|
||||
model_has_failures = False
|
||||
requested_model: Final = data.get("model", None)
|
||||
|
||||
if (
|
||||
self._is_dynamic_rate_limiting_enabled(
|
||||
rpm_limit_type=rpm_limit_type,
|
||||
tpm_limit_type=tpm_limit_type,
|
||||
)
|
||||
and requested_model
|
||||
):
|
||||
model_has_failures = await self._check_model_has_recent_failures(
|
||||
model=requested_model,
|
||||
parent_otel_span=user_api_key_dict.parent_otel_span,
|
||||
)
|
||||
|
||||
# Create rate limit descriptors
|
||||
descriptors: Final = self._create_rate_limit_descriptors(
|
||||
request_data: Final = _REQUEST_RATE_LIMIT_DATA.validate_python(data)
|
||||
model_value: Final = request_data.get("model")
|
||||
requested_model: Final = model_value if isinstance(model_value, str) else None
|
||||
descriptors: Final = await self._build_request_rate_limit_descriptors(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
data=data,
|
||||
rpm_limit_type=rpm_limit_type,
|
||||
tpm_limit_type=tpm_limit_type,
|
||||
model_has_failures=model_has_failures,
|
||||
data=request_data,
|
||||
call_type=call_type,
|
||||
)
|
||||
|
||||
# Add team model rate limits from team_metadata
|
||||
self._add_team_model_rate_limit_descriptor_from_metadata(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
requested_model=requested_model,
|
||||
descriptors=descriptors,
|
||||
)
|
||||
|
||||
# Project Level Rate Limits
|
||||
self._add_project_model_rate_limit_descriptor_from_metadata(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
requested_model=requested_model,
|
||||
descriptors=descriptors,
|
||||
)
|
||||
self.add_project_io_token_rate_limit_descriptors_from_metadata(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
requested_model=requested_model,
|
||||
descriptors=descriptors,
|
||||
)
|
||||
|
||||
# Org Level Rate Limits
|
||||
descriptors.extend(self.create_organization_rate_limit_descriptor(user_api_key_dict, requested_model))
|
||||
|
||||
# Only check rate limits if we have descriptors with actual limits
|
||||
if descriptors:
|
||||
# First pass: RPM and max_parallel_requests sliding-window check.
|
||||
|
|
|
|||
142
litellm/proxy/hooks/prompt_cache_prediction.py
Normal file
142
litellm/proxy/hooks/prompt_cache_prediction.py
Normal 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
|
||||
|
|
@ -6,11 +6,10 @@ from typing import Final
|
|||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import PROXY_LLM_PROVIDER_FALLBACK
|
||||
from litellm.types.router import ModelGroupInfo
|
||||
from litellm.types.utils import PriorityReservationDict
|
||||
|
||||
PROXY_LLM_PROVIDER_FALLBACK: Final = "litellm_proxy"
|
||||
|
||||
|
||||
def resolve_llm_provider_for_rate_limit(
|
||||
model: str | None,
|
||||
|
|
|
|||
|
|
@ -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
278
litellm/proxy/management_endpoints/prompt_cache_prediction.py
Normal file
|
|
@ -0,0 +1,278 @@
|
|||
import time
|
||||
from collections.abc import Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Annotated, Final
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request
|
||||
from pydantic import BaseModel, JsonValue, TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm._internal_context import current_billing_time, pinned_billing_time
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.llms.anthropic.prompt_cache_prediction import (
|
||||
PromptPrefix,
|
||||
TokenCounter,
|
||||
UnsupportedPredictionTarget,
|
||||
cache_scope,
|
||||
count_prompt_tokens,
|
||||
parse_prompt,
|
||||
resolve_prediction_target,
|
||||
supported_prediction_headers,
|
||||
)
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.auth.auth_checks import can_key_call_resolved_model
|
||||
from litellm.proxy.auth.auth_utils import get_cache_prediction_deployments
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.common_utils.http_parsing_utils import (
|
||||
_read_request_body, # pyright: ignore[reportPrivateUsage, reportUnknownVariableType] # canonical parsed-body owner; validate its legacy result at the endpoint boundary
|
||||
)
|
||||
from litellm.proxy.common_utils.prompt_cache_pricing import price_cache_tokens
|
||||
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
|
||||
_PROXY_MaxParallelRequestsHandler_v3, # pyright: ignore[reportPrivateUsage] # use the configured proxy limiter's shared capacity owner
|
||||
)
|
||||
from litellm.proxy.hooks.prompt_cache_prediction import lookup
|
||||
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
|
||||
from litellm.types.management_endpoints.prompt_cache_prediction import (
|
||||
CacheCostScenario,
|
||||
CacheEvidence,
|
||||
CachePredictionArm,
|
||||
CachePredictionRequest,
|
||||
CachePredictionResponse,
|
||||
CacheTokenBuckets,
|
||||
)
|
||||
from litellm.types.router import Deployment
|
||||
from litellm.utils import get_prompt_cache_min_tokens
|
||||
|
||||
router: Final = APIRouter()
|
||||
_REQUEST_DATA: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
class _CallerSettings(BaseModel):
|
||||
config: Mapping[str, object] | None = None
|
||||
|
||||
|
||||
def has_request_transforms() -> bool:
|
||||
from litellm.proxy.hooks import PROXY_HOOKS
|
||||
|
||||
builtins: Final = frozenset(PROXY_HOOKS.values())
|
||||
hooks: Final = ("async_pre_call_hook", "async_pre_request_hook", "async_pre_call_deployment_hook")
|
||||
callbacks: Final = litellm.logging_callback_manager.get_custom_loggers_for_type(callback_type=CustomLogger)
|
||||
return any(
|
||||
type(callback) not in builtins
|
||||
and any(getattr(type(callback), hook) is not getattr(CustomLogger, hook) for hook in hooks)
|
||||
for callback in callbacks
|
||||
)
|
||||
|
||||
|
||||
def _buckets(prefix_tokens: int, suffix_tokens: int, read_tokens: int, ttl_seconds: int) -> CacheTokenBuckets:
|
||||
return CacheTokenBuckets(
|
||||
uncached_input_tokens=suffix_tokens,
|
||||
cache_read_input_tokens=read_tokens,
|
||||
cache_creation_5m_input_tokens=prefix_tokens - read_tokens if ttl_seconds == 300 else 0,
|
||||
cache_creation_1h_input_tokens=prefix_tokens - read_tokens if ttl_seconds == 3600 else 0,
|
||||
)
|
||||
|
||||
|
||||
def _scenario(model: str, deployment_id: str, tokens: CacheTokenBuckets) -> CacheCostScenario | None:
|
||||
cost: Final = price_cache_tokens(model=model, deployment_id=deployment_id, tokens=tokens)
|
||||
return CacheCostScenario(tokens=tokens, input_cost=cost) if cost is not None else None
|
||||
|
||||
|
||||
def _capacity_counter(
|
||||
limiter: _PROXY_MaxParallelRequestsHandler_v3,
|
||||
caller: UserAPIKeyAuth,
|
||||
model_name: str,
|
||||
request_data: Mapping[str, object],
|
||||
) -> TokenCounter:
|
||||
async def count(model: str, api_key: str, body: Mapping[str, JsonValue]) -> int | None:
|
||||
async with limiter.request_capacity(caller, model_name, request_data=request_data):
|
||||
return await count_prompt_tokens(model, api_key, body)
|
||||
|
||||
return count
|
||||
|
||||
|
||||
def _capacity_request_data(
|
||||
http_request: Request, caller: UserAPIKeyAuth, request_data: Mapping[str, object]
|
||||
) -> Mapping[str, object]:
|
||||
# The parsed-body cache retains only original top-level keys. Replay the
|
||||
# shared idempotent tag merges on limiter-only data when auth added metadata.
|
||||
data: Final = dict(request_data) # mutable-ok: the existing tag merge owners accept a dictionary out-param
|
||||
LiteLLMProxyRequestSetup.apply_client_tag_policy_pre_auth(http_request, data, caller) # pyright: ignore[reportUnknownMemberType] # legacy tag owner takes the validated capacity dictionary
|
||||
LiteLLMProxyRequestSetup.apply_key_tags_pre_auth(data, caller) # pyright: ignore[reportUnknownMemberType] # legacy tag owner merges trusted key tags into capacity metadata
|
||||
return MappingProxyType(data)
|
||||
|
||||
|
||||
async def predict_arm(
|
||||
deployment: Deployment,
|
||||
body: Mapping[str, JsonValue],
|
||||
prefix: PromptPrefix,
|
||||
caller_key_hash: str,
|
||||
cache: DualCache,
|
||||
token_counter: TokenCounter,
|
||||
) -> CachePredictionArm:
|
||||
deployment_id: Final = deployment.model_info.id or ""
|
||||
params: Final = deployment.litellm_params
|
||||
unknown: Final = CachePredictionArm(deployment_id=deployment_id, model=params.model)
|
||||
if deployment.model_info.blocked:
|
||||
return unknown.model_copy(update=MappingProxyType({"reason": "unsupported_deployment_configuration"}))
|
||||
target: Final = resolve_prediction_target(params)
|
||||
if isinstance(target, UnsupportedPredictionTarget):
|
||||
return unknown.model_copy(update=MappingProxyType({"reason": target.reason}))
|
||||
model: Final = target.model
|
||||
api_key: Final = target.api_key
|
||||
total_count: Final = await token_counter(model, api_key, body)
|
||||
prefix_count: Final = await token_counter(model, api_key, prefix.prefix_body)
|
||||
if total_count is None or prefix_count is None or total_count < prefix_count:
|
||||
return unknown.model_copy(update=MappingProxyType({"reason": "token_count_unavailable"}))
|
||||
scope: Final = cache_scope(caller_key_hash, deployment_id, api_key, model)
|
||||
observation: Final = await lookup(cache, scope, prefix)
|
||||
exact: Final = observation is not None and observation.fingerprint == prefix.fingerprint
|
||||
cacheable: Final = observation.cached_tokens if exact and observation is not None else prefix_count
|
||||
if cacheable > total_count or (observation is not None and observation.cached_tokens > cacheable):
|
||||
return unknown.model_copy(update=MappingProxyType({"reason": "inconsistent_prefix_token_count"}))
|
||||
suffix: Final = total_count - cacheable
|
||||
evidence: Final = (
|
||||
CacheEvidence(observed_at=observation.observed_at, expires_at=observation.expires_at)
|
||||
if observation is not None
|
||||
else None
|
||||
)
|
||||
if cacheable < get_prompt_cache_min_tokens(params.model):
|
||||
disabled: Final = _scenario(model, deployment_id, CacheTokenBuckets(uncached_input_tokens=total_count))
|
||||
if disabled is None:
|
||||
return unknown.model_copy(update=MappingProxyType({"reason": "pricing_unavailable"}))
|
||||
return CachePredictionArm(
|
||||
deployment_id=deployment_id,
|
||||
model=model,
|
||||
cache_state="disabled",
|
||||
reason="below_cache_minimum",
|
||||
estimate=disabled,
|
||||
cold=disabled,
|
||||
warm=disabled,
|
||||
token_count_source="anthropic_count_tokens",
|
||||
)
|
||||
fresh: Final = observation is not None and observation.expires_at > time.time()
|
||||
read: Final = observation.cached_tokens if fresh and observation is not None else 0
|
||||
with pinned_billing_time(current_billing_time()):
|
||||
cold: Final = _scenario(model, deployment_id, _buckets(cacheable, suffix, 0, prefix.ttl_seconds))
|
||||
warm: Final = _scenario(model, deployment_id, _buckets(cacheable, suffix, cacheable, prefix.ttl_seconds))
|
||||
estimate: Final = _scenario(model, deployment_id, _buckets(cacheable, suffix, read, prefix.ttl_seconds))
|
||||
if cold is None or warm is None or estimate is None:
|
||||
return unknown.model_copy(update=MappingProxyType({"reason": "pricing_unavailable"}))
|
||||
return CachePredictionArm(
|
||||
deployment_id=deployment_id,
|
||||
model=model,
|
||||
cache_state="warm" if fresh and exact else "partial" if fresh else "stale" if observation else "unknown",
|
||||
reason=None if fresh else "observation_expired" if observation else "no_compatible_observation",
|
||||
estimate=estimate,
|
||||
cold=cold,
|
||||
warm=warm,
|
||||
evidence=evidence,
|
||||
token_count_source="anthropic_count_tokens",
|
||||
)
|
||||
|
||||
|
||||
@router.post(
|
||||
"/cost/predict-cache",
|
||||
tags=["Cost Tracking"], # mutable-ok: FastAPI requires a list for OpenAPI tags
|
||||
response_model=CachePredictionResponse,
|
||||
)
|
||||
async def predict_cache_cost(
|
||||
request: CachePredictionRequest,
|
||||
http_request: Request,
|
||||
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
|
||||
) -> CachePredictionResponse:
|
||||
"""Compare the next native Anthropic request on two configured deployment IDs.
|
||||
|
||||
Estimates use provider token counting and recent successful cache telemetry for this key.
|
||||
Unknown cache state uses the cold scenario when prices/counts are available. Cache observations
|
||||
do not guarantee retention. v0 supports one message-content breakpoint, text and client tools;
|
||||
system/tool-only breakpoints, thinking, images, nondefault Anthropic versions, beta headers and
|
||||
request transforms are unknown.
|
||||
Each provider count consumes one RPM unit and holds concurrency capacity; a comparison uses
|
||||
up to four counts. The legacy rate limiter returns unknown without contacting the provider.
|
||||
This endpoint does not generate tokens, prewarm caches, choose a model or alter routing.
|
||||
"""
|
||||
from litellm.proxy.proxy_server import llm_router, proxy_logging_obj
|
||||
|
||||
if llm_router is None:
|
||||
raise HTTPException(status_code=503, detail="Model router is unavailable")
|
||||
deployments: Final = get_cache_prediction_deployments(
|
||||
current_deployment_id=request.current_deployment_id,
|
||||
candidate_deployment_id=request.candidate_deployment_id,
|
||||
llm_router=llm_router,
|
||||
team_id=user_api_key_dict.team_id,
|
||||
)
|
||||
if deployments is None:
|
||||
raise HTTPException(status_code=404, detail="Deployment not found")
|
||||
current, candidate = deployments
|
||||
for deployment in (current, candidate):
|
||||
await can_key_call_resolved_model(
|
||||
model=deployment.model_name,
|
||||
llm_model_list=llm_router.get_model_list(),
|
||||
valid_token=user_api_key_dict,
|
||||
llm_router=llm_router,
|
||||
)
|
||||
prefix: Final = parse_prompt(request.request)
|
||||
caller: Final = user_api_key_dict.api_key
|
||||
caller_settings: Final = _CallerSettings.model_validate(user_api_key_dict, from_attributes=True)
|
||||
unsupported_transform: Final = bool(caller_settings.config) or has_request_transforms()
|
||||
unsupported_headers: Final = not supported_prediction_headers(http_request.headers)
|
||||
limiter: Final = proxy_logging_obj.get_proxy_hook("parallel_request_limiter")
|
||||
if (
|
||||
prefix is None
|
||||
or not caller
|
||||
or unsupported_transform
|
||||
or unsupported_headers
|
||||
or not isinstance(limiter, _PROXY_MaxParallelRequestsHandler_v3)
|
||||
):
|
||||
reason: Final = (
|
||||
"unsupported_provider_headers"
|
||||
if unsupported_headers
|
||||
else "unsupported_request_transform"
|
||||
if unsupported_transform
|
||||
else "unsupported_prompt_shape"
|
||||
if prefix is None
|
||||
else "caller_identity_unavailable"
|
||||
if not caller
|
||||
else "limiter_unavailable"
|
||||
)
|
||||
return CachePredictionResponse(
|
||||
stay=CachePredictionArm(deployment_id=request.current_deployment_id, reason=reason),
|
||||
switch=CachePredictionArm(deployment_id=request.candidate_deployment_id, reason=reason),
|
||||
switch_delta=None,
|
||||
cache_rebuild_penalty=None,
|
||||
)
|
||||
request_data: Final = _capacity_request_data(
|
||||
http_request, user_api_key_dict, _REQUEST_DATA.validate_python(await _read_request_body(http_request))
|
||||
)
|
||||
stay: Final = await predict_arm(
|
||||
current,
|
||||
request.request,
|
||||
prefix,
|
||||
caller,
|
||||
proxy_logging_obj.internal_usage_cache.dual_cache,
|
||||
_capacity_counter(limiter, user_api_key_dict, current.model_name, request_data),
|
||||
)
|
||||
switch: Final = (
|
||||
stay
|
||||
if current.model_info.id == candidate.model_info.id
|
||||
else await predict_arm(
|
||||
candidate,
|
||||
request.request,
|
||||
prefix,
|
||||
caller,
|
||||
proxy_logging_obj.internal_usage_cache.dual_cache,
|
||||
_capacity_counter(limiter, user_api_key_dict, candidate.model_name, request_data),
|
||||
)
|
||||
)
|
||||
return CachePredictionResponse(
|
||||
stay=stay,
|
||||
switch=switch,
|
||||
switch_delta=(switch.estimate.input_cost - stay.estimate.input_cost)
|
||||
if switch.estimate is not None and stay.estimate is not None
|
||||
else None,
|
||||
cache_rebuild_penalty=(switch.estimate.input_cost - switch.warm.input_cost)
|
||||
if switch.estimate is not None and switch.warm is not None
|
||||
else None,
|
||||
)
|
||||
|
|
@ -156,6 +156,7 @@ from litellm.repositories.verification_token_repository import (
|
|||
VerificationTokenRepository,
|
||||
)
|
||||
from litellm.router import Router
|
||||
from litellm.types.proxy.auth.auth_checks import UserNotFoundError
|
||||
from litellm.types.proxy.management_endpoints.common_daily_activity import (
|
||||
SpendAnalyticsPaginatedResponse,
|
||||
)
|
||||
|
|
@ -179,6 +180,7 @@ from litellm.types.proxy.management_endpoints.team_endpoints import (
|
|||
if TYPE_CHECKING:
|
||||
from prisma import Prisma
|
||||
from prisma import models as prisma_models
|
||||
from prisma import types as prisma_types
|
||||
|
||||
router: Final = APIRouter()
|
||||
|
||||
|
|
@ -429,27 +431,26 @@ async def _refresh_cached_team(
|
|||
)
|
||||
|
||||
|
||||
async def _can_manage_team(
|
||||
team_obj: LiteLLM_TeamTable,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
) -> bool:
|
||||
"""True for a proxy admin, an admin of this team, or an org admin for the team's organization."""
|
||||
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN:
|
||||
return True
|
||||
|
||||
if _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team_obj):
|
||||
return True
|
||||
|
||||
return await _is_user_org_admin_for_team(user_api_key_dict=user_api_key_dict, team_obj=team_obj)
|
||||
|
||||
|
||||
async def _verify_team_access(
|
||||
team_obj: LiteLLM_TeamTable,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
) -> None:
|
||||
"""
|
||||
Verify the caller is authorized to manage the given team.
|
||||
|
||||
Access is granted if:
|
||||
- Caller is a proxy admin, OR
|
||||
- Caller is an org admin for the team's organization, OR
|
||||
- Caller is a team admin of this team
|
||||
|
||||
Raises HTTPException(403) otherwise.
|
||||
"""
|
||||
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN:
|
||||
return
|
||||
|
||||
if _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team_obj):
|
||||
return
|
||||
|
||||
if await _is_user_org_admin_for_team(user_api_key_dict=user_api_key_dict, team_obj=team_obj):
|
||||
"""Raise HTTPException(403) unless the caller can manage the given team."""
|
||||
if await _can_manage_team(team_obj=team_obj, user_api_key_dict=user_api_key_dict):
|
||||
return
|
||||
|
||||
raise HTTPException(
|
||||
|
|
@ -4368,6 +4369,20 @@ async def _hydrate_member_user_details(
|
|||
return tuple(hydrate(m) for m in members)
|
||||
|
||||
|
||||
class _OrganizationModelsRow(BaseModel):
|
||||
models: list[str] = [] # mutable-ok: pydantic field default
|
||||
|
||||
|
||||
class _TeamRowWithOrganization(BaseModel):
|
||||
litellm_organization_table: _OrganizationModelsRow | None = None
|
||||
|
||||
|
||||
def _parent_organization_models(team_row: BaseModel) -> list[str] | None:
|
||||
"""Return the parent org's model allow-list, or None when the team has no org."""
|
||||
organization: Final = _TeamRowWithOrganization.model_validate(team_row.model_dump()).litellm_organization_table
|
||||
return organization.models if organization is not None else None
|
||||
|
||||
|
||||
async def _resolve_team_access_group_resources(
|
||||
_team_info: TeamInfoResponseObjectTeamTable,
|
||||
) -> TeamInfoResponseObjectTeamTable:
|
||||
|
|
@ -4439,7 +4454,11 @@ async def team_info(
|
|||
try:
|
||||
team_info: BaseModel | None = await _team_db(prisma_client).find_unique(
|
||||
where={"team_id": team_id},
|
||||
include={"litellm_model_table": True, "object_permission": True},
|
||||
include={
|
||||
"litellm_model_table": True,
|
||||
"object_permission": True,
|
||||
"litellm_organization_table": True,
|
||||
},
|
||||
)
|
||||
if team_info is None:
|
||||
raise Exception
|
||||
|
|
@ -4448,9 +4467,12 @@ async def team_info(
|
|||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail={"message": f"Team not found, passed team id: {team_id}."},
|
||||
)
|
||||
await validate_membership(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
team_table=LiteLLM_TeamTable.model_validate(team_info.model_dump()),
|
||||
team_table: Final = LiteLLM_TeamTable.model_validate(team_info.model_dump())
|
||||
await validate_membership(user_api_key_dict=user_api_key_dict, team_table=team_table)
|
||||
organization_models: Final[list[str] | None] = (
|
||||
_parent_organization_models(team_info)
|
||||
if await _can_manage_team(team_obj=team_table, user_api_key_dict=user_api_key_dict)
|
||||
else None
|
||||
)
|
||||
|
||||
## GET ALL KEYS ##
|
||||
|
|
@ -4510,7 +4532,10 @@ async def team_info(
|
|||
members=resolved_team_info.members_with_roles,
|
||||
)
|
||||
hydrated_team_info: Final = resolved_team_info.model_copy(
|
||||
update={"members_with_roles": hydrated_members} # mutable-ok: pydantic update payload
|
||||
update={ # mutable-ok: pydantic update payload
|
||||
"members_with_roles": hydrated_members,
|
||||
"organization_models": organization_models,
|
||||
}
|
||||
)
|
||||
|
||||
response_object: Final = TeamInfoResponseObject(
|
||||
|
|
@ -4857,6 +4882,26 @@ async def _get_org_admin_org_ids(
|
|||
return org_ids if org_ids else None
|
||||
|
||||
|
||||
async def _get_user_team_ids_from_db(
|
||||
user_id: str,
|
||||
prisma_client: PrismaClient,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
) -> tuple[str, ...]:
|
||||
try:
|
||||
user: Final = await get_user_object(
|
||||
user_id=user_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
user_id_upsert=False,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
check_db_only=True,
|
||||
)
|
||||
except UserNotFoundError:
|
||||
return ()
|
||||
return tuple(user.teams or ()) if user is not None else ()
|
||||
|
||||
|
||||
async def _build_team_list_where_conditions(
|
||||
prisma_client: PrismaClient,
|
||||
team_id: str | None,
|
||||
|
|
@ -4867,12 +4912,16 @@ async def _build_team_list_where_conditions(
|
|||
search: str | None = None,
|
||||
search_team_id_match: TeamIdSearchMatch = "exact",
|
||||
org_admin_org_ids: list[str] | None = None,
|
||||
own_team_ids: tuple[str, ...] = (),
|
||||
user_api_key_cache: UserApiKeyCache | None = None,
|
||||
proxy_logging_obj: ProxyLogging | None = None,
|
||||
) -> dict[str, object] | None:
|
||||
"""
|
||||
Build where conditions for team list query.
|
||||
|
||||
An org admin listing their own teams sees the union of the teams in the
|
||||
orgs they administer and `own_team_ids`, the teams they are a member of.
|
||||
|
||||
Returns None when the query is guaranteed to yield no results (e.g. user
|
||||
has no team memberships), allowing the caller to skip the DB round-trip.
|
||||
"""
|
||||
|
|
@ -4895,6 +4944,11 @@ async def _build_team_list_where_conditions(
|
|||
|
||||
if organization_id:
|
||||
where_conditions["organization_id"] = organization_id
|
||||
elif org_admin_org_ids is not None and own_team_ids:
|
||||
org_or_membership_scope: Final[prisma_types.LiteLLM_TeamTableWhereInput] = {
|
||||
"OR": [{"organization_id": {"in": org_admin_org_ids}}, {"team_id": {"in": list(own_team_ids)}}]
|
||||
}
|
||||
where_conditions["AND"] = [org_or_membership_scope]
|
||||
elif org_admin_org_ids is not None:
|
||||
# Org admin: always scope to their orgs, even when filtering by user_id.
|
||||
where_conditions["organization_id"] = {"in": org_admin_org_ids}
|
||||
|
|
@ -5026,66 +5080,72 @@ async def _enforce_list_team_v2_access(
|
|||
prisma_client: PrismaClient,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
) -> tuple[str | None, list[str] | None]:
|
||||
) -> tuple[str | None, list[str] | None, tuple[str, ...]]:
|
||||
"""Enforce access control for list_team_v2.
|
||||
|
||||
- Proxy admins and admin viewers can query any teams.
|
||||
- Org admins can query teams within their organizations.
|
||||
- Org admins can query teams within their organizations, plus the teams
|
||||
they are a member of when listing their own teams.
|
||||
- Regular users can only query their own teams.
|
||||
|
||||
Returns the (possibly overridden) user_id and org_admin_org_ids.
|
||||
Returns the (possibly overridden) user_id, org_admin_org_ids and, for an
|
||||
org admin's own query, the caller's own team ids.
|
||||
"""
|
||||
is_proxy_admin: Final = _user_has_admin_view(user_api_key_dict)
|
||||
org_admin_org_ids: list[str] | None = None
|
||||
caller_user_id: Final = user_api_key_dict.user_id
|
||||
|
||||
if is_proxy_admin:
|
||||
return user_id, org_admin_org_ids
|
||||
return user_id, None, ()
|
||||
|
||||
# Always check org admin status so that even own-queries see
|
||||
# the full set of organisation teams, not just direct memberships.
|
||||
if user_api_key_dict.user_id:
|
||||
org_admin_org_ids = await _get_org_admin_org_ids(
|
||||
user_id=user_api_key_dict.user_id,
|
||||
org_admin_org_ids: Final = (
|
||||
await _get_org_admin_org_ids(
|
||||
user_id=caller_user_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
if caller_user_id
|
||||
else None
|
||||
)
|
||||
|
||||
if org_admin_org_ids is not None:
|
||||
if caller_user_id and org_admin_org_ids is not None:
|
||||
# Org admin: validate org_id filter if provided
|
||||
if organization_id and organization_id not in org_admin_org_ids:
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail={"error": "You can only view teams within your organizations."},
|
||||
)
|
||||
# When the caller is an org admin querying their own teams (or no
|
||||
# specific user), null out user_id so that
|
||||
# _build_team_list_where_conditions scopes only by organization_id
|
||||
# — org admins should see all teams in their orgs, not just teams
|
||||
# they are a direct member of. Keep user_id when the org admin
|
||||
# explicitly queries a *different* user's teams.
|
||||
if user_id is None or user_id == user_api_key_dict.user_id:
|
||||
user_id = None
|
||||
is_own_query: Final = user_id is None or user_id == caller_user_id
|
||||
own_team_ids: Final = (
|
||||
await _get_user_team_ids_from_db(
|
||||
user_id=caller_user_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
if is_own_query
|
||||
else ()
|
||||
)
|
||||
verbose_proxy_logger.debug(
|
||||
"list_team_v2: org admin access for user=%s, org_ids=%s, user_id_filter=%s",
|
||||
user_api_key_dict.user_id,
|
||||
_sanitize_for_log(caller_user_id),
|
||||
org_admin_org_ids,
|
||||
user_id,
|
||||
_sanitize_for_log(None if is_own_query else user_id),
|
||||
)
|
||||
else:
|
||||
# Not an org admin — fall back to standard route check
|
||||
if not allowed_route_check_inside_route(user_api_key_dict=user_api_key_dict, requested_user_id=user_id):
|
||||
raise HTTPException(
|
||||
status_code=401,
|
||||
detail={
|
||||
"error": f"Only admin users can query all teams/other teams. Your user role={user_api_key_dict.user_role}"
|
||||
},
|
||||
)
|
||||
# Regular user — auto-inject caller's user_id
|
||||
if user_id is None:
|
||||
user_id = user_api_key_dict.user_id
|
||||
return None if is_own_query else user_id, org_admin_org_ids, own_team_ids
|
||||
|
||||
return user_id, org_admin_org_ids
|
||||
# Not an org admin — fall back to standard route check
|
||||
if not allowed_route_check_inside_route(user_api_key_dict=user_api_key_dict, requested_user_id=user_id):
|
||||
raise HTTPException(
|
||||
status_code=401,
|
||||
detail={
|
||||
"error": f"Only admin users can query all teams/other teams. Your user role={user_api_key_dict.user_role}"
|
||||
},
|
||||
)
|
||||
# Regular user — auto-inject caller's user_id
|
||||
return user_id if user_id is not None else caller_user_id, None, ()
|
||||
|
||||
|
||||
@router.get(
|
||||
|
|
@ -5163,7 +5223,7 @@ async def list_team_v2(
|
|||
)
|
||||
|
||||
# --- Access control ---
|
||||
user_id, org_admin_org_ids = await _enforce_list_team_v2_access(
|
||||
user_id, org_admin_org_ids, own_team_ids = await _enforce_list_team_v2_access(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
user_id=user_id,
|
||||
organization_id=organization_id,
|
||||
|
|
@ -5195,6 +5255,7 @@ async def list_team_v2(
|
|||
search=search,
|
||||
search_team_id_match=search_team_id_match,
|
||||
org_admin_org_ids=org_admin_org_ids,
|
||||
own_team_ids=own_team_ids,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
|
@ -5291,17 +5352,16 @@ async def _authorize_and_filter_teams(
|
|||
|
||||
- Proxy admins: all teams (or filtered by user_id if provided).
|
||||
- Org admins: teams from their orgs (scoped to user_id if provided).
|
||||
- Own query (user_id matches caller): teams the user is a member of.
|
||||
- Own query (user_id matches caller): teams the user is a member of, across all orgs.
|
||||
- Others: 401.
|
||||
"""
|
||||
is_proxy_admin: Final = _user_has_admin_view(user_api_key_dict)
|
||||
is_own_query: Final = (
|
||||
user_id is not None and user_api_key_dict.user_id is not None and user_api_key_dict.user_id == user_id
|
||||
)
|
||||
allowed_org_ids: list[str] | None = None
|
||||
|
||||
if not is_proxy_admin:
|
||||
is_own_query: Final = (
|
||||
user_id is not None and user_api_key_dict.user_id is not None and user_api_key_dict.user_id == user_id
|
||||
)
|
||||
|
||||
# Check if user is an org admin (even for own queries, so they see org teams)
|
||||
if user_api_key_dict.user_id is not None:
|
||||
caller_user: Final = await get_user_object(
|
||||
|
|
@ -5328,33 +5388,30 @@ async def _authorize_and_filter_teams(
|
|||
},
|
||||
)
|
||||
|
||||
if allowed_org_ids is not None:
|
||||
# Org admin: query DB for teams in their orgs
|
||||
if allowed_org_ids is not None and not is_own_query:
|
||||
org_teams: Final = await _raw_team_db(TeamRepository(prisma_client)).find_many(
|
||||
where={"organization_id": {"in": allowed_org_ids}},
|
||||
include={"litellm_model_table": True},
|
||||
)
|
||||
if not user_id:
|
||||
return list(org_teams)
|
||||
# Filter org teams to only those where the target user is a member
|
||||
return [
|
||||
team
|
||||
for team in org_teams
|
||||
if team.members_with_roles and any(m.get("user_id") == user_id for m in team.members_with_roles)
|
||||
]
|
||||
elif user_id:
|
||||
# Regular user: fetch all and filter by membership (Prisma can't filter JSON arrays)
|
||||
response: Final = await _raw_team_db(TeamRepository(prisma_client)).find_many(
|
||||
include={"litellm_model_table": True}
|
||||
)
|
||||
return [
|
||||
team
|
||||
for team in response
|
||||
if team.members_with_roles and any(m.get("user_id") == user_id for m in team.members_with_roles)
|
||||
]
|
||||
else:
|
||||
|
||||
response: Final = await _raw_team_db(TeamRepository(prisma_client)).find_many(include={"litellm_model_table": True})
|
||||
if not user_id:
|
||||
# Proxy admin: all teams
|
||||
return list(await _raw_team_db(TeamRepository(prisma_client)).find_many(include={"litellm_model_table": True}))
|
||||
return list(response)
|
||||
|
||||
# Prisma can't filter JSON arrays, so membership is filtered in Python
|
||||
return [
|
||||
team
|
||||
for team in response
|
||||
if team.members_with_roles and any(m.get("user_id") == user_id for m in team.members_with_roles)
|
||||
]
|
||||
|
||||
|
||||
@router.get("/team/list", tags=["team management"], dependencies=[Depends(user_api_key_auth)])
|
||||
|
|
|
|||
|
|
@ -101,6 +101,7 @@ from litellm.proxy.common_utils.admin_ui_utils import (
|
|||
admin_ui_disabled,
|
||||
show_missing_vars_in_env,
|
||||
)
|
||||
from litellm.proxy.common_utils.html_forms.default_credentials_hint import should_hide_default_credentials_hint
|
||||
from litellm.proxy.common_utils.html_forms.jwt_display_template import (
|
||||
jwt_display_template,
|
||||
)
|
||||
|
|
@ -1110,10 +1111,7 @@ async def google_login(
|
|||
|
||||
from fastapi.responses import HTMLResponse
|
||||
|
||||
hide_default_credentials_hint: Final = (
|
||||
os.getenv("LITELLM_HIDE_DEFAULT_CREDENTIALS_HINT", "false").lower() == "true"
|
||||
or general_settings.get("hide_default_credentials_hint", False) is True
|
||||
)
|
||||
hide_default_credentials_hint: Final = should_hide_default_credentials_hint(general_settings)
|
||||
form_response: Final = HTMLResponse(
|
||||
content=build_ui_login_form(
|
||||
show_deprecation_banner=True,
|
||||
|
|
|
|||
|
|
@ -115,15 +115,15 @@ async def run_team_metadata_validation(
|
|||
"error": f"custom_team_metadata_validate is an Enterprise feature. {CommonProxyErrors.not_premium_user.value}"
|
||||
},
|
||||
)
|
||||
if not (
|
||||
inspect.iscoroutinefunction(validator) or inspect.iscoroutinefunction(getattr(validator, "__call__", None))
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
detail={ # mutable-ok: HTTPException.detail has no immutable form
|
||||
"error": "custom_team_metadata_validate must be an async function"
|
||||
},
|
||||
)
|
||||
if not inspect.iscoroutinefunction(validator):
|
||||
validator_call: Final = getattr(validator, "__call__", None) # noqa: B004 # value unwrap for the functor check
|
||||
if not inspect.iscoroutinefunction(validator_call):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
detail={ # mutable-ok: HTTPException.detail has no immutable form
|
||||
"error": "custom_team_metadata_validate must be an async function"
|
||||
},
|
||||
)
|
||||
|
||||
try:
|
||||
raw_result: Final = await asyncio.wait_for(validator(payload), timeout=timeout_seconds)
|
||||
|
|
|
|||
|
|
@ -78,6 +78,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import (
|
|||
)
|
||||
from litellm.proxy.openai_files_endpoints.general_upload_validation import (
|
||||
MB,
|
||||
check_allowed_extension,
|
||||
check_blocked_extension,
|
||||
check_unsafe_filename,
|
||||
check_upload_file_size,
|
||||
|
|
@ -473,6 +474,11 @@ async def create_file(
|
|||
if general_size_failure is not None:
|
||||
raise_upload_validation_failure(general_size_failure)
|
||||
|
||||
allowed_extensions: Final = coerce_optional_str_list_setting(general_settings.get("allowed_file_extensions"))
|
||||
allowed_extension_failure: Final = check_allowed_extension(file.filename, allowed_extensions)
|
||||
if allowed_extension_failure is not None:
|
||||
raise_upload_validation_failure(allowed_extension_failure)
|
||||
|
||||
blocked_extensions: Final = coerce_optional_str_list_setting(general_settings.get("blocked_file_extensions"))
|
||||
blocked_extension_failure: Final = check_blocked_extension(file.filename, blocked_extensions)
|
||||
if blocked_extension_failure is not None:
|
||||
|
|
|
|||
|
|
@ -2,8 +2,8 @@
|
|||
Upload validation applied to every purpose at POST /v1/files.
|
||||
|
||||
batch_file_validation.py checks the JSONL shape of purpose="batch" uploads; this
|
||||
module applies the same fast-fail-before-forwarding shape (size cap, blocked
|
||||
extensions, path-traversal filenames) regardless of purpose.
|
||||
module applies the same fast-fail-before-forwarding shape (size cap, allowed and
|
||||
blocked extensions, path-traversal filenames) regardless of purpose.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
|
@ -31,10 +31,9 @@ def coerce_optional_int_setting(raw: object) -> int | None:
|
|||
raise TypeError(f"expected an integer, got {raw!r}")
|
||||
|
||||
|
||||
def coerce_optional_str_list_setting(raw: object) -> tuple[str, ...]:
|
||||
"""A general_settings value declared as an optional list of strings, e.g. blocked_file_extensions."""
|
||||
def coerce_optional_str_list_setting(raw: object) -> tuple[str, ...] | None:
|
||||
if raw is None:
|
||||
return ()
|
||||
return None
|
||||
if not isinstance(raw, list) or not all(isinstance(item, str) for item in raw):
|
||||
raise TypeError(f"expected a list of strings, got {raw!r}")
|
||||
return tuple(raw)
|
||||
|
|
@ -46,6 +45,11 @@ class UploadedFileTooLarge:
|
|||
limit_mb: int
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class UploadedFileExtensionNotAllowed:
|
||||
extension: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class UploadedFileBlockedExtension:
|
||||
extension: str
|
||||
|
|
@ -56,7 +60,9 @@ class UploadedFileUnsafeFilename:
|
|||
filename: str
|
||||
|
||||
|
||||
UploadValidationFailure = UploadedFileTooLarge | UploadedFileBlockedExtension | UploadedFileUnsafeFilename
|
||||
UploadValidationFailure = (
|
||||
UploadedFileTooLarge | UploadedFileExtensionNotAllowed | UploadedFileBlockedExtension | UploadedFileUnsafeFilename
|
||||
)
|
||||
|
||||
|
||||
def _file_size_bytes(file_source: bytes | BinaryIO) -> int:
|
||||
|
|
@ -81,19 +87,35 @@ def check_upload_file_size(
|
|||
return None
|
||||
|
||||
|
||||
def _normalized_extension(filename: str | None) -> str:
|
||||
if not filename:
|
||||
return ""
|
||||
try:
|
||||
return Path(safe_filename(filename)).suffix.lower()
|
||||
except ValueError:
|
||||
return ""
|
||||
|
||||
|
||||
def check_allowed_extension(
|
||||
filename: str | None,
|
||||
allowed_extensions: tuple[str, ...] | None,
|
||||
) -> UploadedFileExtensionNotAllowed | None:
|
||||
if allowed_extensions is None:
|
||||
return None
|
||||
extension: Final = _normalized_extension(filename)
|
||||
normalized_allowed: Final = frozenset(item.lower() for item in allowed_extensions)
|
||||
if extension and extension in normalized_allowed:
|
||||
return None
|
||||
return UploadedFileExtensionNotAllowed(extension=extension)
|
||||
|
||||
|
||||
def check_blocked_extension(
|
||||
filename: str | None,
|
||||
blocked_extensions: tuple[str, ...],
|
||||
blocked_extensions: tuple[str, ...] | None,
|
||||
) -> UploadedFileBlockedExtension | None:
|
||||
if not blocked_extensions or not filename:
|
||||
if not blocked_extensions:
|
||||
return None
|
||||
try:
|
||||
extension: Final = Path(safe_filename(filename)).suffix.lower()
|
||||
except ValueError:
|
||||
return None
|
||||
# The uploaded name's extension is normalized above; blocked_extensions comes
|
||||
# straight from config.yaml or the DB and is normalized here too, so a
|
||||
# differently-cased entry (".EXE") still catches a lowercase upload.
|
||||
extension: Final = _normalized_extension(filename)
|
||||
normalized_blocked: Final = frozenset(item.lower() for item in blocked_extensions)
|
||||
if extension and extension in normalized_blocked:
|
||||
return UploadedFileBlockedExtension(extension=extension)
|
||||
|
|
@ -128,6 +150,17 @@ def raise_upload_validation_failure(failure: UploadValidationFailure) -> NoRetur
|
|||
param="file",
|
||||
code=413,
|
||||
)
|
||||
case UploadedFileExtensionNotAllowed(extension=extension):
|
||||
raise ProxyException(
|
||||
message=(
|
||||
(f"File extension '{extension}'" if extension else "A file without an extension")
|
||||
+ " is not in this proxy's allowed_file_extensions setting. "
|
||||
"The file was not forwarded to the provider."
|
||||
),
|
||||
type="invalid_request_error",
|
||||
param="file",
|
||||
code=400,
|
||||
)
|
||||
case UploadedFileBlockedExtension(extension=extension):
|
||||
raise ProxyException(
|
||||
message=(
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -58,15 +58,6 @@ class UndeliverableStreamRewrite(Exception):
|
|||
self.guardrail_name: Final = guardrail_name
|
||||
|
||||
|
||||
class UnappliableRequestRewrite(Exception):
|
||||
def __init__(self, guardrail_name: str) -> None:
|
||||
super().__init__(
|
||||
f"Guardrail '{guardrail_name}' rewrote the request in a way this endpoint cannot apply, "
|
||||
"so the request was rejected rather than sent unrewritten"
|
||||
)
|
||||
self.guardrail_name: Final = guardrail_name
|
||||
|
||||
|
||||
def _tool_call_shape(tool_call: object) -> tuple[object, object]:
|
||||
plain: Final = tool_call.model_dump() if isinstance(tool_call, BaseModel) else tool_call
|
||||
function: Final = plain.get("function") if isinstance(plain, Mapping) else None
|
||||
|
|
|
|||
|
|
@ -250,7 +250,7 @@ import litellm._redis
|
|||
from litellm import Router
|
||||
from litellm._logging import _redact_string, verbose_proxy_logger, verbose_router_logger
|
||||
from litellm.caching.caching import DualCache, RedisCache
|
||||
from litellm.caching.redis_cache import RedisCircuitBreakerOpenError
|
||||
from litellm.caching.redis_cache import RedisCircuitBreakerOpenError, is_redis_timeout_failure
|
||||
from litellm.caching.redis_cluster_cache import RedisClusterCache
|
||||
from litellm.constants import (
|
||||
_REALTIME_BODY_CACHE_SIZE,
|
||||
|
|
@ -274,6 +274,8 @@ from litellm.constants import (
|
|||
PROXY_BUDGET_RESCHEDULER_MAX_TIME,
|
||||
PROXY_BUDGET_RESCHEDULER_MIN_TIME,
|
||||
PROXY_CONFIG_RELOAD_INTERVAL_SECONDS,
|
||||
REALTIME_SESSION_FAILURE_LOGGED_KEY,
|
||||
REALTIME_SESSION_SUCCESS_LOGGED_KEY,
|
||||
ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG,
|
||||
USER_SPEND_ALERTS_JOB_ID,
|
||||
WEEKLY_SPEND_REPORT_JOB_ID,
|
||||
|
|
@ -365,6 +367,7 @@ from litellm.proxy.common_utils.healthy_model_filter import (
|
|||
get_hidden_unhealthy_model_names,
|
||||
is_healthy_only_listing_default,
|
||||
)
|
||||
from litellm.proxy.common_utils.html_forms.default_credentials_hint import should_hide_default_credentials_hint
|
||||
from litellm.proxy.common_utils.html_forms.ui_login import build_ui_login_form
|
||||
from litellm.proxy.common_utils.http_parsing_utils import (
|
||||
_read_request_body,
|
||||
|
|
@ -3450,8 +3453,10 @@ async def _invalidate_spend_counter(counter_key: str):
|
|||
async def _apply_spend_counter_increments(pending: Sequence[PendingSpendIncrement]) -> None:
|
||||
try:
|
||||
await increment_spend_counters_pipeline(pending=pending)
|
||||
except RedisCircuitBreakerOpenError:
|
||||
return
|
||||
except Exception as e:
|
||||
if isinstance(e, RedisCircuitBreakerOpenError) or is_redis_timeout_failure(e):
|
||||
return
|
||||
raise
|
||||
|
||||
|
||||
async def increment_spend_counters_pipeline(pending: Sequence[PendingSpendIncrement]) -> None:
|
||||
|
|
@ -7065,6 +7070,9 @@ class ProxyConfig:
|
|||
if "max_file_size_mb" not in self._yaml_general_settings_keys:
|
||||
general_settings["max_file_size_mb"] = _general_settings.get("max_file_size_mb")
|
||||
|
||||
if "allowed_file_extensions" not in self._yaml_general_settings_keys:
|
||||
general_settings["allowed_file_extensions"] = _general_settings.get("allowed_file_extensions")
|
||||
|
||||
if "blocked_file_extensions" not in self._yaml_general_settings_keys:
|
||||
general_settings["blocked_file_extensions"] = _general_settings.get("blocked_file_extensions")
|
||||
|
||||
|
|
@ -11893,6 +11901,13 @@ async def _release_realtime_budget_reservation(user_api_key_dict: UserAPIKeyAuth
|
|||
)
|
||||
|
||||
|
||||
async def _release_realtime_max_parallel_slot(user_api_key_dict: UserAPIKeyAuth) -> None:
|
||||
release_like_http_disconnect: Final = (
|
||||
proxy_logging_obj._arelease_max_parallel_requests_on_disconnect # pyright: ignore[reportPrivateUsage] # shared
|
||||
)
|
||||
await release_like_http_disconnect(user_api_key_dict)
|
||||
|
||||
|
||||
async def _reject_realtime_session(
|
||||
websocket: WebSocket,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
|
|
@ -11912,6 +11927,7 @@ async def _reject_realtime_session(
|
|||
await websocket.close(code=code, reason=reason)
|
||||
finally:
|
||||
await _release_realtime_budget_reservation(user_api_key_dict)
|
||||
await _release_realtime_max_parallel_slot(user_api_key_dict)
|
||||
|
||||
|
||||
@app.websocket("/openai/v1/realtime")
|
||||
|
|
@ -12015,6 +12031,9 @@ async def realtime_websocket_endpoint(
|
|||
websocket, user_api_key_dict, code=1011, reason="Pre-call error", error_message=str(e)
|
||||
)
|
||||
return
|
||||
except BaseException:
|
||||
await _release_realtime_max_parallel_slot(user_api_key_dict)
|
||||
raise
|
||||
|
||||
# Phase 2: route to upstream LLM.
|
||||
try:
|
||||
|
|
@ -12044,12 +12063,10 @@ async def realtime_websocket_endpoint(
|
|||
except Exception: # noqa: BLE001 # the lower layer may have closed the socket already; closing twice is not an error
|
||||
verbose_proxy_logger.debug("Could not close realtime client websocket; it is already gone")
|
||||
finally:
|
||||
from litellm.litellm_core_utils.realtime_streaming import (
|
||||
REALTIME_SESSION_SUCCESS_LOGGED_KEY,
|
||||
)
|
||||
|
||||
if not litellm_logging_obj.model_call_details.get(REALTIME_SESSION_SUCCESS_LOGGED_KEY):
|
||||
await _release_realtime_budget_reservation(user_api_key_dict)
|
||||
if not litellm_logging_obj.model_call_details.get(REALTIME_SESSION_FAILURE_LOGGED_KEY):
|
||||
await _release_realtime_max_parallel_slot(user_api_key_dict)
|
||||
|
||||
|
||||
######################################################################
|
||||
|
|
@ -15033,6 +15050,13 @@ def _get_proxy_model_info(model: dict) -> dict:
|
|||
return _translate_model_name_for_response(model)
|
||||
|
||||
|
||||
def _model_info_json_response(data: Sequence[Mapping[str, object]] | Mapping[str, object]) -> Response:
|
||||
return Response(
|
||||
content=orjson.dumps({"data": data}, default=jsonable_encoder, option=orjson.OPT_NON_STR_KEYS),
|
||||
media_type="application/json",
|
||||
)
|
||||
|
||||
|
||||
@router.get(
|
||||
"/model/info",
|
||||
tags=["model management"],
|
||||
|
|
@ -15080,7 +15104,7 @@ async def model_info_v1(
|
|||
`model_info.direct_access` when the proxy database is connected.
|
||||
|
||||
Returns:
|
||||
Returns a dictionary containing information about each model.
|
||||
A JSON response whose `data` list holds one entry per model.
|
||||
|
||||
Example Response:
|
||||
```json
|
||||
|
|
@ -15128,7 +15152,7 @@ async def model_info_v1(
|
|||
deployment_dict=_deployment_info_dict,
|
||||
excluded_keys={"litellm_credential_name"},
|
||||
)
|
||||
return {"data": _deployment_info_dict}
|
||||
return _model_info_json_response(_deployment_info_dict)
|
||||
|
||||
if llm_model_list is None:
|
||||
raise HTTPException(
|
||||
|
|
@ -15179,7 +15203,7 @@ async def model_info_v1(
|
|||
llm_router=llm_router,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
return {"data": single_model_list}
|
||||
return _model_info_json_response(single_model_list)
|
||||
|
||||
# Return router deployments (same source as /v2/model/info), not wildcard-
|
||||
# expanded model names from get_complete_model_list(). Team-scoped rows
|
||||
|
|
@ -15247,7 +15271,7 @@ async def model_info_v1(
|
|||
visible_models: Final = [model for model in all_models if model.get("model_name") not in hidden_names]
|
||||
|
||||
verbose_proxy_logger.debug("all_models: %s", visible_models)
|
||||
return {"data": visible_models}
|
||||
return _model_info_json_response(visible_models)
|
||||
|
||||
|
||||
@router.get(
|
||||
|
|
@ -15816,10 +15840,7 @@ async def fallback_login(request: Request):
|
|||
|
||||
from fastapi.responses import HTMLResponse
|
||||
|
||||
hide_default_credentials_hint: Final = (
|
||||
os.getenv("LITELLM_HIDE_DEFAULT_CREDENTIALS_HINT", "false").lower() == "true"
|
||||
or general_settings.get("hide_default_credentials_hint", False) is True
|
||||
)
|
||||
hide_default_credentials_hint: Final = should_hide_default_credentials_hint(general_settings)
|
||||
return HTMLResponse(
|
||||
content=build_ui_login_form(
|
||||
show_deprecation_banner=False,
|
||||
|
|
@ -17036,6 +17057,7 @@ _GENERAL_SETTINGS_CONFIG_LIST_FIELD_TYPES: Final[Mapping[str, str]] = MappingPro
|
|||
"max_request_size_mb": "Integer",
|
||||
"max_batch_file_size_mb": "Integer",
|
||||
"max_file_size_mb": "Integer",
|
||||
"allowed_file_extensions": "List",
|
||||
"blocked_file_extensions": "List",
|
||||
"max_response_size_mb": "Integer",
|
||||
"proxy_config_reload_interval_seconds": "Integer",
|
||||
|
|
|
|||
|
|
@ -159,7 +159,7 @@ def _get_spend_logs_metadata(
|
|||
requester_ip_address=None,
|
||||
additional_usage_values=None,
|
||||
applied_guardrails=None,
|
||||
status=None or "success",
|
||||
status="success",
|
||||
error_information=None,
|
||||
proxy_server_request=None,
|
||||
batch_models=None,
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ import copy
|
|||
import hashlib
|
||||
import inspect
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import smtplib
|
||||
import ssl
|
||||
|
|
@ -6404,7 +6405,7 @@ class PrismaClient:
|
|||
return None
|
||||
try:
|
||||
value: Final = float(response_time_ms)
|
||||
return value if value == value and value not in (float("inf"), float("-inf")) else None
|
||||
return value if math.isfinite(value) else None
|
||||
except (ValueError, TypeError):
|
||||
verbose_proxy_logger.warning("Invalid response_time_ms value: %s", response_time_ms)
|
||||
return None
|
||||
|
|
|
|||
|
|
@ -221,6 +221,13 @@ def _status_code_for_error_fields(error_type: str | None, error_code: str | None
|
|||
)
|
||||
|
||||
|
||||
def _mid_stream_fallback_eligible(mapped_exception: Exception) -> bool:
|
||||
if isinstance(mapped_exception, litellm.ContentPolicyViolationError):
|
||||
return True
|
||||
status_code: Final = getattr(mapped_exception, "status_code", None)
|
||||
return not isinstance(status_code, int) or status_code >= 500 or status_code == 429
|
||||
|
||||
|
||||
class BaseResponsesAPIStreamingIterator:
|
||||
"""
|
||||
Base class for streaming iterators that process responses from the Responses API.
|
||||
|
|
@ -521,15 +528,8 @@ class BaseResponsesAPIStreamingIterator:
|
|||
getattr(self.completed_response, "response", None) if self.completed_response else None
|
||||
)
|
||||
error_info: Final = getattr(response_obj, "error", None) if response_obj else None
|
||||
error_message, error_type, error_code = _error_event_fields(error_info)
|
||||
self._record_failed_response_usage(response_obj)
|
||||
exception: Final = litellm.APIError(
|
||||
status_code=_status_code_for_error_fields(error_type, error_code),
|
||||
message=error_message,
|
||||
llm_provider=self.custom_llm_provider or "",
|
||||
model=self.model or "",
|
||||
)
|
||||
self._handle_failure(exception)
|
||||
self._handle_failure(self._map_error_event_exception(error_info))
|
||||
|
||||
def _record_failed_response_usage(self, response_obj: ResponsesAPIResponse | None) -> None:
|
||||
if response_obj is None or self.logging_obj is None:
|
||||
|
|
@ -551,6 +551,28 @@ class BaseResponsesAPIStreamingIterator:
|
|||
self.logging_obj._response_cost_calculator(result=response_obj) or 0.0
|
||||
)
|
||||
|
||||
def _map_error_event_exception(self, error_obj: object) -> Exception:
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
|
||||
error_message, error_type, error_code = _error_event_fields(error_obj)
|
||||
status_code: Final = _status_code_for_error_fields(error_type, error_code)
|
||||
error_body: Final = {"message": error_message, "type": error_type, "code": error_code}
|
||||
provider_exception: Final = BaseLLMException(
|
||||
status_code=status_code,
|
||||
message=f"Error code: {status_code} - {{'error': {error_body}}}",
|
||||
body=error_body,
|
||||
)
|
||||
try:
|
||||
return litellm.exception_type(
|
||||
model=self.model or "",
|
||||
custom_llm_provider=self.custom_llm_provider or "",
|
||||
original_exception=provider_exception,
|
||||
completion_kwargs={},
|
||||
extra_kwargs={},
|
||||
)
|
||||
except Exception as mapped_exception:
|
||||
return mapped_exception
|
||||
|
||||
def _maybe_raise_for_error_event(self, result: object) -> None:
|
||||
chunk_type: Final = getattr(result, "type", None)
|
||||
if chunk_type not in ("error", "response.failed"):
|
||||
|
|
@ -562,15 +584,8 @@ class BaseResponsesAPIStreamingIterator:
|
|||
else getattr(result, "error", None)
|
||||
)
|
||||
|
||||
error_message, error_type, error_code = _error_event_fields(error_obj)
|
||||
status_code: Final = _status_code_for_error_fields(error_type, error_code)
|
||||
mapped_exception: Final = litellm.APIError(
|
||||
status_code=status_code,
|
||||
message=error_message,
|
||||
llm_provider=self.custom_llm_provider or "",
|
||||
model=self.model or "",
|
||||
)
|
||||
if 400 <= status_code < 500 and status_code != 429:
|
||||
mapped_exception: Final = self._map_error_event_exception(error_obj)
|
||||
if not _mid_stream_fallback_eligible(mapped_exception):
|
||||
raise mapped_exception
|
||||
raise MidStreamFallbackError(
|
||||
message=str(mapped_exception),
|
||||
|
|
|
|||
|
|
@ -96,7 +96,6 @@ from litellm.litellm_core_utils.request_timeout_resolver import (
|
|||
from litellm.litellm_core_utils.secret_redaction import redact_string
|
||||
from litellm.litellm_core_utils.sensitive_data_masker import (
|
||||
SensitiveDataMasker,
|
||||
mask_credentials_in_payload,
|
||||
mask_sensitive_structure,
|
||||
)
|
||||
from litellm.litellm_core_utils.token_counter import offload_token_count
|
||||
|
|
@ -623,20 +622,6 @@ def _replay_live_router_model_cost() -> None:
|
|||
set_live_deployment_replay(_replay_live_router_model_cost)
|
||||
|
||||
|
||||
# Kwargs that carry no signal about the failed attempt, so log_retry drops them from a
|
||||
# breadcrumb entirely: the request payload, the proxy's snapshot of the inbound request (its body
|
||||
# aliases the live request metadata, earlier breadcrumbs included, so copying it would nest every
|
||||
# breadcrumb inside the next one), and the router-internal walk state. Credentials are handled
|
||||
# separately by mask_credentials_in_payload, which scrubs credential-named values from whatever
|
||||
# kwargs remain rather than trying to enumerate every credential-bearing key here.
|
||||
RETRY_BREADCRUMB_EXCLUDED_KWARGS: Final = frozenset(
|
||||
(
|
||||
"messages",
|
||||
"original_function",
|
||||
"attempted_targets",
|
||||
"proxy_server_request",
|
||||
)
|
||||
)
|
||||
RETRY_BREADCRUMB_LIMIT: Final = 4
|
||||
|
||||
|
||||
|
|
@ -1553,6 +1538,18 @@ class Router:
|
|||
return False
|
||||
return sum(len(self.model_name_to_deployment_indices.get(member) or ()) for member in group.models) > 1
|
||||
|
||||
def team_model_has_alternatives(self, deployment_id: str) -> bool:
|
||||
deployment: Final = self.get_deployment(model_id=deployment_id)
|
||||
if deployment is None:
|
||||
return False
|
||||
team_id: Final = deployment.model_info.team_id
|
||||
public_model_name: Final = deployment.model_info.team_public_model_name
|
||||
if team_id is None or public_model_name is None:
|
||||
return False
|
||||
sibling_indices: Final = self.team_model_to_deployment_indices.get((team_id, public_model_name)) or ()
|
||||
routable_siblings: Final = self._filter_blocked_deployments([self.model_list[idx] for idx in sibling_indices])
|
||||
return len(routable_siblings) > 1
|
||||
|
||||
_OVERRIDABLE_ROUTING_STRATEGIES: frozenset[str] = frozenset({"simple-shuffle", *_DEFAULT_SELECTOR_ATTR_BY_STRATEGY})
|
||||
|
||||
def _get_request_routing_strategy_override(self, request_kwargs: dict | None) -> str | None:
|
||||
|
|
@ -3271,8 +3268,15 @@ class Router:
|
|||
kwargs=initial_kwargs,
|
||||
metadata_variable_name="litellm_metadata",
|
||||
)
|
||||
# The content-policy dispatch branch matches on the trigger's own type, so a refusal's
|
||||
# MidStreamFallbackError envelope is unwrapped here or the wrong fallback list is consulted.
|
||||
fallback_trigger: Final[Exception] = (
|
||||
e.original_exception
|
||||
if isinstance(e.original_exception, litellm.ContentPolicyViolationError)
|
||||
else e
|
||||
)
|
||||
fallback_response = await self.async_function_with_fallbacks_common_utils(
|
||||
e=e,
|
||||
e=fallback_trigger,
|
||||
disable_fallbacks=False,
|
||||
fallbacks=fallbacks,
|
||||
context_window_fallbacks=context_window_fallbacks,
|
||||
|
|
@ -4108,16 +4112,16 @@ class Router:
|
|||
models: Final = [m.strip() for m in model.split(",")]
|
||||
|
||||
async def _async_completion_no_exceptions(
|
||||
model: str, messages: list[dict[str, str]], stream: bool, **kwargs: Any
|
||||
model_name: str, messages: list[dict[str, str]], stream: bool, **kwargs: Any
|
||||
) -> ModelResponse | CustomStreamWrapper | Exception:
|
||||
"""
|
||||
Wrapper around self.acompletion that catches exceptions and returns them as a result
|
||||
"""
|
||||
try:
|
||||
result = await self.acompletion(model=model, messages=messages, stream=stream, **kwargs)
|
||||
result = await self.acompletion(model=model_name, messages=messages, stream=stream, **kwargs)
|
||||
return result
|
||||
except asyncio.CancelledError:
|
||||
verbose_router_logger.debug("Received 'task.cancel'. Cancelling call w/ model=%s.", model)
|
||||
verbose_router_logger.debug("Received 'task.cancel'. Cancelling call w/ model=%s.", model_name)
|
||||
raise
|
||||
except Exception as e:
|
||||
return e
|
||||
|
|
@ -4144,9 +4148,9 @@ class Router:
|
|||
except KeyError:
|
||||
pass
|
||||
|
||||
for model in models:
|
||||
for model_name in models:
|
||||
task = asyncio.create_task(
|
||||
_async_completion_no_exceptions(model=model, messages=messages, stream=stream, **kwargs)
|
||||
_async_completion_no_exceptions(model_name=model_name, messages=messages, stream=stream, **kwargs)
|
||||
)
|
||||
pending_tasks.append(task)
|
||||
|
||||
|
|
@ -8374,31 +8378,30 @@ class Router:
|
|||
|
||||
def log_retry(self, kwargs: dict, e: Exception) -> dict:
|
||||
"""
|
||||
When a retry or fallback happens, log the details of the just failed model call - similar to Sentry breadcrumbing
|
||||
When a retry or fallback happens, record which model group, deployment and attempt just failed and why
|
||||
"""
|
||||
from litellm.types.router import RetryAttemptRecord
|
||||
|
||||
_metadata_var: Final = "litellm_metadata" if "litellm_metadata" in kwargs else "metadata"
|
||||
request_metadata: Final[Mapping[str, object]] = kwargs[_metadata_var]
|
||||
attempt_kwargs: Final = MappingProxyType(
|
||||
{k: v for k, v in kwargs.items() if k != _metadata_var and k not in RETRY_BREADCRUMB_EXCLUDED_KWARGS}
|
||||
)
|
||||
attempt_metadata: Final = MappingProxyType(
|
||||
{k: v for k, v in request_metadata.items() if k != "previous_models"}
|
||||
)
|
||||
previous_model: Final = MappingProxyType(
|
||||
{
|
||||
"exception_type": type(e).__name__,
|
||||
"exception_string": str(e),
|
||||
**attempt_kwargs,
|
||||
_metadata_var: attempt_metadata,
|
||||
}
|
||||
)
|
||||
model_group: Final = kwargs.get("model")
|
||||
model_info: Final = request_metadata.get("model_info")
|
||||
deployment_id: Final = model_info.get("id") if isinstance(model_info, Mapping) else None
|
||||
attempted_retries: Final = request_metadata.get("attempted_retries")
|
||||
attempt_record: Final[RetryAttemptRecord] = {
|
||||
"model_group": model_group if isinstance(model_group, str) else None,
|
||||
"deployment_id": deployment_id if isinstance(deployment_id, str) else None,
|
||||
"exception_type": type(e).__name__,
|
||||
"exception_string": str(e),
|
||||
"attempted_retries": attempted_retries if type(attempted_retries) is int else None,
|
||||
}
|
||||
earlier_breadcrumbs: Final = request_metadata.get("previous_models")
|
||||
kept_breadcrumbs: Final[tuple[object, ...]] = (
|
||||
tuple(earlier_breadcrumbs)[-(RETRY_BREADCRUMB_LIMIT - 1) :]
|
||||
if isinstance(earlier_breadcrumbs, (list, tuple))
|
||||
else ()
|
||||
)
|
||||
breadcrumbs: Final = (*kept_breadcrumbs, mask_credentials_in_payload(previous_model))
|
||||
breadcrumbs: Final = (*kept_breadcrumbs, attempt_record)
|
||||
kwargs[_metadata_var]["previous_models"] = breadcrumbs # rebind-ok: the logging object already holds this dict
|
||||
return kwargs
|
||||
|
||||
|
|
@ -13878,6 +13881,7 @@ class Router:
|
|||
cooldown_time=_cooldown_time,
|
||||
enable_pre_call_checks=self.enable_pre_call_checks,
|
||||
cooldown_list=_cooldown_list,
|
||||
model_ids=model_ids,
|
||||
)
|
||||
|
||||
if strategy == "simple-shuffle":
|
||||
|
|
@ -13910,6 +13914,7 @@ class Router:
|
|||
cooldown_time=_cooldown_time,
|
||||
enable_pre_call_checks=self.enable_pre_call_checks,
|
||||
cooldown_list=_cooldown_list,
|
||||
model_ids=model_ids,
|
||||
)
|
||||
self._override_selector_pre_call_check(strategy, strategy_selector, deployment)
|
||||
verbose_router_logger.info(
|
||||
|
|
@ -13987,6 +13992,11 @@ class Router:
|
|||
model=model,
|
||||
llm_provider="",
|
||||
)
|
||||
pass_through_model_ids: Final = tuple(
|
||||
deployment["model_info"]["id"]
|
||||
for deployment in pass_through_deployments
|
||||
if "id" in deployment.get("model_info", {})
|
||||
)
|
||||
|
||||
# 4. Apply health-check and cooldown filtering
|
||||
parent_otel_span: Final[Span | None] = _get_parent_otel_span_from_kwargs(request_kwargs)
|
||||
|
|
@ -14024,6 +14034,7 @@ class Router:
|
|||
cooldown_time=_cooldown_time,
|
||||
enable_pre_call_checks=self.enable_pre_call_checks,
|
||||
cooldown_list=_cooldown_list,
|
||||
model_ids=pass_through_model_ids,
|
||||
)
|
||||
|
||||
# 6. Apply load balancing strategy
|
||||
|
|
@ -14057,6 +14068,7 @@ class Router:
|
|||
cooldown_time=_cooldown_time,
|
||||
enable_pre_call_checks=self.enable_pre_call_checks,
|
||||
cooldown_list=_cooldown_list,
|
||||
model_ids=model_ids,
|
||||
)
|
||||
self._override_selector_pre_call_check(strategy, strategy_selector, deployment)
|
||||
|
||||
|
|
|
|||
|
|
@ -343,8 +343,9 @@ def _should_cooldown_deployment(
|
|||
model_group: Final = litellm_router_instance.get_model_group(id=deployment)
|
||||
is_single_deployment_model_group = False
|
||||
if model_group is not None and len(model_group) == 1:
|
||||
is_single_deployment_model_group = not litellm_router_instance.routing_group_has_alternatives(
|
||||
requested_model_group
|
||||
is_single_deployment_model_group = not (
|
||||
litellm_router_instance.routing_group_has_alternatives(requested_model_group)
|
||||
or litellm_router_instance.team_model_has_alternatives(deployment)
|
||||
)
|
||||
|
||||
## CHECK DEPLOYMENT-LEVEL POLICY FIRST (overrides router-level)
|
||||
|
|
|
|||
|
|
@ -93,4 +93,5 @@ async def async_raise_no_deployment_exception(
|
|||
cooldown_time=_cooldown_time,
|
||||
enable_pre_call_checks=litellm_router_instance.enable_pre_call_checks,
|
||||
cooldown_list=cooldown_list_ids,
|
||||
model_ids=model_ids,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -99,23 +99,13 @@ def setup(
|
|||
|
||||
def check_limits(kwargs: Mapping[str, object]) -> None:
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.core_helpers import max_retries_per_request_hit
|
||||
|
||||
current_cost: Final = litellm._current_cost # pyright: ignore[reportPrivateUsage] # shared SDK budget counter has no public accessor
|
||||
if litellm.max_budget and current_cost > litellm.max_budget:
|
||||
raise litellm.BudgetExceededError(current_cost=current_cost, max_budget=litellm.max_budget)
|
||||
metadata: Final = kwargs.get("metadata")
|
||||
if isinstance(metadata, Mapping):
|
||||
typed_metadata: Final = cast( # cast-ok: runtime Mapping check establishes read-only metadata
|
||||
Mapping[str, object], metadata
|
||||
)
|
||||
previous: Final = typed_metadata.get("previous_models")
|
||||
if (
|
||||
isinstance(previous, list)
|
||||
and litellm.num_retries_per_request is not None
|
||||
and len(cast(list[object], previous)) # cast-ok: runtime list check establishes the retry history
|
||||
>= litellm.num_retries_per_request
|
||||
):
|
||||
raise RuntimeError("Max retries per request hit!")
|
||||
if max_retries_per_request_hit(kwargs, litellm.num_retries_per_request):
|
||||
raise RuntimeError("Max retries per request hit!")
|
||||
|
||||
|
||||
def finalize(
|
||||
|
|
|
|||
|
|
@ -552,6 +552,16 @@ class BedrockGuardrailConfigModel(BaseModel):
|
|||
"still rejects is bisected automatically, so this value only trades round trips against "
|
||||
"batch size and cannot fail a request on its own.",
|
||||
)
|
||||
contextual_grounding_from_messages: bool = Field(
|
||||
default=False,
|
||||
description="ApplyGuardrail: when True, post-call scans of a request with no grounding_source / "
|
||||
"query content parts send the system and developer messages as the grounding source and "
|
||||
"the latest user message as the query, so the guardrail's contextual grounding policy can "
|
||||
"score the response. Bedrock bills contextual grounding units for these scans and rejects "
|
||||
"queries, sources and responses over its contextual grounding length limits, so leave this "
|
||||
"off for guardrails without a contextual grounding policy. Default False: plain messages "
|
||||
"are never sent as grounding context.",
|
||||
)
|
||||
|
||||
|
||||
class BedrockGuardrailStreamingParams(BaseModel):
|
||||
|
|
|
|||
|
|
@ -0,0 +1,67 @@
|
|||
from collections.abc import Mapping
|
||||
from typing import Annotated, Literal, TypeAlias
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, JsonValue, StrictInt
|
||||
|
||||
TokenCount: TypeAlias = Annotated[StrictInt, Field(ge=0)]
|
||||
|
||||
|
||||
class CacheTokenBuckets(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
uncached_input_tokens: TokenCount = 0
|
||||
cache_read_input_tokens: TokenCount = 0
|
||||
cache_creation_5m_input_tokens: TokenCount = 0
|
||||
cache_creation_1h_input_tokens: TokenCount = 0
|
||||
|
||||
@property
|
||||
def total_tokens(self) -> int:
|
||||
return (
|
||||
self.uncached_input_tokens
|
||||
+ self.cache_read_input_tokens
|
||||
+ self.cache_creation_5m_input_tokens
|
||||
+ self.cache_creation_1h_input_tokens
|
||||
)
|
||||
|
||||
|
||||
class CacheEvidence(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
observed_at: float
|
||||
expires_at: float
|
||||
source: Literal["provider_usage"] = "provider_usage"
|
||||
confidence: Literal["observed"] = "observed"
|
||||
|
||||
|
||||
class CacheCostScenario(BaseModel):
|
||||
tokens: CacheTokenBuckets
|
||||
input_cost: float
|
||||
|
||||
|
||||
class CachePredictionArm(BaseModel):
|
||||
deployment_id: str
|
||||
model: str | None = None
|
||||
cache_state: Literal["warm", "partial", "stale", "unknown", "disabled"] = "unknown"
|
||||
reason: str | None = None
|
||||
estimate: CacheCostScenario | None = None
|
||||
cold: CacheCostScenario | None = None
|
||||
warm: CacheCostScenario | None = None
|
||||
evidence: CacheEvidence | None = None
|
||||
token_count_source: Literal["anthropic_count_tokens"] | None = None
|
||||
|
||||
|
||||
class CachePredictionRequest(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
current_deployment_id: str = Field(min_length=1, max_length=256)
|
||||
candidate_deployment_id: str = Field(min_length=1, max_length=256)
|
||||
request: Mapping[str, JsonValue]
|
||||
|
||||
|
||||
class CachePredictionResponse(BaseModel):
|
||||
stay: CachePredictionArm
|
||||
switch: CachePredictionArm
|
||||
switch_delta: float | None
|
||||
cache_rebuild_penalty: float | None
|
||||
pricing_basis: Literal["input_before_discounts_and_margins"] = "input_before_discounts_and_margins"
|
||||
cache_guarantee: Literal[False] = False
|
||||
8
litellm/types/proxy/auth/auth_checks.py
Normal file
8
litellm/types/proxy/auth/auth_checks.py
Normal file
|
|
@ -0,0 +1,8 @@
|
|||
"""Failure values raised by `litellm/proxy/auth/auth_checks.py`. Kept free of `litellm` imports so any proxy module can import them without joining the `litellm.proxy` import cycle."""
|
||||
|
||||
|
||||
class UserNotFoundError(ValueError):
|
||||
"""The user row is provably absent, as opposed to merely unreadable, so a caller that reads a missing row as no user-level limits can key on it without also swallowing a database that would not answer."""
|
||||
|
||||
def __init__(self, user_id: str) -> None:
|
||||
super().__init__(f"User doesn't exist in db. 'user_id'={user_id}. Create user via `/user/new` call.")
|
||||
|
|
@ -1,4 +1,5 @@
|
|||
from typing import Any, Final, Literal
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Any, Final, Literal, cast # noqa: TID251 # JSON chat rows have no typed constructor across roles
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
from typing_extensions import TypedDict
|
||||
|
|
@ -158,12 +159,21 @@ def coerce_stream_holdback_value(value: Any) -> int:
|
|||
return 0
|
||||
|
||||
|
||||
def structured_messages_from_response(value: object) -> Sequence[AllMessageValues] | None:
|
||||
if not isinstance(value, list):
|
||||
return None
|
||||
if not all(isinstance(message, Mapping) and isinstance(message.get("role"), str) for message in value):
|
||||
return None
|
||||
return cast("Sequence[AllMessageValues]", value) # cast-ok: JSON rows checked for a role, the same trust texts get
|
||||
|
||||
|
||||
class GenericGuardrailAPIResponse:
|
||||
"""Response model for the Generic Guardrail API"""
|
||||
|
||||
texts: list[str] | None
|
||||
images: list[str] | None
|
||||
tools: list[GuardrailToolParam] | None
|
||||
structured_messages: Sequence[AllMessageValues] | None
|
||||
action: str
|
||||
blocked_reason: str | None
|
||||
stream_holdback_chars: list[int] | None
|
||||
|
|
@ -176,12 +186,14 @@ class GenericGuardrailAPIResponse:
|
|||
images: list[str] | None = None,
|
||||
tools: list[GuardrailToolParam] | None = None,
|
||||
stream_holdback_chars: list[int] | None = None,
|
||||
structured_messages: Sequence[AllMessageValues] | None = None,
|
||||
) -> None:
|
||||
self.action = action
|
||||
self.blocked_reason = blocked_reason
|
||||
self.texts = texts
|
||||
self.images = images
|
||||
self.tools = tools
|
||||
self.structured_messages = structured_messages
|
||||
# Number of trailing chars, indexed the same as ``texts``, that the
|
||||
# framework must withhold from streaming emission until the next
|
||||
# processing round (word-boundary safety for text transformations).
|
||||
|
|
@ -200,4 +212,5 @@ class GenericGuardrailAPIResponse:
|
|||
images=data.get("images"),
|
||||
tools=data.get("tools"),
|
||||
stream_holdback_chars=stream_holdback_chars,
|
||||
structured_messages=structured_messages_from_response(data.get("structured_messages")),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -645,6 +645,7 @@ class RouterErrors(enum.Enum):
|
|||
|
||||
user_defined_ratelimit_error = "Deployment over user-defined ratelimit."
|
||||
no_deployments_available = "No deployments available for selected model"
|
||||
all_deployments_in_cooldown = "All deployments for selected model are in cooldown"
|
||||
no_deployments_with_tag_routing = "Not allowed to access model due to tags configuration"
|
||||
no_deployments_with_provider_budget_routing = "No deployments available - crossed budget"
|
||||
no_healthy_deployments = "There are no healthy deployments for this model"
|
||||
|
|
@ -868,6 +869,11 @@ class RouterRateLimitErrorBasic(ValueError):
|
|||
super().__init__(_message)
|
||||
|
||||
|
||||
class RouterErrorTypes(str, enum.Enum):
|
||||
rate_limit_error = "rate_limit_error"
|
||||
all_deployments_in_cooldown = "all_deployments_in_cooldown"
|
||||
|
||||
|
||||
class RouterRateLimitError(ValueError):
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -875,12 +881,25 @@ class RouterRateLimitError(ValueError):
|
|||
cooldown_time: float,
|
||||
enable_pre_call_checks: bool,
|
||||
cooldown_list: list,
|
||||
model_ids: Sequence[str] = (),
|
||||
) -> None:
|
||||
self.model = model
|
||||
self.cooldown_time = cooldown_time
|
||||
self.enable_pre_call_checks = enable_pre_call_checks
|
||||
self.cooldown_list = cooldown_list
|
||||
_message = f"{RouterErrors.no_deployments_available.value}, Try again in {cooldown_time} seconds. Passed model={model}. pre-call-checks={enable_pre_call_checks}, cooldown_list={cooldown_list}"
|
||||
self.all_deployments_in_cooldown = bool(model_ids) and frozenset(model_ids) <= frozenset(cooldown_list)
|
||||
self.type = (
|
||||
RouterErrorTypes.all_deployments_in_cooldown.value
|
||||
if self.all_deployments_in_cooldown
|
||||
else RouterErrorTypes.rate_limit_error.value
|
||||
)
|
||||
_reason: Final = (
|
||||
f" {RouterErrors.all_deployments_in_cooldown.value}." if self.all_deployments_in_cooldown else ""
|
||||
)
|
||||
_message: Final = (
|
||||
f"{RouterErrors.no_deployments_available.value}, Try again in {cooldown_time} seconds.{_reason} "
|
||||
f"Passed model={model}. pre-call-checks={enable_pre_call_checks}, cooldown_list={cooldown_list}"
|
||||
)
|
||||
super().__init__(_message)
|
||||
|
||||
|
||||
|
|
@ -889,6 +908,14 @@ class RouterModelGroupAliasItem(TypedDict):
|
|||
hidden: bool # if 'True', don't return on `.get_model_list`
|
||||
|
||||
|
||||
class RetryAttemptRecord(TypedDict):
|
||||
model_group: ReadOnly[str | None]
|
||||
deployment_id: ReadOnly[str | None]
|
||||
exception_type: ReadOnly[str]
|
||||
exception_string: ReadOnly[str]
|
||||
attempted_retries: ReadOnly[int | None]
|
||||
|
||||
|
||||
VALID_LITELLM_ENVIRONMENTS = [
|
||||
"development",
|
||||
"staging",
|
||||
|
|
|
|||
|
|
@ -283,8 +283,11 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
|
|||
input_cost_per_video_token: float | None # for gemini omni models with video input
|
||||
input_cost_per_audio_per_second: float | None # only for vertex ai models
|
||||
input_cost_per_video_per_second: float | None # only for vertex ai models
|
||||
input_cost_per_audio_token_batches: ReadOnly[float | None]
|
||||
input_cost_per_image_token_batches: ReadOnly[float | None]
|
||||
input_cost_per_second: float | None # for OpenAI Speech models
|
||||
input_cost_per_token_batches: float | None
|
||||
input_cost_per_video_token_batches: ReadOnly[float | None]
|
||||
output_cost_per_token_batches: float | None
|
||||
output_cost_per_token: Required[float | None]
|
||||
output_cost_per_token_flex: float | None # OpenAI flex service tier pricing
|
||||
|
|
@ -3583,7 +3586,10 @@ class CustomPricingLiteLLMParams(MirroredPricingParams):
|
|||
input_cost_per_video_per_second_above_128k_tokens: float | None = None
|
||||
input_cost_per_video_per_second_above_15s_interval: float | None = None
|
||||
input_cost_per_video_per_second_above_8s_interval: float | None = None
|
||||
input_cost_per_audio_token_batches: float | None = None
|
||||
input_cost_per_image_token_batches: float | None = None
|
||||
input_cost_per_token_batches: float | None = None
|
||||
input_cost_per_video_token_batches: float | None = None
|
||||
output_cost_per_token_batches: float | None = None
|
||||
output_cost_per_token_flex: float | None = None
|
||||
output_cost_per_token_priority: float | None = None
|
||||
|
|
@ -3761,6 +3767,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",
|
||||
|
|
|
|||
|
|
@ -84,6 +84,7 @@ from litellm.constants import (
|
|||
from litellm.litellm_core_utils.core_helpers import normalize_drop_params
|
||||
from litellm.litellm_core_utils.fallback_generalizations import (
|
||||
match_capability_generalizations,
|
||||
match_fill_missing_generalizations,
|
||||
)
|
||||
from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload
|
||||
|
||||
|
|
@ -254,6 +255,7 @@ from litellm.types.utils import (
|
|||
)
|
||||
|
||||
_CALL_TYPE_ENUM_MAP: Final[dict] = {ct.value: ct for ct in CallTypes}
|
||||
_BACKFILL_MODES: Final = frozenset({"chat", "responses"})
|
||||
|
||||
# +-----------------------------------------------+
|
||||
# | |
|
||||
|
|
@ -1260,15 +1262,6 @@ async def _client_async_logging_helper(
|
|||
async_coroutine=logging_obj.async_success_handler(result=result, start_time=start_time, end_time=end_time)
|
||||
)
|
||||
|
||||
################################################
|
||||
# Sync Logging Worker
|
||||
################################################
|
||||
logging_obj.handle_sync_success_callbacks_for_async_calls(
|
||||
result=result,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
|
||||
|
||||
def _get_wrapper_num_retries(kwargs: dict[str, Any], exception: Exception) -> tuple[int | None, dict[str, Any]]:
|
||||
"""
|
||||
|
|
@ -1500,6 +1493,8 @@ def post_call_processing(
|
|||
|
||||
|
||||
def client(original_function):
|
||||
from litellm.litellm_core_utils.core_helpers import max_retries_per_request_hit
|
||||
|
||||
Rules: Final = litellm_utils.Rules
|
||||
rules_obj: Final = Rules()
|
||||
|
||||
|
|
@ -1510,12 +1505,8 @@ def client(original_function):
|
|||
call_type = original_function.__name__
|
||||
if _is_async_request(kwargs):
|
||||
# [OPTIONAL] CHECK MAX RETRIES / REQUEST
|
||||
if litellm.num_retries_per_request is not None:
|
||||
# check if previous_models passed in as ['litellm_params']['metadata]['previous_models']
|
||||
previous_models = (kwargs.get("metadata") or {}).get("previous_models", None)
|
||||
if previous_models is not None:
|
||||
if litellm.num_retries_per_request <= len(previous_models):
|
||||
raise Exception("Max retries per request hit!")
|
||||
if max_retries_per_request_hit(kwargs, litellm.num_retries_per_request):
|
||||
raise Exception("Max retries per request hit!")
|
||||
|
||||
# MODEL CALL
|
||||
result = original_function(*args, **kwargs)
|
||||
|
|
@ -1574,12 +1565,8 @@ def client(original_function):
|
|||
)
|
||||
|
||||
# [OPTIONAL] CHECK MAX RETRIES / REQUEST
|
||||
if litellm.num_retries_per_request is not None:
|
||||
# check if previous_models passed in as ['litellm_params']['metadata]['previous_models']
|
||||
previous_models = (kwargs.get("metadata") or {}).get("previous_models", None)
|
||||
if previous_models is not None:
|
||||
if litellm.num_retries_per_request <= len(previous_models):
|
||||
raise Exception("Max retries per request hit!")
|
||||
if max_retries_per_request_hit(kwargs, litellm.num_retries_per_request):
|
||||
raise Exception("Max retries per request hit!")
|
||||
|
||||
# [OPTIONAL] CHECK CACHE
|
||||
print_verbose(
|
||||
|
|
@ -5819,6 +5806,14 @@ def _get_model_info_helper(
|
|||
):
|
||||
_model_info = None
|
||||
|
||||
if _model_info is not None and key is not None and _model_info.get("mode", "chat") in _BACKFILL_MODES:
|
||||
fill_missing: Final = match_fill_missing_generalizations(key, _model_info.get("litellm_provider", ""))
|
||||
if fill_missing is not None:
|
||||
_model_info = {
|
||||
**{k: v for k, v in fill_missing.items() if k not in _model_info},
|
||||
**_model_info,
|
||||
}
|
||||
|
||||
if _model_info is None:
|
||||
generalization: Final = _get_model_info_from_generalization(
|
||||
model=model,
|
||||
|
|
@ -5928,10 +5923,13 @@ def _get_model_info_helper(
|
|||
input_cost_per_audio_token=_model_info.get("input_cost_per_audio_token", None),
|
||||
input_cost_per_image_token=_model_info.get("input_cost_per_image_token", None),
|
||||
input_cost_per_video_token=_model_info.get("input_cost_per_video_token", None),
|
||||
input_cost_per_audio_token_batches=_model_info.get("input_cost_per_audio_token_batches", None),
|
||||
input_cost_per_image_token_batches=_model_info.get("input_cost_per_image_token_batches", None),
|
||||
input_cost_per_image=_model_info.get("input_cost_per_image", None),
|
||||
input_cost_per_audio_per_second=_model_info.get("input_cost_per_audio_per_second", None),
|
||||
input_cost_per_video_per_second=_model_info.get("input_cost_per_video_per_second", None),
|
||||
input_cost_per_token_batches=_model_info.get("input_cost_per_token_batches"),
|
||||
input_cost_per_video_token_batches=_model_info.get("input_cost_per_video_token_batches", None),
|
||||
output_cost_per_token_batches=_model_info.get("output_cost_per_token_batches"),
|
||||
output_cost_per_token=_output_cost_per_token,
|
||||
output_cost_per_token_flex=_model_info.get("output_cost_per_token_flex", None),
|
||||
|
|
@ -6265,7 +6263,7 @@ def function_to_dict(input_function) -> dict:
|
|||
"enum": param_enum,
|
||||
}
|
||||
|
||||
parameters[param_name] = dict([(k, v) for k, v in param_dict.items() if isinstance(v, str)])
|
||||
parameters[param_name] = {k: v for k, v in param_dict.items() if isinstance(v, str)}
|
||||
|
||||
# Check if the parameter has no default value (i.e., it's required)
|
||||
if param.default == param.empty:
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -249,6 +249,10 @@
|
|||
"type": "number",
|
||||
"minimum": 0
|
||||
},
|
||||
"input_cost_per_audio_token_batches": {
|
||||
"type": "number",
|
||||
"minimum": 0
|
||||
},
|
||||
"input_cost_per_audio_token_priority": {
|
||||
"type": "number",
|
||||
"minimum": 0,
|
||||
|
|
@ -276,6 +280,10 @@
|
|||
"type": "number",
|
||||
"minimum": 0
|
||||
},
|
||||
"input_cost_per_image_token_batches": {
|
||||
"type": "number",
|
||||
"minimum": 0
|
||||
},
|
||||
"input_cost_per_pixel": {
|
||||
"type": "number",
|
||||
"minimum": 0
|
||||
|
|
@ -375,6 +383,14 @@
|
|||
"minimum": 0,
|
||||
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
|
||||
},
|
||||
"input_cost_per_video_token": {
|
||||
"type": "number",
|
||||
"minimum": 0
|
||||
},
|
||||
"input_cost_per_video_token_batches": {
|
||||
"type": "number",
|
||||
"minimum": 0
|
||||
},
|
||||
"input_dbu_cost_per_token": {
|
||||
"type": "number",
|
||||
"minimum": 0
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue