mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-16 23:41:43 +00:00
Merge remote-tracking branch 'origin/main' into litellm_bedrock_sanitize_tool_use_id
This commit is contained in:
commit
fa116a0ecc
596 changed files with 42956 additions and 6852 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[@]}"
|
||||
|
|
|
|||
|
|
@ -25,6 +25,8 @@ Same thing for bug fixes. The tests should make it so that this specific bug can
|
|||
|
||||
Never test structure of code only function of it
|
||||
|
||||
A test must only fail when litellm code changes. Never pin facts we don't own (a vendor's price, a third party's field, an upstream default, today's date) as literals or as "X must be absent"; assert the invariant our code guarantees instead, e.g. two rows agree, a value is within range, a field is derived from another. If an outside fact is truly load-bearing, cite its source and date next to the assertion so a reader can tell stale from broken
|
||||
|
||||
`tests/test_litellm/` mirrors `litellm/` in a parallel path (see `tests/test_litellm/readme.md`). Name tests `test_<filename>.py`, but always match the existing test file in the directory you touch — many provider dirs use longer descriptive names (e.g. `test_anthropic_chat_transformation.py`) to avoid ambiguity across sibling folders. For bug fixes, extend the existing mapped test file rather than creating a new one. Only create a new test file for a new feature (provider, endpoint, or transformation module) that has no mapped test yet, following that directory's naming convention (or `test_<filename>.py` if you're the first test there). One focused regression test beats many shallow ones
|
||||
|
||||
End-to-end tests belong in `tests/e2e/` and must follow the harness conventions documented in that directory's `CLAUDE.md`
|
||||
|
|
|
|||
3
Makefile
3
Makefile
|
|
@ -299,6 +299,9 @@ test-rust-extension:
|
|||
[ "$$#" -eq 1 ] && \
|
||||
UV_PROJECT_ENVIRONMENT="$$temporary/venv" $(UV) sync --python 3.12 --frozen --no-install-project --all-groups --all-extras && \
|
||||
$(UV) pip install --python "$$temporary/venv/bin/python" --no-deps "$$1" && \
|
||||
"$$temporary/venv/bin/python" -I -m mypy.stubtest \
|
||||
--mypy-config-file tests/test_litellm/rust_bridge/stubtest.ini \
|
||||
litellm.rust_bridge._native && \
|
||||
LITELLM_RUST=1 LITELLM_LOCAL_MODEL_COST_MAP=True \
|
||||
"$$temporary/venv/bin/python" -I -m pytest --import-mode=importlib -m requires_rust_extension tests/test_litellm_rust
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,14 @@
|
|||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_BudgetTable" ADD COLUMN IF NOT EXISTS "tpd_limit" BIGINT;
|
||||
|
||||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN IF NOT EXISTS "tpd_limit" BIGINT;
|
||||
|
||||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_DeletedTeamTable" ADD COLUMN IF NOT EXISTS "tpd_limit" BIGINT;
|
||||
|
||||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN IF NOT EXISTS "tpd_limit" BIGINT;
|
||||
|
||||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_DeletedVerificationToken" ADD COLUMN IF NOT EXISTS "tpd_limit" BIGINT;
|
||||
|
|
@ -17,6 +17,7 @@ model LiteLLM_BudgetTable {
|
|||
max_parallel_requests Int?
|
||||
tpm_limit BigInt?
|
||||
rpm_limit BigInt?
|
||||
tpd_limit BigInt?
|
||||
model_max_budget Json?
|
||||
budget_duration String?
|
||||
budget_reset_at DateTime?
|
||||
|
|
@ -133,6 +134,7 @@ model LiteLLM_TeamTable {
|
|||
max_parallel_requests Int?
|
||||
tpm_limit BigInt?
|
||||
rpm_limit BigInt?
|
||||
tpd_limit BigInt?
|
||||
budget_duration String?
|
||||
budget_reset_at DateTime?
|
||||
blocked Boolean @default(false)
|
||||
|
|
@ -203,6 +205,7 @@ model LiteLLM_DeletedTeamTable {
|
|||
max_parallel_requests Int?
|
||||
tpm_limit BigInt?
|
||||
rpm_limit BigInt?
|
||||
tpd_limit BigInt?
|
||||
budget_duration String?
|
||||
budget_reset_at DateTime?
|
||||
blocked Boolean @default(false)
|
||||
|
|
@ -438,6 +441,7 @@ model LiteLLM_VerificationToken {
|
|||
blocked Boolean?
|
||||
tpm_limit BigInt?
|
||||
rpm_limit BigInt?
|
||||
tpd_limit BigInt?
|
||||
max_budget Float?
|
||||
budget_duration String?
|
||||
budget_reset_at DateTime?
|
||||
|
|
@ -534,6 +538,7 @@ model LiteLLM_DeletedVerificationToken {
|
|||
blocked Boolean?
|
||||
tpm_limit BigInt?
|
||||
rpm_limit BigInt?
|
||||
tpd_limit BigInt?
|
||||
max_budget Float?
|
||||
budget_duration String?
|
||||
budget_reset_at DateTime?
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ use serde_json::Value;
|
|||
|
||||
#[derive(Deserialize)]
|
||||
struct Input {
|
||||
path: String,
|
||||
model_alias: String,
|
||||
provider_model: String,
|
||||
api_base: String,
|
||||
|
|
@ -21,7 +22,8 @@ async fn main() {
|
|||
Ok(input) => input,
|
||||
Err(error) => fail(error),
|
||||
};
|
||||
let result = litellm_ai_gateway::trace_parity::traced_messages_request(
|
||||
let result = litellm_ai_gateway::trace_parity::traced_request(
|
||||
input.path,
|
||||
input.model_alias,
|
||||
input.provider_model,
|
||||
input.api_base,
|
||||
|
|
|
|||
|
|
@ -29,14 +29,15 @@ pub struct TracedGatewayResponse {
|
|||
pub trace: Vec<litellm_core::observability::FunctionTraceEvent>,
|
||||
}
|
||||
|
||||
pub async fn traced_messages_request(
|
||||
pub async fn traced_request(
|
||||
path: String,
|
||||
model_alias: String,
|
||||
provider_model: String,
|
||||
api_base: String,
|
||||
body: Value,
|
||||
) -> TracedGatewayResponse {
|
||||
let trace = litellm_core::observability::FunctionTrace::default();
|
||||
let result = messages_request(model_alias, provider_model, api_base, body)
|
||||
let result = request(path, model_alias, provider_model, api_base, body)
|
||||
.with_subscriber(trace.dispatcher())
|
||||
.await;
|
||||
let events = trace.events();
|
||||
|
|
@ -54,7 +55,8 @@ pub async fn traced_messages_request(
|
|||
}
|
||||
}
|
||||
|
||||
pub async fn messages_request(
|
||||
pub async fn request(
|
||||
path: String,
|
||||
model_alias: String,
|
||||
provider_model: String,
|
||||
api_base: String,
|
||||
|
|
@ -75,7 +77,7 @@ pub async fn messages_request(
|
|||
};
|
||||
let request = Request::builder()
|
||||
.method("POST")
|
||||
.uri("/v1/messages")
|
||||
.uri(path)
|
||||
.header(AUTHORIZATION, "Bearer trace-master-key")
|
||||
.header(CONTENT_TYPE, "application/json")
|
||||
.body(Body::from(body.to_string()))
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import ast
|
||||
import contextvars
|
||||
import functools
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
|
|
@ -225,6 +226,35 @@ class AccessLogRedactionFilter(logging.Filter):
|
|||
_access_log_filter: Final = AccessLogRedactionFilter()
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=1)
|
||||
def _parse_disabled_access_log_paths(raw: str) -> frozenset[str]:
|
||||
return frozenset(stripped for path in raw.split(",") if (stripped := path.strip()))
|
||||
|
||||
|
||||
def _disabled_access_log_paths() -> frozenset[str]:
|
||||
"""Read the variable per record so a value loaded later via proxy config
|
||||
environment_variables or dotenv is honored."""
|
||||
return _parse_disabled_access_log_paths(os.getenv("LITELLM_DISABLE_ACCESS_LOG_PATHS", ""))
|
||||
|
||||
|
||||
class AccessLogPathFilter(logging.Filter):
|
||||
"""Drops uvicorn.access records for request paths listed in LITELLM_DISABLE_ACCESS_LOG_PATHS.
|
||||
|
||||
uvicorn passes record.args as (client_addr, method, full_path, http_version, status_code).
|
||||
"""
|
||||
|
||||
def filter(self, record: logging.LogRecord) -> bool:
|
||||
if not isinstance(record.args, tuple) or len(record.args) < 3:
|
||||
return True
|
||||
full_path: Final = record.args[2]
|
||||
if not isinstance(full_path, str):
|
||||
return True
|
||||
return full_path.partition("?")[0] not in _disabled_access_log_paths()
|
||||
|
||||
|
||||
_access_log_path_filter: Final = AccessLogPathFilter()
|
||||
|
||||
|
||||
def _get_max_string_length_stdout_log() -> int:
|
||||
"""Read the limit per record so a value loaded later via proxy config
|
||||
environment_variables is honored."""
|
||||
|
|
@ -663,6 +693,7 @@ def _redact_third_party_loggers() -> None:
|
|||
for name in _REDACTED_THIRD_PARTY_LOGGERS:
|
||||
logging.getLogger(name).addFilter(_secret_filter)
|
||||
for name in _REDACTED_ACCESS_LOGGERS:
|
||||
logging.getLogger(name).addFilter(_access_log_path_filter)
|
||||
logging.getLogger(name).addFilter(_access_log_filter)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -21,6 +21,7 @@ from litellm._logging import verbose_logger, verbose_proxy_logger
|
|||
from litellm.a2a_protocol.streaming_iterator import A2AStreamingIterator
|
||||
from litellm.a2a_protocol.utils import A2ARequestUtils
|
||||
from litellm.constants import DEFAULT_A2A_AGENT_TIMEOUT
|
||||
from litellm.litellm_core_utils.asyncify import asyncify
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
get_async_httpx_client,
|
||||
|
|
@ -507,7 +508,7 @@ async def asend_message(
|
|||
prompt_tokens,
|
||||
completion_tokens,
|
||||
_,
|
||||
) = A2ARequestUtils.calculate_usage_from_request_response(
|
||||
) = await asyncify(A2ARequestUtils.calculate_usage_from_request_response)(
|
||||
request=request,
|
||||
response_dict=response_dict,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ import litellm
|
|||
from litellm._logging import verbose_logger
|
||||
from litellm.a2a_protocol.cost_calculator import A2ACostCalculator
|
||||
from litellm.a2a_protocol.utils import A2ARequestUtils
|
||||
from litellm.litellm_core_utils.asyncify import asyncify
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -99,11 +100,11 @@ class A2AStreamingIterator:
|
|||
# Calculate tokens from collected text
|
||||
input_message: Final = A2ARequestUtils.get_input_message_from_request(self.request)
|
||||
input_text: Final = A2ARequestUtils.extract_text_from_message(input_message)
|
||||
prompt_tokens: Final = A2ARequestUtils.count_tokens(input_text)
|
||||
prompt_tokens: Final = await asyncify(A2ARequestUtils.count_tokens)(input_text)
|
||||
|
||||
# Use the last (most complete) text from chunks
|
||||
output_text: Final = self.collected_text_parts[-1] if self.collected_text_parts else ""
|
||||
completion_tokens: Final = A2ARequestUtils.count_tokens(output_text)
|
||||
completion_tokens: Final = await asyncify(A2ARequestUtils.count_tokens)(output_text)
|
||||
|
||||
total_tokens: Final = prompt_tokens + completion_tokens
|
||||
|
||||
|
|
|
|||
|
|
@ -22,6 +22,7 @@
|
|||
"mcp-servers-2025-12-04": null,
|
||||
"oauth-2025-04-20": "oauth-2025-04-20",
|
||||
"output-128k-2025-02-19": "output-128k-2025-02-19",
|
||||
"per-turn-control-2026-07-01": "per-turn-control-2026-07-01",
|
||||
"prompt-caching-scope-2026-01-05": "prompt-caching-scope-2026-01-05",
|
||||
"skills-2025-10-02": "skills-2025-10-02",
|
||||
"structured-outputs-2025-11-13": "structured-outputs-2025-11-13",
|
||||
|
|
@ -52,6 +53,7 @@
|
|||
"mcp-servers-2025-12-04": null,
|
||||
"output-128k-2025-02-19": null,
|
||||
"structured-output-2024-03-01": null,
|
||||
"per-turn-control-2026-07-01": null,
|
||||
"prompt-caching-scope-2026-01-05": "prompt-caching-scope-2026-01-05",
|
||||
"skills-2025-10-02": "skills-2025-10-02",
|
||||
"structured-outputs-2025-11-13": "structured-outputs-2025-11-13",
|
||||
|
|
@ -82,6 +84,7 @@
|
|||
"mcp-servers-2025-12-04": null,
|
||||
"output-128k-2025-02-19": null,
|
||||
"structured-output-2024-03-01": null,
|
||||
"per-turn-control-2026-07-01": null,
|
||||
"prompt-caching-scope-2026-01-05": null,
|
||||
"skills-2025-10-02": null,
|
||||
"structured-outputs-2025-11-13": "structured-outputs-2025-11-13",
|
||||
|
|
@ -113,6 +116,7 @@
|
|||
"mcp-servers-2025-12-04": null,
|
||||
"output-128k-2025-02-19": null,
|
||||
"structured-output-2024-03-01": null,
|
||||
"per-turn-control-2026-07-01": null,
|
||||
"prompt-caching-scope-2026-01-05": null,
|
||||
"skills-2025-10-02": null,
|
||||
"structured-outputs-2025-11-13": null,
|
||||
|
|
@ -144,6 +148,7 @@
|
|||
"mcp-servers-2025-12-04": null,
|
||||
"output-128k-2025-02-19": null,
|
||||
"structured-output-2024-03-01": null,
|
||||
"per-turn-control-2026-07-01": null,
|
||||
"prompt-caching-scope-2026-01-05": null,
|
||||
"skills-2025-10-02": null,
|
||||
"structured-outputs-2025-11-13": null,
|
||||
|
|
@ -176,6 +181,7 @@
|
|||
"mcp-servers-2025-12-04": null,
|
||||
"oauth-2025-04-20": "oauth-2025-04-20",
|
||||
"output-128k-2025-02-19": "output-128k-2025-02-19",
|
||||
"per-turn-control-2026-07-01": null,
|
||||
"prompt-caching-scope-2026-01-05": "prompt-caching-scope-2026-01-05",
|
||||
"skills-2025-10-02": "skills-2025-10-02",
|
||||
"structured-outputs-2025-11-13": "structured-outputs-2025-11-13",
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
@ -679,7 +679,14 @@ class Cache:
|
|||
cache_key, cached_data, kwargs = self._add_cache_logic(result=result, **kwargs)
|
||||
self.cache.set_cache(cache_key, cached_data, **kwargs)
|
||||
except Exception as e:
|
||||
log_redis_failure(verbose_logger, logging.ERROR, "LiteLLM Cache: exception in add_cache", e)
|
||||
self._log_add_cache_failure(e)
|
||||
|
||||
def _log_add_cache_failure(self, exc: Exception) -> None:
|
||||
message: Final = "LiteLLM Cache: exception in add_cache"
|
||||
if isinstance(self.cache, RedisCache):
|
||||
log_redis_failure(verbose_logger, logging.ERROR, message, exc)
|
||||
return
|
||||
verbose_logger.error("%s: %s", message, exc)
|
||||
|
||||
async def async_add_cache(self, result, dynamic_cache_object: BaseCache | None = None, **kwargs):
|
||||
"""
|
||||
|
|
@ -698,7 +705,7 @@ class Cache:
|
|||
else:
|
||||
await self.cache.async_set_cache(cache_key, cached_data, **kwargs)
|
||||
except Exception as e:
|
||||
log_redis_failure(verbose_logger, logging.ERROR, "LiteLLM Cache: exception in add_cache", e)
|
||||
self._log_add_cache_failure(e)
|
||||
|
||||
def _convert_to_cached_embedding(
|
||||
self,
|
||||
|
|
@ -877,7 +884,7 @@ class Cache:
|
|||
else:
|
||||
await self.cache.async_set_cache_pipeline(cache_list=cache_list, **kwargs)
|
||||
except Exception as e:
|
||||
log_redis_failure(verbose_logger, logging.ERROR, "LiteLLM Cache: exception in add_cache", e)
|
||||
self._log_add_cache_failure(e)
|
||||
|
||||
def should_use_cache(self, **kwargs):
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -21,6 +21,7 @@ from litellm.constants import (
|
|||
QDRANT_VECTOR_SIZE,
|
||||
SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS,
|
||||
)
|
||||
from litellm.litellm_core_utils.asyncify import asyncify
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
get_str_from_messages,
|
||||
)
|
||||
|
|
@ -255,7 +256,7 @@ class QdrantSemanticCache(BaseCache):
|
|||
llm_router = None
|
||||
|
||||
router: Final = resolve_embedding_router(self.embedding_model, llm_router, llm_model_list)
|
||||
embedding_input: Final = self._embedding_input(prompt, router)
|
||||
embedding_input: Final = await asyncify(self._embedding_input)(prompt, router)
|
||||
embedding_call: Final = (
|
||||
router.aembedding(
|
||||
model=self.embedding_model,
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ import hashlib
|
|||
import inspect
|
||||
import json
|
||||
import logging
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Awaitable, Callable, Iterator, Sequence
|
||||
from contextvars import ContextVar
|
||||
|
|
@ -32,6 +33,7 @@ from litellm.constants import (
|
|||
REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD,
|
||||
REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT,
|
||||
REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION,
|
||||
REDIS_TIMEOUT_LOG_INTERVAL,
|
||||
)
|
||||
from litellm.litellm_core_utils.core_helpers import _get_parent_otel_span_from_kwargs
|
||||
from litellm.litellm_core_utils.coroutine_checker import coroutine_checker
|
||||
|
|
@ -340,7 +342,7 @@ def _explicit_causes(exc: BaseException) -> Iterator[BaseException]:
|
|||
current = current.__cause__
|
||||
|
||||
|
||||
def _is_redis_timeout_failure(exc: BaseException) -> bool:
|
||||
def is_redis_timeout_failure(exc: BaseException) -> bool:
|
||||
"""True when ``exc`` or any exception it was explicitly raised ``from`` is a timeout.
|
||||
|
||||
redis-py's blocking pool reports a pool wait timeout as ``ConnectionError`` chained from
|
||||
|
|
@ -414,7 +416,7 @@ def _record_swallowed_redis_failure(breaker: RedisCircuitBreaker, exc: BaseExcep
|
|||
"""
|
||||
if not _is_redis_health_failure(exc):
|
||||
return
|
||||
breaker.record_failure(is_timeout=_is_redis_timeout_failure(exc))
|
||||
breaker.record_failure(is_timeout=is_redis_timeout_failure(exc))
|
||||
_swallowed_redis_failures.set(_swallowed_redis_failures.get() + 1)
|
||||
|
||||
|
||||
|
|
@ -422,13 +424,58 @@ class RedisCircuitBreakerOpenError(Exception):
|
|||
pass
|
||||
|
||||
|
||||
class _RedisTimeoutLogThrottle:
|
||||
"""Admits one Redis timeout log line per interval and counts the timeouts it suppressed in between."""
|
||||
|
||||
def __init__(self, interval: float, clock: Callable[[], float] = time.monotonic) -> None:
|
||||
self.interval = interval
|
||||
self._clock = clock
|
||||
self._lock = threading.Lock()
|
||||
self._last_logged_at: float | None = None
|
||||
self._suppressed = 0
|
||||
|
||||
def admit(self) -> int | None:
|
||||
"""Return the number of timeouts suppressed since the last admitted line, or None to suppress this one."""
|
||||
with self._lock:
|
||||
now: Final = self._clock()
|
||||
if self._last_logged_at is not None and now - self._last_logged_at < self.interval:
|
||||
self._suppressed += 1
|
||||
return None
|
||||
suppressed: Final = self._suppressed
|
||||
self._suppressed = 0
|
||||
self._last_logged_at = now
|
||||
return suppressed
|
||||
|
||||
|
||||
_redis_timeout_log_throttle: Final = _RedisTimeoutLogThrottle(REDIS_TIMEOUT_LOG_INTERVAL)
|
||||
|
||||
|
||||
def log_redis_failure(
|
||||
logger: logging.Logger, level: int, message: str, exc: BaseException, with_traceback: bool = False
|
||||
) -> None:
|
||||
if isinstance(exc, RedisCircuitBreakerOpenError):
|
||||
logger.debug("%s: %s", message, exc)
|
||||
logger.debug("%s: %s", message, exc, stacklevel=2)
|
||||
return
|
||||
logger.log(level, "%s: %s", message, exc, exc_info=exc if with_traceback else None)
|
||||
exc_info: Final = exc if with_traceback else None
|
||||
if not is_redis_timeout_failure(exc):
|
||||
logger.log(level, "%s: %s", message, exc, exc_info=exc_info, stacklevel=2)
|
||||
return
|
||||
suppressed: Final = _redis_timeout_log_throttle.admit()
|
||||
if suppressed is None:
|
||||
logger.debug("%s: %s", message, exc, stacklevel=2)
|
||||
return
|
||||
if suppressed == 0:
|
||||
logger.log(level, "%s: %s", message, exc, exc_info=exc_info, stacklevel=2)
|
||||
return
|
||||
logger.log(
|
||||
level,
|
||||
"%s: %s (%d more Redis timeouts since the previous Redis timeout line were logged at DEBUG)",
|
||||
message,
|
||||
exc,
|
||||
suppressed,
|
||||
exc_info=exc_info,
|
||||
stacklevel=2,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -475,7 +522,7 @@ async def _run_under_circuit_breaker(
|
|||
result: Final = await call()
|
||||
except Exception as e:
|
||||
if _is_redis_health_failure(e):
|
||||
breaker.record_failure(is_timeout=_is_redis_timeout_failure(e))
|
||||
breaker.record_failure(is_timeout=is_redis_timeout_failure(e))
|
||||
raise
|
||||
_exit_circuit_breaker(breaker, admission)
|
||||
return result
|
||||
|
|
@ -492,7 +539,7 @@ def _run_under_circuit_breaker_sync(
|
|||
result: Final = call()
|
||||
except Exception as e:
|
||||
if _is_redis_health_failure(e):
|
||||
breaker.record_failure(is_timeout=_is_redis_timeout_failure(e))
|
||||
breaker.record_failure(is_timeout=is_redis_timeout_failure(e))
|
||||
raise
|
||||
_exit_circuit_breaker(breaker, admission)
|
||||
return result
|
||||
|
|
@ -801,10 +848,8 @@ class RedisCache(BaseCache):
|
|||
## LOGGING ##
|
||||
end_time = time.time()
|
||||
_duration = end_time - start_time
|
||||
verbose_logger.error(
|
||||
"LiteLLM Redis Caching: increment_cache() - Got exception from REDIS %s, Writing value=%s",
|
||||
str(e),
|
||||
value,
|
||||
log_redis_failure(
|
||||
verbose_logger, logging.ERROR, "LiteLLM Redis Caching: increment_cache() - Got exception from REDIS", e
|
||||
)
|
||||
raise e
|
||||
|
||||
|
|
@ -1010,11 +1055,8 @@ class RedisCache(BaseCache):
|
|||
call_type=f"async_set_cache <- {_get_call_stack_info()}",
|
||||
)
|
||||
)
|
||||
verbose_logger.error(
|
||||
"LiteLLM Redis Caching: async set() - Got exception from REDIS %s, key=%r, value=%r",
|
||||
str(e),
|
||||
key,
|
||||
value,
|
||||
log_redis_failure(
|
||||
verbose_logger, logging.ERROR, "LiteLLM Redis Caching: async set() - Got exception from REDIS", e
|
||||
)
|
||||
raise e
|
||||
|
||||
|
|
@ -1062,10 +1104,8 @@ class RedisCache(BaseCache):
|
|||
event_metadata={"key": key},
|
||||
)
|
||||
)
|
||||
verbose_logger.error(
|
||||
"LiteLLM Redis Caching: async set() - Got exception from REDIS %s, Writing value=%s",
|
||||
str(e),
|
||||
value,
|
||||
log_redis_failure(
|
||||
verbose_logger, logging.ERROR, "LiteLLM Redis Caching: async set() - Got exception from REDIS", e
|
||||
)
|
||||
_record_swallowed_redis_failure(self._circuit_breaker, e)
|
||||
|
||||
|
|
@ -1112,7 +1152,6 @@ class RedisCache(BaseCache):
|
|||
start_time: Final = time.time()
|
||||
|
||||
print_verbose(f"Set Async Redis Cache: key list: {cache_list}\nttl={ttl}, redis_version={self.redis_version}")
|
||||
cache_value: Final = None
|
||||
try:
|
||||
async with _redis_client.pipeline(transaction=False) as pipe:
|
||||
results: Final = await self._pipeline_helper(pipe, cache_list, ttl)
|
||||
|
|
@ -1149,10 +1188,11 @@ class RedisCache(BaseCache):
|
|||
)
|
||||
)
|
||||
|
||||
verbose_logger.error(
|
||||
"LiteLLM Redis Caching: async set_cache_pipeline() - Got exception from REDIS %s, Writing value=%s",
|
||||
str(e),
|
||||
cache_value,
|
||||
log_redis_failure(
|
||||
verbose_logger,
|
||||
logging.ERROR,
|
||||
"LiteLLM Redis Caching: async set_cache_pipeline() - Got exception from REDIS",
|
||||
e,
|
||||
)
|
||||
_record_swallowed_redis_failure(self._circuit_breaker, e)
|
||||
|
||||
|
|
@ -1191,8 +1231,11 @@ class RedisCache(BaseCache):
|
|||
end_time=time.time(),
|
||||
)
|
||||
)
|
||||
verbose_logger.error(
|
||||
"LiteLLM Redis Caching: async_set_cache_pipeline_with_ttls() - Got exception from REDIS %s", str(e)
|
||||
log_redis_failure(
|
||||
verbose_logger,
|
||||
logging.ERROR,
|
||||
"LiteLLM Redis Caching: async_set_cache_pipeline_with_ttls() - Got exception from REDIS",
|
||||
e,
|
||||
)
|
||||
_record_swallowed_redis_failure(self._circuit_breaker, e)
|
||||
|
||||
|
|
@ -1235,10 +1278,8 @@ class RedisCache(BaseCache):
|
|||
)
|
||||
)
|
||||
# NON blocking - notify users Redis is throwing an exception
|
||||
verbose_logger.error(
|
||||
"LiteLLM Redis Caching: async set() - Got exception from REDIS %s, Writing value=%s",
|
||||
str(e),
|
||||
value,
|
||||
log_redis_failure(
|
||||
verbose_logger, logging.ERROR, "LiteLLM Redis Caching: async set() - Got exception from REDIS", e
|
||||
)
|
||||
raise e
|
||||
|
||||
|
|
@ -1274,10 +1315,11 @@ class RedisCache(BaseCache):
|
|||
)
|
||||
)
|
||||
# NON blocking - notify users Redis is throwing an exception
|
||||
verbose_logger.error(
|
||||
"LiteLLM Redis Caching: async set_cache_sadd() - Got exception from REDIS %s, Writing value=%s",
|
||||
str(e),
|
||||
value,
|
||||
log_redis_failure(
|
||||
verbose_logger,
|
||||
logging.ERROR,
|
||||
"LiteLLM Redis Caching: async set_cache_sadd() - Got exception from REDIS",
|
||||
e,
|
||||
)
|
||||
_record_swallowed_redis_failure(self._circuit_breaker, e)
|
||||
|
||||
|
|
@ -1359,10 +1401,11 @@ class RedisCache(BaseCache):
|
|||
parent_otel_span=parent_otel_span,
|
||||
)
|
||||
)
|
||||
verbose_logger.error(
|
||||
"LiteLLM Redis Caching: async async_increment() - Got exception from REDIS %s, Writing value=%s",
|
||||
str(e),
|
||||
value,
|
||||
log_redis_failure(
|
||||
verbose_logger,
|
||||
logging.ERROR,
|
||||
"LiteLLM Redis Caching: async async_increment() - Got exception from REDIS",
|
||||
e,
|
||||
)
|
||||
raise e
|
||||
|
||||
|
|
@ -1448,7 +1491,9 @@ class RedisCache(BaseCache):
|
|||
print_verbose(f"Got Redis Cache: key: {key}, cached_response {cached_response}")
|
||||
return self._get_cache_logic(cached_response=cached_response)
|
||||
except Exception as e:
|
||||
verbose_logger.error("litellm.caching.caching: get() - Got exception from REDIS: %s", e)
|
||||
log_redis_failure(
|
||||
verbose_logger, logging.ERROR, "litellm.caching.caching: get() - Got exception from REDIS", e
|
||||
)
|
||||
_record_swallowed_redis_failure(self._circuit_breaker, e)
|
||||
|
||||
def _run_redis_mget_operation(self, keys: list[str]) -> Sequence[bytes | str | None]:
|
||||
|
|
@ -1526,7 +1571,7 @@ class RedisCache(BaseCache):
|
|||
end_time=failed_at,
|
||||
parent_otel_span=parent_otel_span,
|
||||
)
|
||||
verbose_logger.error("Error occurred in batch get cache - %s", e)
|
||||
log_redis_failure(verbose_logger, logging.ERROR, "Error occurred in batch get cache", e)
|
||||
_record_swallowed_redis_failure(self._circuit_breaker, e)
|
||||
return key_value_dict
|
||||
|
||||
|
|
@ -1645,7 +1690,7 @@ class RedisCache(BaseCache):
|
|||
parent_otel_span=parent_otel_span,
|
||||
)
|
||||
)
|
||||
verbose_logger.error("Error occurred in async batch get cache - %s", e)
|
||||
log_redis_failure(verbose_logger, logging.ERROR, "Error occurred in async batch get cache", e)
|
||||
_record_swallowed_redis_failure(self._circuit_breaker, e)
|
||||
return key_value_dict
|
||||
|
||||
|
|
@ -1870,9 +1915,11 @@ class RedisCache(BaseCache):
|
|||
parent_otel_span=_get_parent_otel_span_from_kwargs(kwargs),
|
||||
)
|
||||
)
|
||||
verbose_logger.error(
|
||||
"LiteLLM Redis Caching: async increment_pipeline() - Got exception from REDIS %s",
|
||||
str(e),
|
||||
log_redis_failure(
|
||||
verbose_logger,
|
||||
logging.ERROR,
|
||||
"LiteLLM Redis Caching: async increment_pipeline() - Got exception from REDIS",
|
||||
e,
|
||||
)
|
||||
raise e
|
||||
|
||||
|
|
@ -1949,7 +1996,7 @@ class RedisCache(BaseCache):
|
|||
call_type=f"async_rpush <- {_get_call_stack_info()}",
|
||||
)
|
||||
)
|
||||
verbose_logger.error("LiteLLM Redis Cache RPUSH: - Got exception from REDIS : %s", e)
|
||||
log_redis_failure(verbose_logger, logging.ERROR, "LiteLLM Redis Cache RPUSH: - Got exception from REDIS", e)
|
||||
raise e
|
||||
|
||||
async def _pipeline_rpush_helper(
|
||||
|
|
@ -2017,9 +2064,11 @@ class RedisCache(BaseCache):
|
|||
call_type=f"async_rpush_pipeline <- {_get_call_stack_info()}",
|
||||
)
|
||||
)
|
||||
verbose_logger.error(
|
||||
"LiteLLM Redis Caching: async_rpush_pipeline() - Got exception from REDIS %s",
|
||||
str(e),
|
||||
log_redis_failure(
|
||||
verbose_logger,
|
||||
logging.ERROR,
|
||||
"LiteLLM Redis Caching: async_rpush_pipeline() - Got exception from REDIS",
|
||||
e,
|
||||
)
|
||||
raise e
|
||||
|
||||
|
|
@ -2095,7 +2144,7 @@ class RedisCache(BaseCache):
|
|||
call_type=f"async_lpop <- {_get_call_stack_info()}",
|
||||
)
|
||||
)
|
||||
verbose_logger.error("LiteLLM Redis Cache LPOP: - Got exception from REDIS : %s", e)
|
||||
log_redis_failure(verbose_logger, logging.ERROR, "LiteLLM Redis Cache LPOP: - Got exception from REDIS", e)
|
||||
raise e
|
||||
|
||||
async def _pipeline_lpop_helper(
|
||||
|
|
@ -2206,8 +2255,10 @@ class RedisCache(BaseCache):
|
|||
call_type=f"async_lpop_pipeline <- {_get_call_stack_info()}",
|
||||
)
|
||||
)
|
||||
verbose_logger.error(
|
||||
"LiteLLM Redis Caching: async_lpop_pipeline() - Got exception from REDIS %s",
|
||||
str(e),
|
||||
log_redis_failure(
|
||||
verbose_logger,
|
||||
logging.ERROR,
|
||||
"LiteLLM Redis Caching: async_lpop_pipeline() - Got exception from REDIS",
|
||||
e,
|
||||
)
|
||||
raise e
|
||||
|
|
|
|||
|
|
@ -19,6 +19,7 @@ from typing import TYPE_CHECKING, Any, Final, cast
|
|||
import litellm
|
||||
from litellm._logging import print_verbose, verbose_logger
|
||||
from litellm.constants import SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS
|
||||
from litellm.litellm_core_utils.asyncify import asyncify
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
get_str_from_messages,
|
||||
)
|
||||
|
|
@ -522,7 +523,7 @@ class RedisSemanticCache(BaseCache):
|
|||
llm_router = None
|
||||
|
||||
router: Final = resolve_embedding_router(self.embedding_model, llm_router, llm_model_list)
|
||||
embedding_input: Final = self._embedding_input(prompt, router)
|
||||
embedding_input: Final = await asyncify(self._embedding_input)(prompt, router)
|
||||
embedding_call: Final = (
|
||||
router.aembedding(
|
||||
model=self.embedding_model,
|
||||
|
|
|
|||
|
|
@ -205,21 +205,42 @@ 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 _cached_prefix_indices(messages: Sequence[Mapping[str, object]]) -> tuple[int, ...]:
|
||||
last_breakpoint: Final = max(
|
||||
(index for index, msg in enumerate(messages) if _message_has_cache_control(msg)),
|
||||
default=-1,
|
||||
)
|
||||
return tuple(range(last_breakpoint + 1))
|
||||
|
||||
|
||||
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
|
||||
- Every message up to and including the last one 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 the exact bytes of every
|
||||
row up to it, so rewriting any row inside that prefix 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")
|
||||
return tuple(dict.fromkeys(system_indices + last_user + assistant_indices[-1:] + _cached_prefix_indices(messages)))
|
||||
|
||||
|
||||
def _combine_scores(
|
||||
|
|
@ -421,7 +442,7 @@ def compress(
|
|||
combined_scores = bm25_scores
|
||||
|
||||
# Protected messages are never compressed
|
||||
protected_indices: Final = get_protected_indices(normalized_messages)
|
||||
protected_indices: Final = get_protected_indices(original_messages)
|
||||
kept_indices: set[int] = set(protected_indices)
|
||||
|
||||
tool_exchange_spans: list[set[int]] = []
|
||||
|
|
|
|||
|
|
@ -227,6 +227,9 @@ PRE_CALL_EXECUTED_GUARDRAILS_KEY: Final = "_pre_call_executed_guardrails"
|
|||
# Attribute stamped on log_guardrail_information wrappers so __init_subclass__ does not wrap them again
|
||||
LOGS_GUARDRAIL_INFORMATION_MARKER: Final = "_litellm_logs_guardrail_information"
|
||||
|
||||
# llm_provider stamped on proxy-side rate limit errors when the model resolves to no deployment
|
||||
PROXY_LLM_PROVIDER_FALLBACK: Final = "litellm_proxy"
|
||||
|
||||
# Generic fallback for unknown models
|
||||
DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET: Final = int(
|
||||
os.getenv("DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET", 128)
|
||||
|
|
@ -311,6 +314,12 @@ REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS: Final = float(
|
|||
# RFC 6455 caps the close frame payload at 125 bytes, 2 of which carry the status code
|
||||
WEBSOCKET_CLOSE_REASON_MAX_BYTES: Final = 123
|
||||
|
||||
BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY: Final = "litellm.bedrock_realtime.pending_session_update"
|
||||
BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY: Final = "litellm.bedrock_realtime.session_committed"
|
||||
BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY: Final = "litellm.bedrock_realtime.committed_failure"
|
||||
REALTIME_SESSION_SUCCESS_LOGGED_KEY: Final = "realtime_session_success_logged"
|
||||
REALTIME_SESSION_FAILURE_LOGGED_KEY: Final = "realtime_session_failure_logged"
|
||||
|
||||
# SSL/TLS cipher configuration for faster handshakes
|
||||
# Strategy: Strongly prefer fast modern ciphers, but allow fallback to commonly supported ones
|
||||
# This balances performance with broad compatibility
|
||||
|
|
@ -461,6 +470,7 @@ REDIS_CIRCUIT_BREAKER_ENABLED: Final = os.getenv("REDIS_CIRCUIT_BREAKER_ENABLED"
|
|||
# minimum seconds a timeout-only failure streak must span before it can open the breaker,
|
||||
# so one event-loop stall timing out many queued calls at once does not trip it
|
||||
REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION: Final = float(os.getenv("REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION", 5.0))
|
||||
REDIS_TIMEOUT_LOG_INTERVAL: Final = float(os.getenv("REDIS_TIMEOUT_LOG_INTERVAL", "5.0"))
|
||||
# Seconds of idle before a Redis cluster connection is validated with a PING and
|
||||
# reconnected if dead, so a connection silently dropped by a cluster restart
|
||||
# (e.g. ElastiCache Serverless maintenance) is not reused while broken
|
||||
|
|
@ -1556,6 +1566,8 @@ BASE_MCP_ROUTE: Final = "/mcp"
|
|||
|
||||
BATCH_STATUS_POLL_INTERVAL_SECONDS: Final = int(os.getenv("BATCH_STATUS_POLL_INTERVAL_SECONDS", 3600)) # 1 hour
|
||||
BATCH_STATUS_POLL_MAX_ATTEMPTS: Final = int(os.getenv("BATCH_STATUS_POLL_MAX_ATTEMPTS", 24)) # for 24 hours
|
||||
BATCH_TPD_WINDOW_SECONDS: Final = 86400
|
||||
BATCH_TPD_DESCRIPTOR_SUFFIX: Final = "_tpd"
|
||||
|
||||
HEALTH_CHECK_TIMEOUT_SECONDS: Final = int(os.getenv("HEALTH_CHECK_TIMEOUT_SECONDS", 60)) # 60 seconds
|
||||
_background_health_check_max_tokens_env: Final = os.getenv("BACKGROUND_HEALTH_CHECK_MAX_TOKENS")
|
||||
|
|
@ -1966,6 +1978,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(
|
||||
|
|
|
|||
|
|
@ -97,6 +97,7 @@ from litellm.llms.vertex_ai.cost_calculator import cost_router as google_cost_ro
|
|||
from litellm.llms.xai.cost_calculator import cost_per_token as xai_cost_per_token
|
||||
from litellm.responses.utils import ResponseAPILoggingUtils
|
||||
from litellm.types.agents import LiteLLMSendMessageResponse
|
||||
from litellm.types.llms.base import CachedTokensDetails
|
||||
from litellm.types.llms.openai import (
|
||||
HttpxBinaryResponseContent,
|
||||
ImageGenerationRequestQuality,
|
||||
|
|
@ -813,6 +814,7 @@ def _select_model_name_for_cost_calc(
|
|||
if (
|
||||
entry.get("input_cost_per_token") is not None
|
||||
or entry.get("input_cost_per_second") is not None
|
||||
or entry.get("input_cost_per_query") is not None
|
||||
or entry.get("tiered_pricing") is not None
|
||||
):
|
||||
return_model = router_model_id
|
||||
|
|
@ -2277,6 +2279,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,
|
||||
|
|
@ -2336,7 +2351,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"]
|
||||
|
|
@ -2381,6 +2418,46 @@ def _summable_prompt_token_fields(prompt_tokens_details: BaseModel) -> list[str]
|
|||
return [attr for attr in field_names if attr != "cache_creation_tokens"]
|
||||
|
||||
|
||||
def _combine_cached_tokens_details(
|
||||
current: CachedTokensDetails | None, new: CachedTokensDetails
|
||||
) -> CachedTokensDetails:
|
||||
def _sum_optional(current_value: int | None, new_value: int | None) -> int | None:
|
||||
if current_value is None and new_value is None:
|
||||
return None
|
||||
return (current_value or 0) + (new_value or 0)
|
||||
|
||||
return CachedTokensDetails(
|
||||
text_tokens=_sum_optional(current.text_tokens if current is not None else None, new.text_tokens),
|
||||
audio_tokens=_sum_optional(current.audio_tokens if current is not None else None, new.audio_tokens),
|
||||
image_tokens=_sum_optional(current.image_tokens if current is not None else None, new.image_tokens),
|
||||
)
|
||||
|
||||
|
||||
def _combine_prompt_tokens_details(
|
||||
current: PromptTokensDetailsWrapper | None, new: PromptTokensDetailsWrapper
|
||||
) -> PromptTokensDetailsWrapper:
|
||||
base: Final = current if current is not None else PromptTokensDetailsWrapper()
|
||||
base_values: Final = MappingProxyType(
|
||||
{attr: getattr(base, attr) for attr in type(base).model_fields if hasattr(base, attr)}
|
||||
)
|
||||
summed: Final = MappingProxyType(
|
||||
{
|
||||
attr: (getattr(base, attr, 0) or 0) + (getattr(new, attr) or 0)
|
||||
for attr in _summable_prompt_token_fields(new)
|
||||
if hasattr(new, attr) and isinstance(getattr(new, attr) or 0, (int, float))
|
||||
}
|
||||
)
|
||||
new_cached_tokens_details: Final = getattr(new, "cached_tokens_details", None)
|
||||
cached_tokens_details: Final = (
|
||||
_combine_cached_tokens_details(getattr(base, "cached_tokens_details", None), new_cached_tokens_details)
|
||||
if isinstance(new_cached_tokens_details, CachedTokensDetails)
|
||||
else getattr(base, "cached_tokens_details", None)
|
||||
)
|
||||
return PromptTokensDetailsWrapper(
|
||||
**MappingProxyType({**base_values, **summed, "cached_tokens_details": cached_tokens_details})
|
||||
)
|
||||
|
||||
|
||||
class BaseTokenUsageProcessor:
|
||||
@staticmethod
|
||||
def combine_usage_objects(usage_objects: list[Usage]) -> Usage:
|
||||
|
|
@ -2389,7 +2466,6 @@ class BaseTokenUsageProcessor:
|
|||
"""
|
||||
from litellm.types.utils import (
|
||||
CompletionTokensDetailsWrapper,
|
||||
PromptTokensDetailsWrapper,
|
||||
Usage,
|
||||
)
|
||||
|
||||
|
|
@ -2408,27 +2484,10 @@ class BaseTokenUsageProcessor:
|
|||
and isinstance(current_val, (int, float))
|
||||
):
|
||||
setattr(combined, attr, current_val + new_val)
|
||||
# Handle nested prompt_tokens_details
|
||||
if hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details:
|
||||
if not hasattr(combined, "prompt_tokens_details") or not combined.prompt_tokens_details:
|
||||
combined.prompt_tokens_details = PromptTokensDetailsWrapper()
|
||||
|
||||
# Check what keys exist in the model's prompt_tokens_details
|
||||
# Access model_fields on the class, not the instance, to avoid Pydantic 2.11+ deprecation warnings
|
||||
for attr in _summable_prompt_token_fields(usage.prompt_tokens_details):
|
||||
if (
|
||||
hasattr(usage.prompt_tokens_details, attr)
|
||||
and not attr.startswith("_")
|
||||
and not callable(_attribute_value(usage.prompt_tokens_details, attr))
|
||||
):
|
||||
current_val = getattr(combined.prompt_tokens_details, attr, 0) or 0
|
||||
new_val = getattr(usage.prompt_tokens_details, attr, 0) or 0
|
||||
if new_val is not None and isinstance(new_val, (int, float)):
|
||||
setattr(
|
||||
combined.prompt_tokens_details,
|
||||
attr,
|
||||
current_val + new_val,
|
||||
)
|
||||
combined.prompt_tokens_details = _combine_prompt_tokens_details(
|
||||
getattr(combined, "prompt_tokens_details", None), usage.prompt_tokens_details
|
||||
)
|
||||
|
||||
# Handle nested completion_tokens_details
|
||||
if hasattr(usage, "completion_tokens_details") and usage.completion_tokens_details:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ from typing_extensions import ReadOnly, TypedDict
|
|||
from litellm._logging import verbose_logger
|
||||
from litellm.compression import compress
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.asyncify import asyncify
|
||||
from litellm.types.integrations.compression_interception import (
|
||||
CompressionInterceptionConfig,
|
||||
CompressionSavingsMetadata,
|
||||
|
|
@ -153,7 +154,7 @@ class CompressionInterceptionLogger(CustomLogger):
|
|||
|
||||
self._prune_expired_cache()
|
||||
|
||||
compressed: Final = compress(
|
||||
compressed: Final = await asyncify(compress)(
|
||||
messages=messages,
|
||||
model=model,
|
||||
call_type=CallTypes.anthropic_messages,
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -822,46 +822,33 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
|||
def truncate_standard_logging_payload_content(
|
||||
self,
|
||||
standard_logging_object: StandardLoggingPayload,
|
||||
):
|
||||
) -> StandardLoggingPayload:
|
||||
"""
|
||||
Truncate error strings and message content in logging payload
|
||||
Return a copy of the logging payload with error_str, messages, and response truncated
|
||||
|
||||
Some loggers like DataDog/ GCS Bucket have a limit on the size of the payload. (1MB)
|
||||
|
||||
This function truncates the error string and the message content if they exceed a certain length.
|
||||
Every callback of a request shares one standard logging object, so the payload passed in is left
|
||||
untouched and the callbacks that run later (the prompt caching router check, spend logs) still see
|
||||
the original fields.
|
||||
"""
|
||||
MAX_STR_LENGTH: Final = 10_000
|
||||
max_str_length: Final = 10_000
|
||||
candidates: Final = {
|
||||
field: self._truncate_field(field_value=standard_logging_object.get(field), max_length=max_str_length)
|
||||
for field in ("error_str", "messages", "response")
|
||||
}
|
||||
truncated_fields: Final = {field: text for field, text in candidates.items() if text is not None}
|
||||
return {**standard_logging_object, **truncated_fields}
|
||||
|
||||
# Truncate fields that might exceed max length
|
||||
fields_to_truncate: Final = ["error_str", "messages", "response"]
|
||||
for field in fields_to_truncate:
|
||||
self._truncate_field(
|
||||
standard_logging_object=standard_logging_object,
|
||||
field_name=field,
|
||||
max_length=MAX_STR_LENGTH,
|
||||
)
|
||||
|
||||
def _truncate_field(
|
||||
self,
|
||||
standard_logging_object: StandardLoggingPayload,
|
||||
field_name: str,
|
||||
max_length: int,
|
||||
) -> None:
|
||||
def _truncate_field(self, field_value: object, max_length: int) -> str | None:
|
||||
"""
|
||||
Helper function to truncate a field in the logging payload
|
||||
Return the truncated text of a field that exceeds max_length, or None when the field fits
|
||||
|
||||
This converts the field to a string and then truncates it if it exceeds the max length.
|
||||
|
||||
Why convert to string ?
|
||||
1. User was sending a poorly formatted list for `messages` field, we could not predict where they would send content
|
||||
- Converting to string and then truncating the logged content catches this
|
||||
2. We want to avoid modifying the original `messages`, `response`, and `error_str` in the logging payload since these are in kwargs and could be returned to the user
|
||||
The field is measured as a string because users send poorly formatted lists for `messages`, so there is
|
||||
no fixed place the content would be.
|
||||
"""
|
||||
field_value: Final[object] = standard_logging_object.get(field_name)
|
||||
if field_value:
|
||||
str_value: Final = str(field_value)
|
||||
if len(str_value) > max_length:
|
||||
standard_logging_object[field_name] = self._truncate_text(text=str_value, max_length=max_length)
|
||||
text: Final = str(field_value or "")
|
||||
return self._truncate_text(text=text, max_length=max_length) if len(text) > max_length else None
|
||||
|
||||
def _truncate_text(self, text: str, max_length: int) -> str:
|
||||
"""Truncate text if it exceeds max_length"""
|
||||
|
|
|
|||
|
|
@ -563,11 +563,10 @@ class DataDogLogger(
|
|||
if standard_logging_object.get("status") == "failure":
|
||||
status = DataDogStatus.ERROR
|
||||
|
||||
# Build the initial payload
|
||||
self.truncate_standard_logging_payload_content(standard_logging_object)
|
||||
truncated_payload: Final = self.truncate_standard_logging_payload_content(standard_logging_object)
|
||||
|
||||
dd_payload: Final = self._create_datadog_logging_payload_helper(
|
||||
standard_logging_object=standard_logging_object,
|
||||
standard_logging_object=truncated_payload,
|
||||
status=status,
|
||||
)
|
||||
return dd_payload
|
||||
|
|
|
|||
|
|
@ -39,6 +39,8 @@ def is_serializable(value):
|
|||
|
||||
|
||||
class LangsmithLogger(CustomBatchLogger):
|
||||
preserve_events_added_during_flush = True
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
langsmith_api_key: str | None = None,
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ from typing import Final
|
|||
|
||||
from litellm.integrations.otel.mappers.base import AttributeMap, AttrValue, SpanData
|
||||
from litellm.integrations.otel.mappers.utils import (
|
||||
MAX_MESSAGE_ATTRS_PER_SPAN,
|
||||
MAX_TOOL_DEFINITION_ATTRS_PER_SPAN,
|
||||
collect,
|
||||
drop_none,
|
||||
|
|
@ -26,6 +27,8 @@ from litellm.integrations.otel.model.payloads import (
|
|||
ToolDefinition,
|
||||
)
|
||||
|
||||
_MAX_INDEXED_MESSAGES: Final = MAX_MESSAGE_ATTRS_PER_SPAN // 2
|
||||
|
||||
|
||||
class OpenInferenceMapper:
|
||||
"""Emits OpenInference attributes for LLM_CALL spans.
|
||||
|
|
@ -84,27 +87,44 @@ class OpenInferenceMapper:
|
|||
return {}
|
||||
|
||||
def _llm_call(self, data: LLMCallSpanData) -> AttributeMap:
|
||||
outputs: Final = output_messages(data)
|
||||
indexed_in, indexed_out = self._indexed_split(len(data.messages_in), len(outputs))
|
||||
return {
|
||||
**collect(self._LLM_CALL_ATTRS, data),
|
||||
**collect(self._BLOB_ATTRS, data),
|
||||
**self._messages("llm.input_messages", "input.value", data.messages_in),
|
||||
**self._messages("llm.output_messages", "output.value", output_messages(data)),
|
||||
**self._messages(
|
||||
"llm.input_messages",
|
||||
"input.value",
|
||||
data.messages_in,
|
||||
self._prompt_positions(len(data.messages_in), indexed_in),
|
||||
),
|
||||
**self._messages("llm.output_messages", "output.value", outputs, range(indexed_out)),
|
||||
**self._tools(data),
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _messages(prefix: str, value_key: str, messages: Sequence[object]) -> AttributeMap:
|
||||
"""Per-message ``{prefix}.{idx}.message.*`` keys + the ``value_key`` blob."""
|
||||
def _indexed_split(inputs: int, outputs: int) -> tuple[int, int]:
|
||||
"""Prompt and response share one allowance; the response is reserved at least half of it."""
|
||||
indexed_out: Final = min(outputs, max(_MAX_INDEXED_MESSAGES // 2, _MAX_INDEXED_MESSAGES - inputs))
|
||||
return _MAX_INDEXED_MESSAGES - indexed_out, indexed_out
|
||||
|
||||
@staticmethod
|
||||
def _prompt_positions(total: int, indexed: int) -> tuple[int, ...]:
|
||||
"""Prompt messages that get per-index attributes: message 0 and the most recent turns."""
|
||||
if total <= indexed:
|
||||
return tuple(range(total))
|
||||
return (0, *range(total - indexed + 1, total))
|
||||
|
||||
@staticmethod
|
||||
def _messages(prefix: str, value_key: str, messages: Sequence[object], positions: Sequence[int]) -> AttributeMap:
|
||||
"""``{prefix}.{idx}.message.*`` keys for the messages at ``positions`` + the ``value_key`` blob of all."""
|
||||
parsed: Final = [(m.get("role") if isinstance(m, dict) else None, message_content(m)) for m in messages]
|
||||
attrs: Final = drop_none(
|
||||
{
|
||||
key: value
|
||||
for idx, (role, content) in enumerate(parsed)
|
||||
for idx, (role, content) in ((idx, parsed[idx]) for idx in positions)
|
||||
for key, value in (
|
||||
(
|
||||
f"{prefix}.{idx}.message.role",
|
||||
role if isinstance(role, str) else None,
|
||||
),
|
||||
(f"{prefix}.{idx}.message.role", role if isinstance(role, str) else None),
|
||||
(f"{prefix}.{idx}.message.content", content),
|
||||
)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -32,6 +32,14 @@ core telemetry no matter how many vocabularies are configured.
|
|||
"""
|
||||
|
||||
|
||||
MAX_MESSAGE_ATTRS_PER_SPAN: Final = DEFAULT_SPAN_ATTRIBUTE_LIMIT // 8
|
||||
"""Span-wide ceiling on per-index chat message attributes, prompt and response together.
|
||||
|
||||
An eighth is the largest share that still fits beside the tool ceiling and the core
|
||||
of every vocabulary at once. The complete conversation still rides the JSON blobs.
|
||||
"""
|
||||
|
||||
|
||||
def tool_attr_budget(vocabularies: int) -> int:
|
||||
"""Split the span-wide tool-definition ceiling across active vocabularies."""
|
||||
return MAX_TOOL_DEFINITION_ATTRS_PER_SPAN // max(vocabularies, 1)
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ from pydantic import BaseModel
|
|||
|
||||
import litellm
|
||||
from litellm._logging import print_verbose, verbose_logger
|
||||
from litellm.constants import PROXY_LLM_PROVIDER_FALLBACK
|
||||
from litellm.exceptions import (
|
||||
validate_rate_limit_category,
|
||||
validate_rate_limit_type,
|
||||
|
|
@ -2581,6 +2582,15 @@ class PrometheusLogger(CustomLogger):
|
|||
)
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _extract_api_provider_from_exception(exception: Exception) -> str | None:
|
||||
if not isinstance(exception, litellm.exceptions.RateLimitError):
|
||||
return None
|
||||
llm_provider: Final = exception.llm_provider
|
||||
if not llm_provider or llm_provider == PROXY_LLM_PROVIDER_FALLBACK:
|
||||
return None
|
||||
return llm_provider
|
||||
|
||||
async def async_post_call_failure_hook(
|
||||
self,
|
||||
request_data: dict,
|
||||
|
|
@ -2600,12 +2610,6 @@ class PrometheusLogger(CustomLogger):
|
|||
StandardLoggingPayloadSetup,
|
||||
)
|
||||
|
||||
if self._should_skip_metrics_for_invalid_key(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
exception=original_exception,
|
||||
):
|
||||
return
|
||||
|
||||
status_code: Final = self._extract_status_code(exception=original_exception)
|
||||
|
||||
try:
|
||||
|
|
@ -2616,12 +2620,14 @@ class PrometheusLogger(CustomLogger):
|
|||
_metadata: Final = request_data.get("metadata", {}) or {}
|
||||
model_id: Final = _metadata.get("model_info", {}).get("id") or request_data.get("model_info", {}).get("id")
|
||||
rate_limit_category, rate_limit_type = self._extract_rate_limit_labels(original_exception)
|
||||
api_provider: Final = self._extract_api_provider_from_request_data(request_data)
|
||||
api_provider: Final = self._extract_api_provider_from_request_data(
|
||||
request_data
|
||||
) or self._extract_api_provider_from_exception(original_exception)
|
||||
enum_values: Final = UserAPIKeyLabelValues(
|
||||
end_user=user_api_key_dict.end_user_id,
|
||||
user=user_api_key_dict.user_id,
|
||||
user_email=user_api_key_dict.user_email,
|
||||
hashed_api_key=user_api_key_dict.api_key,
|
||||
hashed_api_key=None if status_code == 401 else user_api_key_dict.api_key,
|
||||
api_key_alias=user_api_key_dict.key_alias,
|
||||
team=user_api_key_dict.team_id,
|
||||
team_alias=user_api_key_dict.team_alias,
|
||||
|
|
|
|||
|
|
@ -4,18 +4,10 @@ imported_openAIResponse = True
|
|||
try:
|
||||
import io
|
||||
import logging
|
||||
import sys
|
||||
from typing import Any, TypeVar
|
||||
from typing import Any, Literal, Protocol, TypeVar
|
||||
|
||||
from wandb.sdk.data_types import trace_tree
|
||||
|
||||
if sys.version_info >= (3, 8):
|
||||
from typing import Literal, Protocol
|
||||
else:
|
||||
from typing import Literal
|
||||
|
||||
from typing_extensions import Protocol
|
||||
|
||||
logger: Final = logging.getLogger(__name__)
|
||||
|
||||
K = TypeVar("K", bound=str)
|
||||
|
|
|
|||
|
|
@ -303,6 +303,16 @@ def get_metadata_variable_name_from_kwargs(
|
|||
return "litellm_metadata" if "litellm_metadata" in kwargs else "metadata"
|
||||
|
||||
|
||||
def max_retries_per_request_hit(kwargs: Mapping[str, object], num_retries_per_request: int | None) -> bool:
|
||||
if num_retries_per_request is None:
|
||||
return False
|
||||
metadata: Final = kwargs.get(get_metadata_variable_name_from_kwargs(kwargs))
|
||||
if not isinstance(metadata, Mapping):
|
||||
return False
|
||||
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(
|
||||
request_data: dict,
|
||||
) -> tuple[Literal["metadata", "litellm_metadata"], dict]:
|
||||
|
|
|
|||
|
|
@ -36,7 +36,7 @@ class CoroutineChecker:
|
|||
target = callback
|
||||
if not inspect.isfunction(target) and not inspect.ismethod(target):
|
||||
try:
|
||||
call_attr: Final = getattr(target, "__call__", None)
|
||||
call_attr: Final = getattr(target, "__call__", None) # noqa: B004 # value unwrap so iscoroutinefunction sees through functors
|
||||
if call_attr is not None:
|
||||
target = call_attr
|
||||
except Exception:
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -34,6 +34,11 @@ rules never mix the two and never use ``extends``. A rule whose
|
|||
Rules are only consulted after exact and case-insensitive lookups miss, so an
|
||||
exact cost-map entry always takes precedence over any rule.
|
||||
|
||||
Rules flagged with ``fill_missing_for_providers: [..]`` also fill only keys
|
||||
missing from an exact cost-map entry when the entry's ``litellm_provider`` is
|
||||
listed, while values already present on the entry win on conflict. Only flagged
|
||||
capability rules participate in this fill; routing rules never do.
|
||||
|
||||
Patterns are matched case-insensitively with ``re.search`` and are not implicitly
|
||||
anchored: a rule must include ``^`` and ``$`` to bind to the whole model name,
|
||||
otherwise it matches as a substring. Keeping anchoring in the regex makes the rule
|
||||
|
|
@ -46,17 +51,19 @@ Rules are compiled and classified once, at install time. The match functions are
|
|||
O(number of rules); callers must only invoke them on a cache miss.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import re
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
|
||||
verbose_logger: Final = logging.getLogger("LiteLLM")
|
||||
NAME_FIELD: Final = "name"
|
||||
PATTERN_FIELD: Final = "pattern"
|
||||
MODEL_INFO_FIELD: Final = "model_info"
|
||||
PROVIDER_KEY: Final = "litellm_provider"
|
||||
LEGACY_EXTENDS_FIELD: Final = "extends"
|
||||
FILL_MISSING_FOR_PROVIDERS_FIELD: Final = "fill_missing_for_providers"
|
||||
|
||||
|
||||
def _resolve_legacy_extends(rules: list) -> list:
|
||||
|
|
@ -98,11 +105,28 @@ class _RoutingRule:
|
|||
class _CapabilityRule:
|
||||
pattern: re.Pattern
|
||||
model_info: dict
|
||||
fill_missing_for_providers: frozenset[str]
|
||||
|
||||
|
||||
_CompiledRule = _RoutingRule | _CapabilityRule
|
||||
|
||||
|
||||
def _parse_fill_missing_for_providers(rule: Mapping[str, object], pattern_label: object) -> frozenset[str] | None:
|
||||
if FILL_MISSING_FOR_PROVIDERS_FIELD not in rule:
|
||||
return frozenset()
|
||||
raw_fill_missing_for_providers: Final = rule.get(FILL_MISSING_FOR_PROVIDERS_FIELD)
|
||||
if not isinstance(raw_fill_missing_for_providers, (list, tuple)) or not all(
|
||||
isinstance(provider, str) for provider in raw_fill_missing_for_providers
|
||||
):
|
||||
verbose_logger.warning(
|
||||
"LiteLLM: skipping malformed fallback generalization rule %s ('%s' must be a list of provider strings).",
|
||||
rule.get(NAME_FIELD, pattern_label),
|
||||
FILL_MISSING_FOR_PROVIDERS_FIELD,
|
||||
)
|
||||
return None
|
||||
return frozenset(raw_fill_missing_for_providers)
|
||||
|
||||
|
||||
def _compile_rule(rule: object) -> tuple[_CompiledRule, ...]:
|
||||
if not isinstance(rule, dict):
|
||||
return ()
|
||||
|
|
@ -125,8 +149,17 @@ def _compile_rule(rule: object) -> tuple[_CompiledRule, ...]:
|
|||
e,
|
||||
)
|
||||
return ()
|
||||
fill_missing_for_providers: Final = _parse_fill_missing_for_providers(rule, pattern)
|
||||
if fill_missing_for_providers is None:
|
||||
return ()
|
||||
if PROVIDER_KEY not in model_info:
|
||||
return (_CapabilityRule(pattern=compiled, model_info=model_info),)
|
||||
return (
|
||||
_CapabilityRule(
|
||||
pattern=compiled,
|
||||
model_info=model_info,
|
||||
fill_missing_for_providers=fill_missing_for_providers,
|
||||
),
|
||||
)
|
||||
provider: Final = model_info[PROVIDER_KEY]
|
||||
if not isinstance(provider, str):
|
||||
verbose_logger.warning(
|
||||
|
|
@ -140,7 +173,11 @@ def _compile_rule(rule: object) -> tuple[_CompiledRule, ...]:
|
|||
return (_RoutingRule(pattern=compiled, provider=provider),)
|
||||
return (
|
||||
_RoutingRule(pattern=compiled, provider=provider),
|
||||
_CapabilityRule(pattern=compiled, model_info=model_info),
|
||||
_CapabilityRule(
|
||||
pattern=compiled,
|
||||
model_info=model_info,
|
||||
fill_missing_for_providers=fill_missing_for_providers,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -151,6 +188,7 @@ class _FallbackGeneralizations:
|
|||
self.rules: list = []
|
||||
self.routing_rules: tuple = ()
|
||||
self.capability_rules: tuple = ()
|
||||
self.fill_missing_rules: tuple[_CapabilityRule, ...] = ()
|
||||
|
||||
def set_rules(self, rules: list | None) -> None:
|
||||
installed: Final = rules if isinstance(rules, list) else []
|
||||
|
|
@ -158,6 +196,7 @@ class _FallbackGeneralizations:
|
|||
self.rules = installed
|
||||
self.routing_rules = tuple(rule for rule in compiled if isinstance(rule, _RoutingRule))
|
||||
self.capability_rules = tuple(rule for rule in compiled if isinstance(rule, _CapabilityRule))
|
||||
self.fill_missing_rules = tuple(rule for rule in self.capability_rules if rule.fill_missing_for_providers)
|
||||
|
||||
def match_routing(self, model: str) -> str | None:
|
||||
if not model:
|
||||
|
|
@ -175,6 +214,21 @@ class _FallbackGeneralizations:
|
|||
return None
|
||||
return {key: value for model_info in matched for key, value in model_info.items()}
|
||||
|
||||
def match_fill_missing(self, model: str, provider: str) -> Mapping[str, object] | None:
|
||||
if not model or not provider:
|
||||
return None
|
||||
matched = tuple(
|
||||
rule.model_info
|
||||
for rule in self.fill_missing_rules
|
||||
if provider in rule.fill_missing_for_providers and rule.pattern.search(model) is not None
|
||||
)
|
||||
if not matched:
|
||||
return None
|
||||
fill_missing: Final[Mapping[str, object]] = {
|
||||
key: value for model_info in matched for key, value in model_info.items() if key != PROVIDER_KEY
|
||||
}
|
||||
return fill_missing or None
|
||||
|
||||
|
||||
_registry: Final = _FallbackGeneralizations()
|
||||
|
||||
|
|
@ -210,3 +264,14 @@ def match_capability_generalizations(model: str) -> dict | None:
|
|||
capability rule matches. O(number of rules); only call once exact lookups have missed.
|
||||
"""
|
||||
return _registry.match_capabilities(model)
|
||||
|
||||
|
||||
def match_fill_missing_generalizations(model: str, provider: str) -> Mapping[str, object] | None:
|
||||
"""Return flagged capability rules matching ``model`` for ``provider``.
|
||||
|
||||
Later rules override earlier ones on key conflicts. Only rules listing
|
||||
``provider`` in ``fill_missing_for_providers`` contribute. Returns ``None``
|
||||
when no flagged rule matches. O(number of rules); only call once exact
|
||||
lookups have matched.
|
||||
"""
|
||||
return _registry.match_fill_missing(model, provider)
|
||||
|
|
|
|||
|
|
@ -41,6 +41,7 @@ OPTIONAL_KWARGS_KEYS: Final = (
|
|||
"azure_password",
|
||||
"azure_scope",
|
||||
"timeout",
|
||||
"client_side_timeout",
|
||||
"gcs_bucket_name",
|
||||
"bucket_name",
|
||||
"vertex_credentials",
|
||||
|
|
|
|||
|
|
@ -238,6 +238,8 @@ def get_llm_provider(
|
|||
if dynamic_api_key is not None and not isinstance(dynamic_api_key, str):
|
||||
raise Exception(f"dynamic_api_key needs to be a string. Got type={type(dynamic_api_key).__name__}")
|
||||
return model, custom_llm_provider, dynamic_api_key, api_base
|
||||
if "/" in model and is_registered_custom_provider(provider_prefix):
|
||||
return model.split("/", 1)[1], provider_prefix, dynamic_api_key, api_base
|
||||
# check if api base is a known openai compatible endpoint
|
||||
if api_base:
|
||||
for endpoint in litellm.openai_compatible_endpoints:
|
||||
|
|
@ -536,6 +538,10 @@ def get_llm_provider(
|
|||
)
|
||||
|
||||
|
||||
def is_registered_custom_provider(custom_llm_provider: str | None) -> bool:
|
||||
return any(item["provider"] == custom_llm_provider for item in litellm.custom_provider_map)
|
||||
|
||||
|
||||
def _dashscope_family_chat_config(custom_llm_provider: str) -> "litellm.DashScopeChatConfig":
|
||||
if custom_llm_provider == "qwencloud":
|
||||
return litellm.QwenCloudChatConfig()
|
||||
|
|
|
|||
|
|
@ -556,7 +556,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
# ids leaking into a different, later request on the same thread. Sync
|
||||
# support is deferred to a follow-up PR with its own safe-restore
|
||||
# mechanism; async calls (the proxy's only call path) are unaffected.
|
||||
if supports_correlation_logging:
|
||||
if supports_correlation_logging and litellm.request_correlation_in_logs:
|
||||
set_trace_id(self.litellm_trace_id)
|
||||
set_session_id(self.litellm_session_id)
|
||||
# set_trace_id()/set_session_id() sanitize (strip control chars, bound
|
||||
|
|
@ -644,6 +644,24 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
"""Keep ``_response_ms`` / ``litellm_overhead_time_ms`` for a result that has no ``_hidden_params``."""
|
||||
self.response_timing_metrics = dict(timing_metrics) # mutable-ok: kept deep-copyable
|
||||
|
||||
def add_dynamic_callback(self, callback: CustomLogger) -> None:
|
||||
self.dynamic_input_callbacks = self._with_dynamic_callback(self.dynamic_input_callbacks, callback)
|
||||
self.dynamic_success_callbacks = self._with_dynamic_callback(self.dynamic_success_callbacks, callback)
|
||||
self.dynamic_async_success_callbacks = self._with_dynamic_callback(
|
||||
self.dynamic_async_success_callbacks, callback
|
||||
)
|
||||
self.dynamic_failure_callbacks = self._with_dynamic_callback(self.dynamic_failure_callbacks, callback)
|
||||
self.dynamic_async_failure_callbacks = self._with_dynamic_callback(
|
||||
self.dynamic_async_failure_callbacks, callback
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _with_dynamic_callback(
|
||||
callbacks: Sequence[str | Callable | CustomLogger] | None, callback: CustomLogger
|
||||
) -> list[str | Callable | CustomLogger]:
|
||||
existing: Final = tuple(callbacks or ())
|
||||
return [*existing, *(() if callback in existing else (callback,))]
|
||||
|
||||
def process_dynamic_callbacks(self):
|
||||
"""
|
||||
Initializes CustomLogger compatible callbacks in self.dynamic_* callbacks
|
||||
|
|
@ -1973,6 +1991,12 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
self.model_call_details["combined_usage_object"] = usage
|
||||
self.model_call_details["response_cost"] = response_cost
|
||||
|
||||
def record_assembled_response_for_failure(self, assembled: ModelResponse) -> None:
|
||||
"""Bill a fully streamed response on the failure log when a post-call hook rejects it."""
|
||||
usage: Final = getattr(assembled, "usage", None)
|
||||
if isinstance(usage, Usage):
|
||||
self.record_partial_usage_for_failure(usage, self._response_cost_calculator(result=assembled) or 0.0)
|
||||
|
||||
async def dispatch_failure_handlers(
|
||||
self,
|
||||
exception: Exception,
|
||||
|
|
@ -2442,7 +2466,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
call) would leave the outer request's subsequent log lines stamped with
|
||||
the nested call's trace_id/session_id instead of its own.
|
||||
|
||||
Uses a plain set() of the captured pre-call value rather than
|
||||
Uses a plain contextvar set() of the captured pre-call value rather than
|
||||
contextvars.Token-based reset(), since this can end up called from a
|
||||
different asyncio Task/context than __init__ ran in (e.g. the request
|
||||
task's own wrapper() finally block, plus async_success_handler
|
||||
|
|
@ -2453,8 +2477,8 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
that Task's view of the contextvars, so calling it multiple times
|
||||
(once per Task involved in this attempt) is required, not just safe.
|
||||
"""
|
||||
set_trace_id(self._pre_call_trace_id)
|
||||
set_session_id(self._pre_call_session_id)
|
||||
trace_id_var.set(self._pre_call_trace_id)
|
||||
session_id_var.set(self._pre_call_session_id)
|
||||
|
||||
def _restore_correlation_context_if_unclaimed(self) -> None:
|
||||
"""Guarded variant for __del__-triggered cleanup only.
|
||||
|
|
|
|||
|
|
@ -9,6 +9,8 @@ from types import MappingProxyType
|
|||
from typing import Any, Final, Literal, TypedDict, cast
|
||||
from zoneinfo import ZoneInfo, ZoneInfoNotFoundError
|
||||
|
||||
from typing_extensions import ReadOnly
|
||||
|
||||
import litellm
|
||||
from litellm._internal_context import current_billing_time
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -482,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(
|
||||
|
|
@ -504,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,
|
||||
)
|
||||
|
||||
|
|
@ -530,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)
|
||||
"""
|
||||
|
|
@ -552,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)
|
||||
|
||||
|
|
@ -637,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)
|
||||
|
|
@ -663,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,
|
||||
|
|
@ -680,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,
|
||||
),
|
||||
|
|
@ -772,6 +767,7 @@ def calculate_cache_writing_cost(
|
|||
|
||||
class PromptTokensDetailsResult(TypedDict):
|
||||
cache_hit_tokens: int
|
||||
cache_hit_audio_tokens: ReadOnly[int]
|
||||
cache_creation_tokens: int
|
||||
cache_creation_token_details: CacheCreationTokenDetails | None
|
||||
text_tokens: int
|
||||
|
|
@ -802,12 +798,34 @@ def parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult:
|
|||
)
|
||||
or None
|
||||
)
|
||||
text_tokens: Final = (
|
||||
cast(int | None, getattr(usage.prompt_tokens_details, "text_tokens", None))
|
||||
or 0 # default to prompt tokens, if this field is not set
|
||||
cached_tokens_details: Final = getattr(usage.prompt_tokens_details, "cached_tokens_details", None)
|
||||
cached_audio_tokens: Final = min(
|
||||
_get_token_detail_value(cached_tokens_details, "audio_tokens") or 0, cache_hit_tokens
|
||||
)
|
||||
cached_text_tokens: Final = min(
|
||||
_get_token_detail_value(cached_tokens_details, "text_tokens") or 0,
|
||||
cache_hit_tokens - cached_audio_tokens,
|
||||
)
|
||||
cached_image_tokens: Final = min(
|
||||
_get_token_detail_value(cached_tokens_details, "image_tokens") or 0,
|
||||
cache_hit_tokens - cached_audio_tokens - cached_text_tokens,
|
||||
)
|
||||
text_tokens: Final = max(
|
||||
(
|
||||
cast(int | None, getattr(usage.prompt_tokens_details, "text_tokens", None))
|
||||
or 0 # default to prompt tokens, if this field is not set
|
||||
)
|
||||
- cached_text_tokens,
|
||||
0,
|
||||
)
|
||||
audio_tokens: Final = max(
|
||||
(cast(int | None, getattr(usage.prompt_tokens_details, "audio_tokens", 0)) or 0) - cached_audio_tokens,
|
||||
0,
|
||||
)
|
||||
image_tokens: Final = max(
|
||||
(cast(int | None, getattr(usage.prompt_tokens_details, "image_tokens", 0)) or 0) - cached_image_tokens,
|
||||
0,
|
||||
)
|
||||
audio_tokens: Final = cast(int | None, getattr(usage.prompt_tokens_details, "audio_tokens", 0)) or 0
|
||||
image_tokens: Final = cast(int | None, getattr(usage.prompt_tokens_details, "image_tokens", 0)) or 0
|
||||
video_tokens: Final = _coerce_token_count(getattr(usage.prompt_tokens_details, "video_tokens", 0))
|
||||
character_count: Final = (
|
||||
cast(
|
||||
|
|
@ -835,6 +853,7 @@ def parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult:
|
|||
|
||||
return PromptTokensDetailsResult(
|
||||
cache_hit_tokens=cache_hit_tokens,
|
||||
cache_hit_audio_tokens=cached_audio_tokens,
|
||||
cache_creation_tokens=cache_creation_tokens,
|
||||
cache_creation_token_details=cache_creation_token_details,
|
||||
text_tokens=text_tokens,
|
||||
|
|
@ -918,15 +937,28 @@ def _calculate_input_cost(
|
|||
prompt_cost = float(prompt_tokens_details["text_tokens"]) * prompt_base_cost
|
||||
|
||||
### CACHE READ COST - Now uses tiered pricing
|
||||
prompt_cost += float(prompt_tokens_details["cache_hit_tokens"]) * cache_read_cost
|
||||
cache_hit_audio_tokens: Final = prompt_tokens_details["cache_hit_audio_tokens"]
|
||||
audio_cache_read_rate: Final = _get_cost_per_unit(
|
||||
model_info,
|
||||
_get_service_tier_cost_key("cache_read_input_audio_token_cost", service_tier),
|
||||
None,
|
||||
)
|
||||
prompt_cost += float(prompt_tokens_details["cache_hit_tokens"] - cache_hit_audio_tokens) * cache_read_cost
|
||||
prompt_cost += float(cache_hit_audio_tokens) * (
|
||||
audio_cache_read_rate if audio_cache_read_rate is not None else cache_read_cost
|
||||
)
|
||||
|
||||
### AUDIO COST
|
||||
if prompt_tokens_details["audio_tokens"]:
|
||||
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"
|
||||
|
|
@ -935,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"
|
||||
|
|
@ -1149,6 +1183,7 @@ def generic_cost_per_token(
|
|||
### PROCESSING COST
|
||||
prompt_tokens_details = PromptTokensDetailsResult(
|
||||
cache_hit_tokens=0,
|
||||
cache_hit_audio_tokens=0,
|
||||
cache_creation_tokens=0,
|
||||
cache_creation_token_details=None,
|
||||
text_tokens=usage.prompt_tokens,
|
||||
|
|
@ -1319,6 +1354,7 @@ class BilledTokenRates:
|
|||
input_cost_per_token: float
|
||||
output_cost_per_token: float
|
||||
cache_read_input_token_cost: float
|
||||
cache_read_input_audio_token_cost: float
|
||||
cache_creation_input_token_cost: float
|
||||
cache_creation_input_token_cost_above_1hr: float
|
||||
output_cost_per_reasoning_token: float
|
||||
|
|
@ -1330,6 +1366,7 @@ class BilledTokenRates:
|
|||
input_cost_per_token=self.input_cost_per_token * multiplier,
|
||||
output_cost_per_token=self.output_cost_per_token * multiplier,
|
||||
cache_read_input_token_cost=self.cache_read_input_token_cost * multiplier,
|
||||
cache_read_input_audio_token_cost=self.cache_read_input_audio_token_cost * multiplier,
|
||||
cache_creation_input_token_cost=self.cache_creation_input_token_cost * multiplier,
|
||||
cache_creation_input_token_cost_above_1hr=self.cache_creation_input_token_cost_above_1hr * multiplier,
|
||||
output_cost_per_reasoning_token=self.output_cost_per_reasoning_token * multiplier,
|
||||
|
|
@ -1353,15 +1390,16 @@ def _reasoning_token_count(usage: Usage) -> int:
|
|||
return parsed or _coerce_token_count(getattr(usage, "reasoning_tokens", 0))
|
||||
|
||||
|
||||
def _cache_token_counts(usage: Usage) -> tuple[int, int, CacheCreationTokenDetails | None]:
|
||||
"""(cache read tokens, cache creation tokens, cache creation details): read from prompt_tokens_details
|
||||
first, then the private top-level counters the Usage constructor mirrors cache tokens onto for
|
||||
providers/callers that bypass the details."""
|
||||
def _cache_token_counts(usage: Usage) -> tuple[int, int, int, CacheCreationTokenDetails | None]:
|
||||
"""(cache read tokens, cached audio tokens, cache creation tokens, cache creation details): read from
|
||||
prompt_tokens_details first, then the private top-level counters the Usage constructor mirrors cache
|
||||
tokens onto for providers/callers that bypass the details."""
|
||||
parsed: Final = parse_prompt_tokens_details(usage) if usage.prompt_tokens_details is not None else None
|
||||
parsed_read: Final = parsed["cache_hit_tokens"] if parsed is not None else 0
|
||||
parsed_creation: Final = parsed["cache_creation_tokens"] if parsed is not None else 0
|
||||
return (
|
||||
parsed_read or _coerce_token_count(getattr(usage, "_cache_read_input_tokens", 0)),
|
||||
parsed["cache_hit_audio_tokens"] if parsed is not None else 0,
|
||||
parsed_creation or _coerce_token_count(getattr(usage, "_cache_creation_input_tokens", 0)),
|
||||
parsed["cache_creation_token_details"] if parsed is not None else None,
|
||||
)
|
||||
|
|
@ -1372,11 +1410,13 @@ def _custom_pricing_rates(custom_cost_per_token: CostPerToken) -> BilledTokenRat
|
|||
cache rates (else the input rate) and reasoning at the output rate, as _cost_per_token_custom_pricing_helper does."""
|
||||
input_rate: Final = custom_cost_per_token["input_cost_per_token"]
|
||||
output_rate: Final = custom_cost_per_token["output_cost_per_token"]
|
||||
cache_read_rate: Final = custom_cost_per_token.get("cache_read_input_token_cost", input_rate)
|
||||
cache_creation_rate: Final = custom_cost_per_token.get("cache_creation_input_token_cost", input_rate)
|
||||
return BilledTokenRates(
|
||||
input_cost_per_token=input_rate,
|
||||
output_cost_per_token=output_rate,
|
||||
cache_read_input_token_cost=custom_cost_per_token.get("cache_read_input_token_cost", input_rate),
|
||||
cache_read_input_token_cost=cache_read_rate,
|
||||
cache_read_input_audio_token_cost=cache_read_rate,
|
||||
cache_creation_input_token_cost=cache_creation_rate,
|
||||
cache_creation_input_token_cost_above_1hr=cache_creation_rate,
|
||||
output_cost_per_reasoning_token=output_rate,
|
||||
|
|
@ -1413,6 +1453,11 @@ def _cost_map_billed_rates(
|
|||
completion_base_cost=completion_base_cost,
|
||||
current_time=billing_time,
|
||||
)
|
||||
audio_cache_read_rate: Final = _get_cost_per_unit(
|
||||
model_info,
|
||||
_get_service_tier_cost_key("cache_read_input_audio_token_cost", service_tier),
|
||||
None,
|
||||
)
|
||||
multiplier: Final = (
|
||||
_get_regional_uplift_multiplier(model_info, data_residency)
|
||||
* get_vertex_regional_endpoint_uplift(model_info, vertex_location)
|
||||
|
|
@ -1422,6 +1467,9 @@ def _cost_map_billed_rates(
|
|||
input_cost_per_token=prompt_base_cost,
|
||||
output_cost_per_token=completion_base_cost,
|
||||
cache_read_input_token_cost=cache_read_cost_rate,
|
||||
cache_read_input_audio_token_cost=(
|
||||
audio_cache_read_rate if audio_cache_read_rate is not None else cache_read_cost_rate
|
||||
),
|
||||
cache_creation_input_token_cost=cache_creation_cost_rate,
|
||||
cache_creation_input_token_cost_above_1hr=cache_creation_cost_above_1hr_rate,
|
||||
output_cost_per_reasoning_token=reasoning_rate,
|
||||
|
|
@ -1494,7 +1542,9 @@ def get_token_type_cost_breakdown(
|
|||
if rates is None:
|
||||
return TokenTypeCostBreakdown(0.0, 0.0, 0.0)
|
||||
|
||||
cache_read_tokens, cache_creation_tokens, cache_creation_token_details = _cache_token_counts(usage)
|
||||
cache_read_tokens, cached_audio_tokens, cache_creation_tokens, cache_creation_token_details = _cache_token_counts(
|
||||
usage
|
||||
)
|
||||
cache_creation_cost: Final = (
|
||||
float(cache_creation_tokens) * rates.cache_creation_input_token_cost
|
||||
if custom_cost_per_token is not None
|
||||
|
|
@ -1507,7 +1557,10 @@ def get_token_type_cost_breakdown(
|
|||
)
|
||||
return TokenTypeCostBreakdown(
|
||||
reasoning_cost=float(_reasoning_token_count(usage)) * rates.output_cost_per_reasoning_token,
|
||||
cache_read_cost=float(cache_read_tokens) * rates.cache_read_input_token_cost,
|
||||
cache_read_cost=(
|
||||
float(cache_read_tokens - cached_audio_tokens) * rates.cache_read_input_token_cost
|
||||
+ float(cached_audio_tokens) * rates.cache_read_input_audio_token_cost
|
||||
),
|
||||
cache_creation_cost=cache_creation_cost,
|
||||
rates=rates,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1771,7 +1771,7 @@ def convert_to_anthropic_tool_invoke(
|
|||
anthropic_tool_invoke: Final[list[AnthropicMessagesToolUseParam | dict[str, object]]] = []
|
||||
|
||||
for tool in tool_calls:
|
||||
if not get_attribute_or_key(tool, "type") == "function":
|
||||
if get_attribute_or_key(tool, "type") != "function":
|
||||
continue
|
||||
|
||||
tool_id = cast(str, get_attribute_or_key(tool, "id"))
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ from typing_extensions import ReadOnly
|
|||
|
||||
import litellm
|
||||
from litellm._logging import redact_internal_details_from_client_message, verbose_logger
|
||||
from litellm.constants import REALTIME_SESSION_FAILURE_LOGGED_KEY, REALTIME_SESSION_SUCCESS_LOGGED_KEY
|
||||
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
||||
from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig
|
||||
from litellm.types.llms.openai import (
|
||||
|
|
@ -35,9 +36,6 @@ else:
|
|||
CLIENT_CONNECTION_CLASS = Any
|
||||
|
||||
|
||||
REALTIME_SESSION_SUCCESS_LOGGED_KEY: Final = "realtime_session_success_logged"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class BackendClose:
|
||||
code: int
|
||||
|
|
@ -1153,6 +1151,7 @@ class RealTimeStreaming:
|
|||
self._logging_worker.ensure_initialized_and_enqueue(
|
||||
self.logging_obj.dispatch_failure_handlers(error, traceback.format_exc(), prefer_async_handlers=True)
|
||||
)
|
||||
self.logging_obj.model_call_details[REALTIME_SESSION_FAILURE_LOGGED_KEY] = True
|
||||
|
||||
@staticmethod
|
||||
def _detect_beta_header(websocket: ScopedWebSocket) -> bool:
|
||||
|
|
|
|||
|
|
@ -19,6 +19,7 @@ from typing_extensions import NotRequired, TypedDict
|
|||
import litellm
|
||||
from litellm import verbose_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.litellm_core_utils.asyncify import asyncify
|
||||
from litellm.litellm_core_utils.model_response_utils import (
|
||||
is_model_response_stream_empty,
|
||||
)
|
||||
|
|
@ -2247,7 +2248,7 @@ class CustomStreamWrapper:
|
|||
if self.sent_last_chunk is True:
|
||||
# log the final chunk with accurate streaming values
|
||||
try:
|
||||
complete_streaming_response = litellm.stream_chunk_builder(
|
||||
complete_streaming_response = await asyncify(litellm.stream_chunk_builder)(
|
||||
chunks=self.chunks,
|
||||
messages=self.messages,
|
||||
logging_obj=self.logging_obj,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ from collections.abc import Awaitable, Callable
|
|||
from typing import TYPE_CHECKING, Final, TypeAlias
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.litellm_core_utils.asyncify import asyncify
|
||||
from litellm.types.llms.anthropic import AppliedEdit
|
||||
|
||||
from .constants import CLEAR_TOOL_USES_EDIT_TYPE, COMPACT_EDIT_TYPE
|
||||
|
|
@ -82,9 +83,9 @@ async def apply_context_management(
|
|||
"""Run edits in order; return a single ``PolyfillResult``.
|
||||
|
||||
The dispatcher is async so async editors (``compact_20260112``) can
|
||||
``await`` the configured summarization model. Sync editors are called
|
||||
inline — ``inspect.iscoroutinefunction`` decides how each editor is
|
||||
invoked.
|
||||
``await`` the configured summarization model. Sync editors run in a
|
||||
worker thread so their token counts stay off the event loop;
|
||||
``inspect.iscoroutinefunction`` decides how each editor is invoked.
|
||||
"""
|
||||
edits: Final = _normalize_spec(context_management_spec)
|
||||
if not edits:
|
||||
|
|
@ -121,7 +122,7 @@ async def apply_context_management(
|
|||
user_api_key_auth=user_api_key_auth,
|
||||
)
|
||||
if editor_is_async
|
||||
else editor(
|
||||
else await asyncify(editor)(
|
||||
model=model,
|
||||
messages=current_messages,
|
||||
tools=tools,
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict, Unpack
|
|||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.litellm_core_utils.asyncify import asyncify
|
||||
from litellm.types.llms.anthropic import (
|
||||
AppliedEdit,
|
||||
CompactionBlock,
|
||||
|
|
@ -1157,7 +1158,7 @@ async def apply_compact_20260112(
|
|||
|
||||
# Phase B: threshold check.
|
||||
try:
|
||||
current_tokens = _count_effective_tokens(
|
||||
current_tokens = await asyncify(_count_effective_tokens)(
|
||||
model=model,
|
||||
effective_messages=effective_messages,
|
||||
# ``augmented_system`` already carries the prior compaction summary
|
||||
|
|
|
|||
|
|
@ -678,7 +678,7 @@ class BaseAnthropicMessagesStreamingIterator:
|
|||
"""
|
||||
from litellm.proxy.pass_through_endpoints.streaming_handler import PassThroughStreamingHandler
|
||||
|
||||
PassThroughStreamingHandler.schedule_stream_failure_logging(
|
||||
await PassThroughStreamingHandler.schedule_stream_failure_logging(
|
||||
litellm_logging_obj=self.litellm_logging_obj,
|
||||
endpoint_type=EndpointType.ANTHROPIC,
|
||||
request_body=self.request_body,
|
||||
|
|
|
|||
|
|
@ -42,6 +42,10 @@ DROP_UNFITTING_REASONING_EFFORT_WARNING: Final = (
|
|||
)
|
||||
|
||||
|
||||
def _messages_carry_output_config(messages: Sequence[object]) -> bool:
|
||||
return any(isinstance(message, Mapping) and "output_config" in message for message in messages)
|
||||
|
||||
|
||||
class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
|
||||
@property
|
||||
def custom_llm_provider(self) -> str | None:
|
||||
|
|
@ -331,6 +335,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
|
|||
headers = self._update_headers_with_anthropic_beta(
|
||||
headers=headers,
|
||||
optional_params=optional_params,
|
||||
messages=messages,
|
||||
)
|
||||
|
||||
return headers, api_base
|
||||
|
|
@ -664,6 +669,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
|
|||
headers: dict,
|
||||
optional_params: dict,
|
||||
custom_llm_provider: str = "anthropic",
|
||||
messages: Sequence[object] = (),
|
||||
) -> dict:
|
||||
"""
|
||||
Auto-inject anthropic-beta headers based on features used.
|
||||
|
|
@ -673,24 +679,30 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
|
|||
- tool_search: adds provider-specific tool search header
|
||||
- output_format: adds 'structured-outputs-2025-11-13'
|
||||
- speed: adds 'fast-mode-2026-02-01'
|
||||
- a message carrying output_config: adds 'per-turn-control-2026-07-01'
|
||||
|
||||
Args:
|
||||
headers: Request headers dict
|
||||
optional_params: Optional parameters including tools, context_management, output_format, speed
|
||||
custom_llm_provider: Provider name for looking up correct tool search header
|
||||
messages: Request messages, scanned for per-message output_config
|
||||
"""
|
||||
beta_values: Final[set] = set()
|
||||
|
||||
# Get existing beta headers if any
|
||||
existing_beta: Final = headers.get("anthropic-beta")
|
||||
if existing_beta:
|
||||
beta_values.update(b.strip() for b in existing_beta.split(","))
|
||||
existing_beta: Final = tuple(
|
||||
piece.strip()
|
||||
for key, value in headers.items()
|
||||
if key.lower() == "anthropic-beta"
|
||||
for piece in value.split(",")
|
||||
if piece.strip()
|
||||
)
|
||||
beta_values.update(existing_beta)
|
||||
|
||||
# Check for context management
|
||||
context_management_param: Final = optional_params.get("context_management")
|
||||
if context_management_param is not None:
|
||||
# Check edits array for compact_20260112 type
|
||||
edits: Final = context_management_param.get("edits", [])
|
||||
edits: Final = context_management_param.get("edits", ())
|
||||
has_compact = False
|
||||
has_other = False
|
||||
|
||||
|
|
@ -722,24 +734,18 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
|
|||
if optional_params.get("speed") == "fast":
|
||||
beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.FAST_MODE_2026_02_01.value)
|
||||
|
||||
# Check for advisor tool
|
||||
tools = optional_params.get("tools")
|
||||
if tools:
|
||||
for tool in tools:
|
||||
if isinstance(tool, dict) and tool.get("type") == ANTHROPIC_ADVISOR_TOOL_TYPE:
|
||||
beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.ADVISOR_TOOL_2026_03_01.value)
|
||||
break
|
||||
if _messages_carry_output_config(messages):
|
||||
beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.PER_TURN_CONTROL_2026_07_01.value)
|
||||
|
||||
# Check for tool search tools
|
||||
tools = optional_params.get("tools")
|
||||
if tools:
|
||||
anthropic_model_info: Final = AnthropicModelInfo()
|
||||
if anthropic_model_info.is_tool_search_used(tools):
|
||||
# Use provider-specific tool search header
|
||||
tool_search_header: Final = get_tool_search_beta_header(custom_llm_provider)
|
||||
beta_values.add(tool_search_header)
|
||||
tools: Final = optional_params.get("tools")
|
||||
if any(isinstance(tool, dict) and tool.get("type") == ANTHROPIC_ADVISOR_TOOL_TYPE for tool in tools or ()):
|
||||
beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.ADVISOR_TOOL_2026_03_01.value)
|
||||
|
||||
if beta_values:
|
||||
headers["anthropic-beta"] = ",".join(sorted(beta_values))
|
||||
if AnthropicModelInfo().is_tool_search_used(tools):
|
||||
beta_values.add(get_tool_search_beta_header(custom_llm_provider))
|
||||
|
||||
return headers
|
||||
if not beta_values:
|
||||
return headers
|
||||
merged: Final = {key: value for key, value in headers.items() if key.lower() != "anthropic-beta"}
|
||||
merged["anthropic-beta"] = ",".join(sorted(beta_values))
|
||||
return merged
|
||||
|
|
|
|||
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,
|
||||
)
|
||||
|
|
@ -68,6 +68,7 @@ class AzureAnthropicMessagesConfig(AnthropicMessagesConfig):
|
|||
headers = self._update_headers_with_anthropic_beta(
|
||||
headers=headers,
|
||||
optional_params=optional_params,
|
||||
messages=messages,
|
||||
)
|
||||
|
||||
return headers, api_base
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -46,6 +46,7 @@ class BedrockClaudePlatformMessagesConfig(BedrockClaudePlatformMixin, AnthropicM
|
|||
headers = self._update_headers_with_anthropic_beta(
|
||||
headers=headers,
|
||||
optional_params=optional_params,
|
||||
messages=messages,
|
||||
)
|
||||
|
||||
return headers, api_base
|
||||
|
|
|
|||
|
|
@ -7,13 +7,21 @@ This uses aws_sdk_bedrock_runtime for bidirectional streaming with Nova Sonic.
|
|||
import asyncio
|
||||
import contextlib
|
||||
import json
|
||||
from collections.abc import AsyncIterator, Mapping
|
||||
from typing import Final, Protocol
|
||||
from collections.abc import AsyncIterator, Mapping, MutableMapping
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
from typing import Final, NoReturn, Protocol
|
||||
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm._logging import _redact_string, verbose_proxy_logger
|
||||
from litellm.constants import (
|
||||
BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY,
|
||||
BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY,
|
||||
BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY,
|
||||
REALTIME_SESSION_SUCCESS_LOGGED_KEY,
|
||||
)
|
||||
from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
|
||||
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
||||
|
|
@ -28,6 +36,32 @@ from .transformation import BedrockRealtimeConfig
|
|||
_CLIENT_MODALITIES_ADAPTER: Final[TypeAdapter["list[str] | None"]] = TypeAdapter(list[str] | None)
|
||||
_CLIENT_MESSAGE_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
|
||||
|
||||
_EMPTY_JSON_OBJECT: Final[Mapping[str, JsonValue]] = MappingProxyType({})
|
||||
|
||||
_BEDROCK_STREAM_ERROR_STATUS: Final[Mapping[str, int]] = MappingProxyType(
|
||||
{
|
||||
"AccessDeniedException": 403,
|
||||
"ConflictException": 400,
|
||||
"InternalServerException": 500,
|
||||
"ModelErrorException": 424,
|
||||
"ModelNotReadyException": 429,
|
||||
"ModelStreamErrorException": 424,
|
||||
"ModelTimeoutException": 408,
|
||||
"ResourceNotFoundException": 404,
|
||||
"ServiceQuotaExceededException": 400,
|
||||
"ServiceUnavailableException": 503,
|
||||
"ThrottlingException": 429,
|
||||
"ValidationException": 400,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _as_bedrock_error(error: BaseException) -> BaseException:
|
||||
status_code: Final = _BEDROCK_STREAM_ERROR_STATUS.get(type(error).__name__)
|
||||
if status_code is None:
|
||||
return error
|
||||
return BedrockError(status_code=status_code, message=f"{type(error).__name__}: {error}")
|
||||
|
||||
|
||||
def _json_dict(value: JsonValue) -> dict[str, JsonValue]:
|
||||
return value if isinstance(value, dict) else {}
|
||||
|
|
@ -51,6 +85,8 @@ def _should_log_event(openai_message: Mapping[str, object]) -> bool:
|
|||
class RealtimeClientWebSocket(Protocol):
|
||||
"""The client-facing websocket surface the realtime bridge talks to."""
|
||||
|
||||
scope: MutableMapping[str, object] # mutable-ok: the ASGI scope is the per-connection state store
|
||||
|
||||
async def receive_text(self) -> str: ...
|
||||
|
||||
async def send_text(self, data: str) -> None: ...
|
||||
|
|
@ -85,6 +121,81 @@ class BedrockBidirectionalStream(Protocol):
|
|||
async def await_output(self) -> tuple[object, BedrockOutputStream]: ...
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _BridgeOutcome:
|
||||
logged_events: tuple[OpenAIRealtimeEvents, ...]
|
||||
provider_failure: BaseException | None
|
||||
client_disconnected: bool
|
||||
|
||||
|
||||
async def _client_messages(client_ws: RealtimeClientWebSocket, initial_message: str | None) -> AsyncIterator[str]:
|
||||
if initial_message is not None:
|
||||
yield initial_message
|
||||
while True:
|
||||
try:
|
||||
yield await client_ws.receive_text()
|
||||
except Exception as e: # noqa: BLE001 # any receive failure means the client is gone
|
||||
verbose_proxy_logger.debug("Client to Bedrock forwarding ended: %s", e, exc_info=True)
|
||||
return
|
||||
|
||||
|
||||
def _pending_session_update(scope: Mapping[str, object]) -> str | None:
|
||||
"""A fallback attempt on the same websocket replays the session.update the failed attempt never acked."""
|
||||
if scope.get(BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY) is True:
|
||||
committed_failure: Final = scope.get(BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY)
|
||||
raise BedrockError(
|
||||
status_code=400,
|
||||
message=(
|
||||
"Bedrock realtime session already committed to a provider stream; it cannot be replayed"
|
||||
+ (f". The committed stream failed with: {committed_failure}" if committed_failure else "")
|
||||
),
|
||||
)
|
||||
pending: Final = scope.get(BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY)
|
||||
return pending if isinstance(pending, str) else None
|
||||
|
||||
|
||||
def _raise_provider_failure(scope: MutableMapping[str, object], failure: BaseException) -> NoReturn:
|
||||
error: Final = _as_bedrock_error(failure)
|
||||
verbose_proxy_logger.error("Bedrock Realtime: provider stream failed: %s", _redact_string(str(error)))
|
||||
if scope.get(BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY) is True:
|
||||
scope[BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY] = _redact_string(str(error))
|
||||
raise error from failure
|
||||
|
||||
|
||||
def _parse_client_message(message: str) -> Mapping[str, JsonValue]:
|
||||
try:
|
||||
return _json_dict(_CLIENT_MESSAGE_ADAPTER.validate_json(message))
|
||||
except ValueError:
|
||||
return _EMPTY_JSON_OBJECT
|
||||
|
||||
|
||||
async def _ack_session_update(
|
||||
client_ws: RealtimeClientWebSocket,
|
||||
bedrock_stream: BedrockBidirectionalStream,
|
||||
transformation_config: BedrockRealtimeConfig,
|
||||
model: str,
|
||||
logging_obj: LiteLLMLogging | None,
|
||||
parsed_client_message: Mapping[str, JsonValue],
|
||||
) -> bool:
|
||||
"""Ack the client's session.update once Bedrock accepted the stream; False means the client is gone."""
|
||||
await bedrock_stream.await_output()
|
||||
client_ws.scope.pop(BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY, None)
|
||||
client_ws.scope[BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY] = True # rebind-ok: scope outlives the attempt
|
||||
if logging_obj is None:
|
||||
return True
|
||||
requested_modalities: Final = _CLIENT_MODALITIES_ADAPTER.validate_python(
|
||||
_json_dict(parsed_client_message.get("session")).get("modalities")
|
||||
)
|
||||
try:
|
||||
await client_ws.send_text(
|
||||
json.dumps(transformation_config.session_updated_event(model, logging_obj, requested_modalities))
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # any send failure means the client is gone
|
||||
verbose_proxy_logger.debug("Client to Bedrock forwarding ended: %s", e, exc_info=True)
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
class BedrockRealtime(BaseAWSLLM):
|
||||
"""Handler for Bedrock Nova Sonic realtime speech-to-speech API."""
|
||||
|
||||
|
|
@ -132,6 +243,8 @@ class BedrockRealtime(BaseAWSLLM):
|
|||
except ImportError:
|
||||
raise ImportError("Missing aws_sdk_bedrock_runtime. Install with: pip install aws-sdk-bedrock-runtime")
|
||||
|
||||
pending_session_update: Final = _pending_session_update(websocket.scope)
|
||||
|
||||
# Get AWS region
|
||||
if aws_region_name is None:
|
||||
optional_params: Final = {
|
||||
|
|
@ -190,90 +303,106 @@ class BedrockRealtime(BaseAWSLLM):
|
|||
|
||||
transformation_config: Final = BedrockRealtimeConfig()
|
||||
|
||||
try:
|
||||
# Initialize the bidirectional stream
|
||||
bedrock_stream: Final = await open_bidirectional_stream()
|
||||
bedrock_stream: Final = await open_bidirectional_stream()
|
||||
|
||||
verbose_proxy_logger.debug("Bedrock Realtime: Bidirectional stream established")
|
||||
verbose_proxy_logger.debug("Bedrock Realtime: Bidirectional stream established")
|
||||
|
||||
if pending_session_update is None:
|
||||
await websocket.send_text(json.dumps(transformation_config.session_created_event(model, logging_obj)))
|
||||
verbose_proxy_logger.debug("Bedrock Realtime: sent session.created to client on connect")
|
||||
|
||||
# Track state for transformation
|
||||
session_state: Final[RealtimeResponseTransformInput] = {
|
||||
"current_output_item_id": None,
|
||||
"current_response_id": None,
|
||||
"current_conversation_id": None,
|
||||
"current_delta_chunks": None,
|
||||
"current_item_chunks": None,
|
||||
"current_delta_type": None,
|
||||
"session_configuration_request": None,
|
||||
}
|
||||
# Track state for transformation
|
||||
session_state: Final[RealtimeResponseTransformInput] = {
|
||||
"current_output_item_id": None,
|
||||
"current_response_id": None,
|
||||
"current_conversation_id": None,
|
||||
"current_delta_chunks": None,
|
||||
"current_item_chunks": None,
|
||||
"current_delta_type": None,
|
||||
"session_configuration_request": None,
|
||||
}
|
||||
|
||||
# Create tasks for bidirectional forwarding
|
||||
client_to_bedrock_task: Final = asyncio.create_task(
|
||||
self._forward_client_to_bedrock(
|
||||
websocket,
|
||||
bedrock_stream,
|
||||
transformation_config,
|
||||
model,
|
||||
session_state,
|
||||
logging_obj,
|
||||
outcome: Final = await self._bridge(
|
||||
websocket,
|
||||
bedrock_stream,
|
||||
transformation_config,
|
||||
model,
|
||||
session_state,
|
||||
logging_obj,
|
||||
initial_message=pending_session_update,
|
||||
)
|
||||
|
||||
logged_events: Final = (
|
||||
*outcome.logged_events,
|
||||
*(
|
||||
leftover_event
|
||||
for leftover_event in transformation_config.leftover_usage_done_events()
|
||||
if _should_log_event(leftover_event)
|
||||
),
|
||||
)
|
||||
if logged_events:
|
||||
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(
|
||||
logging_obj.dispatch_success_handlers(
|
||||
list(logged_events), # mutable-ok: realtime spend logging requires a list result
|
||||
prefer_async_handlers=True,
|
||||
)
|
||||
)
|
||||
logging_obj.model_call_details[REALTIME_SESSION_SUCCESS_LOGGED_KEY] = True
|
||||
|
||||
async def forward_bedrock_and_collect_logged_events() -> tuple[OpenAIRealtimeEvents, ...]:
|
||||
return tuple(
|
||||
[
|
||||
event
|
||||
async for event in self._forward_bedrock_to_client(
|
||||
bedrock_stream,
|
||||
websocket,
|
||||
transformation_config,
|
||||
model,
|
||||
logging_obj,
|
||||
session_state,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
bedrock_to_client_task: Final = asyncio.create_task(forward_bedrock_and_collect_logged_events())
|
||||
|
||||
# Wait for both tasks to complete
|
||||
await asyncio.gather(
|
||||
client_to_bedrock_task,
|
||||
bedrock_to_client_task,
|
||||
return_exceptions=True,
|
||||
if outcome.provider_failure is None:
|
||||
return
|
||||
if outcome.client_disconnected:
|
||||
verbose_proxy_logger.debug(
|
||||
"Bedrock Realtime: stream failed after the client disconnected: %s", outcome.provider_failure
|
||||
)
|
||||
return
|
||||
_raise_provider_failure(websocket.scope, outcome.provider_failure)
|
||||
|
||||
forwarded_logged_events: Final = (
|
||||
bedrock_to_client_task.result()
|
||||
if not bedrock_to_client_task.cancelled() and bedrock_to_client_task.exception() is None
|
||||
else ()
|
||||
)
|
||||
logged_events: Final = (
|
||||
*forwarded_logged_events,
|
||||
*(
|
||||
leftover_event
|
||||
for leftover_event in transformation_config.leftover_usage_done_events()
|
||||
if _should_log_event(leftover_event)
|
||||
),
|
||||
)
|
||||
if logged_events:
|
||||
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(
|
||||
logging_obj.dispatch_success_handlers(
|
||||
list(logged_events), # mutable-ok: realtime spend logging requires a list result
|
||||
prefer_async_handlers=True,
|
||||
)
|
||||
)
|
||||
async def _bridge(
|
||||
self,
|
||||
websocket: RealtimeClientWebSocket,
|
||||
bedrock_stream: BedrockBidirectionalStream,
|
||||
transformation_config: BedrockRealtimeConfig,
|
||||
model: str,
|
||||
session_state: RealtimeResponseTransformInput,
|
||||
logging_obj: LiteLLMLogging,
|
||||
initial_message: str | None,
|
||||
) -> _BridgeOutcome:
|
||||
"""Run both forwarding directions until the client leaves or either side fails."""
|
||||
logged: Final[list[OpenAIRealtimeEvents]] = [] # mutable-ok: events forwarded before a failure are still spend
|
||||
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception("Error in BedrockRealtime.async_realtime: %s", e)
|
||||
try:
|
||||
await websocket.close(code=1011, reason=_redact_string(f"Internal error: {e}"))
|
||||
except Exception:
|
||||
pass
|
||||
raise
|
||||
async def collect_logged_events() -> None:
|
||||
async for event in self._forward_bedrock_to_client(
|
||||
bedrock_stream, websocket, transformation_config, model, logging_obj, session_state
|
||||
):
|
||||
logged.append(event)
|
||||
|
||||
client_task: Final = asyncio.create_task(
|
||||
self._forward_client_to_bedrock(
|
||||
websocket, bedrock_stream, transformation_config, model, session_state, logging_obj, initial_message
|
||||
)
|
||||
)
|
||||
bedrock_task: Final = asyncio.create_task(collect_logged_events())
|
||||
|
||||
await asyncio.wait((client_task, bedrock_task), return_when=asyncio.FIRST_COMPLETED)
|
||||
client_disconnected: Final = (
|
||||
client_task.done() and not client_task.cancelled() and client_task.exception() is None
|
||||
)
|
||||
client_task.cancel()
|
||||
bedrock_task.cancel()
|
||||
client_outcome, bedrock_outcome = await asyncio.gather(client_task, bedrock_task, return_exceptions=True)
|
||||
|
||||
return _BridgeOutcome(
|
||||
logged_events=tuple(logged),
|
||||
provider_failure=(
|
||||
client_outcome
|
||||
if isinstance(client_outcome, Exception)
|
||||
else bedrock_outcome
|
||||
if isinstance(bedrock_outcome, Exception)
|
||||
else None
|
||||
),
|
||||
client_disconnected=client_disconnected,
|
||||
)
|
||||
|
||||
async def _forward_client_to_bedrock(
|
||||
self,
|
||||
|
|
@ -283,8 +412,12 @@ class BedrockRealtime(BaseAWSLLM):
|
|||
model: str,
|
||||
session_state: RealtimeResponseTransformInput,
|
||||
logging_obj: LiteLLMLogging | None = None,
|
||||
):
|
||||
"""Forward messages from client WebSocket to Bedrock stream."""
|
||||
initial_message: str | None = None,
|
||||
) -> None:
|
||||
"""Forward messages from client WebSocket to Bedrock stream.
|
||||
|
||||
Returns once the client is gone; provider failures (input stream or readiness) propagate to the caller.
|
||||
"""
|
||||
from aws_sdk_bedrock_runtime.models import (
|
||||
BidirectionalInputPayloadPart,
|
||||
InvokeModelWithBidirectionalStreamInputChunk,
|
||||
|
|
@ -299,41 +432,28 @@ class BedrockRealtime(BaseAWSLLM):
|
|||
verbose_proxy_logger.debug("Bedrock Realtime: Sent to Bedrock: %s", bedrock_message[:200])
|
||||
|
||||
try:
|
||||
while True:
|
||||
# Receive message from client
|
||||
message = await client_ws.receive_text()
|
||||
async for message in _client_messages(client_ws, initial_message):
|
||||
verbose_proxy_logger.debug("Bedrock Realtime: Received from client: %s", message[:200])
|
||||
parsed_client_message = _parse_client_message(message)
|
||||
is_session_update = _json_str(parsed_client_message.get("type")) == "session.update"
|
||||
if is_session_update:
|
||||
client_ws.scope[BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY] = (
|
||||
message # rebind-ok: scope outlives the attempt
|
||||
)
|
||||
|
||||
# Transform OpenAI format to Bedrock format
|
||||
transformed_messages = transformation_config.transform_realtime_request(
|
||||
message=message,
|
||||
model=model,
|
||||
session_configuration_request=session_state.get("session_configuration_request"),
|
||||
)
|
||||
|
||||
# Send transformed messages to Bedrock
|
||||
for bedrock_message in transformed_messages:
|
||||
await send_to_bedrock(bedrock_message)
|
||||
|
||||
if logging_obj is not None:
|
||||
client_message_type: str | None = None
|
||||
requested_modalities: list[str] | None = None
|
||||
with contextlib.suppress(Exception):
|
||||
parsed_client_message = _json_dict(_CLIENT_MESSAGE_ADAPTER.validate_json(message))
|
||||
client_message_type = _json_str(parsed_client_message.get("type"))
|
||||
if client_message_type == "session.update":
|
||||
requested_modalities = _CLIENT_MODALITIES_ADAPTER.validate_python(
|
||||
_json_dict(parsed_client_message.get("session")).get("modalities")
|
||||
)
|
||||
if client_message_type == "session.update":
|
||||
await client_ws.send_text(
|
||||
json.dumps(
|
||||
transformation_config.session_updated_event(model, logging_obj, requested_modalities)
|
||||
)
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.debug("Client to Bedrock forwarding ended: %s", e, exc_info=True)
|
||||
if is_session_update and not await _ack_session_update(
|
||||
client_ws, bedrock_stream, transformation_config, model, logging_obj, parsed_client_message
|
||||
):
|
||||
break
|
||||
finally:
|
||||
for close_message in transformation_config.session_close_messages():
|
||||
with contextlib.suppress(Exception):
|
||||
await send_to_bedrock(close_message)
|
||||
|
|
@ -349,68 +469,71 @@ class BedrockRealtime(BaseAWSLLM):
|
|||
logging_obj: LiteLLMLogging,
|
||||
session_state: RealtimeResponseTransformInput,
|
||||
) -> AsyncIterator[OpenAIRealtimeEvents]:
|
||||
"""Forward messages from Bedrock to the client, yielding the ones to record for spend logging."""
|
||||
try:
|
||||
while True:
|
||||
# Receive from Bedrock
|
||||
output = await bedrock_stream.await_output()
|
||||
result = await output[1].receive()
|
||||
"""Forward messages from Bedrock to the client, yielding the ones to record for spend logging.
|
||||
|
||||
if result is None:
|
||||
verbose_proxy_logger.debug("Bedrock Realtime: Bedrock stream ended")
|
||||
break
|
||||
Provider failures propagate to the caller; the client websocket is only closed on a normal stream end.
|
||||
"""
|
||||
|
||||
payload_bytes = result.value.bytes_ if result.value else None
|
||||
if payload_bytes:
|
||||
bedrock_response = payload_bytes.decode("utf-8")
|
||||
verbose_proxy_logger.debug("Bedrock Realtime: Received from Bedrock: %s", bedrock_response[:200])
|
||||
|
||||
# Transform Bedrock format to OpenAI format
|
||||
realtime_response_transform_input: RealtimeResponseTransformInput = {
|
||||
"current_output_item_id": session_state.get("current_output_item_id"),
|
||||
"current_response_id": session_state.get("current_response_id"),
|
||||
"current_conversation_id": session_state.get("current_conversation_id"),
|
||||
"current_delta_chunks": session_state.get("current_delta_chunks"),
|
||||
"current_item_chunks": session_state.get("current_item_chunks"),
|
||||
"current_delta_type": session_state.get("current_delta_type"),
|
||||
"session_configuration_request": session_state.get("session_configuration_request"),
|
||||
}
|
||||
|
||||
transformed_response = transformation_config.transform_realtime_response(
|
||||
message=bedrock_response,
|
||||
model=model,
|
||||
logging_obj=logging_obj,
|
||||
realtime_response_transform_input=realtime_response_transform_input,
|
||||
)
|
||||
|
||||
# Update session state
|
||||
session_state.update(
|
||||
{
|
||||
"current_output_item_id": transformed_response.get("current_output_item_id"),
|
||||
"current_response_id": transformed_response.get("current_response_id"),
|
||||
"current_conversation_id": transformed_response.get("current_conversation_id"),
|
||||
"current_delta_chunks": transformed_response.get("current_delta_chunks"),
|
||||
"current_item_chunks": transformed_response.get("current_item_chunks"),
|
||||
"current_delta_type": transformed_response.get("current_delta_type"),
|
||||
"session_configuration_request": transformed_response.get("session_configuration_request"),
|
||||
}
|
||||
)
|
||||
|
||||
# Send transformed messages to client
|
||||
response_value = transformed_response["response"]
|
||||
openai_messages = response_value if isinstance(response_value, list) else (response_value,)
|
||||
for openai_message in openai_messages:
|
||||
message_json = json.dumps(openai_message)
|
||||
await client_ws.send_text(message_json)
|
||||
verbose_proxy_logger.debug("Bedrock Realtime: Sent to client: %s", message_json[:200])
|
||||
if _should_log_event(openai_message):
|
||||
yield openai_message
|
||||
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.debug("Bedrock to client forwarding ended: %s", e, exc_info=True)
|
||||
finally:
|
||||
# Close the client WebSocket
|
||||
async def send_to_client(message_json: str) -> bool:
|
||||
try:
|
||||
await client_ws.close()
|
||||
except Exception:
|
||||
pass
|
||||
await client_ws.send_text(message_json)
|
||||
except Exception as e: # noqa: BLE001 # any send failure means the client is gone
|
||||
verbose_proxy_logger.debug("Bedrock to client forwarding ended: %s", e, exc_info=True)
|
||||
return False
|
||||
verbose_proxy_logger.debug("Bedrock Realtime: Sent to client: %s", message_json[:200])
|
||||
return True
|
||||
|
||||
output: Final = await bedrock_stream.await_output()
|
||||
while True:
|
||||
result = await output[1].receive()
|
||||
|
||||
if result is None:
|
||||
verbose_proxy_logger.debug("Bedrock Realtime: Bedrock stream ended")
|
||||
with contextlib.suppress(Exception):
|
||||
await client_ws.close()
|
||||
return
|
||||
|
||||
payload_bytes = result.value.bytes_ if result.value else None
|
||||
if payload_bytes:
|
||||
bedrock_response = payload_bytes.decode("utf-8")
|
||||
verbose_proxy_logger.debug("Bedrock Realtime: Received from Bedrock: %s", bedrock_response[:200])
|
||||
|
||||
# Transform Bedrock format to OpenAI format
|
||||
realtime_response_transform_input: RealtimeResponseTransformInput = {
|
||||
"current_output_item_id": session_state.get("current_output_item_id"),
|
||||
"current_response_id": session_state.get("current_response_id"),
|
||||
"current_conversation_id": session_state.get("current_conversation_id"),
|
||||
"current_delta_chunks": session_state.get("current_delta_chunks"),
|
||||
"current_item_chunks": session_state.get("current_item_chunks"),
|
||||
"current_delta_type": session_state.get("current_delta_type"),
|
||||
"session_configuration_request": session_state.get("session_configuration_request"),
|
||||
}
|
||||
|
||||
transformed_response = transformation_config.transform_realtime_response(
|
||||
message=bedrock_response,
|
||||
model=model,
|
||||
logging_obj=logging_obj,
|
||||
realtime_response_transform_input=realtime_response_transform_input,
|
||||
)
|
||||
|
||||
# Update session state
|
||||
session_state.update(
|
||||
{
|
||||
"current_output_item_id": transformed_response.get("current_output_item_id"),
|
||||
"current_response_id": transformed_response.get("current_response_id"),
|
||||
"current_conversation_id": transformed_response.get("current_conversation_id"),
|
||||
"current_delta_chunks": transformed_response.get("current_delta_chunks"),
|
||||
"current_item_chunks": transformed_response.get("current_item_chunks"),
|
||||
"current_delta_type": transformed_response.get("current_delta_type"),
|
||||
"session_configuration_request": transformed_response.get("session_configuration_request"),
|
||||
}
|
||||
)
|
||||
|
||||
# Send transformed messages to client
|
||||
response_value = transformed_response["response"]
|
||||
openai_messages = response_value if isinstance(response_value, list) else (response_value,)
|
||||
for openai_message in openai_messages:
|
||||
if not await send_to_client(json.dumps(openai_message)):
|
||||
return
|
||||
if _should_log_event(openai_message):
|
||||
yield openai_message
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -66,6 +66,7 @@ class DeepSeekAnthropicMessagesConfig(AnthropicMessagesConfig):
|
|||
headers=headers,
|
||||
optional_params=optional_params,
|
||||
custom_llm_provider=self.custom_llm_provider or "deepseek",
|
||||
messages=messages,
|
||||
)
|
||||
|
||||
return headers, api_base
|
||||
|
|
|
|||
|
|
@ -250,8 +250,7 @@ class FireworksAIRerankConfig(FireworksAIMixin, BaseRerankConfig):
|
|||
|
||||
rerank_results.append(rerank_result)
|
||||
|
||||
# Use model name as id if no id is provided
|
||||
response_id: Final = raw_response_json.get("id") or raw_response_json.get("model") or str(uuid.uuid4())
|
||||
response_id: Final = raw_response_json.get("id") or str(uuid.uuid4())
|
||||
|
||||
return RerankResponse(
|
||||
id=response_id,
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
||||
|
|
|
|||
|
|
@ -92,7 +92,7 @@ class GithubCopilotAnthropicMessagesConfig(AnthropicMessagesConfig):
|
|||
headers["anthropic-version"] = "2023-06-01"
|
||||
|
||||
headers = self._update_headers_with_anthropic_beta(
|
||||
headers, optional_params, custom_llm_provider="github_copilot"
|
||||
headers, optional_params, custom_llm_provider="github_copilot", messages=messages
|
||||
)
|
||||
|
||||
return headers, dynamic_api_base
|
||||
|
|
|
|||
|
|
@ -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 ()
|
||||
|
|
|
|||
|
|
@ -56,6 +56,7 @@ class OpenAILikeAnthropicMessagesConfig(AnthropicMessagesConfig):
|
|||
merged: Final = self._update_headers_with_anthropic_beta(
|
||||
headers=normalized,
|
||||
optional_params=optional_params,
|
||||
messages=messages,
|
||||
)
|
||||
return merged, api_base
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -4,6 +4,8 @@ Translates from Cohere's `/v1/rerank` input format to Vertex AI Discovery Engine
|
|||
Why separate file? Make it easy to see how transformation works
|
||||
"""
|
||||
|
||||
import math
|
||||
import uuid
|
||||
from collections.abc import Mapping
|
||||
from typing import Any, Final
|
||||
|
||||
|
|
@ -32,6 +34,8 @@ class VertexAIRerankConfig(BaseRerankConfig, VertexBase):
|
|||
Reference: https://cloud.google.com/generative-ai-app-builder/docs/ranking#rank_or_rerank_a_set_of_records_according_to_a_query
|
||||
"""
|
||||
|
||||
MAX_RECORDS_PER_SEARCH_UNIT = 100
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
|
||||
|
|
@ -208,10 +212,11 @@ class VertexAIRerankConfig(BaseRerankConfig, VertexBase):
|
|||
RerankResponseResult(index=result["index"], relevance_score=result["relevance_score"])
|
||||
)
|
||||
|
||||
# Create meta object
|
||||
meta: Final = RerankResponseMeta(billed_units=RerankBilledUnits(search_units=len(records)))
|
||||
input_record_count: Final = len(request_data.get("records", ()))
|
||||
search_units: Final = math.ceil(input_record_count / self.MAX_RECORDS_PER_SEARCH_UNIT)
|
||||
meta: Final = RerankResponseMeta(billed_units=RerankBilledUnits(search_units=search_units))
|
||||
|
||||
return RerankResponse(id=f"vertex_ai_rerank_{model}", results=rerank_results, meta=meta)
|
||||
return RerankResponse(id=f"vertex_ai_rerank_{uuid.uuid4()}", results=rerank_results, meta=meta)
|
||||
|
||||
def get_supported_cohere_rerank_params(self, model: str) -> list:
|
||||
return [
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ from typing import Any, Final
|
|||
|
||||
import httpx
|
||||
|
||||
from litellm._uuid import uuid
|
||||
from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj
|
||||
from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
|
|
@ -127,7 +128,7 @@ class VoyageRerankConfig(BaseRerankConfig):
|
|||
rerank_meta: Final = RerankResponseMeta(billed_units=_billed_units, tokens=_tokens)
|
||||
|
||||
return RerankResponse(
|
||||
id=_json_response.get("id", f"voyage-rerank-{model}"),
|
||||
id=_json_response.get("id") or str(uuid.uuid4()),
|
||||
results=transformed_results,
|
||||
meta=rerank_meta,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -191,7 +191,7 @@ class IBMWatsonXRerankConfig(IBMWatsonXMixin, BaseRerankConfig):
|
|||
|
||||
transformed_results.append(transformed_result)
|
||||
|
||||
response_id: Final = raw_response_json.get("id") or raw_response_json.get("model_id") or str(uuid.uuid4())
|
||||
response_id: Final = raw_response_json.get("id") or str(uuid.uuid4())
|
||||
|
||||
# Extract usage information
|
||||
_tokens: Final = RerankTokens(
|
||||
|
|
|
|||
|
|
@ -219,13 +219,19 @@ class XAIChatConfig(OpenAIGPTConfig):
|
|||
litellm_params: dict,
|
||||
headers: dict,
|
||||
) -> dict:
|
||||
"""
|
||||
Handle https://github.com/BerriAI/litellm/issues/9720
|
||||
"""Handle https://github.com/BerriAI/litellm/issues/9720"""
|
||||
if "web_search_options" in optional_params:
|
||||
verbose_logger.warning(
|
||||
"XAI no longer supports web search on /chat/completions (Live Search is deprecated). "
|
||||
"Dropping 'web_search_options'. Use the Responses API for XAI web search."
|
||||
)
|
||||
|
||||
Filter out 'name' from messages
|
||||
"""
|
||||
messages = strip_name_from_messages(messages)
|
||||
return super().transform_request(model, messages, optional_params, litellm_params, headers)
|
||||
chat_params: Final = { # mutable-ok: base transform_request takes a plain dict of optional params
|
||||
key: value for key, value in optional_params.items() if key != "web_search_options"
|
||||
}
|
||||
return super().transform_request(
|
||||
model, strip_name_from_messages(messages), chat_params, litellm_params, headers
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _fix_choice_finish_reason_for_tool_calls(choice: Choices) -> None:
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ from types import MappingProxyType
|
|||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
import httpx
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -32,6 +33,8 @@ if TYPE_CHECKING:
|
|||
else:
|
||||
LiteLLMLoggingObj = Any
|
||||
|
||||
_STR_MAPPING_ADAPTER: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def _usage_restated_from_xai_ticks(usage: ResponseAPIUsage | None) -> ResponseAPIUsage | None:
|
||||
reported_cost: Final = xai_reported_cost_in_usd(getattr(usage, "cost_in_usd_ticks", None))
|
||||
|
|
@ -81,30 +84,25 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
|
|||
- enable_image_understanding
|
||||
|
||||
XAI does NOT support search_context_size (OpenAI-specific).
|
||||
|
||||
Domains may come nested under 'filters' (the OpenAI/XAI documented shape) or flat on the tool.
|
||||
"""
|
||||
xai_tool: Final[dict[str, object]] = {"type": "web_search"}
|
||||
|
||||
# Remove search_context_size if present (not supported by XAI)
|
||||
if "search_context_size" in tool:
|
||||
verbose_logger.info(
|
||||
"XAI does not support 'search_context_size' parameter. Removing it from web_search tool."
|
||||
)
|
||||
|
||||
# Handle filters (XAI-specific structure)
|
||||
filters: Final = {}
|
||||
if "allowed_domains" in tool:
|
||||
allowed_domains: Final = tool["allowed_domains"]
|
||||
filters["allowed_domains"] = allowed_domains
|
||||
nested_filters: Final = tool.get("filters")
|
||||
domains: Final = (
|
||||
_STR_MAPPING_ADAPTER.validate_python(nested_filters) if isinstance(nested_filters, Mapping) else tool
|
||||
)
|
||||
filters: Final = {key: domains[key] for key in ("allowed_domains", "excluded_domains") if key in domains}
|
||||
|
||||
if "excluded_domains" in tool:
|
||||
excluded_domains: Final = tool["excluded_domains"]
|
||||
filters["excluded_domains"] = excluded_domains
|
||||
|
||||
# Add filters if any were specified
|
||||
if filters:
|
||||
xai_tool["filters"] = filters
|
||||
|
||||
# Handle enable_image_understanding (top-level in XAI format)
|
||||
if "enable_image_understanding" in tool:
|
||||
xai_tool["enable_image_understanding"] = tool["enable_image_understanding"]
|
||||
|
||||
|
|
|
|||
|
|
@ -67,7 +67,7 @@ from litellm.constants import (
|
|||
)
|
||||
from litellm.exceptions import LiteLLMUnknownProvider
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.asyncify import run_async_function
|
||||
from litellm.litellm_core_utils.asyncify import asyncify, run_async_function
|
||||
from litellm.litellm_core_utils.audio_utils.utils import (
|
||||
calculate_request_duration,
|
||||
get_audio_file_for_health_check,
|
||||
|
|
@ -1072,10 +1072,6 @@ def responses_api_bridge_check(
|
|||
mode = "responses"
|
||||
model_info["mode"] = mode
|
||||
|
||||
if web_search_options is not None and custom_llm_provider == "xai":
|
||||
model_info["mode"] = "responses"
|
||||
model = model.replace("responses/", "")
|
||||
|
||||
except Exception as e:
|
||||
verbose_logger.debug("Error getting model info: %s", e)
|
||||
|
||||
|
|
@ -1084,6 +1080,10 @@ def responses_api_bridge_check(
|
|||
mode = "responses"
|
||||
model_info["mode"] = mode
|
||||
|
||||
if web_search_options is not None and custom_llm_provider == "xai":
|
||||
model_info["mode"] = "responses"
|
||||
model = model.replace("responses/", "")
|
||||
|
||||
# OpenAI/Azure GPT-5 chat-completions that need Responses-only fields (e.g.
|
||||
# ``reasoningSummary`` in ``extra_body``) must be bridged; Chat Completions rejects
|
||||
# those keys.
|
||||
|
|
@ -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(
|
||||
|
|
@ -9127,7 +9127,7 @@ async def acount_tokens(
|
|||
fallback_messages = messages or []
|
||||
if system and fallback_messages:
|
||||
fallback_messages = [{"role": "system", "content": system}] + fallback_messages
|
||||
local_count: Final = litellm.token_counter(
|
||||
local_count: Final = await asyncify(litellm.token_counter)(
|
||||
model=model,
|
||||
messages=fallback_messages,
|
||||
tools=tools,
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -26,6 +26,7 @@ class LiteLLM_BudgetTable(LiteLLMPydanticObjectBase):
|
|||
max_parallel_requests: int | None = None
|
||||
tpm_limit: int | None = None
|
||||
rpm_limit: int | None = None
|
||||
tpd_limit: int | None = None
|
||||
model_max_budget: dict | None = None
|
||||
budget_duration: str | None = None
|
||||
allowed_models: list[str] | None = None # per-member model scope; empty = inherit team models
|
||||
|
|
|
|||
|
|
@ -5,6 +5,8 @@ These are the canonical credential types for the proxy. They live in the model
|
|||
layer; ``litellm.types.utils`` re-exports them for backwards compatibility.
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping
|
||||
|
||||
from pydantic import BaseModel, model_validator
|
||||
|
||||
|
||||
|
|
@ -27,3 +29,10 @@ class CreateCredentialItem(CredentialBase):
|
|||
if not values.get("credential_values") and not values.get("model_id"):
|
||||
raise ValueError("Either credential_values or model_id must be set")
|
||||
return values
|
||||
|
||||
|
||||
class UpdateCredentialItem(BaseModel):
|
||||
credential_name: str
|
||||
credential_info: Mapping[str, object]
|
||||
credential_values: Mapping[str, object] | None = None
|
||||
model_id: str | None = None
|
||||
|
|
|
|||
|
|
@ -71,6 +71,7 @@ class TeamBase(LiteLLMPydanticObjectBase):
|
|||
metadata: dict | None = None
|
||||
tpm_limit: int | None = None
|
||||
rpm_limit: int | None = None
|
||||
tpd_limit: int | None = None
|
||||
max_budget: float | None = None
|
||||
soft_budget: float | None = None
|
||||
budget_duration: str | None = None
|
||||
|
|
|
|||
|
|
@ -31,6 +31,7 @@ class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase):
|
|||
metadata: dict = {}
|
||||
tpm_limit: int | None = None
|
||||
rpm_limit: int | None = None
|
||||
tpd_limit: int | None = None
|
||||
budget_duration: str | None = None
|
||||
budget_reset_at: datetime | None = None
|
||||
allowed_cache_controls: list | None = []
|
||||
|
|
|
|||
|
|
@ -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": [
|
||||
{
|
||||
|
|
@ -12962,18 +12968,24 @@
|
|||
"PHONE_NUMBER",
|
||||
"MEDICAL_LICENSE",
|
||||
"URL",
|
||||
"MAC_ADDRESS",
|
||||
"UUID",
|
||||
"US_BANK_NUMBER",
|
||||
"US_DRIVER_LICENSE",
|
||||
"US_ITIN",
|
||||
"US_PASSPORT",
|
||||
"US_SSN",
|
||||
"US_MBI",
|
||||
"US_NPI",
|
||||
"UK_NHS",
|
||||
"UK_NINO",
|
||||
"UK_PASSPORT",
|
||||
"UK_POSTCODE",
|
||||
"UK_VEHICLE_REGISTRATION",
|
||||
"UK_DRIVING_LICENCE",
|
||||
"ES_NIF",
|
||||
"ES_NIE",
|
||||
"ES_PASSPORT",
|
||||
"IT_FISCAL_CODE",
|
||||
"IT_DRIVER_LICENSE",
|
||||
"IT_VAT_CODE",
|
||||
|
|
@ -12991,7 +13003,38 @@
|
|||
"IN_VEHICLE_REGISTRATION",
|
||||
"IN_VOTER",
|
||||
"IN_PASSPORT",
|
||||
"FI_PERSONAL_IDENTITY_CODE"
|
||||
"IN_GSTIN",
|
||||
"FI_PERSONAL_IDENTITY_CODE",
|
||||
"DE_TAX_ID",
|
||||
"DE_TAX_NUMBER",
|
||||
"DE_VAT_ID",
|
||||
"DE_PASSPORT",
|
||||
"DE_ID_CARD",
|
||||
"DE_FUEHRERSCHEIN",
|
||||
"DE_SOCIAL_SECURITY",
|
||||
"DE_HEALTH_INSURANCE",
|
||||
"DE_LANR",
|
||||
"DE_BSNR",
|
||||
"DE_KFZ",
|
||||
"DE_HANDELSREGISTER",
|
||||
"DE_PLZ",
|
||||
"KR_RRN",
|
||||
"KR_FRN",
|
||||
"KR_PASSPORT",
|
||||
"KR_DRIVER_LICENSE",
|
||||
"KR_BRN",
|
||||
"CA_SIN",
|
||||
"SE_PERSONNUMMER",
|
||||
"SE_ORGANISATIONSNUMMER",
|
||||
"TH_TNIN",
|
||||
"TR_NATIONAL_ID",
|
||||
"TR_LICENSE_PLATE",
|
||||
"NG_NIN",
|
||||
"NG_VEHICLE_REGISTRATION",
|
||||
"PH_TIN",
|
||||
"PH_UMID",
|
||||
"PH_PASSPORT",
|
||||
"ZA_ID_NUMBER"
|
||||
],
|
||||
"title": "PiiEntityType",
|
||||
"type": "string"
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
@ -284,6 +286,7 @@ class KeyManagementRoutes(str, enum.Enum):
|
|||
# team's `team_member_permissions`, non-admin members of that team may set
|
||||
# `access_group_ids` on keys they create/update. Default-deny.
|
||||
KEY_ACCESS_GROUP_ASSIGNMENT = "/key/access_group_assignment"
|
||||
AUTO_ROUTER_MANAGE = "/auto_router/manage"
|
||||
|
||||
# info and health routes
|
||||
KEY_INFO = "/key/info"
|
||||
|
|
@ -650,15 +653,18 @@ class LiteLLMRoutes(enum.Enum):
|
|||
KeyManagementRoutes.KEY_RESET_SPEND.value,
|
||||
KeyManagementRoutes.KEY_ALIASES.value,
|
||||
KeyManagementRoutes.KEY_ACCESS_GROUP_ASSIGNMENT.value,
|
||||
KeyManagementRoutes.AUTO_ROUTER_MANAGE.value,
|
||||
]
|
||||
|
||||
management_routes = (
|
||||
[
|
||||
# user
|
||||
"/user/new",
|
||||
"/management/v1/users/bulk",
|
||||
"/user/update",
|
||||
"/user/bulk_update",
|
||||
"/user/delete",
|
||||
"/management/v1/users/bulk_delete",
|
||||
"/user/info",
|
||||
"/user/list",
|
||||
"/user/daily/activity",
|
||||
|
|
@ -838,6 +844,7 @@ class LiteLLMRoutes(enum.Enum):
|
|||
self_managed_routes = [
|
||||
"/team/member_add",
|
||||
"/team/member_delete",
|
||||
"/management/v1/teams/{team_id}/members/bulk_delete",
|
||||
"/team/member_update",
|
||||
"/team/{team_id}/member/{user_id}/reset_spend",
|
||||
"/team/permissions_list",
|
||||
|
|
@ -864,6 +871,7 @@ class LiteLLMRoutes(enum.Enum):
|
|||
"/organization/daily/activity",
|
||||
"/user/available_roles", # read-only role metadata; any authenticated user may read
|
||||
"/user/list", # org admins checked in endpoint; non-admins get 403
|
||||
"/management/v1/users/bulk_delete", # proxy admins delete anyone, org admins only their orgs' users; others 403
|
||||
"/model/{model_id}/update",
|
||||
"/prompt/list",
|
||||
"/prompt/info",
|
||||
|
|
@ -887,6 +895,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,
|
||||
|
|
@ -1197,6 +1206,7 @@ class AllowedVectorStoreIndexItem(LiteLLMPydanticObjectBase):
|
|||
|
||||
class KeyRequestBase(GenerateRequestBase):
|
||||
key: str | None = None
|
||||
tpd_limit: int | None = None
|
||||
default_estimated_output_tokens: PositiveInt | None = None
|
||||
default_estimated_output_tokens_per_model: Mapping[str, PositiveInt] | None = None
|
||||
budget_id: str | None = None
|
||||
|
|
@ -1882,6 +1892,9 @@ class BudgetNewRequest(LiteLLMPydanticObjectBase):
|
|||
)
|
||||
tpm_limit: int | None = Field(default=None, description="Max tokens per minute, allowed for this budget id.")
|
||||
rpm_limit: int | None = Field(default=None, description="Max requests per minute, allowed for this budget id.")
|
||||
tpd_limit: int | None = Field(
|
||||
default=None, description="Max tokens per day, charged by batch submissions, allowed for this budget id."
|
||||
)
|
||||
budget_duration: str | None = Field(
|
||||
default=None,
|
||||
description="Max duration budget should be set for (e.g. '1hr', '1d', '28d')",
|
||||
|
|
@ -1980,8 +1993,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
|
||||
|
|
@ -2052,6 +2071,7 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase):
|
|||
metadata: dict | None = None
|
||||
tpm_limit: int | None = None
|
||||
rpm_limit: int | None = None
|
||||
tpd_limit: int | None = None
|
||||
max_budget: float | None = None
|
||||
soft_budget: float | None = None
|
||||
models: list | None = None
|
||||
|
|
@ -2079,7 +2099,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
|
||||
|
|
@ -2635,9 +2655,13 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
|
|||
None,
|
||||
description="max file size in MB for /v1/files uploads, for any purpose, if a file is larger than this size it will be rejected before being forwarded to the provider",
|
||||
)
|
||||
allowed_file_extensions: tuple[str, ...] | None = Field(
|
||||
None,
|
||||
description="the only file extensions (e.g. ['.jsonl', '.pdf', '.txt']) accepted on /v1/files uploads, for any purpose, matched case-insensitively against the uploaded filename. Files with any other extension, or none, are rejected. An empty list rejects every upload. Unset means no allowlist is applied",
|
||||
)
|
||||
blocked_file_extensions: tuple[str, ...] | None = Field(
|
||||
None,
|
||||
description="file extensions (e.g. ['.exe', '.sh']) rejected on /v1/files uploads, for any purpose, matched case-insensitively against the uploaded filename",
|
||||
description="file extensions (e.g. ['.exe', '.sh']) rejected on /v1/files uploads, for any purpose, matched case-insensitively against the uploaded filename. Deprecated in favour of allowed_file_extensions; still enforced, after the allowlist, when set",
|
||||
)
|
||||
max_response_size_mb: int | None = Field(
|
||||
None,
|
||||
|
|
@ -3003,6 +3027,7 @@ class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken):
|
|||
team_alias: str | None = None
|
||||
team_tpm_limit: int | None = None
|
||||
team_rpm_limit: int | None = None
|
||||
team_tpd_limit: int | None = None
|
||||
team_max_budget: float | None = None
|
||||
team_soft_budget: float | None = None
|
||||
team_models: list = []
|
||||
|
|
@ -3022,6 +3047,7 @@ class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken):
|
|||
end_user_id: str | None = None
|
||||
end_user_tpm_limit: int | None = None
|
||||
end_user_rpm_limit: int | None = None
|
||||
end_user_tpd_limit: int | None = None
|
||||
end_user_max_budget: float | None = None
|
||||
end_user_model_max_budget: dict | None = None
|
||||
|
||||
|
|
@ -3820,6 +3846,7 @@ class SpendLogsMetadata(TypedDict):
|
|||
user_api_key_team_alias: str | None
|
||||
spend_logs_metadata: dict | None # special param to log k,v pairs to spendlogs for a call
|
||||
requester_ip_address: str | None
|
||||
user_agent: ReadOnly[str | None]
|
||||
litellm_call_id: str | None
|
||||
applied_guardrails: list[str] | None
|
||||
mcp_tool_call_metadata: StandardLoggingMCPToolCall | None
|
||||
|
|
@ -4414,6 +4441,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):
|
||||
|
|
@ -4694,6 +4724,7 @@ class JWTAuthBuilderResult(TypedDict):
|
|||
org_id: str | None
|
||||
team_membership: LiteLLM_TeamMembership | None
|
||||
jwt_claims: dict # Decoded JWT token claims (avoids re-decoding)
|
||||
agent_id: ReadOnly[str | None]
|
||||
|
||||
|
||||
class ClientSideFallbackModel(TypedDict, total=False):
|
||||
|
|
@ -4837,6 +4868,14 @@ class JWTIssuerConfig(BaseModel):
|
|||
default=None,
|
||||
description="Issuer-specific claim path to normalize into LiteLLM's end-user id.",
|
||||
)
|
||||
virtual_key_claim_field: str | None = Field(
|
||||
default=None,
|
||||
description="Issuer-specific claim path used for the virtual key mapping lookup. Falls back to the global field.",
|
||||
)
|
||||
unregistered_jwt_client_behavior: UnregisteredJWTClientBehavior | None = Field(
|
||||
default=None,
|
||||
description="Issuer-specific policy when the virtual key claim has no mapping. Falls back to the global policy.",
|
||||
)
|
||||
|
||||
model_config = {
|
||||
"extra": "forbid",
|
||||
|
|
@ -4924,6 +4963,14 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase):
|
|||
user_allowed_roles: list[str] | None = None
|
||||
user_id_upsert: bool = Field(default=False, description="If user doesn't exist, upsert them into the db.")
|
||||
end_user_id_jwt_field: str | None = None
|
||||
agent_id_jwt_field: str | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"The field in the JWT token that identifies the calling agent (e.g. 'azp' for a Microsoft Entra ID "
|
||||
"app token). Supports dot notation. The value is matched against a registered agent's agent_id, "
|
||||
"then agent_name, and the request is rejected when it matches neither."
|
||||
),
|
||||
)
|
||||
public_key_ttl: float = 600
|
||||
public_key_stale_ttl: float = Field(
|
||||
default=DEFAULT_JWKS_STALE_TTL,
|
||||
|
|
@ -5063,6 +5110,28 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase):
|
|||
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def get_issuer_config(self, issuer: str | None) -> JWTIssuerConfig | None:
|
||||
if issuer is None or self.issuers is None:
|
||||
return None
|
||||
return next((config for config in self.issuers if config.issuer == issuer), None)
|
||||
|
||||
def is_virtual_key_mapping_configured(self) -> bool:
|
||||
if self.virtual_key_claim_field is not None:
|
||||
return True
|
||||
return any(config.virtual_key_claim_field is not None for config in self.issuers or ())
|
||||
|
||||
def get_virtual_key_claim_field(self, issuer: str | None) -> str | None:
|
||||
issuer_config: Final = self.get_issuer_config(issuer)
|
||||
if issuer_config is not None and issuer_config.virtual_key_claim_field is not None:
|
||||
return issuer_config.virtual_key_claim_field
|
||||
return self.virtual_key_claim_field
|
||||
|
||||
def get_unregistered_jwt_client_behavior(self, issuer: str | None) -> UnregisteredJWTClientBehavior:
|
||||
issuer_config: Final = self.get_issuer_config(issuer)
|
||||
if issuer_config is not None and issuer_config.unregistered_jwt_client_behavior is not None:
|
||||
return issuer_config.unregistered_jwt_client_behavior
|
||||
return self.unregistered_jwt_client_behavior
|
||||
|
||||
|
||||
class PrismaCompatibleUpdateDBModel(TypedDict, total=False):
|
||||
model_name: str
|
||||
|
|
|
|||
|
|
@ -39,6 +39,7 @@ from litellm.constants import (
|
|||
from litellm.litellm_core_utils.dd_tracing import tracer
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
|
||||
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
|
||||
from litellm.models.project import LiteLLM_ProjectTable
|
||||
from litellm.proxy._types import (
|
||||
RBAC_ROLES,
|
||||
CallInfo,
|
||||
|
|
@ -109,7 +110,7 @@ from litellm.proxy.utils import PrismaClient, ProxyLogging, log_db_metrics
|
|||
from litellm.repositories.budget_repository import BudgetRepository
|
||||
from litellm.repositories.object_permission_repository import ObjectPermissionRepository
|
||||
from litellm.repositories.organization_repository import OrganizationRepository
|
||||
from litellm.repositories.prisma_protocols import RowT_co
|
||||
from litellm.repositories.prisma_protocols import DatabaseClient, RowT_co
|
||||
from litellm.repositories.project_repository import ProjectRepository
|
||||
from litellm.repositories.table_repositories import (
|
||||
AccessGroupRepository,
|
||||
|
|
@ -123,6 +124,7 @@ from litellm.repositories.table_repositories import (
|
|||
from litellm.repositories.team_repository import TeamRepository
|
||||
from litellm.repositories.user_repository import UserRepository
|
||||
from litellm.router import Router
|
||||
from litellm.types.proxy.auth.auth_checks import UserNotFoundError
|
||||
from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget
|
||||
from litellm.utils import get_utc_datetime
|
||||
|
||||
|
|
@ -327,9 +329,23 @@ _safe_json_loads_obj: Final = _typed_json_loads(safe_json_loads)
|
|||
last_db_access_time: Final = LimitedSizeOrderedDict(max_size=100)
|
||||
db_cache_expiry: Final = DEFAULT_IN_MEMORY_TTL # refresh every 5s
|
||||
|
||||
_TEAM_MEMBERSHIP_INFLIGHT_MAX: Final = 10000
|
||||
_team_membership_inflight: Final = LimitedSizeOrderedDict(max_size=_TEAM_MEMBERSHIP_INFLIGHT_MAX)
|
||||
|
||||
|
||||
class _TeamMembershipCacheMiss:
|
||||
__slots__ = ()
|
||||
|
||||
|
||||
_TEAM_MEMBERSHIP_CACHE_MISS: Final = _TeamMembershipCacheMiss()
|
||||
|
||||
all_routes: Final = LiteLLMRoutes.openai_routes.value + LiteLLMRoutes.management_routes.value
|
||||
|
||||
|
||||
def _membership_from_shared_load(result: object) -> LiteLLM_TeamMembership | None:
|
||||
return result if isinstance(result, LiteLLM_TeamMembership) else None
|
||||
|
||||
|
||||
def _log_budget_lookup_failure(entity: str, error: Exception) -> None:
|
||||
"""
|
||||
Log a warning when budget lookup fails; cache will not be populated.
|
||||
|
|
@ -832,6 +848,7 @@ BUDGET_ENFORCED_SIDE_EFFECT_ROUTES: Final = frozenset(
|
|||
"/health",
|
||||
"/health/services",
|
||||
"/health/test_connection",
|
||||
"/auto_router/test_routing",
|
||||
}
|
||||
)
|
||||
|
||||
|
|
@ -880,6 +897,7 @@ async def common_checks(
|
|||
request_query_params=_safe_get_request_query_params(request=request),
|
||||
llm_router=llm_router,
|
||||
request=request,
|
||||
team_id=valid_token.team_id if valid_token is not None else None,
|
||||
)
|
||||
|
||||
skip_all_budget_checks: Final = skip_budget_checks or (
|
||||
|
|
@ -887,6 +905,22 @@ async def common_checks(
|
|||
and (route in MODEL_DISCOVERY_ROUTES or not RouteChecks.is_llm_api_route(route=route))
|
||||
)
|
||||
|
||||
membership_user_id: Final = (
|
||||
valid_token.user_id if valid_token is not None and (bool(_model) or not skip_all_budget_checks) else None
|
||||
)
|
||||
team_membership_loaded: Final = team_object is not None and membership_user_id is not None
|
||||
loaded_team_membership: Final = (
|
||||
await get_team_membership(
|
||||
user_id=membership_user_id,
|
||||
team_id=team_object.team_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
if team_object is not None and membership_user_id is not None
|
||||
else None
|
||||
)
|
||||
|
||||
unpriced_models: Final = (
|
||||
_unpriced_models_in_request(model=_model, llm_router=llm_router)
|
||||
if litellm.block_requests_for_models_without_pricing and RouteChecks.is_llm_api_route(route=route)
|
||||
|
|
@ -936,6 +970,8 @@ async def common_checks(
|
|||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
team_membership=loaded_team_membership,
|
||||
team_membership_loaded=team_membership_loaded,
|
||||
)
|
||||
|
||||
# Require trace id for agent keys when agent has require_trace_id_on_calls_by_agent
|
||||
|
|
@ -987,6 +1023,8 @@ async def common_checks(
|
|||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
team_membership=loaded_team_membership,
|
||||
team_membership_loaded=team_membership_loaded,
|
||||
)
|
||||
|
||||
# Run before apply_key_tags_pre_auth injects key metadata.tags into request_body.
|
||||
|
|
@ -1096,6 +1134,8 @@ async def common_checks(
|
|||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
team_membership=loaded_team_membership,
|
||||
team_membership_loaded=team_membership_loaded,
|
||||
),
|
||||
_check_end_user_budget(end_user_obj=end_user_object, route=route)
|
||||
if end_user_object is not None and end_user_object.litellm_budget_table is not None
|
||||
|
|
@ -2141,7 +2181,76 @@ async def get_tag_object(
|
|||
return tag_objects.get(tag_name)
|
||||
|
||||
|
||||
def _membership_from_cached_payload(
|
||||
cached: object,
|
||||
) -> LiteLLM_TeamMembership | None | _TeamMembershipCacheMiss:
|
||||
if cached is None:
|
||||
return _TEAM_MEMBERSHIP_CACHE_MISS
|
||||
if cached == NO_TEAM_MEMBERSHIP_SENTINEL:
|
||||
return None
|
||||
cached_membership: Final = CacheCodec.deserialize(cached, model_type=LiteLLM_TeamMembership)
|
||||
return cached_membership if cached_membership is not None else _TEAM_MEMBERSHIP_CACHE_MISS
|
||||
|
||||
|
||||
@log_db_metrics
|
||||
async def _fetch_team_membership_from_db(
|
||||
user_id: str,
|
||||
team_id: str,
|
||||
prisma_client: PrismaClient,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
parent_otel_span: Span | None = None,
|
||||
proxy_logging_obj: ProxyLogging | None = None,
|
||||
) -> LiteLLM_TeamMembership | None:
|
||||
_ = parent_otel_span, proxy_logging_obj
|
||||
response: Final = await _dictable_table(TeamMembershipRepository(prisma_client)).find_unique(
|
||||
where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}},
|
||||
include={"litellm_budget_table": True},
|
||||
)
|
||||
membership: Final = None if response is None else LiteLLM_TeamMembership.model_validate(response.dict())
|
||||
_key: Final = team_membership_reservation_cache_key(user_id=user_id, team_id=team_id)
|
||||
if membership is None:
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=_key,
|
||||
value=NO_TEAM_MEMBERSHIP_SENTINEL,
|
||||
ttl=get_management_object_ttl(user_api_key_cache),
|
||||
)
|
||||
else:
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=_key,
|
||||
value=membership,
|
||||
model_type=LiteLLM_TeamMembership,
|
||||
)
|
||||
return membership
|
||||
|
||||
|
||||
async def _load_team_membership_on_cache_miss(
|
||||
user_id: str,
|
||||
team_id: str,
|
||||
cache_key: str,
|
||||
prisma_client: PrismaClient,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
parent_otel_span: Span | None,
|
||||
proxy_logging_obj: ProxyLogging | None,
|
||||
) -> LiteLLM_TeamMembership | None:
|
||||
try:
|
||||
redis_cached: Final[object] = await user_api_key_cache.async_get_cache(key=cache_key)
|
||||
redis_membership: Final = _membership_from_cached_payload(redis_cached)
|
||||
if not isinstance(redis_membership, _TeamMembershipCacheMiss):
|
||||
return redis_membership
|
||||
|
||||
return await _fetch_team_membership_from_db(
|
||||
user_id=user_id,
|
||||
team_id=team_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
except Exception:
|
||||
verbose_proxy_logger.exception("Error getting team membership")
|
||||
return None
|
||||
|
||||
|
||||
async def get_team_membership(
|
||||
user_id: str,
|
||||
team_id: str,
|
||||
|
|
@ -2155,54 +2264,42 @@ async def get_team_membership(
|
|||
|
||||
Do a isolated check for team membership vs. doing a combined key + team + user + team-membership check, as key might come in frequently for different users/teams. Larger call will slowdown query time. This way we get to cache the constant (key/team/user info) and only update based on the changing value (team membership).
|
||||
"""
|
||||
from litellm.proxy._types import LiteLLM_TeamMembership
|
||||
|
||||
if prisma_client is None:
|
||||
raise Exception("No db connected")
|
||||
|
||||
if user_id is None or team_id is None:
|
||||
return None
|
||||
|
||||
_key: Final = team_membership_reservation_cache_key(user_id=user_id, team_id=team_id)
|
||||
|
||||
# check if in cache
|
||||
cached: Final[object] = await user_api_key_cache.async_get_cache(key=_key)
|
||||
if cached == NO_TEAM_MEMBERSHIP_SENTINEL:
|
||||
return None
|
||||
cached_membership_obj: Final = CacheCodec.deserialize(cached, model_type=LiteLLM_TeamMembership)
|
||||
if cached_membership_obj is not None:
|
||||
return cached_membership_obj
|
||||
l1_cached: Final[object] = await user_api_key_cache.async_get_cache(key=_key, local_only=True)
|
||||
l1_membership: Final = _membership_from_cached_payload(l1_cached)
|
||||
if not isinstance(l1_membership, _TeamMembershipCacheMiss):
|
||||
return l1_membership
|
||||
|
||||
# else, check db
|
||||
try:
|
||||
response: Final = await _dictable_table(TeamMembershipRepository(prisma_client)).find_unique(
|
||||
where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}},
|
||||
include={"litellm_budget_table": True},
|
||||
inflight: Final[object] = _team_membership_inflight.get(_key)
|
||||
if isinstance(inflight, asyncio.Task):
|
||||
return _membership_from_shared_load(await asyncio.shield(inflight))
|
||||
|
||||
if prisma_client is None:
|
||||
raise Exception("No db connected")
|
||||
|
||||
task: Final = asyncio.ensure_future(
|
||||
_load_team_membership_on_cache_miss(
|
||||
user_id=user_id,
|
||||
team_id=team_id,
|
||||
cache_key=_key,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
)
|
||||
_team_membership_inflight[_key] = task
|
||||
|
||||
if response is None:
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=_key,
|
||||
value=NO_TEAM_MEMBERSHIP_SENTINEL,
|
||||
ttl=get_management_object_ttl(user_api_key_cache),
|
||||
)
|
||||
return None
|
||||
def _clear_inflight(_done: object) -> None:
|
||||
if _team_membership_inflight.get(_key) is task:
|
||||
_team_membership_inflight.pop(_key, None)
|
||||
|
||||
_response: Final = LiteLLM_TeamMembership.model_validate(response.dict())
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=_key,
|
||||
value=_response,
|
||||
model_type=LiteLLM_TeamMembership,
|
||||
)
|
||||
|
||||
return _response
|
||||
except Exception:
|
||||
verbose_proxy_logger.exception(
|
||||
"Error getting team membership for user_id: %s, team_id: %s",
|
||||
user_id,
|
||||
team_id,
|
||||
)
|
||||
return None
|
||||
task.add_done_callback(_clear_inflight)
|
||||
return _membership_from_shared_load(await asyncio.shield(task))
|
||||
|
||||
|
||||
def model_in_access_group(model: str, team_models: list[str] | None, llm_router: Router | None) -> bool:
|
||||
|
|
@ -2375,13 +2472,6 @@ async def _backfill_null_user_email(
|
|||
return updated_row
|
||||
|
||||
|
||||
class UserNotFoundError(ValueError):
|
||||
"""The user row is provably absent, as opposed to merely unreadable, so a caller that reads a missing row as no user-level limits can key on it without also swallowing a database that would not answer."""
|
||||
|
||||
def __init__(self, user_id: str) -> None:
|
||||
super().__init__(f"User doesn't exist in db. 'user_id'={user_id}. Create user via `/user/new` call.")
|
||||
|
||||
|
||||
@log_db_metrics
|
||||
async def get_user_object(
|
||||
user_id: str | None,
|
||||
|
|
@ -2668,6 +2758,12 @@ async def invalidate_team_member_spend_state(
|
|||
publish_auth_cache_invalidation,
|
||||
)
|
||||
|
||||
inflight: Final[object] = _team_membership_inflight.pop(
|
||||
team_membership_reservation_cache_key(user_id=user_id, team_id=team_id), None
|
||||
)
|
||||
if isinstance(inflight, asyncio.Task) and inflight is not asyncio.current_task():
|
||||
await asyncio.wait((inflight,))
|
||||
|
||||
if new_spend is not None:
|
||||
from litellm.proxy.proxy_server import SPEND_DB_FLOOR_CACHE_TTL_SECONDS, spend_counter_cache
|
||||
|
||||
|
|
@ -3078,7 +3174,7 @@ async def _delete_cache_access_object(
|
|||
@log_db_metrics
|
||||
async def get_access_object(
|
||||
access_group_id: str,
|
||||
prisma_client: PrismaClient | None,
|
||||
prisma_client: DatabaseClient | None,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
proxy_logging_obj: ProxyLogging | None = None,
|
||||
) -> LiteLLM_AccessGroupTable:
|
||||
|
|
@ -3824,7 +3920,7 @@ async def get_org_object(
|
|||
async def _get_resources_from_access_groups(
|
||||
access_group_ids: Sequence[str],
|
||||
resource_field: Literal["access_model_names", "access_mcp_server_ids", "access_agent_ids"],
|
||||
prisma_client: PrismaClient | None = None,
|
||||
prisma_client: DatabaseClient | None = None,
|
||||
user_api_key_cache: UserApiKeyCache | None = None,
|
||||
proxy_logging_obj: ProxyLogging | None = None,
|
||||
) -> list[str]:
|
||||
|
|
@ -3882,7 +3978,7 @@ async def _get_resources_from_access_groups(
|
|||
|
||||
async def _get_models_from_access_groups(
|
||||
access_group_ids: Sequence[str],
|
||||
prisma_client: PrismaClient | None = None,
|
||||
prisma_client: DatabaseClient | None = None,
|
||||
user_api_key_cache: UserApiKeyCache | None = None,
|
||||
proxy_logging_obj: ProxyLogging | None = None,
|
||||
) -> list[str]:
|
||||
|
|
@ -4122,18 +4218,21 @@ async def _team_member_granted_models(
|
|||
prisma_client: PrismaClient,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
team_membership: LiteLLM_TeamMembership | None = None,
|
||||
team_membership_loaded: bool = False,
|
||||
) -> Sequence[str]:
|
||||
"""The member's own ``allowed_models`` scope; empty when the member is not narrowed below the team."""
|
||||
if team_object is None or valid_token.user_id is None:
|
||||
return ()
|
||||
|
||||
team_membership: Final = await get_team_membership(
|
||||
user_id=valid_token.user_id,
|
||||
team_id=team_object.team_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
if not team_membership_loaded:
|
||||
team_membership = await get_team_membership(
|
||||
user_id=valid_token.user_id,
|
||||
team_id=team_object.team_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
return () if team_membership is None else _member_allowed_models(team_membership)
|
||||
|
||||
|
||||
|
|
@ -4169,6 +4268,8 @@ async def _granted_model_lists(
|
|||
prisma_client: PrismaClient,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
team_membership: LiteLLM_TeamMembership | None = None,
|
||||
team_membership_loaded: bool = False,
|
||||
) -> tuple[Sequence[str], ...]:
|
||||
"""One model allowlist per level that participates in authorizing the request."""
|
||||
return (
|
||||
|
|
@ -4180,6 +4281,8 @@ async def _granted_model_lists(
|
|||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
team_membership=team_membership,
|
||||
team_membership_loaded=team_membership_loaded,
|
||||
),
|
||||
project_object.models if project_object is not None else (),
|
||||
await _org_granted_models(
|
||||
|
|
@ -4274,6 +4377,8 @@ async def collect_matched_model_access_groups(
|
|||
prisma_client: PrismaClient | None,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
team_membership: LiteLLM_TeamMembership | None = None,
|
||||
team_membership_loaded: bool = False,
|
||||
) -> tuple[str, ...]:
|
||||
"""
|
||||
The budgeted model access groups that authorized this request, sorted and deduplicated.
|
||||
|
|
@ -4319,6 +4424,8 @@ async def collect_matched_model_access_groups(
|
|||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
team_membership=team_membership,
|
||||
team_membership_loaded=team_membership_loaded,
|
||||
)
|
||||
for granted_model in granted_models
|
||||
)
|
||||
|
|
@ -4334,6 +4441,8 @@ async def stamp_matched_model_access_groups(
|
|||
prisma_client: PrismaClient | None,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
team_membership: LiteLLM_TeamMembership | None = None,
|
||||
team_membership_loaded: bool = False,
|
||||
) -> tuple[str, ...]:
|
||||
"""Record the groups that authorized this request on its auth object, for the post-call spend
|
||||
writer and the reservation counters, and hand them back for the budget check."""
|
||||
|
|
@ -4350,6 +4459,8 @@ async def stamp_matched_model_access_groups(
|
|||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
team_membership=team_membership,
|
||||
team_membership_loaded=team_membership_loaded,
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # fail-safe: attribution is spend telemetry, it must never break auth
|
||||
verbose_proxy_logger.debug("model access group attribution failed: %s", e)
|
||||
|
|
@ -4363,9 +4474,10 @@ 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,
|
||||
prisma_client: DatabaseClient | None = None,
|
||||
) -> Literal[True]:
|
||||
"""
|
||||
Checks if token can call a given model
|
||||
|
|
@ -4395,6 +4507,7 @@ async def can_key_call_model(
|
|||
if key_access_group_ids:
|
||||
models_from_groups: Final = await _get_models_from_access_groups(
|
||||
access_group_ids=key_access_group_ids,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
if models_from_groups:
|
||||
return _can_object_call_model(
|
||||
|
|
@ -4410,7 +4523,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:
|
||||
|
|
@ -4523,6 +4636,7 @@ async def can_team_access_model(
|
|||
team_object: LiteLLM_TeamTable | None,
|
||||
llm_router: Router | None,
|
||||
team_model_aliases: dict[str, str] | None = None,
|
||||
prisma_client: DatabaseClient | None = None,
|
||||
) -> Literal[True]:
|
||||
"""
|
||||
Returns True if the team can access a specific model.
|
||||
|
|
@ -4545,12 +4659,13 @@ async def can_team_access_model(
|
|||
if team_access_group_ids:
|
||||
models_from_groups: Final = await _get_models_from_access_groups(
|
||||
access_group_ids=team_access_group_ids,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
if models_from_groups:
|
||||
return _can_object_call_model(
|
||||
model=model,
|
||||
llm_router=llm_router,
|
||||
models=models_from_groups,
|
||||
models=list(dict.fromkeys([*(team_object.models if team_object else []), *models_from_groups])),
|
||||
team_model_aliases=team_model_aliases,
|
||||
team_id=team_object.team_id if team_object else None,
|
||||
object_type="team",
|
||||
|
|
@ -4640,7 +4755,7 @@ async def _key_access_group_grants_model(
|
|||
|
||||
def can_project_access_model(
|
||||
model: str | list[str],
|
||||
project_object: LiteLLM_ProjectTableCachedObj,
|
||||
project_object: LiteLLM_ProjectTable,
|
||||
llm_router: Router | None,
|
||||
) -> Literal[True]:
|
||||
"""
|
||||
|
|
@ -5152,6 +5267,8 @@ async def _check_team_member_budget(
|
|||
prisma_client: PrismaClient | None,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
team_membership: LiteLLM_TeamMembership | None = None,
|
||||
team_membership_loaded: bool = False,
|
||||
):
|
||||
"""Check if team member is over their max budget within the team."""
|
||||
if (
|
||||
|
|
@ -5160,23 +5277,25 @@ async def _check_team_member_budget(
|
|||
and valid_token is not None
|
||||
and valid_token.user_id is not None
|
||||
):
|
||||
team_membership: Final = await get_team_membership(
|
||||
user_id=valid_token.user_id,
|
||||
team_id=team_object.team_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
if not team_membership_loaded:
|
||||
team_membership = await get_team_membership(
|
||||
user_id=valid_token.user_id,
|
||||
team_id=team_object.team_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
loaded_membership = team_membership
|
||||
|
||||
# Per-member override wins; otherwise fall back to the team-level
|
||||
# default configured via team.metadata["team_member_budget_id"].
|
||||
team_member_budget: float | None = None
|
||||
if (
|
||||
team_membership is not None
|
||||
and team_membership.litellm_budget_table is not None
|
||||
and team_membership.litellm_budget_table.max_budget is not None
|
||||
loaded_membership is not None
|
||||
and loaded_membership.litellm_budget_table is not None
|
||||
and loaded_membership.litellm_budget_table.max_budget is not None
|
||||
):
|
||||
team_member_budget = team_membership.litellm_budget_table.max_budget
|
||||
team_member_budget = loaded_membership.litellm_budget_table.max_budget
|
||||
else:
|
||||
default_budget_id: Final = (team_object.metadata or {}).get("team_member_budget_id")
|
||||
if isinstance(default_budget_id, str):
|
||||
|
|
@ -5195,7 +5314,7 @@ async def _check_team_member_budget(
|
|||
team_member_budget = default_budget.max_budget
|
||||
|
||||
if team_member_budget is not None:
|
||||
team_member_spend = (team_membership.spend if team_membership is not None else 0.0) or 0.0
|
||||
team_member_spend = (loaded_membership.spend if loaded_membership is not None else 0.0) or 0.0
|
||||
|
||||
# Read from cross-pod counter (Redis-first) if available
|
||||
from litellm.proxy.proxy_server import get_current_spend
|
||||
|
|
@ -5224,6 +5343,8 @@ async def _check_team_member_model_access(
|
|||
prisma_client: Optional["PrismaClient"],
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
team_membership: LiteLLM_TeamMembership | None = None,
|
||||
team_membership_loaded: bool = False,
|
||||
) -> None:
|
||||
"""
|
||||
Check if a team member's per-member model scope allows access to the requested model.
|
||||
|
|
@ -5234,22 +5355,24 @@ async def _check_team_member_model_access(
|
|||
if valid_token.user_id is None or team_object.team_id is None:
|
||||
return
|
||||
|
||||
team_membership: Final = await get_team_membership(
|
||||
user_id=valid_token.user_id,
|
||||
team_id=team_object.team_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
if not team_membership_loaded:
|
||||
team_membership = await get_team_membership(
|
||||
user_id=valid_token.user_id,
|
||||
team_id=team_object.team_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
loaded_membership = team_membership
|
||||
|
||||
if (
|
||||
team_membership is None
|
||||
or team_membership.litellm_budget_table is None
|
||||
or not team_membership.litellm_budget_table.allowed_models
|
||||
loaded_membership is None
|
||||
or loaded_membership.litellm_budget_table is None
|
||||
or not loaded_membership.litellm_budget_table.allowed_models
|
||||
):
|
||||
return # no per-member restriction — inherit team-level check
|
||||
|
||||
member_allowed_models: Final[list[str]] = team_membership.litellm_budget_table.allowed_models
|
||||
member_allowed_models: Final[list[str]] = loaded_membership.litellm_budget_table.allowed_models
|
||||
try:
|
||||
_can_object_call_model(
|
||||
model=model,
|
||||
|
|
@ -5650,8 +5773,7 @@ async def _organization_max_budget_check(
|
|||
if org_table.litellm_budget_table is not None:
|
||||
org_max_budget = org_table.litellm_budget_table.max_budget
|
||||
|
||||
# Only check if organization has a valid max_budget set
|
||||
if org_max_budget is None or org_max_budget <= 0:
|
||||
if org_max_budget is None:
|
||||
return
|
||||
|
||||
# Read spend from cross-pod counter (Redis-first) or cached object (fallback)
|
||||
|
|
|
|||
|
|
@ -23,6 +23,7 @@ from litellm.proxy.auth.auth_utils import (
|
|||
_get_request_ip_address,
|
||||
is_invalid_virtual_key_error,
|
||||
mark_invalid_virtual_key_error,
|
||||
normalize_request_route,
|
||||
)
|
||||
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
|
||||
from litellm.types.services import ServiceTypes
|
||||
|
|
@ -74,17 +75,28 @@ def _as_proxy_exception(e: Exception) -> ProxyException:
|
|||
)
|
||||
|
||||
|
||||
def _with_requester_ip_address(request_data: dict[str, object], requester_ip: str | None) -> dict[str, object]:
|
||||
def _get_user_agent(request: Request) -> str | None:
|
||||
if "headers" not in request.scope:
|
||||
return None
|
||||
return request.headers.get("user-agent")
|
||||
|
||||
|
||||
def _with_client_context(
|
||||
request_data: dict[str, object], requester_ip: str | None, user_agent: str | None
|
||||
) -> dict[str, object]:
|
||||
"""Auth gate rejections are raised before `add_litellm_data_to_request` records the
|
||||
caller IP, so their failure logs would otherwise carry no IP nor key/user identity."""
|
||||
if not requester_ip:
|
||||
return request_data
|
||||
caller IP and User-Agent, so their failure logs would otherwise carry neither."""
|
||||
key: Final = "litellm_metadata" if "litellm_metadata" in request_data else "metadata"
|
||||
metadata: Final = request_data.get(key)
|
||||
base: Final[Mapping[str, object]] = metadata if isinstance(metadata, Mapping) else EMPTY_MAPPING
|
||||
if base.get("requester_ip_address"):
|
||||
stamped: Final = {
|
||||
name: value
|
||||
for name, value in (("requester_ip_address", requester_ip), ("user_agent", user_agent))
|
||||
if value and not base.get(name)
|
||||
}
|
||||
if not stamped:
|
||||
return request_data
|
||||
return {**request_data, key: {**base, "requester_ip_address": requester_ip}} # mutable-ok: logging needs dicts
|
||||
return {**request_data, key: {**base, **stamped}} # mutable-ok: logging needs dicts
|
||||
|
||||
|
||||
class UserAPIKeyAuthExceptionHandler:
|
||||
|
|
@ -148,6 +160,7 @@ class UserAPIKeyAuthExceptionHandler:
|
|||
request=request,
|
||||
use_x_forwarded_for=general_settings.get("use_x_forwarded_for") is True,
|
||||
)
|
||||
user_agent: Final = _get_user_agent(request)
|
||||
|
||||
# Log authentication failures before identity seeding and callbacks, so the log
|
||||
# survives a raising callback pipeline. Classify and route malformed virtual-key
|
||||
|
|
@ -172,7 +185,7 @@ class UserAPIKeyAuthExceptionHandler:
|
|||
# so the handler is side-effect-free for the caller's identity object.
|
||||
user_api_key_dict = resolved_identity.model_copy() if resolved_identity is not None else UserAPIKeyAuth()
|
||||
user_api_key_dict.parent_otel_span = parent_otel_span
|
||||
user_api_key_dict.request_route = route
|
||||
user_api_key_dict.request_route = normalize_request_route(route)
|
||||
user_api_key_dict.api_key = user_api_key_dict.api_key or UserAPIKeyAuth(api_key=api_key).api_key
|
||||
|
||||
# Stamp identity onto the request's server span now, before the request
|
||||
|
|
@ -200,7 +213,7 @@ class UserAPIKeyAuthExceptionHandler:
|
|||
|
||||
# Allow callbacks to transform the error response
|
||||
transformed_exception: Final = await proxy_logging_obj.post_call_failure_hook(
|
||||
request_data=_with_requester_ip_address(request_data, requester_ip),
|
||||
request_data=_with_client_context(request_data, requester_ip, user_agent),
|
||||
original_exception=e,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
error_type=ProxyErrorTypes.auth_error,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
136
litellm/proxy/auth/auto_router_checks.py
Normal file
136
litellm/proxy/auth/auto_router_checks.py
Normal file
|
|
@ -0,0 +1,136 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.router import Router
|
||||
|
||||
_MAPPING_ADAPTER: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def _mapping(value: object) -> Mapping[str, object] | None:
|
||||
try:
|
||||
return _MAPPING_ADAPTER.validate_python(value)
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
async def authorize_member_auto_router_inference(
|
||||
*,
|
||||
deployment: Mapping[str, object] | None,
|
||||
request_kwargs: Mapping[str, object],
|
||||
llm_router: Router,
|
||||
) -> None:
|
||||
if deployment is None:
|
||||
return
|
||||
model_info: Final = _mapping(deployment.get("model_info"))
|
||||
if model_info is None or model_info.get("member_auto_router") is not True:
|
||||
return
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.auth.auth_checks import (
|
||||
OrganizationNotFoundError,
|
||||
TeamNotFoundError,
|
||||
get_org_object,
|
||||
get_project_object,
|
||||
get_team_membership,
|
||||
get_team_object,
|
||||
)
|
||||
from litellm.proxy.management_helpers.auto_router_permissions import (
|
||||
MemberAutoRouterDependencyObjects,
|
||||
authorize_member_auto_router_dependencies,
|
||||
validate_member_auto_router_config,
|
||||
)
|
||||
|
||||
metadata: Final = _mapping(request_kwargs.get(get_metadata_variable_name_from_kwargs(request_kwargs)))
|
||||
actor: Final = metadata.get("user_api_key_auth") if metadata is not None else None
|
||||
team_id: Final = model_info.get("team_id")
|
||||
if not isinstance(actor, UserAPIKeyAuth) or not isinstance(team_id, str) or not team_id:
|
||||
raise HTTPException(status_code=403, detail="Member auto-routers require authenticated team access")
|
||||
if actor.team_id != team_id and actor.user_role != LitellmUserRoles.PROXY_ADMIN:
|
||||
raise HTTPException(status_code=403, detail="This auto-router belongs to a different team")
|
||||
|
||||
from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache
|
||||
|
||||
if prisma_client is None:
|
||||
raise HTTPException(status_code=503, detail="Cannot verify auto-router model access without a database")
|
||||
try:
|
||||
team: Final = await get_team_object(
|
||||
team_id=team_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=actor.parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
except TeamNotFoundError as error:
|
||||
raise HTTPException(status_code=403, detail="The auto-router team no longer exists") from error
|
||||
if (
|
||||
actor.user_role != LitellmUserRoles.PROXY_ADMIN
|
||||
and actor.user_id is not None
|
||||
and (not actor.user_id or not any(member.user_id == actor.user_id for member in team.members_with_roles))
|
||||
):
|
||||
raise HTTPException(status_code=403, detail="You are no longer a member of this auto-router's team")
|
||||
if team.blocked:
|
||||
raise HTTPException(status_code=403, detail="This auto router's team is blocked.")
|
||||
params: Final = _mapping(deployment.get("litellm_params"))
|
||||
if params is None:
|
||||
raise HTTPException(status_code=403, detail="The member auto-router configuration is invalid")
|
||||
raw_config: Final = _mapping(params.get("complexity_router_config"))
|
||||
if raw_config is None:
|
||||
raise HTTPException(status_code=403, detail="The member auto-router configuration is invalid")
|
||||
default_model: Final = params.get("complexity_router_default_model")
|
||||
config: Final = validate_member_auto_router_config(raw_config)
|
||||
membership: Final = (
|
||||
await get_team_membership(
|
||||
user_id=actor.user_id,
|
||||
team_id=team_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=actor.parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
if actor.user_id
|
||||
else None
|
||||
)
|
||||
try:
|
||||
organization: Final = (
|
||||
await get_org_object(
|
||||
org_id=team.organization_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=actor.parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
if team.organization_id
|
||||
else None
|
||||
)
|
||||
except OrganizationNotFoundError as error:
|
||||
raise HTTPException(status_code=403, detail="The auto router's organization is unavailable.") from error
|
||||
project: Final = (
|
||||
await get_project_object(
|
||||
project_id=actor.project_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
if actor.project_id
|
||||
else None
|
||||
)
|
||||
await authorize_member_auto_router_dependencies(
|
||||
config=config,
|
||||
default_model=default_model if isinstance(default_model, str) else None,
|
||||
user_api_key_dict=actor,
|
||||
team=team,
|
||||
prisma_client=None,
|
||||
llm_router=llm_router,
|
||||
dependency_objects=MemberAutoRouterDependencyObjects(
|
||||
membership=membership, organization=organization, project=project
|
||||
),
|
||||
)
|
||||
|
|
@ -14,7 +14,7 @@ import hashlib
|
|||
import os
|
||||
import re
|
||||
import time
|
||||
from collections.abc import Awaitable, Callable, Sequence
|
||||
from collections.abc import Awaitable, Callable, Mapping, Sequence
|
||||
from typing import Any, Final, Literal, NoReturn, Protocol, TypeVar, cast
|
||||
|
||||
import httpx
|
||||
|
|
@ -61,6 +61,7 @@ from litellm.proxy.common_utils.user_api_key_cache import (
|
|||
)
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
||||
from litellm.repositories.user_repository import UserRepository
|
||||
from litellm.types.agents import AgentResponse
|
||||
|
||||
from .auth_checks import (
|
||||
_allowed_routes_check,
|
||||
|
|
@ -127,6 +128,26 @@ class _UserInfoResponse(Protocol):
|
|||
def json(self) -> dict[str, object]: ...
|
||||
|
||||
|
||||
class AgentLookup(Protocol):
|
||||
"""The registered-agent lookups a JWT agent claim is matched against."""
|
||||
|
||||
def get_agent_by_id(self, agent_id: str) -> AgentResponse | None:
|
||||
"""The agent registered under ``agent_id``, if any."""
|
||||
|
||||
def get_agent_by_name(self, agent_name: str) -> AgentResponse | None:
|
||||
"""The agent registered under ``agent_name``, if any."""
|
||||
|
||||
|
||||
class _NoRegisteredAgents:
|
||||
"""The lookup in force until the proxy binds its agent registry: no agent is registered, so no claim matches."""
|
||||
|
||||
def get_agent_by_id(self, agent_id: str) -> None:
|
||||
return None
|
||||
|
||||
def get_agent_by_name(self, agent_name: str) -> None:
|
||||
return None
|
||||
|
||||
|
||||
def _discovery_document(response: _OIDCDiscoveryResponse) -> _OIDCDiscoveryBody:
|
||||
"""Decode an OIDC discovery response body."""
|
||||
return response.json()
|
||||
|
|
@ -198,6 +219,10 @@ class JWTHandler:
|
|||
self.leeway = 0
|
||||
# Per-cache-key locks so a TTL lapse triggers one refresh instead of one per in-flight request.
|
||||
self._refresh_locks: dict[str, asyncio.Lock] = {} # mutable-ok: lock registry, keyed by JWKS url
|
||||
self.agent_lookup: AgentLookup = _NoRegisteredAgents()
|
||||
|
||||
def bind_agent_lookup(self, agent_lookup: AgentLookup) -> None:
|
||||
self.agent_lookup = agent_lookup
|
||||
|
||||
def update_environment(
|
||||
self,
|
||||
|
|
@ -623,6 +648,12 @@ class JWTHandler:
|
|||
object_id = default_value
|
||||
return object_id
|
||||
|
||||
def get_agent_claim(self, token: Mapping[str, object]) -> str | None:
|
||||
if self.litellm_jwtauth.agent_id_jwt_field is None:
|
||||
return None
|
||||
claim: Final[object] = get_nested_value(data=token, key_path=self.litellm_jwtauth.agent_id_jwt_field)
|
||||
return claim if isinstance(claim, str) and claim else None
|
||||
|
||||
def get_org_id(self, token: dict, default_value: str | None) -> str | None:
|
||||
if self._has_trusted_issuer_normalized_claim(token=token, claim=self.LITELLM_ORG_ID_CLAIM):
|
||||
return token.get(self.LITELLM_ORG_ID_CLAIM)
|
||||
|
|
@ -1380,6 +1411,7 @@ class JWTAuthManager:
|
|||
api_key: str,
|
||||
jwt_valid_token: dict | None = None,
|
||||
user_email: str | None = None,
|
||||
agent_id: str | None = None,
|
||||
) -> JWTAuthBuilderResult | None:
|
||||
"""Check admin status and route access permissions"""
|
||||
if not jwt_handler.is_admin(scopes=scopes):
|
||||
|
|
@ -1409,8 +1441,28 @@ class JWTAuthManager:
|
|||
org_id=org_id,
|
||||
team_membership=None,
|
||||
jwt_claims=jwt_valid_token or {},
|
||||
agent_id=agent_id,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def resolve_agent_id(
|
||||
jwt_handler: JWTHandler,
|
||||
jwt_valid_token: Mapping[str, object],
|
||||
agent_registry: AgentLookup,
|
||||
) -> str | None:
|
||||
agent_claim: Final = jwt_handler.get_agent_claim(token=jwt_valid_token)
|
||||
if agent_claim is None:
|
||||
return None
|
||||
agent: Final = agent_registry.get_agent_by_id(agent_id=agent_claim) or agent_registry.get_agent_by_name(
|
||||
agent_name=agent_claim
|
||||
)
|
||||
if agent is None:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail=f"No registered agent matches JWT claim {jwt_handler.litellm_jwtauth.agent_id_jwt_field}={agent_claim}",
|
||||
)
|
||||
return agent.agent_id
|
||||
|
||||
@staticmethod
|
||||
async def find_and_validate_specific_team_id(
|
||||
jwt_handler: JWTHandler,
|
||||
|
|
@ -2268,9 +2320,23 @@ class JWTAuthManager:
|
|||
elif rbac_role == LitellmUserRoles.INTERNAL_USER:
|
||||
user_id = object_id
|
||||
|
||||
agent_id: Final = JWTAuthManager.resolve_agent_id(
|
||||
jwt_handler=jwt_handler,
|
||||
jwt_valid_token=jwt_valid_token,
|
||||
agent_registry=jwt_handler.agent_lookup,
|
||||
)
|
||||
|
||||
# Check admin access
|
||||
admin_result: Final = await JWTAuthManager.check_admin_access(
|
||||
jwt_handler, scopes, route, user_id, org_id, api_key, jwt_valid_token, user_email=user_email
|
||||
jwt_handler,
|
||||
scopes,
|
||||
route,
|
||||
user_id,
|
||||
org_id,
|
||||
api_key,
|
||||
jwt_valid_token,
|
||||
user_email=user_email,
|
||||
agent_id=agent_id,
|
||||
)
|
||||
if admin_result:
|
||||
await JWTAuthManager._attach_team_from_header_for_admin(
|
||||
|
|
@ -2514,4 +2580,5 @@ class JWTAuthManager:
|
|||
token=api_key,
|
||||
team_membership=team_membership_object,
|
||||
jwt_claims=jwt_valid_token,
|
||||
agent_id=agent_id,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -28,11 +28,11 @@ from litellm.proxy._types import (
|
|||
)
|
||||
from litellm.proxy.auth.auth_checks import (
|
||||
TeamNotFoundError,
|
||||
UserNotFoundError,
|
||||
get_team_membership,
|
||||
get_team_object,
|
||||
get_user_object,
|
||||
)
|
||||
from litellm.types.proxy.auth.auth_checks import UserNotFoundError
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy._types import Span
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
import re
|
||||
from collections.abc import Sequence
|
||||
from collections.abc import Collection
|
||||
from typing import Final
|
||||
|
||||
from fastapi import HTTPException, Request, status
|
||||
|
|
@ -24,10 +24,13 @@ _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",
|
||||
# team
|
||||
"/team/new",
|
||||
"/management/v1/teams/{team_id}/members/bulk_delete",
|
||||
"/team/update",
|
||||
"/team/delete",
|
||||
"/team/block",
|
||||
|
|
@ -136,6 +139,9 @@ class RouteChecks:
|
|||
# For llm_api_routes, also check registered pass-through endpoints
|
||||
################################################
|
||||
if allowed_route == "llm_api_routes":
|
||||
if route == "/auto_router/session" and RouteChecks._get_request_method(request) == "GET":
|
||||
return True
|
||||
|
||||
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
|
||||
InitPassThroughEndpointHelpers,
|
||||
)
|
||||
|
|
@ -584,7 +590,7 @@ class RouteChecks:
|
|||
return False
|
||||
|
||||
@staticmethod
|
||||
def check_route_access(route: str, allowed_routes: Sequence[str]) -> bool:
|
||||
def check_route_access(route: str, allowed_routes: Collection[str]) -> bool:
|
||||
"""
|
||||
Check if a route has access by checking both exact matches and patterns
|
||||
|
||||
|
|
@ -755,9 +761,12 @@ class RouteChecks:
|
|||
_ADMIN_VIEWER_BLOCKED_WRITE_ROUTES = frozenset(
|
||||
[
|
||||
"/user/new",
|
||||
"/management/v1/users/bulk",
|
||||
"/user/delete",
|
||||
"/management/v1/users/bulk_delete",
|
||||
"/user/bulk_update",
|
||||
"/team/new",
|
||||
"/management/v1/teams/{team_id}/members/bulk_delete",
|
||||
"/team/update",
|
||||
"/team/delete",
|
||||
"/model/new",
|
||||
|
|
@ -821,7 +830,7 @@ class RouteChecks:
|
|||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail=f"user not allowed to access this route, role= {_user_role}. Trying to access: {route} and updating invalid param: {param}. only user_email and password can be updated",
|
||||
)
|
||||
elif route in _PROXY_ADMIN_VIEW_ONLY_BLOCKED_ROUTES or (
|
||||
elif RouteChecks.check_route_access(route=route, allowed_routes=_PROXY_ADMIN_VIEW_ONLY_BLOCKED_ROUTES) or (
|
||||
route.startswith("/key/") and route.endswith(_PROXY_ADMIN_VIEW_ONLY_BLOCKED_KEY_SUFFIXES)
|
||||
):
|
||||
# Block write operations for PROXY_ADMIN_VIEW_ONLY
|
||||
|
|
@ -856,9 +865,9 @@ class RouteChecks:
|
|||
# Hard-block known write routes regardless of HTTP method (defensive
|
||||
# — these are POSTs in practice, but pinning them here protects
|
||||
# against future GET-shaped writes).
|
||||
if route in RouteChecks._ADMIN_VIEWER_BLOCKED_WRITE_ROUTES or (
|
||||
route.startswith("/key/") and route.endswith("/regenerate")
|
||||
):
|
||||
if RouteChecks.check_route_access(
|
||||
route=route, allowed_routes=RouteChecks._ADMIN_VIEWER_BLOCKED_WRITE_ROUTES
|
||||
) or (route.startswith("/key/") and route.endswith("/regenerate")):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail=f"user not allowed to access this route, role= {_user_role}. Trying to access: {route}",
|
||||
|
|
|
|||
|
|
@ -56,6 +56,7 @@ class TeamGrants(TypedDict, total=False):
|
|||
team_alias: ReadOnly[str | None]
|
||||
team_tpm_limit: ReadOnly[int | None]
|
||||
team_rpm_limit: ReadOnly[int | None]
|
||||
team_tpd_limit: ReadOnly[int | None]
|
||||
team_max_budget: ReadOnly[float | None]
|
||||
team_soft_budget: ReadOnly[float | None]
|
||||
team_spend: ReadOnly[float | None]
|
||||
|
|
@ -97,6 +98,7 @@ def team_grants(
|
|||
team_alias=team_object.team_alias,
|
||||
team_tpm_limit=team_object.tpm_limit,
|
||||
team_rpm_limit=team_object.rpm_limit,
|
||||
team_tpd_limit=team_object.tpd_limit,
|
||||
team_max_budget=team_object.max_budget,
|
||||
team_soft_budget=team_object.soft_budget,
|
||||
team_spend=team_object.spend,
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue