mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-20 00:11:50 +00:00
Merge branch 'main' into litellm_bulk_user_delete
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
commit
6ea1085bc3
267 changed files with 21335 additions and 2482 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$"
|
||||
|
|
|
|||
3
.github/pull_request_template.md
vendored
3
.github/pull_request_template.md
vendored
|
|
@ -101,7 +101,8 @@ If you're seeing a delay in your PR being merged, ping the LiteLLM Team on [Slac
|
|||
For bug fixes: Before shows the reproduction, After shows the same steps passing
|
||||
For new features: Before shows the capability missing, After shows it working end-to-end
|
||||
If the change applies to all three LLM endpoints (/v1/responses, /v1/chat/completions, /v1/messages), make each endpoint its own case, not just one
|
||||
For UI changes: before/after screenshots under the same headings -->
|
||||
For UI changes: before/after screenshots under the same headings
|
||||
If the main use case runs through a coding tool like Claude Code or Codex, drive that tool interactively the way the user does (never `claude -p`, `codex exec`, or curl on its own) and embed before/after screenshots of its pane under the same headings; curl replays and headless runs can follow as extra cases, never as the only proof -->
|
||||
|
||||
## Type
|
||||
|
||||
|
|
|
|||
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[@]}"
|
||||
|
|
|
|||
|
|
@ -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 # cap on Router retries of one model group; resets per fallback hop
|
||||
num_retries_per_request: Optional[int] = None # for the request overall (incl. fallbacks + model retries)
|
||||
####### SECRET MANAGERS #####################
|
||||
secret_manager_client: Optional[Any] = (
|
||||
None # list of instantiated key management clients - e.g. azure kv, infisical, etc.
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
125
litellm/caching/affinity_cache.py
Normal file
125
litellm/caching/affinity_cache.py
Normal file
|
|
@ -0,0 +1,125 @@
|
|||
"""Atomic affinity claims shared by deployment and tier-model selection."""
|
||||
|
||||
import json
|
||||
from collections.abc import Mapping
|
||||
from typing import (
|
||||
Final,
|
||||
cast, # noqa: TID251 # Redis script results are narrowed only to object, then validated
|
||||
)
|
||||
|
||||
from pydantic import JsonValue, TypeAdapter, ValidationError
|
||||
|
||||
from litellm._logging import verbose_router_logger
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
|
||||
_PIN_JSON_ADAPTER: Final = TypeAdapter[JsonValue](JsonValue)
|
||||
|
||||
_CLAIM_PIN_SCRIPT: Final = """
|
||||
local current = redis.call('GET', KEYS[1])
|
||||
if current == false then
|
||||
redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2])
|
||||
return ARGV[1]
|
||||
end
|
||||
if ARGV[3] then
|
||||
local decoded, stored = pcall(cjson.decode, current)
|
||||
if decoded and type(stored) == 'table' then
|
||||
for _, eligible in ipairs(cjson.decode(ARGV[3])) do
|
||||
local matches = true
|
||||
for key, value in pairs(eligible) do
|
||||
if stored[key] ~= value then matches = false; break end
|
||||
end
|
||||
for key, _ in pairs(stored) do
|
||||
if eligible[key] == nil then matches = false; break end
|
||||
end
|
||||
if matches then
|
||||
redis.call('EXPIRE', KEYS[1], ARGV[2])
|
||||
return current
|
||||
end
|
||||
end
|
||||
end
|
||||
redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2])
|
||||
return ARGV[1]
|
||||
end
|
||||
if current == ARGV[1] then
|
||||
redis.call('EXPIRE', KEYS[1], ARGV[2])
|
||||
end
|
||||
return current
|
||||
"""
|
||||
|
||||
|
||||
def set_local_affinity_pin(cache: DualCache, cache_key: str, value: object, ttl_seconds: int) -> None:
|
||||
"""Replace the entry because InMemoryCache.set_cache preserves a live key's expiry."""
|
||||
cache.in_memory_cache.delete_cache(cache_key)
|
||||
cache.in_memory_cache.set_cache(cache_key, value, ttl=ttl_seconds)
|
||||
|
||||
|
||||
def _legacy_pin_matches(stored: object, pin_value: Mapping[str, str]) -> bool:
|
||||
if isinstance(stored, dict):
|
||||
return all(stored.get(key) is not None and str(stored[key]) == value for key, value in pin_value.items())
|
||||
return isinstance(stored, str) and len(pin_value) == 1 and stored in pin_value.values()
|
||||
|
||||
|
||||
def claim_affinity_pin_in_memory(
|
||||
cache: DualCache,
|
||||
cache_key: str,
|
||||
pin_value: Mapping[str, str],
|
||||
ttl_seconds: int,
|
||||
*,
|
||||
eligible_values: tuple[Mapping[str, str], ...] | None = None,
|
||||
) -> object:
|
||||
"""No await between read and write, so same-loop claims agree during a Redis outage."""
|
||||
existing: Final[object] = cache.in_memory_cache.get_cache(cache_key)
|
||||
if existing is not None and eligible_values is None:
|
||||
if _legacy_pin_matches(existing, pin_value):
|
||||
set_local_affinity_pin(cache, cache_key, pin_value, ttl_seconds)
|
||||
return existing
|
||||
winner: Final = existing if existing is not None and existing in (eligible_values or ()) else pin_value
|
||||
set_local_affinity_pin(cache, cache_key, winner, ttl_seconds)
|
||||
return winner
|
||||
|
||||
|
||||
def _decode_pin(value: str) -> object:
|
||||
try:
|
||||
return _PIN_JSON_ADAPTER.validate_json(value)
|
||||
except ValidationError:
|
||||
return value
|
||||
|
||||
|
||||
async def claim_affinity_pin(
|
||||
cache: DualCache,
|
||||
cache_key: str,
|
||||
pin_value: Mapping[str, str],
|
||||
ttl_seconds: int,
|
||||
*,
|
||||
eligible_values: tuple[Mapping[str, str], ...] | None = None,
|
||||
) -> object:
|
||||
"""Return the authoritative first writer, replacing it only when it becomes ineligible.
|
||||
|
||||
Eligible claims refresh the returned winner. Legacy deployment claims only refresh
|
||||
a matching candidate. Resolve Redis per call because the proxy attaches it lazily.
|
||||
"""
|
||||
redis_cache: Final = cache.redis_cache
|
||||
if redis_cache is not None:
|
||||
try:
|
||||
claim_script: Final = redis_cache.async_register_script(_CLAIM_PIN_SCRIPT)
|
||||
args: Final = (
|
||||
json.dumps(dict(pin_value)), # mutable-ok: JSON serialization requires dict, not a generic Mapping
|
||||
int(ttl_seconds),
|
||||
*(
|
||||
(json.dumps(tuple(dict(value) for value in eligible_values)),) # mutable-ok: JSON requires dict
|
||||
if eligible_values is not None
|
||||
else ()
|
||||
),
|
||||
)
|
||||
raw: Final = cast( # cast-ok: Redis scripts return heterogeneous values; only object is asserted here
|
||||
object, await claim_script(keys=(cache_key,), args=args)
|
||||
)
|
||||
decoded: Final = raw.decode("utf-8") if isinstance(raw, bytes) else raw
|
||||
if not isinstance(decoded, str):
|
||||
return pin_value
|
||||
winner: Final = _decode_pin(decoded)
|
||||
set_local_affinity_pin(cache, cache_key, winner, ttl_seconds)
|
||||
return winner
|
||||
except Exception as error: # noqa: BLE001 # Redis/Lua faults retain same-pod affinity through local claims
|
||||
verbose_router_logger.debug("Affinity cache: Redis claim failed, using pod-local claim. error=%s", error)
|
||||
return claim_affinity_pin_in_memory(cache, cache_key, pin_value, ttl_seconds, eligible_values=eligible_values)
|
||||
|
|
@ -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]] = []
|
||||
|
|
|
|||
|
|
@ -1976,6 +1976,8 @@ BROWSER_SECURITY_HEADERS: Final[frozenset[str]] = frozenset(
|
|||
|
||||
UNSAFE_PROXY_RESPONSE_HEADERS: Final[frozenset[str]] = HTTP_FRAMING_HEADERS | BROWSER_SECURITY_HEADERS
|
||||
|
||||
STRINGIFIED_NONE: Final[str] = "None"
|
||||
|
||||
# A retrieved response replays the usage of the call that created it, so pricing these
|
||||
# read/management routes like inference bills the same tokens twice.
|
||||
NON_INFERENCE_CALL_TYPES: Final[frozenset[str]] = frozenset(
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -309,8 +309,8 @@ def max_retries_per_request_hit(kwargs: Mapping[str, object], num_retries_per_re
|
|||
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
|
||||
retry_count: Final = metadata.get("request_retry_count")
|
||||
return type(retry_count) is int and 0 < retry_count and num_retries_per_request <= retry_count
|
||||
|
||||
|
||||
def get_or_create_metadata_bucket(
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -1,6 +1,9 @@
|
|||
import inspect
|
||||
import json
|
||||
import re
|
||||
import traceback
|
||||
from collections.abc import Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final, Protocol, cast
|
||||
|
||||
import httpx
|
||||
|
|
@ -202,11 +205,17 @@ def _get_response_headers(original_exception: Exception) -> httpx.Headers | None
|
|||
return _response_headers
|
||||
|
||||
|
||||
def _accepted_init_kwargs(exception_class: type[Exception], candidates: Mapping[str, object]) -> Mapping[str, object]:
|
||||
accepted: Final = inspect.signature(exception_class).parameters
|
||||
return MappingProxyType({name: value for name, value in candidates.items() if name in accepted})
|
||||
|
||||
|
||||
def extract_and_raise_litellm_exception(
|
||||
response: Any | None,
|
||||
error_str: str,
|
||||
model: str,
|
||||
custom_llm_provider: str,
|
||||
body: object | None = None,
|
||||
):
|
||||
"""
|
||||
Covers scenario where litellm sdk calling proxy.
|
||||
|
|
@ -216,32 +225,19 @@ def extract_and_raise_litellm_exception(
|
|||
Relevant Issue: https://github.com/BerriAI/litellm/issues/7259
|
||||
"""
|
||||
pattern: Final = r"litellm\.\w+Error"
|
||||
|
||||
# Search for the exception in the error string
|
||||
match: Final = re.search(pattern, error_str)
|
||||
|
||||
# Extract the exception if found
|
||||
if match:
|
||||
exception_name = match.group(0)
|
||||
exception_name = exception_name.strip().replace("litellm.", "")
|
||||
raised_exception_obj: Final = getattr(litellm, exception_name, None)
|
||||
if raised_exception_obj:
|
||||
# Try with response parameter first, fall back to without it
|
||||
# Some exceptions (e.g., APIConnectionError) don't accept response param
|
||||
try:
|
||||
raise raised_exception_obj(
|
||||
message=error_str,
|
||||
llm_provider=custom_llm_provider,
|
||||
model=model,
|
||||
response=response,
|
||||
)
|
||||
except TypeError:
|
||||
# Exception doesn't accept response parameter
|
||||
raise raised_exception_obj(
|
||||
message=error_str,
|
||||
llm_provider=custom_llm_provider,
|
||||
model=model,
|
||||
)
|
||||
if match is None:
|
||||
return
|
||||
exception_name: Final = match.group(0).removeprefix("litellm.")
|
||||
raised_exception_obj: Final = getattr(litellm, exception_name, None)
|
||||
if not raised_exception_obj:
|
||||
return
|
||||
raise raised_exception_obj(
|
||||
message=error_str,
|
||||
llm_provider=custom_llm_provider,
|
||||
model=model,
|
||||
**_accepted_init_kwargs(raised_exception_obj, MappingProxyType({"response": response, "body": body})),
|
||||
)
|
||||
|
||||
|
||||
class _ProviderHTTPException(Protocol):
|
||||
|
|
@ -254,6 +250,23 @@ class _ProviderHTTPException(Protocol):
|
|||
llm_provider: str
|
||||
|
||||
|
||||
def _litellm_proxy_response(
|
||||
original_exception: _ProviderHTTPException, custom_llm_provider: str
|
||||
) -> httpx.Response | None:
|
||||
response: Final = getattr(original_exception, "response", None)
|
||||
if custom_llm_provider != "litellm_proxy" or not isinstance(response, httpx.Response) or response.headers:
|
||||
return response
|
||||
headers: Final = getattr(original_exception, "headers", None)
|
||||
if not isinstance(headers, Mapping) or not headers:
|
||||
return response
|
||||
pairs: Final = headers.multi_items() if isinstance(headers, httpx.Headers) else headers.items()
|
||||
return httpx.Response(
|
||||
status_code=response.status_code,
|
||||
headers=[(str(k), str(v)) for k, v in pairs],
|
||||
request=getattr(original_exception, "request", None),
|
||||
)
|
||||
|
||||
|
||||
def _map_openai_exception(
|
||||
*,
|
||||
model: str,
|
||||
|
|
@ -264,6 +277,7 @@ def _map_openai_exception(
|
|||
exception_provider: str,
|
||||
extra_information: str,
|
||||
) -> None:
|
||||
response: Final = _litellm_proxy_response(original_exception, custom_llm_provider)
|
||||
# custom_llm_provider is openai, make it OpenAI
|
||||
message = get_error_message(error_obj=original_exception)
|
||||
if message is None:
|
||||
|
|
@ -292,14 +306,14 @@ def _map_openai_exception(
|
|||
message=f"RateLimitError: {exception_provider} - {message}",
|
||||
model=model,
|
||||
llm_provider=custom_llm_provider,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
)
|
||||
elif ExceptionCheckers.is_error_str_context_window_exceeded(error_str):
|
||||
raise ContextWindowExceededError(
|
||||
message=f"ContextWindowExceededError: {exception_provider} - {message}",
|
||||
llm_provider=custom_llm_provider,
|
||||
model=model,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
)
|
||||
elif "invalid_request_error" in error_str and "model_not_found" in error_str:
|
||||
|
|
@ -307,7 +321,7 @@ def _map_openai_exception(
|
|||
message=f"{exception_provider} - {message}",
|
||||
llm_provider=custom_llm_provider,
|
||||
model=model,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
)
|
||||
elif "A timeout occurred" in error_str:
|
||||
|
|
@ -326,8 +340,9 @@ def _map_openai_exception(
|
|||
message=f"ContentPolicyViolationError: {exception_provider} - {message}",
|
||||
llm_provider=custom_llm_provider,
|
||||
model=model,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
body=getattr(original_exception, "body", None),
|
||||
)
|
||||
elif "invalid_encrypted_content" in error_str or "could not be verified" in error_str:
|
||||
helpful_message: Final = (
|
||||
|
|
@ -345,7 +360,7 @@ def _map_openai_exception(
|
|||
message=helpful_message,
|
||||
llm_provider=custom_llm_provider,
|
||||
model=model,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
body=getattr(original_exception, "body", None),
|
||||
)
|
||||
|
|
@ -354,7 +369,7 @@ def _map_openai_exception(
|
|||
message=f"{exception_provider} - {message}",
|
||||
llm_provider=custom_llm_provider,
|
||||
model=model,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
body=getattr(original_exception, "body", None),
|
||||
)
|
||||
|
|
@ -372,7 +387,7 @@ def _map_openai_exception(
|
|||
message=f"RateLimitError: {exception_provider} - {message}",
|
||||
model=model,
|
||||
llm_provider=custom_llm_provider,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
)
|
||||
elif (
|
||||
|
|
@ -383,7 +398,7 @@ def _map_openai_exception(
|
|||
message=f"AuthenticationError: {exception_provider} - {message}",
|
||||
llm_provider=custom_llm_provider,
|
||||
model=model,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
)
|
||||
elif "Mistral API raised a streaming error" in error_str:
|
||||
|
|
@ -402,15 +417,16 @@ def _map_openai_exception(
|
|||
message=f"{exception_provider} - {message}",
|
||||
llm_provider=custom_llm_provider,
|
||||
model=model,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
body=getattr(original_exception, "body", None),
|
||||
)
|
||||
elif original_exception.status_code == 401:
|
||||
raise AuthenticationError(
|
||||
message=f"AuthenticationError: {exception_provider} - {message}",
|
||||
llm_provider=custom_llm_provider,
|
||||
model=model,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
)
|
||||
elif original_exception.status_code == 404:
|
||||
|
|
@ -418,7 +434,7 @@ def _map_openai_exception(
|
|||
message=f"NotFoundError: {exception_provider} - {message}",
|
||||
model=model,
|
||||
llm_provider=custom_llm_provider,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
)
|
||||
elif original_exception.status_code == 408:
|
||||
|
|
@ -433,7 +449,7 @@ def _map_openai_exception(
|
|||
message=f"{exception_provider} - {message}",
|
||||
model=model,
|
||||
llm_provider=custom_llm_provider,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
body=getattr(original_exception, "body", None),
|
||||
)
|
||||
|
|
@ -442,7 +458,7 @@ def _map_openai_exception(
|
|||
message=f"RateLimitError: {exception_provider} - {message}",
|
||||
model=model,
|
||||
llm_provider=custom_llm_provider,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
)
|
||||
elif original_exception.status_code == 500:
|
||||
|
|
@ -450,7 +466,7 @@ def _map_openai_exception(
|
|||
message=f"InternalServerError: {exception_provider} - {message}",
|
||||
model=model,
|
||||
llm_provider=custom_llm_provider,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
)
|
||||
elif original_exception.status_code == 502:
|
||||
|
|
@ -458,7 +474,7 @@ def _map_openai_exception(
|
|||
message=f"BadGatewayError: {exception_provider} - {message}",
|
||||
model=model,
|
||||
llm_provider=custom_llm_provider,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
)
|
||||
elif original_exception.status_code == 503:
|
||||
|
|
@ -466,7 +482,7 @@ def _map_openai_exception(
|
|||
message=f"ServiceUnavailableError: {exception_provider} - {message}",
|
||||
model=model,
|
||||
llm_provider=custom_llm_provider,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
)
|
||||
elif original_exception.status_code == 504: # gateway timeout error
|
||||
|
|
@ -2423,10 +2439,11 @@ def exception_type(
|
|||
custom_llm_provider == "litellm_proxy"
|
||||
): # handle special case where calling litellm proxy + exception str contains error message
|
||||
extract_and_raise_litellm_exception(
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=_litellm_proxy_response(mappable_exception, custom_llm_provider),
|
||||
error_str=error_str,
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
body=getattr(original_exception, "body", None),
|
||||
)
|
||||
if (
|
||||
custom_llm_provider == "openai"
|
||||
|
|
|
|||
|
|
@ -484,11 +484,12 @@ def apply_off_peak_pricing(model_info: ModelInfo, current_time: datetime | None,
|
|||
def _apply_off_peak_to_base_costs(
|
||||
model_info: ModelInfo,
|
||||
current_time: datetime | None,
|
||||
base_costs: tuple[float, float, float, float, float],
|
||||
base_costs: tuple[float, float, float, float | None, float],
|
||||
) -> tuple[float, float, float, float, float]:
|
||||
"""Apply off-peak rates to an already-resolved set of base costs, whichever pricing path
|
||||
produced them. The one-hour cache-creation rate passes through untouched, since
|
||||
off_peak_pricing has no field for it, and reasoning is left to _resolve_billed_reasoning_rate.
|
||||
produced them. off_peak_pricing has no field for the one-hour cache-creation rate, so a
|
||||
present one passes through untouched and an absent one resolves to the applied
|
||||
cache-creation rate. Reasoning is left to _resolve_billed_reasoning_rate.
|
||||
"""
|
||||
prompt, completion, cache_creation, cache_creation_above_1hr, cache_read = base_costs
|
||||
rates: Final = apply_off_peak_pricing(
|
||||
|
|
@ -506,7 +507,7 @@ def _apply_off_peak_to_base_costs(
|
|||
rates.input_rate,
|
||||
rates.output_rate,
|
||||
rates.cache_creation_rate,
|
||||
cache_creation_above_1hr,
|
||||
rates.cache_creation_rate if cache_creation_above_1hr is None else cache_creation_above_1hr,
|
||||
rates.cache_read_rate,
|
||||
)
|
||||
|
||||
|
|
@ -532,6 +533,11 @@ def _get_token_base_cost(
|
|||
`missing_cache_read_uses_input` resolves an absent cache-read rate to the resolved
|
||||
input rate instead of 0.0; an explicit 0.0 rate stays a real price either way.
|
||||
|
||||
An absent cache-creation rate always resolves to the resolved input rate, the way the
|
||||
tiered table and custom deployment pricing already do, since a provider that publishes
|
||||
no write price bills cache writes as ordinary input. An absent 1h write rate resolves
|
||||
to the cache-creation rate, off-peak included. An explicit 0.0 stays a real price for both.
|
||||
|
||||
Returns:
|
||||
Tuple[float, float, float, float] - (prompt_cost, completion_cost, cache_creation_cost, cache_read_cost)
|
||||
"""
|
||||
|
|
@ -554,10 +560,9 @@ def _get_token_base_cost(
|
|||
output_image_cost: Final = _get_cost_per_unit(model_info, "output_cost_per_image_token", None)
|
||||
if output_image_cost is not None:
|
||||
completion_base_cost = cast(float, output_image_cost)
|
||||
cache_creation_cost = cast(float, _get_cost_per_unit(model_info, cache_creation_cost_key))
|
||||
cache_creation_cost_above_1hr = cast(
|
||||
float,
|
||||
_get_cost_per_unit(model_info, "cache_creation_input_token_cost_above_1hr"),
|
||||
cache_creation_cost = _get_cost_per_unit(model_info, cache_creation_cost_key, default_value=None)
|
||||
cache_creation_cost_above_1hr = _get_cost_per_unit(
|
||||
model_info, "cache_creation_input_token_cost_above_1hr", default_value=None
|
||||
)
|
||||
cache_read_cost = _get_cost_per_unit(model_info, cache_read_cost_key, default_value=None)
|
||||
|
||||
|
|
@ -639,22 +644,10 @@ def _get_token_base_cost(
|
|||
else f"cache_read_input_token_cost_above_{threshold_str}_tokens"
|
||||
)
|
||||
|
||||
cache_creation_cost = cast(
|
||||
float,
|
||||
_get_cost_per_unit(
|
||||
model_info,
|
||||
cache_creation_tiered_key,
|
||||
cache_creation_cost,
|
||||
),
|
||||
)
|
||||
cache_creation_cost = _get_cost_per_unit(model_info, cache_creation_tiered_key, cache_creation_cost)
|
||||
|
||||
cache_creation_cost_above_1hr = cast(
|
||||
float,
|
||||
_get_cost_per_unit(
|
||||
model_info,
|
||||
cache_creation_1hr_tiered_key,
|
||||
cache_creation_cost_above_1hr,
|
||||
),
|
||||
cache_creation_cost_above_1hr = _get_cost_per_unit(
|
||||
model_info, cache_creation_1hr_tiered_key, cache_creation_cost_above_1hr
|
||||
)
|
||||
|
||||
cache_read_cost = _get_cost_per_unit(model_info, cache_read_tiered_key, cache_read_cost)
|
||||
|
|
@ -665,16 +658,16 @@ def _get_token_base_cost(
|
|||
except Exception:
|
||||
continue
|
||||
|
||||
input_rate_for_missing_cache_rates: Final = _off_peak_rate(
|
||||
_open_off_peak_block(model_info, current_time) or MappingProxyType({}),
|
||||
"input_cost_per_token",
|
||||
prompt_base_cost,
|
||||
)
|
||||
if cache_read_cost is None:
|
||||
cache_read_cost = (
|
||||
_off_peak_rate(
|
||||
_open_off_peak_block(model_info, current_time) or MappingProxyType({}),
|
||||
"input_cost_per_token",
|
||||
prompt_base_cost,
|
||||
)
|
||||
if missing_cache_read_uses_input
|
||||
else 0.0
|
||||
)
|
||||
cache_read_cost = input_rate_for_missing_cache_rates if missing_cache_read_uses_input else 0.0
|
||||
resolved_cache_creation_cost: Final = (
|
||||
input_rate_for_missing_cache_rates if cache_creation_cost is None else cache_creation_cost
|
||||
)
|
||||
|
||||
return _apply_off_peak_to_base_costs(
|
||||
model_info,
|
||||
|
|
@ -682,7 +675,7 @@ def _get_token_base_cost(
|
|||
(
|
||||
prompt_base_cost,
|
||||
completion_base_cost,
|
||||
cache_creation_cost,
|
||||
resolved_cache_creation_cost,
|
||||
cache_creation_cost_above_1hr,
|
||||
cache_read_cost,
|
||||
),
|
||||
|
|
@ -956,12 +949,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 +967,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"))
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ from itertools import chain, repeat
|
|||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Protocol, cast, overload, runtime_checkable
|
||||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
from typing_extensions import ReadOnly, TypedDict, assert_never
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
|
@ -44,6 +45,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,
|
||||
|
|
@ -103,9 +105,24 @@ class ToolResultBlockTextTarget:
|
|||
block_idx: int
|
||||
|
||||
|
||||
InputWriteBackTarget = (
|
||||
MessageContentTarget | ContentBlockTextTarget | ToolResultStringTarget | ToolResultBlockTextTarget
|
||||
)
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class SystemStringTarget:
|
||||
pass
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class SystemBlockTextTarget:
|
||||
block_idx: int
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ToolUseInputTarget:
|
||||
msg_idx: int
|
||||
content_idx: int
|
||||
|
||||
|
||||
MessageTextTarget = MessageContentTarget | ContentBlockTextTarget | ToolResultStringTarget | ToolResultBlockTextTarget
|
||||
InputWriteBackTarget = SystemStringTarget | SystemBlockTextTarget | MessageTextTarget
|
||||
|
||||
|
||||
def _as_str_mapping(value: Mapping[str, object]) -> Mapping[str, object]:
|
||||
|
|
@ -146,10 +163,17 @@ class ScannedText:
|
|||
target: InputWriteBackTarget
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ScannedToolCall:
|
||||
tool_call: ChatCompletionToolCallChunk
|
||||
target: ToolUseInputTarget
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ExtractedInput:
|
||||
scanned: tuple[ScannedText, ...]
|
||||
images: tuple[str, ...]
|
||||
tool_calls: tuple[ScannedToolCall, ...] = ()
|
||||
|
||||
|
||||
EMPTY_EXTRACTED_INPUT: Final = ExtractedInput(scanned=(), images=())
|
||||
|
|
@ -161,6 +185,74 @@ class _ToolCallShape:
|
|||
arguments: str
|
||||
|
||||
|
||||
def _is_client_tool_use(block: Mapping[str, object]) -> bool:
|
||||
return (
|
||||
block.get("type") == "tool_use"
|
||||
and isinstance(block.get("id"), str)
|
||||
and isinstance(block.get("name"), str)
|
||||
and isinstance(block.get("input"), dict)
|
||||
)
|
||||
|
||||
|
||||
def _write_back_system_block(system: object, block_idx: int, response: str) -> None:
|
||||
if not isinstance(system, list):
|
||||
return
|
||||
text_blocks: Final = tuple(block for block in system if isinstance(block, dict) and block.get("type") == "text")
|
||||
if block_idx < len(text_blocks):
|
||||
text_blocks[block_idx]["text"] = (
|
||||
response # mutable-ok: guardrails rewrite the caller's request payload in place
|
||||
)
|
||||
|
||||
|
||||
def _write_back_message_text(message: _WritableMessage, target: MessageTextTarget, response: str) -> None:
|
||||
content: Final = message.get("content", None)
|
||||
if content is None:
|
||||
return
|
||||
match target:
|
||||
case MessageContentTarget():
|
||||
if isinstance(content, str):
|
||||
message["content"] = response # mutable-ok: guardrails rewrite the caller's request payload in place
|
||||
case ContentBlockTextTarget(content_idx=content_idx):
|
||||
if isinstance(content, list):
|
||||
content[content_idx]["text"] = (
|
||||
response # mutable-ok: guardrails rewrite the caller's request payload in place
|
||||
)
|
||||
case ToolResultStringTarget(content_idx=content_idx):
|
||||
if isinstance(content, list):
|
||||
content[content_idx]["content"] = (
|
||||
response # mutable-ok: guardrails rewrite the caller's request payload in place
|
||||
)
|
||||
case ToolResultBlockTextTarget(content_idx=content_idx, block_idx=block_idx):
|
||||
if isinstance(content, list):
|
||||
content[content_idx]["content"][block_idx]["text"] = (
|
||||
response # mutable-ok: guardrails rewrite the caller's request payload in place
|
||||
)
|
||||
case _:
|
||||
assert_never(target)
|
||||
|
||||
|
||||
_TOOL_USE_INPUT_ADAPTER: Final = TypeAdapter(dict[str, object])
|
||||
|
||||
|
||||
def _rewritten_tool_use_input(arguments: str) -> Mapping[str, object] | None:
|
||||
try:
|
||||
return _TOOL_USE_INPUT_ADAPTER.validate_json(arguments)
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
def _write_back_tool_use(
|
||||
message: _WritableMessage, target: ToolUseInputTarget, shape: _ToolCallShape, rewritten_input: Mapping[str, object]
|
||||
) -> None:
|
||||
content: Final = message.get("content", None)
|
||||
block: Final = content[target.content_idx] if isinstance(content, list) else None
|
||||
if not isinstance(block, dict):
|
||||
return
|
||||
block["input"] = rewritten_input # mutable-ok: guardrails rewrite the caller's request payload in place
|
||||
if shape.name is not None and shape.name != block.get("name"):
|
||||
block["name"] = shape.name # mutable-ok: guardrails rewrite the caller's request payload in place
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _SSEFieldRewrite:
|
||||
"""One field of one nested section of a buffered SSE event, rewritten."""
|
||||
|
|
@ -452,9 +544,8 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
skip_tool: Final = effective_skip_tool_message_for_guardrail(guardrail_to_apply)
|
||||
scan_only_tool_results: Final = effective_scan_only_tool_results_for_guardrail(guardrail_to_apply)
|
||||
|
||||
# Exclude only the trusted top-level prompt. In-sequence system entries are untrusted
|
||||
# and must stay aligned with texts_to_check for positional masking. When the top-level
|
||||
# prompt is included, the pre-existing count mismatch disables positional masking.
|
||||
# The top-level prompt is translated on its own below so it can be hoisted in front of
|
||||
# any mid-turn system entries and scanned first, aligned with that structured position.
|
||||
translation_source: Final = { # mutable-ok: API message payload
|
||||
key: value for key, value in data.items() if key != "system"
|
||||
}
|
||||
|
|
@ -490,7 +581,12 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
]
|
||||
)
|
||||
|
||||
# Step 1: Extract all text content and images
|
||||
# Step 1: Extract all text content, images, and tool calls
|
||||
top_level_system_scanned: Final = (
|
||||
()
|
||||
if hoisted_system_message is None or scan_only_tool_results
|
||||
else self._extract_top_level_system_text(hoisted_system_message)
|
||||
)
|
||||
extracted: Final = tuple(
|
||||
self._extract_input_text_and_images(
|
||||
message=message,
|
||||
|
|
@ -501,17 +597,27 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
)
|
||||
for msg_idx, message in enumerate(messages)
|
||||
)
|
||||
scanned: Final = tuple(item for one_message in extracted for item in one_message.scanned)
|
||||
scanned: Final = (
|
||||
*top_level_system_scanned,
|
||||
*(item for one_message in extracted for item in one_message.scanned),
|
||||
)
|
||||
texts_to_check: Final = [item.text for item in scanned] # mutable-ok: GenericGuardrailAPIInputs takes list[str]
|
||||
images_to_check: Final = [
|
||||
image for one_message in extracted for image in one_message.images
|
||||
] # mutable-ok: GenericGuardrailAPIInputs takes list[str]
|
||||
scanned_tool_calls: Final = tuple(item for one_message in extracted for item in one_message.tool_calls)
|
||||
tool_calls_to_check: Final = [
|
||||
item.tool_call for item in scanned_tool_calls
|
||||
] # mutable-ok: GenericGuardrailAPIInputs takes list[ChatCompletionToolCallChunk]
|
||||
pre_guardrail_tool_calls: Final = _tool_call_shapes(tool_calls_to_check)
|
||||
|
||||
# Step 2: Apply guardrail to all texts in batch
|
||||
if texts_to_check:
|
||||
# Step 2: Apply guardrail to all texts and tool calls in batch
|
||||
if texts_to_check or tool_calls_to_check:
|
||||
inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check)
|
||||
if images_to_check:
|
||||
inputs["images"] = images_to_check
|
||||
if tool_calls_to_check:
|
||||
inputs["tool_calls"] = tool_calls_to_check
|
||||
if tools_to_check:
|
||||
inputs["tools"] = tools_to_check
|
||||
original_structured_messages: Final = structured_messages
|
||||
|
|
@ -570,9 +676,18 @@ 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)
|
||||
self._apply_guardrail_tool_calls_to_input(
|
||||
messages=messages,
|
||||
scanned_tool_calls=scanned_tool_calls,
|
||||
pre_guardrail_tool_calls=pre_guardrail_tool_calls,
|
||||
returned_tool_calls=guardrailed_inputs.get("tool_calls"),
|
||||
guardrail_name=guardrail_to_apply.guardrail_name,
|
||||
)
|
||||
# Step 3: Map guardrail responses back to original message structure
|
||||
await self._apply_guardrail_responses_to_input(
|
||||
messages=messages,
|
||||
data=data,
|
||||
responses=guardrailed_texts,
|
||||
scanned=scanned,
|
||||
)
|
||||
|
|
@ -598,6 +713,19 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
hoisted: Final = probe.get("messages") or [] # mutable-ok: API message payload
|
||||
return hoisted[0] if hoisted else None
|
||||
|
||||
@staticmethod
|
||||
def _extract_top_level_system_text(hoisted_system_message: AllMessageValues) -> tuple[ScannedText, ...]:
|
||||
content: Final = hoisted_system_message.get("content")
|
||||
if isinstance(content, str):
|
||||
return (ScannedText(content, SystemStringTarget()),)
|
||||
if not isinstance(content, list):
|
||||
return ()
|
||||
return tuple(
|
||||
ScannedText(text_str, SystemBlockTextTarget(block_idx))
|
||||
for block_idx, block in enumerate(content)
|
||||
if isinstance(block, dict) and isinstance(text_str := block.get("text"), str)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _openai_system_message_to_anthropic(
|
||||
message: Mapping[str, object],
|
||||
|
|
@ -852,9 +980,25 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
for content_idx, content_item in enumerate(content)
|
||||
if isinstance(content_item, dict)
|
||||
)
|
||||
tool_use_blocks: Final = (
|
||||
()
|
||||
if scan_only_tool_results
|
||||
else tuple(
|
||||
(content_idx, content_item)
|
||||
for content_idx, content_item in enumerate(content)
|
||||
if isinstance(content_item, dict) and _is_client_tool_use(content_item)
|
||||
)
|
||||
)
|
||||
return ExtractedInput(
|
||||
scanned=tuple(item for block in blocks for item in block.scanned),
|
||||
images=tuple(image for block in blocks for image in block.images),
|
||||
tool_calls=tuple(
|
||||
ScannedToolCall(
|
||||
tool_call=AnthropicConfig.convert_tool_use_to_openai_format(content_item, tool_call_idx),
|
||||
target=ToolUseInputTarget(msg_idx, content_idx),
|
||||
)
|
||||
for tool_call_idx, (content_idx, content_item) in enumerate(tool_use_blocks)
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
|
|
@ -940,43 +1084,59 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
|
||||
async def _apply_guardrail_responses_to_input(
|
||||
self,
|
||||
messages: Sequence[_WritableMessage],
|
||||
responses: list[str],
|
||||
data: dict[str, object], # mutable-ok: API message payload
|
||||
responses: Sequence[str],
|
||||
scanned: tuple[ScannedText, ...],
|
||||
) -> None:
|
||||
"""
|
||||
Apply guardrail responses back to input messages.
|
||||
Apply guardrail responses back to the top-level system prompt and the input messages.
|
||||
"""
|
||||
raw_messages: Final = data.get("messages")
|
||||
messages: Final[Sequence[_WritableMessage]] = raw_messages if isinstance(raw_messages, list) else ()
|
||||
for item, guardrail_response in zip(scanned, responses):
|
||||
target = item.target
|
||||
message = messages[target.msg_idx]
|
||||
content = message.get("content", None)
|
||||
if content is None:
|
||||
continue
|
||||
|
||||
match target:
|
||||
case MessageContentTarget():
|
||||
if isinstance(content, str):
|
||||
message["content"] = (
|
||||
guardrail_response # mutable-ok: guardrails rewrite the caller's request payload in place
|
||||
)
|
||||
case ContentBlockTextTarget(content_idx=content_idx):
|
||||
if isinstance(content, list):
|
||||
content[content_idx]["text"] = (
|
||||
guardrail_response # mutable-ok: guardrails rewrite the caller's request payload in place
|
||||
)
|
||||
case ToolResultStringTarget(content_idx=content_idx):
|
||||
if isinstance(content, list):
|
||||
content[content_idx]["content"] = (
|
||||
guardrail_response # mutable-ok: guardrails rewrite the caller's request payload in place
|
||||
)
|
||||
case ToolResultBlockTextTarget(content_idx=content_idx, block_idx=block_idx):
|
||||
if isinstance(content, list):
|
||||
content[content_idx]["content"][block_idx]["text"] = (
|
||||
match item.target:
|
||||
case SystemStringTarget():
|
||||
if isinstance(data.get("system"), str):
|
||||
data["system"] = (
|
||||
guardrail_response # mutable-ok: guardrails rewrite the caller's request payload in place
|
||||
)
|
||||
case SystemBlockTextTarget(block_idx=block_idx):
|
||||
_write_back_system_block(data.get("system"), block_idx, guardrail_response)
|
||||
case (
|
||||
MessageContentTarget()
|
||||
| ContentBlockTextTarget()
|
||||
| ToolResultStringTarget()
|
||||
| ToolResultBlockTextTarget() as message_target
|
||||
):
|
||||
_write_back_message_text(messages[message_target.msg_idx], message_target, guardrail_response)
|
||||
case _:
|
||||
assert_never(target)
|
||||
assert_never(item.target)
|
||||
|
||||
@staticmethod
|
||||
def _apply_guardrail_tool_calls_to_input(
|
||||
messages: Sequence[_WritableMessage],
|
||||
scanned_tool_calls: tuple[ScannedToolCall, ...],
|
||||
pre_guardrail_tool_calls: tuple[_ToolCallShape, ...],
|
||||
returned_tool_calls: Sequence[object] | None,
|
||||
guardrail_name: str | None,
|
||||
) -> None:
|
||||
post_guardrail_tool_calls: Final = _tool_call_shapes(
|
||||
returned_tool_calls
|
||||
if returned_tool_calls is not None and len(returned_tool_calls) == len(pre_guardrail_tool_calls)
|
||||
else tuple(item.tool_call for item in scanned_tool_calls)
|
||||
)
|
||||
rewritten: Final = tuple(
|
||||
(item, after, _rewritten_tool_use_input(after.arguments))
|
||||
for item, before, after in zip(scanned_tool_calls, pre_guardrail_tool_calls, post_guardrail_tool_calls)
|
||||
if before != after
|
||||
)
|
||||
applicable: Final = tuple(
|
||||
(item, after, rewritten_input) for item, after, rewritten_input in rewritten if rewritten_input is not None
|
||||
)
|
||||
if len(applicable) != len(rewritten):
|
||||
raise unappliable_request_rewrite(guardrail_name)
|
||||
for item, after, rewritten_input in applicable:
|
||||
_write_back_tool_use(messages[item.target.msg_idx], item.target, after, rewritten_input)
|
||||
|
||||
async def process_output_response(
|
||||
self,
|
||||
|
|
|
|||
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")
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ from concurrent.futures import ThreadPoolExecutor
|
|||
from datetime import datetime
|
||||
from functools import partial
|
||||
from threading import Lock
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, ParamSpec, TypeVar, cast, get_args, overload
|
||||
|
||||
import httpx
|
||||
|
|
@ -96,6 +97,77 @@ def _assume_role_params(
|
|||
)
|
||||
|
||||
|
||||
_SecureTransportBool = TypedDict("_SecureTransportBool", {"aws:SecureTransport": ReadOnly[Literal["true"]]})
|
||||
|
||||
|
||||
class _SecureTransportCondition(TypedDict):
|
||||
Bool: ReadOnly[_SecureTransportBool]
|
||||
|
||||
|
||||
class _SessionPolicyStatement(TypedDict):
|
||||
Sid: ReadOnly[str]
|
||||
Effect: ReadOnly[Literal["Allow"]]
|
||||
Action: ReadOnly[tuple[str, ...]]
|
||||
Resource: ReadOnly[Literal["*"]]
|
||||
Condition: ReadOnly[_SecureTransportCondition]
|
||||
|
||||
|
||||
class WebIdentitySessionPolicy(TypedDict):
|
||||
Version: ReadOnly[Literal["2012-10-17"]]
|
||||
Statement: ReadOnly[tuple[_SessionPolicyStatement, ...]]
|
||||
|
||||
|
||||
_WEB_IDENTITY_SESSION_POLICY_ACTIONS: Final[Mapping[str, tuple[str, ...]]] = MappingProxyType(
|
||||
{
|
||||
"BedrockLiteLLM": (
|
||||
"bedrock:InvokeModel",
|
||||
"bedrock:InvokeModelWithResponseStream",
|
||||
"bedrock:CountTokens",
|
||||
"bedrock:Rerank",
|
||||
"bedrock:Retrieve",
|
||||
"bedrock:ListKnowledgeBases",
|
||||
"bedrock:InvokeAgent",
|
||||
"bedrock:ApplyGuardrail",
|
||||
"bedrock:GetGuardrail",
|
||||
"bedrock:ListGuardrails",
|
||||
),
|
||||
"BedrockAgentCoreLiteLLM": (
|
||||
"bedrock-agentcore:InvokeAgentRuntime",
|
||||
"bedrock-agentcore:InvokeAgentRuntimeForUser",
|
||||
"bedrock-agentcore:InvokeGateway",
|
||||
),
|
||||
"ClaudePlatformLiteLLM": (
|
||||
"aws-external-anthropic:CreateInference",
|
||||
"aws-external-anthropic:CreateBatchInference",
|
||||
"aws-external-anthropic:CancelBatchInference",
|
||||
"aws-external-anthropic:DeleteBatchInference",
|
||||
"aws-external-anthropic:CountTokens",
|
||||
"aws-external-anthropic:Get*",
|
||||
"aws-external-anthropic:List*",
|
||||
),
|
||||
"BedrockMantleLiteLLM": ("bedrock-mantle:CreateInference",),
|
||||
}
|
||||
)
|
||||
|
||||
_SECURE_TRANSPORT_ONLY: Final = _SecureTransportCondition(Bool=_SecureTransportBool({"aws:SecureTransport": "true"}))
|
||||
|
||||
|
||||
def build_web_identity_session_policy() -> WebIdentitySessionPolicy:
|
||||
return WebIdentitySessionPolicy(
|
||||
Version="2012-10-17",
|
||||
Statement=tuple(
|
||||
_SessionPolicyStatement(
|
||||
Sid=sid,
|
||||
Effect="Allow",
|
||||
Action=actions,
|
||||
Resource="*",
|
||||
Condition=_SECURE_TRANSPORT_ONLY,
|
||||
)
|
||||
for sid, actions in _WEB_IDENTITY_SESSION_POLICY_ACTIONS.items()
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class BedrockRequestTarget(BaseModel):
|
||||
aws_region_name: str
|
||||
aws_bedrock_runtime_endpoint: str | None
|
||||
|
|
@ -940,60 +1012,12 @@ class BaseAWSLLM(SignsRequestsWithAWS):
|
|||
# auth only (static creds + IRSA take other code paths).
|
||||
# https://docs.aws.amazon.com/STS/latest/APIReference/API_AssumeRoleWithWebIdentity.html
|
||||
# https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sts/client/assume_role_with_web_identity.html
|
||||
bedrock_session_policy: Final = {
|
||||
"Version": "2012-10-17",
|
||||
"Statement": [
|
||||
{
|
||||
"Sid": "BedrockLiteLLM",
|
||||
"Effect": "Allow",
|
||||
"Action": [
|
||||
"bedrock:InvokeModel",
|
||||
"bedrock:InvokeModelWithResponseStream",
|
||||
"bedrock:CountTokens",
|
||||
"bedrock:ApplyGuardrail",
|
||||
"bedrock:GetGuardrail",
|
||||
"bedrock:ListGuardrails",
|
||||
],
|
||||
"Resource": "*",
|
||||
"Condition": {"Bool": {"aws:SecureTransport": "true"}},
|
||||
},
|
||||
# Claude Platform on AWS (added by #27678 for the
|
||||
# ``bedrock/claude_platform/<model>`` route) lives under
|
||||
# a separate IAM action namespace; without these entries
|
||||
# the OIDC path 403s on every claude_platform request
|
||||
# even with a fully permissive identity policy (#30200).
|
||||
{
|
||||
"Sid": "ClaudePlatformLiteLLM",
|
||||
"Effect": "Allow",
|
||||
"Action": [
|
||||
"aws-external-anthropic:CreateInference",
|
||||
"aws-external-anthropic:CreateBatchInference",
|
||||
"aws-external-anthropic:CancelBatchInference",
|
||||
"aws-external-anthropic:DeleteBatchInference",
|
||||
"aws-external-anthropic:CountTokens",
|
||||
"aws-external-anthropic:Get*",
|
||||
"aws-external-anthropic:List*",
|
||||
],
|
||||
"Resource": "*",
|
||||
"Condition": {"Bool": {"aws:SecureTransport": "true"}},
|
||||
},
|
||||
{
|
||||
"Sid": "BedrockMantleLiteLLM",
|
||||
"Effect": "Allow",
|
||||
"Action": [
|
||||
"bedrock-mantle:CreateInference",
|
||||
],
|
||||
"Resource": "*",
|
||||
"Condition": {"Bool": {"aws:SecureTransport": "true"}},
|
||||
},
|
||||
],
|
||||
}
|
||||
assume_role_params: Final = {
|
||||
"RoleArn": aws_role_name,
|
||||
"RoleSessionName": aws_session_name,
|
||||
"WebIdentityToken": oidc_token,
|
||||
"DurationSeconds": 3600,
|
||||
"Policy": json.dumps(bedrock_session_policy, separators=(",", ":")),
|
||||
"Policy": json.dumps(build_web_identity_session_policy(), separators=(",", ":")),
|
||||
}
|
||||
|
||||
# Add ExternalId parameter if provided
|
||||
|
|
|
|||
|
|
@ -17,7 +17,7 @@ BaseAWSLLM._sign_request after the request body is finalized.
|
|||
|
||||
import json
|
||||
from collections.abc import Mapping
|
||||
from typing import Any, Final
|
||||
from typing import Any, Final, cast # noqa: TID251 # map_openai_params returns the filtered params as a bare dict
|
||||
|
||||
import httpx
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
|
@ -32,6 +32,7 @@ from litellm.llms.bedrock_mantle.common_utils import (
|
|||
BedrockMantleAuthMixin,
|
||||
)
|
||||
from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
|
||||
from litellm.responses.additional_tools import HoistedAdditionalTools, hoist_additional_tools
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.llms.openai import (
|
||||
ResponseInputParam,
|
||||
|
|
@ -58,8 +59,6 @@ _BEDROCK_MANTLE_SUPPORTED_RESPONSE_TOOL_TYPES: Final = frozenset(
|
|||
_BEDROCK_MANTLE_SUPPORTED_SERVICE_TIERS: Final = frozenset({"auto", "default"})
|
||||
_BEDROCK_MANTLE_OPENAI_PATH_SUPPORTED_REASONING_SUMMARIES: Final = frozenset({"auto"})
|
||||
|
||||
_CODEX_ADDITIONAL_TOOLS_INPUT_ITEM_TYPE: Final = "additional_tools"
|
||||
|
||||
_CODEX_AGENT_MESSAGE_INPUT_ITEM_TYPE: Final = "agent_message"
|
||||
_CODEX_CONTEXT_COMPACTION_INPUT_ITEM_TYPE: Final = "context_compaction"
|
||||
_CODEX_LOCAL_SHELL_CALL_INPUT_ITEM_TYPE: Final = "local_shell_call"
|
||||
|
|
@ -233,17 +232,14 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI
|
|||
litellm_params: GenericLiteLLMParams,
|
||||
headers: dict,
|
||||
) -> dict:
|
||||
remaining_input, hoisted_tools = self._hoist_codex_additional_tools(input)
|
||||
normalized_input: Final = self._normalize_codex_input_items(remaining_input)
|
||||
params: Final = cast( # cast-ok: the base signature leaves the params dict untyped
|
||||
"ResponsesAPIOptionalRequestParams", response_api_optional_request_params
|
||||
)
|
||||
hoisted: Final = hoist_additional_tools(input, params.get("tools"))
|
||||
normalized_input: Final = self._normalize_codex_input_items(hoisted.input)
|
||||
request_params: Final = (
|
||||
{
|
||||
**response_api_optional_request_params,
|
||||
"tools": [
|
||||
*(response_api_optional_request_params.get("tools") or []),
|
||||
*hoisted_tools,
|
||||
],
|
||||
}
|
||||
if hoisted_tools
|
||||
self._params_with_hoisted_tools(params, hoisted)
|
||||
if hoisted.hoisted
|
||||
else response_api_optional_request_params
|
||||
)
|
||||
return super().transform_responses_api_request(
|
||||
|
|
@ -254,41 +250,14 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI
|
|||
headers=headers,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _is_codex_additional_tools_item(item: Any) -> bool:
|
||||
return isinstance(item, dict) and item.get("type") == _CODEX_ADDITIONAL_TOOLS_INPUT_ITEM_TYPE
|
||||
|
||||
@staticmethod
|
||||
def _tools_of_additional_tools_item(item: "dict[str, Any]") -> "list[Any]":
|
||||
tools: Final = item.get("tools")
|
||||
return tools if isinstance(tools, list) else []
|
||||
|
||||
@classmethod
|
||||
def _hoist_codex_additional_tools(
|
||||
cls,
|
||||
input: "str | ResponseInputParam",
|
||||
) -> "tuple[str | ResponseInputParam, list[Any]]":
|
||||
"""Codex's "responses lite" wire mode ships tool definitions inside
|
||||
`input` as {"type": "additional_tools", "role": "developer",
|
||||
"tools": [...]} items. api.openai.com accepts that item type; Mantle
|
||||
rejects the whole request with 400 "Invalid 'input': value did not
|
||||
match any expected variant" but accepts the same tools at the top
|
||||
level, so move them there and strip the items from `input`.
|
||||
"""
|
||||
if not isinstance(input, list):
|
||||
return input, []
|
||||
additional_tools_items: Final = [item for item in input if cls._is_codex_additional_tools_item(item)]
|
||||
if not additional_tools_items:
|
||||
return input, []
|
||||
remaining_input: Final = [item for item in input if not cls._is_codex_additional_tools_item(item)]
|
||||
hoisted_tools = [tool for item in additional_tools_items for tool in cls._tools_of_additional_tools_item(item)]
|
||||
verbose_logger.debug(
|
||||
"Bedrock Mantle Responses API: hoisting %d tool(s) out of %d 'additional_tools' input item(s) "
|
||||
"into the top-level tools param (Mantle rejects that input item type).",
|
||||
len(hoisted_tools),
|
||||
len(additional_tools_items),
|
||||
)
|
||||
return remaining_input, cls._filter_unsupported_tools(hoisted_tools)
|
||||
def _params_with_hoisted_tools(
|
||||
cls, params: Mapping[str, object], hoisted: HoistedAdditionalTools
|
||||
) -> dict[str, object]:
|
||||
supported_tools: Final = cls._filter_unsupported_tools(list(hoisted.tools))
|
||||
if supported_tools:
|
||||
return {**params, "tools": supported_tools}
|
||||
return {key: value for key, value in params.items() if key != "tools"}
|
||||
|
||||
@staticmethod
|
||||
def _agent_message_text(item: "Mapping[str, object]") -> str:
|
||||
|
|
|
|||
|
|
@ -115,6 +115,19 @@ def _parse_setup(session_configuration_request: str) -> BidiGenerateContentSetup
|
|||
return envelope.get("setup", empty_setup)
|
||||
|
||||
|
||||
def _grounding_metadata_from_frame(frame: Mapping[str, object]) -> tuple[Mapping[str, object], ...]:
|
||||
"""Read ``serverContent.groundingMetadata`` off the frame that carries the turn's usage.
|
||||
|
||||
Live reports grounding in the server frames rather than in ``usageMetadata``, and it emits both
|
||||
on the same frame, so the per-query charge is countable at the point usage is built.
|
||||
"""
|
||||
server_content: Final = frame.get("serverContent")
|
||||
if not isinstance(server_content, Mapping):
|
||||
return ()
|
||||
metadata: Final = server_content.get("groundingMetadata")
|
||||
return (metadata,) if isinstance(metadata, Mapping) else ()
|
||||
|
||||
|
||||
# Google bills Live transcription at an estimated 25 audio tokens/sec of input and
|
||||
# 175 text tokens/min of output (ai.google.dev/gemini-api/docs/pricing).
|
||||
GEMINI_LIVE_TRANSCRIBE_AUDIO_TOKENS_PER_SECOND: Final = 25
|
||||
|
|
@ -323,7 +336,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig):
|
|||
)
|
||||
elif key == "input_audio_transcription" and value is not None:
|
||||
optional_params["inputAudioTranscription"] = {}
|
||||
elif key == "turn_detection":
|
||||
elif key == "turn_detection" and value is not None:
|
||||
value_typed = cast(OpenAIRealtimeTurnDetection, value)
|
||||
if (
|
||||
isinstance(value_typed, dict)
|
||||
|
|
@ -1049,6 +1062,11 @@ class GeminiRealtimeConfig(BaseRealtimeConfig):
|
|||
{**cast(dict, message), "usageMetadata": resolved_usage_metadata},
|
||||
),
|
||||
)
|
||||
grounding_metadata: Final = _grounding_metadata_from_frame(message)
|
||||
if grounding_metadata:
|
||||
VertexGeminiConfig._set_grounding_usage_counters( # pyright: ignore[reportPrivateUsage] # shared with the chat path; no public alias exists yet
|
||||
_chat_completion_usage, grounding_metadata
|
||||
)
|
||||
else:
|
||||
_chat_completion_usage = get_empty_usage()
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
@ -210,6 +213,17 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
task_mappings=tool_call_task_mappings,
|
||||
)
|
||||
|
||||
elif (
|
||||
not images_to_check
|
||||
and not guardrail_to_apply.records_own_guardrail_information
|
||||
and (not_run_reason := self._not_run_reason(messages)) is not None
|
||||
):
|
||||
guardrail_to_apply.add_standard_logging_guardrail_information_to_request_data(
|
||||
guardrail_json_response=not_run_reason,
|
||||
request_data=data,
|
||||
guardrail_status="not_run",
|
||||
)
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
"OpenAI Chat Completions: Processed input messages: %s",
|
||||
data.get("messages"),
|
||||
|
|
@ -217,6 +231,28 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
|
||||
return data
|
||||
|
||||
def _not_run_reason(
|
||||
self,
|
||||
messages: Sequence[dict[str, Any]], # mutable-ok: raw request messages consumed by _extract_inputs
|
||||
) -> str | None:
|
||||
"""Why nothing was scanned, or None when the only unscoped content is images, which this handler never scans."""
|
||||
texts: Final[list[str]] = [] # mutable-ok: filled by _extract_inputs
|
||||
images: Final[list[str]] = [] # mutable-ok: filled by _extract_inputs
|
||||
tool_calls: Final[list[ChatCompletionToolParam]] = [] # mutable-ok: filled by _extract_inputs
|
||||
for msg_idx, message in enumerate(messages):
|
||||
self._extract_inputs(
|
||||
message=message,
|
||||
msg_idx=msg_idx,
|
||||
texts_to_check=texts,
|
||||
images_to_check=images,
|
||||
tool_calls_to_check=tool_calls,
|
||||
text_task_mappings=[], # mutable-ok: required by _extract_inputs, unused here
|
||||
tool_call_task_mappings=[], # mutable-ok: required by _extract_inputs, unused here
|
||||
)
|
||||
if texts or tool_calls:
|
||||
return "no scannable content after message scoping"
|
||||
return None if images else "no scannable content"
|
||||
|
||||
def extract_request_tool_names(self, data: dict) -> list[str]:
|
||||
"""Extract tool names from OpenAI chat completions request (tools[].function.name, functions[].name)."""
|
||||
names: Final[list[str]] = []
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -6,8 +6,10 @@ from typing import Final, TypeAlias
|
|||
from pydantic import BaseModel, TypeAdapter, ValidationError
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.responses.litellm_completion_transformation.custom_tools import custom_tool_grammar_suffix
|
||||
from litellm.responses.litellm_completion_transformation.transformation import (
|
||||
NAMESPACE_DESCRIPTION_SEPARATOR,
|
||||
NAMESPACE_MEMBER_TYPES_WITH_CHAT_TOOLS,
|
||||
LiteLLMCompletionResponsesConfig,
|
||||
)
|
||||
|
||||
|
|
@ -34,8 +36,8 @@ def _validated_tools(values: Iterable[object]) -> tuple[Tool, ...]:
|
|||
return tuple(tool for tool in validated if tool is not None)
|
||||
|
||||
|
||||
def _is_function(tool: Tool) -> bool:
|
||||
return tool.get("type") == "function"
|
||||
def _has_chat_tool(member: Tool) -> bool:
|
||||
return member.get("type") in NAMESPACE_MEMBER_TYPES_WITH_CHAT_TOOLS
|
||||
|
||||
|
||||
def _chat_tool_key(tool: Tool) -> str:
|
||||
|
|
@ -67,18 +69,19 @@ def _function_fields(tool: Tool) -> Tool:
|
|||
return function if function is not None else MappingProxyType({})
|
||||
|
||||
|
||||
def _without_namespace_prefix(key: str, value: object, prefix: str) -> object:
|
||||
if key != "description" or not isinstance(value, str) or not value.startswith(prefix):
|
||||
def _member_description(key: str, value: object, prefix: str, suffix: str) -> object:
|
||||
if key != "description" or not isinstance(value, str):
|
||||
return value
|
||||
return value[len(prefix) :]
|
||||
return value.replace(prefix, "", 1).replace(suffix, "", 1)
|
||||
|
||||
|
||||
def _rebuilt_member(member: Tool, flattened: Tool, guardrailed: Tool, namespace_description: str) -> Tool:
|
||||
flattened_function: Final = _function_fields(flattened)
|
||||
prefix: Final = f"{namespace_description}{NAMESPACE_DESCRIPTION_SEPARATOR}" if namespace_description else ""
|
||||
suffix: Final = custom_tool_grammar_suffix(member.get("format")) if member.get("type") == "custom" else ""
|
||||
changed_function: Final = MappingProxyType(
|
||||
{
|
||||
key: _without_namespace_prefix(key, value, prefix)
|
||||
key: _member_description(key, value, prefix, suffix)
|
||||
for key, value in _function_fields(guardrailed).items()
|
||||
if flattened_function.get(key) != value
|
||||
}
|
||||
|
|
@ -93,8 +96,8 @@ def _rebuilt_member(member: Tool, flattened: Tool, guardrailed: Tool, namespace_
|
|||
return {**member, **changed_extras, **changed_function} # mutable-ok: json.dumps rejects MappingProxyType
|
||||
|
||||
|
||||
def _rebuilt_function_members(
|
||||
function_members: Sequence[Tool],
|
||||
def _rebuilt_flattened_members(
|
||||
flattened_members: Sequence[Tool],
|
||||
flattened_group: Sequence[Tool],
|
||||
group_keys: Sequence[IndexedKey],
|
||||
guardrailed_by_key: Mapping[IndexedKey, Tool],
|
||||
|
|
@ -106,7 +109,7 @@ def _rebuilt_function_members(
|
|||
else member
|
||||
if guardrailed_by_key[key] == flattened
|
||||
else _rebuilt_member(member, flattened, guardrailed_by_key[key], namespace_description)
|
||||
for member, flattened, key in zip(function_members, flattened_group, group_keys)
|
||||
for member, flattened, key in zip(flattened_members, flattened_group, group_keys)
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -118,9 +121,9 @@ def _rebuilt_namespace(
|
|||
guardrailed_by_key: Mapping[IndexedKey, Tool],
|
||||
) -> tuple[Tool, ...]:
|
||||
namespace_description: Final = str(original.get("description") or "")
|
||||
rebuilt_functions: Final = iter(
|
||||
_rebuilt_function_members(
|
||||
tuple(member for member in members if _is_function(member)),
|
||||
rebuilt_flattened: Final = iter(
|
||||
_rebuilt_flattened_members(
|
||||
tuple(member for member in members if _has_chat_tool(member)),
|
||||
flattened_group,
|
||||
group_keys,
|
||||
guardrailed_by_key,
|
||||
|
|
@ -129,7 +132,7 @@ def _rebuilt_namespace(
|
|||
)
|
||||
rebuilt_members: Final = tuple(
|
||||
rebuilt
|
||||
for rebuilt in (next(rebuilt_functions) if _is_function(member) else member for member in members)
|
||||
for rebuilt in (next(rebuilt_flattened) if _has_chat_tool(member) else member for member in members)
|
||||
if rebuilt is not None
|
||||
)
|
||||
if not rebuilt_members:
|
||||
|
|
@ -149,7 +152,7 @@ def _merged_original(
|
|||
if guardrailed_group == tuple(flattened_group):
|
||||
return (original,)
|
||||
members: Final = _namespace_members(original) if original.get("type") == "namespace" else ()
|
||||
if members and sum(map(_is_function, members)) == len(flattened_group):
|
||||
if members and sum(map(_has_chat_tool, members)) == len(flattened_group):
|
||||
return _rebuilt_namespace(original, members, flattened_group, group_keys, guardrailed_by_key)
|
||||
if not guardrailed_group:
|
||||
return ()
|
||||
|
|
|
|||
|
|
@ -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": [
|
||||
{
|
||||
|
|
|
|||
|
|
@ -4,11 +4,12 @@ import os
|
|||
from collections.abc import Callable, Mapping
|
||||
from datetime import datetime
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple
|
||||
from typing import TYPE_CHECKING, Annotated, Any, Final, Literal, NamedTuple
|
||||
|
||||
import httpx
|
||||
from pydantic import (
|
||||
BaseModel,
|
||||
BeforeValidator,
|
||||
ConfigDict,
|
||||
Field,
|
||||
Json,
|
||||
|
|
@ -47,6 +48,7 @@ from litellm.types.proxy.carried_budget_state import (
|
|||
)
|
||||
from litellm.types.proxy.control_plane_endpoints import WorkerRegistryEntry
|
||||
from litellm.types.router import RouterErrors, UpdateRouterConfig
|
||||
from litellm.types.router_weights import validate_router_settings_dict
|
||||
from litellm.types.secret_managers.main import KeyManagementSystem
|
||||
from litellm.types.utils import (
|
||||
CallTypes,
|
||||
|
|
@ -656,6 +658,7 @@ class LiteLLMRoutes(enum.Enum):
|
|||
[
|
||||
# user
|
||||
"/user/new",
|
||||
"/management/v1/users/bulk",
|
||||
"/user/update",
|
||||
"/user/bulk_update",
|
||||
"/user/delete",
|
||||
|
|
@ -890,6 +893,7 @@ class LiteLLMRoutes(enum.Enum):
|
|||
"/auto_router/validate_complexity_router_config",
|
||||
# Per-session auto-router read - the endpoint scopes the row to the caller's own key hash
|
||||
"/auto_router/session",
|
||||
"/cost/predict-cache",
|
||||
# Agent registry - reads are role-scoped and writes are proxy-admin-gated
|
||||
# inside agent_endpoints/endpoints.py
|
||||
*agent_management_routes,
|
||||
|
|
@ -1983,8 +1987,14 @@ class OrgMember(MemberBase):
|
|||
|
||||
from litellm.models.team import TeamBase as TeamBase # noqa: E402
|
||||
|
||||
RouterSettingsDict = Annotated[
|
||||
dict[str, object],
|
||||
BeforeValidator(validate_router_settings_dict, json_schema_input_type=UpdateRouterConfig),
|
||||
]
|
||||
|
||||
|
||||
class NewTeamRequest(TeamBase):
|
||||
router_settings: RouterSettingsDict | None = None
|
||||
model_aliases: dict | None = None
|
||||
tags: list | None = None
|
||||
guardrails: list[str] | None = None
|
||||
|
|
@ -2082,7 +2092,7 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase):
|
|||
allowed_vector_store_indexes: list[AllowedVectorStoreIndexItem] | None = None
|
||||
enforced_batch_output_expires_after: dict | None = None
|
||||
enforced_file_expires_after: dict | None = None
|
||||
router_settings: dict | None = None
|
||||
router_settings: RouterSettingsDict | None = None
|
||||
access_group_ids: list[str] | None = None
|
||||
budget_limits: list[BudgetLimitEntry] | None = None # multiple concurrent budget windows
|
||||
default_team_member_models: list[str] | None = None # default allowed_models seeded onto new team members
|
||||
|
|
@ -4421,6 +4431,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):
|
||||
|
|
|
|||
|
|
@ -895,6 +895,7 @@ async def common_checks(
|
|||
request_query_params=_safe_get_request_query_params(request=request),
|
||||
llm_router=llm_router,
|
||||
request=request,
|
||||
team_id=valid_token.team_id if valid_token is not None else None,
|
||||
)
|
||||
|
||||
skip_all_budget_checks: Final = skip_budget_checks or (
|
||||
|
|
@ -4471,7 +4472,7 @@ async def stamp_matched_model_access_groups(
|
|||
|
||||
async def can_key_call_model(
|
||||
model: str | list[str],
|
||||
llm_model_list: list | None,
|
||||
llm_model_list: Sequence[object] | None,
|
||||
valid_token: UserAPIKeyAuth,
|
||||
llm_router: litellm.Router | None,
|
||||
) -> Literal[True]:
|
||||
|
|
@ -4518,7 +4519,7 @@ async def can_key_call_model(
|
|||
|
||||
async def can_key_call_resolved_model(
|
||||
model: str,
|
||||
llm_model_list: list | None,
|
||||
llm_model_list: Sequence[object] | None,
|
||||
valid_token: UserAPIKeyAuth,
|
||||
llm_router: litellm.Router | None,
|
||||
) -> None:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -249,6 +249,7 @@ async def authenticate_user(
|
|||
|
||||
if os.getenv("DATABASE_URL") is not None:
|
||||
response = await generate_key_helper_fn(
|
||||
llm_router=None,
|
||||
request_type="key",
|
||||
**{
|
||||
"user_role": LitellmUserRoles.PROXY_ADMIN,
|
||||
|
|
@ -324,6 +325,7 @@ async def authenticate_user(
|
|||
await _rehash_password_if_needed(_user_row.user_id, password, _password)
|
||||
if os.getenv("DATABASE_URL") is not None:
|
||||
response = await generate_key_helper_fn(
|
||||
llm_router=None,
|
||||
request_type="key",
|
||||
**{
|
||||
"user_role": user_role,
|
||||
|
|
|
|||
|
|
@ -24,6 +24,7 @@ _PROXY_ADMIN_VIEW_ONLY_BLOCKED_ROUTES: Final = frozenset(
|
|||
[
|
||||
# user
|
||||
"/user/new",
|
||||
"/management/v1/users/bulk",
|
||||
"/user/delete",
|
||||
"/management/v1/users/bulk_delete",
|
||||
"/user/bulk_update",
|
||||
|
|
@ -760,6 +761,7 @@ class RouteChecks:
|
|||
_ADMIN_VIEWER_BLOCKED_WRITE_ROUTES = frozenset(
|
||||
[
|
||||
"/user/new",
|
||||
"/management/v1/users/bulk",
|
||||
"/user/delete",
|
||||
"/management/v1/users/bulk_delete",
|
||||
"/user/bulk_update",
|
||||
|
|
|
|||
|
|
@ -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]"):
|
||||
|
|
@ -876,6 +878,7 @@ async def _auto_register_jwt_mapping(
|
|||
# the NOT NULL @id constraint. Every successful key-creation caller (e.g.
|
||||
# /key/generate) passes table_name="key" explicitly.
|
||||
key_data: Final = await generate_key_helper_fn(
|
||||
llm_router=None,
|
||||
request_type="key",
|
||||
table_name="key",
|
||||
team_id=team_id,
|
||||
|
|
@ -1652,6 +1655,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 +1696,7 @@ async def _user_api_key_auth_builder(
|
|||
route=route,
|
||||
request=request,
|
||||
llm_router=llm_router,
|
||||
team_id=valid_token.team_id,
|
||||
)
|
||||
),
|
||||
)
|
||||
|
|
@ -2091,6 +2096,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 +2215,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 +2246,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 +2742,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 +2859,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 +3312,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 +3420,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 +3462,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)
|
||||
|
||||
|
|
|
|||
|
|
@ -580,8 +580,8 @@ What the command changed is recorded in `~/.litellm/claude_configure_state.json`
|
|||
`lite configure claude`, `lite login --config-claude`, `lite up` and `lite autoroute up` also install a status line (`~/.litellm/statusline.py`, registered as `statusLine` in `~/.claude/settings.json` unless you already run one) that shows which model the auto-router actually served the last turn and, once the proxy has recorded the session, what the session cost against the router's savings baseline:
|
||||
|
||||
```
|
||||
claude-auto · Routed to: claude-haiku-4-5 -63% vs Claude Opus 5
|
||||
LiteLLM ████████░░░░░░░░░░░░░░░░ $0.14
|
||||
Routed to: claude-haiku-4-5 -63% vs Claude Opus 5
|
||||
claude-auto ████████░░░░░░░░░░░░░░░░ $0.14
|
||||
Claude Opus 5 ████████████████████████ $0.38
|
||||
```
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -28,6 +28,7 @@ import os
|
|||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
import unicodedata
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from collections.abc import Callable, Mapping
|
||||
|
|
@ -42,7 +43,6 @@ FETCH_TIMEOUT_SECONDS: Final = 3
|
|||
BAR_WIDTH: Final = 24
|
||||
BAR_FULL: Final = "\u2588"
|
||||
BAR_EMPTY: Final = "\u2591"
|
||||
SEPARATOR: Final = " \u00b7 "
|
||||
TRANSCRIPT_SCAN_LIMIT_BYTES: Final = 4 * 1024 * 1024
|
||||
CLAUDE_BASE_URL_ENV_KEYS: Final = ("ANTHROPIC_BASE_URL",)
|
||||
CLAUDE_API_KEY_ENV_KEYS: Final = ("ANTHROPIC_AUTH_TOKEN", "ANTHROPIC_API_KEY")
|
||||
|
|
@ -50,7 +50,6 @@ CODEX_BASE_URL_ENV_KEYS: Final = ("OPENAI_BASE_URL",)
|
|||
CODEX_API_KEY_ENV_KEYS: Final = ("OPENAI_API_KEY",)
|
||||
CODEX_STOP_EVENT: Final = "Stop"
|
||||
SYNTHETIC_MODEL: Final = "<synthetic>"
|
||||
LITELLM_LABEL: Final = "LiteLLM"
|
||||
RESET: Final = "\033[0m"
|
||||
BOLD: Final = "\033[1m"
|
||||
DIM: Final = "\033[90m"
|
||||
|
|
@ -302,31 +301,37 @@ def _bar(fraction: float, color: str, width: int, use_color: bool) -> str:
|
|||
return f"{color}{BAR_FULL * filled}{DIM}{BAR_EMPTY * (width - filled)}{RESET}"
|
||||
|
||||
|
||||
def _display_width(label: str) -> int:
|
||||
return sum(
|
||||
2 if unicodedata.east_asian_width(character) in ("W", "F") else 1
|
||||
for character in label
|
||||
if unicodedata.category(character) not in ("Mn", "Me")
|
||||
)
|
||||
|
||||
|
||||
def render(model: str, session: Session | None, config_dir: Path, use_color: bool, bar_width: int = BAR_WIDTH) -> str:
|
||||
def paint(code: str, text: str) -> str:
|
||||
return f"{code}{text}{RESET}" if use_color else text
|
||||
|
||||
routed: Final = paint(BOLD, f"Routed to: {model}")
|
||||
if session is None:
|
||||
if session is None or session.baseline_model is None or session.baseline_spend <= 0:
|
||||
return routed
|
||||
header: Final = f"{session.router_name}{SEPARATOR}{routed}"
|
||||
if session.baseline_model is None or session.baseline_spend <= 0:
|
||||
return header
|
||||
reference: Final = baseline_label(session.baseline_model, config_dir)
|
||||
pct: Final = (session.baseline_spend - session.spend) / session.baseline_spend * 100
|
||||
delta: Final = paint(LITELLM_COLOR, f"{'-' if pct >= 0 else '+'}{abs(round(pct))}% vs {reference}")
|
||||
peak: Final = max(session.spend, session.baseline_spend)
|
||||
label_width: Final = max(len(LITELLM_LABEL), len(reference))
|
||||
label_width: Final = max(_display_width(session.router_name), _display_width(reference))
|
||||
rows: Final = (
|
||||
(LITELLM_LABEL, session.spend, LITELLM_COLOR),
|
||||
(session.router_name, session.spend, LITELLM_COLOR),
|
||||
(reference, session.baseline_spend, BASELINE_COLOR),
|
||||
)
|
||||
lines: Final = (
|
||||
f"{paint(DIM, label.ljust(label_width))} {_bar(amount / peak, color, bar_width, use_color)} "
|
||||
f"{paint(DIM, label + ' ' * (label_width - _display_width(label)))} "
|
||||
f"{_bar(amount / peak, color, bar_width, use_color)} "
|
||||
f"{paint(DIM, f'${amount:.2f}')}"
|
||||
for label, amount, color in rows
|
||||
)
|
||||
return "\n".join((f"{header} {delta}", *lines))
|
||||
return "\n".join((f"{routed} {delta}", *lines))
|
||||
|
||||
|
||||
def color_enabled(env: Mapping[str, str]) -> bool:
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ import httpx
|
|||
import orjson
|
||||
from fastapi import HTTPException, Request, status
|
||||
from fastapi.responses import JSONResponse, Response, StreamingResponse
|
||||
from pydantic import ValidationError
|
||||
from starlette.types import Receive, Scope, Send
|
||||
|
||||
import litellm
|
||||
|
|
@ -76,6 +77,7 @@ from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_di
|
|||
from litellm.router_utils.common_utils import resolve_model_group_alias
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.router import RouterRateLimitError
|
||||
from litellm.types.router_weights import validate_router_weights
|
||||
|
||||
_LateResponseT = TypeVar("_LateResponseT", bound=Response)
|
||||
_LlmCallT = TypeVar("_LlmCallT")
|
||||
|
|
@ -1571,6 +1573,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 +1603,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,
|
||||
|
|
@ -1936,6 +1941,13 @@ class ProxyBaseLLMRequestProcessing:
|
|||
# This avoids expensive Router instantiation on each request
|
||||
if router_settings is not None:
|
||||
self.data["router_settings_override"] = router_settings
|
||||
try:
|
||||
self.data["_router_weights"] = validate_router_weights(router_settings.get("weights"))
|
||||
except ValidationError:
|
||||
self.data["_router_weights"] = None
|
||||
verbose_proxy_logger.warning(
|
||||
"Ignoring invalid saved router weights; update team/key router_settings"
|
||||
)
|
||||
alias_target: Final = await _resolve_per_request_model_group_alias(
|
||||
requested_model=self.data.get("model"),
|
||||
router_settings=router_settings,
|
||||
|
|
|
|||
|
|
@ -8,6 +8,8 @@ from typing import Final
|
|||
|
||||
from fastapi import status
|
||||
|
||||
from litellm.constants import STRINGIFIED_NONE
|
||||
|
||||
_OPENAI_ERROR_TYPE_BY_STATUS: Final[Mapping[int, str]] = MappingProxyType(
|
||||
{
|
||||
status.HTTP_401_UNAUTHORIZED: "authentication_error",
|
||||
|
|
@ -35,7 +37,7 @@ def openai_error_type(exc: object, status_code: int) -> str:
|
|||
"""OpenAI types ``error.type`` as a required string, so an exception carrying none
|
||||
falls back to the type its status code stands for."""
|
||||
carried: Final = attribute_of(exc, "type")
|
||||
if isinstance(carried, str):
|
||||
if isinstance(carried, str) and carried != STRINGIFIED_NONE:
|
||||
return carried
|
||||
mapped: Final = _OPENAI_ERROR_TYPE_BY_STATUS.get(status_code)
|
||||
if mapped is not None:
|
||||
|
|
@ -49,4 +51,4 @@ def openai_error_param(exc: object) -> str | None:
|
|||
"""OpenAI types ``error.param`` as nullable, so an exception carrying none
|
||||
serializes as JSON ``null``."""
|
||||
carried: Final = attribute_of(exc, "param")
|
||||
return carried if isinstance(carried, str) else None
|
||||
return carried if isinstance(carried, str) and carried != STRINGIFIED_NONE else None
|
||||
|
|
|
|||
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
|
||||
|
|
@ -26,7 +26,7 @@ class ComplianceChecker:
|
|||
|
||||
def __init__(self, data: ComplianceCheckRequest):
|
||||
self.data = data
|
||||
self.guardrails = data.guardrail_information or []
|
||||
self.guardrails = tuple(g for g in data.guardrail_information or () if g.get("guardrail_status") != "not_run")
|
||||
|
||||
def _get_guardrails_by_mode(self, mode: str) -> list[dict]:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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 ()
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -1600,8 +1600,8 @@ class PanwPrismaAirsHandler(CustomGuardrail):
|
|||
|
||||
Args:
|
||||
texts: Flattened text entries from the framework.
|
||||
messages: Original request messages (request_data["messages"]),
|
||||
NOT structured_messages (which may have injected system content).
|
||||
messages: The structured messages the framework flattened into ``texts``,
|
||||
hoisted top-level system prompt included, so positions line up.
|
||||
|
||||
Returns a set of scannable indices, or None on count mismatch or no user/developer
|
||||
message (safety fallback to existing role-filter behavior).
|
||||
|
|
@ -1788,15 +1788,10 @@ class PanwPrismaAirsHandler(CustomGuardrail):
|
|||
structured_messages: Final = inputs.get("structured_messages")
|
||||
if structured_messages:
|
||||
# For Anthropic /v1/messages: default to latest-user-only scanning.
|
||||
# Uses request_data["messages"] (original format), NOT structured_messages
|
||||
# (which has injected system content from adapter translation).
|
||||
if self._use_latest_user_only(request_data, logging_obj):
|
||||
original_messages: Final = request_data.get("messages")
|
||||
if original_messages:
|
||||
scannable_indices = self._get_latest_user_text_indices(texts, original_messages)
|
||||
scannable_indices = self._get_latest_user_text_indices(texts, structured_messages)
|
||||
# Fall through to existing role filtering if:
|
||||
# - not Anthropic, OR flag explicitly False, OR
|
||||
# - no original messages, OR
|
||||
# - latest-user extraction returned None (no user / count mismatch)
|
||||
if scannable_indices is None:
|
||||
scannable_indices = self._get_scannable_text_indices(texts, structured_messages)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -42,7 +42,7 @@ if TYPE_CHECKING:
|
|||
router: Final = APIRouter()
|
||||
|
||||
_EMPTY_UNITS: Final[Mapping[str, int]] = MappingProxyType({})
|
||||
_ACTION_SEVERITY: Final[Mapping[str, int]] = MappingProxyType({"passed": 0, "flagged": 1, "blocked": 2})
|
||||
_ACTION_SEVERITY: Final[Mapping[str, int]] = MappingProxyType({"not_run": 0, "passed": 1, "flagged": 2, "blocked": 3})
|
||||
|
||||
_T = TypeVar("_T")
|
||||
|
||||
|
|
@ -325,7 +325,7 @@ class UsageDetailResponse(BaseModel):
|
|||
class UsageLogEntry(BaseModel):
|
||||
id: str
|
||||
timestamp: str
|
||||
action: str # blocked | passed | flagged
|
||||
action: str # blocked | passed | flagged | not_run
|
||||
score: float | None
|
||||
latency_ms: float | None
|
||||
model: str | None
|
||||
|
|
|
|||
|
|
@ -193,10 +193,12 @@ async def _upsert_rows_with_retry(
|
|||
|
||||
|
||||
def guardrail_status_to_action(status: str | None) -> str:
|
||||
"""Map StandardLogging guardrail_status to blocked/passed/flagged."""
|
||||
"""Map StandardLogging guardrail_status to blocked/passed/flagged/not_run."""
|
||||
if not status:
|
||||
return "passed"
|
||||
s: Final = (status or "").lower()
|
||||
if s == "not_run":
|
||||
return "not_run"
|
||||
if "intervened" in s or "block" in s:
|
||||
return "blocked"
|
||||
if "flagged" in s or "fail" in s or "error" in s:
|
||||
|
|
@ -354,37 +356,49 @@ async def process_spend_logs_guardrail_usage(
|
|||
"flagged_count": 0,
|
||||
}
|
||||
)
|
||||
index_rows: Final[list[dict[str, object]]] = []
|
||||
index_rows_by_key: Final[dict[tuple[str, str], dict[str, object]]] = {}
|
||||
|
||||
for payload in logs_to_process:
|
||||
request_id = payload.get("request_id")
|
||||
start_time = _parse_payload_start_time(payload)
|
||||
if not request_id or start_time is None:
|
||||
if not isinstance(request_id, str) or not request_id or start_time is None:
|
||||
continue
|
||||
date_key = _date_str(start_time)
|
||||
|
||||
for entry in _parse_guardrail_info_from_payload(payload):
|
||||
guardrail_id = entry.get("guardrail_id") or entry.get("guardrail_name") or ""
|
||||
if not guardrail_id:
|
||||
entries = _parse_guardrail_info_from_payload(payload)
|
||||
ids_by_name = MappingProxyType(
|
||||
{
|
||||
e["guardrail_name"]: e["guardrail_id"]
|
||||
for e in entries
|
||||
if e.get("guardrail_id") and isinstance(e.get("guardrail_name"), str) and e["guardrail_name"]
|
||||
}
|
||||
)
|
||||
for entry in entries:
|
||||
raw_name = entry.get("guardrail_name")
|
||||
guardrail_name = raw_name if isinstance(raw_name, str) else ""
|
||||
guardrail_id = entry.get("guardrail_id") or ids_by_name.get(guardrail_name) or guardrail_name
|
||||
if not isinstance(guardrail_id, str) or not guardrail_id:
|
||||
continue
|
||||
key = _MetricsKey(guardrail_id, date_key)
|
||||
daily_guardrail[key]["requests_evaluated"] += 1
|
||||
action = guardrail_status_to_action(entry.get("guardrail_status"))
|
||||
if action == "passed":
|
||||
daily_guardrail[key]["passed_count"] += 1
|
||||
elif action == "blocked":
|
||||
daily_guardrail[key]["blocked_count"] += 1
|
||||
else:
|
||||
daily_guardrail[key]["flagged_count"] += 1
|
||||
if action != "not_run":
|
||||
key = _MetricsKey(guardrail_id, date_key)
|
||||
daily_guardrail[key]["requests_evaluated"] += 1
|
||||
if action == "passed":
|
||||
daily_guardrail[key]["passed_count"] += 1
|
||||
elif action == "blocked":
|
||||
daily_guardrail[key]["blocked_count"] += 1
|
||||
else:
|
||||
daily_guardrail[key]["flagged_count"] += 1
|
||||
policy_id = entry.get("policy_id")
|
||||
index_rows.append(
|
||||
{
|
||||
prior = index_rows_by_key.get((request_id, guardrail_id))
|
||||
if prior is None or (prior["policy_id"] is None and policy_id is not None):
|
||||
index_rows_by_key[(request_id, guardrail_id)] = {
|
||||
"request_id": request_id,
|
||||
"guardrail_id": guardrail_id,
|
||||
"policy_id": policy_id,
|
||||
"start_time": start_time,
|
||||
}
|
||||
)
|
||||
index_rows: Final = tuple(index_rows_by_key.values())
|
||||
|
||||
async with pending.lock:
|
||||
pending_metrics: Final = pending.metrics
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -1,5 +1,6 @@
|
|||
"""Contract machinery shared by every LiteLLM-defined list route, on any surface."""
|
||||
|
||||
from collections.abc import Sequence
|
||||
from typing import Final
|
||||
from urllib.parse import urlencode
|
||||
|
||||
|
|
@ -7,6 +8,7 @@ from fastapi import Request
|
|||
from fastapi.dependencies.utils import get_flat_params
|
||||
from fastapi.params import ParamTypes
|
||||
from fastapi.responses import JSONResponse
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm.types.proxy.management_endpoints.management_v1 import (
|
||||
ListLinks,
|
||||
|
|
@ -56,6 +58,40 @@ def escape_like(value: str) -> str:
|
|||
return value.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_")
|
||||
|
||||
|
||||
class ValidationErrorDetail(TypedDict):
|
||||
"""The keys of a pydantic/FastAPI validation error a problem document needs."""
|
||||
|
||||
type: ReadOnly[str]
|
||||
loc: ReadOnly[tuple[int | str, ...]]
|
||||
msg: ReadOnly[str]
|
||||
|
||||
|
||||
def _is_length_error_of_rejected_items(error: ValidationErrorDetail, errors: Sequence[ValidationErrorDetail]) -> bool:
|
||||
"""pydantic counts only items that validated, so a bad item also trips the parent's min_length."""
|
||||
return error["type"] == "too_short" and any(
|
||||
len(other["loc"]) > len(error["loc"]) and other["loc"][: len(error["loc"])] == error["loc"] for other in errors
|
||||
)
|
||||
|
||||
|
||||
def request_validation_problem(raw_errors: Sequence[ValidationErrorDetail]) -> ProblemDetail:
|
||||
"""A body that fails validation (an unknown field included) is 422; a bad query parameter is 400."""
|
||||
errors: Final = tuple(error for error in raw_errors if not _is_length_error_of_rejected_items(error, raw_errors))
|
||||
detail: Final = "; ".join(f"{'.'.join(str(part) for part in error['loc'][1:])}: {error['msg']}" for error in errors)
|
||||
if any(error["loc"] and error["loc"][0] == "body" for error in errors):
|
||||
return ProblemDetail(
|
||||
type=f"{PROBLEM_TYPE_BASE}invalid-request-body",
|
||||
title="Invalid request body",
|
||||
status=422,
|
||||
detail=detail or "The request body is invalid.",
|
||||
)
|
||||
return ProblemDetail(
|
||||
type=f"{PROBLEM_TYPE_BASE}invalid-query-parameter",
|
||||
title="Invalid query parameter",
|
||||
status=400,
|
||||
detail=detail or "The request query parameters are invalid.",
|
||||
)
|
||||
|
||||
|
||||
def unknown_query_param_problem(unknown: tuple[str, ...], allowed: tuple[str, ...]) -> ProblemDetail:
|
||||
return ProblemDetail(
|
||||
type=f"{PROBLEM_TYPE_BASE}unknown-query-parameter",
|
||||
|
|
|
|||
|
|
@ -221,6 +221,8 @@ LITELLM_TRACE_CONTROL_METADATA_FIELDS: Final = frozenset(
|
|||
)
|
||||
|
||||
_UNTRUSTED_ROOT_CONTROL_FIELDS: Final = (
|
||||
"weights",
|
||||
"_router_weights",
|
||||
"proxy_server_request",
|
||||
"standard_logging_object",
|
||||
"secret_fields",
|
||||
|
|
@ -334,7 +336,7 @@ _CLIENT_PRICING_METADATA_FIELDS: Final = frozenset({"model_info", "standard_logg
|
|||
# and read by spend logs as fact; a client value has no legitimate meaning and no
|
||||
# key or team setting keeps it, so the strip is never gated.
|
||||
_ROUTER_RESERVED_METADATA_FIELDS: Final = frozenset(
|
||||
{"attempted_fallbacks", "original_model_group", CLIENT_OUTPUT_CEILING_METADATA_KEY}
|
||||
{"attempted_fallbacks", "original_model_group", "request_retry_count", CLIENT_OUTPUT_CEILING_METADATA_KEY}
|
||||
)
|
||||
_ALLOW_CLIENT_PRICING_OVERRIDE_METADATA_KEY: Final = "allow_client_pricing_override"
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -567,7 +567,7 @@ async def new_user(
|
|||
teams = check_if_default_team_set()
|
||||
organization_ids: Final = cast(list[str] | None, data_json.pop("organizations", None))
|
||||
|
||||
response: Final = await generate_key_helper_fn(request_type="user", **data_json)
|
||||
response: Final = await generate_key_helper_fn(request_type="user", **data_json, llm_router=None)
|
||||
# Admin UI Logic
|
||||
# Add User to Team and Organization
|
||||
# if team_id passed add this user to the team
|
||||
|
|
|
|||
|
|
@ -96,6 +96,7 @@ from litellm.proxy.management_endpoints.common_utils import (
|
|||
from litellm.proxy.management_endpoints.model_management_endpoints import (
|
||||
_add_model_to_db,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.router_weights import validate_router_settings_weights
|
||||
from litellm.proxy.management_helpers.access_group_key_sync import (
|
||||
sync_key_access_group_membership,
|
||||
sync_key_regeneration_access_group_membership,
|
||||
|
|
@ -148,6 +149,7 @@ from litellm.types.proxy.management_endpoints.key_management_endpoints import (
|
|||
BulkUpdateKeyRequest,
|
||||
BulkUpdateKeyResponse,
|
||||
BulkUpdateTeamKeysRequest,
|
||||
CustomKeyPolicyRequest,
|
||||
FailedKeyUpdate,
|
||||
KeySearchWhere,
|
||||
SuccessfulKeyUpdate,
|
||||
|
|
@ -201,6 +203,10 @@ class _KeyUpdateResult(TypedDict):
|
|||
data: ReadOnly[Mapping[str, object]]
|
||||
|
||||
|
||||
class _StoredKeyRouterSettings(BaseModel):
|
||||
router_settings: Mapping[str, object] | None = None
|
||||
|
||||
|
||||
class _KeyRowWhere(TypedDict):
|
||||
token: ReadOnly[str]
|
||||
|
||||
|
|
@ -280,6 +286,7 @@ def _config_table(prisma_client: PrismaClient) -> _ConfigTableActions:
|
|||
class _CustomKeyHooksModule(Protocol):
|
||||
user_custom_key_generate: Callable[..., Awaitable[Mapping[str, object]]] | None
|
||||
user_custom_key_update: Callable[..., Awaitable[Mapping[str, object]]] | None
|
||||
user_custom_key_policy: Callable[..., Awaitable[Mapping[str, object]]] | None
|
||||
|
||||
|
||||
def _custom_key_generate_hook(
|
||||
|
|
@ -294,6 +301,161 @@ def _custom_key_update_hook(
|
|||
return hooks.user_custom_key_update
|
||||
|
||||
|
||||
def _custom_key_policy_hook(
|
||||
hooks: _CustomKeyHooksModule,
|
||||
) -> Callable[..., Awaitable[Mapping[str, object]]] | None:
|
||||
return hooks.user_custom_key_policy
|
||||
|
||||
|
||||
async def _enforce_custom_key_update_policy(
|
||||
hook: Callable[..., Awaitable[Mapping[str, object]]] | None,
|
||||
data: UpdateKeyRequest,
|
||||
) -> None:
|
||||
if hook is None:
|
||||
return
|
||||
if not inspect.iscoroutinefunction(hook):
|
||||
raise ValueError("user_custom_key_update must be a coroutine")
|
||||
result: Final = await hook(data)
|
||||
if not result.get("decision", True):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail=result.get("message", "Authentication Failed - Custom Auth Rule"),
|
||||
)
|
||||
|
||||
|
||||
async def _enforce_custom_key_policy(
|
||||
hook: Callable[..., Awaitable[Mapping[str, object]]] | None,
|
||||
build_policy_request: Callable[[], CustomKeyPolicyRequest],
|
||||
) -> None:
|
||||
if hook is None:
|
||||
return
|
||||
if not inspect.iscoroutinefunction(hook):
|
||||
raise ValueError("user_custom_key_policy must be a coroutine")
|
||||
result: Final = await hook(build_policy_request())
|
||||
if not result.get("decision", True):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail=result.get("message", "Authentication Failed - Custom Auth Rule"),
|
||||
)
|
||||
|
||||
|
||||
_KEY_UPDATE_JSON_STRING_COLUMNS: Final = frozenset({"router_settings", "budget_limits"})
|
||||
|
||||
_KEY_METADATA_REQUEST_FIELDS: Final = frozenset(
|
||||
(*LiteLLM_ManagementEndpoint_MetadataFields_Premium, *LiteLLM_ManagementEndpoint_MetadataFields)
|
||||
)
|
||||
|
||||
|
||||
def _decode_json_string_column(column: str, value: object) -> object:
|
||||
if column in _KEY_UPDATE_JSON_STRING_COLUMNS and isinstance(value, str):
|
||||
return json.loads(value)
|
||||
return value
|
||||
|
||||
|
||||
def _verification_token_from_row(row: Mapping[str, object]) -> LiteLLM_VerificationToken:
|
||||
org_id: Final = row["organization_id"] if "organization_id" in row else row.get("org_id")
|
||||
return LiteLLM_VerificationToken.model_validate(MappingProxyType({**row, "org_id": org_id}))
|
||||
|
||||
|
||||
def _effective_key_after_update(
|
||||
existing_key_row: LiteLLM_VerificationToken,
|
||||
non_default_values: Mapping[str, object],
|
||||
) -> LiteLLM_VerificationToken:
|
||||
overlay: Final = MappingProxyType(
|
||||
{column: _decode_json_string_column(column, value) for column, value in non_default_values.items()}
|
||||
)
|
||||
return _verification_token_from_row(
|
||||
MappingProxyType({**existing_key_row.model_dump(), **overlay, "object_permission": None})
|
||||
)
|
||||
|
||||
|
||||
def _update_policy_request(
|
||||
operation: Literal["update", "regenerate"],
|
||||
existing_key_row: LiteLLM_VerificationToken,
|
||||
non_default_values: Mapping[str, object],
|
||||
request: UpdateKeyRequest | RegenerateKeyRequest,
|
||||
) -> CustomKeyPolicyRequest:
|
||||
return CustomKeyPolicyRequest(
|
||||
operation=operation,
|
||||
existing_key=_verification_token_from_row(existing_key_row.model_dump()),
|
||||
effective_key=_effective_key_after_update(
|
||||
existing_key_row=existing_key_row, non_default_values=non_default_values
|
||||
),
|
||||
request=request,
|
||||
)
|
||||
|
||||
|
||||
def _generate_budget_windows(
|
||||
budget_limits: Sequence[BudgetLimitEntry] | None,
|
||||
) -> tuple[Mapping[str, object], ...] | None:
|
||||
if not budget_limits:
|
||||
return None
|
||||
return tuple(
|
||||
MappingProxyType(
|
||||
{
|
||||
**window.model_dump(),
|
||||
"reset_at": get_budget_reset_time(budget_duration=window.budget_duration).isoformat(),
|
||||
}
|
||||
)
|
||||
for window in budget_limits
|
||||
)
|
||||
|
||||
|
||||
def _effective_key_for_generate(data: GenerateKeyRequest, now: datetime) -> LiteLLM_VerificationToken:
|
||||
requested: Final = data.model_dump(exclude_unset=True, exclude_none=True)
|
||||
metadata_fields: Final = MappingProxyType(
|
||||
{field: value for field, value in requested.items() if field in _KEY_METADATA_REQUEST_FIELDS}
|
||||
)
|
||||
column_fields: Final = MappingProxyType(
|
||||
{field: value for field, value in requested.items() if field not in _KEY_METADATA_REQUEST_FIELDS}
|
||||
)
|
||||
metadata: Final = data.metadata or MappingProxyType({})
|
||||
folded_metadata: Final = {**metadata, **metadata_fields} # mutable-ok: encrypt_callback_vars needs a dict
|
||||
columns: Final = handle_key_type(data, {**column_fields}) # mutable-ok: handle_key_type mutates in place
|
||||
expires: Final = (
|
||||
now + timedelta(seconds=duration_in_seconds(duration=data.duration)) if data.duration is not None else None
|
||||
)
|
||||
budget_reset_at: Final = (
|
||||
get_budget_reset_time(budget_duration=data.budget_duration) if data.budget_duration is not None else None
|
||||
)
|
||||
key_rotation_at: Final = (
|
||||
now + timedelta(seconds=duration_in_seconds(duration=data.rotation_interval))
|
||||
if data.auto_rotate and data.rotation_interval
|
||||
else None
|
||||
)
|
||||
return _verification_token_from_row(
|
||||
MappingProxyType(
|
||||
{
|
||||
**columns,
|
||||
"metadata": encrypt_callback_vars(folded_metadata),
|
||||
"expires": expires,
|
||||
"budget_reset_at": budget_reset_at,
|
||||
"key_rotation_at": key_rotation_at,
|
||||
"budget_limits": _generate_budget_windows(data.budget_limits),
|
||||
"object_permission": None,
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
_EMPTY_DURATION_MEANS_UNCHANGED: Final = frozenset({"duration", "budget_duration"})
|
||||
|
||||
|
||||
def _regenerate_request_as_update_request(key: str, data: RegenerateKeyRequest) -> UpdateKeyRequest | None:
|
||||
changed_fields: Final = MappingProxyType(
|
||||
{
|
||||
field: value
|
||||
for field, value in data.model_dump(exclude_unset=True).items()
|
||||
if field in UpdateKeyRequest.model_fields
|
||||
and field != "key"
|
||||
and not (field in _EMPTY_DURATION_MEANS_UNCHANGED and value == "")
|
||||
}
|
||||
)
|
||||
if not changed_fields:
|
||||
return None
|
||||
return UpdateKeyRequest(key=key, **changed_fields)
|
||||
|
||||
|
||||
class _LegacyDumpable(Protocol):
|
||||
def dict(self) -> Mapping[str, object]: ...
|
||||
|
||||
|
|
@ -987,6 +1149,7 @@ async def _common_key_generation_helper(
|
|||
litellm_changed_by: str | None,
|
||||
team_table: LiteLLM_TeamTableCachedObj | None,
|
||||
) -> GenerateKeyResponse:
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy.proxy_server import (
|
||||
litellm_proxy_admin_name,
|
||||
llm_router,
|
||||
|
|
@ -1135,6 +1298,16 @@ async def _common_key_generation_helper(
|
|||
"litellm.proxy.proxy_server.generate_key_fn(): Enterprise key management params not applied - %s", e
|
||||
)
|
||||
|
||||
await _enforce_custom_key_policy(
|
||||
hook=_custom_key_policy_hook(proxy_server),
|
||||
build_policy_request=lambda: CustomKeyPolicyRequest(
|
||||
operation="generate",
|
||||
existing_key=None,
|
||||
effective_key=_effective_key_for_generate(data=data, now=datetime.now(timezone.utc)),
|
||||
request=data,
|
||||
),
|
||||
)
|
||||
|
||||
# TODO: @ishaan-jaff: Migrate all budget tracking to use LiteLLM_BudgetTable
|
||||
_budget_id = data.budget_id
|
||||
if prisma_client is not None and data.soft_budget is not None:
|
||||
|
|
@ -1330,7 +1503,7 @@ async def _common_key_generation_helper(
|
|||
prisma_client=prisma_client,
|
||||
)
|
||||
|
||||
response = await generate_key_helper_fn(request_type="key", **data_json, table_name="key")
|
||||
response = await generate_key_helper_fn(request_type="key", **data_json, table_name="key", llm_router=llm_router)
|
||||
|
||||
response["soft_budget"] = data.soft_budget # include the user-input soft budget in the response
|
||||
|
||||
|
|
@ -2234,7 +2407,26 @@ async def _update_key_row_with_soft_budget(
|
|||
async def prepare_key_update_data(
|
||||
data: UpdateKeyRequest | RegenerateKeyRequest,
|
||||
existing_key_row: LiteLLM_VerificationToken,
|
||||
*,
|
||||
prisma_client: PrismaClient | None = None,
|
||||
llm_router: Router | None = None,
|
||||
):
|
||||
if data.router_settings is not None or (
|
||||
"router_settings" not in data.model_fields_set
|
||||
and "team_id" in data.model_fields_set
|
||||
and data.team_id != existing_key_row.team_id
|
||||
):
|
||||
effective_settings: Final = (
|
||||
data.router_settings
|
||||
if data.router_settings is not None
|
||||
else _StoredKeyRouterSettings.model_validate(existing_key_row, from_attributes=True).router_settings
|
||||
)
|
||||
await validate_router_settings_weights(
|
||||
effective_settings,
|
||||
team_id=data.team_id if "team_id" in data.model_fields_set else existing_key_row.team_id,
|
||||
prisma_client=prisma_client,
|
||||
llm_router=llm_router,
|
||||
)
|
||||
data_json: Final[dict] = data.model_dump(exclude_unset=True)
|
||||
data_json.pop("key", None)
|
||||
data_json.pop("new_key", None)
|
||||
|
|
@ -2301,12 +2493,6 @@ async def prepare_key_update_data(
|
|||
# sentinel for Json? columns, so store the JSON literal null
|
||||
non_default_values["budget_limits"] = json.dumps(None)
|
||||
|
||||
if "object_permission" in non_default_values:
|
||||
non_default_values = await _handle_update_object_permission(
|
||||
data_json=non_default_values,
|
||||
existing_key_row=existing_key_row,
|
||||
)
|
||||
|
||||
_metadata: Final = existing_key_row.metadata or {}
|
||||
|
||||
# validate model_max_budget
|
||||
|
|
@ -2327,13 +2513,12 @@ async def prepare_key_update_data(
|
|||
async def _handle_update_object_permission(
|
||||
data_json: dict,
|
||||
existing_key_row: LiteLLM_VerificationToken,
|
||||
prisma_client: PrismaClient,
|
||||
) -> dict:
|
||||
"""
|
||||
Handle the update of object permission.
|
||||
"""
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
"""Persist the requested object permission row and swap it for its id, only after the key policy allowed the write."""
|
||||
if "object_permission" not in data_json:
|
||||
return data_json
|
||||
|
||||
# Use the common helper to handle the object permission update
|
||||
object_permission_id: Final = await handle_update_object_permission_common(
|
||||
data_json=data_json,
|
||||
existing_object_permission_id=existing_key_row.object_permission_id,
|
||||
|
|
@ -2467,6 +2652,7 @@ async def _process_single_key_update(
|
|||
llm_router: Router | None,
|
||||
user_custom_key_update: Callable | None = None,
|
||||
existing_key_row: LiteLLM_VerificationToken | None = None,
|
||||
user_custom_key_policy: Callable[..., Awaitable[Mapping[str, object]]] | None = None,
|
||||
) -> dict[str, object]:
|
||||
"""
|
||||
Process a single key update with all validations and checks.
|
||||
|
|
@ -2575,7 +2761,19 @@ async def _process_single_key_update(
|
|||
)
|
||||
|
||||
# Prepare update data
|
||||
non_default_values = await prepare_key_update_data(data=update_key_request, existing_key_row=existing_key_row)
|
||||
non_default_values = await prepare_key_update_data(
|
||||
data=update_key_request, existing_key_row=existing_key_row, prisma_client=prisma_client, llm_router=llm_router
|
||||
)
|
||||
|
||||
await _enforce_custom_key_policy(
|
||||
hook=user_custom_key_policy,
|
||||
build_policy_request=lambda: _update_policy_request(
|
||||
operation="update",
|
||||
existing_key_row=existing_key_row,
|
||||
non_default_values=non_default_values,
|
||||
request=update_key_request,
|
||||
),
|
||||
)
|
||||
|
||||
# Update key in database
|
||||
if prisma_client is None:
|
||||
|
|
@ -2584,7 +2782,12 @@ async def _process_single_key_update(
|
|||
detail={"error": "Database not connected"},
|
||||
)
|
||||
|
||||
_data: Final = {**non_default_values, "token": update_key_request.key}
|
||||
update_values: Final = await _handle_update_object_permission(
|
||||
data_json=non_default_values,
|
||||
existing_key_row=existing_key_row,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
_data: Final = {**update_values, "token": update_key_request.key}
|
||||
response: Final[Mapping[str, object] | None] = cast( # cast-ok: every update_data branch returns a str-keyed dict
|
||||
"Mapping[str, object] | None",
|
||||
await prisma_client.update_data(token=update_key_request.key, data=_data),
|
||||
|
|
@ -3077,23 +3280,13 @@ async def update_key_fn(
|
|||
user_api_key_cache=user_api_key_cache,
|
||||
)
|
||||
|
||||
# Custom key update hook
|
||||
custom_key_update_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = _custom_key_update_hook(
|
||||
proxy_server
|
||||
)
|
||||
if custom_key_update_hook is not None:
|
||||
if inspect.iscoroutinefunction(custom_key_update_hook):
|
||||
result: Final = await custom_key_update_hook(data)
|
||||
else:
|
||||
raise ValueError("user_custom_key_update must be a coroutine")
|
||||
decision: Final = result.get("decision", True)
|
||||
message: Final = result.get("message", "Authentication Failed - Custom Auth Rule")
|
||||
if not decision:
|
||||
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail=message)
|
||||
await _enforce_custom_key_update_policy(hook=_custom_key_update_hook(proxy_server), data=data)
|
||||
|
||||
# Enforce upperbound key params on update (don't fill defaults)
|
||||
_enforce_upperbound_key_params(data, fill_defaults=False)
|
||||
non_default_values: Final = await prepare_key_update_data(data=data, existing_key_row=existing_key_row)
|
||||
non_default_values: Final = await prepare_key_update_data(
|
||||
data=data, existing_key_row=existing_key_row, prisma_client=prisma_client, llm_router=llm_router
|
||||
)
|
||||
|
||||
# Only validate key_alias format if it's actually being changed
|
||||
new_key_alias: Final = non_default_values.get("key_alias", None)
|
||||
|
|
@ -3114,21 +3307,36 @@ async def update_key_fn(
|
|||
existing_key_alias=existing_key_row.key_alias,
|
||||
)
|
||||
|
||||
await _enforce_custom_key_policy(
|
||||
hook=_custom_key_policy_hook(proxy_server),
|
||||
build_policy_request=lambda: _update_policy_request(
|
||||
operation="update",
|
||||
existing_key_row=existing_key_row,
|
||||
non_default_values=non_default_values,
|
||||
request=data,
|
||||
),
|
||||
)
|
||||
|
||||
if prisma_client is None:
|
||||
raise Exception("Not connected to DB!")
|
||||
|
||||
update_values: Final = await _handle_update_object_permission(
|
||||
data_json=non_default_values,
|
||||
existing_key_row=existing_key_row,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
changed_by: Final = user_api_key_dict.user_id or litellm_proxy_admin_name
|
||||
response: Final = (
|
||||
await _update_key_row_with_soft_budget(
|
||||
prisma_client=prisma_client,
|
||||
key=key,
|
||||
data=data,
|
||||
non_default_values=non_default_values,
|
||||
non_default_values=update_values,
|
||||
existing_key_row=existing_key_row,
|
||||
changed_by=changed_by,
|
||||
)
|
||||
if "soft_budget" in data.model_fields_set
|
||||
else await prisma_client.update_data(token=key, data=MappingProxyType({**non_default_values, "token": key}))
|
||||
else await prisma_client.update_data(token=key, data=MappingProxyType({**update_values, "token": key}))
|
||||
)
|
||||
|
||||
# Delete - key from cache, since it's been updated!
|
||||
|
|
@ -3263,6 +3471,7 @@ async def bulk_update_keys(
|
|||
)
|
||||
|
||||
custom_key_update_hook: Final = _custom_key_update_hook(proxy_server)
|
||||
custom_key_policy_hook: Final = _custom_key_policy_hook(proxy_server)
|
||||
|
||||
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value:
|
||||
raise HTTPException(
|
||||
|
|
@ -3310,6 +3519,7 @@ async def bulk_update_keys(
|
|||
proxy_logging_obj=proxy_logging_obj,
|
||||
llm_router=llm_router,
|
||||
user_custom_key_update=custom_key_update_hook,
|
||||
user_custom_key_policy=custom_key_policy_hook,
|
||||
)
|
||||
|
||||
successful_updates.append(
|
||||
|
|
@ -3427,6 +3637,7 @@ async def bulk_update_team_keys(
|
|||
)
|
||||
|
||||
custom_key_update_hook: Final = _custom_key_update_hook(proxy_server)
|
||||
custom_key_policy_hook: Final = _custom_key_policy_hook(proxy_server)
|
||||
|
||||
if prisma_client is None:
|
||||
raise HTTPException(
|
||||
|
|
@ -3557,6 +3768,7 @@ async def bulk_update_team_keys(
|
|||
proxy_logging_obj=proxy_logging_obj,
|
||||
llm_router=llm_router,
|
||||
user_custom_key_update=custom_key_update_hook,
|
||||
user_custom_key_policy=custom_key_policy_hook,
|
||||
existing_key_row=existing_by_token[db_token],
|
||||
)
|
||||
|
||||
|
|
@ -4082,6 +4294,40 @@ def _check_model_access_group(models: list[str] | None, llm_router: Router | Non
|
|||
return True
|
||||
|
||||
|
||||
_NO_METADATA: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
|
||||
|
||||
def metadata_json_with_limits(
|
||||
metadata: Mapping[str, object] | None,
|
||||
*,
|
||||
model_rpm_limit: Mapping[str, object] | None,
|
||||
model_tpm_limit: Mapping[str, object] | None,
|
||||
mcp_rpm_limit: Mapping[str, int] | None,
|
||||
tag_rpm_limit: Mapping[str, int] | None,
|
||||
guardrails: Sequence[str] | None,
|
||||
policies: Sequence[str] | None,
|
||||
prompts: Sequence[str] | None,
|
||||
) -> str:
|
||||
"""Serialize the stored metadata blob with the per-model, MCP, tag, guardrail, policy and prompt settings folded in."""
|
||||
limits: Final = tuple(
|
||||
(name, value)
|
||||
for name, value in (
|
||||
("model_rpm_limit", model_rpm_limit),
|
||||
("model_tpm_limit", model_tpm_limit),
|
||||
("mcp_rpm_limit", mcp_rpm_limit),
|
||||
("tag_rpm_limit", tag_rpm_limit),
|
||||
("guardrails", guardrails),
|
||||
("policies", policies),
|
||||
("prompts", prompts),
|
||||
)
|
||||
if value is not None
|
||||
)
|
||||
if metadata is None and not limits:
|
||||
return json.dumps(None)
|
||||
merged: Final = {**(metadata or _NO_METADATA), **dict(limits)} # mutable-ok: encrypt_callback_vars takes a dict
|
||||
return json.dumps(encrypt_callback_vars(merged))
|
||||
|
||||
|
||||
async def generate_key_helper_fn(
|
||||
request_type: Literal["user", "key"], # identifies if this request is from /user/new or /key/generate
|
||||
duration: str | None = None,
|
||||
|
|
@ -4137,15 +4383,24 @@ async def generate_key_helper_fn(
|
|||
object_permission: LiteLLM_ObjectPermissionBase | None = None,
|
||||
auto_rotate: bool | None = None,
|
||||
rotation_interval: str | None = None,
|
||||
router_settings: dict | None = None,
|
||||
router_settings: dict[str, object] | None = None,
|
||||
access_group_ids: list[str] | None = None,
|
||||
budget_limits: list | None = None, # multiple concurrent budget windows
|
||||
*,
|
||||
llm_router: Router | None = None,
|
||||
):
|
||||
from litellm.proxy.proxy_server import premium_user, prisma_client
|
||||
|
||||
if prisma_client is None:
|
||||
raise Exception("Connect Proxy to database to generate keys - https://docs.litellm.ai/docs/proxy/virtual_keys ")
|
||||
|
||||
await validate_router_settings_weights(
|
||||
router_settings,
|
||||
team_id=team_id,
|
||||
prisma_client=prisma_client,
|
||||
llm_router=llm_router,
|
||||
)
|
||||
|
||||
if token is None:
|
||||
if key is not None:
|
||||
token = key
|
||||
|
|
@ -4184,31 +4439,16 @@ async def generate_key_helper_fn(
|
|||
permissions_json: Final = json.dumps(permissions)
|
||||
router_settings_json: Final = safe_dumps(router_settings) if router_settings is not None else safe_dumps({})
|
||||
|
||||
# Add model_rpm_limit and model_tpm_limit to metadata
|
||||
if model_rpm_limit is not None:
|
||||
metadata = metadata or {}
|
||||
metadata["model_rpm_limit"] = model_rpm_limit
|
||||
if model_tpm_limit is not None:
|
||||
metadata = metadata or {}
|
||||
metadata["model_tpm_limit"] = model_tpm_limit
|
||||
if mcp_rpm_limit is not None:
|
||||
metadata = metadata or {}
|
||||
metadata["mcp_rpm_limit"] = mcp_rpm_limit
|
||||
if tag_rpm_limit is not None:
|
||||
metadata = metadata or {}
|
||||
metadata["tag_rpm_limit"] = tag_rpm_limit
|
||||
if guardrails is not None:
|
||||
metadata = metadata or {}
|
||||
metadata["guardrails"] = guardrails
|
||||
if policies is not None:
|
||||
metadata = metadata or {}
|
||||
metadata["policies"] = policies
|
||||
if prompts is not None:
|
||||
metadata = metadata or {}
|
||||
metadata["prompts"] = prompts
|
||||
|
||||
metadata = encrypt_callback_vars(metadata)
|
||||
metadata_json: Final = json.dumps(metadata)
|
||||
metadata_json: Final = metadata_json_with_limits(
|
||||
metadata,
|
||||
model_rpm_limit=model_rpm_limit,
|
||||
model_tpm_limit=model_tpm_limit,
|
||||
mcp_rpm_limit=mcp_rpm_limit,
|
||||
tag_rpm_limit=tag_rpm_limit,
|
||||
guardrails=guardrails,
|
||||
policies=policies,
|
||||
prompts=prompts,
|
||||
)
|
||||
validate_model_max_budget(model_max_budget)
|
||||
model_max_budget_json: Final = json.dumps(model_max_budget)
|
||||
budget_fallbacks_json: Final = json.dumps(budget_fallbacks or {})
|
||||
|
|
@ -5070,6 +5310,7 @@ async def _insert_deprecated_key(
|
|||
async def _execute_virtual_key_regeneration(
|
||||
*,
|
||||
prisma_client: PrismaClient,
|
||||
llm_router: Router | None = None,
|
||||
key_in_db: LiteLLM_VerificationToken,
|
||||
hashed_api_key: str,
|
||||
key: str,
|
||||
|
|
@ -5080,6 +5321,7 @@ async def _execute_virtual_key_regeneration(
|
|||
proxy_logging_obj: ProxyLogging,
|
||||
) -> GenerateKeyResponse:
|
||||
"""Generate new token, update DB, invalidate cache, and return response."""
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy.proxy_server import hash_token
|
||||
|
||||
# Mirror the /key/update ownership rebind guard. See helper docstring.
|
||||
|
|
@ -5127,15 +5369,34 @@ async def _execute_virtual_key_regeneration(
|
|||
|
||||
non_default_values = {}
|
||||
if data is not None:
|
||||
update_request: Final = _regenerate_request_as_update_request(key=hashed_api_key, data=data)
|
||||
if update_request is not None:
|
||||
await _enforce_custom_key_update_policy(hook=_custom_key_update_hook(proxy_server), data=update_request)
|
||||
# Enforce upperbound key params on regenerate (don't fill defaults)
|
||||
_enforce_upperbound_key_params(data, fill_defaults=False)
|
||||
non_default_values = await prepare_key_update_data(data=data, existing_key_row=key_in_db)
|
||||
non_default_values = await prepare_key_update_data(
|
||||
data=data, existing_key_row=key_in_db, prisma_client=prisma_client, llm_router=llm_router
|
||||
)
|
||||
# Only validate key_alias format if it's actually being changed
|
||||
new_key_alias: Final = non_default_values.get("key_alias")
|
||||
if new_key_alias != key_in_db.key_alias:
|
||||
_validate_key_alias_format(key_alias=new_key_alias)
|
||||
verbose_proxy_logger.debug("non_default_values: %s", non_default_values)
|
||||
update_data.update(non_default_values)
|
||||
await _enforce_custom_key_policy(
|
||||
hook=_custom_key_policy_hook(proxy_server),
|
||||
build_policy_request=lambda: _update_policy_request(
|
||||
operation="regenerate",
|
||||
existing_key_row=key_in_db,
|
||||
non_default_values=non_default_values,
|
||||
request=data if data is not None else RegenerateKeyRequest(),
|
||||
),
|
||||
)
|
||||
update_values: Final = await _handle_update_object_permission(
|
||||
data_json=non_default_values,
|
||||
existing_key_row=key_in_db,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
update_data.update(update_values)
|
||||
jsonified_update_data: Final[Mapping[str, object]] = prisma_client.jsonify_object(data=update_data)
|
||||
|
||||
# Snapshot before the token update: the FK cascade rewrites mapping rows to the new hash,
|
||||
|
|
@ -5145,6 +5406,13 @@ async def _execute_virtual_key_regeneration(
|
|||
prisma_client=prisma_client,
|
||||
)
|
||||
|
||||
await _persist_deleted_verification_tokens(
|
||||
keys=[key_in_db],
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_changed_by=litellm_changed_by,
|
||||
)
|
||||
|
||||
# If grace period set, insert deprecated key so old key remains valid
|
||||
await _insert_deprecated_key(
|
||||
prisma_client=prisma_client,
|
||||
|
|
@ -5268,6 +5536,7 @@ async def regenerate_key_fn(
|
|||
try:
|
||||
from litellm.proxy.proxy_server import (
|
||||
hash_token,
|
||||
llm_router,
|
||||
master_key,
|
||||
premium_user,
|
||||
prisma_client,
|
||||
|
|
@ -5443,19 +5712,9 @@ async def regenerate_key_fn(
|
|||
if litellm_changed_by is not None and not isinstance(litellm_changed_by, str):
|
||||
litellm_changed_by = None
|
||||
|
||||
# Save the old key record to deleted table before regeneration.
|
||||
# This preserves key_alias and team_id metadata for historical spend records.
|
||||
# If this fails, abort the regeneration to avoid permanently losing the
|
||||
# old hash→metadata mapping.
|
||||
await _persist_deleted_verification_tokens(
|
||||
keys=[_key_in_db],
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_changed_by=litellm_changed_by,
|
||||
)
|
||||
|
||||
return await _execute_virtual_key_regeneration(
|
||||
prisma_client=prisma_client,
|
||||
llm_router=llm_router,
|
||||
key_in_db=_key_in_db,
|
||||
hashed_api_key=hashed_api_key,
|
||||
key=key,
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
"""`POST /management/v1/users/bulk_delete`."""
|
||||
"""`POST /management/v1/users/bulk` and `POST /management/v1/users/bulk_delete`."""
|
||||
|
||||
from typing import Annotated, Final
|
||||
|
||||
|
|
@ -9,19 +9,105 @@ from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth
|
|||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.list_api.common import PROBLEM_TYPE_BASE, ManagementProblem, reject_unknown_query_params
|
||||
from litellm.proxy.management_endpoints.management_v1.common import MANAGEMENT_V1_PREFIX
|
||||
from litellm.proxy.management_helpers.bulk_user_creation import bulk_create_users
|
||||
from litellm.proxy.management_helpers.bulk_user_deletion import bulk_delete_users
|
||||
from litellm.proxy.management_helpers.utils import (
|
||||
management_endpoint_wrapper, # pyright: ignore[reportUnknownVariableType] # legacy decorator is untyped
|
||||
management_endpoint_wrapper, # pyright: ignore[reportUnknownVariableType] # legacy untyped decorator
|
||||
)
|
||||
from litellm.types.proxy.management_endpoints.internal_user_endpoints import (
|
||||
BulkDeleteUserRequest,
|
||||
BulkDeleteUsersResponse,
|
||||
BulkNewUserRequest,
|
||||
BulkNewUserResponse,
|
||||
)
|
||||
from litellm.types.proxy.management_endpoints.management_v1 import ProblemDetail
|
||||
|
||||
router: Final = APIRouter(prefix=MANAGEMENT_V1_PREFIX)
|
||||
|
||||
|
||||
@router.post(
|
||||
"/users/bulk",
|
||||
tags=["Internal User management"], # mutable-ok: fastapi types tags as list[str | Enum]
|
||||
dependencies=(Depends(user_api_key_auth),),
|
||||
response_model=BulkNewUserResponse,
|
||||
)
|
||||
@management_endpoint_wrapper
|
||||
async def bulk_create_users_route(
|
||||
data: BulkNewUserRequest,
|
||||
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
|
||||
) -> BulkNewUserResponse:
|
||||
"""
|
||||
Create up to 500 internal users in one request, optionally adding each one to teams.
|
||||
|
||||
Every entry in `users` takes the same fields as `/user/new`, with two differences: `auto_create_key`
|
||||
defaults to `false` (opt in per user to also get a virtual key back) and `send_invite_email` is not
|
||||
supported. Unknown fields are rejected with 422. Rows are validated together (duplicate ids or emails,
|
||||
unknown teams, roles the caller may not grant), inserted in one statement, and each referenced team is
|
||||
written once for all of its new members.
|
||||
|
||||
Rows fail independently: a bad row is reported in `data` with `success: false` and an `error`, and the
|
||||
other rows still get created. A user that was created but could not be added to one of its teams is
|
||||
reported with `success: true`, `teams` listing where they did land, and `error` naming the failed team.
|
||||
The whole request is refused with a 403 problem document only if creating the valid rows would exceed
|
||||
the license seat limit.
|
||||
|
||||
Example curl:
|
||||
```
|
||||
curl -X POST "http://localhost:4000/management/v1/users/bulk" \\
|
||||
-H "Content-Type: application/json" \\
|
||||
-H "Authorization: Bearer sk-1234" \\
|
||||
-d '{
|
||||
"users": [
|
||||
{"user_email": "a@example.com", "user_role": "internal_user", "teams": ["team-1"]},
|
||||
{"user_email": "b@example.com", "user_role": "internal_user", "auto_create_key": true}
|
||||
]
|
||||
}'
|
||||
```
|
||||
|
||||
Returns `data` (one entry per input row, in order, with `user_id`, `user_email`, `success`, `teams`,
|
||||
`key`, `error`) and `meta` with `total_requested`, `created` and `failed`.
|
||||
"""
|
||||
try:
|
||||
from litellm.proxy.proxy_server import (
|
||||
_license_check, # pyright: ignore[reportPrivateUsage] # same proxy license singleton /user/new reads
|
||||
litellm_proxy_admin_name,
|
||||
prisma_client,
|
||||
user_api_key_cache,
|
||||
)
|
||||
|
||||
if prisma_client is None:
|
||||
raise ManagementProblem(
|
||||
ProblemDetail(
|
||||
type=f"{PROBLEM_TYPE_BASE}database-not-connected",
|
||||
title="Database not connected",
|
||||
status=503,
|
||||
detail=CommonProxyErrors.db_not_connected_error.value,
|
||||
)
|
||||
)
|
||||
|
||||
return await bulk_create_users(
|
||||
users=data.users,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
prisma_client=prisma_client,
|
||||
license_check=_license_check,
|
||||
litellm_proxy_admin_name=litellm_proxy_admin_name,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
)
|
||||
|
||||
except ManagementProblem:
|
||||
raise
|
||||
except Exception: # noqa: BLE001 # a driver error answers as a problem document, not the OpenAI error shape
|
||||
verbose_proxy_logger.exception("/management/v1/users/bulk: Exception occurred")
|
||||
raise ManagementProblem(
|
||||
ProblemDetail(
|
||||
type=f"{PROBLEM_TYPE_BASE}internal-server-error",
|
||||
title="Internal server error",
|
||||
status=500,
|
||||
detail="Failed to create users.",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@router.post(
|
||||
"/users/bulk_delete",
|
||||
tags=["Internal User management"], # mutable-ok: FastAPI types `tags` as list[str], not Sequence
|
||||
|
|
|
|||
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,
|
||||
)
|
||||
129
litellm/proxy/management_endpoints/router_weights.py
Normal file
129
litellm/proxy/management_endpoints/router_weights.py
Normal file
|
|
@ -0,0 +1,129 @@
|
|||
from abc import abstractmethod
|
||||
from collections.abc import Mapping
|
||||
from typing import Annotated, Final, Protocol
|
||||
|
||||
from fastapi import HTTPException
|
||||
from pydantic import BaseModel, BeforeValidator, ValidationError
|
||||
|
||||
from litellm.repositories.prisma_protocols import TableActions
|
||||
from litellm.types.router_weights import RouterWeights
|
||||
|
||||
|
||||
class _StoredModel(Protocol):
|
||||
@property
|
||||
@abstractmethod
|
||||
def model_id(self) -> str:
|
||||
pass
|
||||
|
||||
|
||||
class _ModelDb(Protocol):
|
||||
@property
|
||||
@abstractmethod
|
||||
def litellm_proxymodeltable(self) -> TableActions[_StoredModel]:
|
||||
pass
|
||||
|
||||
|
||||
class _PrismaClient(Protocol):
|
||||
@property
|
||||
@abstractmethod
|
||||
def db(self) -> _ModelDb:
|
||||
pass
|
||||
|
||||
|
||||
class _Router(Protocol):
|
||||
@abstractmethod
|
||||
def get_deployment(self, model_id: str) -> object | None:
|
||||
pass
|
||||
|
||||
|
||||
class _RouterWeightSettings(BaseModel):
|
||||
weights: RouterWeights | None = None
|
||||
|
||||
|
||||
class _RouterWeightModelInfo(BaseModel):
|
||||
team_id: str | None = None
|
||||
db_model: bool | None = None
|
||||
team_public_model_name: str | None = None
|
||||
|
||||
|
||||
def _router_weight_model_info(value: object) -> _RouterWeightModelInfo:
|
||||
if isinstance(value, str):
|
||||
return _RouterWeightModelInfo.model_validate_json(value)
|
||||
return _RouterWeightModelInfo.model_validate(value or {}, from_attributes=True)
|
||||
|
||||
|
||||
class _RouterWeightDeployment(BaseModel):
|
||||
model_name: str
|
||||
model_info: Annotated[_RouterWeightModelInfo, BeforeValidator(_router_weight_model_info)]
|
||||
|
||||
|
||||
def _validate_router_weight_reference(
|
||||
model_group: str,
|
||||
deployment_id: str,
|
||||
team_id: str | None,
|
||||
stored: _RouterWeightDeployment | None,
|
||||
configured: object | None,
|
||||
) -> None:
|
||||
reference: Final = (
|
||||
stored
|
||||
if stored is not None
|
||||
else (
|
||||
_RouterWeightDeployment.model_validate(configured, from_attributes=True) if configured is not None else None
|
||||
)
|
||||
)
|
||||
if (
|
||||
reference is None
|
||||
or (stored is None and reference.model_info.db_model)
|
||||
or (reference.model_info.team_id is not None and reference.model_info.team_id != team_id)
|
||||
):
|
||||
raise HTTPException(status_code=400, detail=f"Unknown deployment ID in router weights: {deployment_id}")
|
||||
canonical_group: Final = (
|
||||
reference.model_info.team_public_model_name if reference.model_info.team_id is not None else None
|
||||
) or reference.model_name
|
||||
if model_group != canonical_group:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=f"Deployment {deployment_id} does not belong to model group {model_group}",
|
||||
)
|
||||
|
||||
|
||||
async def validate_router_settings_weights(
|
||||
router_settings: BaseModel | Mapping[str, object] | None,
|
||||
*,
|
||||
team_id: str | None,
|
||||
prisma_client: _PrismaClient | None,
|
||||
llm_router: _Router | None,
|
||||
) -> None:
|
||||
try:
|
||||
weights: Final = (
|
||||
_RouterWeightSettings.model_validate(router_settings, from_attributes=True).weights
|
||||
if router_settings is not None
|
||||
else None
|
||||
)
|
||||
except ValidationError:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="Invalid router weights. Replace or clear router_settings.weights.",
|
||||
) from None
|
||||
if not weights:
|
||||
return
|
||||
deployment_ids: Final = frozenset(deployment_id for group in weights.values() for deployment_id in group)
|
||||
if not deployment_ids:
|
||||
return
|
||||
if prisma_client is None:
|
||||
raise HTTPException(status_code=503, detail="Database unavailable while validating router weights")
|
||||
stored_models: Final = await prisma_client.db.litellm_proxymodeltable.find_many(
|
||||
where={"model_id": {"in": list(deployment_ids)}}
|
||||
)
|
||||
stored_by_id: Final = {
|
||||
row.model_id: _RouterWeightDeployment.model_validate(row, from_attributes=True) for row in stored_models
|
||||
}
|
||||
for model_group, group_weights in weights.items():
|
||||
for deployment_id in group_weights:
|
||||
_validate_router_weight_reference(
|
||||
model_group,
|
||||
deployment_id,
|
||||
team_id,
|
||||
stored_by_id.get(deployment_id),
|
||||
llm_router.get_deployment(model_id=deployment_id) if llm_router is not None else None,
|
||||
)
|
||||
|
|
@ -112,6 +112,7 @@ from litellm.proxy.management_endpoints.common_utils import (
|
|||
from litellm.proxy.management_endpoints.organization_endpoints import (
|
||||
add_member_to_organization,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.router_weights import validate_router_settings_weights
|
||||
from litellm.proxy.management_endpoints.tag_management_endpoints import (
|
||||
get_daily_activity,
|
||||
)
|
||||
|
|
@ -431,27 +432,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(
|
||||
|
|
@ -1289,6 +1289,7 @@ async def new_team(
|
|||
create_audit_log_for_update,
|
||||
general_settings,
|
||||
litellm_proxy_admin_name,
|
||||
llm_router,
|
||||
prisma_client,
|
||||
user_api_key_cache,
|
||||
)
|
||||
|
|
@ -1463,6 +1464,13 @@ async def new_team(
|
|||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
await validate_router_settings_weights(
|
||||
data.router_settings,
|
||||
team_id=data.team_id,
|
||||
prisma_client=prisma_client,
|
||||
llm_router=llm_router,
|
||||
)
|
||||
|
||||
## ADD TO MODEL TABLE
|
||||
_model_id = None
|
||||
if data.model_aliases is not None and isinstance(data.model_aliases, dict):
|
||||
|
|
@ -2076,6 +2084,13 @@ async def update_team(
|
|||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
await validate_router_settings_weights(
|
||||
data.router_settings,
|
||||
team_id=data.team_id,
|
||||
prisma_client=prisma_client,
|
||||
llm_router=llm_router,
|
||||
)
|
||||
|
||||
_existing_team_metadata: Final[object] = getattr(existing_team_row, "metadata", None)
|
||||
enforce_output_token_estimates_are_admin_only(
|
||||
data=data,
|
||||
|
|
@ -4370,6 +4385,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:
|
||||
|
|
@ -4441,7 +4470,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
|
||||
|
|
@ -4450,9 +4483,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 ##
|
||||
|
|
@ -4512,7 +4548,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(
|
||||
|
|
|
|||
|
|
@ -3592,6 +3592,7 @@ class SSOAuthenticationHandler:
|
|||
verbose_proxy_logger.info("user_defined_values for creating ui key: %s", user_defined_values)
|
||||
|
||||
response: Final = await generate_key_helper_fn(
|
||||
llm_router=None,
|
||||
request_type="key",
|
||||
duration=LITELLM_UI_SESSION_DURATION,
|
||||
key_max_budget=litellm.max_ui_session_budget,
|
||||
|
|
|
|||
871
litellm/proxy/management_helpers/bulk_user_creation.py
Normal file
871
litellm/proxy/management_helpers/bulk_user_creation.py
Normal file
|
|
@ -0,0 +1,871 @@
|
|||
"""Batched internal user creation behind `POST /management/v1/users/bulk`.
|
||||
|
||||
The batch is validated with set queries, user rows land in one `create_many`, and every
|
||||
referenced team is written once under its advisory lock instead of once per user.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from collections.abc import Awaitable, Callable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, Literal, TypeAlias, TypeVar
|
||||
|
||||
from fastapi import HTTPException, Request
|
||||
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.integrations.prometheus import PrometheusLogger
|
||||
from litellm.proxy._types import (
|
||||
LiteLLM_TeamTable,
|
||||
LitellmUserRoles,
|
||||
Member,
|
||||
NewUserRequestTeam,
|
||||
OrganizationMemberAddRequest,
|
||||
OrgMember,
|
||||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.auth.auth_checks import invalidate_team_member_spend_state
|
||||
from litellm.proxy.auth.litellm_license import LicenseCheck
|
||||
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
|
||||
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
|
||||
from litellm.proxy.hooks.user_management_event_hooks import UserManagementEventHooks
|
||||
from litellm.proxy.list_api.common import PROBLEM_TYPE_BASE, ManagementProblem
|
||||
from litellm.proxy.management_endpoints.common_utils import (
|
||||
_is_user_org_admin_for_team, # pyright: ignore[reportPrivateUsage] # same team-admin check /user/new uses
|
||||
_is_user_team_admin, # pyright: ignore[reportPrivateUsage] # same team-admin check /user/new uses
|
||||
validate_budget_duration,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.internal_user_endpoints import (
|
||||
_update_internal_new_user_params, # pyright: ignore[reportPrivateUsage, reportUnknownVariableType] # /user/new defaults; result validated below
|
||||
check_if_default_team_set,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
_check_permissions_caller_permission, # pyright: ignore[reportPrivateUsage] # same permission check /user/new uses
|
||||
generate_key_helper_fn, # pyright: ignore[reportUnknownVariableType] # legacy untyped helper; result validated by _KEY_RESPONSE
|
||||
metadata_json_with_limits,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.organization_endpoints import organization_member_add
|
||||
from litellm.proxy.management_helpers.access_group_team_sync import TEAM_ADVISORY_LOCK_SQL
|
||||
from litellm.proxy.management_helpers.object_permission_utils import (
|
||||
_set_object_permission, # pyright: ignore[reportPrivateUsage, reportUnknownVariableType] # shared with /user/new; result validated below
|
||||
)
|
||||
from litellm.proxy.management_helpers.utils import (
|
||||
_resolve_member_budget_id, # pyright: ignore[reportPrivateUsage] # shared with /team/member_add
|
||||
)
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
from litellm.repositories.prisma_protocols import TableActions
|
||||
from litellm.repositories.team_repository import TeamRepository
|
||||
from litellm.repositories.user_repository import UserRepository
|
||||
from litellm.types.proxy.management_endpoints.internal_user_endpoints import (
|
||||
BulkNewUserItem,
|
||||
BulkNewUserMeta,
|
||||
BulkNewUserResponse,
|
||||
UserCreateResult,
|
||||
)
|
||||
from litellm.types.proxy.management_endpoints.management_v1 import ProblemDetail
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from prisma import Prisma
|
||||
from prisma import models as prisma_models
|
||||
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
|
||||
BULK_NEW_USER_CONCURRENCY: Final = 10
|
||||
|
||||
TeamRole: TypeAlias = Literal["user", "admin"]
|
||||
KeyGenerator: TypeAlias = Callable[..., Awaitable[object]]
|
||||
_T: Final = TypeVar("_T")
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _RowFailure:
|
||||
index: int
|
||||
user_id: str | None
|
||||
user_email: str | None
|
||||
error: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _PendingUser:
|
||||
index: int
|
||||
request: BulkNewUserItem
|
||||
user_id: str
|
||||
teams: tuple[NewUserRequestTeam, ...]
|
||||
|
||||
|
||||
class _UserRow(BaseModel):
|
||||
"""The `/user/new` body after defaults and object permission were applied."""
|
||||
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
user_id: str
|
||||
user_email: str | None = None
|
||||
user_alias: str | None = None
|
||||
user_role: str | None = None
|
||||
team_id: str | None = None
|
||||
max_budget: float | None = None
|
||||
spend: float | None = 0.0
|
||||
models: tuple[str, ...] | None = None
|
||||
metadata: Mapping[str, object] | None = None
|
||||
max_parallel_requests: int | None = None
|
||||
tpm_limit: int | None = None
|
||||
rpm_limit: int | None = None
|
||||
budget_duration: str | None = None
|
||||
allowed_cache_controls: tuple[str, ...] | None = None
|
||||
sso_user_id: str | None = None
|
||||
object_permission_id: str | None = None
|
||||
model_max_budget: Mapping[str, object] | None = None
|
||||
model_rpm_limit: Mapping[str, object] | None = None
|
||||
model_tpm_limit: Mapping[str, object] | None = None
|
||||
mcp_rpm_limit: Mapping[str, int] | None = None
|
||||
tag_rpm_limit: Mapping[str, int] | None = None
|
||||
guardrails: tuple[str, ...] | None = None
|
||||
policies: tuple[str, ...] | None = None
|
||||
prompts: tuple[str, ...] | None = None
|
||||
duration: str | None = None
|
||||
key_alias: str | None = None
|
||||
aliases: Mapping[str, object] | None = None
|
||||
config: Mapping[str, object] | None = None
|
||||
permissions: Mapping[str, object] | None = None
|
||||
blocked: bool | None = None
|
||||
agent_id: str | None = None
|
||||
budget_fallbacks: Mapping[str, tuple[str, ...]] | None = None
|
||||
budget_limits: tuple[Mapping[str, object], ...] | None = None
|
||||
organizations: tuple[str, ...] | None = None
|
||||
|
||||
|
||||
_USER_ROW: Final = TypeAdapter(_UserRow)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _PreparedUser:
|
||||
pending: _PendingUser
|
||||
row: _UserRow
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _TeamAssignment:
|
||||
user_id: str
|
||||
user_email: str | None
|
||||
role: TeamRole
|
||||
max_budget_in_team: float | None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _TeamWrite:
|
||||
"""Outcome of one locked roster write. `failed` maps user ids to the reason they were not added."""
|
||||
|
||||
team_id: str
|
||||
after: tuple[Member, ...]
|
||||
added: frozenset[str]
|
||||
failed: Mapping[str, str]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _CreatedUser:
|
||||
prepared: _PreparedUser
|
||||
teams: tuple[str, ...]
|
||||
key: str | None
|
||||
errors: tuple[str, ...]
|
||||
|
||||
|
||||
_ERROR_DETAIL: Final = TypeAdapter(Mapping[str, object])
|
||||
_JSON_OBJECT: Final = TypeAdapter(dict[str, object])
|
||||
|
||||
|
||||
class _KeyResponse(BaseModel):
|
||||
token: str
|
||||
|
||||
|
||||
_KEY_RESPONSE: Final = TypeAdapter(_KeyResponse)
|
||||
|
||||
|
||||
def _error_message(exc: BaseException) -> str:
|
||||
if not isinstance(exc, HTTPException):
|
||||
return str(exc)
|
||||
try:
|
||||
detail: Final = _ERROR_DETAIL.validate_python(exc.detail)
|
||||
except ValidationError:
|
||||
return str(exc.detail)
|
||||
return str(detail.get("error", detail))
|
||||
|
||||
|
||||
def _requested_teams(item: BulkNewUserItem) -> tuple[NewUserRequestTeam, ...]:
|
||||
if item.team_id is not None:
|
||||
return (NewUserRequestTeam(team_id=item.team_id),)
|
||||
teams: Final = item.teams if item.teams is not None else check_if_default_team_set()
|
||||
if teams is None:
|
||||
return ()
|
||||
return tuple(team if isinstance(team, NewUserRequestTeam) else NewUserRequestTeam(team_id=team) for team in teams)
|
||||
|
||||
|
||||
def _row_error(item: BulkNewUserItem, user_api_key_dict: UserAPIKeyAuth) -> str | None:
|
||||
if (
|
||||
item.user_role in (LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY)
|
||||
and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN
|
||||
):
|
||||
return (
|
||||
"Only proxy admins can create administrative users (proxy_admin, proxy_admin_viewer). "
|
||||
f"Attempted to create user with role: {item.user_role}. Your role: {user_api_key_dict.user_role}"
|
||||
)
|
||||
try:
|
||||
validate_budget_duration(item.budget_duration)
|
||||
_check_permissions_caller_permission(data=item, user_api_key_dict=user_api_key_dict)
|
||||
except Exception as exc: # noqa: BLE001 # any validation failure is reported on this row only
|
||||
return _error_message(exc)
|
||||
return None
|
||||
|
||||
|
||||
def _normalized_email(email: str | None) -> str | None:
|
||||
return email.strip().lower() if email else None
|
||||
|
||||
|
||||
def _partition_rows(
|
||||
users: Sequence[BulkNewUserItem], user_api_key_dict: UserAPIKeyAuth
|
||||
) -> tuple[tuple[_PendingUser, ...], tuple[_RowFailure, ...]]:
|
||||
"""Assign ids, run the per-row checks and fail later rows that repeat an earlier row's id or email."""
|
||||
user_ids: Final = tuple(item.user_id or str(uuid.uuid4()) for item in users)
|
||||
first_index_by_id: Final = MappingProxyType(
|
||||
{user_id: index for index, user_id in reversed(tuple(enumerate(user_ids)))}
|
||||
)
|
||||
first_index_by_email: Final = MappingProxyType(
|
||||
{
|
||||
email: index
|
||||
for index, email in reversed(tuple(enumerate(_normalized_email(item.user_email) for item in users)))
|
||||
if email is not None
|
||||
}
|
||||
)
|
||||
|
||||
def classify(index: int, item: BulkNewUserItem) -> _PendingUser | _RowFailure:
|
||||
user_id: Final = user_ids[index]
|
||||
email: Final = _normalized_email(item.user_email)
|
||||
if first_index_by_id[user_id] != index:
|
||||
return _RowFailure(index, user_id, item.user_email, f"Duplicate user_id in request: {user_id}")
|
||||
if email is not None and first_index_by_email[email] != index:
|
||||
return _RowFailure(index, user_id, item.user_email, f"Duplicate user_email in request: {item.user_email}")
|
||||
error: Final = _row_error(item, user_api_key_dict)
|
||||
if error is not None:
|
||||
return _RowFailure(index, user_id, item.user_email, error)
|
||||
return _PendingUser(index, item, user_id, _requested_teams(item))
|
||||
|
||||
outcomes: Final = tuple(classify(index, item) for index, item in enumerate(users))
|
||||
return (
|
||||
tuple(outcome for outcome in outcomes if isinstance(outcome, _PendingUser)),
|
||||
tuple(outcome for outcome in outcomes if isinstance(outcome, _RowFailure)),
|
||||
)
|
||||
|
||||
|
||||
def _user_table(prisma_client: PrismaClient) -> "TableActions[prisma_models.LiteLLM_UserTable]":
|
||||
return UserRepository(prisma_client).table
|
||||
|
||||
|
||||
async def _existing_user_conflicts(
|
||||
prisma_client: PrismaClient, pending: Sequence[_PendingUser]
|
||||
) -> tuple[frozenset[str], frozenset[str]]:
|
||||
"""Return the requested user ids and (lowercased) emails that already exist, using one query each."""
|
||||
user_ids: Final = sorted(user.user_id for user in pending)
|
||||
emails: Final = sorted(frozenset(user.request.user_email for user in pending if user.request.user_email))
|
||||
if not user_ids:
|
||||
return frozenset(), frozenset()
|
||||
table: Final = _user_table(prisma_client)
|
||||
id_filter: Final = {"user_id": {"in": user_ids}} # mutable-ok: Prisma query filters are dict-shaped
|
||||
email_filter: Final = {"user_email": {"in": emails, "mode": "insensitive"}} # mutable-ok: Prisma filter
|
||||
id_rows: Final = await table.find_many(where=id_filter)
|
||||
email_rows: Final = await table.find_many(where=email_filter) if emails else ()
|
||||
return (
|
||||
frozenset(row.user_id for row in id_rows),
|
||||
frozenset(lowered for row in email_rows if (lowered := _normalized_email(row.user_email)) is not None),
|
||||
)
|
||||
|
||||
|
||||
async def _load_teams(prisma_client: PrismaClient, team_ids: frozenset[str]) -> Mapping[str, LiteLLM_TeamTable]:
|
||||
if not team_ids:
|
||||
return MappingProxyType({})
|
||||
rows: Final = await TeamRepository(prisma_client).table.find_many(
|
||||
where={"team_id": {"in": sorted(team_ids)}} # mutable-ok: Prisma query filters are dict-shaped
|
||||
)
|
||||
return MappingProxyType({row.team_id: LiteLLM_TeamTable.model_validate(row.model_dump()) for row in rows})
|
||||
|
||||
|
||||
async def _team_permission_error(team: LiteLLM_TeamTable, user_api_key_dict: UserAPIKeyAuth) -> str | None:
|
||||
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
|
||||
return None
|
||||
if _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team):
|
||||
return None
|
||||
if await _is_user_org_admin_for_team(user_api_key_dict=user_api_key_dict, team_obj=team):
|
||||
return None
|
||||
return f"Call not allowed. User not proxy admin OR team admin. team_id={team.team_id}"
|
||||
|
||||
|
||||
async def _unusable_teams(
|
||||
prisma_client: PrismaClient,
|
||||
pending: Sequence[_PendingUser],
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
) -> tuple[Mapping[str, LiteLLM_TeamTable], Mapping[str, str]]:
|
||||
"""Load every referenced team once and explain, per team id, why rows naming it cannot proceed."""
|
||||
team_ids: Final = frozenset(team.team_id for user in pending for team in user.teams)
|
||||
teams: Final = await _load_teams(prisma_client, team_ids)
|
||||
permission_errors: Final = await asyncio.gather(
|
||||
*(_team_permission_error(team, user_api_key_dict) for team in teams.values())
|
||||
)
|
||||
missing: Final = tuple(
|
||||
(team_id, f"Team id={team_id} does not exist") for team_id in team_ids if team_id not in teams
|
||||
)
|
||||
denied: Final = tuple(
|
||||
(team.team_id, error)
|
||||
for team, error in zip(teams.values(), permission_errors, strict=True)
|
||||
if error is not None
|
||||
)
|
||||
return teams, MappingProxyType({team_id: error for team_id, error in (*missing, *denied)})
|
||||
|
||||
|
||||
def _db_failure(
|
||||
user: _PendingUser,
|
||||
existing_ids: frozenset[str],
|
||||
existing_emails: frozenset[str],
|
||||
team_errors: Mapping[str, str],
|
||||
) -> _RowFailure | None:
|
||||
email: Final = _normalized_email(user.request.user_email)
|
||||
if user.user_id in existing_ids:
|
||||
return _RowFailure(user.index, user.user_id, user.request.user_email, f"User id={user.user_id} already exists")
|
||||
if email is not None and email in existing_emails:
|
||||
return _RowFailure(
|
||||
user.index, user.user_id, user.request.user_email, f"User email={user.request.user_email} already exists"
|
||||
)
|
||||
errors: Final = tuple(team_errors[team.team_id] for team in user.teams if team.team_id in team_errors)
|
||||
if errors:
|
||||
return _RowFailure(user.index, user.user_id, user.request.user_email, "; ".join(errors))
|
||||
return None
|
||||
|
||||
|
||||
async def _prepare_user(user: _PendingUser, prisma_client: PrismaClient) -> _PreparedUser | _RowFailure:
|
||||
try:
|
||||
dumped: Final = user.request.model_dump(exclude={"user_id"}) # mutable-ok: pydantic IncEx takes a set
|
||||
data: Final = {**dumped, "user_id": user.user_id} # mutable-ok: /user/new defaults helper mutates in place
|
||||
data_json: Final = _JSON_OBJECT.validate_python(_update_internal_new_user_params(data, user.request))
|
||||
with_permission: Final = _JSON_OBJECT.validate_python(
|
||||
await _set_object_permission(data_json=data_json, prisma_client=prisma_client) # pyright: ignore[reportUnknownArgumentType] # validated by the adapter
|
||||
)
|
||||
return _PreparedUser(user, _USER_ROW.validate_python(with_permission))
|
||||
except Exception as exc: # noqa: BLE001 # any preparation failure is reported on this row only
|
||||
verbose_proxy_logger.warning("/user/bulk_new: could not prepare row %d - %s", user.index, type(exc).__name__)
|
||||
return _RowFailure(user.index, user.user_id, user.request.user_email, _error_message(exc))
|
||||
|
||||
|
||||
class _UserCreateData(TypedDict):
|
||||
"""One `LiteLLM_UserTable` row as `create_many` takes it; JSON columns are pre-serialized."""
|
||||
|
||||
user_id: ReadOnly[str]
|
||||
user_email: ReadOnly[str | None]
|
||||
user_alias: ReadOnly[str | None]
|
||||
user_role: ReadOnly[str | None]
|
||||
team_id: ReadOnly[str | None]
|
||||
max_budget: ReadOnly[float | None]
|
||||
spend: ReadOnly[float]
|
||||
models: ReadOnly[tuple[str, ...]]
|
||||
metadata: ReadOnly[str]
|
||||
max_parallel_requests: ReadOnly[int | None]
|
||||
tpm_limit: ReadOnly[int | None]
|
||||
rpm_limit: ReadOnly[int | None]
|
||||
budget_duration: ReadOnly[str | None]
|
||||
budget_reset_at: ReadOnly[datetime | None]
|
||||
allowed_cache_controls: ReadOnly[tuple[str, ...]]
|
||||
sso_user_id: ReadOnly[str | None]
|
||||
object_permission_id: ReadOnly[str | None]
|
||||
teams: ReadOnly[tuple[str, ...]]
|
||||
model_max_budget: ReadOnly[str]
|
||||
|
||||
|
||||
def _user_create_payload(prepared: _PreparedUser) -> _UserCreateData:
|
||||
row: Final = prepared.row
|
||||
metadata_json: Final = metadata_json_with_limits(
|
||||
row.metadata,
|
||||
model_rpm_limit=row.model_rpm_limit,
|
||||
model_tpm_limit=row.model_tpm_limit,
|
||||
mcp_rpm_limit=row.mcp_rpm_limit,
|
||||
tag_rpm_limit=row.tag_rpm_limit,
|
||||
guardrails=row.guardrails,
|
||||
policies=row.policies,
|
||||
prompts=row.prompts,
|
||||
)
|
||||
payload: Final[_UserCreateData] = {
|
||||
"user_id": row.user_id,
|
||||
"user_email": row.user_email,
|
||||
"user_alias": row.user_alias,
|
||||
"user_role": row.user_role,
|
||||
"team_id": row.team_id,
|
||||
"max_budget": row.max_budget,
|
||||
"spend": row.spend or 0.0,
|
||||
"models": row.models or (),
|
||||
"metadata": metadata_json,
|
||||
"max_parallel_requests": row.max_parallel_requests,
|
||||
"tpm_limit": row.tpm_limit,
|
||||
"rpm_limit": row.rpm_limit,
|
||||
"budget_duration": row.budget_duration,
|
||||
"budget_reset_at": get_budget_reset_time(row.budget_duration) if row.budget_duration else None,
|
||||
"allowed_cache_controls": row.allowed_cache_controls or (),
|
||||
"sso_user_id": row.sso_user_id,
|
||||
"object_permission_id": row.object_permission_id,
|
||||
"teams": tuple(team.team_id for team in prepared.pending.teams),
|
||||
"model_max_budget": json.dumps(row.model_max_budget) if row.model_max_budget else "{}",
|
||||
}
|
||||
return payload
|
||||
|
||||
|
||||
async def _bounded(limit: int, awaitables: Sequence[Awaitable[_T]]) -> tuple[_T | BaseException, ...]:
|
||||
semaphore: Final = asyncio.Semaphore(limit)
|
||||
|
||||
async def run(awaitable: Awaitable[_T]) -> _T:
|
||||
async with semaphore:
|
||||
return await awaitable
|
||||
|
||||
return tuple(await asyncio.gather(*(run(awaitable) for awaitable in awaitables), return_exceptions=True))
|
||||
|
||||
|
||||
async def _insert_users(
|
||||
prisma_client: PrismaClient, prepared: Sequence[_PreparedUser]
|
||||
) -> tuple[tuple[_PreparedUser, ...], tuple[_RowFailure, ...]]:
|
||||
"""Insert every row in one statement. If that fails, retry rows one at a time so the error lands on its row."""
|
||||
if not prepared:
|
||||
return (), ()
|
||||
table: Final = _user_table(prisma_client)
|
||||
payloads: Final = tuple(_user_create_payload(user) for user in prepared)
|
||||
try:
|
||||
await table.create_many(data=payloads)
|
||||
return tuple(prepared), ()
|
||||
except Exception as exc: # noqa: BLE001 # fall back to per-row inserts so the failing row can be identified
|
||||
verbose_proxy_logger.warning("/user/bulk_new: create_many failed, retrying rows individually", exc_info=True)
|
||||
outcome_unknown: Final = PrismaDBExceptionHandler.is_database_infrastructure_error(exc)
|
||||
requested: Final = frozenset(payload["user_id"] for payload in payloads)
|
||||
landed_rows: Final = await table.find_many(where={"user_id": {"in": list(requested)}}) # mutable-ok: Prisma filter
|
||||
landed: Final = frozenset(row.user_id for row in landed_rows)
|
||||
# create_many is one INSERT: after a lost response the full set is ours, any partial set belongs to another request
|
||||
if outcome_unknown and landed == requested:
|
||||
return tuple(prepared), ()
|
||||
taken: Final = tuple(user for user in prepared if user.row.user_id in landed)
|
||||
retried: Final = tuple(user for user in prepared if user.row.user_id not in landed)
|
||||
outcomes: Final = await _bounded(
|
||||
BULK_NEW_USER_CONCURRENCY, tuple(table.create(data=_user_create_payload(user)) for user in retried)
|
||||
)
|
||||
failed: Final = MappingProxyType(
|
||||
{
|
||||
**{
|
||||
user.row.user_id: _RowFailure(
|
||||
user.pending.index,
|
||||
user.pending.user_id,
|
||||
user.row.user_email,
|
||||
f"User id={user.row.user_id} already exists",
|
||||
)
|
||||
for user in taken
|
||||
},
|
||||
**{
|
||||
user.row.user_id: _RowFailure(
|
||||
user.pending.index, user.pending.user_id, user.row.user_email, _error_message(outcome)
|
||||
)
|
||||
for user, outcome in zip(retried, outcomes, strict=True)
|
||||
if isinstance(outcome, BaseException)
|
||||
},
|
||||
}
|
||||
)
|
||||
return (
|
||||
tuple(user for user in prepared if user.row.user_id not in failed),
|
||||
tuple(failed.values()),
|
||||
)
|
||||
|
||||
|
||||
def _assignments_by_team(created: Sequence[_PreparedUser]) -> Mapping[str, tuple[_TeamAssignment, ...]]:
|
||||
team_ids: Final = tuple(dict.fromkeys(team.team_id for user in created for team in user.pending.teams))
|
||||
return MappingProxyType(
|
||||
{
|
||||
team_id: tuple(
|
||||
_TeamAssignment(user.pending.user_id, user.row.user_email, team.user_role, team.max_budget_in_team)
|
||||
for user in created
|
||||
for team in user.pending.teams
|
||||
if team.team_id == team_id
|
||||
)
|
||||
for team_id in team_ids
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
class _MembershipData(TypedDict):
|
||||
team_id: ReadOnly[str]
|
||||
user_id: ReadOnly[str]
|
||||
budget_id: ReadOnly[str | None]
|
||||
|
||||
|
||||
class _RosterData(TypedDict):
|
||||
members_with_roles: ReadOnly[str]
|
||||
|
||||
|
||||
class _TeamsData(TypedDict):
|
||||
teams: ReadOnly[tuple[str, ...]]
|
||||
|
||||
|
||||
def _default_member_budget_id(team: LiteLLM_TeamTable) -> str | None:
|
||||
metadata: Final = (
|
||||
_JSON_OBJECT.validate_python(
|
||||
team.metadata # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # LiteLLM_TeamTable.metadata is a bare dict; validated by the adapter
|
||||
)
|
||||
if team.metadata # pyright: ignore[reportUnknownMemberType] # same bare dict
|
||||
else None
|
||||
)
|
||||
budget_id: Final = metadata.get("team_member_budget_id") if metadata is not None else None
|
||||
return budget_id if isinstance(budget_id, str) else None
|
||||
|
||||
|
||||
def _team_tx_db(tx: "Prisma") -> "TableActions[prisma_models.LiteLLM_TeamTable]":
|
||||
return tx.litellm_teamtable # pyright: ignore[reportReturnType] # TableActions widens the generated inputs to Mapping, as the repositories do
|
||||
|
||||
|
||||
def _membership_tx_db(tx: "Prisma") -> "TableActions[prisma_models.LiteLLM_TeamMembership]":
|
||||
return tx.litellm_teammembership # pyright: ignore[reportReturnType] # TableActions widens the generated inputs to Mapping, as the repositories do
|
||||
|
||||
|
||||
async def _write_team_roster(
|
||||
prisma_client: PrismaClient,
|
||||
team: LiteLLM_TeamTable,
|
||||
members: Sequence[_TeamAssignment],
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
litellm_proxy_admin_name: str,
|
||||
) -> _TeamWrite:
|
||||
"""Add every new member to one team under its advisory lock: one roster rewrite and one membership insert."""
|
||||
try:
|
||||
async with prisma_client.tx() as tx:
|
||||
await tx.query_raw(TEAM_ADVISORY_LOCK_SQL, team.team_id)
|
||||
roster: Final = await TeamRepository(prisma_client).get_members_with_roles_locked(tx, team.team_id)
|
||||
if roster is None:
|
||||
raise ValueError(f"Team id={team.team_id} does not exist")
|
||||
already_present: Final = frozenset(member.user_id for member in roster if member.user_id)
|
||||
new_members: Final = tuple(member for member in members if member.user_id not in already_present)
|
||||
budget_ids: Final = tuple(
|
||||
[ # mutable-ok: budgets are created one at a time on the transaction's single connection
|
||||
await _resolve_member_budget_id(
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_proxy_admin_name=litellm_proxy_admin_name,
|
||||
max_budget_in_team=member.max_budget_in_team,
|
||||
allowed_models=team.default_team_member_models or None,
|
||||
budget_duration=None,
|
||||
default_team_budget_id=_default_member_budget_id(team),
|
||||
tx=tx, # pyright: ignore[reportArgumentType] # MemberWriteTx lags the generated Prisma signatures, same as /team/member_add
|
||||
)
|
||||
for member in new_members
|
||||
]
|
||||
)
|
||||
await _membership_tx_db(tx).create_many(
|
||||
data=tuple(
|
||||
_MembershipData(team_id=team.team_id, user_id=member.user_id, budget_id=budget_id)
|
||||
for member, budget_id in zip(new_members, budget_ids, strict=True)
|
||||
),
|
||||
skip_duplicates=True,
|
||||
)
|
||||
after: Final = (
|
||||
*roster,
|
||||
*(Member(user_id=m.user_id, user_email=m.user_email, role=m.role) for m in new_members),
|
||||
)
|
||||
await _team_tx_db(tx).update(
|
||||
where={"team_id": team.team_id}, # mutable-ok: Prisma query filters are dict-shaped
|
||||
data=_RosterData(members_with_roles=json.dumps(tuple(member.model_dump() for member in after))),
|
||||
)
|
||||
return _TeamWrite(
|
||||
team_id=team.team_id,
|
||||
after=after,
|
||||
added=frozenset(member.user_id for member in members),
|
||||
failed=MappingProxyType({}),
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001 # the team write failure is reported on each affected row
|
||||
verbose_proxy_logger.exception("/user/bulk_new: failed to add %d members to a team", len(members))
|
||||
message: Final = f"Failed to add user to team {team.team_id}: {_error_message(exc)}"
|
||||
return _TeamWrite(
|
||||
team_id=team.team_id,
|
||||
after=(),
|
||||
added=frozenset(),
|
||||
failed=MappingProxyType({member.user_id: message for member in members}),
|
||||
)
|
||||
|
||||
|
||||
async def _detach_failed_teams(
|
||||
prisma_client: PrismaClient, created: Sequence[_PreparedUser], writes: Mapping[str, _TeamWrite]
|
||||
) -> None:
|
||||
"""Users are inserted with `teams` already set; drop the teams whose roster write did not take them."""
|
||||
table: Final = _user_table(prisma_client)
|
||||
updates: Final = tuple(
|
||||
table.update(
|
||||
where={"user_id": user.row.user_id}, # mutable-ok: Prisma query filters are dict-shaped
|
||||
data=_TeamsData(teams=landed),
|
||||
)
|
||||
for user in created
|
||||
if (landed := _row_teams(user, writes)[0]) != tuple(team.team_id for team in user.pending.teams)
|
||||
)
|
||||
for outcome in await _bounded(BULK_NEW_USER_CONCURRENCY, updates):
|
||||
if isinstance(outcome, BaseException):
|
||||
verbose_proxy_logger.warning(
|
||||
"/user/bulk_new: could not detach failed teams from user - %s", type(outcome).__name__
|
||||
)
|
||||
|
||||
|
||||
async def _publish_team_writes(writes: Sequence[_TeamWrite], user_api_key_cache: "UserApiKeyCache") -> None:
|
||||
prometheus_logger: Final = PrometheusLogger.get_instance()
|
||||
for write in writes:
|
||||
if prometheus_logger is None or not write.added:
|
||||
continue
|
||||
try:
|
||||
prometheus_logger.set_team_members_metric(
|
||||
LiteLLM_TeamTable(
|
||||
team_id=write.team_id,
|
||||
members_with_roles=write.after, # pyright: ignore[reportArgumentType] # pydantic coerces the tuple into the declared list
|
||||
)
|
||||
)
|
||||
except Exception: # noqa: BLE001 # metrics are best-effort and must not fail the request
|
||||
verbose_proxy_logger.debug("Prometheus: failed to emit team members metric", exc_info=True)
|
||||
evictions: Final = await _bounded(
|
||||
BULK_NEW_USER_CONCURRENCY,
|
||||
tuple(
|
||||
invalidate_team_member_spend_state(
|
||||
user_id=user_id, team_id=write.team_id, user_api_key_cache=user_api_key_cache
|
||||
)
|
||||
for write in writes
|
||||
for user_id in write.added
|
||||
),
|
||||
)
|
||||
for eviction in evictions:
|
||||
if isinstance(eviction, BaseException):
|
||||
verbose_proxy_logger.warning("/user/bulk_new: cache eviction failed - %s", type(eviction).__name__)
|
||||
|
||||
|
||||
_KEY_FIELDS: Final = MappingProxyType(
|
||||
{
|
||||
name: True
|
||||
for name in (
|
||||
"user_id",
|
||||
"team_id",
|
||||
"agent_id",
|
||||
"duration",
|
||||
"key_alias",
|
||||
"models",
|
||||
"aliases",
|
||||
"config",
|
||||
"permissions",
|
||||
"blocked",
|
||||
"spend",
|
||||
"budget_fallbacks",
|
||||
"budget_limits",
|
||||
"metadata",
|
||||
"max_parallel_requests",
|
||||
"tpm_limit",
|
||||
"rpm_limit",
|
||||
"allowed_cache_controls",
|
||||
"model_max_budget",
|
||||
"model_rpm_limit",
|
||||
"model_tpm_limit",
|
||||
"mcp_rpm_limit",
|
||||
"tag_rpm_limit",
|
||||
"guardrails",
|
||||
"policies",
|
||||
"prompts",
|
||||
"object_permission_id",
|
||||
)
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
async def _generate_key(prepared: _PreparedUser, generate_key: KeyGenerator) -> str:
|
||||
response: Final = _KEY_RESPONSE.validate_python(
|
||||
await generate_key(
|
||||
request_type="key", table_name="key", **prepared.row.model_dump(include=_KEY_FIELDS, exclude_none=True)
|
||||
)
|
||||
)
|
||||
return response.token
|
||||
|
||||
|
||||
async def _add_to_organizations(
|
||||
prepared: _PreparedUser, organizations: Sequence[str], user_api_key_dict: UserAPIKeyAuth
|
||||
) -> None:
|
||||
for organization_id in organizations:
|
||||
await organization_member_add(
|
||||
data=OrganizationMemberAddRequest(
|
||||
organization_id=organization_id,
|
||||
member=OrgMember(user_id=prepared.row.user_id, role=LitellmUserRoles.INTERNAL_USER),
|
||||
),
|
||||
http_request=Request(scope={"type": "http", "path": "/user/bulk_new"}), # mutable-ok: ASGI scopes are dicts
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
|
||||
async def _run_per_user(
|
||||
created: Sequence[_PreparedUser],
|
||||
select: Callable[[_PreparedUser], bool],
|
||||
action: Callable[[_PreparedUser], Awaitable[_T]],
|
||||
) -> Mapping[str, _T | BaseException]:
|
||||
chosen: Final = tuple(user for user in created if select(user))
|
||||
outcomes: Final = await _bounded(BULK_NEW_USER_CONCURRENCY, tuple(action(user) for user in chosen))
|
||||
return MappingProxyType({user.row.user_id: outcome for user, outcome in zip(chosen, outcomes, strict=True)})
|
||||
|
||||
|
||||
async def _write_audit_logs(
|
||||
prisma_client: PrismaClient,
|
||||
created: Sequence[_PreparedUser],
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
litellm_proxy_admin_name: str,
|
||||
) -> None:
|
||||
if not created:
|
||||
return
|
||||
created_ids: Final = sorted(user.row.user_id for user in created)
|
||||
created_filter: Final = {"user_id": {"in": created_ids}} # mutable-ok: Prisma query filters are dict-shaped
|
||||
rows: Final = await _user_table(prisma_client).find_many(where=created_filter)
|
||||
outcomes: Final = await _bounded(
|
||||
BULK_NEW_USER_CONCURRENCY,
|
||||
tuple(
|
||||
UserManagementEventHooks.create_internal_user_audit_log(
|
||||
user_id=row.user_id,
|
||||
action="created",
|
||||
litellm_changed_by=user_api_key_dict.user_id,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_proxy_admin_name=litellm_proxy_admin_name,
|
||||
before_value=None,
|
||||
after_value=row.model_dump_json(exclude_none=True),
|
||||
)
|
||||
for row in rows
|
||||
),
|
||||
)
|
||||
for outcome in outcomes:
|
||||
if isinstance(outcome, BaseException):
|
||||
verbose_proxy_logger.warning(
|
||||
"Unable to create audit log for user on `/user/bulk_new` - %s", type(outcome).__name__
|
||||
)
|
||||
|
||||
|
||||
def _row_teams(prepared: _PreparedUser, writes: Mapping[str, _TeamWrite]) -> tuple[tuple[str, ...], tuple[str, ...]]:
|
||||
"""Split a user's requested teams into the ones they landed in and the errors for the ones they did not."""
|
||||
requested: Final = tuple(team.team_id for team in prepared.pending.teams)
|
||||
return (
|
||||
tuple(team_id for team_id in requested if prepared.row.user_id in writes[team_id].added),
|
||||
tuple(
|
||||
writes[team_id].failed[prepared.row.user_id]
|
||||
for team_id in requested
|
||||
if prepared.row.user_id in writes[team_id].failed
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _to_result(created: _CreatedUser) -> UserCreateResult:
|
||||
return UserCreateResult(
|
||||
user_id=created.prepared.row.user_id,
|
||||
user_email=created.prepared.row.user_email,
|
||||
success=True,
|
||||
teams=created.teams,
|
||||
key=created.key,
|
||||
error="; ".join(created.errors) if created.errors else None,
|
||||
)
|
||||
|
||||
|
||||
def _failure_result(failure: _RowFailure) -> UserCreateResult:
|
||||
return UserCreateResult(user_id=failure.user_id, user_email=failure.user_email, success=False, error=failure.error)
|
||||
|
||||
|
||||
async def bulk_create_users(
|
||||
users: Sequence[BulkNewUserItem],
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
prisma_client: PrismaClient,
|
||||
license_check: LicenseCheck,
|
||||
litellm_proxy_admin_name: str,
|
||||
user_api_key_cache: "UserApiKeyCache",
|
||||
generate_key: KeyGenerator = generate_key_helper_fn,
|
||||
) -> BulkNewUserResponse:
|
||||
"""Create every valid row in `users`; rows that fail validation or a write are reported, not raised.
|
||||
|
||||
Raises a 403 `ManagementProblem` only when the whole batch would push the deployment over its license seat
|
||||
limit.
|
||||
"""
|
||||
pending, request_failures = _partition_rows(users, user_api_key_dict)
|
||||
existing_ids, existing_emails = await _existing_user_conflicts(prisma_client, pending)
|
||||
teams, team_errors = await _unusable_teams(prisma_client, pending, user_api_key_dict)
|
||||
db_failures: Final = tuple(
|
||||
failure
|
||||
for user in pending
|
||||
if (failure := _db_failure(user, existing_ids, existing_emails, team_errors)) is not None
|
||||
)
|
||||
failed_indexes: Final = frozenset(failure.index for failure in db_failures)
|
||||
creatable: Final = tuple(user for user in pending if user.index not in failed_indexes)
|
||||
|
||||
billable_users: Final = await UserRepository(prisma_client).count_billable_users()
|
||||
if creatable and license_check.is_over_limit(total_users=billable_users + len(creatable)):
|
||||
raise ManagementProblem(
|
||||
ProblemDetail(
|
||||
type=f"{PROBLEM_TYPE_BASE}license-limit-exceeded",
|
||||
title="License limit exceeded",
|
||||
status=403,
|
||||
detail="License is over limit. Please contact support@berri.ai to upgrade your license.",
|
||||
)
|
||||
)
|
||||
|
||||
prepared_outcomes: Final = tuple([await _prepare_user(user, prisma_client) for user in creatable])
|
||||
prepare_failures: Final = tuple(o for o in prepared_outcomes if isinstance(o, _RowFailure))
|
||||
created, insert_failures = await _insert_users(
|
||||
prisma_client, tuple(o for o in prepared_outcomes if isinstance(o, _PreparedUser))
|
||||
)
|
||||
|
||||
team_writes: Final = MappingProxyType(
|
||||
{
|
||||
team_id: await _write_team_roster(
|
||||
prisma_client, teams[team_id], members, user_api_key_dict, litellm_proxy_admin_name
|
||||
)
|
||||
for team_id, members in _assignments_by_team(created).items()
|
||||
}
|
||||
)
|
||||
await _detach_failed_teams(prisma_client, created, team_writes)
|
||||
await _publish_team_writes(tuple(team_writes.values()), user_api_key_cache)
|
||||
|
||||
keys: Final = await _run_per_user(
|
||||
created, lambda user: user.pending.request.auto_create_key, lambda user: _generate_key(user, generate_key)
|
||||
)
|
||||
org_outcomes: Final = await _run_per_user(
|
||||
created,
|
||||
lambda user: bool(user.row.organizations),
|
||||
lambda user: _add_to_organizations(user, user.row.organizations or (), user_api_key_dict),
|
||||
)
|
||||
await _write_audit_logs(prisma_client, created, user_api_key_dict, litellm_proxy_admin_name)
|
||||
|
||||
def finish(prepared: _PreparedUser) -> _CreatedUser:
|
||||
landed, team_failures = _row_teams(prepared, team_writes)
|
||||
key_outcome: Final = keys.get(prepared.row.user_id)
|
||||
org_outcome: Final = org_outcomes.get(prepared.row.user_id)
|
||||
return _CreatedUser(
|
||||
prepared=prepared,
|
||||
teams=landed,
|
||||
key=key_outcome if isinstance(key_outcome, str) else None,
|
||||
errors=(
|
||||
*team_failures,
|
||||
*(
|
||||
(f"Failed to create key: {_error_message(key_outcome)}",)
|
||||
if isinstance(key_outcome, BaseException)
|
||||
else ()
|
||||
),
|
||||
*(
|
||||
(f"Failed to add user to organizations: {_error_message(org_outcome)}",)
|
||||
if isinstance(org_outcome, BaseException)
|
||||
else ()
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
failures: Final = MappingProxyType(
|
||||
{
|
||||
failure.index: _failure_result(failure)
|
||||
for failure in (*request_failures, *db_failures, *prepare_failures, *insert_failures)
|
||||
}
|
||||
)
|
||||
successes_by_index: Final = MappingProxyType({user.pending.index: _to_result(finish(user)) for user in created})
|
||||
results: Final = tuple(
|
||||
failures[index] if index in failures else successes_by_index[index] for index in range(len(users))
|
||||
)
|
||||
successes: Final = sum(1 for result in results if result.success)
|
||||
return BulkNewUserResponse(
|
||||
data=results,
|
||||
meta=BulkNewUserMeta(total_requested=len(users), created=successes, failed=len(users) - successes),
|
||||
)
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -5,18 +5,105 @@ Handles cost tracking and logging for Vertex AI Live API WebSocket passthrough e
|
|||
Supports different modalities: text, audio, video, and web search.
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping, Sequence
|
||||
from datetime import datetime
|
||||
from typing import Any, Final
|
||||
from itertools import chain, pairwise
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal, TypeAlias
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.vertex_ai.gemini.grounding_requests import GroundingRequests, calculate_grounding_requests
|
||||
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.base_passthrough_logging_handler import (
|
||||
BasePassthroughLoggingHandler,
|
||||
)
|
||||
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler import (
|
||||
PassThroughEndpointLoggingTypedDict,
|
||||
)
|
||||
from litellm.types.utils import LlmProviders, ModelResponse, Usage
|
||||
from litellm.utils import get_model_info
|
||||
from litellm.types.utils import (
|
||||
CompletionTokensDetailsWrapper,
|
||||
CostBreakdown,
|
||||
LlmProviders,
|
||||
ModelResponse,
|
||||
PromptTokensDetailsWrapper,
|
||||
Usage,
|
||||
)
|
||||
|
||||
_NO_GROUNDING: Final = GroundingRequests(web_search_requests=None, google_maps_grounding_requests=None)
|
||||
|
||||
_AGGREGATED_FIELDS: Final = frozenset(
|
||||
{
|
||||
"promptTokenCount",
|
||||
"candidatesTokenCount",
|
||||
"totalTokenCount",
|
||||
"toolUsePromptTokenCount",
|
||||
"promptTokensDetails",
|
||||
"candidatesTokensDetails",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _detail_entries(raw: object) -> tuple[Mapping[str, object], ...]:
|
||||
"""Narrow one turn's ``*TokensDetails`` value to the entries that are actually shaped like one."""
|
||||
return tuple(entry for entry in raw if isinstance(entry, Mapping)) if isinstance(raw, Sequence) else ()
|
||||
|
||||
|
||||
def _grounding_metadata(websocket_messages: Sequence[object]) -> tuple[Mapping[str, object], ...]:
|
||||
"""Collect every ``serverContent.groundingMetadata`` a session emitted.
|
||||
|
||||
Live reports grounding in the server frames, never in ``usageMetadata``, so the per-query
|
||||
charge has to be counted here rather than derived from the token totals.
|
||||
"""
|
||||
return tuple(
|
||||
metadata
|
||||
for message in websocket_messages
|
||||
if isinstance(message, Mapping)
|
||||
for server_content in (message.get("serverContent"),)
|
||||
if isinstance(server_content, Mapping)
|
||||
for metadata in (server_content.get("groundingMetadata"),)
|
||||
if isinstance(metadata, Mapping)
|
||||
)
|
||||
|
||||
|
||||
def _turns(websocket_messages: Sequence[object]) -> tuple[tuple[object, ...], ...]:
|
||||
"""Split a session at every ``usageMetadata`` frame; frames after the last one never got their usage."""
|
||||
closes: Final = tuple(
|
||||
index + 1
|
||||
for index, message in enumerate(websocket_messages)
|
||||
if isinstance(message, Mapping) and isinstance(message.get("usageMetadata"), dict)
|
||||
)
|
||||
return tuple(tuple(websocket_messages[start:end]) for start, end in pairwise((0, *closes)))
|
||||
|
||||
|
||||
def _session_grounding_requests(websocket_messages: Sequence[object]) -> GroundingRequests:
|
||||
per_turn: Final = tuple(
|
||||
calculate_grounding_requests(_grounding_metadata(turn)) for turn in _turns(websocket_messages)
|
||||
)
|
||||
web_search_requests: Final = sum(requests.web_search_requests or 0 for requests in per_turn)
|
||||
google_maps_grounding_requests: Final = sum(requests.google_maps_grounding_requests or 0 for requests in per_turn)
|
||||
return GroundingRequests(
|
||||
web_search_requests=web_search_requests or None,
|
||||
google_maps_grounding_requests=google_maps_grounding_requests or None,
|
||||
)
|
||||
|
||||
|
||||
_SummedField: TypeAlias = Literal[
|
||||
"input_cost",
|
||||
"output_cost",
|
||||
"tool_usage_cost",
|
||||
"cache_read_cost",
|
||||
"cache_creation_cost",
|
||||
"reasoning_cost",
|
||||
"original_cost",
|
||||
"discount_amount",
|
||||
"margin_fixed_amount",
|
||||
"margin_total_amount",
|
||||
]
|
||||
|
||||
|
||||
def _summed(breakdowns: Sequence[CostBreakdown], field: _SummedField) -> float | None:
|
||||
values: Final = tuple(value for breakdown in breakdowns if (value := breakdown.get(field)) is not None)
|
||||
return sum(values) if values else None
|
||||
|
||||
|
||||
class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler):
|
||||
|
|
@ -48,186 +135,110 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler):
|
|||
"""Return the LLM provider name."""
|
||||
return LlmProviders.VERTEX_AI
|
||||
|
||||
@staticmethod
|
||||
def _resolve_detail_counts(
|
||||
details: Sequence[Mapping[str, object]],
|
||||
declared_total: object,
|
||||
) -> tuple[tuple[str, int], ...]:
|
||||
"""
|
||||
Pair each of one turn's ``*TokensDetails`` entries with its token count.
|
||||
|
||||
Live sometimes names the modality that carries the rest of a turn without a
|
||||
``tokenCount``, and reading the absent key as zero drops those tokens from the
|
||||
breakdown, so real audio ends up priced as text. A lone unpriced entry therefore takes
|
||||
whatever the turn's declared count leaves over. Two or more cannot be told apart, so
|
||||
they are left out and the cost calculator charges the remainder as text.
|
||||
"""
|
||||
priced: Final = tuple(
|
||||
(str(detail.get("modality", "TEXT")), count)
|
||||
for detail in details
|
||||
if isinstance(count := detail.get("tokenCount"), int)
|
||||
)
|
||||
unpriced: Final = tuple(
|
||||
str(detail.get("modality", "TEXT")) for detail in details if not isinstance(detail.get("tokenCount"), int)
|
||||
)
|
||||
if len(unpriced) != 1 or not isinstance(declared_total, int):
|
||||
return priced
|
||||
residual: Final = declared_total - sum(count for _, count in priced)
|
||||
return priced if residual <= 0 else (*priced, (unpriced[0], residual))
|
||||
|
||||
@staticmethod
|
||||
def _sum_by_modality(counts: Sequence[tuple[str, int]]) -> Mapping[str, int]:
|
||||
"""Total the (modality, tokenCount) pairs of one or more turns per modality."""
|
||||
return MappingProxyType({modality: sum(c for m, c in counts if m == modality) for modality, _ in counts})
|
||||
|
||||
@staticmethod
|
||||
def _merged_modality_totals(
|
||||
snapshots: Sequence[Mapping[str, object]],
|
||||
count_key: str,
|
||||
details_key: str,
|
||||
) -> Mapping[str, int]:
|
||||
"""Total every turn's per-modality counts, so the breakdown adds up the way the totals do."""
|
||||
return VertexAILivePassthroughLoggingHandler._sum_by_modality(
|
||||
tuple(
|
||||
chain.from_iterable(
|
||||
VertexAILivePassthroughLoggingHandler._resolve_detail_counts(
|
||||
_detail_entries(snapshot.get(details_key)), snapshot.get(count_key)
|
||||
)
|
||||
for snapshot in snapshots
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _extract_usage_metadata_from_websocket_messages(
|
||||
websocket_messages: list[dict],
|
||||
websocket_messages: Sequence[object],
|
||||
) -> dict | None:
|
||||
"""
|
||||
Extract and aggregate usage metadata from a list of WebSocket messages.
|
||||
|
||||
Live emits one ``usageMetadata`` per turn and Google charges per turn for every token in
|
||||
the session context window, which is the current turn's tokens plus all accumulated
|
||||
tokens from previous turns, so the turns add up rather than restating each other. See
|
||||
the Live API note under https://cloud.google.com/vertex-ai/generative-ai/pricing.
|
||||
|
||||
Args:
|
||||
websocket_messages: List of WebSocket messages from the Live API
|
||||
|
||||
Returns:
|
||||
Dictionary containing aggregated usage metadata, or None if not found
|
||||
"""
|
||||
all_usage_metadata: Final = []
|
||||
snapshots: Final = tuple(
|
||||
metadata
|
||||
for message in websocket_messages
|
||||
if isinstance(message, Mapping)
|
||||
for metadata in (message.get("usageMetadata"),)
|
||||
if isinstance(metadata, dict)
|
||||
)
|
||||
|
||||
# Collect all usage metadata messages
|
||||
for message in websocket_messages:
|
||||
if isinstance(message, dict) and "usageMetadata" in message:
|
||||
all_usage_metadata.append(message["usageMetadata"])
|
||||
|
||||
if not all_usage_metadata:
|
||||
if not snapshots:
|
||||
return None
|
||||
|
||||
# If only one usage metadata, return it as-is
|
||||
if len(all_usage_metadata) == 1:
|
||||
return all_usage_metadata[0]
|
||||
|
||||
# Aggregate multiple usage metadata messages
|
||||
aggregated: Final[dict[str, Any]] = {
|
||||
"promptTokenCount": 0,
|
||||
"candidatesTokenCount": 0,
|
||||
"totalTokenCount": 0,
|
||||
"promptTokensDetails": [],
|
||||
"candidatesTokensDetails": [],
|
||||
prompt_totals: Final = VertexAILivePassthroughLoggingHandler._merged_modality_totals(
|
||||
snapshots, "promptTokenCount", "promptTokensDetails"
|
||||
)
|
||||
candidate_totals: Final = VertexAILivePassthroughLoggingHandler._merged_modality_totals(
|
||||
snapshots, "candidatesTokenCount", "candidatesTokensDetails"
|
||||
)
|
||||
return {
|
||||
**{key: value for key, value in snapshots[0].items() if key not in _AGGREGATED_FIELDS},
|
||||
"promptTokenCount": sum(snapshot.get("promptTokenCount", 0) for snapshot in snapshots),
|
||||
"candidatesTokenCount": sum(snapshot.get("candidatesTokenCount", 0) for snapshot in snapshots),
|
||||
"totalTokenCount": sum(snapshot.get("totalTokenCount", 0) for snapshot in snapshots),
|
||||
"toolUsePromptTokenCount": sum(snapshot.get("toolUsePromptTokenCount", 0) for snapshot in snapshots),
|
||||
"promptTokensDetails": [
|
||||
{"modality": modality, "tokenCount": count} for modality, count in prompt_totals.items() if count > 0
|
||||
],
|
||||
"candidatesTokensDetails": [
|
||||
{"modality": modality, "tokenCount": count} for modality, count in candidate_totals.items() if count > 0
|
||||
],
|
||||
}
|
||||
|
||||
# Aggregate token counts
|
||||
for usage in all_usage_metadata:
|
||||
aggregated["promptTokenCount"] += usage.get("promptTokenCount", 0)
|
||||
aggregated["candidatesTokenCount"] += usage.get("candidatesTokenCount", 0)
|
||||
aggregated["totalTokenCount"] += usage.get("totalTokenCount", 0)
|
||||
|
||||
# Aggregate token details by modality
|
||||
modality_totals: Final = {}
|
||||
|
||||
for usage in all_usage_metadata:
|
||||
# Process prompt tokens details
|
||||
for detail in usage.get("promptTokensDetails", []):
|
||||
modality = detail.get("modality", "TEXT")
|
||||
token_count = detail.get("tokenCount", 0)
|
||||
|
||||
if modality not in modality_totals:
|
||||
modality_totals[modality] = {"prompt": 0, "candidate": 0}
|
||||
modality_totals[modality]["prompt"] += token_count
|
||||
|
||||
# Process candidate tokens details
|
||||
for detail in usage.get("candidatesTokensDetails", []):
|
||||
modality = detail.get("modality", "TEXT")
|
||||
token_count = detail.get("tokenCount", 0)
|
||||
|
||||
if modality not in modality_totals:
|
||||
modality_totals[modality] = {"prompt": 0, "candidate": 0}
|
||||
modality_totals[modality]["candidate"] += token_count
|
||||
|
||||
# Convert aggregated modality totals back to details format
|
||||
for modality, totals in modality_totals.items():
|
||||
if totals["prompt"] > 0:
|
||||
aggregated["promptTokensDetails"].append({"modality": modality, "tokenCount": totals["prompt"]})
|
||||
if totals["candidate"] > 0:
|
||||
aggregated["candidatesTokensDetails"].append({"modality": modality, "tokenCount": totals["candidate"]})
|
||||
|
||||
# Add any additional fields from the first usage metadata
|
||||
first_usage: Final = all_usage_metadata[0]
|
||||
for key, value in first_usage.items():
|
||||
if key not in aggregated:
|
||||
aggregated[key] = value
|
||||
|
||||
return aggregated
|
||||
|
||||
@staticmethod
|
||||
def _calculate_live_api_cost(
|
||||
model: str,
|
||||
usage_metadata: dict,
|
||||
custom_llm_provider: str = "vertex_ai",
|
||||
) -> float:
|
||||
"""
|
||||
Calculate cost for Vertex AI Live API based on usage metadata.
|
||||
|
||||
Args:
|
||||
model: The model name (e.g., "gemini-2.0-flash-live-preview-04-09")
|
||||
usage_metadata: Usage metadata from the Live API response
|
||||
custom_llm_provider: The LLM provider (default: "vertex_ai")
|
||||
|
||||
Returns:
|
||||
Total cost in USD
|
||||
"""
|
||||
try:
|
||||
# Get model pricing information
|
||||
model_info: Final = get_model_info(model=model, custom_llm_provider=custom_llm_provider)
|
||||
|
||||
verbose_proxy_logger.debug("Vertex AI Live API model info for '%s': %s", model, model_info)
|
||||
|
||||
# Check if pricing info is available
|
||||
if not model_info or not model_info.get("input_cost_per_token"):
|
||||
verbose_proxy_logger.error("No pricing info found for %s in local model pricing database", model)
|
||||
return 0.0
|
||||
|
||||
total_cost = 0.0
|
||||
|
||||
# Extract token counts from usage metadata
|
||||
prompt_token_count: Final = usage_metadata.get("promptTokenCount", 0)
|
||||
candidates_token_count: Final = usage_metadata.get("candidatesTokenCount", 0)
|
||||
|
||||
# Calculate base text token costs
|
||||
input_cost_per_token: Final = model_info.get("input_cost_per_token", 0.0)
|
||||
output_cost_per_token: Final = model_info.get("output_cost_per_token", 0.0)
|
||||
|
||||
total_cost += prompt_token_count * input_cost_per_token
|
||||
total_cost += candidates_token_count * output_cost_per_token
|
||||
|
||||
# Handle modality-specific costs if present
|
||||
prompt_tokens_details: Final = usage_metadata.get("promptTokensDetails", [])
|
||||
candidates_tokens_details: Final = usage_metadata.get("candidatesTokensDetails", [])
|
||||
|
||||
# Process prompt tokens by modality
|
||||
for detail in prompt_tokens_details:
|
||||
modality = detail.get("modality", "TEXT")
|
||||
token_count = detail.get("tokenCount", 0)
|
||||
|
||||
if modality == "AUDIO":
|
||||
audio_cost_per_token = model_info.get("input_cost_per_audio_token", 0.0)
|
||||
total_cost += token_count * audio_cost_per_token
|
||||
elif modality == "VIDEO":
|
||||
# Video tokens are typically per second, but we'll treat as per token for now
|
||||
video_cost_per_token = model_info.get("input_cost_per_video_per_second", 0.0)
|
||||
total_cost += token_count * video_cost_per_token
|
||||
# TEXT tokens are already handled above
|
||||
|
||||
# Process candidate tokens by modality
|
||||
for detail in candidates_tokens_details:
|
||||
modality = detail.get("modality", "TEXT")
|
||||
token_count = detail.get("tokenCount", 0)
|
||||
|
||||
if modality == "AUDIO":
|
||||
audio_cost_per_token = model_info.get("output_cost_per_audio_token", 0.0)
|
||||
total_cost += token_count * audio_cost_per_token
|
||||
elif modality == "VIDEO":
|
||||
# Video tokens are typically per second, but we'll treat as per token for now
|
||||
video_cost_per_token = model_info.get("output_cost_per_video_per_second", 0.0)
|
||||
total_cost += token_count * video_cost_per_token
|
||||
# TEXT tokens are already handled above
|
||||
|
||||
# Handle web search costs if present
|
||||
tool_use_prompt_token_count: Final = usage_metadata.get("toolUsePromptTokenCount", 0)
|
||||
if tool_use_prompt_token_count > 0:
|
||||
# Web search typically has a fixed cost per request
|
||||
web_search_cost: Final = model_info.get("web_search_cost_per_request", 0.0)
|
||||
if isinstance(web_search_cost, (int, float)) and web_search_cost > 0:
|
||||
total_cost += web_search_cost
|
||||
else:
|
||||
# Fallback to token-based pricing for tool use
|
||||
total_cost += tool_use_prompt_token_count * input_cost_per_token
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
f"Vertex AI Live API cost calculation - Model: {model}, "
|
||||
f"Prompt tokens: {prompt_token_count}, "
|
||||
f"Candidate tokens: {candidates_token_count}, "
|
||||
f"Total cost: ${total_cost:.6f}"
|
||||
)
|
||||
|
||||
return total_cost
|
||||
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error("Error calculating Vertex AI Live API cost: %s", e)
|
||||
return 0.0
|
||||
|
||||
@staticmethod
|
||||
def _create_usage_object_from_metadata(
|
||||
usage_metadata: dict,
|
||||
model: str,
|
||||
grounding_requests: GroundingRequests = _NO_GROUNDING,
|
||||
) -> Usage:
|
||||
"""
|
||||
Create a LiteLLM Usage object from Live API usage metadata.
|
||||
|
|
@ -235,48 +246,124 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler):
|
|||
Args:
|
||||
usage_metadata: Usage metadata from the Live API response
|
||||
model: The model name
|
||||
grounding_requests: The Search and Maps grounding requests summed over the session's
|
||||
turns, matching the per-turn charge
|
||||
|
||||
Returns:
|
||||
LiteLLM Usage object
|
||||
"""
|
||||
prompt_tokens: Final = usage_metadata.get("promptTokenCount", 0)
|
||||
completion_tokens: Final = usage_metadata.get("candidatesTokenCount", 0)
|
||||
total_tokens: Final = usage_metadata.get("totalTokenCount", 0)
|
||||
prompt_by_modality: Final = VertexAILivePassthroughLoggingHandler._sum_by_modality(
|
||||
VertexAILivePassthroughLoggingHandler._resolve_detail_counts(
|
||||
_detail_entries(usage_metadata.get("promptTokensDetails")), usage_metadata.get("promptTokenCount")
|
||||
)
|
||||
)
|
||||
candidates_by_modality: Final = VertexAILivePassthroughLoggingHandler._sum_by_modality(
|
||||
VertexAILivePassthroughLoggingHandler._resolve_detail_counts(
|
||||
_detail_entries(usage_metadata.get("candidatesTokensDetails")),
|
||||
usage_metadata.get("candidatesTokenCount"),
|
||||
)
|
||||
)
|
||||
|
||||
# Create modality-specific token details if available
|
||||
prompt_tokens_details: Final = usage_metadata.get("promptTokensDetails", [])
|
||||
candidates_tokens_details: Final = usage_metadata.get("candidatesTokensDetails", [])
|
||||
|
||||
# Extract text tokens from details
|
||||
text_prompt_tokens = 0
|
||||
text_completion_tokens = 0
|
||||
|
||||
for detail in prompt_tokens_details:
|
||||
if detail.get("modality") == "TEXT":
|
||||
text_prompt_tokens = detail.get("tokenCount", 0)
|
||||
break
|
||||
|
||||
for detail in candidates_tokens_details:
|
||||
if detail.get("modality") == "TEXT":
|
||||
text_completion_tokens = detail.get("tokenCount", 0)
|
||||
break
|
||||
|
||||
# If no text tokens found in details, use total counts
|
||||
if text_prompt_tokens == 0:
|
||||
text_prompt_tokens = prompt_tokens
|
||||
if text_completion_tokens == 0:
|
||||
text_completion_tokens = completion_tokens
|
||||
prompt_tokens: Final = usage_metadata.get("promptTokenCount", 0) or sum(prompt_by_modality.values())
|
||||
completion_tokens: Final = usage_metadata.get("candidatesTokenCount", 0) or sum(candidates_by_modality.values())
|
||||
|
||||
return Usage(
|
||||
prompt_tokens=text_prompt_tokens,
|
||||
completion_tokens=text_completion_tokens,
|
||||
total_tokens=total_tokens,
|
||||
prompt_tokens=prompt_tokens,
|
||||
completion_tokens=completion_tokens,
|
||||
total_tokens=usage_metadata.get("totalTokenCount", 0) or (prompt_tokens + completion_tokens),
|
||||
prompt_tokens_details=PromptTokensDetailsWrapper(
|
||||
text_tokens=prompt_by_modality.get("TEXT"),
|
||||
audio_tokens=prompt_by_modality.get("AUDIO"),
|
||||
image_tokens=prompt_by_modality.get("IMAGE"),
|
||||
video_tokens=prompt_by_modality.get("VIDEO"),
|
||||
tool_use_tokens=usage_metadata.get("toolUsePromptTokenCount") or None,
|
||||
web_search_requests=grounding_requests.web_search_requests,
|
||||
google_maps_grounding_requests=grounding_requests.google_maps_grounding_requests,
|
||||
),
|
||||
completion_tokens_details=CompletionTokensDetailsWrapper(
|
||||
text_tokens=candidates_by_modality.get("TEXT"),
|
||||
audio_tokens=candidates_by_modality.get("AUDIO"),
|
||||
image_tokens=candidates_by_modality.get("IMAGE"),
|
||||
video_tokens=candidates_by_modality.get("VIDEO"),
|
||||
),
|
||||
)
|
||||
|
||||
def _session_usage(self, websocket_messages: Sequence[object], model: str) -> Usage | None:
|
||||
usage_metadata: Final = self._extract_usage_metadata_from_websocket_messages(websocket_messages)
|
||||
if usage_metadata is None:
|
||||
return None
|
||||
return self._create_usage_object_from_metadata(
|
||||
usage_metadata=usage_metadata,
|
||||
grounding_requests=_session_grounding_requests(websocket_messages),
|
||||
model=model,
|
||||
)
|
||||
|
||||
def _turn_cost(
|
||||
self,
|
||||
turn: Sequence[object],
|
||||
model: str,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
) -> tuple[float, CostBreakdown] | None:
|
||||
usage: Final = self._session_usage(turn, model)
|
||||
if usage is None:
|
||||
return None
|
||||
cost: Final = logging_obj._response_cost_calculator( # pyright: ignore[reportPrivateUsage] # the call's own calculator keeps custom pricing and the deployment's region in step with the spend row
|
||||
result=ModelResponse(model=model, usage=usage),
|
||||
litellm_model_name=model,
|
||||
)
|
||||
if cost is None:
|
||||
return None
|
||||
breakdown: Final = logging_obj.cost_breakdown
|
||||
return None if breakdown is None else (cost, breakdown)
|
||||
|
||||
def _session_cost(
|
||||
self,
|
||||
websocket_messages: Sequence[object],
|
||||
model: str,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
) -> float | None:
|
||||
"""Price each turn on its own tokens and grounding, so two grounded turns pay the query fee twice.
|
||||
|
||||
The fixed cost margin is a flat per-request fee, so the session's single spend row carries it once
|
||||
rather than once per turn.
|
||||
"""
|
||||
turn_costs: Final = tuple(self._turn_cost(turn, model, logging_obj) for turn in _turns(websocket_messages))
|
||||
priced: Final = tuple(turn_cost for turn_cost in turn_costs if turn_cost is not None)
|
||||
if not priced or len(priced) != len(turn_costs):
|
||||
return None
|
||||
breakdowns: Final = tuple(breakdown for _, breakdown in priced)
|
||||
first: Final = breakdowns[0]
|
||||
fixed_margin: Final = first.get("margin_fixed_amount") or 0.0
|
||||
duplicated_fixed_margin: Final = fixed_margin * (len(priced) - 1)
|
||||
total_cost: Final = sum(cost for cost, _ in priced) - duplicated_fixed_margin
|
||||
summed_margin_total: Final = _summed(breakdowns, "margin_total_amount")
|
||||
margin_total_amount: Final = (
|
||||
None if summed_margin_total is None else summed_margin_total - duplicated_fixed_margin
|
||||
)
|
||||
logging_obj.set_cost_breakdown(
|
||||
input_cost=_summed(breakdowns, "input_cost") or 0.0,
|
||||
output_cost=_summed(breakdowns, "output_cost") or 0.0,
|
||||
total_cost=total_cost,
|
||||
cost_for_built_in_tools_cost_usd_dollar=_summed(breakdowns, "tool_usage_cost") or 0.0,
|
||||
original_cost=_summed(breakdowns, "original_cost"),
|
||||
discount_percent=first.get("discount_percent"),
|
||||
discount_amount=_summed(breakdowns, "discount_amount"),
|
||||
margin_percent=first.get("margin_percent"),
|
||||
margin_fixed_amount=first.get("margin_fixed_amount"),
|
||||
margin_total_amount=margin_total_amount,
|
||||
cache_read_cost=_summed(breakdowns, "cache_read_cost"),
|
||||
cache_creation_cost=_summed(breakdowns, "cache_creation_cost"),
|
||||
reasoning_cost=_summed(breakdowns, "reasoning_cost"),
|
||||
service_tier=first.get("service_tier"),
|
||||
data_residency=first.get("data_residency"),
|
||||
vertex_location=first.get("vertex_location"),
|
||||
)
|
||||
return total_cost
|
||||
|
||||
def vertex_ai_live_passthrough_handler(
|
||||
self,
|
||||
websocket_messages: list[dict],
|
||||
logging_obj,
|
||||
websocket_messages: Sequence[object],
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
url_route: str,
|
||||
start_time: datetime,
|
||||
end_time: datetime,
|
||||
|
|
@ -300,34 +387,25 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler):
|
|||
"""
|
||||
try:
|
||||
# Extract model from request body or kwargs
|
||||
model: Final = kwargs.get("model", "gemini-2.0-flash-live-preview-04-09")
|
||||
requested_model: Final = kwargs.get("model")
|
||||
model: Final = (
|
||||
requested_model if isinstance(requested_model, str) else "gemini-2.0-flash-live-preview-04-09"
|
||||
)
|
||||
custom_llm_provider: Final = kwargs.get("custom_llm_provider", "vertex_ai")
|
||||
verbose_proxy_logger.debug(
|
||||
"Vertex AI Live API model: %s, custom_llm_provider: %s", model, custom_llm_provider
|
||||
)
|
||||
|
||||
# Extract usage metadata from WebSocket messages
|
||||
usage_metadata: Final = self._extract_usage_metadata_from_websocket_messages(websocket_messages)
|
||||
usage: Final = self._session_usage(websocket_messages, model)
|
||||
|
||||
if not usage_metadata:
|
||||
if usage is None:
|
||||
verbose_proxy_logger.warning("No usage metadata found in Vertex AI Live API WebSocket messages")
|
||||
return {
|
||||
"result": None,
|
||||
"kwargs": kwargs,
|
||||
}
|
||||
|
||||
# Calculate cost using Live API specific pricing
|
||||
response_cost: Final = self._calculate_live_api_cost(
|
||||
model=model,
|
||||
usage_metadata=usage_metadata,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
)
|
||||
|
||||
# Create Usage object for standard LiteLLM logging
|
||||
usage: Final = self._create_usage_object_from_metadata(
|
||||
usage_metadata=usage_metadata,
|
||||
model=model,
|
||||
)
|
||||
response_cost: Final = self._session_cost(websocket_messages, model, logging_obj)
|
||||
|
||||
# Create a mock ModelResponse for standard logging
|
||||
litellm_model_response: Final = ModelResponse(
|
||||
|
|
@ -338,9 +416,9 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler):
|
|||
usage=usage,
|
||||
choices=[],
|
||||
)
|
||||
if response_cost is not None:
|
||||
litellm_model_response._hidden_params["response_cost"] = response_cost # pyright: ignore[reportPrivateUsage] # the logger reads the cost off the response's hidden params; the constructor's hidden_params kwarg is reset by pydantic
|
||||
|
||||
# Update kwargs with cost information
|
||||
kwargs["response_cost"] = response_cost
|
||||
kwargs["model"] = model
|
||||
kwargs["custom_llm_provider"] = custom_llm_provider
|
||||
|
||||
|
|
@ -348,12 +426,15 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler):
|
|||
import re
|
||||
|
||||
allowed_pattern: Final = re.compile(r"^[A-Za-z0-9._\-:]+$")
|
||||
safe_model: Final = model if isinstance(model, str) and allowed_pattern.match(model) else "[REDACTED]"
|
||||
safe_model: Final = model if allowed_pattern.match(model) else "[REDACTED]"
|
||||
verbose_proxy_logger.debug(
|
||||
f"Vertex AI Live API passthrough cost tracking - "
|
||||
f"Model: {safe_model}, Cost: ${response_cost:.6f}, "
|
||||
f"Prompt tokens: {usage.prompt_tokens}, "
|
||||
f"Completion tokens: {usage.completion_tokens}"
|
||||
"Vertex AI Live API passthrough cost tracking - Model: %s, "
|
||||
"Prompt tokens: %s %s, Completion tokens: %s %s",
|
||||
safe_model,
|
||||
usage.prompt_tokens,
|
||||
usage.prompt_tokens_details,
|
||||
usage.completion_tokens,
|
||||
usage.completion_tokens_details,
|
||||
)
|
||||
|
||||
return {
|
||||
|
|
|
|||
|
|
@ -2090,6 +2090,22 @@ def _rewrite_vertex_live_setup_model(text_data: str, setup_model_rewriter: Calla
|
|||
return json.dumps({**message, "setup": {**setup, "model": rewritten_model}}) # mutable-ok: one-shot json payload
|
||||
|
||||
|
||||
def _resolved_vertex_live_setup(
|
||||
setup_data: Mapping[str, object], setup_model_rewriter: Callable[[str], str] | None
|
||||
) -> Mapping[str, object]:
|
||||
"""
|
||||
Give the model extractor the same fully qualified path the upstream will receive.
|
||||
|
||||
Clients may name a bare gateway alias, which the rewriter turns into a ``projects/...`` path before
|
||||
it reaches Vertex. The extractor only reads a path containing ``/models/``, so running it on the raw
|
||||
frame logs the session as ``unknown`` at no cost, which is precisely the supported client form
|
||||
"""
|
||||
setup_model: Final = setup_data.get("model")
|
||||
if setup_model_rewriter is None or not isinstance(setup_model, str):
|
||||
return setup_data
|
||||
return {**setup_data, "model": setup_model_rewriter(setup_model)}
|
||||
|
||||
|
||||
def _truncated_close_reason(reason: str) -> str:
|
||||
"""
|
||||
Fit a close reason inside the byte budget a WebSocket close frame allows, without splitting a character
|
||||
|
|
@ -2314,7 +2330,9 @@ async def websocket_passthrough_request(
|
|||
setup_data,
|
||||
)
|
||||
if isinstance(setup_data, dict) and "model" in setup_data:
|
||||
extracted_model = _extract_model_from_vertex_ai_setup(setup_data)
|
||||
extracted_model = _extract_model_from_vertex_ai_setup(
|
||||
_resolved_vertex_live_setup(setup_data, setup_model_rewriter)
|
||||
)
|
||||
if extracted_model:
|
||||
kwargs["model"] = extracted_model
|
||||
kwargs["custom_llm_provider"] = "vertex_ai-language-models"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -476,9 +476,10 @@ from litellm.proxy.hooks.prompt_injection_detection import (
|
|||
from litellm.proxy.hooks.proxy_track_cost_callback import _ProxyDBLogger, run_spend_event
|
||||
from litellm.proxy.image_endpoints.endpoints import router as image_router
|
||||
from litellm.proxy.list_api.common import (
|
||||
PROBLEM_TYPE_BASE,
|
||||
ManagementProblem,
|
||||
ValidationErrorDetail,
|
||||
problem_response,
|
||||
request_validation_problem,
|
||||
)
|
||||
from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
|
||||
from litellm.proxy.logging_endpoints.callback_logs_endpoints import (
|
||||
|
|
@ -601,7 +602,6 @@ from litellm.proxy.spend_tracking.spend_event_producer import (
|
|||
SpendEventProducer,
|
||||
build_spend_event_producer,
|
||||
)
|
||||
from litellm.types.proxy.management_endpoints.management_v1 import ProblemDetail
|
||||
|
||||
try:
|
||||
from litellm.proxy.enterprise_billing.billing_metrics import (
|
||||
|
|
@ -928,6 +928,7 @@ def cleanup_router_config_variables():
|
|||
user_custom_auth_path, \
|
||||
user_custom_key_generate, \
|
||||
user_custom_key_update, \
|
||||
user_custom_key_policy, \
|
||||
user_custom_sso, \
|
||||
user_custom_ui_sso_sign_in_handler, \
|
||||
use_background_health_checks, \
|
||||
|
|
@ -945,6 +946,7 @@ def cleanup_router_config_variables():
|
|||
user_custom_auth_path = None
|
||||
user_custom_key_generate = None
|
||||
user_custom_key_update = None
|
||||
user_custom_key_policy = None
|
||||
TEAM_METADATA_VALIDATOR_REGISTRY.set(None)
|
||||
TEAM_METADATA_SCHEMA_REGISTRY.set(())
|
||||
user_custom_sso = None
|
||||
|
|
@ -1787,40 +1789,13 @@ class _ExceptionRow(TypedDict, total=False):
|
|||
exception_counts: Mapping[str, int]
|
||||
|
||||
|
||||
class _ValidationErrorDetail(TypedDict):
|
||||
type: ReadOnly[str]
|
||||
loc: ReadOnly[tuple[int | str, ...]]
|
||||
msg: ReadOnly[str]
|
||||
|
||||
|
||||
def _is_length_error_of_rejected_items(error: _ValidationErrorDetail, errors: Sequence[_ValidationErrorDetail]) -> bool:
|
||||
"""pydantic counts only items that validated, so a bad item also trips the parent's min_length."""
|
||||
return error["type"] == "too_short" and any(
|
||||
len(other["loc"]) > len(error["loc"]) and other["loc"][: len(error["loc"])] == error["loc"] for other in errors
|
||||
)
|
||||
|
||||
|
||||
@app.exception_handler(RequestValidationError)
|
||||
async def otel_request_validation_exception_handler(request: Request, exc: RequestValidationError):
|
||||
if request.url.path.startswith(MANAGEMENT_V1_PREFIX):
|
||||
raw_errors: Final[Sequence[_ValidationErrorDetail]] = exc.errors()
|
||||
validation_errors: Final = tuple(
|
||||
error for error in raw_errors if not _is_length_error_of_rejected_items(error, raw_errors)
|
||||
)
|
||||
in_body: Final = any(error["loc"] and error["loc"][0] == "body" for error in validation_errors)
|
||||
status: Final = 422 if in_body else 400
|
||||
_close_dangling_otel_server_span(request, status, exc=exc)
|
||||
return problem_response(
|
||||
ProblemDetail(
|
||||
type=f"{PROBLEM_TYPE_BASE}{'invalid-request-body' if in_body else 'invalid-query-parameter'}",
|
||||
title="Invalid request body" if in_body else "Invalid query parameter",
|
||||
status=status,
|
||||
detail="; ".join(
|
||||
f"{'.'.join(str(part) for part in error['loc'][1:])}: {error['msg']}" for error in validation_errors
|
||||
)
|
||||
or "The request is invalid.",
|
||||
)
|
||||
)
|
||||
validation_errors: Final[Sequence[ValidationErrorDetail]] = exc.errors()
|
||||
problem: Final = request_validation_problem(validation_errors)
|
||||
_close_dangling_otel_server_span(request, problem.status, exc=exc)
|
||||
return problem_response(problem)
|
||||
_close_dangling_otel_server_span(request, 422, exc=exc)
|
||||
return JSONResponse(
|
||||
status_code=422,
|
||||
|
|
@ -2382,6 +2357,7 @@ user_custom_key_generate = None
|
|||
_pkce_no_redis_warning_emitted: bool = False
|
||||
_cp_no_redis_warning_emitted: bool = False
|
||||
user_custom_key_update = None
|
||||
user_custom_key_policy = None
|
||||
user_custom_sso = None
|
||||
user_custom_ui_sso_sign_in_handler = None
|
||||
use_background_health_checks = None
|
||||
|
|
@ -4269,6 +4245,7 @@ _DB_OVERLAY_REMOTE_MODULE_STR_FIELDS: Final[dict[str, tuple[str, ...]]] = {
|
|||
"custom_auth",
|
||||
"custom_key_generate",
|
||||
"custom_key_update",
|
||||
"custom_key_policy",
|
||||
"custom_team_metadata_validate",
|
||||
"custom_sso",
|
||||
"custom_ui_sso_sign_in_handler",
|
||||
|
|
@ -5418,6 +5395,7 @@ class ProxyConfig:
|
|||
user_custom_auth_path, \
|
||||
user_custom_key_generate, \
|
||||
user_custom_key_update, \
|
||||
user_custom_key_policy, \
|
||||
user_custom_sso, \
|
||||
user_custom_ui_sso_sign_in_handler, \
|
||||
use_background_health_checks, \
|
||||
|
|
@ -5955,6 +5933,10 @@ class ProxyConfig:
|
|||
if custom_key_update is not None:
|
||||
user_custom_key_update = get_instance_fn(value=custom_key_update, config_file_path=config_file_path)
|
||||
|
||||
custom_key_policy: Final = general_settings.get("custom_key_policy", None)
|
||||
if custom_key_policy is not None:
|
||||
user_custom_key_policy = get_instance_fn(value=custom_key_policy, config_file_path=config_file_path)
|
||||
|
||||
custom_team_metadata_validate: Final = general_settings.get("custom_team_metadata_validate", None)
|
||||
TEAM_METADATA_VALIDATOR_REGISTRY.set(
|
||||
get_instance_fn(value=custom_team_metadata_validate, config_file_path=config_file_path)
|
||||
|
|
@ -9559,6 +9541,7 @@ class ProxyStartupEvent:
|
|||
gate the first duration window.
|
||||
"""
|
||||
await generate_key_helper_fn(
|
||||
llm_router=llm_router,
|
||||
request_type="user",
|
||||
table_name="user",
|
||||
user_id=LITELLM_PROXY_BUDGET_NAME,
|
||||
|
|
@ -16303,6 +16286,7 @@ async def _generate_onboarding_ui_session_token(user_obj: _UserTableRow) -> str:
|
|||
global master_key, general_settings
|
||||
|
||||
response: Final = await generate_key_helper_fn(
|
||||
llm_router=llm_router,
|
||||
request_type="key",
|
||||
**{
|
||||
"user_role": user_obj.user_role,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -6405,7 +6406,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
|
||||
|
|
|
|||
|
|
@ -30,6 +30,7 @@ from litellm.types.router import GenericLiteLLMParams
|
|||
from litellm.types.utils import CallTypes, LlmProviders
|
||||
from litellm.utils import ProviderConfigManager
|
||||
|
||||
from ..litellm_core_utils.credential_accessor import CredentialAccessor
|
||||
from ..litellm_core_utils.get_litellm_params import get_litellm_params
|
||||
from ..litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
|
||||
from ..llms.azure.common_utils import get_azure_ad_token
|
||||
|
|
@ -54,6 +55,17 @@ xai_realtime: Final = XAIRealtime()
|
|||
vertex_llm_base: Final = VertexBase()
|
||||
base_llm_http_handler = BaseLLMHTTPHandler()
|
||||
_EMPTY_MODEL_PARAMS: Final[Mapping[str, Any]] = MappingProxyType({})
|
||||
_EMPTY_AUTH_HEADERS: Final[Mapping[str, str]] = MappingProxyType({})
|
||||
|
||||
|
||||
def _model_params_with_stored_credentials(model_params: Mapping[str, Any]) -> Mapping[str, Any]:
|
||||
credential_name: Final = model_params.get("litellm_credential_name")
|
||||
credential_values: Final = (
|
||||
CredentialAccessor.get_credential_values(credential_name)
|
||||
if isinstance(credential_name, str)
|
||||
else _EMPTY_MODEL_PARAMS
|
||||
)
|
||||
return MappingProxyType({**credential_values, **model_params})
|
||||
|
||||
|
||||
def _with_resolved_session_model(session: dict[str, object], model_name: str) -> dict[str, object]:
|
||||
|
|
@ -591,13 +603,15 @@ def _azure_realtime_health_protocol(
|
|||
|
||||
def _realtime_health_check_auth_headers(
|
||||
custom_llm_provider: str, api_key: str | None, model_params: Mapping[str, Any]
|
||||
) -> Mapping[str, str | None]:
|
||||
if custom_llm_provider != "azure":
|
||||
return MappingProxyType({"api-key": api_key})
|
||||
return azure_realtime.get_auth_headers(
|
||||
api_key=api_key,
|
||||
azure_ad_token=(None if api_key else get_azure_ad_token(GenericLiteLLMParams(**model_params))),
|
||||
)
|
||||
) -> Mapping[str, str]:
|
||||
if custom_llm_provider == "azure":
|
||||
return azure_realtime.get_auth_headers(
|
||||
api_key=api_key,
|
||||
azure_ad_token=(None if api_key else get_azure_ad_token(GenericLiteLLMParams(**model_params))),
|
||||
)
|
||||
if api_key is None:
|
||||
return _EMPTY_AUTH_HEADERS
|
||||
return MappingProxyType({"Authorization": f"Bearer {api_key}"})
|
||||
|
||||
|
||||
async def _realtime_health_check(
|
||||
|
|
@ -629,34 +643,46 @@ async def _realtime_health_check(
|
|||
"""
|
||||
import websockets
|
||||
|
||||
resolved_params: Final = _model_params_with_stored_credentials(model_params or _EMPTY_MODEL_PARAMS)
|
||||
resolved_api_key: Final = cast( # cast-ok: provider parameters expose optional string credentials
|
||||
str | None, api_key or resolved_params.get("api_key")
|
||||
)
|
||||
resolved_api_base: Final = cast( # cast-ok: provider parameters expose optional string endpoints
|
||||
str | None, api_base or resolved_params.get("api_base")
|
||||
)
|
||||
resolved_api_version: Final = cast( # cast-ok: provider parameters expose optional string versions
|
||||
str | None, api_version or resolved_params.get("api_version")
|
||||
)
|
||||
url: str | None = None
|
||||
auth_headers: Final = _realtime_health_check_auth_headers(
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
api_key=api_key,
|
||||
model_params=model_params or _EMPTY_MODEL_PARAMS,
|
||||
api_key=resolved_api_key,
|
||||
model_params=resolved_params,
|
||||
)
|
||||
if custom_llm_provider == "azure":
|
||||
resolved_protocol, azure_query_params = _azure_realtime_health_protocol(
|
||||
model=model,
|
||||
realtime_protocol=realtime_protocol,
|
||||
model_params=model_params or _EMPTY_MODEL_PARAMS,
|
||||
model_params=resolved_params,
|
||||
)
|
||||
url = azure_realtime._construct_url(
|
||||
api_base=api_base or "",
|
||||
api_base=resolved_api_base or "",
|
||||
model=model,
|
||||
api_version=api_version or "2024-10-01-preview",
|
||||
api_version=resolved_api_version or "2024-10-01-preview",
|
||||
realtime_protocol=resolved_protocol,
|
||||
query_params=azure_query_params,
|
||||
)
|
||||
elif custom_llm_provider == "openai":
|
||||
url = openai_realtime._construct_url(
|
||||
api_base=api_base or "https://api.openai.com/",
|
||||
api_base=resolved_api_base or "https://api.openai.com/",
|
||||
query_params={"model": model},
|
||||
)
|
||||
elif custom_llm_provider == "xai":
|
||||
url = xai_realtime._construct_url(api_base=api_base or "https://api.x.ai/v1", query_params={"model": model})
|
||||
url = xai_realtime._construct_url(
|
||||
api_base=resolved_api_base or "https://api.x.ai/v1", query_params={"model": model}
|
||||
)
|
||||
elif custom_llm_provider == "vertex_ai":
|
||||
vertex_model_params: Final = model_params or {}
|
||||
vertex_model_params: Final = dict(resolved_params)
|
||||
resolved_location: Final = vertex_llm_base.get_vertex_region(
|
||||
vertex_region=VertexBase.safe_get_vertex_ai_location(vertex_model_params),
|
||||
model=model,
|
||||
|
|
@ -675,19 +701,19 @@ async def _realtime_health_check(
|
|||
project=resolved_project,
|
||||
location=resolved_location,
|
||||
)
|
||||
url = vertex_realtime_config.get_complete_url(api_base=api_base, model=model)
|
||||
ssl_context = get_shared_realtime_ssl_context()
|
||||
url = vertex_realtime_config.get_complete_url(api_base=resolved_api_base, model=model)
|
||||
vertex_ssl_context: Final = get_shared_realtime_ssl_context()
|
||||
headers: Final = vertex_realtime_config.validate_environment(headers={}, model=model, api_key=None)
|
||||
async with websockets.connect(
|
||||
url,
|
||||
additional_headers=headers,
|
||||
max_size=REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES,
|
||||
ssl=ssl_context,
|
||||
ssl=vertex_ssl_context,
|
||||
):
|
||||
return True
|
||||
else:
|
||||
raise ValueError(f"Unsupported model: {model}")
|
||||
ssl_context = get_shared_realtime_ssl_context()
|
||||
ssl_context: Final = get_shared_realtime_ssl_context()
|
||||
async with websockets.connect(
|
||||
url,
|
||||
additional_headers=auth_headers,
|
||||
|
|
|
|||
65
litellm/responses/additional_tools.py
Normal file
65
litellm/responses/additional_tools.py
Normal file
|
|
@ -0,0 +1,65 @@
|
|||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import Final, cast # noqa: TID251 # validating the openai tool union strips vendor keys from raw tools
|
||||
|
||||
from pydantic import BaseModel, ValidationError
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.types.llms.openai import ALL_RESPONSES_API_TOOL_PARAMS, ResponseInputParam
|
||||
|
||||
ADDITIONAL_TOOLS_INPUT_ITEM_TYPE: Final = "additional_tools"
|
||||
|
||||
|
||||
class _InputItemType(BaseModel):
|
||||
type: str = ""
|
||||
|
||||
|
||||
class _AdditionalToolsItem(BaseModel):
|
||||
tools: tuple[dict[str, object], ...] = ()
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class HoistedAdditionalTools:
|
||||
input: str | ResponseInputParam
|
||||
tools: tuple[ALL_RESPONSES_API_TOOL_PARAMS, ...]
|
||||
hoisted: tuple[ALL_RESPONSES_API_TOOL_PARAMS, ...]
|
||||
|
||||
|
||||
def _is_additional_tools_item(item: object) -> bool:
|
||||
try:
|
||||
return _InputItemType.model_validate(item).type == ADDITIONAL_TOOLS_INPUT_ITEM_TYPE
|
||||
except ValidationError:
|
||||
return False
|
||||
|
||||
|
||||
def _tools_of_item(item: object) -> tuple[ALL_RESPONSES_API_TOOL_PARAMS, ...]:
|
||||
try:
|
||||
parsed: Final = _AdditionalToolsItem.model_validate(item)
|
||||
except ValidationError:
|
||||
return ()
|
||||
return tuple(
|
||||
cast(
|
||||
"ALL_RESPONSES_API_TOOL_PARAMS", tool
|
||||
) # cast-ok: nested tools carry the same raw tool JSON as top-level tools
|
||||
for tool in parsed.tools
|
||||
)
|
||||
|
||||
|
||||
def hoist_additional_tools(
|
||||
input: str | ResponseInputParam,
|
||||
tools: Sequence[ALL_RESPONSES_API_TOOL_PARAMS] | None,
|
||||
) -> HoistedAdditionalTools:
|
||||
existing: Final = tuple(tools or ())
|
||||
if isinstance(input, str):
|
||||
return HoistedAdditionalTools(input=input, tools=existing, hoisted=())
|
||||
items: Final = tuple(item for item in input if _is_additional_tools_item(item))
|
||||
if not items:
|
||||
return HoistedAdditionalTools(input=input, tools=existing, hoisted=())
|
||||
hoisted: Final = tuple(tool for item in items for tool in _tools_of_item(item))
|
||||
verbose_logger.debug(
|
||||
"Responses API: hoisting %d tool(s) out of %d 'additional_tools' input item(s) into the top-level tools param.",
|
||||
len(hoisted),
|
||||
len(items),
|
||||
)
|
||||
remaining_input: Final = [item for item in input if not _is_additional_tools_item(item)]
|
||||
return HoistedAdditionalTools(input=remaining_input, tools=(*existing, *hoisted), hoisted=hoisted)
|
||||
|
|
@ -39,15 +39,38 @@ def openai_shaped_tool_call_item_id(item_type: str, tool_id: str) -> str:
|
|||
return f"{prefix}_{tool_id}"
|
||||
|
||||
|
||||
class _ToolNameFields(BaseModel):
|
||||
type: str = ""
|
||||
name: str = ""
|
||||
tools: tuple[object, ...] = ()
|
||||
|
||||
|
||||
def _tool_name_fields_of(tool: object) -> _ToolNameFields | None:
|
||||
try:
|
||||
return _ToolNameFields.model_validate(tool)
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
def _custom_tool_name_of(tool: object) -> str | None:
|
||||
parsed: Final = _tool_name_fields_of(tool)
|
||||
if parsed is None or parsed.type != "custom" or not parsed.name:
|
||||
return None
|
||||
return parsed.name
|
||||
|
||||
|
||||
def _nested_tools_of(tool: object) -> tuple[object, ...]:
|
||||
parsed: Final = _tool_name_fields_of(tool)
|
||||
if parsed is None or parsed.type != "namespace":
|
||||
return ()
|
||||
return parsed.tools
|
||||
|
||||
|
||||
def extract_custom_tool_names(tools: Sequence[object] | None) -> set[str]:
|
||||
"""Extract names of tools originally defined as ``type: "custom"``."""
|
||||
if not tools:
|
||||
return set()
|
||||
names: Final[set[str]] = set()
|
||||
for tool in tools:
|
||||
if isinstance(tool, dict) and tool.get("type") == "custom" and "name" in tool:
|
||||
names.add(tool["name"])
|
||||
return names
|
||||
"""Extract names of ``type: "custom"`` tools, at the top level or one level inside a ``namespace`` tool."""
|
||||
top_level: Final = tuple(tools or ())
|
||||
nested: Final = tuple(nested_tool for tool in top_level for nested_tool in _nested_tools_of(tool))
|
||||
return {name for tool in (*top_level, *nested) if (name := _custom_tool_name_of(tool)) is not None}
|
||||
|
||||
|
||||
def is_custom_tool_call(tool_name: str, custom_tool_names: set[str]) -> bool:
|
||||
|
|
@ -143,7 +166,7 @@ def validated_allowed_callers(value: object) -> list[str] | None:
|
|||
raise ValueError("allowed_callers must be a list of strings") from exc
|
||||
|
||||
|
||||
def _grammar_suffix(fmt: object) -> str:
|
||||
def custom_tool_grammar_suffix(fmt: object) -> str:
|
||||
try:
|
||||
parsed: Final = _CustomToolFormat.model_validate(fmt)
|
||||
except ValidationError:
|
||||
|
|
@ -167,7 +190,9 @@ def convert_custom_tool_to_function_tool(tool: Mapping[str, object]) -> ChatComp
|
|||
raw_name: Final = tool.get("name")
|
||||
name: Final = raw_name if isinstance(raw_name, str) else ""
|
||||
raw_description: Final = tool.get("description")
|
||||
description = (raw_description if isinstance(raw_description, str) else "") + _grammar_suffix(tool.get("format"))
|
||||
description: Final = (raw_description if isinstance(raw_description, str) else "") + custom_tool_grammar_suffix(
|
||||
tool.get("format")
|
||||
)
|
||||
allowed_callers: Final = validated_allowed_callers(tool.get("allowed_callers"))
|
||||
function_chunk: Final = ChatCompletionToolParamFunctionChunk(
|
||||
name=name,
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ from collections.abc import Coroutine, Mapping
|
|||
from typing import Final
|
||||
|
||||
import litellm
|
||||
from litellm.responses.additional_tools import hoist_additional_tools
|
||||
from litellm.responses.litellm_completion_transformation.streaming_iterator import (
|
||||
LiteLLMCompletionStreamingIterator,
|
||||
)
|
||||
|
|
@ -37,11 +38,16 @@ class LiteLLMCompletionTransformationHandler:
|
|||
| BaseResponsesAPIStreamingIterator
|
||||
| Coroutine[object, object, ResponsesAPIResponse | BaseResponsesAPIStreamingIterator]
|
||||
):
|
||||
hoisted: Final = hoist_additional_tools(input, responses_api_request.get("tools"))
|
||||
bridged_input: Final = hoisted.input
|
||||
bridged_request: Final[ResponsesAPIOptionalRequestParams] = (
|
||||
{**responses_api_request, "tools": list(hoisted.tools)} if hoisted.hoisted else responses_api_request
|
||||
)
|
||||
litellm_completion_request: Final[dict] = (
|
||||
LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request(
|
||||
model=model,
|
||||
input=input,
|
||||
responses_api_request=responses_api_request,
|
||||
input=bridged_input,
|
||||
responses_api_request=bridged_request,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
stream=stream,
|
||||
extra_headers=extra_headers,
|
||||
|
|
@ -52,8 +58,8 @@ class LiteLLMCompletionTransformationHandler:
|
|||
if _is_async:
|
||||
return self.async_response_api_handler(
|
||||
litellm_completion_request=litellm_completion_request,
|
||||
request_input=input,
|
||||
responses_api_request=responses_api_request,
|
||||
request_input=bridged_input,
|
||||
responses_api_request=bridged_request,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
|
@ -70,8 +76,8 @@ class LiteLLMCompletionTransformationHandler:
|
|||
responses_api_response: Final[ResponsesAPIResponse] = (
|
||||
LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response(
|
||||
chat_completion_response=litellm_completion_response,
|
||||
request_input=input,
|
||||
responses_api_request=responses_api_request,
|
||||
request_input=bridged_input,
|
||||
responses_api_request=bridged_request,
|
||||
)
|
||||
)
|
||||
|
||||
|
|
@ -81,8 +87,8 @@ class LiteLLMCompletionTransformationHandler:
|
|||
return LiteLLMCompletionStreamingIterator(
|
||||
model=model,
|
||||
litellm_custom_stream_wrapper=litellm_completion_response,
|
||||
request_input=input,
|
||||
responses_api_request=responses_api_request,
|
||||
request_input=bridged_input,
|
||||
responses_api_request=bridged_request,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
litellm_metadata=kwargs.get("litellm_metadata", {}),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ from litellm.main import stream_chunk_builder
|
|||
from litellm.responses.litellm_completion_transformation.custom_tools import (
|
||||
build_tool_call_item_kwargs,
|
||||
extract_custom_tool_names,
|
||||
is_custom_tool_call,
|
||||
serialize_tool_call_arguments,
|
||||
)
|
||||
from litellm.responses.litellm_completion_transformation.transformation import (
|
||||
|
|
@ -166,6 +167,14 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
|
|||
return tool_name, namespace
|
||||
return fn_name, None
|
||||
|
||||
def _tool_call_item_kwargs(self, call_id: str, fn_name: str, arguments: str, status: str) -> dict[str, str]:
|
||||
item_kwargs: Final = build_tool_call_item_kwargs(call_id, fn_name, arguments, status, self._custom_tool_names)
|
||||
if is_custom_tool_call(fn_name, self._custom_tool_names):
|
||||
return item_kwargs
|
||||
tool_name, tool_namespace = self._responses_namespace_tool_call_fields(fn_name)
|
||||
namespace_kwargs: Final = {"namespace": tool_namespace} if tool_namespace else {}
|
||||
return {**item_kwargs, "name": tool_name, **namespace_kwargs}
|
||||
|
||||
def _is_reasoning_end(self, chunk):
|
||||
delta: Final = chunk.choices[0].delta
|
||||
|
||||
|
|
@ -244,17 +253,13 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
|
|||
else:
|
||||
fn_name = str(getattr(fn, "name", "") or "")
|
||||
fn_args_delta = serialize_tool_call_arguments(getattr(fn, "arguments", ""))
|
||||
tool_name, tool_namespace = self._responses_namespace_tool_call_fields(fn_name)
|
||||
output_index = self._get_or_assign_tool_output_index(call_id)
|
||||
|
||||
if call_id not in self._tool_args_by_call_id:
|
||||
self._tool_args_by_call_id[call_id] = ""
|
||||
self._sequence_number += 1
|
||||
names = self._custom_tool_names
|
||||
item_kwargs = build_tool_call_item_kwargs(call_id, tool_name, "", "in_progress", names)
|
||||
item_kwargs = self._tool_call_item_kwargs(call_id, fn_name, "", "in_progress")
|
||||
self._tool_item_id_by_call_id[call_id] = item_kwargs["id"]
|
||||
if tool_namespace:
|
||||
item_kwargs["namespace"] = tool_namespace
|
||||
event = OutputItemAddedEvent(
|
||||
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED,
|
||||
output_index=output_index,
|
||||
|
|
@ -315,7 +320,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
|
|||
else:
|
||||
fn_name = str(getattr(fn, "name", "") or "")
|
||||
fn_args = serialize_tool_call_arguments(getattr(fn, "arguments", ""))
|
||||
tool_name, tool_namespace = self._responses_namespace_tool_call_fields(fn_name)
|
||||
web_search_call = self._web_search_calls.get(call_id)
|
||||
if web_search_call is not None:
|
||||
if call_id not in self._queued_web_search_call_ids:
|
||||
|
|
@ -330,11 +334,8 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
|
|||
if is_new_tool_call:
|
||||
self._tool_args_by_call_id[call_id] = ""
|
||||
self._sequence_number += 1
|
||||
names = self._custom_tool_names
|
||||
item_kwargs = build_tool_call_item_kwargs(call_id, tool_name, "", "in_progress", names)
|
||||
item_kwargs = self._tool_call_item_kwargs(call_id, fn_name, "", "in_progress")
|
||||
self._tool_item_id_by_call_id[call_id] = item_kwargs["id"]
|
||||
if tool_namespace:
|
||||
item_kwargs["namespace"] = tool_namespace
|
||||
event = OutputItemAddedEvent(
|
||||
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED,
|
||||
output_index=output_index,
|
||||
|
|
@ -376,11 +377,8 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
|
|||
self._pending_tool_events.append(done_event)
|
||||
|
||||
self._sequence_number += 1
|
||||
names = self._custom_tool_names
|
||||
item_kwargs = build_tool_call_item_kwargs(call_id, tool_name, final_args, "completed", names)
|
||||
item_kwargs = self._tool_call_item_kwargs(call_id, fn_name, final_args, "completed")
|
||||
item_kwargs["id"] = self._tool_item_id_by_call_id.setdefault(call_id, item_kwargs["id"])
|
||||
if tool_namespace:
|
||||
item_kwargs["namespace"] = tool_namespace
|
||||
item_done_event = OutputItemDoneEvent(
|
||||
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE,
|
||||
output_index=output_index,
|
||||
|
|
|
|||
|
|
@ -110,6 +110,7 @@ NamespaceTool: TypeAlias = Mapping[str, object]
|
|||
ResponseTools: TypeAlias = Sequence[Mapping[str, object]] | None
|
||||
ChatToolParam: TypeAlias = ChatCompletionToolParam | OpenAIMcpServerTool
|
||||
NAMESPACE_DESCRIPTION_SEPARATOR: Final = "\n\n"
|
||||
NAMESPACE_MEMBER_TYPES_WITH_CHAT_TOOLS: Final = frozenset({"function", "custom"})
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -1891,9 +1892,21 @@ class LiteLLMCompletionResponsesConfig:
|
|||
namespace_tool: NamespaceTool,
|
||||
nested: bool,
|
||||
) -> ChatCompletionToolParam | None:
|
||||
if nested and namespace_tool.get("type") != "function":
|
||||
tool_type: Final = namespace_tool.get("type")
|
||||
if nested and tool_type not in NAMESPACE_MEMBER_TYPES_WITH_CHAT_TOOLS:
|
||||
return None
|
||||
|
||||
raw_description: Final = str(namespace_tool.get("description") or "")
|
||||
description: Final = (
|
||||
f"{namespace_description}{NAMESPACE_DESCRIPTION_SEPARATOR}{raw_description}"
|
||||
if nested and namespace_description and raw_description
|
||||
else namespace_description
|
||||
if nested and namespace_description
|
||||
else raw_description
|
||||
)
|
||||
if nested and tool_type == "custom":
|
||||
return convert_custom_tool_to_function_tool({**namespace_tool, "description": description})
|
||||
|
||||
raw_parameters: Final = namespace_tool.get("parameters")
|
||||
parameters: Final = (
|
||||
MappingProxyType(raw_parameters) if isinstance(raw_parameters, Mapping) else MappingProxyType({})
|
||||
|
|
@ -1902,14 +1915,6 @@ class LiteLLMCompletionResponsesConfig:
|
|||
parameters if parameters and "type" in parameters else MappingProxyType({**parameters, "type": "object"})
|
||||
)
|
||||
tool_name: Final = str(namespace_tool.get("name") or "")
|
||||
raw_description: Final = str(namespace_tool.get("description") or "")
|
||||
description: Final = (
|
||||
f"{namespace_description}{NAMESPACE_DESCRIPTION_SEPARATOR}{raw_description}"
|
||||
if nested and namespace_description and raw_description
|
||||
else namespace_description
|
||||
if nested and namespace_description
|
||||
else raw_description
|
||||
)
|
||||
chat_tool_name: Final = f"{namespace}__{tool_name}" if nested else tool_name
|
||||
function: Final = ChatCompletionToolParamFunctionChunk(
|
||||
name=chat_tool_name,
|
||||
|
|
@ -2826,6 +2831,22 @@ class LiteLLMCompletionResponsesConfig:
|
|||
if cache_write_tokens is not None
|
||||
else MappingProxyType({})
|
||||
)
|
||||
# The cost path reads the grounding counters off the input details, and a realtime
|
||||
# session's usage is rebuilt from its own response.done, so dropping them here bills
|
||||
# no per-query grounding fee at all.
|
||||
grounding_request_counts: Final[Mapping[str, int]] = MappingProxyType(
|
||||
{
|
||||
counter: count
|
||||
for counter, count in (
|
||||
("web_search_requests", getattr(prompt_details, "web_search_requests", None)),
|
||||
(
|
||||
"google_maps_grounding_requests",
|
||||
getattr(prompt_details, "google_maps_grounding_requests", None),
|
||||
),
|
||||
)
|
||||
if count is not None
|
||||
}
|
||||
)
|
||||
response_usage.input_tokens_details = InputTokensDetails(
|
||||
cached_tokens=prompt_details.cached_tokens if prompt_details.cached_tokens is not None else 0,
|
||||
text_tokens=prompt_details.text_tokens,
|
||||
|
|
@ -2834,6 +2855,7 @@ class LiteLLMCompletionResponsesConfig:
|
|||
cached_tokens_details if isinstance(cached_tokens_details, CachedTokensDetails) else None
|
||||
),
|
||||
**cache_write_extra,
|
||||
**grounding_request_counts,
|
||||
)
|
||||
|
||||
# Translate completion_tokens_details to output_tokens_details
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -1183,6 +1183,10 @@ class ResponseAPILoggingUtils:
|
|||
response_api_usage.input_tokens_details, "cached_tokens_details", None
|
||||
),
|
||||
cache_write_tokens=getattr(response_api_usage.input_tokens_details, "cache_write_tokens", None),
|
||||
web_search_requests=getattr(response_api_usage.input_tokens_details, "web_search_requests", None),
|
||||
google_maps_grounding_requests=getattr(
|
||||
response_api_usage.input_tokens_details, "google_maps_grounding_requests", None
|
||||
),
|
||||
)
|
||||
completion_tokens_details: CompletionTokensDetailsWrapper | None = None
|
||||
output_tokens_details: Final[OutputTokensDetails | None] = getattr(
|
||||
|
|
|
|||
|
|
@ -3268,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,
|
||||
|
|
@ -4105,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
|
||||
|
|
@ -4141,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)
|
||||
|
||||
|
|
@ -4842,6 +4849,7 @@ class Router:
|
|||
model=model,
|
||||
messages=messages,
|
||||
specific_deployment=kwargs.pop("specific_deployment", None),
|
||||
request_kwargs=kwargs,
|
||||
)
|
||||
|
||||
data: Final = deployment["litellm_params"].copy()
|
||||
|
|
@ -5156,13 +5164,11 @@ class Router:
|
|||
return healthy_deployments[0]
|
||||
|
||||
# Use simple_shuffle for weighted selection
|
||||
return cast(
|
||||
GuardrailTypedDict,
|
||||
simple_shuffle(
|
||||
llm_router_instance=self,
|
||||
healthy_deployments=healthy_deployments,
|
||||
model=guardrail_name,
|
||||
),
|
||||
return simple_shuffle(
|
||||
resolve_model_alias=self._get_model_from_alias,
|
||||
healthy_deployments=healthy_deployments,
|
||||
model=guardrail_name,
|
||||
request_kwargs=None,
|
||||
)
|
||||
|
||||
async def _ageneric_api_call_with_fallbacks(self, model: str, original_function: Callable, **kwargs):
|
||||
|
|
@ -8371,7 +8377,8 @@ class Router:
|
|||
|
||||
def log_retry(self, kwargs: dict, e: Exception) -> dict:
|
||||
"""
|
||||
When a retry or fallback happens, record which model group, deployment and attempt just failed and why
|
||||
When a retry or fallback happens, record which model group, deployment and attempt just failed and why,
|
||||
and count it toward the request-wide num_retries_per_request cap
|
||||
"""
|
||||
from litellm.types.router import RetryAttemptRecord
|
||||
|
||||
|
|
@ -8395,7 +8402,10 @@ class Router:
|
|||
else ()
|
||||
)
|
||||
breadcrumbs: Final = (*kept_breadcrumbs, attempt_record)
|
||||
earlier: Final = request_metadata.get("request_retry_count")
|
||||
request_retry_count: Final = (earlier if type(earlier) is int and 0 <= earlier else 0) + 1
|
||||
kwargs[_metadata_var]["previous_models"] = breadcrumbs # rebind-ok: the logging object already holds this dict
|
||||
kwargs[_metadata_var]["request_retry_count"] = request_retry_count # rebind-ok: same dict, read by the cap
|
||||
return kwargs
|
||||
|
||||
def _update_usage(self, deployment_id: str, parent_otel_span: Span | None) -> int:
|
||||
|
|
@ -13038,9 +13048,10 @@ class Router:
|
|||
start_time: Final = time.time()
|
||||
if strategy == "simple-shuffle":
|
||||
return simple_shuffle(
|
||||
llm_router_instance=self,
|
||||
resolve_model_alias=self._get_model_from_alias,
|
||||
healthy_deployments=healthy_deployments,
|
||||
model=model,
|
||||
request_kwargs=request_kwargs,
|
||||
)
|
||||
deployment: Final = await self._select_deployment_async(
|
||||
strategy=strategy,
|
||||
|
|
@ -13183,9 +13194,10 @@ class Router:
|
|||
start_time: Final = time.perf_counter()
|
||||
if strategy == "simple-shuffle":
|
||||
return simple_shuffle(
|
||||
llm_router_instance=self,
|
||||
resolve_model_alias=self._get_model_from_alias,
|
||||
healthy_deployments=pass_through_deployments,
|
||||
model=model,
|
||||
request_kwargs=request_kwargs,
|
||||
)
|
||||
deployment: Final = await self._select_deployment_async(
|
||||
strategy=strategy,
|
||||
|
|
@ -13881,9 +13893,10 @@ class Router:
|
|||
# if users pass rpm or tpm, we do a random weighted pick - based on rpm/tpm
|
||||
############## Check 'weight' param set for weighted pick #################
|
||||
return simple_shuffle(
|
||||
llm_router_instance=self,
|
||||
resolve_model_alias=self._get_model_from_alias,
|
||||
healthy_deployments=healthy_deployments,
|
||||
model=model,
|
||||
request_kwargs=request_kwargs,
|
||||
)
|
||||
deployment: Final = self._select_deployment_sync(
|
||||
strategy=strategy,
|
||||
|
|
@ -13951,6 +13964,7 @@ class Router:
|
|||
messages=messages,
|
||||
input=input,
|
||||
specific_deployment=specific_deployment,
|
||||
request_kwargs=request_kwargs,
|
||||
)
|
||||
|
||||
strategy, strategy_selector = self._get_routing_context(model, request_kwargs)
|
||||
|
|
@ -14033,9 +14047,10 @@ class Router:
|
|||
# 6. Apply load balancing strategy
|
||||
if strategy == "simple-shuffle":
|
||||
return simple_shuffle(
|
||||
llm_router_instance=self,
|
||||
resolve_model_alias=self._get_model_from_alias,
|
||||
healthy_deployments=pass_through_deployments,
|
||||
model=model,
|
||||
request_kwargs=request_kwargs,
|
||||
)
|
||||
deployment: Final = self._select_deployment_sync(
|
||||
strategy=strategy,
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue