mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
Merge branch 'main' of https://github.com/BerriAI/litellm into fix/databricks-gpt5-function-tools-reasoning
# Conflicts: # litellm/responses/bridge_check.py
This commit is contained in:
commit
ca7deee3d5
2279 changed files with 310831 additions and 36677 deletions
|
|
@ -38,7 +38,7 @@ commands:
|
|||
parameters:
|
||||
category:
|
||||
type: enum
|
||||
enum: ["backend", "client", "provider-harness", "redis-compat"]
|
||||
enum: ["backend", "client"]
|
||||
default: "backend"
|
||||
steps:
|
||||
- run:
|
||||
|
|
@ -339,8 +339,15 @@ commands:
|
|||
Add-Content -Path $PROFILE -Value "`$env:Path = `"$cargoBin;`$env:Path`""
|
||||
}
|
||||
setup_litellm_test_deps:
|
||||
parameters:
|
||||
checkout:
|
||||
type: boolean
|
||||
default: true
|
||||
steps:
|
||||
- checkout
|
||||
- when:
|
||||
condition: << parameters.checkout >>
|
||||
steps:
|
||||
- checkout
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- install_rust
|
||||
|
|
@ -376,13 +383,28 @@ jobs:
|
|||
shell: bash.exe
|
||||
command: bash .circleci/scripts/path_filter.sh windows-release
|
||||
- run:
|
||||
name: Install Python
|
||||
name: Install Python 3.11.9 from python.org
|
||||
no_output_timeout: 6m
|
||||
command: |
|
||||
choco install python --version=3.11.0 -y --no-progress --force
|
||||
refreshenv
|
||||
python --version
|
||||
environment:
|
||||
CHOCOLATEY_CONFIRM_ALL: "true"
|
||||
$ProgressPreference = "SilentlyContinue"
|
||||
$installer = Join-Path $env:TEMP "python-3.11.9-amd64.exe"
|
||||
Invoke-WebRequest -Uri https://www.python.org/ftp/python/3.11.9/python-3.11.9-amd64.exe -OutFile $installer -TimeoutSec 120
|
||||
$expected = "5ee42c4eee1e6b4464bb23722f90b45303f79442df63083f05322f1785f5fdde"
|
||||
$actual = (Get-FileHash -Path $installer -Algorithm SHA256).Hash.ToLower()
|
||||
if ($actual -ne $expected) {
|
||||
throw "Python installer hash mismatch: expected $expected got $actual"
|
||||
}
|
||||
$arguments = "/quiet InstallAllUsers=1 PrependPath=1 Include_test=0 Include_doc=0 Include_launcher=0 TargetDir=C:\Python311"
|
||||
$process = Start-Process -FilePath $installer -ArgumentList $arguments -PassThru
|
||||
if (-not $process.WaitForExit(300000)) {
|
||||
$process.Kill()
|
||||
throw "Python installer did not finish within 5 minutes"
|
||||
}
|
||||
if ($process.ExitCode -ne 0) {
|
||||
throw "Python installer exited with $($process.ExitCode)"
|
||||
}
|
||||
Remove-Item $installer
|
||||
& C:\Python311\python.exe --version
|
||||
- install_windows_toolchain
|
||||
- run:
|
||||
name: Record the Rust build environment for the release cargo cache key
|
||||
|
|
@ -405,7 +427,7 @@ jobs:
|
|||
environment:
|
||||
UV_HTTP_TIMEOUT: "300"
|
||||
command: |
|
||||
$env:Path = "$HOME\.cargo\bin;$HOME\.local\bin;$env:Path"
|
||||
$env:Path = "C:\Python311;C:\Python311\Scripts;$HOME\.cargo\bin;$HOME\.local\bin;$env:Path"
|
||||
uv build --wheel --out-dir dist
|
||||
if ($LASTEXITCODE -ne 0) {
|
||||
exit $LASTEXITCODE
|
||||
|
|
@ -512,11 +534,12 @@ jobs:
|
|||
--set image.repository=litellm-ci \
|
||||
--set image.tag=${IMAGE_TAG} \
|
||||
--set image.pullPolicy=Never
|
||||
# Wait for pod to be ready
|
||||
echo "Waiting 30 seconds for pod to be ready..."
|
||||
sleep 30
|
||||
if ! kubectl rollout status deployment/litellm --timeout=180s; then
|
||||
kubectl describe pods -l app.kubernetes.io/name=litellm
|
||||
kubectl logs -l app.kubernetes.io/name=litellm --tail=200
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Print pod logs before running tests
|
||||
echo "Printing pod logs..."
|
||||
kubectl logs $(kubectl get pods -l app.kubernetes.io/name=litellm -o jsonpath="{.items[0].metadata.name}")
|
||||
|
||||
|
|
@ -619,7 +642,7 @@ jobs:
|
|||
name: Run Basic Proxy Startup Tests (Health Readiness and Chat Completion)
|
||||
command: |
|
||||
mkdir -p test-results
|
||||
TEST_FILES=$(circleci tests glob "tests/basic_proxy_startup_tests/**/test_*.py")
|
||||
TEST_FILES=$(circleci tests glob "tests/integration/configuration/test_basic_proxy_startup.py")
|
||||
echo "$TEST_FILES" | circleci tests run \
|
||||
--verbose \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest --tb=short \
|
||||
|
|
@ -696,7 +719,7 @@ jobs:
|
|||
name: Run second round of tests
|
||||
command: |
|
||||
mkdir -p test-results
|
||||
TEST_FILES=$(circleci tests glob "tests/basic_proxy_startup_tests/**/test_*.py")
|
||||
TEST_FILES=$(circleci tests glob "tests/integration/configuration/test_basic_proxy_startup.py")
|
||||
echo "$TEST_FILES" | circleci tests run \
|
||||
--verbose \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest --tb=short \
|
||||
|
|
@ -708,84 +731,46 @@ jobs:
|
|||
# Store test results
|
||||
- store_test_results:
|
||||
path: test-results
|
||||
proxy_pass_through_endpoint_tests:
|
||||
machine:
|
||||
image: ubuntu-2204:2024.04.1
|
||||
e2e_ui_build:
|
||||
docker:
|
||||
- image: cimg/python:3.12-browsers@sha256:b432899af01c9a311bf74f4f22e9ada2e5306d4b1b4383f8d29e1228a5844ef2
|
||||
auth:
|
||||
username: ${DOCKERHUB_USERNAME}
|
||||
password: ${DOCKERHUB_PASSWORD}
|
||||
resource_class: large
|
||||
working_directory: ~/project
|
||||
steps:
|
||||
- checkout
|
||||
- run:
|
||||
name: Generate LiteLLM master key
|
||||
command: |
|
||||
key="$(openssl rand -hex 16)"
|
||||
printf 'export LITELLM_MASTER_KEY=sk-%s\n' "$key" >> "$BASH_ENV"
|
||||
- skip_if_unrelated_changes
|
||||
- skip_if_unrelated_changes:
|
||||
category: client
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- install_rust
|
||||
- run:
|
||||
name: Install Dependencies
|
||||
command: |
|
||||
uv sync --frozen --all-groups --all-extras --python 3.12
|
||||
- start_postgres
|
||||
- attach_workspace:
|
||||
at: ~/project
|
||||
- run:
|
||||
name: Load Docker Database Image
|
||||
command: |
|
||||
zstd -d litellm-docker-database.tar.zst --stdout | docker load
|
||||
docker images | grep litellm-docker-database
|
||||
- run:
|
||||
name: Run Docker container
|
||||
command: |
|
||||
docker run -d \
|
||||
-p 4000:4000 \
|
||||
-e DATABASE_URL=postgresql://postgres:postgres@host.docker.internal:5432/circle_test \
|
||||
-e LITELLM_MASTER_KEY="$LITELLM_MASTER_KEY" \
|
||||
-e OPENAI_API_KEY=$OPENAI_API_KEY \
|
||||
-e GEMINI_API_KEY=$GEMINI_API_KEY \
|
||||
-e ANTHROPIC_API_KEY=$ANTHROPIC_API_KEY \
|
||||
-e ASSEMBLYAI_API_KEY=$ASSEMBLYAI_API_KEY \
|
||||
-e AZURE_API_KEY=$AZURE_API_KEY \
|
||||
-e AZURE_API_BASE=$AZURE_API_BASE \
|
||||
-e USE_DDTRACE=True \
|
||||
-e DD_API_KEY=$DD_API_KEY \
|
||||
-e DD_SITE=$DD_SITE \
|
||||
-e LITELLM_LICENSE=$LITELLM_LICENSE \
|
||||
-e LITELLM_USE_CHAT_COMPLETIONS_URL_FOR_ANTHROPIC_MESSAGES=true \
|
||||
-e LITELLM_LOG=ERROR \
|
||||
--add-host host.docker.internal:host-gateway \
|
||||
--name my-app \
|
||||
-v $(pwd)/litellm/proxy/example_config_yaml/pass_through_config.yaml:/app/config.yaml \
|
||||
-v $(pwd)/litellm/proxy/example_config_yaml/custom_auth_basic.py:/app/custom_auth_basic.py \
|
||||
litellm-docker-database:ci \
|
||||
--config /app/config.yaml \
|
||||
--port 4000
|
||||
- run:
|
||||
name: Start outputting logs
|
||||
command: docker logs -f my-app
|
||||
background: true
|
||||
- wait_for_service:
|
||||
url: http://localhost:4000
|
||||
timeout: "300"
|
||||
# Install Node.js directly from nodejs.org with SHA256 verification,
|
||||
# instead of piping NodeSource's setup_24.x apt-repo installer into
|
||||
# sudo bash (which runs a mutable upstream script unattended).
|
||||
- install_node
|
||||
|
||||
- restore_cache:
|
||||
keys:
|
||||
- ui-dashboard-node-deps-v1-{{ checksum "ui/litellm-dashboard/package-lock.json" }}
|
||||
- run:
|
||||
name: Install Node.js test dependencies
|
||||
name: Install dashboard dependencies
|
||||
command: |
|
||||
cd tests/pass_through_tests
|
||||
cd ui/litellm-dashboard
|
||||
npm ci
|
||||
|
||||
- save_cache:
|
||||
key: ui-dashboard-node-deps-v1-{{ checksum "ui/litellm-dashboard/package-lock.json" }}
|
||||
paths:
|
||||
- ui/litellm-dashboard/node_modules
|
||||
- run:
|
||||
name: Run Vertex AI, Google AI Studio Node.js tests
|
||||
name: Build UI bundle
|
||||
command: |
|
||||
cd tests/pass_through_tests
|
||||
NODE_OPTIONS=--experimental-vm-modules npx jest . --verbose
|
||||
no_output_timeout: 30m
|
||||
cd ui/litellm-dashboard
|
||||
npm run build
|
||||
mkdir -p ../../ui-bundle
|
||||
mv out ../../ui-bundle/out
|
||||
find ../../ui-bundle/out -name '*.html' ! -name 'index.html' | while read -r f; do
|
||||
d="${f%.html}"; mkdir -p "$d"; mv "$f" "$d/index.html"
|
||||
done
|
||||
- persist_to_workspace:
|
||||
root: ~/project
|
||||
paths:
|
||||
- ui-bundle/out
|
||||
|
||||
e2e_ui_testing:
|
||||
docker:
|
||||
|
|
@ -798,7 +783,7 @@ jobs:
|
|||
POSTGRES_USER: e2euser
|
||||
POSTGRES_PASSWORD: e2epassword
|
||||
POSTGRES_DB: litellm_e2e
|
||||
resource_class: large
|
||||
resource_class: xlarge
|
||||
working_directory: ~/project
|
||||
environment:
|
||||
DATABASE_URL: "postgresql://e2euser:e2epassword@localhost:5432/litellm_e2e"
|
||||
|
|
@ -817,6 +802,8 @@ jobs:
|
|||
printf 'export LITELLM_MASTER_KEY=sk-%s\n' "$key" >> "$BASH_ENV"
|
||||
- skip_if_unrelated_changes:
|
||||
category: client
|
||||
- attach_workspace:
|
||||
at: ~/project
|
||||
- setup_google_dns
|
||||
- install_node
|
||||
- install_uv
|
||||
|
|
@ -835,7 +822,7 @@ jobs:
|
|||
- ~/.cache/uv
|
||||
- restore_cache:
|
||||
keys:
|
||||
- ui-e2e-node-deps-v4-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }}
|
||||
- ui-playwright-node-deps-v1-{{ checksum "tests/e2e/ui/package-lock.json" }}
|
||||
- run:
|
||||
name: Install Node dependencies and Playwright
|
||||
# The cimg/python:3.12-browsers image already ships the Chromium system
|
||||
|
|
@ -844,36 +831,19 @@ jobs:
|
|||
# 5-10 minutes to the job and frequently stalls on flaky Ubuntu mirrors,
|
||||
# so we install just the browser binary.
|
||||
command: |
|
||||
cd ui/litellm-dashboard
|
||||
npm ci
|
||||
cd ../../tests/e2e/ui
|
||||
cd tests/e2e/ui
|
||||
npm ci
|
||||
npx playwright install chromium
|
||||
- save_cache:
|
||||
key: ui-e2e-node-deps-v4-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }}
|
||||
key: ui-playwright-node-deps-v1-{{ checksum "tests/e2e/ui/package-lock.json" }}
|
||||
paths:
|
||||
- ui/litellm-dashboard/node_modules
|
||||
- tests/e2e/ui/node_modules
|
||||
- ~/.cache/ms-playwright
|
||||
- run:
|
||||
name: Build UI from source
|
||||
# Prior version used `cp -r out/ ../../litellm/proxy/_experimental/out/`.
|
||||
# GNU cp (used on CircleCI's Ubuntu image) interprets that as "copy the
|
||||
# source directory as a child of the destination" when the destination
|
||||
# already exists — silently creating `_experimental/out/out/` instead of
|
||||
# replacing the served bundle. The proxy continued serving whatever was
|
||||
# checked into `_experimental/out/*`, so this job was effectively testing
|
||||
# the pre-build bundle on every run. Replace-and-move guarantees the
|
||||
# freshly built bundle is what the proxy actually serves.
|
||||
name: Use built UI bundle
|
||||
command: |
|
||||
cd ui/litellm-dashboard
|
||||
npm run build
|
||||
rm -rf ../../litellm/proxy/_experimental/out
|
||||
mv out ../../litellm/proxy/_experimental/out
|
||||
# Restructure HTML so extensionless routes work (login.html -> login/index.html)
|
||||
find ../../litellm/proxy/_experimental/out -name '*.html' ! -name 'index.html' | while read -r f; do
|
||||
d="${f%.html}"; mkdir -p "$d"; mv "$f" "$d/index.html"
|
||||
done
|
||||
rm -rf litellm/proxy/_experimental/out
|
||||
cp -r ui-bundle/out litellm/proxy/_experimental/out
|
||||
- wait_for_service:
|
||||
url: tcp://localhost:5432
|
||||
timeout: "30"
|
||||
|
|
@ -940,7 +910,7 @@ jobs:
|
|||
command: |
|
||||
cd tests/e2e/ui
|
||||
LITELLM_LICENSE="$LITELLM_LICENSE" \
|
||||
npx playwright test --config playwright.config.ts
|
||||
npx playwright test --config playwright.config.ts --workers=3 --retries=0
|
||||
no_output_timeout: 10m
|
||||
- store_artifacts:
|
||||
path: tests/e2e/ui/test-results
|
||||
|
|
@ -978,6 +948,8 @@ jobs:
|
|||
printf 'export LITELLM_MASTER_KEY=sk-%s\n' "$key" >> "$BASH_ENV"
|
||||
- skip_if_unrelated_changes:
|
||||
category: client
|
||||
- attach_workspace:
|
||||
at: ~/project
|
||||
- setup_google_dns
|
||||
- install_node
|
||||
- install_uv
|
||||
|
|
@ -996,31 +968,23 @@ jobs:
|
|||
- ~/.cache/uv
|
||||
- restore_cache:
|
||||
keys:
|
||||
- ui-e2e-node-deps-v4-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }}
|
||||
- ui-playwright-node-deps-v1-{{ checksum "tests/e2e/ui/package-lock.json" }}
|
||||
- run:
|
||||
name: Install Node dependencies and Playwright
|
||||
command: |
|
||||
cd ui/litellm-dashboard
|
||||
npm ci
|
||||
cd ../../tests/e2e/ui
|
||||
cd tests/e2e/ui
|
||||
npm ci
|
||||
npx playwright install chromium
|
||||
- save_cache:
|
||||
key: ui-e2e-node-deps-v4-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }}
|
||||
key: ui-playwright-node-deps-v1-{{ checksum "tests/e2e/ui/package-lock.json" }}
|
||||
paths:
|
||||
- ui/litellm-dashboard/node_modules
|
||||
- tests/e2e/ui/node_modules
|
||||
- ~/.cache/ms-playwright
|
||||
- run:
|
||||
name: Build UI from source
|
||||
name: Use built UI bundle
|
||||
command: |
|
||||
cd ui/litellm-dashboard
|
||||
npm run build
|
||||
rm -rf ../../litellm/proxy/_experimental/out
|
||||
mv out ../../litellm/proxy/_experimental/out
|
||||
find ../../litellm/proxy/_experimental/out -name '*.html' ! -name 'index.html' | while read -r f; do
|
||||
d="${f%.html}"; mkdir -p "$d"; mv "$f" "$d/index.html"
|
||||
done
|
||||
rm -rf litellm/proxy/_experimental/out
|
||||
cp -r ui-bundle/out litellm/proxy/_experimental/out
|
||||
- wait_for_service:
|
||||
url: tcp://localhost:5432
|
||||
timeout: "30"
|
||||
|
|
@ -1271,45 +1235,6 @@ jobs:
|
|||
! docker exec my-app true 2>/dev/null || fail "container is still running"
|
||||
echo "Proxy refused to serve (exit $exit_code) and never reached startup. Test passed."
|
||||
|
||||
provider_replay_harness:
|
||||
docker:
|
||||
- &python312_image
|
||||
image: cimg/python:3.12@sha256:9c796c23c84e84a66a964acb508d39dc5433c81a47e07efd56dccbbc2427e07c
|
||||
auth:
|
||||
username: ${DOCKERHUB_USERNAME}
|
||||
password: ${DOCKERHUB_PASSWORD}
|
||||
- image: redis@sha256:e2debfb7956fa12c7ddc79d7e645c8cf26b30c99a6e9161ea9bf4171e1668a5f
|
||||
working_directory: ~/project
|
||||
resource_class: medium
|
||||
environment:
|
||||
E2E_CACHE_TEST_REDIS_URL: redis://127.0.0.1:6379/0
|
||||
E2E_PROVIDER_CACHE: "0"
|
||||
E2E_FIXTURE_MODE: live
|
||||
steps:
|
||||
- checkout
|
||||
- run:
|
||||
name: Generate LiteLLM master key
|
||||
command: |
|
||||
key="$(openssl rand -hex 16)"
|
||||
printf 'export LITELLM_MASTER_KEY=sk-%s\n' "$key" >> "$BASH_ENV"
|
||||
- skip_if_unrelated_changes:
|
||||
category: provider-harness
|
||||
- setup_litellm_test_deps
|
||||
- wait_for_service:
|
||||
url: tcp://localhost:6379
|
||||
- run:
|
||||
name: Test provider capture and replay harness
|
||||
command: |
|
||||
mkdir -p test-results/provider-replay-harness
|
||||
uv run --no-sync pytest --tb=short -q --noconftest -o addopts= -o "pythonpath=tests/e2e tests/e2e_harness" -p no:rerunfailures \
|
||||
--junitxml=test-results/provider-replay-harness/junit.xml \
|
||||
tests/e2e_harness/test_provider_edge.py tests/e2e_harness/test_fixture_bundle.py \
|
||||
tests/e2e_harness/test_fixture_canonical.py tests/e2e_harness/test_fixture_mode.py \
|
||||
tests/code_coverage_tests/test_provider_replay_harness.py \
|
||||
tests/code_coverage_tests/test_provider_cache.py
|
||||
- store_test_results:
|
||||
path: test-results/provider-replay-harness
|
||||
|
||||
integration_contracts:
|
||||
parameters:
|
||||
suite:
|
||||
|
|
@ -1358,9 +1283,19 @@ jobs:
|
|||
- run:
|
||||
name: Build the candidate dashboard
|
||||
command: cd ui/litellm-dashboard && NEXT_TELEMETRY_DISABLED=1 npm run build
|
||||
- when:
|
||||
condition:
|
||||
equal: [mcp, << parameters.suite >>]
|
||||
steps:
|
||||
- run:
|
||||
name: Install MCP SDK1 peer
|
||||
command: |
|
||||
uv venv --python 3.12 .venv-mcp-peer
|
||||
uv pip install --python .venv-mcp-peer 'mcp==1.28.1' 'langchain-mcp-adapters==0.2.1'
|
||||
echo "export MCP_TEST_PEER_PYTHON=$PWD/.venv-mcp-peer/bin/python" >> "$BASH_ENV"
|
||||
- start_postgres:
|
||||
image: postgres:16@sha256:e17e86066e5ef83e0952a9347f5c792b7ece00972e2aa787a6986f471b3dd3d5
|
||||
server_args: "-c shared_preload_libraries=pg_stat_statements -c pg_stat_statements.track=all -c pg_stat_statements.max=20000"
|
||||
server_args: "-c max_connections=300 -c shared_preload_libraries=pg_stat_statements -c pg_stat_statements.track=all -c pg_stat_statements.max=20000"
|
||||
- start_redis
|
||||
- run:
|
||||
name: Run owned integration contracts
|
||||
|
|
@ -1413,12 +1348,27 @@ jobs:
|
|||
default: standard
|
||||
shards:
|
||||
type: integer
|
||||
machine:
|
||||
image: ubuntu-2204:2024.04.1
|
||||
resource_class: medium
|
||||
docker:
|
||||
- image: cimg/python:3.12@sha256:9c796c23c84e84a66a964acb508d39dc5433c81a47e07efd56dccbbc2427e07c
|
||||
auth:
|
||||
username: ${DOCKERHUB_USERNAME}
|
||||
password: ${DOCKERHUB_PASSWORD}
|
||||
resource_class: small
|
||||
working_directory: ~/project
|
||||
environment:
|
||||
UV_PYTHON_DOWNLOADS: never
|
||||
steps:
|
||||
- setup_litellm_test_deps
|
||||
- checkout
|
||||
- install_uv
|
||||
- run:
|
||||
name: Install the locked shard verifier and coverage dependencies
|
||||
command: |
|
||||
uv export --frozen --all-groups --all-extras --no-emit-project --no-emit-workspace --no-header --no-annotate \
|
||||
| awk '/^[^ ]/ { keep = $0 ~ /^(annotated-types|coverage|pydantic|pydantic-core|typing-extensions|typing-inspection)==/ } keep' \
|
||||
> /tmp/fan-in-requirements.txt
|
||||
test "$(grep -c '^[^ ]' /tmp/fan-in-requirements.txt)" -eq 6
|
||||
uv venv .venv --python 3.12
|
||||
uv pip install --no-config --python .venv --require-hashes -r /tmp/fan-in-requirements.txt
|
||||
- attach_workspace:
|
||||
at: /tmp/integration-shards
|
||||
- run:
|
||||
|
|
@ -1534,7 +1484,8 @@ jobs:
|
|||
steps:
|
||||
- checkout
|
||||
- skip_if_unrelated_changes
|
||||
- setup_litellm_test_deps
|
||||
- setup_litellm_test_deps:
|
||||
checkout: false
|
||||
- start_postgres:
|
||||
db_name: litellm_test
|
||||
image: postgres:16@sha256:e17e86066e5ef83e0952a9347f5c792b7ece00972e2aa787a6986f471b3dd3d5
|
||||
|
|
@ -1579,112 +1530,15 @@ jobs:
|
|||
- store_artifacts:
|
||||
path: test-results
|
||||
|
||||
mcp_integration:
|
||||
machine:
|
||||
image: ubuntu-2204:2024.04.1
|
||||
resource_class: large
|
||||
working_directory: ~/project
|
||||
environment:
|
||||
COVERAGE_CORE: sysmon
|
||||
LITELLM_LOCAL_MODEL_COST_MAP: "True"
|
||||
steps:
|
||||
- checkout
|
||||
- skip_if_unrelated_changes
|
||||
- setup_litellm_test_deps
|
||||
- run:
|
||||
name: Generate Prisma client
|
||||
command: uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
|
||||
- run:
|
||||
name: Install MCP SDK1 peer
|
||||
command: |
|
||||
uv venv --python 3.12 .venv-mcp-peer
|
||||
uv pip install --python .venv-mcp-peer 'mcp==1.28.1' 'langchain-mcp-adapters==0.2.1'
|
||||
echo "export MCP_TEST_PEER_PYTHON=$PWD/.venv-mcp-peer/bin/python" >> "$BASH_ENV"
|
||||
- run:
|
||||
name: Run MCP integration tests
|
||||
command: |
|
||||
mkdir -p test-results
|
||||
env -u OPENAI_API_KEY -u ANTHROPIC_API_KEY \
|
||||
timeout --signal=TERM 20m uv run --no-sync pytest \
|
||||
tests/mcp_tests tests/unit/experimental_mcp_client tests/unit/proxy/_experimental/mcp_server \
|
||||
tests/unit/responses/mcp --tb=short -vv --maxfail=10 -n 2 --dist=loadscope --reruns 0 \
|
||||
--reruns-delay 1 --timeout=120 --rerun-except "from pytest-timeout" --durations=20 \
|
||||
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml:coverage.xml \
|
||||
--cov-config=pyproject.toml -o junit_family=xunit1 --junitxml=test-results/junit.xml
|
||||
- install_codecov_cli
|
||||
- run:
|
||||
name: Upload MCP integration coverage
|
||||
when: always
|
||||
command: |
|
||||
[ -f coverage.xml ] || { echo "no coverage.xml produced"; exit 1; }
|
||||
codecov upload-process --disable-search --fail-on-error -f coverage.xml \
|
||||
-F mcp-integration -C "$CIRCLE_SHA1" \
|
||||
-n "mcp-integration-${CIRCLE_BUILD_NUM}" --git-service github
|
||||
- store_test_results:
|
||||
path: test-results
|
||||
- store_artifacts:
|
||||
path: test-results
|
||||
|
||||
redis_compat:
|
||||
parameters:
|
||||
redis_py:
|
||||
type: string
|
||||
machine:
|
||||
image: ubuntu-2204:2024.04.1
|
||||
resource_class: large
|
||||
working_directory: ~/project
|
||||
steps:
|
||||
- checkout
|
||||
- skip_if_unrelated_changes:
|
||||
category: redis-compat
|
||||
- setup_litellm_test_deps
|
||||
- run:
|
||||
name: Pin redis-py version
|
||||
command: |
|
||||
uv pip install "redis==<< parameters.redis_py >>"
|
||||
uv run --no-sync python -c "import redis; assert redis.__version__ == '<< parameters.redis_py >>', redis.__version__; print('redis-py', redis.__version__)"
|
||||
- run:
|
||||
name: Install redis-server 7.2.16
|
||||
command: |
|
||||
mkdir -p "$HOME/.local/bin"
|
||||
cid="$(docker create redis:7.2.16@sha256:0637954999d01b7c9ce9167db2da50656e2590d3b884f1c600c5f63bb6e6773c)"
|
||||
docker cp "${cid}":/usr/local/bin/redis-server "$HOME/.local/bin/redis-server"
|
||||
docker rm "$cid"
|
||||
redis-server --version | grep -q 'v=7.2.16'
|
||||
- run:
|
||||
name: Run Redis compatibility tests
|
||||
command: |
|
||||
mkdir -p test-results
|
||||
env -u CASSETTE_REDIS_URL -u AZURE_CLIENT_ID -u AZURE_CLIENT_SECRET -u AZURE_TENANT_ID \
|
||||
timeout --signal=TERM 15m uv run --no-sync pytest \
|
||||
tests/unit/test_redis.py tests/unit/caching/test_redis_connection_pool.py \
|
||||
tests/unit/caching/test_redis_cluster_cache.py tests/unit/caching/test_evicted_client_closer.py \
|
||||
tests/integration/sdk/test_redis_cluster_iam_auth.py::test_sync_cluster_authenticates_with_azure_credentials \
|
||||
tests/integration/sdk/test_redis_cluster_iam_auth.py::test_sync_cluster_authenticates_with_gcp_credentials \
|
||||
--tb=short -vv --reruns 2 --reruns-delay 1 --durations=20 --cov=./litellm \
|
||||
--cov-report=xml:coverage.xml -o junit_family=xunit1 --junitxml=test-results/junit.xml
|
||||
- when:
|
||||
condition:
|
||||
equal: ["5.3.1", << parameters.redis_py >>]
|
||||
steps:
|
||||
- install_codecov_cli
|
||||
- run:
|
||||
name: Upload Redis compatibility coverage
|
||||
when: always
|
||||
command: |
|
||||
[ -f coverage.xml ] || { echo "no coverage.xml produced"; exit 1; }
|
||||
codecov upload-process --disable-search --fail-on-error -f coverage.xml \
|
||||
-F redis-compat -C "$CIRCLE_SHA1" \
|
||||
-n "redis-compat-${CIRCLE_BUILD_NUM}" --git-service github
|
||||
- store_test_results:
|
||||
path: test-results
|
||||
- store_artifacts:
|
||||
path: test-results
|
||||
|
||||
workflows:
|
||||
migration_startup:
|
||||
when: << pipeline.parameters.run_migration_tests >>
|
||||
jobs: &migration_jobs
|
||||
migration_startup_scheduled:
|
||||
triggers:
|
||||
- schedule:
|
||||
cron: "17 0,6,12,18 * * *"
|
||||
filters:
|
||||
branches:
|
||||
only: main
|
||||
jobs:
|
||||
- build_docker_database_image:
|
||||
migration_qualification: true
|
||||
- migration_startup_tests:
|
||||
|
|
@ -1709,33 +1563,11 @@ workflows:
|
|||
suite: shaped
|
||||
baseline: true
|
||||
requires: [build_docker_database_image]
|
||||
migration_startup_scheduled:
|
||||
triggers:
|
||||
- schedule:
|
||||
cron: "17 0,6,12,18 * * *"
|
||||
filters:
|
||||
branches:
|
||||
only: main
|
||||
jobs: *migration_jobs
|
||||
routing_parity:
|
||||
when:
|
||||
not:
|
||||
equal: ["", << pipeline.parameters.routing_parity_base >>]
|
||||
jobs:
|
||||
- routing_parity:
|
||||
name: routing-parity-<< matrix.suite >>
|
||||
matrix:
|
||||
parameters:
|
||||
suite: [management, accounting, database, providers, extensions, cost, mcp]
|
||||
integration:
|
||||
unless:
|
||||
or:
|
||||
- << pipeline.parameters.run_migration_tests >>
|
||||
- not:
|
||||
equal: ["", << pipeline.parameters.routing_parity_base >>]
|
||||
jobs:
|
||||
- integration_contracts:
|
||||
name: integration-<< matrix.suite >>
|
||||
filters: &integration_default_run not pipeline.parameters.run_migration_tests and pipeline.parameters.routing_parity_base == ""
|
||||
matrix:
|
||||
parameters:
|
||||
suite: [sdk, cost, browser]
|
||||
|
|
@ -1745,150 +1577,178 @@ workflows:
|
|||
suite: security
|
||||
parallelism: 1
|
||||
resource_class: xlarge
|
||||
filters: *integration_default_run
|
||||
- integration_contracts:
|
||||
name: integration-management
|
||||
suite: management
|
||||
parallelism: 4
|
||||
resource_class: xlarge
|
||||
filters: *integration_default_run
|
||||
- integration_shard_coverage:
|
||||
name: integration-management-coverage
|
||||
suite: management
|
||||
shards: 4
|
||||
requires: [integration-management]
|
||||
filters: *integration_default_run
|
||||
- integration_contracts:
|
||||
name: integration-accounting
|
||||
suite: accounting
|
||||
parallelism: 4
|
||||
resource_class: xlarge
|
||||
filters: *integration_default_run
|
||||
- integration_shard_coverage:
|
||||
name: integration-accounting-coverage
|
||||
suite: accounting
|
||||
shards: 4
|
||||
requires: [integration-accounting]
|
||||
filters: *integration_default_run
|
||||
- integration_contracts:
|
||||
name: integration-providers
|
||||
suite: providers
|
||||
parallelism: 4
|
||||
resource_class: xlarge
|
||||
filters: *integration_default_run
|
||||
- integration_shard_coverage:
|
||||
name: integration-providers-coverage
|
||||
suite: providers
|
||||
shards: 4
|
||||
requires:
|
||||
- integration-providers
|
||||
filters: *integration_default_run
|
||||
- integration_contracts:
|
||||
name: integration-database
|
||||
suite: database
|
||||
parallelism: 2
|
||||
filters: *integration_default_run
|
||||
- integration_shard_coverage:
|
||||
name: integration-database-coverage
|
||||
suite: database
|
||||
shards: 2
|
||||
requires: [integration-database]
|
||||
filters: *integration_default_run
|
||||
- integration_contracts:
|
||||
name: integration-mcp
|
||||
suite: mcp
|
||||
resource_class: xlarge
|
||||
parallelism: 2
|
||||
filters: *integration_default_run
|
||||
- integration_shard_coverage:
|
||||
name: integration-mcp-coverage
|
||||
suite: mcp
|
||||
shards: 2
|
||||
requires:
|
||||
- integration-mcp
|
||||
filters: *integration_default_run
|
||||
- integration_contracts:
|
||||
name: integration-extensions
|
||||
suite: extensions
|
||||
parallelism: 8
|
||||
resource_class: xlarge
|
||||
filters: *integration_default_run
|
||||
- integration_shard_coverage:
|
||||
name: integration-extensions-coverage
|
||||
suite: extensions
|
||||
shards: 8
|
||||
requires:
|
||||
- integration-extensions
|
||||
filters: *integration_default_run
|
||||
- integration_contracts:
|
||||
name: integration-management-replica
|
||||
suite: management
|
||||
mode: replica
|
||||
parallelism: 4
|
||||
resource_class: xlarge
|
||||
filters: *integration_default_run
|
||||
- integration_shard_coverage:
|
||||
name: integration-management-replica-coverage
|
||||
suite: management
|
||||
mode: replica
|
||||
shards: 4
|
||||
requires: [integration-management-replica]
|
||||
filters: *integration_default_run
|
||||
- integration_contracts:
|
||||
name: integration-database-replica
|
||||
suite: database
|
||||
mode: replica
|
||||
parallelism: 2
|
||||
filters: *integration_default_run
|
||||
- integration_shard_coverage:
|
||||
name: integration-database-replica-coverage
|
||||
suite: database
|
||||
mode: replica
|
||||
shards: 2
|
||||
requires: [integration-database-replica]
|
||||
- postgres_suite:
|
||||
name: proxy-behavior
|
||||
test_path: tests/proxy_behavior
|
||||
seed: true
|
||||
coverage_flag: lens-postgres
|
||||
timeout_minutes: 25
|
||||
- postgres_suite:
|
||||
name: proxy-security
|
||||
test_path: tests/proxy_security_tests
|
||||
seed: true
|
||||
timeout_minutes: 15
|
||||
- postgres_suite:
|
||||
name: schema-migration
|
||||
test_path: tests/proxy_migration_tests
|
||||
seed: false
|
||||
timeout_minutes: 20
|
||||
workers: 2
|
||||
filters: *integration_default_run
|
||||
- postgres_suite:
|
||||
name: roi-database
|
||||
test_path: tests/integration/database/test_roi_observed.py
|
||||
seed: false
|
||||
coverage_flag: roi-postgres
|
||||
timeout_minutes: 10
|
||||
- mcp_integration:
|
||||
name: mcp-integration
|
||||
- redis_compat:
|
||||
name: redis-compat-<< matrix.redis_py >>
|
||||
matrix:
|
||||
parameters:
|
||||
redis_py:
|
||||
- "5.3.1"
|
||||
- "6.4.0"
|
||||
- "7.4.1"
|
||||
- "8.0.1"
|
||||
- windows_release_wheel
|
||||
- base_sdk_install
|
||||
build_and_test:
|
||||
unless:
|
||||
or:
|
||||
- << pipeline.parameters.run_migration_tests >>
|
||||
- not:
|
||||
equal: ["", << pipeline.parameters.routing_parity_base >>]
|
||||
jobs:
|
||||
- provider_replay_harness
|
||||
- build_docker_database_image
|
||||
- e2e_ui_testing
|
||||
- e2e_ui_testing_server_root_path
|
||||
filters: *integration_default_run
|
||||
- windows_release_wheel:
|
||||
filters: *integration_default_run
|
||||
- base_sdk_install:
|
||||
filters: *integration_default_run
|
||||
- build_docker_database_image:
|
||||
filters: *integration_default_run
|
||||
- proxy_logging_guardrails_model_info_tests:
|
||||
requires:
|
||||
- build_docker_database_image
|
||||
- proxy_pass_through_endpoint_tests:
|
||||
requires:
|
||||
- build_docker_database_image
|
||||
filters: *integration_default_run
|
||||
- db_migration_disable_update_check:
|
||||
requires:
|
||||
- build_docker_database_image
|
||||
- helm_chart_testing:
|
||||
requires:
|
||||
- build_docker_database_image
|
||||
filters: *integration_default_run
|
||||
- test_bad_database_url:
|
||||
requires:
|
||||
- build_docker_database_image
|
||||
filters: *integration_default_run
|
||||
- helm_chart_testing:
|
||||
requires:
|
||||
- build_docker_database_image
|
||||
filters: *integration_default_run
|
||||
- e2e_ui_build:
|
||||
filters: *integration_default_run
|
||||
- e2e_ui_testing:
|
||||
requires: [e2e_ui_build]
|
||||
filters: *integration_default_run
|
||||
- e2e_ui_testing_server_root_path:
|
||||
requires: [e2e_ui_build]
|
||||
filters: *integration_default_run
|
||||
- build_docker_database_image:
|
||||
name: build_docker_database_image-migration
|
||||
migration_qualification: true
|
||||
filters: pipeline.parameters.run_migration_tests
|
||||
- migration_startup_tests:
|
||||
name: migration-startup
|
||||
suite: startup
|
||||
requires: [build_docker_database_image-migration]
|
||||
filters: pipeline.parameters.run_migration_tests
|
||||
- migration_startup_tests:
|
||||
name: migration-recovery
|
||||
suite: recovery
|
||||
requires: [build_docker_database_image-migration]
|
||||
filters: pipeline.parameters.run_migration_tests
|
||||
- migration_startup_tests:
|
||||
name: migration-legacy-and-pooling
|
||||
suite: legacy
|
||||
requires: [build_docker_database_image-migration]
|
||||
filters: pipeline.parameters.run_migration_tests
|
||||
- migration_startup_tests:
|
||||
name: migration-upgrade
|
||||
suite: upgrade
|
||||
baseline: true
|
||||
requires: [build_docker_database_image-migration]
|
||||
filters: pipeline.parameters.run_migration_tests
|
||||
- migration_startup_tests:
|
||||
name: migration-upgrade-shaped
|
||||
suite: shaped
|
||||
baseline: true
|
||||
requires: [build_docker_database_image-migration]
|
||||
filters: pipeline.parameters.run_migration_tests
|
||||
- routing_parity:
|
||||
name: routing-parity-<< matrix.suite >>
|
||||
matrix:
|
||||
parameters:
|
||||
suite: [management, accounting, database, providers, extensions, cost, mcp]
|
||||
filters: pipeline.parameters.routing_parity_base != ""
|
||||
|
|
|
|||
|
|
@ -1,12 +1,11 @@
|
|||
#!/usr/bin/env bash
|
||||
set -uo pipefail
|
||||
|
||||
category="${1:?usage: classify_changes.sh <backend|client|ui|provider-harness|cost-map-only|mcp-dependencies|windows-release|redis-compat>}"
|
||||
category="${1:?usage: classify_changes.sh <backend|client|ui|cost-map-only|mcp-dependencies|windows-release|redis-compat>}"
|
||||
|
||||
has_client=false
|
||||
has_backend=false
|
||||
has_ci=false
|
||||
has_provider_harness=false
|
||||
has_cost_map=false
|
||||
has_mcp_dependencies=false
|
||||
has_windows_release=false
|
||||
|
|
@ -15,7 +14,7 @@ outside_cost_map_set=false
|
|||
while IFS= read -r file || [ -n "$file" ]; do
|
||||
[ -n "$file" ] || continue
|
||||
case "$file" in
|
||||
litellm/_redis.py | litellm/_redis_credential_provider.py | litellm/caching/redis_cache.py | litellm/caching/evicted_client_closer.py | tests/unit/test_redis.py | tests/integration/sdk/test_redis_cluster_iam_auth.py | tests/unit/caching/test_redis_connection_pool.py | tests/unit/caching/test_redis_cluster_cache.py | tests/unit/caching/test_evicted_client_closer.py | .circleci/config.yml | .circleci/scripts/classify_changes.sh | .circleci/scripts/path_filter.sh | pyproject.toml | uv.lock)
|
||||
litellm/_redis.py | litellm/_redis_credential_provider.py | litellm/caching/redis_cache.py | litellm/caching/evicted_client_closer.py | tests/unit/test_redis.py | tests/integration/sdk/test_redis_cluster_iam_auth.py | tests/unit/caching/test_redis_connection_pool.py | tests/unit/caching/test_redis_cluster_cache.py | tests/unit/caching/test_evicted_client_closer.py | .github/workflows/test-unit.yml | .github/actions/detect-changes/* | .github/scripts/detect_changes.sh | .circleci/scripts/classify_changes.sh | tests/integration/run.py | pyproject.toml | uv.lock)
|
||||
has_redis_compat=true ;;
|
||||
esac
|
||||
case "$file" in
|
||||
|
|
@ -23,11 +22,6 @@ while IFS= read -r file || [ -n "$file" ]; do
|
|||
pyproject.toml | */pyproject.toml | uv.lock | uv.toml | .python-version | rust-toolchain.toml | litellm-rust/* | litellm/__init__.py | litellm/proxy/proxy_server.py | litellm/*mcp* | tests/*mcp* | litellm/integrations/arize/* | tests/base_sdk_tests/* | scripts/check_mcp_sdk_install.py | .github/workflows/test-mcp-dependency-resolution.yml | .github/actions/detect-changes/* | .github/actions/setup-uv-with-retries/* | .github/actions/cache-cargo-build/* | .github/scripts/detect_changes.sh | .github/scripts/uv_sync_with_retries.sh | .circleci/scripts/classify_changes.sh | tests/unit/test_circleci_path_filter.py | tests/unit/test_detect_changes.py)
|
||||
has_mcp_dependencies=true ;;
|
||||
esac
|
||||
case "$file" in
|
||||
tests/e2e/*/*.py | tests/e2e_harness/*/*.py) : ;;
|
||||
tests/e2e/*.py | tests/e2e_harness/*.py | tests/code_coverage_tests/test_provider_cache.py | tests/code_coverage_tests/test_provider_replay_harness.py | tests/unit/test_circleci_path_filter.py | .circleci/* | pyproject.toml | uv.lock)
|
||||
has_provider_harness=true ;;
|
||||
esac
|
||||
case "$file" in
|
||||
litellm-rust/* | litellm/rust_bridge/* | rust-toolchain.toml | pyproject.toml | uv.lock | tests/windows_tests/* | .circleci/*)
|
||||
has_windows_release=true ;;
|
||||
|
|
@ -53,9 +47,6 @@ case "$category" in
|
|||
cost-map-only)
|
||||
{ [ "$has_cost_map" = true ] && [ "$outside_cost_map_set" = false ]; } && echo run || echo skip
|
||||
;;
|
||||
provider-harness)
|
||||
[ "$has_provider_harness" = true ] && echo run || echo skip
|
||||
;;
|
||||
windows-release)
|
||||
[ "$has_windows_release" = true ] && echo run || echo skip
|
||||
;;
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
#!/usr/bin/env bash
|
||||
set -uo pipefail
|
||||
|
||||
category="${1:?usage: path_filter.sh <backend|client|provider-harness|redis-compat>}"
|
||||
category="${1:?usage: path_filter.sh <backend|client|windows-release>}"
|
||||
here="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
|
||||
run_full() {
|
||||
|
|
|
|||
|
|
@ -128,20 +128,14 @@ 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"
|
||||
|
||||
isolated_workers=false
|
||||
if [ "$mode" = parity ]; then
|
||||
export INTEGRATION_WORKERS=1
|
||||
else
|
||||
case "$suite" in
|
||||
management|accounting) export INTEGRATION_WORKERS=2; isolated_workers=true ;;
|
||||
providers|extensions) export INTEGRATION_WORKERS=4; isolated_workers=true ;;
|
||||
management|accounting) export INTEGRATION_WORKERS=2 ;;
|
||||
providers|extensions) export INTEGRATION_WORKERS=4 ;;
|
||||
esac
|
||||
fi
|
||||
if [ "$isolated_workers" = false ]; then
|
||||
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=$!
|
||||
fi
|
||||
if [ "$suite" = cost ]; then
|
||||
export INTEGRATION_WORKERS=8
|
||||
fi
|
||||
|
|
@ -149,7 +143,16 @@ if [ "$suite" = security ]; then
|
|||
export INTEGRATION_WORKERS=4
|
||||
fi
|
||||
if [ "$suite" = mcp ]; then
|
||||
export INTEGRATION_WORKERS=4 INTEGRATION_COVERAGE=1
|
||||
export INTEGRATION_WORKERS=2 INTEGRATION_COVERAGE=1
|
||||
fi
|
||||
isolated_workers=false
|
||||
if [ "${INTEGRATION_WORKERS:-1}" -gt 1 ]; then
|
||||
isolated_workers=true
|
||||
fi
|
||||
if [ "$isolated_workers" = false ]; then
|
||||
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=$!
|
||||
fi
|
||||
coverage_data="$PWD/$results/coverage/data"
|
||||
proxy_command=(.venv/bin/python -m integration._support.proxy)
|
||||
|
|
@ -191,9 +194,6 @@ start_proxy() {
|
|||
> "$results/$log_name" 2>&1 &
|
||||
launched_pid=$!
|
||||
}
|
||||
if { [ "$suite" = management ] || [ "$suite" = mcp ]; } && [ "$isolated_workers" = true ]; then
|
||||
export INTEGRATION_PEER_URL=http://127.0.0.1:4001
|
||||
fi
|
||||
if [ "$isolated_workers" = false ]; then
|
||||
start_proxy 4000 proxy.log
|
||||
proxy_pid="$launched_pid"
|
||||
|
|
@ -201,7 +201,7 @@ if [ "$isolated_workers" = false ]; then
|
|||
curl --noproxy '*' -sSf -X POST "$INTEGRATION_PROXY_URL/config/update" \
|
||||
-H "Authorization: Bearer $LITELLM_MASTER_KEY" -H 'Content-Type: application/json' \
|
||||
-d '{"router_settings": {"num_retries": 0}}' > "$results/seed-router-settings.json"
|
||||
if [ "$suite" = management ] || [ "$suite" = mcp ]; then
|
||||
if [ "$suite" = management ] || [ "$suite" = mcp ] || [ "$suite" = accounting ]; then
|
||||
export INTEGRATION_PEER_URL=http://127.0.0.1:4001
|
||||
start_proxy 4001 peer.log
|
||||
peer_pid="$launched_pid"
|
||||
|
|
@ -239,13 +239,24 @@ if [ "${CIRCLE_NODE_TOTAL:-1}" -gt 1 ]; then
|
|||
shard_arguments=(--shard-files "$results/node-files.txt")
|
||||
fi
|
||||
|
||||
shared_services=()
|
||||
worker_coverage=()
|
||||
if [ "$isolated_workers" = false ]; then
|
||||
shared_services=(
|
||||
"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"
|
||||
)
|
||||
elif [ "${INTEGRATION_COVERAGE:-0}" = 1 ]; then
|
||||
worker_coverage=("INTEGRATION_COVERAGE_FILE=$coverage_data")
|
||||
fi
|
||||
|
||||
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" \
|
||||
DATABASE_URL="$DATABASE_URL" \
|
||||
"${shared_services[@]}" "${worker_coverage[@]}" \
|
||||
INTEGRATION_WORKERS="${INTEGRATION_WORKERS:-1}" \
|
||||
INTEGRATION_MASTER_KEY="$INTEGRATION_MASTER_KEY" LITELLM_MODE=PRODUCTION \
|
||||
LITELLM_MODE=PRODUCTION \
|
||||
LITELLM_LICENSE="${LITELLM_LICENSE:-}" \
|
||||
INTEGRATION_SEED="$INTEGRATION_SEED" \
|
||||
INTEGRATION_ORDER_SEED="$INTEGRATION_ORDER_SEED" \
|
||||
|
|
@ -254,6 +265,7 @@ env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" \
|
|||
INTEGRATION_PROXY_DATABASE_URL="$INTEGRATION_PROXY_DATABASE_URL" \
|
||||
INTEGRATION_PROXY_READ_REPLICA_URL="$INTEGRATION_PROXY_READ_REPLICA_URL" \
|
||||
INTEGRATION_ROUTING="$INTEGRATION_ROUTING" \
|
||||
MCP_TEST_PEER_PYTHON="${MCP_TEST_PEER_PYTHON:-}" \
|
||||
.venv/bin/python tests/integration/run.py "$suite" --results "$results" \
|
||||
"${shard_arguments[@]}"
|
||||
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@ from typing import Final
|
|||
from xml.etree import ElementTree
|
||||
|
||||
SUITES: Final = {
|
||||
"startup": (("test_startup.py",), 12),
|
||||
"startup": (("test_startup.py",), 13),
|
||||
"recovery": (("test_recovery.py",), 15),
|
||||
"legacy": (("test_legacy.py", "test_pooling.py"), 11),
|
||||
"upgrade": (("test_upgrade.py", "test_rolling_upgrade.py"), 5),
|
||||
|
|
|
|||
4
.github/CODEOWNERS
vendored
4
.github/CODEOWNERS
vendored
|
|
@ -1,2 +1,2 @@
|
|||
/model_prices_and_context_window.json @mateo-berri @ryan-crabbe-berri @kerrylu-berri
|
||||
/litellm/model_prices_and_context_window_backup.json @mateo-berri @ryan-crabbe-berri @kerrylu-berri
|
||||
/model_prices_and_context_window.json @mateo-berri @ryan-crabbe-berri @kerry-berri
|
||||
/litellm/model_prices_and_context_window_backup.json @mateo-berri @ryan-crabbe-berri @kerry-berri
|
||||
|
|
|
|||
2
.github/PULL_REQUEST_TEMPLATE/general.md
vendored
2
.github/PULL_REQUEST_TEMPLATE/general.md
vendored
|
|
@ -68,7 +68,7 @@ After: the same request comes back with real token counts, so the dashboard show
|
|||
- [ ] The handful of test files covering my change pass locally, e.g. `uv run pytest tests/unit/<your_test_file>.py -v`. Leave the suites (`make test-unit-*`, `make test-unit`) to CI: it finishes in ~15 minutes where a laptop takes an hour or more
|
||||
- [ ] My PR passes all required CI/CD checks (e.g., lint, schema.d.ts sync check, etc.)
|
||||
- [ ] My PR's scope is as isolated as possible; it only solves 1 specific problem
|
||||
- [ ] I have received a Greptile **Confidence Score of at least 4/5** before requesting a maintainer review (Greptile reviews automatically once the PR is opened; only comment `@greptileai` to re-request a review after pushing changes)
|
||||
- [ ] Every automated review on my PR's latest commit has finished, and I have addressed or replied to each of its findings before requesting a maintainer review
|
||||
|
||||
## Delays in PR merge?
|
||||
|
||||
|
|
|
|||
30
.github/workflows/deploy-dev-on-main.yml
vendored
Normal file
30
.github/workflows/deploy-dev-on-main.yml
vendored
Normal file
|
|
@ -0,0 +1,30 @@
|
|||
name: Deploy main to gateway-dev
|
||||
|
||||
on:
|
||||
push:
|
||||
branches: [main]
|
||||
|
||||
permissions: {}
|
||||
|
||||
concurrency:
|
||||
group: deploy-dev-dispatch
|
||||
cancel-in-progress: true
|
||||
|
||||
jobs:
|
||||
dispatch:
|
||||
if: github.repository == 'BerriAI/litellm'
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 5
|
||||
steps:
|
||||
- uses: actions/create-github-app-token@fee1f7d63c2ff003460e3d139729b119787bc349 # v2.2.2
|
||||
id: token
|
||||
with:
|
||||
app-id: ${{ secrets.LITELLM_OPS_DISPATCH_APP_ID }}
|
||||
private-key: ${{ secrets.LITELLM_OPS_DISPATCH_APP_PRIVATE_KEY }}
|
||||
owner: BerriAI
|
||||
repositories: litellm-ops
|
||||
permission-actions: write
|
||||
- name: Dispatch the litellm-ops dev deploy
|
||||
env:
|
||||
GH_TOKEN: ${{ steps.token.outputs.token }}
|
||||
run: gh api --method POST repos/BerriAI/litellm-ops/actions/workflows/deploy-dev.yml/dispatches -f ref=main
|
||||
14
.github/workflows/image-scan.yml
vendored
14
.github/workflows/image-scan.yml
vendored
|
|
@ -27,7 +27,7 @@ on:
|
|||
- backend/routes/allowlist.py
|
||||
- pyproject.toml
|
||||
- litellm-proxy-extras/**
|
||||
- tests/proxy_migration_tests/**
|
||||
- tests/integration/database/image/**
|
||||
- uv.lock
|
||||
- ui/litellm-dashboard/package-lock.json
|
||||
- ui/Dockerfile
|
||||
|
|
@ -92,7 +92,7 @@ jobs:
|
|||
LITELLM_ADMIN_MCP_SCHEMA_IMAGE: litellm-admin-mcp-schema:${{ github.sha }}
|
||||
run: |
|
||||
python -m pip install "pytest==9.0.3"
|
||||
python -m pytest tests/proxy_migration_tests/test_offline_image_migration.py tests/proxy_migration_tests/test_image_bedrock_realtime_extra.py tests/proxy_migration_tests/test_image_admin_mcp.py -v
|
||||
python -m pytest --noconftest tests/integration/database/image/test_offline_image_migration.py tests/integration/database/image/test_image_bedrock_realtime_extra.py tests/integration/database/image/test_image_admin_mcp.py -v
|
||||
|
||||
# Scans the whole shipped artifact: OS/apk plus every language package
|
||||
# baked into the image, including ones no lockfile declares (e.g. prisma's
|
||||
|
|
@ -145,7 +145,7 @@ jobs:
|
|||
LITELLM_ADMIN_MCP_SCHEMA_IMAGE: litellm-admin-mcp-schema:${{ github.sha }}
|
||||
run: |
|
||||
python -m pip install "pytest==9.0.3"
|
||||
python -m pytest tests/proxy_migration_tests/test_offline_image_migration.py tests/proxy_migration_tests/test_image_bedrock_realtime_extra.py tests/proxy_migration_tests/test_image_admin_mcp.py -v
|
||||
python -m pytest --noconftest tests/integration/database/image/test_offline_image_migration.py tests/integration/database/image/test_image_bedrock_realtime_extra.py tests/integration/database/image/test_image_admin_mcp.py -v
|
||||
|
||||
|
||||
migrations-image:
|
||||
|
|
@ -177,7 +177,7 @@ jobs:
|
|||
LITELLM_MIGRATION_SCRIPT: /app/run.py
|
||||
run: |
|
||||
python -m pip install "pytest==9.0.3"
|
||||
python -m pytest tests/proxy_migration_tests/test_offline_image_migration.py -v
|
||||
python -m pytest --noconftest tests/integration/database/image/test_offline_image_migration.py -v
|
||||
|
||||
gateway-image:
|
||||
name: gateway-image
|
||||
|
|
@ -208,7 +208,7 @@ jobs:
|
|||
LITELLM_IMAGE_COMPONENT: gateway
|
||||
run: |
|
||||
python -m pip install "pytest==9.0.3"
|
||||
python -m pytest tests/proxy_migration_tests/test_component_image_serves_offline.py tests/proxy_migration_tests/test_image_bedrock_realtime_extra.py tests/proxy_migration_tests/test_image_admin_mcp.py -v
|
||||
python -m pytest --noconftest tests/integration/database/image/test_component_image_serves_offline.py tests/integration/database/image/test_image_bedrock_realtime_extra.py tests/integration/database/image/test_image_admin_mcp.py -v
|
||||
|
||||
ui-image:
|
||||
name: ui-image
|
||||
|
|
@ -237,7 +237,7 @@ jobs:
|
|||
LITELLM_IMAGE: litellm-ui-scan:${{ github.sha }}
|
||||
run: |
|
||||
python -m pip install "pytest==9.0.3"
|
||||
python -m pytest tests/proxy_migration_tests/test_ui_image_serves_offline.py -v
|
||||
python -m pytest --noconftest tests/integration/database/image/test_ui_image_serves_offline.py -v
|
||||
|
||||
backend-image:
|
||||
name: backend-image
|
||||
|
|
@ -272,4 +272,4 @@ jobs:
|
|||
LITELLM_IMAGE_COMPONENT: backend
|
||||
run: |
|
||||
python -m pip install "pytest==9.0.3"
|
||||
python -m pytest tests/proxy_migration_tests/test_component_image_serves_offline.py tests/proxy_migration_tests/test_image_admin_mcp.py -v
|
||||
python -m pytest --noconftest tests/integration/database/image/test_component_image_serves_offline.py tests/integration/database/image/test_image_admin_mcp.py -v
|
||||
|
|
|
|||
4
.github/workflows/test-e2e-redis-chaos.yml
vendored
4
.github/workflows/test-e2e-redis-chaos.yml
vendored
|
|
@ -91,12 +91,12 @@ jobs:
|
|||
|
||||
- name: Run the Redis chaos load test
|
||||
env:
|
||||
E2E_REDIS_CHAOS: "1"
|
||||
E2E_PERF_TESTS: "1"
|
||||
LITELLM_PROXY_URL: http://localhost:4000
|
||||
REDIS_HOST: 127.0.0.1
|
||||
REDIS_PORT: "6379"
|
||||
run: |
|
||||
uv run --no-sync pytest tests/e2e/load/test_redis_chaos_e2e.py -v --tb=short -rA -s
|
||||
uv run --no-sync pytest tests/e2e/perf/test_redis_chaos_e2e.py -v --tb=short -rA -s
|
||||
|
||||
- name: Show proxy log on failure
|
||||
if: failure()
|
||||
|
|
|
|||
26
.github/workflows/test-linting.yml
vendored
26
.github/workflows/test-linting.yml
vendored
|
|
@ -22,6 +22,16 @@ jobs:
|
|||
actions: read
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 15
|
||||
services:
|
||||
redis:
|
||||
image: redis@sha256:e2debfb7956fa12c7ddc79d7e645c8cf26b30c99a6e9161ea9bf4171e1668a5f
|
||||
ports:
|
||||
- 6379:6379
|
||||
options: >-
|
||||
--health-cmd "redis-cli ping"
|
||||
--health-interval 10s
|
||||
--health-timeout 5s
|
||||
--health-retries 5
|
||||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
|
|
@ -189,8 +199,9 @@ jobs:
|
|||
if: steps.changes.outputs.decision != 'skip'
|
||||
env:
|
||||
LITELLM_MASTER_KEY: sk-e2e-harness-tests-reach-no-proxy
|
||||
E2E_CACHE_TEST_REDIS_URL: redis://127.0.0.1:6379/0
|
||||
run: |
|
||||
if ! git diff --name-only --diff-filter=ACMRD "$GATE_BASE_SHA" HEAD -- tests/e2e tests/e2e_harness ':(exclude)tests/e2e/ui' pyproject.toml uv.lock .github/workflows/test-linting.yml | grep -q .; then
|
||||
if ! git diff --name-only --diff-filter=ACMRD "$GATE_BASE_SHA" HEAD -- tests/e2e tests/e2e_harness ':(exclude)tests/e2e/ui' tests/code_coverage_tests/test_provider_replay_harness.py tests/code_coverage_tests/test_provider_cache.py pyproject.toml uv.lock .github/workflows/test-linting.yml | grep -q .; then
|
||||
echo "No changed e2e harness files; skipping."
|
||||
exit 0
|
||||
fi
|
||||
|
|
@ -198,6 +209,9 @@ jobs:
|
|||
CLAUDE_VERSION="$(retry uv run --no-sync python tests/e2e/claude_code/pr_gate_version_resolver.py)"
|
||||
tests/e2e/claude_code/cron_vm/install_claude_code.sh "$CLAUDE_VERSION" "$RUNNER_TEMP/claude-cli"
|
||||
PATH="$RUNNER_TEMP/claude-cli:$PATH" uv run --no-sync pytest -q tests/e2e_harness
|
||||
uv run --no-sync pytest --tb=short -q --noconftest -o addopts= -o "pythonpath=tests/e2e tests/e2e_harness" -p no:rerunfailures \
|
||||
tests/code_coverage_tests/test_provider_replay_harness.py \
|
||||
tests/code_coverage_tests/test_provider_cache.py
|
||||
|
||||
- name: Check for circular imports
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
|
|
@ -586,6 +600,9 @@ jobs:
|
|||
- name: check_unsafe_enterprise_import
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/check_unsafe_enterprise_import.py
|
||||
|
||||
- name: check_telemetry_imports
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/check_telemetry_imports.py
|
||||
|
||||
- name: ban_copy_deepcopy_kwargs
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/ban_copy_deepcopy_kwargs.py
|
||||
|
||||
|
|
@ -601,11 +618,8 @@ jobs:
|
|||
- name: check_migrations_no_data_rewrites
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/check_migrations_no_data_rewrites.py
|
||||
|
||||
- name: check_greptile_files_json
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/check_greptile_files_json.py
|
||||
|
||||
- name: test_check_greptile_files_json
|
||||
run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_check_greptile_files_json.py
|
||||
- name: check_migrations_lock_timeout
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/check_migrations_lock_timeout.py
|
||||
|
||||
- name: check_no_publicly_known_master_key
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/check_no_publicly_known_master_key.py
|
||||
|
|
|
|||
25
.github/workflows/test-merge-smoke.yml
vendored
25
.github/workflows/test-merge-smoke.yml
vendored
|
|
@ -234,7 +234,30 @@ jobs:
|
|||
exit 1
|
||||
|
||||
- name: Run the startup tests
|
||||
run: .venv/bin/python -m pytest --tb=short -vv tests/basic_proxy_startup_tests
|
||||
shell: bash
|
||||
run: |
|
||||
set -euo pipefail
|
||||
base_url=http://localhost:4000
|
||||
readiness_response="$(curl -sS --fail-with-body --write-out $'\n%{http_code}' "${base_url}/health/readiness")"
|
||||
readiness_status="${readiness_response##*$'\n'}"
|
||||
readiness_body="${readiness_response%$'\n'*}"
|
||||
[[ "${readiness_status}" == "200" ]]
|
||||
jq -e '.status == "healthy" or .status == "connected"' <<< "${readiness_body}" > /dev/null
|
||||
liveness_status="$(curl -sS --fail-with-body --output /dev/null --write-out '%{http_code}' "${base_url}/health/liveness")"
|
||||
[[ "${liveness_status}" == "200" ]]
|
||||
completion_response="$(curl -sS --fail-with-body --write-out $'\n%{http_code}' \
|
||||
--header "Authorization: Bearer ${LITELLM_MASTER_KEY}" \
|
||||
--header 'Content-Type: application/json' \
|
||||
--data '{"model":"gpt-4","messages":[{"role":"system","content":"You are a helpful assistant."},{"role":"user","content":"Hello!"}]}' \
|
||||
"${base_url}/chat/completions")"
|
||||
completion_status="${completion_response##*$'\n'}"
|
||||
completion_body="${completion_response%$'\n'*}"
|
||||
[[ "${completion_status}" == "200" ]]
|
||||
jq -e '
|
||||
(.choices | length) == 1 and
|
||||
.choices[0].message.role == "assistant" and
|
||||
.choices[0].message.content == "Hello! This is a mock response from the fake OpenAI endpoint."
|
||||
' <<< "${completion_body}" > /dev/null
|
||||
|
||||
- name: Show the proxy logs
|
||||
if: always()
|
||||
|
|
|
|||
155
.github/workflows/test-unit.yml
vendored
155
.github/workflows/test-unit.yml
vendored
|
|
@ -233,6 +233,7 @@ jobs:
|
|||
tests/unit/llms/base_llm
|
||||
tests/unit/llms/baseten
|
||||
tests/unit/llms/black_forest_labs
|
||||
tests/unit/llms/brave
|
||||
tests/unit/llms/bytez
|
||||
tests/unit/llms/cerebras
|
||||
tests/unit/llms/chat
|
||||
|
|
@ -268,6 +269,7 @@ jobs:
|
|||
tests/unit/llms/gdc
|
||||
tests/unit/llms/gemini
|
||||
tests/unit/llms/gigachat
|
||||
tests/unit/llms/google_pse
|
||||
tests/unit/llms/github_copilot
|
||||
tests/unit/llms/gradient_ai
|
||||
tests/unit/llms/groq
|
||||
|
|
@ -300,6 +302,7 @@ jobs:
|
|||
tests/unit/llms/nebius
|
||||
tests/unit/llms/neosantara
|
||||
tests/unit/llms/nimble
|
||||
tests/unit/llms/nlp_cloud
|
||||
tests/unit/llms/novita
|
||||
tests/unit/llms/nscale
|
||||
tests/unit/llms/nvidia_nim
|
||||
|
|
@ -310,6 +313,7 @@ jobs:
|
|||
tests/unit/llms/oobabooga
|
||||
tests/unit/llms/openai_like
|
||||
tests/unit/llms/opencode
|
||||
tests/unit/llms/opensandbox
|
||||
tests/unit/llms/openrouter
|
||||
tests/unit/llms/ovhcloud
|
||||
tests/unit/llms/parallel_ai
|
||||
|
|
@ -342,10 +346,12 @@ jobs:
|
|||
tests/unit/llms/tinyfish
|
||||
tests/unit/llms/together_ai
|
||||
tests/unit/llms/tool_loop
|
||||
tests/unit/llms/topaz
|
||||
tests/unit/llms/triton
|
||||
tests/unit/llms/v0
|
||||
tests/unit/llms/valkey
|
||||
tests/unit/llms/vercel_ai_gateway
|
||||
tests/unit/llms/vllm
|
||||
tests/unit/llms/volcengine
|
||||
tests/unit/llms/voyage
|
||||
tests/unit/llms/wandb
|
||||
|
|
@ -412,6 +418,7 @@ jobs:
|
|||
tests/unit/rag
|
||||
tests/unit/realtime_api
|
||||
tests/unit/rerank_api
|
||||
tests/unit/search
|
||||
tests/unit/vector_stores
|
||||
tests/unit/videos
|
||||
python-version: "3.12"
|
||||
|
|
@ -446,6 +453,7 @@ jobs:
|
|||
tests/unit/repositories
|
||||
tests/unit/secret_managers
|
||||
tests/unit/skills/test_skills_main.py
|
||||
tests/unit/telemetry tests/unit/proxy/telemetry
|
||||
tests/unit/tracing
|
||||
python-version: "3.12"
|
||||
workers: 2
|
||||
|
|
@ -794,6 +802,10 @@ jobs:
|
|||
tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py
|
||||
tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py
|
||||
tests/unit/proxy/_experimental/mcp_server/test_db_credentials.py
|
||||
tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py
|
||||
tests/unit/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py
|
||||
tests/unit/proxy/_experimental/mcp_server/test_mcp_max_concurrent_requests.py
|
||||
tests/unit/proxy/_experimental/mcp_server/faults/test_list_outcomes.py
|
||||
tests/unit/proxy/_experimental/mcp_server/outbound_credentials
|
||||
python-version: "3.12"
|
||||
workers: 4
|
||||
|
|
@ -802,6 +814,28 @@ jobs:
|
|||
job-timeout-minutes: 60
|
||||
dist: loadscope
|
||||
|
||||
- shard: mcp-server
|
||||
test-path: |-
|
||||
tests/unit/experimental_mcp_client
|
||||
tests/unit/proxy/_experimental/mcp_server
|
||||
tests/unit/responses/mcp
|
||||
--ignore=tests/unit/experimental_mcp_client/test_mcp_client.py
|
||||
--ignore=tests/unit/proxy/_experimental/mcp_server/test_capabilities.py
|
||||
--ignore=tests/unit/proxy/_experimental/mcp_server/test_interactions.py
|
||||
--ignore=tests/unit/proxy/_experimental/mcp_server/test_mcp_elicitation_handler.py
|
||||
--ignore=tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py
|
||||
--ignore=tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py
|
||||
--ignore=tests/unit/proxy/_experimental/mcp_server/test_operations.py
|
||||
--ignore=tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py
|
||||
--ignore=tests/unit/proxy/_experimental/mcp_server/test_db_credentials.py
|
||||
--ignore=tests/unit/proxy/_experimental/mcp_server/outbound_credentials
|
||||
python-version: "3.12"
|
||||
workers: 4
|
||||
reruns: 0
|
||||
timeout-minutes: 20
|
||||
job-timeout-minutes: 60
|
||||
dist: loadscope
|
||||
|
||||
- shard: custom-logging
|
||||
test-path: |-
|
||||
tests/unit/proxy/test_proxy_server_custom_logger.py
|
||||
|
|
@ -1324,6 +1358,115 @@ jobs:
|
|||
- name: Check consumer types with the installed SDK
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: uv run --no-sync basedpyright --pythonpath .venv/bin/python --project tests/typing/pyrightconfig.json
|
||||
redis-compat:
|
||||
name: redis-py ${{ matrix.redis-py }} compatibility
|
||||
needs: rust-bridge
|
||||
if: ${{ !cancelled() }}
|
||||
permissions:
|
||||
contents: read
|
||||
id-token: write
|
||||
pull-requests: read
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 25
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
redis-py: ["5.3.1", "6.4.0", "7.4.1", "8.0.1"]
|
||||
env:
|
||||
UV_PYTHON: "3.12"
|
||||
LITELLM_LOCAL_MODEL_COST_MAP: "True"
|
||||
REDIS_PY_VERSION: ${{ matrix.redis-py }}
|
||||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
timeout-minutes: 3
|
||||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Detect relevant changes
|
||||
id: changes
|
||||
timeout-minutes: 2
|
||||
uses: ./.github/actions/detect-changes
|
||||
with:
|
||||
category: redis-compat
|
||||
|
||||
- name: Set up Python
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
timeout-minutes: 3
|
||||
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
with:
|
||||
python-version: ${{ env.UV_PYTHON }}
|
||||
|
||||
- name: Set up uv
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
timeout-minutes: 3
|
||||
uses: ./.github/actions/setup-uv-with-retries
|
||||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
- name: Cache uv dependencies
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
timeout-minutes: 5
|
||||
uses: ./.github/actions/cache-uv-downloads
|
||||
|
||||
- name: Set up the Rust build
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
timeout-minutes: 5
|
||||
uses: ./.github/actions/rust-bridge
|
||||
with:
|
||||
artifact: ${{ needs.rust-bridge.outputs.artifact }}
|
||||
|
||||
- name: Install dependencies with redis-py ${{ matrix.redis-py }}
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
timeout-minutes: 8
|
||||
env:
|
||||
RUST_BRIDGE_ARTIFACT: ${{ needs.rust-bridge.outputs.artifact }}
|
||||
run: |
|
||||
if [ -z "$RUST_BRIDGE_ARTIFACT" ]; then
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra caching --extra extra_proxy --extra proxy-runtime --extra utils
|
||||
else
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra caching --extra extra_proxy --extra proxy-runtime --extra utils --no-install-project
|
||||
uv pip install --no-deps --python .venv/bin/python rust-bridge-dist/*.whl
|
||||
cp rust-bridge-dist/litellm/rust_bridge/_native.abi3.so litellm/rust_bridge/_native.abi3.so
|
||||
fi
|
||||
uv pip install --python .venv/bin/python "redis==${REDIS_PY_VERSION:?}"
|
||||
uv run --no-sync python -c 'import os, redis; assert redis.__version__ == os.environ["REDIS_PY_VERSION"], redis.__version__'
|
||||
|
||||
- name: Install redis-server 7.2.16
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
timeout-minutes: 5
|
||||
run: |
|
||||
mkdir -p "$RUNNER_TEMP/redis-bin"
|
||||
cid="$(docker create redis:7.2.16@sha256:0637954999d01b7c9ce9167db2da50656e2590d3b884f1c600c5f63bb6e6773c)"
|
||||
docker cp "${cid}:/usr/local/bin/redis-server" "$RUNNER_TEMP/redis-bin/redis-server"
|
||||
docker rm "$cid"
|
||||
"$RUNNER_TEMP/redis-bin/redis-server" --version | grep -q 'v=7.2.16'
|
||||
echo "$RUNNER_TEMP/redis-bin" >> "$GITHUB_PATH"
|
||||
|
||||
- name: Run Redis compatibility tests
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
timeout-minutes: 15
|
||||
env:
|
||||
COVERAGE_CORE: sysmon
|
||||
run: |
|
||||
uv run --no-sync pytest \
|
||||
tests/unit/test_redis.py \
|
||||
tests/unit/caching/test_redis_connection_pool.py \
|
||||
tests/unit/caching/test_redis_cluster_cache.py \
|
||||
tests/unit/caching/test_evicted_client_closer.py \
|
||||
tests/integration/sdk/test_redis_cluster_iam_auth.py \
|
||||
--tb=short -vv -rs -p no:rerunfailures -p no:pytest-retry --timeout=120 --durations=20 \
|
||||
--cov=./litellm --cov-report=xml:coverage.xml
|
||||
|
||||
- name: Upload redis-py ${{ matrix.redis-py }} coverage to Codecov
|
||||
if: always() && steps.changes.outputs.decision != 'skip' && matrix.redis-py == '5.3.1'
|
||||
continue-on-error: true
|
||||
uses: codecov/codecov-action@0fb7174895f61a3b6b78fc075e0cd60383518dac # v5.5.5
|
||||
with:
|
||||
use_oidc: true
|
||||
files: coverage.xml
|
||||
root_dir: ${{ github.workspace }}
|
||||
flags: redis-compat
|
||||
fail_ci_if_error: false
|
||||
coverage:
|
||||
name: coverage
|
||||
needs: unit
|
||||
|
|
@ -1369,10 +1512,20 @@ jobs:
|
|||
root_dir: ${{ github.workspace }}
|
||||
flags: unit
|
||||
fail_ci_if_error: false
|
||||
|
||||
- name: Upload MCP server coverage under its own flag
|
||||
continue-on-error: true
|
||||
uses: codecov/codecov-action@0fb7174895f61a3b6b78fc075e0cd60383518dac # v5.5.5
|
||||
with:
|
||||
use_oidc: true
|
||||
files: coverage-reports/coverage-mcp-server-${{ github.run_id }}-${{ github.run_attempt }}/coverage.xml
|
||||
root_dir: ${{ github.workspace }}
|
||||
flags: mcp-integration
|
||||
fail_ci_if_error: false
|
||||
unit-passed:
|
||||
name: unit passed
|
||||
permissions: {}
|
||||
needs: [rust-bridge, assert-shard-coverage, unit, ui-unit, docs, helm, openai-sdk-compat]
|
||||
needs: [rust-bridge, assert-shard-coverage, unit, ui-unit, docs, helm, openai-sdk-compat, redis-compat]
|
||||
if: always()
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 2
|
||||
|
|
|
|||
4
.github/workflows/weekly_load_anomaly.yml
vendored
4
.github/workflows/weekly_load_anomaly.yml
vendored
|
|
@ -63,7 +63,7 @@ jobs:
|
|||
|
||||
- name: Start the proxy
|
||||
run: |
|
||||
nohup uv run --no-sync litellm --config tests/e2e/load/weekly_anomaly_config.yml --port 4000 > proxy.log 2>&1 &
|
||||
nohup uv run --no-sync litellm --config tests/e2e/perf/weekly_anomaly_config.yml --port 4000 > proxy.log 2>&1 &
|
||||
for _ in $(seq 1 90); do
|
||||
if curl -fs http://localhost:4000/health/liveliness > /dev/null; then
|
||||
exit 0
|
||||
|
|
@ -78,7 +78,7 @@ jobs:
|
|||
env:
|
||||
E2E_WEEKLY_ANOMALY: "1"
|
||||
run: |
|
||||
uv run --no-sync pytest tests/e2e/load/test_weekly_session_anomaly_e2e.py -v --tb=short -rA
|
||||
uv run --no-sync pytest tests/e2e/perf/test_weekly_session_anomaly_e2e.py -v --tb=short -rA
|
||||
|
||||
- name: Show proxy log on failure
|
||||
if: failure()
|
||||
|
|
|
|||
|
|
@ -1,613 +0,0 @@
|
|||
{
|
||||
"files": [
|
||||
{
|
||||
"path": "litellm/_v2/AGENTS.md",
|
||||
"description": "Conventions for code under litellm/_v2/",
|
||||
"scope": [
|
||||
"litellm/_v2/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm/_v2/cache/AGENTS.md",
|
||||
"description": "Conventions for code under litellm/_v2/cache/",
|
||||
"scope": [
|
||||
"litellm/_v2/cache/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm/proxy/_experimental/mcp_server/AGENTS.md",
|
||||
"description": "Conventions for code under litellm/proxy/_experimental/mcp_server/",
|
||||
"scope": [
|
||||
"litellm/proxy/_experimental/mcp_server/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm/rust_bridge/AGENTS.md",
|
||||
"description": "Conventions for code under litellm/rust_bridge/",
|
||||
"scope": [
|
||||
"litellm/rust_bridge/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm/tracing/AGENTS.md",
|
||||
"description": "Conventions for code under litellm/tracing/",
|
||||
"scope": [
|
||||
"litellm/tracing/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/",
|
||||
"scope": [
|
||||
"litellm-rust/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/cache-response/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/cache-response/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/cache-response/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/callbacks-legacy-python/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/callbacks-legacy-python/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/callbacks-legacy-python/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/coroutine/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/coroutine/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/coroutine/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/gateway/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/gateway/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/gateway/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/gateway-auth/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/gateway-auth/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/gateway-auth/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/gateway-inference/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/gateway-inference/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/gateway-inference/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/gateway-mcp/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/gateway-mcp/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/gateway-mcp/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/host/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/host/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/host/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/host/src/machine/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/host/src/machine/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/host/src/machine/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/host-http/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/host-http/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/host-http/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/host-native/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/host-native/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/host-native/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/host-python/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/host-python/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/host-python/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/http/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/http/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/http/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/inference/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/inference/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/inference/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/inference-chat/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/inference-chat/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/inference-chat/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/inference-messages/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/inference-messages/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/inference-messages/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/inference-ocr/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/inference-ocr/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/inference-ocr/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/inference-responses/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/inference-responses/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/inference-responses/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/inference-transcription/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/inference-transcription/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/inference-transcription/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/llms/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/llms/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/llms/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/llms/src/anthropic/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/llms/src/anthropic/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/llms/src/anthropic/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/llms/src/anthropic/batches/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/llms/src/anthropic/batches/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/llms/src/anthropic/batches/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/llms/src/anthropic/chat/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/llms/src/anthropic/chat/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/llms/src/anthropic/chat/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/llms/src/anthropic/count_tokens/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/llms/src/anthropic/count_tokens/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/llms/src/anthropic/count_tokens/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/llms/src/anthropic/messages/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/llms/src/anthropic/messages/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/llms/src/anthropic/messages/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/llms/src/aws_textract/ocr/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/llms/src/aws_textract/ocr/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/llms/src/aws_textract/ocr/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/llms/src/azure_ai/messages/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/llms/src/azure_ai/messages/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/llms/src/azure_ai/messages/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/llms/src/azure_ai/ocr/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/llms/src/azure_ai/ocr/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/llms/src/azure_ai/ocr/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/llms/src/base_llm/messages/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/llms/src/base_llm/messages/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/llms/src/base_llm/messages/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/llms/src/bedrock/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/llms/src/bedrock/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/llms/src/bedrock/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/llms/src/bedrock/audio_transcription/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/llms/src/bedrock/audio_transcription/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/llms/src/bedrock/audio_transcription/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/llms/src/bedrock/chat/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/llms/src/bedrock/chat/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/llms/src/bedrock/chat/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/llms/src/bedrock/messages/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/llms/src/bedrock/messages/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/llms/src/bedrock/messages/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/llms/src/cohere/ocr/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/llms/src/cohere/ocr/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/llms/src/cohere/ocr/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/llms/src/deepseek/messages/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/llms/src/deepseek/messages/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/llms/src/deepseek/messages/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/llms/src/mistral/ocr/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/llms/src/mistral/ocr/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/llms/src/mistral/ocr/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/llms/src/openai/responses/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/llms/src/openai/responses/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/llms/src/openai/responses/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/llms/src/openai_like/chat/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/llms/src/openai_like/chat/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/llms/src/openai_like/chat/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/llms/src/reducto/ocr/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/llms/src/reducto/ocr/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/llms/src/reducto/ocr/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/llms/src/vertex_ai/messages/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/llms/src/vertex_ai/messages/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/llms/src/vertex_ai/messages/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/llms/src/vertex_ai/ocr/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/llms/src/vertex_ai/ocr/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/llms/src/vertex_ai/ocr/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/llms-types/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/llms-types/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/llms-types/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/llms-types/src/formats/audio_transcription/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/llms-types/src/formats/audio_transcription/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/llms-types/src/formats/audio_transcription/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/llms-types/src/formats/batches/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/llms-types/src/formats/batches/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/llms-types/src/formats/batches/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/llms-types/src/formats/chat_completions/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/llms-types/src/formats/chat_completions/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/llms-types/src/formats/chat_completions/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/llms-types/src/formats/messages/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/llms-types/src/formats/messages/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/llms-types/src/formats/messages/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/llms-types/src/formats/ocr/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/llms-types/src/formats/ocr/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/llms-types/src/formats/ocr/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/llms-types/src/formats/responses/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/llms-types/src/formats/responses/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/llms-types/src/formats/responses/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/llms-types/src/providers/anthropic/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/llms-types/src/providers/anthropic/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/llms-types/src/providers/anthropic/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/llms-types/src/providers/bedrock/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/llms-types/src/providers/bedrock/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/llms-types/src/providers/bedrock/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/llms-types/src/providers/minimax/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/llms-types/src/providers/minimax/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/llms-types/src/providers/minimax/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/llms-types/src/providers/vertex_ai/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/llms-types/src/providers/vertex_ai/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/llms-types/src/providers/vertex_ai/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/model-catalog/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/model-catalog/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/model-catalog/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/python-bridge/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/python-bridge/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/python-bridge/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/python-bridge/src/cache/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/python-bridge/src/cache/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/python-bridge/src/cache/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/python-bridge/src/cache/native/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/python-bridge/src/cache/native/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/python-bridge/src/cache/native/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/python-bridge/src/cache/python/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/python-bridge/src/cache/python/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/python-bridge/src/cache/python/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/python-bridge/src/routes/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/python-bridge/src/routes/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/python-bridge/src/routes/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/python-compat/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/python-compat/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/python-compat/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/secrets-aws/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/secrets-aws/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/secrets-aws/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/secrets-azure/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/secrets-azure/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/secrets-azure/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/secrets-cyberark/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/secrets-cyberark/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/secrets-cyberark/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/secrets-google/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/secrets-google/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/secrets-google/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/secrets-hashicorp/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/secrets-hashicorp/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/secrets-hashicorp/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/crates/storage-clickhouse/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/crates/storage-clickhouse/",
|
||||
"scope": [
|
||||
"litellm-rust/crates/storage-clickhouse/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "litellm-rust/docs/adrs/AGENTS.md",
|
||||
"description": "Conventions for code under litellm-rust/docs/adrs/",
|
||||
"scope": [
|
||||
"litellm-rust/docs/adrs/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "tests/AGENTS.md",
|
||||
"description": "Conventions for code under tests/",
|
||||
"scope": [
|
||||
"tests/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "tests/e2e/AGENTS.md",
|
||||
"description": "Conventions for code under tests/e2e/",
|
||||
"scope": [
|
||||
"tests/e2e/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "tests/e2e_harness/AGENTS.md",
|
||||
"description": "Conventions for code under tests/e2e_harness/",
|
||||
"scope": [
|
||||
"tests/e2e_harness/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "tests/integration/AGENTS.md",
|
||||
"description": "Conventions for code under tests/integration/",
|
||||
"scope": [
|
||||
"tests/integration/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "tests/integration/translation/AGENTS.md",
|
||||
"description": "Conventions for code under tests/integration/translation/",
|
||||
"scope": [
|
||||
"tests/integration/translation/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "tests/rust-python-harness/AGENTS.md",
|
||||
"description": "Conventions for code under tests/rust-python-harness/",
|
||||
"scope": [
|
||||
"tests/rust-python-harness/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "tests/rust-python-harness/strategies/e2e_parity/AGENTS.md",
|
||||
"description": "Conventions for code under tests/rust-python-harness/strategies/e2e_parity/",
|
||||
"scope": [
|
||||
"tests/rust-python-harness/strategies/e2e_parity/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "tests/rust-python-harness/strategies/trace_parity/AGENTS.md",
|
||||
"description": "Conventions for code under tests/rust-python-harness/strategies/trace_parity/",
|
||||
"scope": [
|
||||
"tests/rust-python-harness/strategies/trace_parity/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "tests/rust-python-harness/strategies/unit_tests_parity/AGENTS.md",
|
||||
"description": "Conventions for code under tests/rust-python-harness/strategies/unit_tests_parity/",
|
||||
"scope": [
|
||||
"tests/rust-python-harness/strategies/unit_tests_parity/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "tests/rust-python-harness/strategies/unit_tests_rust/AGENTS.md",
|
||||
"description": "Conventions for code under tests/rust-python-harness/strategies/unit_tests_rust/",
|
||||
"scope": [
|
||||
"tests/rust-python-harness/strategies/unit_tests_rust/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "tests/test_litellm_rust/AGENTS.md",
|
||||
"description": "Conventions for code under tests/test_litellm_rust/",
|
||||
"scope": [
|
||||
"tests/test_litellm_rust/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "tests/unit/AGENTS.md",
|
||||
"description": "Conventions for code under tests/unit/",
|
||||
"scope": [
|
||||
"tests/unit/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "tests/unit/rust_bridge/AGENTS.md",
|
||||
"description": "Conventions for code under tests/unit/rust_bridge/",
|
||||
"scope": [
|
||||
"tests/unit/rust_bridge/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "ui/litellm-dashboard/AGENTS.md",
|
||||
"description": "Conventions for code under ui/litellm-dashboard/",
|
||||
"scope": [
|
||||
"ui/litellm-dashboard/**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"path": "ui/litellm-dashboard/src/components/chat/AGENTS.md",
|
||||
"description": "Conventions for code under ui/litellm-dashboard/src/components/chat/",
|
||||
"scope": [
|
||||
"ui/litellm-dashboard/src/components/chat/**"
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
|
|
@ -23,7 +23,7 @@ When adding new features, add meaningful tests. Don't add tests that don't check
|
|||
|
||||
Same thing for bug fixes. The tests should make it so that this specific bug can never happen again without failing tests (i.e., regression)
|
||||
|
||||
Never test structure of code only function of it
|
||||
Never test structure or implementation details of code, only function of it. A test should catch a change that breaks behavior by accident. If the only way it can fail is someone deliberately undoing the change, don't write 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
|
||||
|
||||
|
|
@ -130,4 +130,4 @@ Before implementing:
|
|||
|
||||
Ask yourself: "Would a senior engineer say this is overcomplicated?" If yes, simplify
|
||||
|
||||
Before requesting maintainer review, verify the current PR tip passes required CI and code coverage, meets Greptile confidence of at least 4/5, and has acceptable Veria and Bugbot reviews. Inspect warnings and findings, fix actionable issues, and rerun the affected checks and reviewers after changes. Record evidence for any false positive or unavailable review. Never treat a pending or missing bot result as a pass. Do not lower coverage thresholds or raise `ANY_CAPS` to satisfy a check
|
||||
Before requesting maintainer review, verify the current PR tip passes required CI and code coverage and has acceptable reviews from the review bots on the PR, such as Veria and Bugbot. Inspect warnings and findings, fix actionable issues, and rerun the affected checks and reviewers after changes. Record evidence for any false positive or unavailable review. Never treat a pending or missing bot result as a pass. Do not lower coverage thresholds or raise `ANY_CAPS` to satisfy a check
|
||||
|
|
|
|||
2
Makefile
2
Makefile
|
|
@ -290,7 +290,7 @@ test-rust-extension:
|
|||
"$$temporary/venv/bin/python" -I -m mypy.stubtest \
|
||||
--mypy-config-file tests/unit/rust_bridge/stubtest.ini \
|
||||
litellm.rust_bridge._native && \
|
||||
LITELLM_RUST=1 LITELLM_LOCAL_MODEL_COST_MAP=True \
|
||||
LITELLM_LOCAL_MODEL_COST_MAP=True \
|
||||
"$$temporary/venv/bin/python" -I -m pytest --import-mode=importlib -m requires_rust_extension tests/test_litellm_rust
|
||||
|
||||
rust-sqlx-prepare:
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@
|
|||
<a href="https://render.com/deploy?repo=https://github.com/BerriAI/litellm" target="_blank" rel="nofollow"><img src="https://render.com/images/deploy-to-render-button.svg" alt="Deploy to Render" height="40"></a>
|
||||
<a href="https://railway.com/deploy/RhvhdC?referralCode=7mRv9K&utm_medium=integration&utm_source=template&utm_campaign=generic"><img src="https://railway.com/button.svg" alt="Deploy on Railway" height="40"></a>
|
||||
<a href="https://console.aws.amazon.com/cloudshell/home" target="_blank" rel="nofollow"><img src="./.github/deploy-on-aws.png" alt="Deploy on AWS" height="40"></a>
|
||||
<a href="https://ssh.cloud.google.com/cloudshell/editor?cloudshell_git_repo=https%3A%2F%2Fgithub.com%2FBerriAI%2Flitellm&cloudshell_workspace=terraform%2Flitellm%2Fgcp%2Fexamples%2Fdefault&cloudshell_tutorial=TUTORIAL.md&cloudshell_image=gcr.io/ds-artifacts-cloudshell/deploystack_custom_image&shellonly=true" target="_blank" rel="nofollow"><img src="./.github/deploy-on-gcp.png" alt="Deploy on GCP" height="40"></a>
|
||||
<a href="https://ssh.cloud.google.com/cloudshell/editor?cloudshell_git_repo=https%3A%2F%2Fgithub.com%2FBerriAI%2Fterraform-google-litellm&cloudshell_git_branch=v1.105.0&cloudshell_workspace=examples%2Fdefault&cloudshell_tutorial=TUTORIAL.md&cloudshell_image=gcr.io/ds-artifacts-cloudshell/deploystack_custom_image&shellonly=true" target="_blank" rel="nofollow"><img src="./.github/deploy-on-gcp.png" alt="Deploy on GCP" height="40"></a>
|
||||
</p>
|
||||
</p>
|
||||
<h4 align="center"><a href="https://docs.litellm.ai/docs/simple_proxy" target="_blank">LiteLLM Proxy Server (AI Gateway)</a> | <a href="https://docs.litellm.ai/docs/enterprise#hosted-litellm-proxy" target="_blank"> Hosted Proxy</a> | <a href="https://litellm.ai/enterprise"target="_blank">Enterprise Tier</a> | <a href="https://www.litellm.ai/ai-gateway" target="_blank">Website</a></h4>
|
||||
|
|
@ -342,6 +342,7 @@ Set `LITELLM_PROXY_API_BASE` and `LITELLM_PROXY_API_KEY` and every model call th
|
|||
| [Cohere Chat (`cohere_chat`)](https://docs.litellm.ai/docs/providers/cohere) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [CometAPI (`cometapi`)](https://docs.litellm.ai/docs/providers/cometapi) | ✅ | ✅ | ✅ | ✅ | | | | | | |
|
||||
| [CompactifAI (`compactifai`)](https://docs.litellm.ai/docs/providers/compactifai) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [CoralBricks (`coralbricks`)](https://docs.litellm.ai/docs/providers/coralbricks) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [Custom (`custom`)](https://docs.litellm.ai/docs/providers/custom_llm_server) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [Custom OpenAI (`custom_openai`)](https://docs.litellm.ai/docs/providers/openai_compatible) | ✅ | ✅ | ✅ | | | ✅ | ✅ | ✅ | ✅ | |
|
||||
| [Dashscope (`dashscope`)](https://docs.litellm.ai/docs/providers/dashscope) | ✅ | ✅ | ✅ | ✅ | | | | | | ✅ |
|
||||
|
|
@ -526,9 +527,9 @@ Provider API keys live in AWS Secrets Manager; reference ARNs via `gateway_extra
|
|||
|
||||
#### GCP — Cloud Run + Cloud SQL + Memorystore + HTTPS LB
|
||||
|
||||
[](https://ssh.cloud.google.com/cloudshell/editor?cloudshell_git_repo=https%3A%2F%2Fgithub.com%2FBerriAI%2Flitellm&cloudshell_workspace=terraform%2Flitellm%2Fgcp%2Fexamples%2Fdefault&cloudshell_tutorial=TUTORIAL.md&cloudshell_image=gcr.io/ds-artifacts-cloudshell/deploystack_custom_image&shellonly=true)
|
||||
[](https://ssh.cloud.google.com/cloudshell/editor?cloudshell_git_repo=https%3A%2F%2Fgithub.com%2FBerriAI%2Fterraform-google-litellm&cloudshell_git_branch=v1.105.0&cloudshell_workspace=examples%2Fdefault&cloudshell_tutorial=TUTORIAL.md&cloudshell_image=gcr.io/ds-artifacts-cloudshell/deploystack_custom_image&shellonly=true)
|
||||
|
||||
Real 1-click. Opens Cloud Shell, clones this repo, and walks you through `terraform apply` via a built-in [DeployStack tutorial](./terraform/litellm/gcp/examples/default/TUTORIAL.md) — pick the project, the tutorial sets up the Artifact Registry remote repo, writes `terraform.tfvars` from your answers, and runs apply.
|
||||
Real 1-click. Opens Cloud Shell, clones the [terraform-google-litellm](https://github.com/BerriAI/terraform-google-litellm) module mirror at the latest stable release (a small download instead of the full litellm repo), and walks you through `terraform apply` via a built-in [DeployStack tutorial](./terraform/litellm/gcp/examples/default/TUTORIAL.md). Pick the project, the tutorial sets up the Artifact Registry remote repo, writes `terraform.tfvars` from your answers, and runs apply
|
||||
|
||||
[Module page →](https://registry.terraform.io/modules/BerriAI/litellm/google/latest)
|
||||
|
||||
|
|
|
|||
|
|
@ -32,6 +32,7 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = (
|
|||
"/token",
|
||||
"/onboarding/",
|
||||
"/audit",
|
||||
"/telemetry/",
|
||||
"/oauth/",
|
||||
"/invitation/",
|
||||
"/jwt/",
|
||||
|
|
|
|||
|
|
@ -56,7 +56,7 @@ def _rendered_schema(cost_map: CostMap) -> str:
|
|||
def _evaluation_output_price_failure(
|
||||
key: str, entry: object, evaluation_output_price_sources: Mapping[str, str]
|
||||
) -> str | None:
|
||||
if not isinstance(entry, dict) or entry.get("mode") != "evaluation":
|
||||
if not isinstance(entry, dict) or entry.get("mode") not in ("evaluation", "decisions"):
|
||||
return None
|
||||
output_cost: Final = entry.get("output_cost_per_token")
|
||||
if (
|
||||
|
|
@ -67,7 +67,7 @@ def _evaluation_output_price_failure(
|
|||
):
|
||||
return None
|
||||
return (
|
||||
f"{key}: mode evaluation rows bill input only; output_cost_per_token must be 0 unless the vendor "
|
||||
f"{key}: mode evaluation and decisions rows bill input only; output_cost_per_token must be 0 unless the vendor "
|
||||
"prices output tokens for this model (add the source to EVALUATION_OUTPUT_PRICE_SOURCES)"
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -322,7 +322,12 @@ def create_migration(
|
|||
|
||||
# Write the SQL to migration.sql
|
||||
migration_file = migration_dir / "migration.sql"
|
||||
migration_file.write_text(diff_sql)
|
||||
migration_sql: Final = (
|
||||
diff_sql
|
||||
if "CONCURRENTLY" in diff_sql.upper()
|
||||
else "SET LOCAL lock_timeout = '10s';\n\n" + diff_sql
|
||||
)
|
||||
migration_file.write_text(migration_sql)
|
||||
|
||||
print(f"Created migration in {migration_dir}")
|
||||
return True
|
||||
|
|
|
|||
|
|
@ -71,6 +71,8 @@ flag_management:
|
|||
carryforward: false
|
||||
- name: lens-python-310
|
||||
carryforward: false
|
||||
- name: lens-postgres
|
||||
carryforward: false
|
||||
- name: proxy-db-key-generation
|
||||
carryforward: false
|
||||
- name: proxy-db-auth-checks
|
||||
|
|
|
|||
|
|
@ -55,6 +55,8 @@ ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_session_id_idx"
|
|||
RENAME TO "LiteLLM_SpendLogs_legacy_session_id_idx";
|
||||
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx"
|
||||
RENAME TO "LiteLLM_SpendLogs_legacy_api_key_startTime_idx";
|
||||
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_litellm_call_id_idx"
|
||||
RENAME TO "LiteLLM_SpendLogs_legacy_litellm_call_id_idx";
|
||||
|
||||
CREATE TABLE "LiteLLM_SpendLogs" (
|
||||
LIKE "LiteLLM_SpendLogs_legacy" INCLUDING DEFAULTS INCLUDING GENERATED
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@
|
|||
# beyond local evaluation, pin the image to a specific release tag.
|
||||
services:
|
||||
litellm:
|
||||
image: docker.litellm.ai/berriai/litellm:main-stable
|
||||
image: docker.litellm.ai/berriai/litellm:latest
|
||||
ports:
|
||||
# LITELLM_BIND is empty by default, so this stays "4000:4000". The quickstart
|
||||
# script sets it to "127.0.0.1:" so new installs listen on this machine only.
|
||||
|
|
|
|||
|
|
@ -8,21 +8,20 @@ from datetime import datetime, timedelta, timezone
|
|||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, List, Literal, Optional, Protocol, Tuple, cast
|
||||
|
||||
from litellm_enterprise.proxy.common_utils.managed_object_attribution import (
|
||||
ManagedObjectCreatorAttribution,
|
||||
ManagedObjectRow,
|
||||
)
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.constants import (
|
||||
CLI_SESSION_KEY_PREFIX,
|
||||
MANAGED_OBJECT_STALENESS_CUTOFF_DAYS,
|
||||
MAX_OBJECTS_PER_POLL_CYCLE,
|
||||
)
|
||||
from litellm.repositories.table_repositories import ManagedObjectRepository
|
||||
from litellm.repositories.team_repository import TeamRepository
|
||||
from litellm.repositories.user_repository import UserRepository
|
||||
from litellm.repositories.verification_token_repository import VerificationTokenRepository
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from prisma import models as prisma_models
|
||||
|
||||
from litellm.integrations.prometheus import PrometheusLogger
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
||||
from litellm.repositories.prisma_protocols import TableActions
|
||||
|
|
@ -47,29 +46,6 @@ TERMINAL_MANAGED_OBJECT_STATUSES: Final[Tuple[str, ...]] = (
|
|||
)
|
||||
|
||||
|
||||
class _ManagedObjectRow(Protocol):
|
||||
@property
|
||||
def id(self) -> str: ...
|
||||
|
||||
@property
|
||||
def unified_object_id(self) -> str: ...
|
||||
|
||||
@property
|
||||
def created_by(self) -> str | None: ...
|
||||
|
||||
@property
|
||||
def file_object(self) -> object: ...
|
||||
|
||||
@property
|
||||
def org_id(self) -> str | None: ...
|
||||
|
||||
@property
|
||||
def api_key(self) -> str | None: ...
|
||||
|
||||
@property
|
||||
def team_id(self) -> str | None: ...
|
||||
|
||||
|
||||
class _ReadableFileContent(Protocol):
|
||||
async def read(self) -> bytes: ...
|
||||
|
||||
|
|
@ -87,22 +63,10 @@ async def _file_content_bytes(file_content: object) -> bytes:
|
|||
return cast(bytes, file_content)
|
||||
|
||||
|
||||
def _managed_object_table(prisma_client: "PrismaClient") -> "TableActions[_ManagedObjectRow]":
|
||||
def _managed_object_table(prisma_client: "PrismaClient") -> "TableActions[ManagedObjectRow]":
|
||||
return ManagedObjectRepository(prisma_client).table
|
||||
|
||||
|
||||
def _user_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.LiteLLM_UserTable]":
|
||||
return UserRepository(prisma_client).table
|
||||
|
||||
|
||||
def _token_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.LiteLLM_VerificationToken]":
|
||||
return VerificationTokenRepository(prisma_client).table
|
||||
|
||||
|
||||
def _team_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.LiteLLM_TeamTable]":
|
||||
return TeamRepository(prisma_client).table
|
||||
|
||||
|
||||
class CheckBatchCost:
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -118,6 +82,7 @@ class CheckBatchCost:
|
|||
self.prisma_client: PrismaClient = prisma_client
|
||||
self.llm_router: Router = llm_router
|
||||
self._track_unmanaged_batch_cost = track_unmanaged_batch_cost
|
||||
self._creator_attribution: Final = ManagedObjectCreatorAttribution(prisma_client, "CheckBatchCost")
|
||||
# Cached after the first poll cycle. Once we know the column is absent we skip
|
||||
# the guaranteed-failing primary query on every subsequent cycle.
|
||||
self._has_batch_processed_column: bool = True
|
||||
|
|
@ -149,133 +114,6 @@ class CheckBatchCost:
|
|||
return
|
||||
self.batch_processed_support_confirmed = True
|
||||
|
||||
async def _get_user_info(self, batch_id: str, user_id: Optional[str]) -> dict[str, str | None]:
|
||||
"""
|
||||
Look up user email and key alias by user_id for enriching the S3 callback metadata.
|
||||
Returns a dict with user_api_key_user_email and user_api_key_alias (both may be None).
|
||||
Returns an empty dict when user_id is None: batches created by a team or service
|
||||
account key carry no user id, and find_unique(where={"user_id": None}) raises.
|
||||
"""
|
||||
if not user_id:
|
||||
return {}
|
||||
try:
|
||||
user_row: prisma_models.LiteLLM_UserTable | None = await _user_table(self.prisma_client).find_unique(
|
||||
where={"user_id": user_id}
|
||||
)
|
||||
if user_row is None:
|
||||
return {}
|
||||
return {
|
||||
"user_api_key_user_email": getattr(user_row, "user_email", None),
|
||||
"user_api_key_alias": getattr(user_row, "user_alias", None),
|
||||
}
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error(f"CheckBatchCost: could not look up user {user_id} for batch {batch_id}: {e}")
|
||||
return {}
|
||||
|
||||
async def _get_key_alias(self, batch_id: str, api_key: str | None, created_by: str | None) -> str | None:
|
||||
"""Resolve the creating virtual key's alias from its hashed token."""
|
||||
if not api_key:
|
||||
return None
|
||||
if created_by and api_key == f"{CLI_SESSION_KEY_PREFIX}-{created_by}":
|
||||
return api_key
|
||||
try:
|
||||
key_row: prisma_models.LiteLLM_VerificationToken | None = await _token_table(
|
||||
self.prisma_client
|
||||
).find_unique(where={"token": api_key})
|
||||
return getattr(key_row, "key_alias", None) if key_row is not None else None
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error(f"CheckBatchCost: could not look up key alias for batch {batch_id}: {e}")
|
||||
return None
|
||||
|
||||
async def _get_team_alias(self, team_id: str | None) -> str | None:
|
||||
"""Resolve a team's alias from its id."""
|
||||
if not team_id:
|
||||
return None
|
||||
try:
|
||||
team_row: prisma_models.LiteLLM_TeamTable | None = await _team_table(self.prisma_client).find_unique(
|
||||
where={"team_id": team_id}
|
||||
)
|
||||
return getattr(team_row, "team_alias", None) if team_row is not None else None
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error(f"CheckBatchCost: could not look up team alias for team {team_id}: {e}")
|
||||
return None
|
||||
|
||||
async def _get_org_id(self, job: "_ManagedObjectRow", batch_id: str) -> str | None:
|
||||
org_id: Final = getattr(job, "org_id", None)
|
||||
if org_id:
|
||||
return org_id
|
||||
api_key: Final = getattr(job, "api_key", None)
|
||||
team_id: Final = getattr(job, "team_id", None)
|
||||
if api_key:
|
||||
try:
|
||||
key_row: prisma_models.LiteLLM_VerificationToken | None = await _token_table(
|
||||
self.prisma_client
|
||||
).find_unique(where={"token": api_key})
|
||||
key_org_id: Final = (
|
||||
cast(str | None, getattr(key_row, "organization_id", None))
|
||||
if key_row is not None
|
||||
else None
|
||||
)
|
||||
if key_org_id:
|
||||
return key_org_id
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error(
|
||||
f"CheckBatchCost: could not resolve the key's org for batch {batch_id}, "
|
||||
f"still trying the team's: {e}"
|
||||
)
|
||||
if not team_id:
|
||||
return None
|
||||
try:
|
||||
team_row: prisma_models.LiteLLM_TeamTable | None = await _team_table(self.prisma_client).find_unique(
|
||||
where={"team_id": team_id}
|
||||
)
|
||||
return cast(str | None, getattr(team_row, "organization_id", None)) if team_row is not None else None
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error(f"CheckBatchCost: could not resolve the team's org for batch {batch_id}: {e}")
|
||||
return None
|
||||
|
||||
async def _build_creator_attribution_metadata(
|
||||
self, job: "_ManagedObjectRow", batch_id: str
|
||||
) -> dict[str, object]:
|
||||
"""
|
||||
Rebuild the spend-tracking metadata for the key, team, and tags that created the
|
||||
batch so the batch-cost spend log is attributed the same way a non-batch request
|
||||
is. Rows created before api_key and request_tags were persisted carry only
|
||||
created_by and team_id, and fall back to those. A named creating key owns
|
||||
user_api_key_alias; when it has no alias, or the key has since been rotated or
|
||||
deleted, the field keeps the creating user's alias that _get_user_info filled in,
|
||||
because a resolvable name is more useful on the spend row than a null.
|
||||
|
||||
user_api_key_org_id must be resolved here too: the spend update writer reads it
|
||||
off this metadata to increment organization spend, so leaving it out silently
|
||||
drops batch cost from org accounting for keys and teams that belong to one.
|
||||
"""
|
||||
api_key = getattr(job, "api_key", None)
|
||||
team_id = getattr(job, "team_id", None)
|
||||
request_tags = getattr(job, "request_tags", None)
|
||||
|
||||
metadata: dict[str, object] = {
|
||||
"user_api_key_user_id": job.created_by,
|
||||
"user_api_key": api_key,
|
||||
"user_api_key_hash": api_key,
|
||||
"user_api_key_team_id": team_id,
|
||||
**(await self._get_user_info(batch_id, job.created_by)),
|
||||
}
|
||||
|
||||
key_alias = await self._get_key_alias(batch_id, api_key, job.created_by)
|
||||
if key_alias is not None:
|
||||
metadata["user_api_key_alias"] = key_alias
|
||||
team_alias = await self._get_team_alias(team_id)
|
||||
if team_alias is not None:
|
||||
metadata["user_api_key_team_alias"] = team_alias
|
||||
org_id: Final = await self._get_org_id(job, batch_id)
|
||||
if org_id is not None:
|
||||
metadata["user_api_key_org_id"] = org_id
|
||||
if isinstance(request_tags, list) and request_tags:
|
||||
metadata["tags"] = [tag for tag in request_tags if isinstance(tag, str)]
|
||||
|
||||
return metadata
|
||||
|
||||
async def _cleanup_stale_managed_objects(self) -> None:
|
||||
"""
|
||||
Mark managed objects older than MANAGED_OBJECT_STALENESS_CUTOFF_DAYS days
|
||||
|
|
@ -317,7 +155,7 @@ class CheckBatchCost:
|
|||
f"{MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days that were never costed"
|
||||
)
|
||||
|
||||
async def _fallback_find_jobs(self) -> "Sequence[_ManagedObjectRow]":
|
||||
async def _fallback_find_jobs(self) -> "Sequence[ManagedObjectRow]":
|
||||
"""Query batch jobs without the batch_processed filter (for older schemas)."""
|
||||
return await _managed_object_table(self.prisma_client).find_many(
|
||||
where={
|
||||
|
|
@ -337,7 +175,7 @@ class CheckBatchCost:
|
|||
order={"created_at": "asc"},
|
||||
)
|
||||
|
||||
async def _retire_job(self, job: "_ManagedObjectRow", reason: str) -> None:
|
||||
async def _retire_job(self, job: "ManagedObjectRow", reason: str) -> None:
|
||||
"""
|
||||
Take a row that can never be costed out of the poll page. Leaving it selectable
|
||||
would burn one of the MAX_OBJECTS_PER_POLL_CYCLE slots on every future cycle, and
|
||||
|
|
@ -364,7 +202,7 @@ class CheckBatchCost:
|
|||
"so it will no longer be polled"
|
||||
)
|
||||
|
||||
async def _claim_job_for_costing(self, job: "_ManagedObjectRow") -> bool:
|
||||
async def _claim_job_for_costing(self, job: "ManagedObjectRow") -> bool:
|
||||
"""
|
||||
Atomically flip batch_processed from false to true, returning whether this pod won
|
||||
the row. Every pod and uvicorn worker schedules its own poller against the shared
|
||||
|
|
@ -390,7 +228,7 @@ class CheckBatchCost:
|
|||
return False
|
||||
return claimed > 0
|
||||
|
||||
async def _release_job_claim(self, job: "_ManagedObjectRow") -> None:
|
||||
async def _release_job_claim(self, job: "ManagedObjectRow") -> None:
|
||||
"""Give a claimed row back once billing it failed, so a later poll cycle retries it.
|
||||
|
||||
Safe to match on batch_processed=True: while this poller is active the retrieve
|
||||
|
|
@ -411,7 +249,7 @@ class CheckBatchCost:
|
|||
)
|
||||
|
||||
@staticmethod
|
||||
def _has_unified_id_without_model(job: "_ManagedObjectRow") -> bool:
|
||||
def _has_unified_id_without_model(job: "ManagedObjectRow") -> bool:
|
||||
"""A unified id that decodes but carries no model_id can never be routed."""
|
||||
from litellm.proxy.openai_files_endpoints.common_utils import (
|
||||
convert_b64_uid_to_unified_uid,
|
||||
|
|
@ -460,7 +298,7 @@ class CheckBatchCost:
|
|||
return isinstance(error, (NotFoundError, openai.NotFoundError)) and output_file_id in str(error)
|
||||
|
||||
async def _finalize_unbilled_terminal_job(
|
||||
self, job: "_ManagedObjectRow", response: "LiteLLMBatch"
|
||||
self, job: "ManagedObjectRow", response: "LiteLLMBatch"
|
||||
) -> None:
|
||||
"""Persist a terminal batch that has nothing billable, converting any raw
|
||||
provider file ids to managed ids, and take it out of the poll page."""
|
||||
|
|
@ -505,7 +343,7 @@ class CheckBatchCost:
|
|||
|
||||
def _resolve_job_routing(
|
||||
self,
|
||||
job: "_ManagedObjectRow",
|
||||
job: "ManagedObjectRow",
|
||||
prom_logger: Optional["PrometheusLogger"],
|
||||
) -> Optional[Tuple[str, str]]:
|
||||
"""
|
||||
|
|
@ -582,7 +420,7 @@ class CheckBatchCost:
|
|||
|
||||
def _resolve_unmanaged_provider_routing(
|
||||
self,
|
||||
job: "_ManagedObjectRow",
|
||||
job: "ManagedObjectRow",
|
||||
prom_logger: Optional["PrometheusLogger"],
|
||||
llm_provider: str,
|
||||
bare_model_name: str,
|
||||
|
|
@ -678,7 +516,7 @@ class CheckBatchCost:
|
|||
@classmethod
|
||||
def _get_managed_file_model_name(
|
||||
cls,
|
||||
job: "_ManagedObjectRow",
|
||||
job: "ManagedObjectRow",
|
||||
deployment_info: "Deployment",
|
||||
) -> Optional[str]:
|
||||
"""
|
||||
|
|
@ -698,7 +536,7 @@ class CheckBatchCost:
|
|||
)
|
||||
|
||||
@staticmethod
|
||||
def _get_input_file_id(job: "_ManagedObjectRow") -> Optional[str]:
|
||||
def _get_input_file_id(job: "ManagedObjectRow") -> Optional[str]:
|
||||
import json
|
||||
|
||||
from litellm.types.utils import LiteLLMBatch
|
||||
|
|
@ -719,7 +557,7 @@ class CheckBatchCost:
|
|||
|
||||
async def _track_completed_batch_cost(
|
||||
self,
|
||||
job: "_ManagedObjectRow",
|
||||
job: "ManagedObjectRow",
|
||||
response: "LiteLLMBatch",
|
||||
model_id: str,
|
||||
batch_id: str,
|
||||
|
|
@ -735,9 +573,9 @@ class CheckBatchCost:
|
|||
|
||||
"""
|
||||
from litellm.batches.batch_utils import (
|
||||
calculate_batch_cost_and_usage,
|
||||
count_error_file_failed_requests,
|
||||
get_file_content_as_dictionary,
|
||||
calculate_batch_cost_and_usage,
|
||||
)
|
||||
from litellm.files.main import afile_content
|
||||
from litellm.files.types import FileContentCallOptions
|
||||
|
|
@ -911,7 +749,7 @@ class CheckBatchCost:
|
|||
},
|
||||
**({"api_base": mask_api_base_credentials(deployment_api_base)} if deployment_api_base else {}),
|
||||
"metadata": {
|
||||
**(await self._build_creator_attribution_metadata(job, batch_id)),
|
||||
**(await self._creator_attribution.batch_metadata(job, batch_id)),
|
||||
# spend logs read the deployment identity off these metadata keys, so
|
||||
# without them the batch cost row carries no model_id or model_group
|
||||
"model_info": {"id": model_id},
|
||||
|
|
|
|||
|
|
@ -1,12 +1,36 @@
|
|||
"""
|
||||
Polls LiteLLM_ManagedObjectTable to check if the response is complete.
|
||||
Cost tracking is handled by the get-responses call, which prices normally only because the
|
||||
poll stamps itself with BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN; user-facing reads of the
|
||||
same route are non-inference and free.
|
||||
A background response is billed exactly once: whichever client read or poll first sees it
|
||||
finished moves its row out of a pending status, then logs the body that read already fetched as
|
||||
one spend event stamped with BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN and the creating key's
|
||||
attribution, priced against the deployment that served it. The poll's status checks are stamped
|
||||
with BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN and client reads are marked managed by the
|
||||
managed files hook, so both stay free. The spend event carries the id the client holds, the row's
|
||||
unified_object_id, so its spend log is found under that id with the _background_cost suffix.
|
||||
A read that finds no row, because the create could not store one, stores it under the reading key
|
||||
and bills it the same way, so a response the create left untracked is still billed once. A streamed
|
||||
background create is free on its own; its stream stores the row under the creating key and bills
|
||||
the finished response through it before the client sees response.completed.
|
||||
A read the caller retargets with client-side credentials, and a delete of a still-pending response,
|
||||
settle the bill from a fetch through the configured deployment instead: the client controls what a
|
||||
retargeted read returns, and a deleted response leaves the poll nothing to fetch.
|
||||
A row stored before this billing model keeps the encrypted create id as its model_object_id. The
|
||||
first lookup by that id rekeys it by the provider id, and the poll retires such a row without
|
||||
billing it once a read has already tracked the response under the provider id.
|
||||
"""
|
||||
|
||||
import uuid
|
||||
from collections.abc import Awaitable, Callable, Mapping
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import TYPE_CHECKING, Dict, Final, Protocol, cast
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, cast
|
||||
|
||||
from litellm_enterprise.proxy.common_utils.managed_object_attribution import (
|
||||
ManagedObjectCreatorAttribution,
|
||||
ManagedObjectRow,
|
||||
)
|
||||
from pydantic import BaseModel, ConfigDict, TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
|
@ -14,36 +38,124 @@ from litellm.constants import (
|
|||
INTERNAL_CALL_ORIGIN_METADATA_KEY,
|
||||
MANAGED_OBJECT_STALENESS_CUTOFF_DAYS,
|
||||
MAX_OBJECTS_PER_POLL_CYCLE,
|
||||
PENDING_RESPONSE_STATUSES,
|
||||
STALE_OBJECT_CLEANUP_BATCH_SIZE,
|
||||
TERMINAL_RESPONSE_STATUSES,
|
||||
)
|
||||
from litellm.proxy.hooks.responses_id_security import decrypted_response_id
|
||||
from litellm.repositories.table_repositories import ManagedObjectRepository
|
||||
from litellm.responses.utils import ResponsesAPIRequestUtils
|
||||
from litellm.types.llms.openai import ResponsesAPIResponse
|
||||
from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN
|
||||
from litellm.types.router import Deployment
|
||||
from litellm.types.utils import (
|
||||
BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN,
|
||||
BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
||||
from litellm.repositories.prisma_protocols import TableActions
|
||||
from litellm.router import Router
|
||||
|
||||
TERMINAL_RESPONSE_STATUSES = frozenset({"completed", "failed", "cancelled", "incomplete"})
|
||||
_STORED_RESPONSE_ADAPTER: Final[TypeAdapter[dict[str, object]]] = TypeAdapter(dict[str, object])
|
||||
_RETRIEVE_RESPONSE_CALL_TYPE: Final = "aget_responses"
|
||||
|
||||
|
||||
class _ManagedObjectRow(Protocol):
|
||||
@property
|
||||
def id(self) -> str: ...
|
||||
|
||||
@property
|
||||
def unified_object_id(self) -> str: ...
|
||||
|
||||
@property
|
||||
def created_by(self) -> str | None: ...
|
||||
|
||||
@property
|
||||
def file_object(self) -> object: ...
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _PollTarget:
|
||||
response_id: str
|
||||
model_metadata: Mapping[str, object]
|
||||
deployment: Deployment | None
|
||||
|
||||
|
||||
def _managed_object_table(prisma_client: "PrismaClient") -> "TableActions[_ManagedObjectRow]":
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class BilledBackgroundResponse:
|
||||
response: ResponsesAPIResponse
|
||||
model: str
|
||||
custom_llm_provider: str | None
|
||||
litellm_params: Mapping[str, object]
|
||||
metadata: Mapping[str, object]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class BilledCharge:
|
||||
cost: float | None
|
||||
logging_obj: "LiteLLMLogging"
|
||||
|
||||
|
||||
BilledResponseLogger = Callable[[BilledBackgroundResponse], Awaitable[BilledCharge]]
|
||||
UntrackedResponseTracker = Callable[[ResponsesAPIResponse], Awaitable[bool]]
|
||||
|
||||
|
||||
def _without_read_price(response: ResponsesAPIResponse) -> ResponsesAPIResponse:
|
||||
unpriced: Final = response.model_copy(deep=True)
|
||||
unpriced.hidden_params = {key: value for key, value in response.hidden_params.items() if key != "response_cost"}
|
||||
return unpriced
|
||||
|
||||
|
||||
async def log_billed_background_response(billed: BilledBackgroundResponse) -> BilledCharge:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
|
||||
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
||||
|
||||
logging_obj: Final = LiteLLMLogging(
|
||||
model=billed.model,
|
||||
messages=[{"role": "user", "content": "<retrieve_response>"}],
|
||||
stream=False,
|
||||
call_type=_RETRIEVE_RESPONSE_CALL_TYPE,
|
||||
start_time=datetime.now(),
|
||||
litellm_call_id=str(uuid.uuid4()),
|
||||
function_id=str(uuid.uuid4()),
|
||||
)
|
||||
logging_obj.update_environment_variables(
|
||||
litellm_params={**billed.litellm_params, "metadata": dict(billed.metadata)},
|
||||
optional_params={},
|
||||
custom_llm_provider=billed.custom_llm_provider,
|
||||
)
|
||||
unpriced: Final = _without_read_price(billed.response)
|
||||
cost: Final = logging_obj.response_cost_calculator(result=unpriced)
|
||||
end_time: Final = datetime.now()
|
||||
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(
|
||||
async_coroutine=logging_obj.async_success_handler(
|
||||
result=unpriced, start_time=logging_obj.start_time, end_time=end_time
|
||||
)
|
||||
)
|
||||
logging_obj.handle_sync_success_callbacks_for_async_calls(
|
||||
result=unpriced, start_time=logging_obj.start_time, end_time=end_time
|
||||
)
|
||||
return BilledCharge(cost=cost, logging_obj=logging_obj)
|
||||
|
||||
|
||||
class _DeploymentBaseModel(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
base_model: str | None = None
|
||||
|
||||
|
||||
def _deployment_litellm_params(deployment: Deployment) -> Mapping[str, object]:
|
||||
from litellm.litellm_core_utils.litellm_logging import mask_api_base_credentials
|
||||
|
||||
api_base: Final = deployment.litellm_params.api_base
|
||||
base_model: Final = _DeploymentBaseModel.model_validate(deployment.litellm_params, from_attributes=True).base_model
|
||||
return MappingProxyType(
|
||||
{
|
||||
**({"api_base": mask_api_base_credentials(api_base)} if api_base else {}),
|
||||
**({"base_model": base_model} if base_model else {}),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _deployment_metadata(deployment: Deployment) -> Mapping[str, object]:
|
||||
return MappingProxyType(
|
||||
{
|
||||
"model_info": deployment.model_info.model_dump(exclude_none=True),
|
||||
"model_group": deployment.model_name,
|
||||
"deployment": deployment.litellm_params.model,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _managed_object_table(prisma_client: "PrismaClient") -> "TableActions[ManagedObjectRow]":
|
||||
return ManagedObjectRepository(prisma_client).table
|
||||
|
||||
|
||||
|
|
@ -51,12 +163,17 @@ def _is_response_gone_at_provider(error: Exception, provider_response_id: str) -
|
|||
return getattr(error, "status_code", None) == 404 and provider_response_id in str(error)
|
||||
|
||||
|
||||
def _keyed_by_issued_id(job: ManagedObjectRow) -> bool:
|
||||
return decrypted_response_id(job.model_object_id) != job.model_object_id
|
||||
|
||||
|
||||
class CheckResponsesCost:
|
||||
def __init__(
|
||||
self,
|
||||
proxy_logging_obj: "ProxyLogging",
|
||||
prisma_client: "PrismaClient",
|
||||
llm_router: "Router",
|
||||
log_billed_response: BilledResponseLogger = log_billed_background_response,
|
||||
):
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
||||
from litellm.router import Router
|
||||
|
|
@ -64,15 +181,17 @@ class CheckResponsesCost:
|
|||
self.proxy_logging_obj: ProxyLogging = proxy_logging_obj
|
||||
self.prisma_client: PrismaClient = prisma_client
|
||||
self.llm_router: Router = llm_router
|
||||
self._creator_attribution: Final = ManagedObjectCreatorAttribution(prisma_client, "CheckResponsesCost")
|
||||
self._log_billed_response: Final = log_billed_response
|
||||
|
||||
def _resolve_deployment(self, response_id: str) -> bool:
|
||||
def _resolve_deployment(self, response_id: str) -> Deployment | None:
|
||||
model_id: str | None = ResponsesAPIRequestUtils.get_model_id_from_response_id(response_id)
|
||||
return model_id is not None and self.llm_router.get_deployment(model_id=model_id) is not None
|
||||
return self.llm_router.get_deployment(model_id=model_id) if model_id is not None else None
|
||||
|
||||
async def _get_response(
|
||||
self,
|
||||
response_id: str,
|
||||
litellm_metadata: Dict[str, str],
|
||||
litellm_metadata: dict[str, object],
|
||||
via_router: bool,
|
||||
) -> ResponsesAPIResponse:
|
||||
"""Fetch the upstream response, using deployment credentials when available.
|
||||
|
|
@ -91,9 +210,17 @@ class CheckResponsesCost:
|
|||
)
|
||||
return cast(ResponsesAPIResponse, router_response)
|
||||
|
||||
async def _expire_stale_rows(
|
||||
self, cutoff: datetime, batch_size: int
|
||||
) -> int:
|
||||
async def _fetch_unbilled(self, target: _PollTarget) -> ResponsesAPIResponse:
|
||||
return await self._get_response(
|
||||
response_id=target.response_id,
|
||||
litellm_metadata={
|
||||
**target.model_metadata,
|
||||
INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN,
|
||||
},
|
||||
via_router=target.deployment is not None,
|
||||
)
|
||||
|
||||
async def _expire_stale_rows(self, cutoff: datetime, batch_size: int) -> int:
|
||||
"""Execute the bounded UPDATE that marks stale rows as 'stale_expired'.
|
||||
|
||||
Isolated so it can be swapped / mocked in tests without touching the
|
||||
|
|
@ -140,14 +267,154 @@ class CheckResponsesCost:
|
|||
f"(older than {MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days) as stale_expired"
|
||||
)
|
||||
|
||||
def _target(self, response_id: str, model_name: str | None) -> _PollTarget:
|
||||
return _PollTarget(
|
||||
response_id=response_id,
|
||||
model_metadata=MappingProxyType({"model": model_name, "model_group": model_name} if model_name else {}),
|
||||
deployment=self._resolve_deployment(response_id),
|
||||
)
|
||||
|
||||
def _poll_target(self, job: ManagedObjectRow) -> _PollTarget:
|
||||
stored_model: Final = _STORED_RESPONSE_ADAPTER.validate_python(job.file_object).get("model")
|
||||
return self._target(
|
||||
decrypted_response_id(job.unified_object_id), stored_model if isinstance(stored_model, str) else None
|
||||
)
|
||||
|
||||
async def _claim_for_billing(self, job: ManagedObjectRow) -> bool:
|
||||
claimed: Final = await _managed_object_table(self.prisma_client).update_many(
|
||||
# bounded-ok: PENDING_RESPONSE_STATUSES is a two-value constant tuple
|
||||
where={"id": job.id, "status": {"in": list(PENDING_RESPONSE_STATUSES)}},
|
||||
data={"status": "completed"},
|
||||
)
|
||||
return claimed > 0
|
||||
|
||||
async def _release_billing_claim(self, job: ManagedObjectRow) -> None:
|
||||
await _managed_object_table(self.prisma_client).update_many(
|
||||
where={"id": job.id, "status": "completed"},
|
||||
data={"status": job.status or PENDING_RESPONSE_STATUSES[0]},
|
||||
)
|
||||
|
||||
async def _billed_response(
|
||||
self,
|
||||
job: ManagedObjectRow,
|
||||
target: _PollTarget,
|
||||
response: ResponsesAPIResponse,
|
||||
) -> BilledBackgroundResponse:
|
||||
deployment: Final = target.deployment
|
||||
return BilledBackgroundResponse(
|
||||
response=response.model_copy(update={"id": job.unified_object_id}),
|
||||
model=deployment.litellm_params.model if deployment is not None else (response.model or ""),
|
||||
custom_llm_provider=ResponsesAPIRequestUtils.decode_responses_api_response_id(target.response_id).get(
|
||||
"custom_llm_provider"
|
||||
),
|
||||
litellm_params=_deployment_litellm_params(deployment) if deployment is not None else MappingProxyType({}),
|
||||
metadata=MappingProxyType(
|
||||
{
|
||||
**(await self._creator_attribution.response_metadata(job, job.unified_object_id)),
|
||||
INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN,
|
||||
**target.model_metadata,
|
||||
**(_deployment_metadata(deployment) if deployment is not None else {}),
|
||||
}
|
||||
),
|
||||
)
|
||||
|
||||
async def _bill_once(
|
||||
self,
|
||||
job: ManagedObjectRow,
|
||||
target: _PollTarget,
|
||||
response: ResponsesAPIResponse,
|
||||
) -> BilledCharge | None:
|
||||
if not await self._claim_for_billing(job):
|
||||
verbose_proxy_logger.info(f"Response {job.unified_object_id} was already billed elsewhere, skipping")
|
||||
return None
|
||||
try:
|
||||
charge: Final = await self._log_billed_response(await self._billed_response(job, target, response))
|
||||
except BaseException:
|
||||
await self._release_billing_claim(job)
|
||||
raise
|
||||
verbose_proxy_logger.info(f"Response {job.unified_object_id} reached a terminal status and was billed")
|
||||
return charge
|
||||
|
||||
async def _job_for_provider_id(self, provider_response_id: str) -> ManagedObjectRow | None:
|
||||
return await _managed_object_table(self.prisma_client).find_first(
|
||||
where={"model_object_id": provider_response_id, "file_purpose": "response"}
|
||||
)
|
||||
|
||||
async def _rekey_by_provider_id(self, job: ManagedObjectRow) -> None:
|
||||
await _managed_object_table(self.prisma_client).update_many(
|
||||
where={"id": job.id}, data={"model_object_id": decrypted_response_id(job.model_object_id)}
|
||||
)
|
||||
|
||||
async def _tracked_job(self, response_id: str) -> ManagedObjectRow | None:
|
||||
if tracked := await self._job_for_provider_id(decrypted_response_id(response_id)):
|
||||
return tracked
|
||||
stored_by_issued_id: Final = await _managed_object_table(self.prisma_client).find_first(
|
||||
where={"unified_object_id": response_id, "file_purpose": "response"}
|
||||
)
|
||||
if stored_by_issued_id is not None and _keyed_by_issued_id(stored_by_issued_id):
|
||||
await self._rekey_by_provider_id(stored_by_issued_id)
|
||||
return stored_by_issued_id
|
||||
|
||||
async def _superseded_by_a_provider_id_row(self, job: ManagedObjectRow) -> bool:
|
||||
if not _keyed_by_issued_id(job):
|
||||
return False
|
||||
if await self._job_for_provider_id(decrypted_response_id(job.model_object_id)) is None:
|
||||
await self._rekey_by_provider_id(job)
|
||||
return False
|
||||
await self._claim_for_billing(job)
|
||||
return True
|
||||
|
||||
async def _track_from_deployment(
|
||||
self, response_id: str, track_untracked: UntrackedResponseTracker | None
|
||||
) -> BilledCharge | None:
|
||||
target: Final = self._target(decrypted_response_id(response_id), None)
|
||||
if track_untracked is None or target.deployment is None:
|
||||
return None
|
||||
return await self.bill_if_finished(response_id, await self._fetch_unbilled(target), track_untracked)
|
||||
|
||||
async def _start_tracking(
|
||||
self, response_id: str, response: ResponsesAPIResponse, track_untracked: UntrackedResponseTracker | None
|
||||
) -> ManagedObjectRow | None:
|
||||
if track_untracked is None or not await track_untracked(response):
|
||||
return None
|
||||
return await self._tracked_job(response_id)
|
||||
|
||||
async def bill_if_finished(
|
||||
self,
|
||||
response_id: str,
|
||||
response: ResponsesAPIResponse,
|
||||
track_untracked: UntrackedResponseTracker | None = None,
|
||||
) -> BilledCharge | None:
|
||||
if response.status not in TERMINAL_RESPONSE_STATUSES:
|
||||
return None
|
||||
job: Final = await self._tracked_job(response_id) or await self._start_tracking(
|
||||
response_id, response, track_untracked
|
||||
)
|
||||
if job is None or job.status not in PENDING_RESPONSE_STATUSES:
|
||||
return None
|
||||
return await self._bill_once(job, self._poll_target(job), response)
|
||||
|
||||
async def bill_from_deployment(
|
||||
self, response_id: str, track_untracked: UntrackedResponseTracker | None = None
|
||||
) -> BilledCharge | None:
|
||||
job: Final = await self._tracked_job(response_id)
|
||||
if job is None:
|
||||
return await self._track_from_deployment(response_id, track_untracked)
|
||||
if job.status not in PENDING_RESPONSE_STATUSES:
|
||||
return None
|
||||
target: Final = self._poll_target(job)
|
||||
response: Final = await self._fetch_unbilled(target)
|
||||
if response.status not in TERMINAL_RESPONSE_STATUSES:
|
||||
return None
|
||||
return await self._bill_once(job, target, response)
|
||||
|
||||
async def check_responses_cost(self):
|
||||
"""
|
||||
Check if background responses are complete and track their cost.
|
||||
- Get all status="queued" or "in_progress" and file_purpose="response" jobs
|
||||
- Query the provider to check if response is complete
|
||||
- Cost is tracked by the get-responses call, billed because the poll is stamped
|
||||
with BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN
|
||||
- Mark responses in a terminal state as complete in the database
|
||||
- Query the provider, unbilled, to check if the response is complete
|
||||
- Claim each terminal job by moving its row to completed, then log the body that check fetched
|
||||
as its one billed event, so two pollers racing over the same job bill it once
|
||||
- Mark responses the provider no longer has (404 through a resolved
|
||||
deployment) as stale_expired
|
||||
"""
|
||||
|
|
@ -160,99 +427,53 @@ class CheckResponsesCost:
|
|||
|
||||
jobs = await _managed_object_table(self.prisma_client).find_many(
|
||||
where={
|
||||
"status": {"in": ["queued", "in_progress"]},
|
||||
# bounded-ok: PENDING_RESPONSE_STATUSES is a two-value constant tuple
|
||||
"status": {"in": list(PENDING_RESPONSE_STATUSES)},
|
||||
"file_purpose": "response",
|
||||
},
|
||||
take=MAX_OBJECTS_PER_POLL_CYCLE,
|
||||
order={"created_at": "asc"},
|
||||
)
|
||||
|
||||
|
||||
verbose_proxy_logger.debug(f"Found {len(jobs)} response jobs to check")
|
||||
completed_jobs: Final[list[_ManagedObjectRow]] = []
|
||||
expired_jobs: Final[list[_ManagedObjectRow]] = []
|
||||
expired_jobs: Final[list[ManagedObjectRow]] = []
|
||||
|
||||
for job in jobs:
|
||||
unified_object_id = job.unified_object_id
|
||||
|
||||
try:
|
||||
from litellm.proxy.hooks.responses_id_security import (
|
||||
ResponsesIDSecurity,
|
||||
)
|
||||
|
||||
# Get the stored response object to extract model information
|
||||
stored_response = job.file_object
|
||||
model_name = stored_response.get("model", None)
|
||||
|
||||
# Decrypt the response ID
|
||||
responses_id_security, _, _ = ResponsesIDSecurity()._decrypt_response_id(unified_object_id)
|
||||
|
||||
# Prepare metadata with model information for cost tracking
|
||||
litellm_metadata = {
|
||||
"user_api_key_user_id": job.created_by or "default-user-id",
|
||||
INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN,
|
||||
}
|
||||
|
||||
# Add model information if available
|
||||
if model_name:
|
||||
litellm_metadata["model"] = model_name
|
||||
litellm_metadata["model_group"] = model_name # Use same value for model_group
|
||||
via_router = self._resolve_deployment(responses_id_security)
|
||||
if await self._superseded_by_a_provider_id_row(job):
|
||||
verbose_proxy_logger.info(f"Response {unified_object_id} is billed under its provider id row")
|
||||
continue
|
||||
target = self._poll_target(job)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.warning(
|
||||
f"Skipping job {unified_object_id} due to error: {e}"
|
||||
)
|
||||
verbose_proxy_logger.warning(f"Skipping job {unified_object_id} due to error: {e}")
|
||||
continue
|
||||
|
||||
provider_response_id = ResponsesAPIRequestUtils.decode_responses_api_response_id(
|
||||
responses_id_security
|
||||
).get("response_id", responses_id_security)
|
||||
provider_response_id = ResponsesAPIRequestUtils.decode_responses_api_response_id(target.response_id).get(
|
||||
"response_id", target.response_id
|
||||
)
|
||||
try:
|
||||
response = await self._get_response(
|
||||
response_id=responses_id_security,
|
||||
litellm_metadata=litellm_metadata,
|
||||
via_router=via_router,
|
||||
)
|
||||
response = await self._fetch_unbilled(target)
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
f"Response {unified_object_id} status: {response.status}, model: {model_name}"
|
||||
)
|
||||
verbose_proxy_logger.debug(f"Response {unified_object_id} status: {response.status}")
|
||||
if response.status in TERMINAL_RESPONSE_STATUSES:
|
||||
await self._bill_once(job, target, response)
|
||||
|
||||
except Exception as e:
|
||||
if via_router and _is_response_gone_at_provider(e, provider_response_id):
|
||||
if target.deployment is not None and _is_response_gone_at_provider(e, provider_response_id):
|
||||
verbose_proxy_logger.info(
|
||||
f"Response {unified_object_id} no longer available at provider (404), marking stale_expired: {e}"
|
||||
)
|
||||
expired_jobs.append(job)
|
||||
continue
|
||||
verbose_proxy_logger.warning(
|
||||
f"Skipping job {unified_object_id} due to error: {e}"
|
||||
)
|
||||
verbose_proxy_logger.warning(f"Skipping job {unified_object_id} due to error: {e}")
|
||||
continue
|
||||
|
||||
if response.status in TERMINAL_RESPONSE_STATUSES:
|
||||
verbose_proxy_logger.info(
|
||||
f"Response {unified_object_id} has terminal status {response.status}, marking as complete"
|
||||
)
|
||||
completed_jobs.append(job)
|
||||
|
||||
# Mark completed jobs in the database
|
||||
if len(completed_jobs) > 0:
|
||||
await _managed_object_table(self.prisma_client).update_many(
|
||||
# bounded-ok: at most MAX_OBJECTS_PER_POLL_CYCLE rows per cycle, the find_many take above
|
||||
where={"id": {"in": [job.id for job in completed_jobs]}},
|
||||
data={"status": "completed"},
|
||||
)
|
||||
verbose_proxy_logger.info(
|
||||
f"Marked {len(completed_jobs)} response jobs as completed"
|
||||
)
|
||||
|
||||
if len(expired_jobs) > 0:
|
||||
await _managed_object_table(self.prisma_client).update_many(
|
||||
# bounded-ok: at most MAX_OBJECTS_PER_POLL_CYCLE rows per cycle, the find_many take above
|
||||
where={"id": {"in": [job.id for job in expired_jobs]}},
|
||||
data={"status": "stale_expired"},
|
||||
)
|
||||
verbose_proxy_logger.info(
|
||||
f"Marked {len(expired_jobs)} response jobs as stale_expired"
|
||||
)
|
||||
|
||||
verbose_proxy_logger.info(f"Marked {len(expired_jobs)} response jobs as stale_expired")
|
||||
|
|
|
|||
|
|
@ -0,0 +1,218 @@
|
|||
"""Creator attribution for the spend row a managed-object cost poll writes.
|
||||
|
||||
A batch or a background response is billed by a poll job that runs long after the creating
|
||||
request returned, so the managed object row's creator columns are the only record of which
|
||||
key, team, organization, user, end user, and tags carry the spend.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Final, Protocol, cast
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import CLI_SESSION_KEY_PREFIX
|
||||
from litellm.repositories.team_repository import TeamRepository
|
||||
from litellm.repositories.user_repository import UserRepository
|
||||
from litellm.repositories.verification_token_repository import VerificationTokenRepository
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from prisma import models as prisma_models
|
||||
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
from litellm.repositories.prisma_protocols import TableActions
|
||||
|
||||
|
||||
class ManagedObjectRow(Protocol):
|
||||
@property
|
||||
def id(self) -> str: ...
|
||||
|
||||
@property
|
||||
def unified_object_id(self) -> str: ...
|
||||
|
||||
@property
|
||||
def model_object_id(self) -> str: ...
|
||||
|
||||
@property
|
||||
def created_by(self) -> str | None: ...
|
||||
|
||||
@property
|
||||
def status(self) -> str | None: ...
|
||||
|
||||
@property
|
||||
def file_object(self) -> object: ...
|
||||
|
||||
@property
|
||||
def org_id(self) -> str | None: ...
|
||||
|
||||
@property
|
||||
def api_key(self) -> str | None: ...
|
||||
|
||||
@property
|
||||
def team_id(self) -> str | None: ...
|
||||
|
||||
@property
|
||||
def request_tags(self) -> object: ...
|
||||
|
||||
@property
|
||||
def end_user_id(self) -> str | None: ...
|
||||
|
||||
|
||||
_TAG_LIST_ADAPTER: Final[TypeAdapter[list[object]]] = TypeAdapter(list[object])
|
||||
_KEY_OWNER_SCOPE_PREFIX: Final = "key:"
|
||||
|
||||
|
||||
def _string_tags(request_tags: object) -> list[str] | None:
|
||||
if not isinstance(request_tags, list):
|
||||
return None
|
||||
return [tag for tag in _TAG_LIST_ADAPTER.validate_python(request_tags) if isinstance(tag, str)] or None
|
||||
|
||||
|
||||
def _creator_user_id(created_by: str | None) -> str | None:
|
||||
if created_by is None or created_by.startswith(_KEY_OWNER_SCOPE_PREFIX):
|
||||
return None
|
||||
return created_by
|
||||
|
||||
|
||||
def _user_table(prisma_client: PrismaClient) -> TableActions[prisma_models.LiteLLM_UserTable]:
|
||||
return UserRepository(prisma_client).table
|
||||
|
||||
|
||||
def _token_table(prisma_client: PrismaClient) -> TableActions[prisma_models.LiteLLM_VerificationToken]:
|
||||
return VerificationTokenRepository(prisma_client).table
|
||||
|
||||
|
||||
def _team_table(prisma_client: PrismaClient) -> TableActions[prisma_models.LiteLLM_TeamTable]:
|
||||
return TeamRepository(prisma_client).table
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _CreatorRows:
|
||||
key: prisma_models.LiteLLM_VerificationToken | None
|
||||
team: prisma_models.LiteLLM_TeamTable | None
|
||||
user: prisma_models.LiteLLM_UserTable | None
|
||||
|
||||
|
||||
class ManagedObjectCreatorAttribution:
|
||||
"""Rebuilds the spend-tracking metadata of whoever created a managed object, so the poll
|
||||
job's spend row is attributed like the creating request: key hash, team, organization,
|
||||
user, aliases, and tags.
|
||||
|
||||
Rows created before api_key and request_tags were persisted carry only created_by and
|
||||
team_id, and fall back to those. A named creating key owns user_api_key_alias; when it
|
||||
has no alias, or the key has since been rotated or deleted, the field keeps the creating
|
||||
user's alias, because a resolvable name is more useful on the spend row than a null.
|
||||
user_api_key_org_id is resolved too: the spend update writer reads it off this metadata
|
||||
to increment organization spend, so leaving it out drops the cost from org accounting.
|
||||
A background response also carries its end user and the key, team, and user budgets, the
|
||||
way the read that billed it before did, and never names a key-scoped creator as its user.
|
||||
"""
|
||||
|
||||
def __init__(self, prisma_client: PrismaClient, poller_name: str) -> None:
|
||||
self._prisma_client: Final = prisma_client
|
||||
self._poller_name: Final = poller_name
|
||||
|
||||
async def batch_metadata(self, job: ManagedObjectRow, object_id: str) -> Mapping[str, object]:
|
||||
rows: Final = await self._rows(job, object_id)
|
||||
return {"user_api_key_user_id": job.created_by, **_shared_metadata(job, rows)}
|
||||
|
||||
async def response_metadata(self, job: ManagedObjectRow, object_id: str) -> Mapping[str, object]:
|
||||
rows: Final = await self._rows(job, object_id)
|
||||
end_user_id: Final = getattr(job, "end_user_id", None)
|
||||
return {
|
||||
"user_api_key_user_id": _creator_user_id(job.created_by),
|
||||
**_shared_metadata(job, rows),
|
||||
**({"user_api_key_end_user_id": end_user_id} if end_user_id else {}),
|
||||
**_budget_fields(rows.key, "user_api_key_spend", "user_api_key_max_budget"),
|
||||
**_budget_fields(rows.team, "user_api_key_team_spend", "user_api_key_team_max_budget"),
|
||||
**_budget_fields(rows.user, "user_api_key_user_spend", "user_api_key_user_max_budget"),
|
||||
}
|
||||
|
||||
async def _rows(self, job: ManagedObjectRow, object_id: str) -> _CreatorRows:
|
||||
return _CreatorRows(
|
||||
key=await self._key_row(object_id, getattr(job, "api_key", None), job.created_by),
|
||||
team=await self._team_row(object_id, getattr(job, "team_id", None)),
|
||||
user=await self._user_row(object_id, _creator_user_id(job.created_by)),
|
||||
)
|
||||
|
||||
async def _user_row(self, object_id: str, user_id: str | None) -> prisma_models.LiteLLM_UserTable | None:
|
||||
"""Objects created by a team or service account key carry no user id, and
|
||||
find_unique(where={"user_id": None}) raises, so those resolve no row."""
|
||||
if not user_id:
|
||||
return None
|
||||
try:
|
||||
return await _user_table(self._prisma_client).find_unique(where={"user_id": user_id})
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error(f"{self._poller_name}: could not look up user {user_id} for {object_id}: {e}")
|
||||
return None
|
||||
|
||||
async def _key_row(
|
||||
self, object_id: str, api_key: str | None, created_by: str | None
|
||||
) -> prisma_models.LiteLLM_VerificationToken | None:
|
||||
if not api_key or _is_cli_session_key(api_key, created_by):
|
||||
return None
|
||||
try:
|
||||
return await _token_table(self._prisma_client).find_unique(where={"token": api_key})
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error(f"{self._poller_name}: could not look up the creating key for {object_id}: {e}")
|
||||
return None
|
||||
|
||||
async def _team_row(self, object_id: str, team_id: str | None) -> prisma_models.LiteLLM_TeamTable | None:
|
||||
if not team_id:
|
||||
return None
|
||||
try:
|
||||
return await _team_table(self._prisma_client).find_unique(where={"team_id": team_id})
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error(f"{self._poller_name}: could not look up team {team_id} for {object_id}: {e}")
|
||||
return None
|
||||
|
||||
|
||||
def _shared_metadata(job: ManagedObjectRow, rows: _CreatorRows) -> Mapping[str, object]:
|
||||
api_key: Final = getattr(job, "api_key", None)
|
||||
team_id: Final = getattr(job, "team_id", None)
|
||||
key_alias: Final = _key_alias(api_key, job.created_by, rows.key)
|
||||
team_alias: Final = getattr(rows.team, "team_alias", None)
|
||||
org_id: Final = getattr(job, "org_id", None) or _organization_id(rows.key) or _organization_id(rows.team)
|
||||
tags: Final = _string_tags(getattr(job, "request_tags", None))
|
||||
return {
|
||||
"user_api_key": api_key,
|
||||
"user_api_key_hash": api_key,
|
||||
"user_api_key_team_id": team_id,
|
||||
**(
|
||||
{
|
||||
"user_api_key_user_email": getattr(rows.user, "user_email", None),
|
||||
"user_api_key_alias": getattr(rows.user, "user_alias", None),
|
||||
}
|
||||
if rows.user is not None
|
||||
else {}
|
||||
),
|
||||
**({"user_api_key_alias": key_alias} if key_alias is not None else {}),
|
||||
**({"user_api_key_team_alias": team_alias} if team_alias is not None else {}),
|
||||
**({"user_api_key_org_id": org_id} if org_id is not None else {}),
|
||||
**({"tags": tags} if tags else {}),
|
||||
}
|
||||
|
||||
|
||||
def _is_cli_session_key(api_key: str, created_by: str | None) -> bool:
|
||||
return bool(created_by) and api_key == f"{CLI_SESSION_KEY_PREFIX}-{created_by}"
|
||||
|
||||
|
||||
def _key_alias(
|
||||
api_key: str | None, created_by: str | None, key_row: prisma_models.LiteLLM_VerificationToken | None
|
||||
) -> str | None:
|
||||
if api_key and _is_cli_session_key(api_key, created_by):
|
||||
return api_key
|
||||
return getattr(key_row, "key_alias", None) if key_row is not None else None
|
||||
|
||||
|
||||
def _organization_id(row: object) -> str | None:
|
||||
return cast("str | None", getattr(row, "organization_id", None)) if row is not None else None
|
||||
|
||||
|
||||
def _budget_fields(row: object, spend_key: str, max_budget_key: str) -> Mapping[str, object]:
|
||||
if row is None:
|
||||
return {}
|
||||
return {spend_key: getattr(row, "spend", None), max_budget_key: getattr(row, "max_budget", None)}
|
||||
|
|
@ -24,7 +24,9 @@ from uuid import NAMESPACE_URL, uuid5
|
|||
|
||||
import httpx
|
||||
from fastapi import HTTPException
|
||||
from pydantic import ValidationError
|
||||
from openai import APIConnectionError, AsyncOpenAI
|
||||
from openai.types.file_deleted import FileDeleted
|
||||
from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError
|
||||
from typing_extensions import ReadOnly, Unpack
|
||||
|
||||
import litellm
|
||||
|
|
@ -35,17 +37,17 @@ from litellm.caching.caching import DualCache
|
|||
from litellm.constants import (
|
||||
BATCH_OUTPUT_FILE_FALLBACK_RETRY_AFTER_SECONDS,
|
||||
BATCH_OUTPUT_FILE_LOOKUP_TIMEOUT_SECONDS,
|
||||
MANAGED_BACKGROUND_RESPONSE_METADATA_KEY,
|
||||
MAX_FILE_LIST_LIMIT,
|
||||
PENDING_RESPONSE_STATUSES,
|
||||
)
|
||||
from litellm.files.types import FileRetrieveCallOptions, FileRetrieveProvider
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR
|
||||
from litellm.litellm_core_utils.internal_call_metadata import requests_background
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
extract_file_metadata,
|
||||
)
|
||||
from openai import APIConnectionError, AsyncOpenAI
|
||||
from openai.types.file_deleted import FileDeleted
|
||||
|
||||
from litellm.llms.base_llm.files.storage_backend_factory import get_storage_backend
|
||||
from litellm.llms.base_llm.files.transformation import BaseFileEndpoints
|
||||
from litellm.llms.base_llm.managed_resources.isolation import (
|
||||
|
|
@ -54,6 +56,7 @@ from litellm.llms.base_llm.managed_resources.isolation import (
|
|||
can_access_resource,
|
||||
resolve_resource_owner_id,
|
||||
)
|
||||
from litellm.llms.xai.batches.transformation import is_xai_batch_results_id
|
||||
from litellm.proxy._types import (
|
||||
CallTypes,
|
||||
LiteLLM_ManagedFileTable,
|
||||
|
|
@ -87,6 +90,7 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attributio
|
|||
request_tags_from_metadata,
|
||||
)
|
||||
from litellm.repositories.managed_file_repository import ManagedFileRepository
|
||||
from litellm.responses.utils import ResponsesAPIRequestUtils
|
||||
from litellm.types.llms.openai import ( # pyright: ignore[reportAttributeAccessIssue]
|
||||
AllMessageValues,
|
||||
AsyncCursorPage,
|
||||
|
|
@ -99,9 +103,11 @@ from litellm.types.llms.openai import ( # pyright: ignore[reportAttributeAccess
|
|||
ResponsesAPIResponse,
|
||||
)
|
||||
from litellm.types.utils import (
|
||||
ADDRESSED_RESPONSE_ID_FIELD,
|
||||
CallTypesLiteral,
|
||||
LiteLLMBatch,
|
||||
LiteLLMFineTuningJob,
|
||||
LlmProviders,
|
||||
LLMResponseTypes,
|
||||
SpecialEnums,
|
||||
)
|
||||
|
|
@ -194,6 +200,32 @@ def _is_transient_file_retrieve_error(error: Exception) -> bool:
|
|||
return isinstance(error, (httpx.TransportError, APIConnectionError, asyncio.TimeoutError))
|
||||
|
||||
|
||||
_MEASURED_SIZE_SAVE_ATTEMPTS: Final = 3
|
||||
|
||||
|
||||
def _stored_json_text(raw_file_object: object) -> str:
|
||||
"""JSON text of a stored file object exactly as read, for a write that must not race a newer one."""
|
||||
if isinstance(raw_file_object, BaseModel):
|
||||
return raw_file_object.model_dump_json()
|
||||
return json.dumps(raw_file_object)
|
||||
|
||||
|
||||
def _with_measured_size(file_object: OpenAIFileObject, size_bytes: int | None) -> OpenAIFileObject:
|
||||
"""Fill a file object whose provider reported no size (0 bytes) with a size measured elsewhere."""
|
||||
if not size_bytes or file_object.bytes:
|
||||
return file_object
|
||||
return file_object.model_copy(update={"bytes": size_bytes})
|
||||
|
||||
|
||||
def _over_stored_object(file_object: OpenAIFileObject, stored_object: OpenAIFileObject | None) -> OpenAIFileObject:
|
||||
"""Refreshed details for a stored row, keeping its size when none is reported and its stored created_at."""
|
||||
if stored_object is None:
|
||||
return file_object
|
||||
return _with_measured_size(file_object, stored_object.bytes).model_copy(
|
||||
update={"created_at": stored_object.created_at}
|
||||
)
|
||||
|
||||
|
||||
def _proxy_llm_router() -> Router | None:
|
||||
import litellm.proxy.proxy_server as proxy_server_module
|
||||
|
||||
|
|
@ -252,6 +284,23 @@ def _has_provider_file_retrieve_route(
|
|||
return bool(router_credentials) or _model_name_file_retrieve_provider(model_name) is not None
|
||||
|
||||
|
||||
def _is_xai_batch_output(
|
||||
*,
|
||||
llm_router: Router | None,
|
||||
model_id: str | None,
|
||||
model_name: str | None,
|
||||
provider_file_id: str,
|
||||
) -> bool:
|
||||
"""Whether a batch output row holds an xAI batch's results, from its batch id and its deployment's provider."""
|
||||
if not is_xai_batch_results_id(provider_file_id):
|
||||
return False
|
||||
router_credentials: Final = _provider_file_retrieve_credentials(llm_router=llm_router, model_id=model_id)
|
||||
provider: Final = (router_credentials or {}).get("custom_llm_provider") or _model_name_file_retrieve_provider(
|
||||
model_name
|
||||
)
|
||||
return provider == LlmProviders.XAI.value
|
||||
|
||||
|
||||
class _ManagedFileRow(Protocol):
|
||||
unified_file_id: str
|
||||
file_object: OpenAIFileObject
|
||||
|
|
@ -301,6 +350,8 @@ class _ManagedObjectTableActions(Protocol):
|
|||
|
||||
async def update_many(self, where: Mapping[str, object], data: Mapping[str, object]) -> int: ...
|
||||
|
||||
async def create(self, data: Mapping[str, object]) -> "PrismaManagedObjectRow": ...
|
||||
|
||||
|
||||
class _ManagedResourceDatabase(Protocol):
|
||||
@property
|
||||
|
|
@ -347,6 +398,39 @@ def _managed_object_table(prisma_client: _ManagedResourcePrismaClient) -> _Manag
|
|||
return prisma_client.db.litellm_managedobjecttable
|
||||
|
||||
|
||||
_JSON_VALUE_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
|
||||
_POSTGRES_UNSTORABLE_CHARACTER: Final = "\x00"
|
||||
|
||||
|
||||
def _storable_text(text: str) -> str:
|
||||
return text.replace(_POSTGRES_UNSTORABLE_CHARACTER, "")
|
||||
|
||||
|
||||
def _storable_json_value(value: JsonValue) -> JsonValue:
|
||||
if isinstance(value, str):
|
||||
return _storable_text(value)
|
||||
if isinstance(value, list):
|
||||
return [_storable_json_value(item) for item in value]
|
||||
if isinstance(value, dict):
|
||||
return {_storable_text(key): _storable_json_value(item) for key, item in value.items()}
|
||||
return value
|
||||
|
||||
|
||||
def _postgres_storable_json(file_object: BaseModel) -> str:
|
||||
dumped: Final = file_object.model_dump_json()
|
||||
if "\\u0000" not in dumped:
|
||||
return dumped
|
||||
return _JSON_VALUE_ADAPTER.dump_json(_storable_json_value(_JSON_VALUE_ADAPTER.validate_json(dumped))).decode()
|
||||
|
||||
|
||||
def _without_model_output(
|
||||
file_object: LiteLLMBatch | LiteLLMFineTuningJob | ResponsesAPIResponse,
|
||||
) -> LiteLLMBatch | LiteLLMFineTuningJob | ResponsesAPIResponse:
|
||||
return (
|
||||
file_object.model_copy(update={"output": []}) if isinstance(file_object, ResponsesAPIResponse) else file_object
|
||||
)
|
||||
|
||||
|
||||
def _storage_metadata_of(file_object: OpenAIFileObject | None) -> Mapping[str, str]:
|
||||
hidden_params: Final = cast( # cast-ok: _hidden_params is an untyped attribute the upload path sets
|
||||
"Mapping[str, object]", getattr(file_object, "_hidden_params", None) or {}
|
||||
|
|
@ -360,6 +444,40 @@ def _storage_metadata_of(file_object: OpenAIFileObject | None) -> Mapping[str, s
|
|||
)
|
||||
|
||||
|
||||
_MANAGED_RESPONSE_READ_CALL_TYPES: Final = frozenset({"aget_responses", "acancel_responses"})
|
||||
_REQUEST_FIELDS: Final = TypeAdapter(dict[str, object])
|
||||
|
||||
|
||||
def _addresses_a_proxy_issued_response(data: Mapping[str, object]) -> bool:
|
||||
from litellm.proxy.hooks.responses_id_security import decrypted_response_id
|
||||
|
||||
addressed_id: Final = data.get(ADDRESSED_RESPONSE_ID_FIELD) or data.get("response_id")
|
||||
return (
|
||||
isinstance(addressed_id, str)
|
||||
and ResponsesAPIRequestUtils.get_model_id_from_response_id(decrypted_response_id(addressed_id)) is not None
|
||||
)
|
||||
|
||||
|
||||
def _is_streamed_background_create(request: Mapping[str, object]) -> bool:
|
||||
return (
|
||||
request.get("stream") is True
|
||||
and isinstance(request.get("litellm_metadata"), dict)
|
||||
and requests_background(request)
|
||||
)
|
||||
|
||||
|
||||
def _with_managed_background_marker(request: Mapping[str, object]) -> dict[str, object]:
|
||||
metadata_field: Final = "litellm_metadata" if isinstance(request.get("litellm_metadata"), dict) else "metadata"
|
||||
existing_metadata: Final = request.get(metadata_field)
|
||||
return {
|
||||
**request,
|
||||
metadata_field: {
|
||||
**(_REQUEST_FIELDS.validate_python(existing_metadata) if isinstance(existing_metadata, dict) else {}),
|
||||
MANAGED_BACKGROUND_RESPONSE_METADATA_KEY: True,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
_MANAGED_FILES_TARGET: Final = "managed_files"
|
||||
_PROVIDER_FILE_RETRIEVE_RETRY_DELAYS_SECONDS: Final = (0.5, 1.0, 2.0)
|
||||
|
||||
|
|
@ -484,11 +602,14 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
) -> None:
|
||||
"""Persist a managed object row, caching it and upserting it in the DB.
|
||||
|
||||
persist_attribution is set only by the batch create, which is the one caller
|
||||
that can speak for the creator; it gates the api_key and request_tags columns
|
||||
that CheckBatchCost bills against, so a later poll or retrieve of the same
|
||||
batch cannot record itself as the paying key. Like created_by and team_id,
|
||||
both are written only in the upsert create branch, never on update.
|
||||
persist_attribution is set only by the batch and background response creates,
|
||||
the callers that can speak for the creator; it gates the api_key, request_tags, and,
|
||||
for a response, end_user_id columns that the cost poll jobs bill against, so a later poll or
|
||||
retrieve of the same object cannot record itself as the paying key. Like created_by and team_id,
|
||||
they are written only in the upsert create branch, never on update.
|
||||
|
||||
A stored response keeps no model output: billing reads only its status and model from
|
||||
the row, so the row never holds generated content.
|
||||
|
||||
create_if_missing is cleared by callers that observe a batch they did not
|
||||
create, such as a poll. They still refresh status and file_object, but a
|
||||
|
|
@ -501,11 +622,12 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
It is written only in the upsert create branch.
|
||||
"""
|
||||
verbose_logger.info(f"Storing LiteLLM Managed {file_purpose} object with id={unified_object_id} in cache")
|
||||
stored_object: Final = _without_model_output(file_object)
|
||||
litellm_managed_object = LiteLLM_ManagedObjectTable(
|
||||
unified_object_id=unified_object_id,
|
||||
model_object_id=model_object_id,
|
||||
file_purpose=file_purpose,
|
||||
file_object=file_object,
|
||||
file_object=stored_object,
|
||||
)
|
||||
await self.internal_usage_cache.async_set_cache(
|
||||
key=unified_object_id,
|
||||
|
|
@ -516,32 +638,35 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
from prisma import Json
|
||||
|
||||
api_key = LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict) or None
|
||||
end_user_id: Final = user_api_key_dict.end_user_id if file_purpose == "response" else None
|
||||
attribution_columns = (
|
||||
{
|
||||
**({"api_key": api_key} if api_key is not None else {}),
|
||||
**({"request_tags": Json(list(request_tags))} if request_tags else {}),
|
||||
**({"request_tags": Json([_storable_text(tag) for tag in request_tags])} if request_tags else {}),
|
||||
**({"end_user_id": end_user_id} if end_user_id else {}),
|
||||
}
|
||||
if persist_attribution
|
||||
else {}
|
||||
)
|
||||
# FIX: Update status and file_object on every operation to keep state in sync
|
||||
update_columns: Final = {
|
||||
"file_object": file_object.model_dump_json(),
|
||||
"file_object": _postgres_storable_json(stored_object),
|
||||
"status": file_object.status,
|
||||
"updated_by": user_api_key_dict.user_id,
|
||||
}
|
||||
managed_object_table: Final = _managed_object_table(self.prisma_client)
|
||||
if not create_if_missing:
|
||||
await _managed_object_table(self.prisma_client).update_many(
|
||||
await managed_object_table.update_many(
|
||||
where={"unified_object_id": unified_object_id},
|
||||
data=update_columns,
|
||||
)
|
||||
return
|
||||
await _managed_object_table(self.prisma_client).upsert(
|
||||
await managed_object_table.upsert(
|
||||
where={"unified_object_id": unified_object_id},
|
||||
data={
|
||||
"create": {
|
||||
"unified_object_id": unified_object_id,
|
||||
"file_object": file_object.model_dump_json(),
|
||||
"file_object": _postgres_storable_json(stored_object),
|
||||
"model_object_id": model_object_id,
|
||||
"file_purpose": file_purpose,
|
||||
"created_by": resolve_resource_owner_id(user_api_key_dict),
|
||||
|
|
@ -556,6 +681,42 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
},
|
||||
)
|
||||
|
||||
async def track_untracked_background_response(
|
||||
self,
|
||||
response_id: str,
|
||||
response: ResponsesAPIResponse,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
request_tags: Sequence[str] | None = None,
|
||||
) -> bool:
|
||||
if not _addresses_a_proxy_issued_response({"response_id": response_id}):
|
||||
return False
|
||||
from prisma import Json
|
||||
from prisma.errors import UniqueViolationError
|
||||
|
||||
from litellm.proxy.hooks.responses_id_security import decrypted_response_id
|
||||
|
||||
api_key: Final = LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict) or None
|
||||
try:
|
||||
await _managed_object_table(self.prisma_client).create(
|
||||
data={
|
||||
"unified_object_id": response_id,
|
||||
"model_object_id": decrypted_response_id(response_id),
|
||||
"file_purpose": "response",
|
||||
"file_object": _postgres_storable_json(_without_model_output(response)),
|
||||
"status": PENDING_RESPONSE_STATUSES[0],
|
||||
"created_by": resolve_resource_owner_id(user_api_key_dict),
|
||||
"team_id": user_api_key_dict.team_id,
|
||||
"org_id": await self._resolve_creator_org_id(user_api_key_dict),
|
||||
"updated_by": user_api_key_dict.user_id,
|
||||
**({"api_key": api_key} if api_key is not None else {}),
|
||||
**({"end_user_id": user_api_key_dict.end_user_id} if user_api_key_dict.end_user_id else {}),
|
||||
**({"request_tags": Json([_storable_text(tag) for tag in request_tags])} if request_tags else {}),
|
||||
}
|
||||
)
|
||||
except UniqueViolationError:
|
||||
return False
|
||||
return True
|
||||
|
||||
@with_service_target(_MANAGED_FILES_TARGET)
|
||||
async def get_unified_file_id(
|
||||
self, file_id: str, litellm_parent_otel_span: Optional[Span] = None
|
||||
|
|
@ -845,6 +1006,57 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
litellm_parent_otel_span=None,
|
||||
)
|
||||
|
||||
async def _save_xai_refreshed_file_object(
|
||||
self,
|
||||
stored: LiteLLM_ManagedFileTable,
|
||||
file_object: OpenAIFileObject,
|
||||
) -> OpenAIFileObject | None:
|
||||
"""Save an xAI batch output's refreshed details, keeping a known size, only over the row as read; drop the cached copy if a newer write won."""
|
||||
stored_object: Final = _parse_managed_file_object(stored.file_object, stored.unified_file_id)
|
||||
sized_file_object: Final = _over_stored_object(file_object, stored_object)
|
||||
if not await ManagedFileRepository(self.prisma_client).update_file_object(
|
||||
stored.unified_file_id,
|
||||
sized_file_object,
|
||||
if_file_object=_stored_json_text(stored.file_object) if stored.file_object is not None else None,
|
||||
):
|
||||
await self.internal_usage_cache.async_set_cache(
|
||||
key=stored.unified_file_id, value=None, litellm_parent_otel_span=None
|
||||
)
|
||||
return None
|
||||
await self.internal_usage_cache.async_set_cache(
|
||||
key=stored.unified_file_id,
|
||||
value=stored.model_copy(update={"file_object": sized_file_object}).model_dump(),
|
||||
litellm_parent_otel_span=None,
|
||||
)
|
||||
return sized_file_object
|
||||
|
||||
async def _save_batch_output_file_object(
|
||||
self,
|
||||
stored: LiteLLM_ManagedFileTable,
|
||||
file_object: OpenAIFileObject,
|
||||
size_bytes: int | None,
|
||||
) -> None:
|
||||
"""Save an xAI batch output's details; after losing to a refresh, size the refreshed details instead of replacing them."""
|
||||
current_row: LiteLLM_ManagedFileTable | None = stored # rebind-ok: the row as last read
|
||||
details: OpenAIFileObject = file_object # rebind-ok: a concurrent refresh's details replace the job's copy
|
||||
for _ in range(_MEASURED_SIZE_SAVE_ATTEMPTS):
|
||||
if current_row is None:
|
||||
return
|
||||
saved = await self._save_xai_refreshed_file_object( # rebind-ok: per attempt
|
||||
current_row, _with_measured_size(details, size_bytes)
|
||||
)
|
||||
if saved is not None or not size_bytes:
|
||||
return
|
||||
latest_row = await ManagedFileRepository( # rebind-ok: per attempt
|
||||
self.prisma_client, use_writer=True
|
||||
).table.find_first(where={"unified_file_id": stored.unified_file_id})
|
||||
current_row = (
|
||||
None if latest_row is None else LiteLLM_ManagedFileTable.model_validate(latest_row.model_dump())
|
||||
)
|
||||
latest_object = None if current_row is None else current_row.file_object # rebind-ok: per attempt
|
||||
if latest_object is not None and not latest_object.litellm_details_fallback:
|
||||
details = latest_object
|
||||
|
||||
async def store_batch_output_file(
|
||||
self,
|
||||
*,
|
||||
|
|
@ -885,6 +1097,18 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
)
|
||||
return
|
||||
if stored_object is not None and not stored_object.litellm_details_fallback:
|
||||
if (
|
||||
stored_file is not None
|
||||
and size_bytes
|
||||
and not stored_object.bytes
|
||||
and _is_xai_batch_output(
|
||||
llm_router=_proxy_llm_router(),
|
||||
model_id=model_id,
|
||||
model_name=model_name,
|
||||
provider_file_id=provider_file_id,
|
||||
)
|
||||
):
|
||||
await self._save_batch_output_file_object(stored_file, stored_object, size_bytes)
|
||||
return
|
||||
|
||||
fallback_written_recently: Final = (
|
||||
|
|
@ -924,13 +1148,22 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
)
|
||||
)
|
||||
|
||||
xai_batch_output: Final = _is_xai_batch_output(
|
||||
llm_router=_proxy_llm_router(),
|
||||
model_id=model_id,
|
||||
model_name=model_name,
|
||||
provider_file_id=provider_file_id,
|
||||
)
|
||||
if stored_file is not None and xai_batch_output:
|
||||
await self._save_batch_output_file_object(stored_file, file_object, size_bytes)
|
||||
return
|
||||
if stored_file is not None:
|
||||
await self._save_refreshed_file_object(stored_file, file_object)
|
||||
return
|
||||
|
||||
await self.store_unified_file_id(
|
||||
file_id=unified_file_id,
|
||||
file_object=file_object,
|
||||
file_object=_with_measured_size(file_object, size_bytes) if xai_batch_output else file_object,
|
||||
litellm_parent_otel_span=litellm_parent_otel_span,
|
||||
model_mappings={model_id: provider_file_id} if model_id else {},
|
||||
user_api_key_dict=owner,
|
||||
|
|
@ -1162,7 +1395,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
)
|
||||
return False
|
||||
|
||||
async def check_file_ids_access(self, file_ids: List[str], user_api_key_dict: UserAPIKeyAuth) -> None:
|
||||
async def check_file_ids_access(self, file_ids: Sequence[str], user_api_key_dict: UserAPIKeyAuth) -> None:
|
||||
"""
|
||||
Check if the user has access to a list of file IDs.
|
||||
Only checks managed (unified) file IDs.
|
||||
|
|
@ -1258,6 +1491,9 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
unified_vs_ids = self.get_vector_store_ids_from_file_search_tools(tools)
|
||||
if unified_vs_ids:
|
||||
await self.check_vector_store_ids_access(unified_vs_ids, user_api_key_dict)
|
||||
create_request: Final = _REQUEST_FIELDS.validate_python(data)
|
||||
if _is_streamed_background_create(create_request):
|
||||
return _with_managed_background_marker(create_request)
|
||||
elif call_type == CallTypes.afile_content.value:
|
||||
retrieve_file_id = cast(Optional[str], data.get("file_id"))
|
||||
potential_file_id = _is_base64_encoded_unified_file_id(retrieve_file_id) if retrieve_file_id else False
|
||||
|
|
@ -1269,6 +1505,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
elif call_type == CallTypes.acreate_batch.value:
|
||||
input_file_id = cast(Optional[str], data.get("input_file_id"))
|
||||
if input_file_id:
|
||||
await self.check_file_ids_access((input_file_id,), user_api_key_dict)
|
||||
model_file_id_mapping = await self.get_model_file_id_mapping(
|
||||
[input_file_id], user_api_key_dict.parent_otel_span
|
||||
)
|
||||
|
|
@ -1314,12 +1551,17 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
data[accessor_key] = get_batch_id_from_unified_batch_id(potential_llm_object_id)
|
||||
elif retrieve_object_id and accessor_key == "batch_id":
|
||||
await self.enforce_batch_object_access(retrieve_object_id, user_api_key_dict)
|
||||
elif call_type in _MANAGED_RESPONSE_READ_CALL_TYPES:
|
||||
request: Final = _REQUEST_FIELDS.validate_python(data)
|
||||
if _addresses_a_proxy_issued_response(request):
|
||||
return _with_managed_background_marker(request)
|
||||
elif call_type == CallTypes.acreate_fine_tuning_job.value:
|
||||
input_file_id = cast(Optional[str], data.get("training_file"))
|
||||
if input_file_id:
|
||||
model_file_id_mapping = await self.get_model_file_id_mapping(
|
||||
[input_file_id], user_api_key_dict.parent_otel_span
|
||||
)
|
||||
fine_tuning_file_ids: Final = tuple(
|
||||
file_id
|
||||
for file_id in (data.get("training_file"), data.get("validation_file"))
|
||||
if isinstance(file_id, str) and file_id
|
||||
)
|
||||
await self.check_file_ids_access(fine_tuning_file_ids, user_api_key_dict)
|
||||
|
||||
return data
|
||||
|
||||
|
|
@ -1941,8 +2183,20 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
)
|
||||
if refreshed_file_object is None:
|
||||
return _public_file_object(file_object, file_id)
|
||||
await self._save_refreshed_file_object(stored_file_object, refreshed_file_object)
|
||||
return _public_file_object(refreshed_file_object, file_id)
|
||||
if not _is_xai_batch_output(
|
||||
llm_router=llm_router if llm_router is not None else _proxy_llm_router(),
|
||||
model_id=model_id,
|
||||
model_name=model_id,
|
||||
provider_file_id=provider_file_id,
|
||||
):
|
||||
await self._save_refreshed_file_object(stored_file_object, refreshed_file_object)
|
||||
return _public_file_object(refreshed_file_object, file_id)
|
||||
saved_file_object: Final = await self._save_xai_refreshed_file_object(
|
||||
stored_file_object, refreshed_file_object
|
||||
)
|
||||
return _public_file_object(
|
||||
saved_file_object or _over_stored_object(refreshed_file_object, file_object), file_id
|
||||
)
|
||||
except Exception as error:
|
||||
verbose_logger.warning(
|
||||
"Failed to refresh batch file object for "
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm-enterprise"
|
||||
version = "0.1.76"
|
||||
version = "0.1.77"
|
||||
description = "Package for LiteLLM Enterprise features"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.9"
|
||||
|
|
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
|
|||
module-root = ""
|
||||
|
||||
[tool.commitizen]
|
||||
version = "0.1.76"
|
||||
version = "0.1.77"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
"../pyproject.toml:litellm-enterprise==",
|
||||
|
|
|
|||
|
|
@ -121,6 +121,8 @@ GATEWAY_PATH_PREFIXES: tuple[str, ...] = (
|
|||
# Realtime / streaming
|
||||
"/v1/realtime",
|
||||
"/realtime",
|
||||
"/v1/live/sessions",
|
||||
"/live/sessions",
|
||||
# Health & ops
|
||||
"/health",
|
||||
"/metrics",
|
||||
|
|
|
|||
|
|
@ -440,10 +440,10 @@ db:
|
|||
# Prefer the native /health/drain preStop hook over a fixed `sleep`: it marks
|
||||
# the pod NotReady and blocks only until in-flight requests actually finish
|
||||
# (bounded by GRACEFUL_SHUTDOWN_TIMEOUT, default 30s), instead of always
|
||||
# waiting the worst-case duration. The drain runs once (the preStop hook and
|
||||
# the SIGTERM handler share it), so set terminationGracePeriodSeconds a few
|
||||
# seconds above GRACEFUL_SHUTDOWN_TIMEOUT to leave room for teardown before
|
||||
# SIGKILL.
|
||||
# waiting the worst-case duration. After SIGTERM the server waits up to
|
||||
# GRACEFUL_SHUTDOWN_TIMEOUT again for connections still open and then cancels
|
||||
# them, so set terminationGracePeriodSeconds a few seconds above twice
|
||||
# GRACEFUL_SHUTDOWN_TIMEOUT to leave room for teardown before SIGKILL.
|
||||
#
|
||||
# /health/drain is off by default; enable it with
|
||||
# general_settings.enable_drain_endpoint: true. The kubelet calls preStop
|
||||
|
|
|
|||
|
|
@ -5,3 +5,21 @@ This chart deploys separate gateway, backend and UI services. Configure Lens thr
|
|||
Bundled mode generates separate service and identity-signing credentials. External mode requires `lensWorker.gateway.secretName` and `lensWorker.serviceTokenSecret.name`; their keys default to `gateway-secret` and `service-token`. Provision distinct values matching the external Lens deployment
|
||||
|
||||
Follow the [Lens Helm connection guide](https://github.com/BerriAI/lens/blob/main/helm/lens/README.md#connect-a-gateway) for the complete values and verification flow. It also links the GitOps credential requirements and existing-data migration. Source-chart installation requires a built Lens image until the first signed Lens release is published
|
||||
|
||||
For an ingress with separate component hostnames, set all three `ingress.componentHosts` values, leave `ingress.host` and `ingress.extraPaths` empty, configure `ui.backendUrl` as the backend's public origin and `backend.gatewayUrl` as the gateway's public origin, and set `ui.publicUrl` to the browser-facing UI origin. The UI Deployment receives `LITELLM_BACKEND_URL` from `ui.backendUrl` only when `ingress.componentHosts` is configured.
|
||||
|
||||
```yaml
|
||||
ingress:
|
||||
enabled: true
|
||||
host: ""
|
||||
componentHosts:
|
||||
ui: ui.example.com
|
||||
backend: backend.example.com
|
||||
gateway: gateway.example.com
|
||||
extraPaths: []
|
||||
ui:
|
||||
publicUrl: https://ui.example.com
|
||||
backendUrl: https://backend.example.com
|
||||
backend:
|
||||
gatewayUrl: https://gateway.example.com
|
||||
```
|
||||
|
|
|
|||
|
|
@ -22,12 +22,14 @@ Reminders:
|
|||
- database.writer.{host,port,dbname} (writer connection pieces)
|
||||
- database.writer.passwordSecret.{name,usernameKey,passwordKey}
|
||||
(Secret holding the writer DB username + password)
|
||||
- database.writer.passwordSecret.{hostKey,urlKey} (optional, read writer host or full URL from that Secret)
|
||||
- database.writer.useIAMAuth: true (optional — chart sets IAM_TOKEN_DB_AUTH=true and
|
||||
omits DATABASE_PASSWORD / DATABASE_URL so the proxy
|
||||
mints the URL from an IAM token at startup)
|
||||
- database.reader.host (optional — enables read-replica routing; reader
|
||||
.passwordSecret.name is required when set, unless
|
||||
.useIAMAuth is true)
|
||||
- database.reader.host / database.reader.passwordSecret.{hostKey,urlKey}
|
||||
(optional, enables read-replica routing. The Secret
|
||||
name is required for hostKey, urlKey, and for host
|
||||
with password auth)
|
||||
- database.reader.useIAMAuth: true (optional, requires database.writer.useIAMAuth: true —
|
||||
chart emits DATABASE_*_READ_REPLICA env vars and
|
||||
omits DATABASE_PASSWORD_READ_REPLICA /
|
||||
|
|
@ -35,6 +37,8 @@ Reminders:
|
|||
reader URL from an IAM token at startup)
|
||||
- redis.passwordSecret.name (optional — set when redis.host is provided and the
|
||||
cache requires auth)
|
||||
- redis.passwordSecret.{hostKey,urlKey} (optional, read Redis host or full URL from that Secret)
|
||||
(the Secret name is required when either key is set)
|
||||
- redis.cluster: true (optional — chart sets REDIS_CLUSTER_NODES from
|
||||
redis.host / redis.port so the proxy's Cache()
|
||||
constructs a RedisClusterCache; the cluster client
|
||||
|
|
|
|||
|
|
@ -27,6 +27,30 @@ Common naming + label helpers shared by gateway, backend, and ui templates.
|
|||
{{- printf "%s-ui" (include "litellm.fullname" .) | trunc 63 | trimSuffix "-" -}}
|
||||
{{- end -}}
|
||||
|
||||
{{- define "litellm.ui.publicOrigin" -}}
|
||||
{{- $origin := .Values.ui.publicUrl | trim | trimSuffix "/" -}}
|
||||
{{- if not (regexMatch "^https?://[^/\\s]+$" $origin) -}}
|
||||
{{- fail "ui.publicUrl must be an origin like https://ui.example.com, without a path" -}}
|
||||
{{- end -}}
|
||||
{{- $origin -}}
|
||||
{{- end -}}
|
||||
|
||||
{{- define "litellm.ui.backendOrigin" -}}
|
||||
{{- $origin := .Values.ui.backendUrl | trim | trimSuffix "/" -}}
|
||||
{{- if not (regexMatch "^https?://[^/\\s]+$" $origin) -}}
|
||||
{{- fail "ui.backendUrl must be an origin like https://backend.example.com, without a path, when ingress.componentHosts is set" -}}
|
||||
{{- end -}}
|
||||
{{- $origin -}}
|
||||
{{- end -}}
|
||||
|
||||
{{- define "litellm.backend.gatewayOrigin" -}}
|
||||
{{- $origin := .Values.backend.gatewayUrl | trim | trimSuffix "/" -}}
|
||||
{{- if not (regexMatch "^https?://[^/\\s]+$" $origin) -}}
|
||||
{{- fail "backend.gatewayUrl must be an origin like https://gateway.example.com, without a path, when ingress.componentHosts is set" -}}
|
||||
{{- end -}}
|
||||
{{- $origin -}}
|
||||
{{- end -}}
|
||||
|
||||
{{- define "litellm.commonLabels" -}}
|
||||
app.kubernetes.io/name: {{ include "litellm.name" . }}
|
||||
app.kubernetes.io/instance: {{ .Release.Name }}
|
||||
|
|
@ -202,7 +226,9 @@ the chart context (needed for .Values), `component` selects which component's
|
|||
Sensitive values (master key, DB username + password, Redis password) come
|
||||
only from referenced Secrets; the chart never accepts inline values for them.
|
||||
|
||||
The chart never assembles DATABASE_URL itself. It emits only the discrete
|
||||
The chart never assembles DATABASE_URL itself. When
|
||||
database.writer.passwordSecret.urlKey is set, the chart passes the operator's
|
||||
URL through from the Secret. Otherwise it emits only the discrete
|
||||
DATABASE_HOST/PORT/USER/NAME/SCHEMA (+ DATABASE_PASSWORD for password auth)
|
||||
vars; the proxy's entrypoint (DatabaseURLSettings in
|
||||
litellm/proxy/db/db_url_settings.py) builds the URL from them and
|
||||
|
|
@ -218,11 +244,11 @@ token instead of a static password. `database.writer.useAzureEntraAuth: true`
|
|||
does the same with AZURE_POSTGRESQL_AUTH=true and a Microsoft Entra ID token,
|
||||
for Azure Database for PostgreSQL. The two are mutually exclusive.
|
||||
|
||||
The read replica is opt-in via `database.reader.host`. The chart emits
|
||||
DATABASE_HOST_READ_REPLICA / DATABASE_PORT_READ_REPLICA /
|
||||
DATABASE_NAME_READ_REPLICA (+ DATABASE_SCHEMA_READ_REPLICA) for both auth
|
||||
modes, plus DATABASE_USER_READ_REPLICA / DATABASE_PASSWORD_READ_REPLICA for
|
||||
password auth. When `database.reader.useIAMAuth: true` (or
|
||||
The read replica is opt-in via `database.reader.host`, `hostKey`, or `urlKey`.
|
||||
The chart emits DATABASE_URL_READ_REPLICA alone for urlKey, or discrete
|
||||
DATABASE_*_READ_REPLICA vars for host and hostKey configurations, plus
|
||||
DATABASE_USER_READ_REPLICA / DATABASE_PASSWORD_READ_REPLICA for password auth.
|
||||
When `database.reader.useIAMAuth: true` (or
|
||||
`database.reader.useAzureEntraAuth: true`) it omits
|
||||
DATABASE_PASSWORD_READ_REPLICA and the entrypoint mints the reader URL the
|
||||
same way. Reader token auth only takes effect when the writer uses the same
|
||||
|
|
@ -242,8 +268,36 @@ IAM_TOKEN_DB_AUTH / AZURE_POSTGRESQL_AUTH toggle that only the writer sets.
|
|||
value: {{ $component.logLevel | quote }}
|
||||
{{- end }}
|
||||
{{- with $root.Values.database.writer }}
|
||||
{{- $hostKey := .passwordSecret.hostKey }}
|
||||
{{- $urlKey := .passwordSecret.urlKey }}
|
||||
{{- if and .host $hostKey }}
|
||||
{{- fail "database.writer.host and database.writer.passwordSecret.hostKey cannot both be set" }}
|
||||
{{- end }}
|
||||
{{- if and $urlKey (or .host $hostKey .dbname .schema .sslMode .sslRootCert .useIAMAuth .useAzureEntraAuth) }}
|
||||
{{- fail "database.writer.passwordSecret.urlKey cannot be combined with host, hostKey, dbname, schema, sslMode, sslRootCert, useIAMAuth, or useAzureEntraAuth" }}
|
||||
{{- end }}
|
||||
{{- if and (not .host) (not $hostKey) (not $urlKey) }}
|
||||
{{- fail "one of database.writer.host, database.writer.passwordSecret.hostKey, or database.writer.passwordSecret.urlKey is required" }}
|
||||
{{- end }}
|
||||
{{- if and (or $hostKey $urlKey) (not .passwordSecret.name) }}
|
||||
{{- fail "database.writer.passwordSecret.name is required when hostKey or urlKey is set" }}
|
||||
{{- end }}
|
||||
{{- if $urlKey }}
|
||||
- name: DATABASE_URL
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: {{ .passwordSecret.name | quote }}
|
||||
key: {{ $urlKey | quote }}
|
||||
{{- else }}
|
||||
- name: DATABASE_HOST
|
||||
value: {{ required "database.writer.host is required" .host | quote }}
|
||||
{{- if $hostKey }}
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: {{ .passwordSecret.name | quote }}
|
||||
key: {{ $hostKey | quote }}
|
||||
{{- else }}
|
||||
value: {{ .host | quote }}
|
||||
{{- end }}
|
||||
- name: DATABASE_PORT
|
||||
value: {{ .port | default 5432 | quote }}
|
||||
- name: DATABASE_USER
|
||||
|
|
@ -282,8 +336,29 @@ IAM_TOKEN_DB_AUTH / AZURE_POSTGRESQL_AUTH toggle that only the writer sets.
|
|||
key: {{ .passwordSecret.passwordKey | default "password" }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
{{- with $root.Values.database.reader }}
|
||||
{{- if .host }}
|
||||
{{- $readerSecret := .passwordSecret | default dict }}
|
||||
{{- $hostKey := $readerSecret.hostKey }}
|
||||
{{- $urlKey := $readerSecret.urlKey }}
|
||||
{{- $readerEnabled := or .host $hostKey $urlKey }}
|
||||
{{- if and .host $hostKey }}
|
||||
{{- fail "database.reader.host and database.reader.passwordSecret.hostKey cannot both be set" }}
|
||||
{{- end }}
|
||||
{{- if and $urlKey (or .host $hostKey .dbname .schema .useIAMAuth .useAzureEntraAuth) }}
|
||||
{{- fail "database.reader.passwordSecret.urlKey cannot be combined with host, hostKey, dbname, schema, useIAMAuth, or useAzureEntraAuth" }}
|
||||
{{- end }}
|
||||
{{- if and (or $hostKey $urlKey) (not $readerSecret.name) }}
|
||||
{{- fail "database.reader.passwordSecret.name is required when hostKey or urlKey is set" }}
|
||||
{{- end }}
|
||||
{{- if $readerEnabled }}
|
||||
{{- if $urlKey }}
|
||||
- name: DATABASE_URL_READ_REPLICA
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: {{ $readerSecret.name | quote }}
|
||||
key: {{ $urlKey | quote }}
|
||||
{{- else }}
|
||||
{{- if and .useIAMAuth (not $root.Values.database.writer.useIAMAuth) }}
|
||||
{{- fail "database.reader.useIAMAuth requires database.writer.useIAMAuth: true (the proxy gates IAM URL minting on IAM_TOKEN_DB_AUTH, which is only set by the writer)" }}
|
||||
{{- end }}
|
||||
|
|
@ -291,11 +366,18 @@ IAM_TOKEN_DB_AUTH / AZURE_POSTGRESQL_AUTH toggle that only the writer sets.
|
|||
{{- fail "database.reader.useAzureEntraAuth requires database.writer.useAzureEntraAuth: true (the proxy gates Entra URL minting on AZURE_POSTGRESQL_AUTH, which is only set by the writer)" }}
|
||||
{{- end }}
|
||||
- name: DATABASE_HOST_READ_REPLICA
|
||||
{{- if $hostKey }}
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: {{ $readerSecret.name | quote }}
|
||||
key: {{ $hostKey | quote }}
|
||||
{{- else }}
|
||||
value: {{ .host | quote }}
|
||||
{{- end }}
|
||||
- name: DATABASE_PORT_READ_REPLICA
|
||||
value: {{ .port | default 5432 | quote }}
|
||||
- name: DATABASE_NAME_READ_REPLICA
|
||||
value: {{ required "database.reader.dbname is required when database.reader.host is set" .dbname | quote }}
|
||||
value: {{ required "database.reader.dbname is required when database.reader.host, hostKey, or urlKey is set" .dbname | quote }}
|
||||
{{- if .schema }}
|
||||
- name: DATABASE_SCHEMA_READ_REPLICA
|
||||
value: {{ .schema | quote }}
|
||||
|
|
@ -325,6 +407,7 @@ IAM_TOKEN_DB_AUTH / AZURE_POSTGRESQL_AUTH toggle that only the writer sets.
|
|||
{{- end }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
{{/*
|
||||
The migrations Job (helm.sh/hook: pre-upgrade) is the single owner of
|
||||
`prisma migrate deploy`. Without this, every gateway/backend pod also runs
|
||||
|
|
@ -341,9 +424,34 @@ harmless no-op for the Job and authoritative for the app pods.
|
|||
tracking, pod lock manager) via its REDIS_* env fallback. An explicit
|
||||
`general_settings.coordination_redis` block in proxy_config takes
|
||||
precedence over anything emitted here. */}}
|
||||
{{- if $root.Values.redis.host }}
|
||||
{{- $redisSecret := $root.Values.redis.passwordSecret | default dict }}
|
||||
{{- $hostKey := $redisSecret.hostKey }}
|
||||
{{- $urlKey := $redisSecret.urlKey }}
|
||||
{{- if and $root.Values.redis.host $hostKey }}
|
||||
{{- fail "redis.host and redis.passwordSecret.hostKey cannot both be set" }}
|
||||
{{- end }}
|
||||
{{- if and $urlKey (or $root.Values.redis.host $hostKey $root.Values.redis.cluster) }}
|
||||
{{- fail "redis.passwordSecret.urlKey cannot be combined with host, hostKey, or cluster" }}
|
||||
{{- end }}
|
||||
{{- if and (or $hostKey $urlKey) (not $redisSecret.name) }}
|
||||
{{- fail "redis.passwordSecret.name is required when hostKey or urlKey is set" }}
|
||||
{{- end }}
|
||||
{{- if $urlKey }}
|
||||
- name: REDIS_URL
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: {{ $redisSecret.name | quote }}
|
||||
key: {{ $urlKey | quote }}
|
||||
{{- else if or $root.Values.redis.host $hostKey }}
|
||||
- name: REDIS_HOST
|
||||
{{- if $hostKey }}
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: {{ $redisSecret.name | quote }}
|
||||
key: {{ $hostKey | quote }}
|
||||
{{- else }}
|
||||
value: {{ $root.Values.redis.host | quote }}
|
||||
{{- end }}
|
||||
- name: REDIS_PORT
|
||||
value: {{ $root.Values.redis.port | quote }}
|
||||
{{- if $root.Values.redis.passwordSecret.name }}
|
||||
|
|
@ -360,7 +468,11 @@ harmless no-op for the Job and authoritative for the app pods.
|
|||
configured endpoint; the cluster client discovers the remaining nodes from
|
||||
CLUSTER SLOTS at startup. */}}
|
||||
- name: REDIS_CLUSTER_NODES
|
||||
{{- if $hostKey }}
|
||||
value: {{ printf "[{\"host\":\"$(REDIS_HOST)\",\"port\":%v}]" (int $root.Values.redis.port) | quote }}
|
||||
{{- else }}
|
||||
value: {{ printf "[{\"host\":%q,\"port\":%v}]" $root.Values.redis.host (int $root.Values.redis.port) | quote }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
{{- with $component.extraEnv }}
|
||||
|
|
|
|||
|
|
@ -1,4 +1,6 @@
|
|||
{{- if .Values.backend.enabled }}
|
||||
{{- $componentHosts := .Values.ingress.componentHosts | default dict -}}
|
||||
{{- $hasComponentHosts := or $componentHosts.ui $componentHosts.backend $componentHosts.gateway -}}
|
||||
apiVersion: apps/v1
|
||||
kind: Deployment
|
||||
metadata:
|
||||
|
|
@ -58,7 +60,17 @@ spec:
|
|||
protocol: TCP
|
||||
env:
|
||||
{{- include "litellm.lensConnectionEnv" . | nindent 12 }}
|
||||
{{- if .Values.ui.publicUrl }}
|
||||
- name: LITELLM_CORS_ORIGINS
|
||||
value: {{ include "litellm.ui.publicOrigin" . | quote }}
|
||||
- name: LITELLM_CORS_ALLOW_CREDENTIALS
|
||||
value: "true"
|
||||
{{- end }}
|
||||
{{- include "litellm.serverEnv" (dict "root" $ "component" .Values.backend) | nindent 12 }}
|
||||
{{- if .Values.backend.gatewayUrl }}
|
||||
- name: LITELLM_GATEWAY_URL
|
||||
value: {{ if $hasComponentHosts }}{{ include "litellm.backend.gatewayOrigin" . | quote }}{{ else }}{{ .Values.backend.gatewayUrl | quote }}{{ end }}
|
||||
{{- end }}
|
||||
{{- if .Values.gateway.config.create }}
|
||||
- name: CONFIG_FILE_PATH
|
||||
value: /app/config/config.yaml
|
||||
|
|
|
|||
|
|
@ -56,6 +56,10 @@ spec:
|
|||
protocol: TCP
|
||||
env:
|
||||
{{- include "litellm.lensConnectionEnv" . | nindent 12 }}
|
||||
{{- if .Values.ui.publicUrl }}
|
||||
- name: LITELLM_CORS_ORIGINS
|
||||
value: {{ include "litellm.ui.publicOrigin" . | quote }}
|
||||
{{- end }}
|
||||
{{- include "litellm.serverEnv" (dict "root" $ "component" .Values.gateway) | nindent 12 }}
|
||||
{{- if .Values.gateway.config.create }}
|
||||
- name: CONFIG_FILE_PATH
|
||||
|
|
|
|||
|
|
@ -1,3 +1,36 @@
|
|||
{{- $componentHosts := .Values.ingress.componentHosts | default dict -}}
|
||||
{{- $hasComponentHosts := or $componentHosts.ui $componentHosts.backend $componentHosts.gateway -}}
|
||||
{{- if $hasComponentHosts }}
|
||||
{{- if not (and $componentHosts.ui $componentHosts.backend $componentHosts.gateway) }}
|
||||
{{- fail "ingress.componentHosts.ui, backend, and gateway must all be set together" }}
|
||||
{{- end }}
|
||||
{{- if not .Values.gateway.enabled }}
|
||||
{{- fail "ingress.componentHosts requires gateway.enabled=true" }}
|
||||
{{- end }}
|
||||
{{- if not .Values.backend.enabled }}
|
||||
{{- fail "ingress.componentHosts requires backend.enabled=true" }}
|
||||
{{- end }}
|
||||
{{- if not .Values.ui.enabled }}
|
||||
{{- fail "ingress.componentHosts requires ui.enabled=true" }}
|
||||
{{- end }}
|
||||
{{- if .Values.ingress.host }}
|
||||
{{- fail "ingress.host must be empty when ingress.componentHosts is configured" }}
|
||||
{{- end }}
|
||||
{{- if .Values.ingress.extraPaths }}
|
||||
{{- fail "ingress.extraPaths must be empty when ingress.componentHosts is configured" }}
|
||||
{{- end }}
|
||||
{{- if not (trim .Values.ui.backendUrl) }}
|
||||
{{- fail "ui.backendUrl is required when ingress.componentHosts is configured" }}
|
||||
{{- end }}
|
||||
{{- $_ := include "litellm.ui.backendOrigin" . }}
|
||||
{{- if not (trim .Values.backend.gatewayUrl) }}
|
||||
{{- fail "backend.gatewayUrl is required when ingress.componentHosts is configured" }}
|
||||
{{- end }}
|
||||
{{- $_ := include "litellm.backend.gatewayOrigin" . }}
|
||||
{{- if not (trim .Values.ui.publicUrl) }}
|
||||
{{- fail "ui.publicUrl is required when ingress.componentHosts is configured" }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
{{- if .Values.ingress.enabled -}}
|
||||
{{- $gatewayName := include "litellm.gateway.fullname" . -}}
|
||||
{{- $backendName := include "litellm.backend.fullname" . -}}
|
||||
|
|
@ -90,6 +123,66 @@
|
|||
adding to it.
|
||||
*/}}
|
||||
{{- $builtinPathKeys := list "/test|Exact" "/debug/memory/summary|Exact" "/|Prefix" -}}
|
||||
{{- if $hasComponentHosts }}
|
||||
apiVersion: networking.k8s.io/v1
|
||||
kind: Ingress
|
||||
metadata:
|
||||
name: {{ include "litellm.fullname" . }}
|
||||
labels:
|
||||
{{- include "litellm.commonLabels" . | nindent 4 }}
|
||||
{{- with .Values.ingress.annotations }}
|
||||
annotations:
|
||||
{{- toYaml . | nindent 4 }}
|
||||
{{- end }}
|
||||
spec:
|
||||
{{- with .Values.ingress.className }}
|
||||
ingressClassName: {{ . | quote }}
|
||||
{{- end }}
|
||||
{{- with .Values.ingress.tls }}
|
||||
tls:
|
||||
{{- toYaml . | nindent 4 }}
|
||||
{{- end }}
|
||||
rules:
|
||||
- host: {{ $componentHosts.ui | quote }}
|
||||
http:
|
||||
paths:
|
||||
- path: /
|
||||
pathType: Prefix
|
||||
backend:
|
||||
service:
|
||||
name: {{ $uiName }}
|
||||
port:
|
||||
number: {{ $uiPort }}
|
||||
- host: {{ $componentHosts.backend | quote }}
|
||||
http:
|
||||
paths:
|
||||
- path: /
|
||||
pathType: Prefix
|
||||
backend:
|
||||
service:
|
||||
name: {{ $backendName }}
|
||||
port:
|
||||
number: {{ $backendPort }}
|
||||
- host: {{ $componentHosts.gateway | quote }}
|
||||
http:
|
||||
paths:
|
||||
{{- if or (eq (include "litellm.lens.mode" .) "bundled") (and (eq (include "litellm.lens.mode" .) "external") .Values.lensWorker.externalServiceName) }}
|
||||
- path: /lens-ingest
|
||||
pathType: Prefix
|
||||
backend:
|
||||
service:
|
||||
name: {{ .Values.lensWorker.externalServiceName | default (printf "%s-lens-worker" (include "litellm.fullname" .)) }}
|
||||
port:
|
||||
number: {{ .Values.lensWorker.service.port }}
|
||||
{{- end }}
|
||||
- path: /
|
||||
pathType: Prefix
|
||||
backend:
|
||||
service:
|
||||
name: {{ $gatewayName }}
|
||||
port:
|
||||
number: {{ $gatewayPort }}
|
||||
{{- else }}
|
||||
apiVersion: networking.k8s.io/v1
|
||||
kind: Ingress
|
||||
metadata:
|
||||
|
|
@ -217,3 +310,4 @@ spec:
|
|||
port:
|
||||
number: {{ $backendPort }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
|
|
|
|||
|
|
@ -1,3 +1,5 @@
|
|||
{{- $componentHosts := .Values.ingress.componentHosts | default dict -}}
|
||||
{{- $hasComponentHosts := or $componentHosts.ui $componentHosts.backend $componentHosts.gateway -}}
|
||||
{{- if .Values.ui.enabled }}
|
||||
apiVersion: apps/v1
|
||||
kind: Deployment
|
||||
|
|
@ -56,9 +58,9 @@ spec:
|
|||
- name: LITELLM_LOG
|
||||
value: {{ .Values.ui.logLevel | quote }}
|
||||
{{- end }}
|
||||
{{- if .Values.ui.backendUrl }}
|
||||
{{- if and .Values.ui.backendUrl $hasComponentHosts }}
|
||||
- name: LITELLM_BACKEND_URL
|
||||
value: {{ .Values.ui.backendUrl | quote }}
|
||||
value: {{ include "litellm.ui.backendOrigin" . | quote }}
|
||||
{{- end }}
|
||||
{{- with .Values.ui.extraEnv }}
|
||||
{{- toYaml . | nindent 12 }}
|
||||
|
|
|
|||
441
helm/litellm/tests/component_hosts_tests.yaml
Normal file
441
helm/litellm/tests/component_hosts_tests.yaml
Normal file
|
|
@ -0,0 +1,441 @@
|
|||
suite: ingress component hosts
|
||||
templates:
|
||||
- ingress.yaml
|
||||
- backend/deployment.yaml
|
||||
- gateway/deployment.yaml
|
||||
- gateway/configmap.yaml
|
||||
- ui/deployment.yaml
|
||||
values:
|
||||
- ./values/required.yaml
|
||||
tests:
|
||||
- it: routes each hostname to its component and injects the external URLs
|
||||
template: ingress.yaml
|
||||
set:
|
||||
ingress.enabled: true
|
||||
ingress.componentHosts.ui: ui.example.com
|
||||
ingress.componentHosts.backend: backend.example.com
|
||||
ingress.componentHosts.gateway: gateway.example.com
|
||||
lensWorker.mode: disabled
|
||||
ingress.tls:
|
||||
- hosts:
|
||||
- ui.example.com
|
||||
- backend.example.com
|
||||
- gateway.example.com
|
||||
secretName: component-hosts
|
||||
ui.backendUrl: https://backend.example.com
|
||||
ui.publicUrl: https://ui.example.com
|
||||
backend.gatewayUrl: https://gateway.example.com
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.rules[0].host
|
||||
value: ui.example.com
|
||||
- equal:
|
||||
path: spec.tls[0]
|
||||
value:
|
||||
hosts:
|
||||
- ui.example.com
|
||||
- backend.example.com
|
||||
- gateway.example.com
|
||||
secretName: component-hosts
|
||||
- equal:
|
||||
path: spec.rules[0].http.paths
|
||||
value:
|
||||
- path: /
|
||||
pathType: Prefix
|
||||
backend:
|
||||
service:
|
||||
name: RELEASE-NAME-litellm-ui
|
||||
port:
|
||||
number: 3000
|
||||
- equal:
|
||||
path: spec.rules[1].host
|
||||
value: backend.example.com
|
||||
- equal:
|
||||
path: spec.rules[1].http.paths
|
||||
value:
|
||||
- path: /
|
||||
pathType: Prefix
|
||||
backend:
|
||||
service:
|
||||
name: RELEASE-NAME-litellm-backend
|
||||
port:
|
||||
number: 4001
|
||||
- equal:
|
||||
path: spec.rules[2].host
|
||||
value: gateway.example.com
|
||||
- equal:
|
||||
path: spec.rules[2].http.paths
|
||||
value:
|
||||
- path: /
|
||||
pathType: Prefix
|
||||
backend:
|
||||
service:
|
||||
name: RELEASE-NAME-litellm-gateway
|
||||
port:
|
||||
number: 4000
|
||||
|
||||
- it: keeps Lens ingestion ahead of the gateway catch-all
|
||||
set:
|
||||
ingress.enabled: true
|
||||
ingress.componentHosts.ui: ui.example.com
|
||||
ingress.componentHosts.backend: backend.example.com
|
||||
ingress.componentHosts.gateway: gateway.example.com
|
||||
ui.backendUrl: https://backend.example.com
|
||||
ui.publicUrl: https://ui.example.com
|
||||
backend.gatewayUrl: https://gateway.example.com
|
||||
lensWorker.mode: external
|
||||
lensWorker.externalServiceName: existing-lens
|
||||
lensWorker.externalUrl: https://lens.example.com
|
||||
lensWorker.publicUrl: https://lens.example.com
|
||||
lensWorker.serviceTokenSecret.name: existing-lens-signing
|
||||
lensWorker.gateway.secretName: existing-lens-signing
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.rules[2].http.paths[0].path
|
||||
value: /lens-ingest
|
||||
- equal:
|
||||
path: spec.rules[2].http.paths[0].backend.service.name
|
||||
value: existing-lens
|
||||
- equal:
|
||||
path: spec.rules[2].http.paths[1].path
|
||||
value: /
|
||||
template: ingress.yaml
|
||||
|
||||
- it: rejects a partial component host configuration
|
||||
set:
|
||||
ingress.enabled: true
|
||||
ingress.componentHosts.ui: ui.example.com
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: ingress.componentHosts.ui, backend, and gateway must all be set together
|
||||
template: ingress.yaml
|
||||
|
||||
- it: rejects component hosts when the gateway is disabled
|
||||
set:
|
||||
ingress.enabled: true
|
||||
ingress.componentHosts.ui: ui.example.com
|
||||
ingress.componentHosts.backend: backend.example.com
|
||||
ingress.componentHosts.gateway: gateway.example.com
|
||||
ui.backendUrl: https://backend.example.com
|
||||
ui.publicUrl: https://ui.example.com
|
||||
backend.gatewayUrl: https://gateway.example.com
|
||||
gateway.enabled: false
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: ingress.componentHosts requires gateway.enabled=true
|
||||
template: ingress.yaml
|
||||
|
||||
- it: rejects component hosts when the backend is disabled
|
||||
set:
|
||||
ingress.enabled: true
|
||||
ingress.componentHosts.ui: ui.example.com
|
||||
ingress.componentHosts.backend: backend.example.com
|
||||
ingress.componentHosts.gateway: gateway.example.com
|
||||
ui.backendUrl: https://backend.example.com
|
||||
ui.publicUrl: https://ui.example.com
|
||||
backend.gatewayUrl: https://gateway.example.com
|
||||
backend.enabled: false
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: ingress.componentHosts requires backend.enabled=true
|
||||
template: ingress.yaml
|
||||
|
||||
- it: rejects component hosts when the UI is disabled
|
||||
set:
|
||||
ingress.enabled: true
|
||||
ingress.componentHosts.ui: ui.example.com
|
||||
ingress.componentHosts.backend: backend.example.com
|
||||
ingress.componentHosts.gateway: gateway.example.com
|
||||
ui.backendUrl: https://backend.example.com
|
||||
ui.publicUrl: https://ui.example.com
|
||||
backend.gatewayUrl: https://gateway.example.com
|
||||
ui.enabled: false
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: ingress.componentHosts requires ui.enabled=true
|
||||
template: ingress.yaml
|
||||
|
||||
- it: rejects a legacy host alongside component hosts
|
||||
set:
|
||||
ingress.enabled: true
|
||||
ingress.host: legacy.example.com
|
||||
ingress.componentHosts.ui: ui.example.com
|
||||
ingress.componentHosts.backend: backend.example.com
|
||||
ingress.componentHosts.gateway: gateway.example.com
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: ingress.host must be empty when ingress.componentHosts is configured
|
||||
template: ingress.yaml
|
||||
|
||||
- it: rejects extra paths alongside component hosts
|
||||
set:
|
||||
ingress.enabled: true
|
||||
ingress.componentHosts.ui: ui.example.com
|
||||
ingress.componentHosts.backend: backend.example.com
|
||||
ingress.componentHosts.gateway: gateway.example.com
|
||||
ingress.extraPaths:
|
||||
- path: /custom
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: ingress.extraPaths must be empty when ingress.componentHosts is configured
|
||||
template: ingress.yaml
|
||||
|
||||
- it: requires the UI backend URL in component host mode
|
||||
set:
|
||||
ingress.enabled: true
|
||||
ingress.componentHosts.ui: ui.example.com
|
||||
ingress.componentHosts.backend: backend.example.com
|
||||
ingress.componentHosts.gateway: gateway.example.com
|
||||
ui.publicUrl: https://ui.example.com
|
||||
backend.gatewayUrl: https://gateway.example.com
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: ui.backendUrl is required when ingress.componentHosts is configured
|
||||
template: ingress.yaml
|
||||
|
||||
- it: requires the backend gateway URL in component host mode
|
||||
set:
|
||||
ingress.enabled: true
|
||||
ingress.componentHosts.ui: ui.example.com
|
||||
ingress.componentHosts.backend: backend.example.com
|
||||
ingress.componentHosts.gateway: gateway.example.com
|
||||
ui.backendUrl: https://backend.example.com
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: backend.gatewayUrl is required when ingress.componentHosts is configured
|
||||
template: ingress.yaml
|
||||
|
||||
- it: requires the public UI origin in component host mode
|
||||
set:
|
||||
ingress.enabled: true
|
||||
ingress.componentHosts.ui: ui.example.com
|
||||
ingress.componentHosts.backend: backend.example.com
|
||||
ingress.componentHosts.gateway: gateway.example.com
|
||||
ui.backendUrl: https://backend.example.com
|
||||
backend.gatewayUrl: https://gateway.example.com
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: ui.publicUrl is required when ingress.componentHosts is configured
|
||||
template: ingress.yaml
|
||||
|
||||
- it: exposes the gateway URL in the backend Deployment
|
||||
set:
|
||||
ingress.enabled: true
|
||||
ingress.componentHosts.ui: ui.example.com
|
||||
ingress.componentHosts.backend: backend.example.com
|
||||
ingress.componentHosts.gateway: gateway.example.com
|
||||
ui.backendUrl: https://backend.example.com
|
||||
ui.publicUrl: https://ui.example.com
|
||||
backend.gatewayUrl: " https://gateway.example.com/ "
|
||||
template: backend/deployment.yaml
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_GATEWAY_URL
|
||||
value: https://gateway.example.com
|
||||
|
||||
- it: preserves the backend gateway URL verbatim in single-host mode
|
||||
set:
|
||||
ingress.enabled: true
|
||||
ingress.host: litellm.example.com
|
||||
backend.gatewayUrl: " https://gateway.example.com/api/ "
|
||||
template: backend/deployment.yaml
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_GATEWAY_URL
|
||||
value: " https://gateway.example.com/api/ "
|
||||
|
||||
- it: exposes the backend URL in the UI Deployment
|
||||
set:
|
||||
ingress.enabled: true
|
||||
ingress.componentHosts.ui: ui.example.com
|
||||
ingress.componentHosts.backend: backend.example.com
|
||||
ingress.componentHosts.gateway: gateway.example.com
|
||||
ui.backendUrl: " https://backend.example.com/ "
|
||||
ui.publicUrl: https://ui.example.com
|
||||
backend.gatewayUrl: https://gateway.example.com
|
||||
template: ui/deployment.yaml
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_BACKEND_URL
|
||||
value: https://backend.example.com
|
||||
|
||||
- it: omits the backend URL from the UI Deployment in single-host mode
|
||||
set:
|
||||
ingress.enabled: true
|
||||
ingress.host: litellm.example.com
|
||||
ui.backendUrl: https://litellm.example.com/api
|
||||
template: ui/deployment.yaml
|
||||
asserts:
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_BACKEND_URL
|
||||
any: true
|
||||
|
||||
- it: enables credentialed CORS from the public UI origin on the backend
|
||||
set:
|
||||
ui.publicUrl: https://ui.example.com
|
||||
template: backend/deployment.yaml
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_CORS_ORIGINS
|
||||
value: https://ui.example.com
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_CORS_ALLOW_CREDENTIALS
|
||||
value: "true"
|
||||
|
||||
- it: lets backend extraEnv override generated CORS values
|
||||
set:
|
||||
ui.publicUrl: https://ui.example.com
|
||||
backend.extraEnv:
|
||||
- name: LITELLM_CORS_ORIGINS
|
||||
value: https://backend-override.example.com
|
||||
- name: LITELLM_CORS_ALLOW_CREDENTIALS
|
||||
value: "false"
|
||||
template: backend/deployment.yaml
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].env[0].name
|
||||
value: LITELLM_CORS_ORIGINS
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].env[0].value
|
||||
value: https://ui.example.com
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].env[1].name
|
||||
value: LITELLM_CORS_ALLOW_CREDENTIALS
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].env[1].value
|
||||
value: "true"
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].env[10].name
|
||||
value: LITELLM_CORS_ORIGINS
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].env[10].value
|
||||
value: https://backend-override.example.com
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].env[11].name
|
||||
value: LITELLM_CORS_ALLOW_CREDENTIALS
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].env[11].value
|
||||
value: "false"
|
||||
|
||||
- it: exposes the public UI origin to the gateway without enabling credentials
|
||||
set:
|
||||
ui.publicUrl: https://ui.example.com
|
||||
template: gateway/deployment.yaml
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_CORS_ORIGINS
|
||||
value: https://ui.example.com
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_CORS_ALLOW_CREDENTIALS
|
||||
value: "true"
|
||||
|
||||
- it: trims the public UI origin before setting backend CORS
|
||||
set:
|
||||
ui.publicUrl: " https://ui.example.com/ "
|
||||
template: backend/deployment.yaml
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_CORS_ORIGINS
|
||||
value: https://ui.example.com
|
||||
|
||||
- it: trims the public UI origin before setting gateway CORS
|
||||
set:
|
||||
ui.publicUrl: " https://ui.example.com/ "
|
||||
template: gateway/deployment.yaml
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_CORS_ORIGINS
|
||||
value: https://ui.example.com
|
||||
|
||||
- it: rejects a UI public URL with a path
|
||||
set:
|
||||
ui.publicUrl: https://ui.example.com/app
|
||||
template: gateway/deployment.yaml
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: ui.publicUrl must be an origin like https://ui.example.com, without a path
|
||||
|
||||
- it: lets gateway extraEnv override the generated CORS origin
|
||||
set:
|
||||
ui.publicUrl: https://ui.example.com
|
||||
gateway.extraEnv:
|
||||
- name: LITELLM_CORS_ORIGINS
|
||||
value: https://gateway-override.example.com
|
||||
template: gateway/deployment.yaml
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].env[0].name
|
||||
value: LITELLM_CORS_ORIGINS
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].env[0].value
|
||||
value: https://ui.example.com
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].env[9].name
|
||||
value: LITELLM_CORS_ORIGINS
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].env[9].value
|
||||
value: https://gateway-override.example.com
|
||||
|
||||
---
|
||||
suite: component host UI backend origin validation
|
||||
templates:
|
||||
- ingress.yaml
|
||||
values:
|
||||
- ./values/required.yaml
|
||||
tests:
|
||||
- it: rejects a pathful UI backend URL in component host mode
|
||||
set:
|
||||
ingress.enabled: true
|
||||
ingress.componentHosts.ui: ui.example.com
|
||||
ingress.componentHosts.backend: backend.example.com
|
||||
ingress.componentHosts.gateway: gateway.example.com
|
||||
ui.backendUrl: https://backend.example.com/api
|
||||
ui.publicUrl: https://ui.example.com
|
||||
backend.gatewayUrl: https://gateway.example.com
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: ui.backendUrl must be an origin like https://backend.example.com, without a path, when ingress.componentHosts is set
|
||||
- it: rejects a pathful backend gateway URL in component host mode
|
||||
set:
|
||||
ingress.enabled: true
|
||||
ingress.componentHosts.ui: ui.example.com
|
||||
ingress.componentHosts.backend: backend.example.com
|
||||
ingress.componentHosts.gateway: gateway.example.com
|
||||
ui.backendUrl: https://backend.example.com
|
||||
ui.publicUrl: https://ui.example.com
|
||||
backend.gatewayUrl: https://gateway.example.com/v1
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: backend.gatewayUrl must be an origin like https://gateway.example.com, without a path, when ingress.componentHosts is set
|
||||
- it: rejects a backend gateway URL without a scheme in component host mode
|
||||
set:
|
||||
ingress.enabled: true
|
||||
ingress.componentHosts.ui: ui.example.com
|
||||
ingress.componentHosts.backend: backend.example.com
|
||||
ingress.componentHosts.gateway: gateway.example.com
|
||||
ui.backendUrl: https://backend.example.com
|
||||
ui.publicUrl: https://ui.example.com
|
||||
backend.gatewayUrl: gateway.example.com
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: backend.gatewayUrl must be an origin like https://gateway.example.com, without a path, when ingress.componentHosts is set
|
||||
|
|
@ -191,3 +191,400 @@ tests:
|
|||
content:
|
||||
name: DATABASE_PASSWORD_READ_REPLICA
|
||||
any: true
|
||||
|
||||
- it: writer hostKey reads the host from the Secret and supports IAM auth
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
database.writer.host: ""
|
||||
database.writer.passwordSecret.hostKey: hostname
|
||||
database.writer.useIAMAuth: true
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_HOST
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: litellm-writer-secret
|
||||
key: hostname
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: IAM_TOKEN_DB_AUTH
|
||||
value: "true"
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_PASSWORD
|
||||
any: true
|
||||
|
||||
- it: writer urlKey supplies only DATABASE_URL to each database client container and keeps pool env
|
||||
set:
|
||||
database.writer.host: ""
|
||||
database.writer.dbname: ""
|
||||
database.writer.schema: ""
|
||||
database.writer.sslMode: ""
|
||||
database.writer.sslRootCert: ""
|
||||
database.writer.passwordSecret.urlKey: connection-url
|
||||
database.connectionPool.enabled: true
|
||||
gateway.collector.enabled: true
|
||||
gateway.metricsServer.enabled: true
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_URL
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: litellm-writer-secret
|
||||
key: connection-url
|
||||
template: gateway/deployment.yaml
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_HOST
|
||||
any: true
|
||||
template: gateway/deployment.yaml
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_PORT
|
||||
any: true
|
||||
template: gateway/deployment.yaml
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_USER
|
||||
any: true
|
||||
template: gateway/deployment.yaml
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_NAME
|
||||
any: true
|
||||
template: gateway/deployment.yaml
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_SCHEMA
|
||||
any: true
|
||||
template: gateway/deployment.yaml
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_SSLMODE
|
||||
any: true
|
||||
template: gateway/deployment.yaml
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_SSLROOTCERT
|
||||
any: true
|
||||
template: gateway/deployment.yaml
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_PASSWORD
|
||||
any: true
|
||||
template: gateway/deployment.yaml
|
||||
- notExists:
|
||||
path: >-
|
||||
spec.template.spec.containers[0].env[?(@.name == "DATABASE_HOST" ||
|
||||
@.name == "DATABASE_PORT" || @.name == "DATABASE_USER" ||
|
||||
@.name == "DATABASE_NAME" || @.name == "DATABASE_SCHEMA" ||
|
||||
@.name == "DATABASE_SSLMODE" || @.name == "DATABASE_SSLROOTCERT" ||
|
||||
@.name == "DATABASE_PASSWORD")]
|
||||
template: gateway/deployment.yaml
|
||||
- notExists:
|
||||
path: >-
|
||||
spec.template.spec.containers[2].env[?(@.name == "DATABASE_HOST" ||
|
||||
@.name == "DATABASE_PORT" || @.name == "DATABASE_USER" ||
|
||||
@.name == "DATABASE_NAME" || @.name == "DATABASE_SCHEMA" ||
|
||||
@.name == "DATABASE_SSLMODE" || @.name == "DATABASE_SSLROOTCERT" ||
|
||||
@.name == "DATABASE_PASSWORD")]
|
||||
template: gateway/deployment.yaml
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_PGBOUNCER_ENABLED
|
||||
value: "true"
|
||||
template: gateway/deployment.yaml
|
||||
- contains:
|
||||
path: spec.template.spec.containers[2].env
|
||||
content:
|
||||
name: DATABASE_URL
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: litellm-writer-secret
|
||||
key: connection-url
|
||||
template: gateway/deployment.yaml
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[2].env
|
||||
content:
|
||||
name: DATABASE_USER
|
||||
any: true
|
||||
template: gateway/deployment.yaml
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[2].env
|
||||
content:
|
||||
name: DATABASE_PASSWORD
|
||||
any: true
|
||||
template: gateway/deployment.yaml
|
||||
- contains:
|
||||
path: spec.template.spec.containers[2].env
|
||||
content:
|
||||
name: DATABASE_URL
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: litellm-writer-secret
|
||||
key: connection-url
|
||||
template: gateway/deployment.yaml
|
||||
- contains:
|
||||
path: spec.template.spec.containers[2].env
|
||||
content:
|
||||
name: LITELLM_PGBOUNCER_ENABLED
|
||||
value: "true"
|
||||
template: gateway/deployment.yaml
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_URL
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: litellm-writer-secret
|
||||
key: connection-url
|
||||
template: backend/deployment.yaml
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_USER
|
||||
any: true
|
||||
template: backend/deployment.yaml
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_PASSWORD
|
||||
any: true
|
||||
template: backend/deployment.yaml
|
||||
- notExists:
|
||||
path: >-
|
||||
spec.template.spec.containers[0].env[?(@.name == "DATABASE_HOST" ||
|
||||
@.name == "DATABASE_PORT" || @.name == "DATABASE_USER" ||
|
||||
@.name == "DATABASE_NAME" || @.name == "DATABASE_SCHEMA" ||
|
||||
@.name == "DATABASE_SSLMODE" || @.name == "DATABASE_SSLROOTCERT" ||
|
||||
@.name == "DATABASE_PASSWORD")]
|
||||
template: backend/deployment.yaml
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_URL
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: litellm-writer-secret
|
||||
key: connection-url
|
||||
template: migrations-job.yaml
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_USER
|
||||
any: true
|
||||
template: migrations-job.yaml
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_PASSWORD
|
||||
any: true
|
||||
template: migrations-job.yaml
|
||||
- notExists:
|
||||
path: >-
|
||||
spec.template.spec.containers[0].env[?(@.name == "DATABASE_HOST" ||
|
||||
@.name == "DATABASE_PORT" || @.name == "DATABASE_USER" ||
|
||||
@.name == "DATABASE_NAME" || @.name == "DATABASE_SCHEMA" ||
|
||||
@.name == "DATABASE_SSLMODE" || @.name == "DATABASE_SSLROOTCERT" ||
|
||||
@.name == "DATABASE_PASSWORD")]
|
||||
template: migrations-job.yaml
|
||||
|
||||
- it: writer rejects host together with hostKey
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
database.writer.passwordSecret.hostKey: hostname
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: "database.writer.host and database.writer.passwordSecret.hostKey cannot both be set"
|
||||
|
||||
- it: writer rejects urlKey with discrete connection settings
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
database.writer.host: ""
|
||||
database.writer.passwordSecret.urlKey: connection-url
|
||||
database.writer.schema: public
|
||||
database.writer.sslMode: require
|
||||
database.writer.sslRootCert: /etc/ssl/ca.crt
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: "database.writer.passwordSecret.urlKey cannot be combined with host, hostKey, dbname, schema, sslMode, sslRootCert, useIAMAuth, or useAzureEntraAuth"
|
||||
|
||||
- it: writer rejects urlKey with IAM auth
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
database.writer.host: ""
|
||||
database.writer.dbname: ""
|
||||
database.writer.passwordSecret.urlKey: connection-url
|
||||
database.writer.useIAMAuth: true
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: "database.writer.passwordSecret.urlKey cannot be combined with host, hostKey, dbname, schema, sslMode, sslRootCert, useIAMAuth, or useAzureEntraAuth"
|
||||
|
||||
- it: writer rejects urlKey with Entra auth
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
database.writer.host: ""
|
||||
database.writer.dbname: ""
|
||||
database.writer.passwordSecret.urlKey: connection-url
|
||||
database.writer.useAzureEntraAuth: true
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: "database.writer.passwordSecret.urlKey cannot be combined with host, hostKey, dbname, schema, sslMode, sslRootCert, useIAMAuth, or useAzureEntraAuth"
|
||||
|
||||
- it: writer requires an endpoint
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
database.writer.host: ""
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: "one of database.writer.host, database.writer.passwordSecret.hostKey, or database.writer.passwordSecret.urlKey is required"
|
||||
|
||||
- it: writer hostKey requires a Secret name
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
database.writer.host: ""
|
||||
database.writer.passwordSecret.hostKey: hostname
|
||||
database.writer.passwordSecret.name: ""
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: "database.writer.passwordSecret.name is required when hostKey or urlKey is set"
|
||||
|
||||
- it: reader renders without replica env vars when passwordSecret is null
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
database.reader.passwordSecret: null
|
||||
asserts:
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_HOST_READ_REPLICA
|
||||
any: true
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_URL_READ_REPLICA
|
||||
any: true
|
||||
|
||||
- it: reader hostKey enables replica routing and reads its host from the Secret
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
database.writer.useIAMAuth: true
|
||||
database.reader.passwordSecret.hostKey: hostname
|
||||
database.reader.dbname: litellm
|
||||
database.reader.useIAMAuth: true
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_HOST_READ_REPLICA
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: litellm-reader-secret
|
||||
key: hostname
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_NAME_READ_REPLICA
|
||||
value: litellm
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: IAM_TOKEN_DB_AUTH
|
||||
value: "true"
|
||||
|
||||
- it: reader urlKey emits only the replica URL
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
database.reader.passwordSecret.urlKey: connection-url
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_URL_READ_REPLICA
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: litellm-reader-secret
|
||||
key: connection-url
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_HOST_READ_REPLICA
|
||||
any: true
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_PORT_READ_REPLICA
|
||||
any: true
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_USER_READ_REPLICA
|
||||
any: true
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_PASSWORD_READ_REPLICA
|
||||
any: true
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_NAME_READ_REPLICA
|
||||
any: true
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_SCHEMA_READ_REPLICA
|
||||
any: true
|
||||
|
||||
- it: reader rejects host together with hostKey
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
database.reader.host: reader.example.com
|
||||
database.reader.passwordSecret.hostKey: hostname
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: "database.reader.host and database.reader.passwordSecret.hostKey cannot both be set"
|
||||
|
||||
- it: reader rejects urlKey with discrete connection settings
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
database.reader.passwordSecret.urlKey: connection-url
|
||||
database.reader.dbname: litellm
|
||||
database.reader.schema: public
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: "database.reader.passwordSecret.urlKey cannot be combined with host, hostKey, dbname, schema, useIAMAuth, or useAzureEntraAuth"
|
||||
|
||||
- it: reader rejects urlKey with token auth
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
database.writer.useIAMAuth: true
|
||||
database.reader.passwordSecret.urlKey: connection-url
|
||||
database.reader.useIAMAuth: true
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: "database.reader.passwordSecret.urlKey cannot be combined with host, hostKey, dbname, schema, useIAMAuth, or useAzureEntraAuth"
|
||||
|
||||
- it: reader urlKey requires a Secret name
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
database.reader.passwordSecret.urlKey: connection-url
|
||||
database.reader.passwordSecret.name: ""
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: "database.reader.passwordSecret.name is required when hostKey or urlKey is set"
|
||||
|
|
|
|||
|
|
@ -21,6 +21,27 @@ tests:
|
|||
name: REDIS_CLUSTER_NODES
|
||||
any: true
|
||||
|
||||
- it: gateway renders without redis env vars when passwordSecret is null
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
redis.passwordSecret: null
|
||||
asserts:
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: REDIS_HOST
|
||||
any: true
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: REDIS_URL
|
||||
any: true
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: REDIS_PASSWORD
|
||||
any: true
|
||||
|
||||
- it: gateway emits host, port and password when redis is configured
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
|
|
@ -68,6 +89,27 @@ tests:
|
|||
name: redis-secret
|
||||
key: redis-password
|
||||
|
||||
- it: backend renders without redis env vars when passwordSecret is null
|
||||
template: backend/deployment.yaml
|
||||
set:
|
||||
redis.passwordSecret: null
|
||||
asserts:
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: REDIS_HOST
|
||||
any: true
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: REDIS_URL
|
||||
any: true
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: REDIS_PASSWORD
|
||||
any: true
|
||||
|
||||
- it: gateway omits REDIS_PASSWORD for an auth-less redis
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
|
|
@ -107,3 +149,174 @@ tests:
|
|||
content:
|
||||
name: REDIS_CLUSTER_NODES
|
||||
any: true
|
||||
|
||||
- it: gateway gets redis host from hostKey and keeps port and password
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
redis.host: ""
|
||||
redis.port: 6380
|
||||
redis.passwordSecret.name: redis-secret
|
||||
redis.passwordSecret.hostKey: hostname
|
||||
redis.passwordSecret.passwordKey: redis-password
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: REDIS_HOST
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: redis-secret
|
||||
key: hostname
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: REDIS_PORT
|
||||
value: "6380"
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: REDIS_PASSWORD
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: redis-secret
|
||||
key: redis-password
|
||||
|
||||
- it: backend gets redis host from hostKey with the same coordination env
|
||||
template: backend/deployment.yaml
|
||||
set:
|
||||
redis.host: ""
|
||||
redis.passwordSecret.name: redis-secret
|
||||
redis.passwordSecret.hostKey: hostname
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: REDIS_HOST
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: redis-secret
|
||||
key: hostname
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: REDIS_PORT
|
||||
value: "6379"
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: REDIS_PASSWORD
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: redis-secret
|
||||
key: password
|
||||
|
||||
- it: gateway expands the redis cluster seed from hostKey at runtime
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
redis.host: ""
|
||||
redis.passwordSecret.name: redis-secret
|
||||
redis.passwordSecret.hostKey: hostname
|
||||
redis.cluster: true
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: REDIS_CLUSTER_NODES
|
||||
value: '[{"host":"$(REDIS_HOST)","port":6379}]'
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].env[8].name
|
||||
value: REDIS_HOST
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].env[11].name
|
||||
value: REDIS_CLUSTER_NODES
|
||||
|
||||
- it: gateway gets REDIS_URL from urlKey without discrete Redis env
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
redis.passwordSecret.name: redis-secret
|
||||
redis.passwordSecret.urlKey: connection-url
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: REDIS_URL
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: redis-secret
|
||||
key: connection-url
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: REDIS_HOST
|
||||
any: true
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: REDIS_PORT
|
||||
any: true
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: REDIS_PASSWORD
|
||||
any: true
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: REDIS_CLUSTER_NODES
|
||||
any: true
|
||||
|
||||
- it: redis rejects host together with hostKey
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
redis.host: redis.example.com
|
||||
redis.passwordSecret.hostKey: hostname
|
||||
redis.passwordSecret.name: redis-secret
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: "redis.host and redis.passwordSecret.hostKey cannot both be set"
|
||||
|
||||
- it: redis urlKey rejects host
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
redis.host: redis.example.com
|
||||
redis.passwordSecret.urlKey: connection-url
|
||||
redis.passwordSecret.name: redis-secret
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: "redis.passwordSecret.urlKey cannot be combined with host, hostKey, or cluster"
|
||||
|
||||
- it: redis urlKey rejects hostKey
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
redis.passwordSecret.hostKey: hostname
|
||||
redis.passwordSecret.urlKey: connection-url
|
||||
redis.passwordSecret.name: redis-secret
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: "redis.passwordSecret.urlKey cannot be combined with host, hostKey, or cluster"
|
||||
|
||||
- it: redis urlKey rejects cluster mode
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
redis.cluster: true
|
||||
redis.passwordSecret.urlKey: connection-url
|
||||
redis.passwordSecret.name: redis-secret
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: "redis.passwordSecret.urlKey cannot be combined with host, hostKey, or cluster"
|
||||
|
||||
- it: redis hostKey requires a Secret name
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
redis.passwordSecret.hostKey: hostname
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: "redis.passwordSecret.name is required when hostKey or urlKey is set"
|
||||
|
||||
- it: redis urlKey requires a Secret name
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
redis.passwordSecret.urlKey: connection-url
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: "redis.passwordSecret.name is required when hostKey or urlKey is set"
|
||||
|
|
|
|||
|
|
@ -23,7 +23,11 @@ ingress:
|
|||
# wildcard pathType, so that rule could never match there.
|
||||
controller: alb
|
||||
annotations: {}
|
||||
host: "" # optional; if set, becomes the rule's host
|
||||
host: "" # optional; if set, becomes the single ingress rule's host
|
||||
componentHosts:
|
||||
ui: ""
|
||||
backend: ""
|
||||
gateway: ""
|
||||
tls: []
|
||||
# Extra HTTP paths appended to the ingress rule. Additive: every built-in
|
||||
# UI / gateway / backend path is still rendered, these entries are placed
|
||||
|
|
@ -208,6 +212,8 @@ database:
|
|||
name: litellm-writer-secret
|
||||
usernameKey: username
|
||||
passwordKey: password
|
||||
hostKey: "" # Secret key containing the writer database host
|
||||
urlKey: "" # Secret key containing the complete writer database URL
|
||||
# libpq sslmode / sslrootcert applied to the writer and reader URLs (Prisma and the
|
||||
# in-container PgBouncer); e.g. verify-full with /etc/ssl/certs/ca-certificates.crt for AWS RDS.
|
||||
# sslRootCert on its own implies sslMode verify-full
|
||||
|
|
@ -229,6 +235,8 @@ database:
|
|||
name: litellm-reader-secret
|
||||
usernameKey: username
|
||||
passwordKey: password
|
||||
hostKey: "" # Secret key containing the reader database host
|
||||
urlKey: "" # Secret key containing the complete reader database URL
|
||||
|
||||
# In-container connection pool (PgBouncer, transaction mode) shared by every
|
||||
# gateway worker in the pod. Without it each of the `gateway.numWorkers`
|
||||
|
|
@ -274,6 +282,8 @@ redis:
|
|||
passwordSecret:
|
||||
name: "" # Leave empty for auth-less Redis
|
||||
passwordKey: password
|
||||
hostKey: "" # Secret key containing the Redis host
|
||||
urlKey: "" # Secret key containing the complete Redis URL
|
||||
|
||||
# ---------- gateway (LLM data plane) ----------
|
||||
gateway:
|
||||
|
|
@ -493,6 +503,7 @@ gateway:
|
|||
backend:
|
||||
enabled: true
|
||||
logLevel: INFO
|
||||
gatewayUrl: ""
|
||||
extraEnv: []
|
||||
envConfigMaps: []
|
||||
envSecrets: []
|
||||
|
|
@ -575,9 +586,10 @@ ui:
|
|||
service:
|
||||
type: ClusterIP
|
||||
port: 3000
|
||||
# The dashboard expects to know where to reach the backend API. Set this to
|
||||
# the externally-routable URL (typically the ingress host + /api or similar).
|
||||
# The backend's public origin, passed to the UI only when ingress.componentHosts is configured.
|
||||
backendUrl: ""
|
||||
# Public origin of the UI, used for backend and gateway CORS.
|
||||
publicUrl: ""
|
||||
resources:
|
||||
requests:
|
||||
cpu: 500m
|
||||
|
|
|
|||
|
|
@ -1,5 +1,4 @@
|
|||
import random
|
||||
import re
|
||||
import time
|
||||
from collections.abc import Generator, Mapping
|
||||
from contextlib import contextmanager
|
||||
|
|
@ -10,12 +9,11 @@ from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
|
|||
from litellm_proxy_extras._logging import logger
|
||||
from litellm_proxy_extras.prisma_toolchain import (
|
||||
MIGRATION_LOCK_TIMEOUT_ENV_VAR,
|
||||
migration_ddl_lock_timeout,
|
||||
migration_lock_timeout,
|
||||
)
|
||||
|
||||
MIGRATION_LOCK_KEY: Final = int.from_bytes(b"llm_mig2", "big")
|
||||
_LOCK_TIMEOUT_PINNED_RE: Final = re.compile(r"(?:-c\s*|--)lock_timeout=")
|
||||
PRISMA_ADVISORY_LOCK_KEY: Final = 72707369
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import psycopg
|
||||
|
|
@ -30,29 +28,12 @@ def _migration_url(database_url: str, direct_url: "str | None") -> str:
|
|||
return urlunsplit(direct._replace(query=urlencode((*parameters, ("schema", schema)))))
|
||||
|
||||
|
||||
def _with_ddl_lock_timeout(url: str) -> str:
|
||||
split: Final = urlsplit(url)
|
||||
pairs: Final = tuple(parse_qsl(split.query))
|
||||
if any(key == "pgbouncer" and value == "true" for key, value in pairs):
|
||||
return url
|
||||
existing: Final = next((value for key, value in pairs if key == "options"), "")
|
||||
if _LOCK_TIMEOUT_PINNED_RE.search(existing):
|
||||
return url
|
||||
timeout_ms: Final = max(1, int(migration_ddl_lock_timeout() * 1000))
|
||||
options: Final = f"{existing} -c lock_timeout={timeout_ms}".strip()
|
||||
rewritten: Final = tuple((key, value) for key, value in pairs if key != "options") + (("options", options),)
|
||||
return urlunsplit(split._replace(query=urlencode(rewritten)))
|
||||
|
||||
|
||||
def migration_environment(environment: Mapping[str, str]) -> Mapping[str, str]:
|
||||
database_url: Final = environment.get("DATABASE_URL")
|
||||
if not database_url:
|
||||
return environment
|
||||
migration_url: Final = _migration_url(database_url, environment.get("DIRECT_URL"))
|
||||
return {
|
||||
**environment,
|
||||
"DATABASE_URL": _with_ddl_lock_timeout(migration_url),
|
||||
}
|
||||
return {**environment, "DATABASE_URL": migration_url, "PRISMA_SCHEMA_DISABLE_ADVISORY_LOCK": "1"}
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -75,10 +56,15 @@ class MigrationCoordinator:
|
|||
def check_connection(self) -> None:
|
||||
self.connection.execute("SELECT 1")
|
||||
|
||||
def release_statement_snapshot(self) -> None:
|
||||
# a parameterized statement leaves its portal snapshot registered on the backend until the
|
||||
# next statement, and CREATE INDEX CONCURRENTLY waits for every such snapshot to go away
|
||||
self.connection.execute("SELECT 1")
|
||||
|
||||
def acquire_prisma_lock(self) -> None:
|
||||
deadline: Final = time.monotonic() + migration_lock_timeout()
|
||||
while time.monotonic() < deadline:
|
||||
if _try_lock(self.connection, 72707369):
|
||||
if _try_lock(self.connection, PRISMA_ADVISORY_LOCK_KEY):
|
||||
return
|
||||
time.sleep(min(random.uniform(0.5, 1.5), max(0.0, deadline - time.monotonic())))
|
||||
raise RuntimeError(
|
||||
|
|
@ -115,6 +101,17 @@ def migration_lock(database_url: str) -> Generator[MigrationCoordinator, None, N
|
|||
)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def prisma_migration_lock(database_url: str) -> Generator[None, None, None]:
|
||||
"""Hold the coordinator lock and Prisma's own lock key as transaction-level locks on one
|
||||
pinned connection while a Prisma command runs with its session advisory lock disabled, so a
|
||||
transaction pooler can never strand that lock on an idle backend after the command exits."""
|
||||
with migration_lock(database_url) as coordinator:
|
||||
coordinator.acquire_prisma_lock()
|
||||
coordinator.release_statement_snapshot()
|
||||
yield
|
||||
|
||||
|
||||
@contextmanager
|
||||
def held_migration_lock(connection: "psycopg.Connection[tuple[object, ...]]") -> Generator[bool, None, None]:
|
||||
"""A session-level, non-blocking hold of the migration coordinator lock on an autocommit
|
||||
|
|
|
|||
|
|
@ -0,0 +1,14 @@
|
|||
SET LOCAL lock_timeout = '10s';
|
||||
|
||||
CREATE TABLE IF NOT EXISTS "LiteLLM_TelemetryReport" (
|
||||
"id" TEXT NOT NULL,
|
||||
"window_start" DOUBLE PRECISION NOT NULL,
|
||||
"window_end" DOUBLE PRECISION NOT NULL,
|
||||
"report" JSONB NOT NULL,
|
||||
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
CONSTRAINT "LiteLLM_TelemetryReport_pkey" PRIMARY KEY ("id")
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_TelemetryReport_window_end_idx" ON "LiteLLM_TelemetryReport"("window_end");
|
||||
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_TelemetryReport_created_at_idx" ON "LiteLLM_TelemetryReport"("created_at");
|
||||
|
|
@ -0,0 +1,4 @@
|
|||
SET LOCAL lock_timeout = '10s';
|
||||
|
||||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_AccessGroupTable" ADD COLUMN IF NOT EXISTS "mcp_tool_permissions" JSONB;
|
||||
|
|
@ -0,0 +1,4 @@
|
|||
SET LOCAL lock_timeout = '10s';
|
||||
|
||||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_AccessGroupTable" ADD COLUMN IF NOT EXISTS "access_skill_names" TEXT[] DEFAULT ARRAY[]::TEXT[];
|
||||
|
|
@ -0,0 +1,4 @@
|
|||
SET LOCAL lock_timeout = '10s';
|
||||
|
||||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_ManagedObjectTable" ADD COLUMN IF NOT EXISTS "end_user_id" TEXT;
|
||||
|
|
@ -60,13 +60,11 @@ PRISMA_COMMAND_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_COMMAND_TIMEOUT"
|
|||
PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_BOOTSTRAP_TIMEOUT"
|
||||
PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_MIGRATE_DEPLOY_TIMEOUT"
|
||||
MIGRATION_LOCK_TIMEOUT_ENV_VAR = "LITELLM_MIGRATION_LOCK_TIMEOUT"
|
||||
MIGRATION_DDL_LOCK_TIMEOUT_ENV_VAR = "LITELLM_MIGRATION_DDL_LOCK_TIMEOUT"
|
||||
NODEENV_CACHE_DIR_ENV_VAR = "PRISMA_NODEENV_CACHE_DIR"
|
||||
|
||||
DEFAULT_PRISMA_COMMAND_TIMEOUT = 60.0
|
||||
DEFAULT_PRISMA_BOOTSTRAP_TIMEOUT = 600.0
|
||||
DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT = 600.0
|
||||
DEFAULT_MIGRATION_DDL_LOCK_TIMEOUT = 10.0
|
||||
|
||||
BOOTSTRAP_ARG = "--version"
|
||||
PRISMA_CONSOLE_SCRIPT = "prisma"
|
||||
|
|
@ -113,11 +111,6 @@ def migration_lock_timeout() -> float:
|
|||
return _timeout_from_env(MIGRATION_LOCK_TIMEOUT_ENV_VAR, 600.0)
|
||||
|
||||
|
||||
def migration_ddl_lock_timeout() -> float:
|
||||
"""Seconds one migration statement may wait on a table lock before Postgres cancels it."""
|
||||
return _timeout_from_env(MIGRATION_DDL_LOCK_TIMEOUT_ENV_VAR, DEFAULT_MIGRATION_DDL_LOCK_TIMEOUT)
|
||||
|
||||
|
||||
def prisma_bootstrap_timeout() -> float:
|
||||
"""Seconds the one-time Node toolchain install may run for."""
|
||||
return _timeout_from_env(
|
||||
|
|
|
|||
|
|
@ -1179,6 +1179,7 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t
|
|||
org_id String? // creating key's organization at submission time; CheckBatchCost bills org spend against it
|
||||
api_key String?
|
||||
request_tags Json? @default("[]")
|
||||
end_user_id String?
|
||||
updated_at DateTime @updatedAt
|
||||
updated_by String?
|
||||
|
||||
|
|
@ -1588,7 +1589,9 @@ model LiteLLM_AccessGroupTable {
|
|||
// Resource memberships - explicit arrays per type
|
||||
access_model_names String[] @default([])
|
||||
access_mcp_server_ids String[] @default([])
|
||||
mcp_tool_permissions Json?
|
||||
access_agent_ids String[] @default([])
|
||||
access_skill_names String[] @default([])
|
||||
|
||||
assigned_team_ids String[] @default([])
|
||||
assigned_key_ids String[] @default([])
|
||||
|
|
@ -2045,3 +2048,14 @@ model LiteLLM_LensTraceSignal {
|
|||
|
||||
@@id([trace_id, trace_ref])
|
||||
}
|
||||
|
||||
model LiteLLM_TelemetryReport {
|
||||
id String @id
|
||||
window_start Float
|
||||
window_end Float
|
||||
report Json
|
||||
created_at DateTime @default(now())
|
||||
|
||||
@@index([window_end])
|
||||
@@index([created_at])
|
||||
}
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ import tempfile
|
|||
import threading
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
from contextlib import AbstractContextManager, nullcontext
|
||||
from dataclasses import dataclass, replace
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Final, Optional, Protocol, Union
|
||||
|
|
@ -18,7 +19,6 @@ from litellm_proxy_extras import prisma_toolchain
|
|||
from litellm_proxy_extras._logging import logger
|
||||
from litellm_proxy_extras.migration_lock import held_migration_lock
|
||||
from litellm_proxy_extras.prisma_toolchain import (
|
||||
MIGRATION_DDL_LOCK_TIMEOUT_ENV_VAR,
|
||||
PRISMA_COMMAND_TIMEOUT_ENV_VAR,
|
||||
PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR,
|
||||
ensure_prisma_toolchain,
|
||||
|
|
@ -1022,7 +1022,11 @@ class ProxyExtrasDBManager:
|
|||
def _setup_database_v2(use_migrate: bool) -> bool:
|
||||
if not use_migrate:
|
||||
return ProxyExtrasDBManager._run_database_v2(False)
|
||||
from litellm_proxy_extras.migration_lock import migration_environment, migration_lock
|
||||
from litellm_proxy_extras.migration_lock import (
|
||||
migration_environment,
|
||||
migration_lock,
|
||||
prisma_migration_lock,
|
||||
)
|
||||
from litellm_proxy_extras.migration_recovery import (
|
||||
baseline_current_schema,
|
||||
recover_completed_migration,
|
||||
|
|
@ -1056,7 +1060,9 @@ class ProxyExtrasDBManager:
|
|||
migration_environment(_get_prisma_env()),
|
||||
)
|
||||
|
||||
while not ProxyExtrasDBManager._run_database_v2(True, recover_completed, baseline_existing):
|
||||
while not ProxyExtrasDBManager._run_database_v2(
|
||||
True, recover_completed, baseline_existing, lambda: prisma_migration_lock(lock_url)
|
||||
):
|
||||
continue
|
||||
return True
|
||||
|
||||
|
|
@ -1065,6 +1071,7 @@ class ProxyExtrasDBManager:
|
|||
use_migrate: bool,
|
||||
recover_completed: Callable[[str], bool] = lambda name: False,
|
||||
baseline_existing: "Callable[[str], None] | None" = None,
|
||||
deploy_lock: Callable[[], AbstractContextManager[None]] = nullcontext,
|
||||
) -> bool:
|
||||
"""
|
||||
v2 migration resolver (opt-in via --use_v2_migration_resolver).
|
||||
|
|
@ -1127,11 +1134,12 @@ class ProxyExtrasDBManager:
|
|||
while not budget.exhausted:
|
||||
attempt_started = time.monotonic()
|
||||
try:
|
||||
result = prisma_toolchain.run_prisma(
|
||||
[_get_prisma_command(), "migrate", "deploy"],
|
||||
timeout=deploy_timeout,
|
||||
env=migration_env,
|
||||
)
|
||||
with deploy_lock():
|
||||
result = prisma_toolchain.run_prisma(
|
||||
[_get_prisma_command(), "migrate", "deploy"],
|
||||
timeout=deploy_timeout,
|
||||
env=migration_env,
|
||||
)
|
||||
logger.info(f"prisma migrate deploy stdout: {result.stdout}")
|
||||
return True
|
||||
|
||||
|
|
@ -1156,6 +1164,7 @@ class ProxyExtrasDBManager:
|
|||
e,
|
||||
budget,
|
||||
time.monotonic() - attempt_started,
|
||||
deploy_lock=deploy_lock,
|
||||
)
|
||||
|
||||
if next_budget.attempts_left < budget.attempts_left:
|
||||
|
|
@ -1165,10 +1174,9 @@ class ProxyExtrasDBManager:
|
|||
raise RuntimeError(
|
||||
f"Database migration failed after {MAX_MIGRATE_DEPLOY_ATTEMPTS} "
|
||||
"attempts that made no progress (timeouts, lock timeouts or deadlock retries). "
|
||||
"Check database connectivity, load, long-running transactions on the migrated tables, "
|
||||
"and _prisma_migrations ledger state. Raise "
|
||||
f"{PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR} if the attempts timed out, or "
|
||||
f"{MIGRATION_DDL_LOCK_TIMEOUT_ENV_VAR} if they timed out waiting for a table lock."
|
||||
"Check database connectivity, load and _prisma_migrations ledger state. "
|
||||
"If a migration hit SET LOCAL lock_timeout, check for long-running transactions on the migrated tables. "
|
||||
f"Raise {PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR} if the attempts timed out."
|
||||
)
|
||||
finally:
|
||||
os.chdir(original_dir)
|
||||
|
|
@ -1189,15 +1197,19 @@ class ProxyExtrasDBManager:
|
|||
return match.group(1) if match else None
|
||||
|
||||
@staticmethod
|
||||
def _v2_roll_back_migration_best_effort(migration_name: str) -> None:
|
||||
def _v2_roll_back_migration_best_effort(
|
||||
migration_name: str,
|
||||
deploy_lock: Callable[[], AbstractContextManager[None]] = nullcontext,
|
||||
) -> None:
|
||||
from litellm_proxy_extras.migration_lock import migration_environment
|
||||
|
||||
try:
|
||||
prisma_toolchain.run_prisma(
|
||||
[_get_prisma_command(), "migrate", "resolve", "--rolled-back", migration_name],
|
||||
timeout=prisma_command_timeout(),
|
||||
env=migration_environment(_get_prisma_env()),
|
||||
)
|
||||
with deploy_lock():
|
||||
prisma_toolchain.run_prisma(
|
||||
[_get_prisma_command(), "migrate", "resolve", "--rolled-back", migration_name],
|
||||
timeout=prisma_command_timeout(),
|
||||
env=migration_environment(_get_prisma_env()),
|
||||
)
|
||||
except (subprocess.CalledProcessError, subprocess.TimeoutExpired):
|
||||
pass
|
||||
|
||||
|
|
@ -1206,6 +1218,7 @@ class ProxyExtrasDBManager:
|
|||
error: subprocess.CalledProcessError,
|
||||
budget: "_MigrateAttemptBudget",
|
||||
attempt_seconds: float = 0.0,
|
||||
deploy_lock: Callable[[], AbstractContextManager[None]] = nullcontext,
|
||||
) -> "_MigrateAttemptBudget":
|
||||
"""Recover from one failed `prisma migrate deploy`, and price the pass.
|
||||
|
||||
|
|
@ -1227,14 +1240,14 @@ class ProxyExtrasDBManager:
|
|||
"deadlock race, rolling its ledger row back and retrying",
|
||||
migration_name,
|
||||
)
|
||||
ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name)
|
||||
ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name, deploy_lock)
|
||||
return budget.spend()
|
||||
if ledger_logs and _MIGRATION_LOCK_TIMEOUT_MARKER in ledger_logs:
|
||||
logger.info(
|
||||
"Migration %s timed out waiting for a table lock, rolling its ledger row back and retrying",
|
||||
migration_name,
|
||||
)
|
||||
ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name)
|
||||
ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name, deploy_lock)
|
||||
return budget.spend()
|
||||
if ProxyExtrasDBManager._failed_migration_recovered(migration_name, started_at):
|
||||
logger.info(
|
||||
|
|
@ -1278,17 +1291,17 @@ class ProxyExtrasDBManager:
|
|||
"rolling its ledger row back and retrying",
|
||||
migration_name,
|
||||
)
|
||||
ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name)
|
||||
ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name, deploy_lock)
|
||||
return budget.spend()
|
||||
|
||||
if migration_name and _MIGRATION_LOCK_TIMEOUT_MARKER in stderr:
|
||||
ProxyExtrasDBManager._log_migration_lock_holders(migration_name)
|
||||
logger.warning(
|
||||
"Migration %s timed out waiting for a table lock, rolling its ledger row back and retrying. "
|
||||
f"Raise {MIGRATION_DDL_LOCK_TIMEOUT_ENV_VAR} if the database needs longer.",
|
||||
"Migration %s timed out waiting for a table lock after hitting SET LOCAL lock_timeout, rolling "
|
||||
"its ledger row back and retrying. Check for long-running transactions on the migrated tables.",
|
||||
migration_name,
|
||||
)
|
||||
ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name)
|
||||
ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name, deploy_lock)
|
||||
return budget.spend()
|
||||
|
||||
raise RuntimeError(
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm-proxy-extras"
|
||||
version = "0.4.108"
|
||||
version = "0.4.109"
|
||||
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.9"
|
||||
|
|
@ -30,7 +30,7 @@ required-version = ">=0.10.9"
|
|||
module-root = ""
|
||||
|
||||
[tool.commitizen]
|
||||
version = "0.4.108"
|
||||
version = "0.4.109"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
"../pyproject.toml:litellm-proxy-extras==",
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ Use the derived conversions directly (`<&'static str>::from(x)` / `.into()`, `st
|
|||
- `crates/<crate>/tests/` is only for tests that drive several modules together end to end, such as a full request through a route. A test of one type or function never goes there
|
||||
- A helper shared by several inline test modules goes in a `#[cfg(test)] mod test_support;` at the crate root, and only once a second module needs it
|
||||
- A test for another crate's item belongs in that crate, not in a downstream one
|
||||
- Tests that call a real provider go in `crates/<crate>/tests/live/<provider>.rs` behind a `tests/live/main.rs`. Each reads the provider's key from the environment through an `Option` fixture and returns early when it is unset, never `#[ignore]`, so `cargo test` stays offline until the key is exported
|
||||
- Never set `autotests = false` or hand-list `[[test]]` targets. Every file directly under `tests/` is discovered by cargo, and a shared helper goes in `tests/<name>/mod.rs` or `tests/<subject>/support.rs` so it is not picked up as a test crate of its own
|
||||
|
||||
## Test fixtures and cases
|
||||
|
|
|
|||
8
litellm-rust/Cargo.lock
generated
8
litellm-rust/Cargo.lock
generated
|
|
@ -3406,6 +3406,7 @@ dependencies = [
|
|||
name = "litellm-auth-gcp"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"base64 0.22.1",
|
||||
"gcp_auth",
|
||||
"google-cloud-auth",
|
||||
"http 1.4.2",
|
||||
|
|
@ -3922,13 +3923,16 @@ version = "0.1.0"
|
|||
dependencies = [
|
||||
"bytes",
|
||||
"futures-util",
|
||||
"litellm-coroutine",
|
||||
"litellm-host",
|
||||
"parking_lot",
|
||||
"pyo3",
|
||||
"pyo3-async-runtimes",
|
||||
"pythonize",
|
||||
"rstest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
|
|
@ -4024,6 +4028,7 @@ dependencies = [
|
|||
"litellm-cache-memory",
|
||||
"litellm-cache-response",
|
||||
"litellm-core-utils",
|
||||
"litellm-framer",
|
||||
"litellm-host",
|
||||
"litellm-host-native",
|
||||
"litellm-http",
|
||||
|
|
@ -4084,6 +4089,7 @@ dependencies = [
|
|||
"litellm-cache",
|
||||
"litellm-cache-memory",
|
||||
"litellm-cache-response",
|
||||
"litellm-core-utils",
|
||||
"litellm-host",
|
||||
"litellm-host-native",
|
||||
"litellm-http",
|
||||
|
|
@ -4107,9 +4113,11 @@ name = "litellm-inference-testing"
|
|||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"futures-util",
|
||||
"litellm-auth",
|
||||
"litellm-http",
|
||||
"litellm-inference",
|
||||
"litellm-secrets",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
|
|||
|
|
@ -80,6 +80,7 @@ http = "1"
|
|||
google-cloud-auth = { version = "1.16.0", default-features = false }
|
||||
jsonwebtoken = { version = "11.1.0", default-features = false }
|
||||
hyper-util = { version = "0.1.20", default-features = false, features = ["client-proxy"] }
|
||||
parking_lot = "0.12.5"
|
||||
proptest = "1.7.0"
|
||||
pyo3 = "0.29.2"
|
||||
pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] }
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ google-sdk = ["dep:google-cloud-auth", "dep:http"]
|
|||
[dependencies]
|
||||
litellm-auth-types.workspace = true
|
||||
|
||||
base64.workspace = true
|
||||
moka.workspace = true
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
use std::{future::Future, path::Path, pin::Pin, sync::Arc};
|
||||
use std::{future::Future, pin::Pin, sync::Arc};
|
||||
|
||||
use gcp_auth::{CustomServiceAccount, TokenProvider};
|
||||
use litellm_auth_types::{CredentialPlacement, Error, ErrorDetail, http::apply_credential};
|
||||
|
|
@ -8,7 +8,8 @@ use sha2::{Digest, Sha256};
|
|||
|
||||
use crate::{
|
||||
config::{
|
||||
CredentialSource, VertexConfig, credential_source, get_vertex_ai_project, non_empty_env,
|
||||
CredentialSource, VertexConfig, credential_source, get_vertex_ai_project,
|
||||
is_base64_encoded_json, names_a_credentials_file, non_empty_env,
|
||||
},
|
||||
constants::{
|
||||
CLOUD_PLATFORM_SCOPE, GOOGLE_OAUTH_TOKEN_ENDPOINT, VERTEX_AI_API_KEY_ENV,
|
||||
|
|
@ -176,14 +177,7 @@ impl VertexProviderLoader for GcpProviderLoader {
|
|||
.map_err(auth_acquisition_error)?,
|
||||
),
|
||||
CredentialSource::Trusted(configured) => {
|
||||
let configured = configured.expose();
|
||||
let service_account = if Path::new(configured).is_file() {
|
||||
CustomServiceAccount::from_file(configured)
|
||||
} else {
|
||||
CustomServiceAccount::from_json(configured)
|
||||
}
|
||||
.map_err(auth_acquisition_error)?;
|
||||
Arc::new(service_account)
|
||||
Arc::new(load_trusted_service_account(configured.expose())?)
|
||||
}
|
||||
CredentialSource::ApplicationCredentials(path) => {
|
||||
Arc::new(CustomServiceAccount::from_file(path).map_err(auth_acquisition_error)?)
|
||||
|
|
@ -240,6 +234,23 @@ enum CredentialCacheKey {
|
|||
Adc,
|
||||
}
|
||||
|
||||
const BASE64_CREDENTIALS: &str = "The `vertex_credentials` value is base64-encoded JSON, not the JSON itself. \
|
||||
Set it to a readable file path, or to the credentials JSON itself.";
|
||||
|
||||
fn load_trusted_service_account(configured: &str) -> Result<CustomServiceAccount, Error> {
|
||||
if is_base64_encoded_json(configured) {
|
||||
return Err(Error::InvalidConfiguration(ErrorDetail::Message(
|
||||
BASE64_CREDENTIALS.into(),
|
||||
)));
|
||||
}
|
||||
if names_a_credentials_file(configured) {
|
||||
CustomServiceAccount::from_file(configured)
|
||||
} else {
|
||||
CustomServiceAccount::from_json(configured)
|
||||
}
|
||||
.map_err(auth_acquisition_error)
|
||||
}
|
||||
|
||||
fn auth_acquisition_error(error: gcp_auth::Error) -> Error {
|
||||
Error::CredentialAcquisition(litellm_auth_types::ErrorDetail::failed(
|
||||
"Vertex AI credentials",
|
||||
|
|
@ -264,6 +275,73 @@ mod tests {
|
|||
assert_eq!(validate_request_credentials(credentials).is_ok(), accepted);
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::missing_file(
|
||||
"/nonexistent/vertex-key.json",
|
||||
"failed to open application credentials file"
|
||||
)]
|
||||
#[case::malformed_json(
|
||||
" {\"type\": \"service_account\",",
|
||||
"failed to deserialize ApplicationCredentials"
|
||||
)]
|
||||
#[case::quoted_json(
|
||||
"'{\"type\": \"service_account\"}'",
|
||||
"failed to deserialize ApplicationCredentials"
|
||||
)]
|
||||
#[tokio::test]
|
||||
async fn trusted_credentials_are_read_as_a_file_unless_they_are_json(
|
||||
#[case] configured: &str,
|
||||
#[case] expected: &str,
|
||||
) {
|
||||
let Err(error) = GcpProviderLoader
|
||||
.load(CredentialSource::Trusted(SecretValue::new(configured)))
|
||||
.await
|
||||
else {
|
||||
panic!("{configured:?} should not load");
|
||||
};
|
||||
assert!(error.to_string().contains(expected), "{error}");
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::padded("eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCJ9")]
|
||||
#[case::quoted_prefixed_unpadded("\"base64:eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCJ9\"")]
|
||||
#[case::url_safe_wrapped("eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCIsICJr\nZXkiOiAiPz8_In0=")]
|
||||
#[tokio::test]
|
||||
async fn base64_encoded_credentials_are_named_as_base64(#[case] configured: &str) {
|
||||
assert!(
|
||||
trusted_load_error(configured)
|
||||
.await
|
||||
.contains("is base64-encoded JSON, not the JSON itself")
|
||||
);
|
||||
}
|
||||
|
||||
async fn trusted_load_error(configured: &str) -> String {
|
||||
match GcpProviderLoader
|
||||
.load(CredentialSource::Trusted(SecretValue::new(configured)))
|
||||
.await
|
||||
{
|
||||
Ok(_) => panic!("{configured:?} should not load"),
|
||||
Err(error) => error.to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[tokio::test]
|
||||
async fn a_relative_key_file_named_like_json_is_read_as_a_file() {
|
||||
let key = tempfile::Builder::new()
|
||||
.prefix("{vertex")
|
||||
.suffix(".json")
|
||||
.tempfile_in(".")
|
||||
.unwrap();
|
||||
std::fs::write(key.path(), r#"{"project_id":"file-project"}"#).unwrap();
|
||||
let relative_name = key.path().file_name().unwrap().to_str().unwrap();
|
||||
let absolute_path = std::fs::canonicalize(key.path()).unwrap();
|
||||
assert_eq!(
|
||||
trusted_load_error(relative_name).await,
|
||||
trusted_load_error(absolute_path.to_str().unwrap()).await
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
fn inline_and_trusted_credentials_never_share_a_cache_entry() {
|
||||
assert_ne!(
|
||||
|
|
|
|||
|
|
@ -1,5 +1,9 @@
|
|||
use std::{collections::BTreeMap, path::Path};
|
||||
|
||||
use base64::{
|
||||
Engine, alphabet,
|
||||
engine::{DecodePaddingMode, GeneralPurpose, GeneralPurposeConfig},
|
||||
};
|
||||
use litellm_auth_types::{Error, ErrorDetail, InputSource, SecretValue, Sourced};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
|
|
@ -80,9 +84,10 @@ pub fn get_vertex_ai_project_from_credentials(
|
|||
CredentialSource::Inline(configured) => configured.expose().to_string(),
|
||||
CredentialSource::Trusted(configured) => {
|
||||
let configured = configured.expose();
|
||||
match std::fs::read_to_string(configured) {
|
||||
Ok(contents) if Path::new(configured).is_file() => contents,
|
||||
_ => configured.to_string(),
|
||||
if names_a_credentials_file(configured) {
|
||||
std::fs::read_to_string(configured).ok()?
|
||||
} else {
|
||||
configured.to_string()
|
||||
}
|
||||
}
|
||||
CredentialSource::ApplicationCredentials(path) => std::fs::read_to_string(path).ok()?,
|
||||
|
|
@ -105,6 +110,32 @@ pub fn get_vertex_ai_location(
|
|||
.or_else(|| VertexParams::LOCATION.resolve(&|_| None, env_lookup))
|
||||
}
|
||||
|
||||
const INLINE_JSON_LEADING_NOISE: [char; 7] = [' ', '\t', '\r', '\n', '\u{feff}', '"', '\''];
|
||||
|
||||
fn is_inline_json(configured: &str) -> bool {
|
||||
configured
|
||||
.trim_start_matches(INLINE_JSON_LEADING_NOISE)
|
||||
.starts_with('{')
|
||||
}
|
||||
|
||||
pub(crate) fn names_a_credentials_file(configured: &str) -> bool {
|
||||
!is_inline_json(configured) || Path::new(configured).is_file()
|
||||
}
|
||||
|
||||
const BASE64_ANY_PADDING: GeneralPurpose = GeneralPurpose::new(
|
||||
&alphabet::STANDARD,
|
||||
GeneralPurposeConfig::new().with_decode_padding_mode(DecodePaddingMode::Indifferent),
|
||||
);
|
||||
|
||||
pub(crate) fn is_base64_encoded_json(configured: &str) -> bool {
|
||||
let compact: String = configured.split_whitespace().collect();
|
||||
let quoted = compact.trim_matches(['"', '\'']);
|
||||
let encoded = quoted.strip_prefix("base64:").unwrap_or(quoted);
|
||||
BASE64_ANY_PADDING
|
||||
.decode(encoded.replace('-', "+").replace('_', "/"))
|
||||
.is_ok_and(|decoded| serde_json::from_slice::<Map<String, Value>>(&decoded).is_ok())
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub enum CredentialSource {
|
||||
Inline(SecretValue),
|
||||
|
|
|
|||
|
|
@ -93,3 +93,22 @@ fn project_is_read_from_the_credentials_that_would_authenticate() {
|
|||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn project_is_read_from_a_relative_key_file_named_like_json() {
|
||||
let key = tempfile::Builder::new()
|
||||
.prefix("{vertex")
|
||||
.suffix(".json")
|
||||
.tempfile_in(".")
|
||||
.unwrap();
|
||||
std::fs::write(key.path(), r#"{"project_id":"file-project"}"#).unwrap();
|
||||
let relative_name = key.path().file_name().unwrap().to_str().unwrap();
|
||||
let from_file = VertexConfig::from_params(&VertexParams {
|
||||
vertex_credentials: Some(relative_name.into()),
|
||||
..VertexParams::default()
|
||||
});
|
||||
assert_eq!(
|
||||
get_vertex_ai_project_from_credentials(&from_file, &|_| None).as_deref(),
|
||||
Some("file-project")
|
||||
);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ use litellm_host::{
|
|||
interceptors::{RawResponse, RequestContext, WireRequest},
|
||||
lifecycle::{FailureOrigin, Timing, epoch_seconds},
|
||||
};
|
||||
use litellm_host_python::{HookStep, from_py, missing_state, to_py};
|
||||
use litellm_host_python::{Error as HostError, HookStep, from_py, to_py};
|
||||
use pyo3::{
|
||||
exceptions::{PyBaseException, PyException},
|
||||
gc::{PyTraverseError, PyVisit},
|
||||
|
|
@ -160,7 +160,7 @@ impl LegacyLogging {
|
|||
self.response
|
||||
.as_ref()
|
||||
.map(|response| HookStep::Ready(response.clone_ref(py)))
|
||||
.ok_or_else(missing_state)
|
||||
.ok_or_else(|| HostError::MissingState.into())
|
||||
}
|
||||
|
||||
fn dispatch_success(&self, py: Python<'_>) -> PyResult<()> {
|
||||
|
|
@ -209,7 +209,7 @@ impl LegacyLogging {
|
|||
|
||||
fn stream_success(&self, py: Python<'_>, stream: &DeliveredStream) -> PyResult<()> {
|
||||
let logger = self.logger()?;
|
||||
let billing = self.stream_billing().ok_or_else(missing_state)?;
|
||||
let billing = self.stream_billing().ok_or(HostError::MissingState)?;
|
||||
let billed = Streaming::Success.call(
|
||||
py,
|
||||
(
|
||||
|
|
@ -442,7 +442,7 @@ impl LegacyLogging {
|
|||
py: Python<'_>,
|
||||
facts: &litellm_host::interceptors::ExecutionFacts,
|
||||
) -> PyResult<HookStep<Self, ()>> {
|
||||
use litellm_host::interceptors::ResultSource;
|
||||
use litellm_host::interceptors::{Cost, ResultSource};
|
||||
|
||||
let logger = self.logger()?.object(py);
|
||||
let params = logger
|
||||
|
|
@ -462,6 +462,12 @@ impl LegacyLogging {
|
|||
),
|
||||
)?;
|
||||
let details = logger.getattr("model_call_details")?;
|
||||
let reported_cost = match &facts.cost {
|
||||
Cost::Reported { amount } => amount.as_f64(),
|
||||
Cost::CacheHit | Cost::Deferred => None,
|
||||
};
|
||||
logger.setattr("_provider_reported_cost", reported_cost)?;
|
||||
details.call_method1("pop", ("response_cost", py.None()))?;
|
||||
self.cache_key = match &facts.source {
|
||||
ResultSource::Provider => None,
|
||||
ResultSource::Cache { key } => Some(key.clone()),
|
||||
|
|
@ -519,7 +525,7 @@ impl LegacyLogging {
|
|||
return self.stream_failure(py);
|
||||
}
|
||||
if origin == FailureOrigin::Call && self.logger.is_some() && self.runs_deployment_hooks() {
|
||||
let error = self.error.as_ref().ok_or_else(missing_state)?;
|
||||
let error = self.error.as_ref().ok_or(HostError::MissingState)?;
|
||||
return Ok(HookStep::Await(
|
||||
DeploymentHooks::after_failure(py, self.call.kwargs(), error, self.call_type())?,
|
||||
Self::resume_deployment_failure,
|
||||
|
|
@ -530,7 +536,7 @@ impl LegacyLogging {
|
|||
|
||||
pub(crate) fn stream_opened(&mut self, py: Python<'_>, head: &Py<PyAny>) -> PyResult<()> {
|
||||
if self.stream_billing().is_none() {
|
||||
return Err(missing_state());
|
||||
return Err(HostError::MissingState.into());
|
||||
}
|
||||
if let Some(key) = &self.cache_key {
|
||||
head.bind(py).set_item("cache_key", key)?;
|
||||
|
|
@ -545,7 +551,7 @@ impl LegacyLogging {
|
|||
}
|
||||
|
||||
pub(crate) fn stream_chunk(&mut self, py: Python<'_>, chunk: &Py<PyAny>) -> PyResult<()> {
|
||||
let stream = self.stream.as_mut().ok_or_else(missing_state)?;
|
||||
let stream = self.stream.as_mut().ok_or(HostError::MissingState)?;
|
||||
if stream.first_chunk.is_none() {
|
||||
stream.first_chunk = Some(datetime(py, epoch_seconds())?);
|
||||
}
|
||||
|
|
@ -958,6 +964,81 @@ on_pre_call = lambda additional_args: None
|
|||
check = lambda: None
|
||||
";
|
||||
|
||||
#[rstest]
|
||||
#[case::reported(Some(0.37), false)]
|
||||
#[case::free(Some(0.0), false)]
|
||||
#[case::deferred(None, false)]
|
||||
#[case::cached(None, true)]
|
||||
fn result_ready_hands_only_a_reported_cost_to_the_spend_logging_record(
|
||||
#[case] reported: Option<f64>,
|
||||
#[case] cache_hit: bool,
|
||||
#[values(0.001, 7.0)] existing_cost: f64,
|
||||
) {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let locals = namespace(py, PAYLOAD_LOGGER);
|
||||
run(
|
||||
py,
|
||||
&locals,
|
||||
c"
|
||||
logger.litellm_params = {}
|
||||
logger.optional_params = {}
|
||||
logger.model_call_details = {'response_cost': 7.0}
|
||||
logger._provider_reported_cost = None
|
||||
",
|
||||
);
|
||||
let mut logging = LegacyLogging {
|
||||
logger: Some(PythonLogger::new(local(&locals, "logger").unbind())),
|
||||
..legacy_call(py, &locals, false)
|
||||
};
|
||||
let source = if cache_hit {
|
||||
litellm_host::interceptors::ResultSource::Cache {
|
||||
key: "cached-test".into(),
|
||||
}
|
||||
} else {
|
||||
litellm_host::interceptors::ResultSource::Provider
|
||||
};
|
||||
let facts = litellm_host::interceptors::ExecutionFacts {
|
||||
provider: litellm_host::interceptors::ProviderIdentity {
|
||||
model: "native-test".into(),
|
||||
provider: "edenai".into(),
|
||||
},
|
||||
cost: litellm_host::interceptors::Cost::settle(
|
||||
&source,
|
||||
reported.and_then(serde_json::Number::from_f64),
|
||||
),
|
||||
source,
|
||||
};
|
||||
let previous = litellm_host::interceptors::ExecutionFacts {
|
||||
provider: facts.provider.clone(),
|
||||
cost: litellm_host::interceptors::Cost::Reported {
|
||||
amount: serde_json::Number::from_f64(9.0).unwrap(),
|
||||
},
|
||||
source: litellm_host::interceptors::ResultSource::Provider,
|
||||
};
|
||||
logging.result_ready(py, &previous).unwrap();
|
||||
local(&locals, "logger")
|
||||
.getattr("model_call_details")
|
||||
.unwrap()
|
||||
.set_item("response_cost", existing_cost)
|
||||
.unwrap();
|
||||
let step = logging.result_ready(py, &facts).unwrap();
|
||||
assert!(matches!(step, HookStep::Ready(())));
|
||||
locals.set_item("reported", reported).unwrap();
|
||||
locals.set_item("cache_hit", cache_hit).unwrap();
|
||||
run(
|
||||
py,
|
||||
&locals,
|
||||
c"
|
||||
assert logger.model_call_details.get('response_cost') is None, logger.model_call_details
|
||||
assert logger._provider_reported_cost == reported, logger._provider_reported_cost
|
||||
assert logger.model_call_details['cache_hit'] is cache_hit, logger.model_call_details
|
||||
assert logger.update['custom_llm_provider'] == 'edenai', logger.update
|
||||
",
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
const DOCUMENT: &str = "data:application/pdf;base64,YWJj";
|
||||
const EDITED: &str = "data:application/pdf;base64,ZWRpdGVk";
|
||||
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
use pyo3::prelude::*;
|
||||
use strum::{IntoStaticStr, VariantArray};
|
||||
|
||||
const MODULE: &str = "litellm.rust_bridge.callbacks_legacy_python";
|
||||
const MODULE: &str = "litellm.rust_bridge.host.callbacks_legacy_python";
|
||||
|
||||
/// Every litellm Python internal the native call still borrows, grouped by the subsystem it
|
||||
/// belongs to. Rust drives the call; these exist only so behaviour that Python owns today
|
||||
|
|
@ -9,7 +9,7 @@ const MODULE: &str = "litellm.rust_bridge.callbacks_legacy_python";
|
|||
/// A group is deleted once Rust owns that subsystem, so this enum only shrinks. Calling a
|
||||
/// user's own callback is not borrowing and does not belong here.
|
||||
///
|
||||
/// `litellm/rust_bridge/callbacks_legacy_python.py` is the only Python module behind it, and
|
||||
/// `litellm/rust_bridge/host/callbacks_legacy_python.py` is the only Python module behind it, and
|
||||
/// `python_contract.json` pins each function's parameters on both sides.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub(crate) enum LegacyPython {
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ use pyo3::types::{PyDict, PyTuple};
|
|||
use crate::{LegacyLogging, PublicCall};
|
||||
|
||||
/// The parameters of every `callbacks_legacy_python` function, as the real module declares them.
|
||||
/// `tests/unit/rust_bridge/test_callbacks_legacy_python.py` pins this file to the Python
|
||||
/// `tests/unit/rust_bridge/host/test_callbacks_legacy_python.py` pins this file to the Python
|
||||
/// signatures, and [`namespace`] binds every fake call against it.
|
||||
pub(crate) const PYTHON_CONTRACT: &str = include_str!("../python_contract.json");
|
||||
|
||||
|
|
@ -23,10 +23,10 @@ import sys
|
|||
import traceback
|
||||
import types
|
||||
|
||||
for name in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.callbacks_legacy_python'):
|
||||
for name in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.host', 'litellm.rust_bridge.host.callbacks_legacy_python'):
|
||||
sys.modules.setdefault(name, types.ModuleType(name))
|
||||
|
||||
legacy = sys.modules['litellm.rust_bridge.callbacks_legacy_python']
|
||||
legacy = sys.modules['litellm.rust_bridge.host.callbacks_legacy_python']
|
||||
CONTRACT = json.loads(python_contract)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -16,18 +16,28 @@ pub enum LlmProviders {
|
|||
AzureAi,
|
||||
#[strum(serialize = "bedrock")]
|
||||
Bedrock,
|
||||
#[strum(serialize = "bedrock_mantle")]
|
||||
BedrockMantle,
|
||||
#[strum(serialize = "cohere")]
|
||||
Cohere,
|
||||
#[strum(serialize = "deepseek")]
|
||||
Deepseek,
|
||||
#[strum(serialize = "edenai")]
|
||||
Edenai,
|
||||
#[strum(serialize = "minimax")]
|
||||
Minimax,
|
||||
#[strum(serialize = "mistral")]
|
||||
Mistral,
|
||||
#[strum(serialize = "openai")]
|
||||
Openai,
|
||||
#[strum(serialize = "openai_like")]
|
||||
OpenaiLike,
|
||||
#[strum(serialize = "openrouter")]
|
||||
Openrouter,
|
||||
#[strum(serialize = "reducto")]
|
||||
Reducto,
|
||||
#[strum(serialize = "tencent")]
|
||||
Tencent,
|
||||
#[strum(serialize = "vertex_ai")]
|
||||
VertexAi,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@
|
|||
Core must pause mid-call to ask the host for things it cannot do itself (Python callbacks, secret and token reads, `before_send` rewrites, stream demand), then continue where it stopped. Any change to this crate must keep every requirement below. The alternatives section says which one each rejected design breaks
|
||||
|
||||
- R1 Core never calls the host: it names an op and waits for the answer, so it stays free of PyO3 and of any other host runtime
|
||||
- R2 Async host work is awaited by the host's own driver in the caller's asyncio task (`litellm/rust_bridge/lifecycle.py`), so `contextvars` writes reach the caller. A Rust-side `into_future` would run it in a copied context
|
||||
- R2 Async host work is awaited by the host's own driver in the caller's asyncio task (`litellm/rust_bridge/execution.py`), so `contextvars` writes reach the caller. A Rust-side `into_future` would run it in a copied context
|
||||
- R3 The body awaits real I/O (HTTP, `spawn_blocking`, timers) between yields, so `resume` is itself a future driven by the caller's runtime
|
||||
- R4 Route code stays straight-line async (`host.route(OcrOp::ReadDocument).await?`) instead of hand-written states
|
||||
- R5 Each op fixes its answer type at compile time: a host cannot answer `ReadDocument` with a token, and core never matches a result variant it did not ask for
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
use std::str;
|
||||
|
||||
use bytes::{Buf, BufMut, BytesMut};
|
||||
use bytes::{Buf, BufMut, Bytes, BytesMut};
|
||||
use tokio_util::codec::{Decoder, Encoder};
|
||||
|
||||
use crate::SseError;
|
||||
|
|
@ -29,38 +29,107 @@ impl Decoder for SseCodec {
|
|||
return Ok(None);
|
||||
}
|
||||
while let Some(end) = block_end(src) {
|
||||
let block = src.split_to(end);
|
||||
let pending = lines(&block)
|
||||
.map(|(line, _)| line)
|
||||
.take_while(|line| !line.is_empty())
|
||||
.try_fold(Pending::default(), Pending::apply)?;
|
||||
if let Some(event) = pending.dispatch() {
|
||||
if let Some(event) = event_in(&src.split_to(end))? {
|
||||
return Ok(Some(event));
|
||||
}
|
||||
}
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
fn decode_eof(&mut self, _pending: &mut BytesMut) -> Result<Option<SseEvent>, SseError> {
|
||||
Ok(None)
|
||||
/// An event the stream ended inside of is dropped, as the SSE spec requires.
|
||||
fn decode_eof(&mut self, pending: &mut BytesMut) -> Result<Option<SseEvent>, SseError> {
|
||||
self.decode(pending)
|
||||
}
|
||||
}
|
||||
|
||||
pub fn decode_all(wire: &[u8]) -> Result<Vec<SseEvent>, SseError> {
|
||||
let mut codec = SseCodec::default();
|
||||
let mut buffer = BytesMut::from(wire);
|
||||
let mut events = Vec::new();
|
||||
while let Some(event) = codec.decode_eof(&mut buffer)? {
|
||||
events.push(event);
|
||||
}
|
||||
Ok(events)
|
||||
}
|
||||
|
||||
impl SseCodec {
|
||||
fn skip_bom(&mut self, src: &mut BytesMut) -> bool {
|
||||
if self.past_bom {
|
||||
return true;
|
||||
}
|
||||
if src.starts_with(BOM) {
|
||||
src.advance(BOM.len());
|
||||
} else if BOM.starts_with(src) {
|
||||
return false;
|
||||
}
|
||||
self.past_bom = true;
|
||||
true
|
||||
skip_bom(&mut self.past_bom, src)
|
||||
}
|
||||
}
|
||||
|
||||
fn event_in(block: &[u8]) -> Result<Option<SseEvent>, SseError> {
|
||||
Ok(lines(block)
|
||||
.map(|(line, _)| line)
|
||||
.take_while(|line| !line.is_empty())
|
||||
.try_fold(Pending::default(), Pending::apply)?
|
||||
.dispatch())
|
||||
}
|
||||
|
||||
/// One event together with the exact wire bytes that carried it, for relays that forward a
|
||||
/// provider's bytes untouched but still need to see each event go by.
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct SseBlock {
|
||||
pub bytes: Bytes,
|
||||
pub event: SseEvent,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Default)]
|
||||
pub struct SseBlockCodec {
|
||||
past_bom: bool,
|
||||
}
|
||||
|
||||
impl Decoder for SseBlockCodec {
|
||||
type Item = SseBlock;
|
||||
type Error = SseError;
|
||||
|
||||
fn decode(&mut self, src: &mut BytesMut) -> Result<Option<SseBlock>, SseError> {
|
||||
self.decode_block(src, false)
|
||||
}
|
||||
|
||||
fn decode_eof(&mut self, pending: &mut BytesMut) -> Result<Option<SseBlock>, SseError> {
|
||||
self.decode_block(pending, true)
|
||||
}
|
||||
}
|
||||
|
||||
impl SseBlockCodec {
|
||||
fn decode_block(
|
||||
&mut self,
|
||||
src: &mut BytesMut,
|
||||
eof: bool,
|
||||
) -> Result<Option<SseBlock>, SseError> {
|
||||
if !skip_bom(&mut self.past_bom, src) {
|
||||
return Ok(None);
|
||||
}
|
||||
while let Some(end) = block_end(src) {
|
||||
if !eof && end == src.len() && src.ends_with(b"\r") {
|
||||
return Ok(None);
|
||||
}
|
||||
let block = src.split_to(end);
|
||||
if let Some(event) = event_in(&block)? {
|
||||
return Ok(Some(SseBlock {
|
||||
bytes: block.freeze(),
|
||||
event,
|
||||
}));
|
||||
}
|
||||
}
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
|
||||
fn skip_bom(past_bom: &mut bool, src: &mut BytesMut) -> bool {
|
||||
if *past_bom {
|
||||
return true;
|
||||
}
|
||||
if src.starts_with(BOM) {
|
||||
src.advance(BOM.len());
|
||||
} else if BOM.starts_with(src) {
|
||||
return false;
|
||||
}
|
||||
*past_bom = true;
|
||||
true
|
||||
}
|
||||
|
||||
fn block_end(bytes: &[u8]) -> Option<usize> {
|
||||
lines(bytes)
|
||||
.find(|(line, _)| line.is_empty())
|
||||
|
|
@ -168,3 +237,179 @@ impl Encoder<SseEvent> for SseCodec {
|
|||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use futures_util::{StreamExt, TryStreamExt, stream};
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
use crate::frames;
|
||||
|
||||
fn pieces(wire: &[u8], size: usize) -> Vec<Bytes> {
|
||||
wire.chunks(size).map(Bytes::copy_from_slice).collect()
|
||||
}
|
||||
|
||||
async fn blocks(wire: &[u8], size: usize) -> Result<Vec<SseBlock>, SseError> {
|
||||
frames(
|
||||
stream::iter(pieces(wire, size).into_iter().map(Ok::<_, std::io::Error>)),
|
||||
SseBlockCodec::default(),
|
||||
)
|
||||
.try_collect()
|
||||
.await
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::lf(
|
||||
b": keepalive\n\nevent: delta\ndata: a\ndata: b\n\n",
|
||||
b"event: delta\ndata: a\ndata: b\n\n"
|
||||
)]
|
||||
#[case::crlf(
|
||||
b": keepalive\r\n\r\nevent: delta\r\ndata: a\r\ndata: b\r\n\r\n",
|
||||
b"event: delta\r\ndata: a\r\ndata: b\r\n\r\n"
|
||||
)]
|
||||
#[case::cr(
|
||||
b": keepalive\r\revent: delta\rdata: a\rdata: b\r\r",
|
||||
b"event: delta\rdata: a\rdata: b\r\r"
|
||||
)]
|
||||
#[case::bom(b"\xEF\xBB\xBFdata: \xe2\x82\xac\n\n", b"data: \xe2\x82\xac\n\n")]
|
||||
#[case::empty_data(b"event: ignored\n\ndata:\n\n", b"data:\n\n")]
|
||||
#[case::extensions(
|
||||
b"id: 1\nretry: 10\nvendor: x\ndata: a\n\n",
|
||||
b"id: 1\nretry: 10\nvendor: x\ndata: a\n\n"
|
||||
)]
|
||||
#[tokio::test]
|
||||
async fn blocks_carry_the_wire_bytes_and_the_event_they_dispatch(
|
||||
#[case] wire: &[u8],
|
||||
#[case] expected: &[u8],
|
||||
#[values(1, 3, 4096)] size: usize,
|
||||
) {
|
||||
let blocks = blocks(wire, size).await.unwrap();
|
||||
let expected_events = decode_all(expected).unwrap();
|
||||
assert_eq!(expected_events.len(), 1);
|
||||
assert_eq!(
|
||||
blocks,
|
||||
vec![SseBlock {
|
||||
bytes: Bytes::copy_from_slice(expected),
|
||||
event: expected_events[0].clone(),
|
||||
}]
|
||||
);
|
||||
assert_eq!(expected_events, decode_all(wire).unwrap());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::empty(b"", vec![])]
|
||||
#[case::unterminated(b"data: partial\n", vec![])]
|
||||
#[case::partial_tail(b"data: complete\n\ndata: partial\n", vec![&b"data: complete\n\n"[..]])]
|
||||
#[case::multiple(b"data: first\n\ndata: second\n\n", vec![&b"data: first\n\n"[..], &b"data: second\n\n"[..]])]
|
||||
#[case::trailing_comment(b"data: first\n\n: bye\n\n", vec![&b"data: first\n\n"[..]])]
|
||||
#[tokio::test]
|
||||
async fn blocks_dispatch_complete_events_and_drop_a_cut_tail(
|
||||
#[case] wire: &[u8],
|
||||
#[case] expected: Vec<&[u8]>,
|
||||
) {
|
||||
let mut blocks = Box::pin(frames(
|
||||
stream::iter(pieces(wire, 1).into_iter().map(Ok::<_, std::io::Error>)),
|
||||
SseBlockCodec::default(),
|
||||
));
|
||||
for expected in expected {
|
||||
assert_eq!(blocks.next().await.unwrap().unwrap().bytes, expected);
|
||||
}
|
||||
assert!(blocks.next().await.is_none());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn bulk_decode_preserves_every_field() {
|
||||
let events = decode_all(b"id: 1\nretry: 20\nevent: delta\ndata: a\ndata: b\n\n").unwrap();
|
||||
assert_eq!(
|
||||
events,
|
||||
vec![SseEvent {
|
||||
event: Some("delta".into()),
|
||||
data: "a\nb".into(),
|
||||
id: Some("1".into()),
|
||||
retry: Some(20),
|
||||
}]
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::only_a_tail(b"data: partial\n", vec![])]
|
||||
#[case::tail_after_an_event(b"data: complete\n\ndata: partial\n", vec!["complete"])]
|
||||
#[case::no_terminator_at_all(b"data: partial", vec![])]
|
||||
fn bulk_decode_drops_the_event_a_stream_ends_inside(
|
||||
#[case] wire: &[u8],
|
||||
#[case] expected: Vec<&str>,
|
||||
) {
|
||||
let data: Vec<String> = decode_all(wire)
|
||||
.unwrap()
|
||||
.into_iter()
|
||||
.map(|e| e.data)
|
||||
.collect();
|
||||
assert_eq!(data, expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::trailing_blank_lines(b"data: x\n\n\n\n")]
|
||||
#[case::trailing_comment(b"data: x\n\n: bye\n\n")]
|
||||
#[case::lone_cr_at_eof(b"data: x\r\r")]
|
||||
fn bulk_decode_accepts_a_stream_that_ends_between_events(#[case] wire: &[u8]) {
|
||||
assert_eq!(decode_all(wire).unwrap().len(), 1);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn bulk_decode_rejects_invalid_utf8() {
|
||||
assert!(matches!(
|
||||
decode_all(b"data: \xff\n\n"),
|
||||
Err(SseError::InvalidUtf8(_))
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn block_framing_preserves_transport_errors() {
|
||||
let input = stream::iter([
|
||||
Ok(Bytes::from_static(b"data: first\n\n")),
|
||||
Err(std::io::Error::new(
|
||||
std::io::ErrorKind::ConnectionReset,
|
||||
"reset",
|
||||
)),
|
||||
]);
|
||||
let mut blocks = Box::pin(frames(input, SseBlockCodec::default()));
|
||||
assert_eq!(
|
||||
blocks.next().await.unwrap().unwrap().bytes,
|
||||
Bytes::from_static(b"data: first\n\n")
|
||||
);
|
||||
let Some(Err(SseError::Body(error))) = blocks.next().await else {
|
||||
panic!("expected body error");
|
||||
};
|
||||
let cause = error
|
||||
.get_ref()
|
||||
.unwrap()
|
||||
.downcast_ref::<std::io::Error>()
|
||||
.unwrap();
|
||||
assert_eq!(cause.kind(), std::io::ErrorKind::ConnectionReset);
|
||||
assert!(blocks.next().await.is_none());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn invalid_utf8_stops_block_framing_after_prior_events() {
|
||||
let mut blocks = Box::pin(frames(
|
||||
stream::iter(
|
||||
pieces(b"data: first\n\ndata: \xff\n\n", 1)
|
||||
.into_iter()
|
||||
.map(Ok::<_, std::io::Error>),
|
||||
),
|
||||
SseBlockCodec::default(),
|
||||
));
|
||||
assert_eq!(
|
||||
blocks.next().await.unwrap().unwrap().bytes,
|
||||
Bytes::from_static(b"data: first\n\n")
|
||||
);
|
||||
assert!(matches!(
|
||||
blocks.next().await,
|
||||
Some(Err(SseError::InvalidUtf8(_)))
|
||||
));
|
||||
assert!(blocks.next().await.is_none());
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -97,12 +97,11 @@ async fn dispatch_follows_the_data_buffer(#[case] wire: &[u8], #[case] expected:
|
|||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::unterminated_single(b"data: partial\n", vec![])]
|
||||
#[case::unterminated_tail_after_complete(b"data: complete\n\ndata: unfinished\n", vec![event(None, "complete")])]
|
||||
#[case::lone_cr_terminates_at_eof(b"data: x\r\r", vec![event(None, "x")])]
|
||||
#[case::lone_cr_line_then_eof(b"data: x\r", vec![])]
|
||||
#[case::trailing_blank_lines(b"data: x\n\n\n", vec![event(None, "x")])]
|
||||
#[case::trailing_comment(b"data: x\n\n: bye\n\n", vec![event(None, "x")])]
|
||||
#[tokio::test]
|
||||
async fn eof_dispatches_only_terminated_events(
|
||||
async fn eof_between_events_ends_the_stream_cleanly(
|
||||
#[case] wire: &[u8],
|
||||
#[case] expected: Vec<SseEvent>,
|
||||
) {
|
||||
|
|
@ -112,6 +111,19 @@ async fn eof_dispatches_only_terminated_events(
|
|||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::unterminated_single(b"data: partial\n", vec![])]
|
||||
#[case::unterminated_tail_after_complete(b"data: complete\n\ndata: unfinished\n", vec![event(None, "complete")])]
|
||||
#[case::lone_cr_line_then_eof(b"data: x\r", vec![])]
|
||||
#[case::partial_bom(b"\xEF\xBB", vec![])]
|
||||
#[tokio::test]
|
||||
async fn eof_inside_an_event_drops_it_after_the_terminated_events(
|
||||
#[case] wire: &[u8],
|
||||
#[case] expected: Vec<SseEvent>,
|
||||
) {
|
||||
assert_eq!(collect(every(wire, 1)).await.unwrap(), expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::inside_the_first_line(vec![&b"data: a\r"[..], &b"\ndata: b\r\n\r\n"[..]])]
|
||||
#[case::inside_the_blank_line(vec![&b"data: a\r\ndata: b\r\n\r"[..], &b"\n"[..]])]
|
||||
|
|
|
|||
|
|
@ -19,7 +19,12 @@ use crate::{Deployment, Error, Gateway, JsonObject, RequestId, request};
|
|||
|
||||
/// Client headers Python forwards to Anthropic-speaking providers on every call.
|
||||
const ANTHROPIC_API_HEADERS: [&str; 2] = ["anthropic-version", "anthropic-beta"];
|
||||
const ANTHROPIC_API_HEADER_PROVIDERS: &str = "anthropic,bedrock,bedrock_mantle,vertex_ai";
|
||||
const OPENROUTER_API_HEADERS: [&str; 2] = [
|
||||
"x-openrouter-metadata",
|
||||
"x-openrouter-experimental-metadata",
|
||||
];
|
||||
const ANTHROPIC_API_HEADER_PROVIDERS: &str =
|
||||
"anthropic,bedrock,bedrock_mantle,vertex_ai,openrouter";
|
||||
|
||||
pub async fn create(
|
||||
State(gateway): State<Arc<Gateway>>,
|
||||
|
|
@ -88,17 +93,28 @@ fn project(
|
|||
}
|
||||
|
||||
fn anthropic_api_headers(headers: &HeaderMap) -> Option<ProviderSpecificHeaders> {
|
||||
let extra_headers: Map<String, Value> = ANTHROPIC_API_HEADERS
|
||||
let scopes = [
|
||||
(
|
||||
ANTHROPIC_API_HEADER_PROVIDERS,
|
||||
ANTHROPIC_API_HEADERS.as_slice(),
|
||||
),
|
||||
("openrouter", OPENROUTER_API_HEADERS.as_slice()),
|
||||
];
|
||||
let forwarded: Vec<ProviderSpecificHeader> = scopes
|
||||
.into_iter()
|
||||
.filter_map(|name| {
|
||||
let value = headers.get(name)?.to_str().ok()?;
|
||||
Some((name.to_owned(), Value::from(value)))
|
||||
.filter_map(|(provider, names)| {
|
||||
let extra_headers: Map<String, Value> = names
|
||||
.iter()
|
||||
.filter_map(|name| {
|
||||
let value = headers.get(*name)?.to_str().ok()?;
|
||||
Some(((*name).to_owned(), Value::from(value)))
|
||||
})
|
||||
.collect();
|
||||
(!extra_headers.is_empty()).then(|| ProviderSpecificHeader {
|
||||
custom_llm_provider: provider.into(),
|
||||
extra_headers,
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
(!extra_headers.is_empty()).then(|| {
|
||||
ProviderSpecificHeaders::One(ProviderSpecificHeader {
|
||||
custom_llm_provider: ANTHROPIC_API_HEADER_PROVIDERS.into(),
|
||||
extra_headers,
|
||||
})
|
||||
})
|
||||
(!forwarded.is_empty()).then_some(ProviderSpecificHeaders::Many(forwarded))
|
||||
}
|
||||
|
|
|
|||
|
|
@ -192,3 +192,52 @@ async fn model_permissions_prevent_provider_calls(#[case] path: &str) {
|
|||
assert_eq!(response.status(), 403);
|
||||
assert!(upstream.received_requests().await.unwrap().is_empty());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::completion(false)]
|
||||
#[case::streaming(true)]
|
||||
#[tokio::test]
|
||||
async fn openrouter_receives_native_messages_headers_and_its_deployment_bearer(
|
||||
#[case] streaming: bool,
|
||||
) {
|
||||
let upstream = MockServer::start().await;
|
||||
let message = json!({"id": "msg_test", "type": "message", "role": "assistant", "model": "anthropic/test-model", "content": [{"type": "text", "text": "hello"}], "stop_reason": "end_turn", "stop_sequence": null, "usage": {"input_tokens": 1, "output_tokens": 1, "cost": 0.001, "cost_details": "unexpected"}, "openrouter_metadata": {"provider": "test-provider"}});
|
||||
let sse = "event: message_delta\ndata: {\"type\":\"message_delta\",\"usage\":{\"cost\":0.001}}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n";
|
||||
let template = if streaming {
|
||||
ResponseTemplate::new(200).set_body_raw(sse, "text/event-stream")
|
||||
} else {
|
||||
ResponseTemplate::new(200).set_body_json(message.clone())
|
||||
};
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/api/v1/messages"))
|
||||
.and(header("authorization", "Bearer test-key"))
|
||||
.and(header("anthropic-beta", "test-feature"))
|
||||
.and(header("anthropic-version", "test-version"))
|
||||
.and(header("x-openrouter-metadata", "enabled"))
|
||||
.and(body_json(json!({"model": "anthropic/test-model", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16, "stream": streaming, "provider": {"order": ["test-provider"]}})))
|
||||
.respond_with(template).expect(1).mount(&upstream).await;
|
||||
let request = Request::post("/v1/messages")
|
||||
.header("content-type", "application/json")
|
||||
.header("anthropic-beta", "test-feature")
|
||||
.header("anthropic-version", "test-version")
|
||||
.header("x-openrouter-metadata", "enabled")
|
||||
.header("x-api-key", "proxy-key")
|
||||
.header("authorization", "Bearer proxy-key")
|
||||
.body(Body::from(json!({"model": "public/model", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16, "stream": streaming, "provider": {"order": ["test-provider"]}}).to_string())).unwrap();
|
||||
let response = support::app(
|
||||
"openrouter/anthropic/test-model",
|
||||
&format!("{}/api/v1", upstream.uri()),
|
||||
)
|
||||
.oneshot(request)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(response.status(), 200);
|
||||
if streaming {
|
||||
assert_eq!(to_bytes(response.into_body(), 4096).await.unwrap(), sse);
|
||||
} else {
|
||||
assert_eq!(support::json(response).await, message);
|
||||
}
|
||||
let requests = upstream.received_requests().await.unwrap();
|
||||
assert_eq!(requests.len(), 1);
|
||||
assert!(!requests[0].headers.contains_key("x-api-key"));
|
||||
}
|
||||
|
|
|
|||
|
|
@ -24,7 +24,7 @@ From the repository root, this loopback-only example runs the existing Python MC
|
|||
|
||||
```sh
|
||||
cargo run --manifest-path litellm-rust/Cargo.toml -p litellm-gateway-mcp --example stdio_proxy -- \
|
||||
.venv/bin/python -c 'import runpy; runpy.run_path("tests/mcp_tests/mcp_e2e_upstream_server.py")["mcp"].run(transport="stdio")'
|
||||
.venv/bin/python -c 'import runpy; runpy.run_path("tests/integration/_support/mcp_e2e_upstream_server.py")["mcp"].run(transport="stdio")'
|
||||
```
|
||||
|
||||
```sh
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ The binary reads YAML from `LITELLM_CONFIG` and mounts configured MCP servers al
|
|||
From the repository root, start the Python fixture in one terminal:
|
||||
|
||||
```sh
|
||||
MCP_HOST=127.0.0.1 MCP_PORT=8090 .venv/bin/python tests/mcp_tests/mcp_e2e_upstream_server.py
|
||||
MCP_HOST=127.0.0.1 MCP_PORT=8090 .venv/bin/python tests/integration/_support/mcp_e2e_upstream_server.py
|
||||
```
|
||||
|
||||
Start the gateway in another terminal with the example configuration:
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ This crate owns CPython execution mechanics for generic `litellm-host` machines
|
|||
|
||||
`src/native.rs` belongs here: it runs a generic machine through the Python runtime and owns its pending execution and abort handle. Keep provider selection, request projection and public exception policy out of it. A rename to `machine_runner.rs` is optional and must not change behavior
|
||||
|
||||
The execution handle receives its Python lifecycle binding from its consumer through `PythonLifecycle` rather than import a fixed `litellm.rust_bridge` module. Generic suspension and execution state validation belong here. Public stream wrappers and `_hidden_params` conventions belong to the consumer
|
||||
`run_call` returns the unstarted `Execution` and never names a Python stream class or driver. The consumer drives it from Python and supplies the step classes through `StepTypes` rather than have this crate import a fixed `litellm.rust_bridge` module. Generic suspension and execution state validation belong here. Public stream wrappers and `_hidden_params` conventions belong to the consumer
|
||||
|
||||
Creating a resolved asyncio Future from an already constructed Python value belongs here, alongside runtime waiting, interpreter detachment and panic containment. Choosing which callable exceptions become a public `RuntimeError` belongs to the consumer. `python-bridge::callable::wrap_failure` owns that policy
|
||||
|
||||
|
|
@ -28,6 +28,7 @@ Terminal notifications share the selected response or exception. An ordinary not
|
|||
- `PythonBinding::decode_request` receives the keyword view returned by `prepare_arguments` and updated by `arguments_prepared`, not the caller's dict. A binding that decodes from it inherits all composed argument rewrites
|
||||
- A native failure, including one a host op returns as `InvokeError::Native`, is classified exactly once through the binding's `map_error`. A Python exception raised inside the call, and a failure in `prepare_arguments` or `transform_response`, is raised as is
|
||||
- A failing `map_error` is raised with the native error's text as its `__context__`, never swallowed
|
||||
- Consumers import this crate's error as `use litellm_host_python::Error as HostError`, so it never shadows their own `Error`
|
||||
- Use standard PyO3 ownership and conversion APIs
|
||||
- Prefer `Bound<'py, T>` for attached operations/results, `Py<T>` for retention. Binding/unbinding does not copy payloads
|
||||
- Use `pythonize` for selected Serde data, never a JSON-text round trip. Share conversion with `Pythonized<T>`
|
||||
|
|
@ -37,7 +38,7 @@ Terminal notifications share the selected response or exception. An ordinary not
|
|||
- Keep diagnostic counters in the consumer. Wrapper invocations do not measure every interpreter release
|
||||
- Release exclusive class borrows/locks before Python calls or decrements that can invoke finalizers. Expose retained Python edges to GC without calling Python during traversal
|
||||
- Keep coroutine driving in the shared Python driver and the native handle
|
||||
- Shared driver implementation: `litellm/rust_bridge/lifecycle.py`, handle: `src/handle.rs`, call driver: `src/driver.rs`, native-backed behavior tests: `tests/lifecycle.py`. The consumer supplies the lifecycle binding
|
||||
- Shared driver implementation: `litellm/rust_bridge/execution.py`, handle: `src/handle.rs`, call driver: `src/driver.rs`, native-backed behavior tests: `tests/lifecycle.py`. The consumer supplies the lifecycle binding
|
||||
- Every lifecycle suspension is awaited inline in the caller's task. `into_future` creates a separate task and cannot satisfy this contract
|
||||
- References: [ownership](https://pyo3.rs/v0.29.2/types.html), [conversions](https://pyo3.rs/v0.29.2/conversions/traits.html), [pythonize errors](https://docs.rs/pythonize/0.29.0/src/pythonize/error.rs.html)
|
||||
- [GC](https://pyo3.rs/v0.29.2/class/protocols.html#garbage-collector-integration), [re-entry](https://pyo3.rs/v0.29.2/class/call.html), [parallelism](https://pyo3.rs/v0.29.2/parallelism.html), [async delivery source](https://docs.rs/pyo3-async-runtimes/0.29.0/src/pyo3_async_runtimes/generic.rs.html)
|
||||
|
|
|
|||
|
|
@ -6,11 +6,14 @@ license.workspace = true
|
|||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
litellm-coroutine.workspace = true
|
||||
litellm-host.workspace = true
|
||||
|
||||
bytes.workspace = true
|
||||
futures-util.workspace = true
|
||||
parking_lot.workspace = true
|
||||
serde.workspace = true
|
||||
thiserror.workspace = true
|
||||
tokio = { workspace = true, features = ["rt", "sync"] }
|
||||
|
||||
pyo3.workspace = true
|
||||
|
|
|
|||
58
litellm-rust/crates/host-python/src/caller.rs
Normal file
58
litellm-rust/crates/host-python/src/caller.rs
Normal file
|
|
@ -0,0 +1,58 @@
|
|||
use std::ops::ControlFlow;
|
||||
|
||||
use litellm_coroutine::{Abandoned, Co, Reply};
|
||||
use pyo3::prelude::*;
|
||||
|
||||
use crate::Error;
|
||||
use crate::handle::Suspension;
|
||||
|
||||
pub(crate) type Answer = PyResult<Py<PyAny>>;
|
||||
|
||||
pub(crate) struct Yielded {
|
||||
pub(crate) suspension: Suspension,
|
||||
pub(crate) reply: Reply<Answer>,
|
||||
}
|
||||
|
||||
/// Hands suspensions to the Python caller and receives its answers.
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct PythonCaller {
|
||||
co: Co<Yielded>,
|
||||
}
|
||||
|
||||
impl PythonCaller {
|
||||
/// Connects the call body to its coroutine's yield handle.
|
||||
pub(crate) fn new(co: Co<Yielded>) -> Self {
|
||||
Self { co }
|
||||
}
|
||||
|
||||
/// Asks the caller to await Python work and return its value or exception.
|
||||
pub(crate) async fn await_python(&self, awaitable: Py<PyAny>) -> Answer {
|
||||
self.suspend(Suspension::Await(awaitable)).await
|
||||
}
|
||||
|
||||
/// Hands the caller a stream head and waits for demand or closure.
|
||||
pub(crate) async fn open_stream(&self, head: Py<PyAny>) -> ControlFlow<()> {
|
||||
self.stream_boundary(Suspension::StreamOpen(head)).await
|
||||
}
|
||||
|
||||
/// Delivers a chunk and waits for the caller to request another or close.
|
||||
pub(crate) async fn yield_chunk(&self, chunk: Py<PyAny>) -> ControlFlow<()> {
|
||||
self.stream_boundary(Suspension::StreamChunk(chunk)).await
|
||||
}
|
||||
|
||||
/// Maps stream demand to continuation and any error to detachment.
|
||||
async fn stream_boundary(&self, suspension: Suspension) -> ControlFlow<()> {
|
||||
match self.suspend(suspension).await {
|
||||
Ok(_) => ControlFlow::Continue(()),
|
||||
Err(_) => ControlFlow::Break(()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Yields a suspension with its reply handle and waits for the answer.
|
||||
async fn suspend(&self, suspension: Suspension) -> Answer {
|
||||
match self.co.yield_(|reply| Yielded { suspension, reply }).await {
|
||||
Ok(answer) => answer,
|
||||
Err(Abandoned) => Err(Error::MissingState.into()),
|
||||
}
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load diff
|
|
@ -1,20 +1,58 @@
|
|||
use pyo3::prelude::*;
|
||||
use pyo3::{
|
||||
PyErr,
|
||||
exceptions::{PyRuntimeError, PyTypeError},
|
||||
};
|
||||
|
||||
/// A failure host-python raises itself, surfacing in Python as the exception `PyErr::from` picks.
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum Error {
|
||||
#[error("missing native call state")]
|
||||
MissingState,
|
||||
#[error("execution is already running")]
|
||||
ExecutionRunning,
|
||||
#[error("execution is closed")]
|
||||
ExecutionClosed,
|
||||
#[error("can't send non-None value to a just-started execution")]
|
||||
SendBeforeStart,
|
||||
#[error("execution requires start before resume and can only start once")]
|
||||
ExecutionOutOfOrder,
|
||||
#[error("execution was closed while running")]
|
||||
ExecutionClosedWhileRunning,
|
||||
#[error("native execution closed")]
|
||||
NativeExecutionClosed,
|
||||
#[error("synchronous native routes cannot run from a Tokio context; use the async route")]
|
||||
SyncInsideTokio,
|
||||
#[error("the context body produced no value")]
|
||||
ContextBodyProducedNothing,
|
||||
#[error("file read must return bytes or str, got {got}")]
|
||||
FileContentType { got: String },
|
||||
}
|
||||
|
||||
impl From<Error> for PyErr {
|
||||
fn from(error: Error) -> Self {
|
||||
match error {
|
||||
Error::SendBeforeStart | Error::FileContentType { .. } => {
|
||||
PyTypeError::new_err(error.to_string())
|
||||
}
|
||||
Error::MissingState
|
||||
| Error::ExecutionRunning
|
||||
| Error::ExecutionClosed
|
||||
| Error::ExecutionOutOfOrder
|
||||
| Error::ExecutionClosedWhileRunning
|
||||
| Error::NativeExecutionClosed
|
||||
| Error::SyncInsideTokio
|
||||
| Error::ContextBodyProducedNothing => PyRuntimeError::new_err(error.to_string()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Why a custom operation the host answered did not produce a result: the route's own code
|
||||
/// rejected it, which the route classifies like any other native failure, or Python code
|
||||
/// raised, which reaches the caller as it was raised.
|
||||
#[derive(Debug)]
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum InvokeError<E> {
|
||||
#[error("{0}")]
|
||||
Native(E),
|
||||
Python(PyErr),
|
||||
}
|
||||
|
||||
impl<E> From<PyErr> for InvokeError<E> {
|
||||
fn from(error: PyErr) -> Self {
|
||||
Self::Python(error)
|
||||
}
|
||||
}
|
||||
|
||||
pub fn missing_state() -> PyErr {
|
||||
pyo3::exceptions::PyRuntimeError::new_err("missing native call state")
|
||||
#[error(transparent)]
|
||||
Python(#[from] PyErr),
|
||||
}
|
||||
|
|
|
|||
|
|
@ -3,13 +3,14 @@
|
|||
|
||||
use bytes::Bytes;
|
||||
use pyo3::{
|
||||
exceptions::PyTypeError,
|
||||
gc::{PyTraverseError, PyVisit},
|
||||
prelude::*,
|
||||
pybacked::PyBackedBytes,
|
||||
types::{PyBytes, PyString},
|
||||
};
|
||||
|
||||
use crate::Error;
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct FileContent {
|
||||
pub bytes: Bytes,
|
||||
|
|
@ -50,10 +51,10 @@ impl PythonFileReader {
|
|||
} else if value.is_instance_of::<PyBytes>() {
|
||||
py_bytes(&value)?
|
||||
} else {
|
||||
return Err(PyTypeError::new_err(format!(
|
||||
"file read must return bytes or str, got {}",
|
||||
value.get_type(),
|
||||
)));
|
||||
return Err(Error::FileContentType {
|
||||
got: value.get_type().to_string(),
|
||||
}
|
||||
.into());
|
||||
};
|
||||
Ok(FileContent {
|
||||
bytes,
|
||||
|
|
|
|||
|
|
@ -3,7 +3,9 @@ use std::sync::{
|
|||
atomic::{AtomicU64, Ordering},
|
||||
};
|
||||
|
||||
use pyo3::{exceptions::PyRuntimeError, prelude::*};
|
||||
use pyo3::prelude::*;
|
||||
|
||||
use crate::Error;
|
||||
|
||||
/// The caller's `contextvars` context, captured at the Python entry point so blocking Python
|
||||
/// work started from Rust sees the same request-local values as the Python caller.
|
||||
|
|
@ -49,7 +51,7 @@ impl PythonContext {
|
|||
.lock()
|
||||
.expect("context value slot poisoned")
|
||||
.take()
|
||||
.ok_or_else(|| PyRuntimeError::new_err("the context body produced no value"))
|
||||
.ok_or_else(|| Error::ContextBodyProducedNothing.into())
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -1,19 +1,30 @@
|
|||
use std::panic::{AssertUnwindSafe, catch_unwind};
|
||||
|
||||
use crate::panic_to_pyerr;
|
||||
use pyo3::exceptions::{PyBaseException, PyRuntimeError};
|
||||
use crate::{Error, panic_to_pyerr};
|
||||
use pyo3::exceptions::PyBaseException;
|
||||
use pyo3::gc::{PyTraverseError, PyVisit};
|
||||
use pyo3::prelude::*;
|
||||
|
||||
pub type PythonLifecycle = for<'py> fn(Python<'py>) -> PyResult<Bound<'py, PyModule>>;
|
||||
pub type StepTypes = for<'py> fn(Python<'py>) -> PyResult<Bound<'py, PyModule>>;
|
||||
|
||||
pub enum ExecutionStep {
|
||||
Return(Py<PyAny>),
|
||||
/// What a suspended execution waits on: an awaitable the caller awaits, or a stream
|
||||
/// boundary the caller's stream reads past.
|
||||
pub enum Suspension {
|
||||
Await(Py<PyAny>),
|
||||
/// The call streams: the caller gets a stream over this execution carrying this head,
|
||||
/// and the execution stays suspended until the stream asks for a chunk.
|
||||
Open(Py<PyAny>),
|
||||
Yield(Py<PyAny>),
|
||||
StreamOpen(Py<PyAny>),
|
||||
StreamChunk(Py<PyAny>),
|
||||
}
|
||||
|
||||
enum Resume {
|
||||
Send(Py<PyAny>),
|
||||
Throw(PyErr),
|
||||
}
|
||||
|
||||
pub enum ExecutionStep {
|
||||
Suspended(Suspension),
|
||||
Return(Py<PyAny>),
|
||||
}
|
||||
|
||||
pub trait ExecutionBody: Send + Sync {
|
||||
|
|
@ -31,77 +42,59 @@ enum ExecutionState {
|
|||
#[pyclass]
|
||||
pub struct Execution {
|
||||
state: ExecutionState,
|
||||
lifecycle: PythonLifecycle,
|
||||
steps: StepTypes,
|
||||
}
|
||||
|
||||
impl Execution {
|
||||
pub fn new(body: impl ExecutionBody + 'static, lifecycle: PythonLifecycle) -> Self {
|
||||
pub fn new(body: impl ExecutionBody + 'static, steps: StepTypes) -> Self {
|
||||
Self {
|
||||
state: ExecutionState::Created(Box::new(body)),
|
||||
lifecycle,
|
||||
steps,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn into_coroutine(self, py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let binding = (self.lifecycle)(py)?;
|
||||
let execution = Py::new(py, self)?;
|
||||
binding.getattr("drive")?.call1((execution,))
|
||||
}
|
||||
|
||||
pub(crate) fn into_sync_stream(
|
||||
self,
|
||||
py: Python<'_>,
|
||||
head: Py<PyAny>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
(self.lifecycle)(py)?
|
||||
.getattr("SyncStream")?
|
||||
.call1((Py::new(py, self)?, head))
|
||||
}
|
||||
|
||||
/// An execution already started elsewhere and now waiting for its next input.
|
||||
pub fn suspended(body: impl ExecutionBody + 'static, lifecycle: PythonLifecycle) -> Self {
|
||||
Self {
|
||||
state: ExecutionState::Suspended(Box::new(body)),
|
||||
lifecycle,
|
||||
}
|
||||
}
|
||||
|
||||
fn advance(
|
||||
slf: &Bound<'_, Self>,
|
||||
py: Python<'_>,
|
||||
result: Option<PyResult<Py<PyAny>>>,
|
||||
) -> PyResult<Py<PyAny>> {
|
||||
let mut body = {
|
||||
fn advance(slf: &Bound<'_, Self>, py: Python<'_>, input: Resume) -> PyResult<Py<PyAny>> {
|
||||
let (mut body, started) = {
|
||||
let mut execution = slf.borrow_mut();
|
||||
match (&execution.state, result.is_some()) {
|
||||
(ExecutionState::Created(_), false) | (ExecutionState::Suspended(_), true) => {}
|
||||
let started = match (&execution.state, &input) {
|
||||
(ExecutionState::Created(_), Resume::Send(value)) if value.is_none(py) => false,
|
||||
(ExecutionState::Created(_), Resume::Send(_)) => {
|
||||
return Err(Error::SendBeforeStart.into());
|
||||
}
|
||||
(ExecutionState::Suspended(_), _) => true,
|
||||
(ExecutionState::Running, _) => {
|
||||
return Err(PyRuntimeError::new_err("execution is already running"));
|
||||
return Err(Error::ExecutionRunning.into());
|
||||
}
|
||||
(ExecutionState::Closed, _) => {
|
||||
return Err(PyRuntimeError::new_err("execution is closed"));
|
||||
return Err(Error::ExecutionClosed.into());
|
||||
}
|
||||
_ => {
|
||||
return Err(PyRuntimeError::new_err(
|
||||
"execution requires start before resume and can only start once",
|
||||
));
|
||||
(ExecutionState::Created(_), Resume::Throw(_)) => {
|
||||
return Err(Error::ExecutionOutOfOrder.into());
|
||||
}
|
||||
}
|
||||
};
|
||||
match std::mem::replace(&mut execution.state, ExecutionState::Running) {
|
||||
ExecutionState::Created(body) | ExecutionState::Suspended(body) => body,
|
||||
ExecutionState::Created(body) | ExecutionState::Suspended(body) => (body, started),
|
||||
_ => unreachable!(),
|
||||
}
|
||||
};
|
||||
let lifecycle = slf.borrow().lifecycle;
|
||||
let result = match input {
|
||||
Resume::Send(value) => started.then_some(Ok(value)),
|
||||
Resume::Throw(error) => Some(Err(error)),
|
||||
};
|
||||
let steps = slf.borrow().steps;
|
||||
let outcome = catch_unwind(AssertUnwindSafe(|| {
|
||||
let step = body.resume(result)?;
|
||||
let (tag, value, suspended) = match step {
|
||||
ExecutionStep::Await(value) => ("Await", value, true),
|
||||
ExecutionStep::Open(head) => ("Open", head, true),
|
||||
ExecutionStep::Yield(value) => ("Yield", value, true),
|
||||
ExecutionStep::Suspended(Suspension::Await(value)) => ("Await", value, true),
|
||||
ExecutionStep::Suspended(Suspension::StreamOpen(head)) => {
|
||||
("StreamOpen", head, true)
|
||||
}
|
||||
ExecutionStep::Suspended(Suspension::StreamChunk(chunk)) => {
|
||||
("StreamChunk", chunk, true)
|
||||
}
|
||||
ExecutionStep::Return(value) => ("Complete", value, false),
|
||||
};
|
||||
let step = lifecycle(py)?.getattr(tag)?.call1((value,))?.unbind();
|
||||
let step = steps(py)?.getattr(tag)?.call1((value,))?.unbind();
|
||||
Ok((step, suspended))
|
||||
}))
|
||||
.map_err(panic_to_pyerr)
|
||||
|
|
@ -116,9 +109,7 @@ impl Execution {
|
|||
drop(body);
|
||||
outcome.and_then(|(step, suspended)| {
|
||||
if suspended {
|
||||
Err(PyRuntimeError::new_err(
|
||||
"execution was closed while running",
|
||||
))
|
||||
Err(Error::ExecutionClosedWhileRunning.into())
|
||||
} else {
|
||||
Ok(step)
|
||||
}
|
||||
|
|
@ -130,24 +121,16 @@ impl Execution {
|
|||
|
||||
#[pymethods]
|
||||
impl Execution {
|
||||
fn start(slf: &Bound<'_, Self>, py: Python<'_>) -> PyResult<Py<PyAny>> {
|
||||
Self::advance(slf, py, None)
|
||||
fn send(slf: &Bound<'_, Self>, py: Python<'_>, value: Py<PyAny>) -> PyResult<Py<PyAny>> {
|
||||
Self::advance(slf, py, Resume::Send(value))
|
||||
}
|
||||
|
||||
fn resume_value(
|
||||
slf: &Bound<'_, Self>,
|
||||
py: Python<'_>,
|
||||
value: Py<PyAny>,
|
||||
) -> PyResult<Py<PyAny>> {
|
||||
Self::advance(slf, py, Some(Ok(value)))
|
||||
}
|
||||
|
||||
fn resume_error(
|
||||
fn throw(
|
||||
slf: &Bound<'_, Self>,
|
||||
py: Python<'_>,
|
||||
error: Bound<'_, PyBaseException>,
|
||||
) -> PyResult<Py<PyAny>> {
|
||||
Self::advance(slf, py, Some(Err(PyErr::from_value(error.into_any()))))
|
||||
Self::advance(slf, py, Resume::Throw(PyErr::from_value(error.into_any())))
|
||||
}
|
||||
|
||||
fn close(slf: &Bound<'_, Self>) {
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ use pyo3::{
|
|||
types::PyDict,
|
||||
};
|
||||
|
||||
use crate::{HookResume, HookStep, PythonCallEvent, PythonCallHooks, PythonOwned, missing_state};
|
||||
use crate::{Error, HookResume, HookStep, PythonCallEvent, PythonCallHooks, PythonOwned};
|
||||
|
||||
pub(super) enum ChainStep<T> {
|
||||
Ready(T),
|
||||
|
|
@ -110,7 +110,7 @@ impl<H: PythonCallHooks> ChainHooks for HookAdapter<H> {
|
|||
result: PyResult<Py<PyAny>>,
|
||||
) -> PyResult<ChainStep<Py<PyDict>>> {
|
||||
let Some(Continuation::Arguments(resume)) = self.continuation.take() else {
|
||||
return Err(missing_state());
|
||||
return Err(Error::MissingState.into());
|
||||
};
|
||||
let step = resume(&mut self.hooks, py, result)?;
|
||||
Ok(self.step(step, Continuation::Arguments))
|
||||
|
|
@ -132,7 +132,7 @@ impl<H: PythonCallHooks> ChainHooks for HookAdapter<H> {
|
|||
result: PyResult<Py<PyAny>>,
|
||||
) -> PyResult<ChainStep<Box<WireRequest>>> {
|
||||
let Some(Continuation::Wire(resume)) = self.continuation.take() else {
|
||||
return Err(missing_state());
|
||||
return Err(Error::MissingState.into());
|
||||
};
|
||||
let step = resume(&mut self.hooks, py, result)?;
|
||||
Ok(self.step(step, Continuation::Wire))
|
||||
|
|
@ -154,7 +154,7 @@ impl<H: PythonCallHooks> ChainHooks for HookAdapter<H> {
|
|||
result: PyResult<Py<PyAny>>,
|
||||
) -> PyResult<ChainStep<Py<PyAny>>> {
|
||||
let Some(Continuation::Response(resume)) = self.continuation.take() else {
|
||||
return Err(missing_state());
|
||||
return Err(Error::MissingState.into());
|
||||
};
|
||||
let step = resume(&mut self.hooks, py, result)?;
|
||||
Ok(self.step(step, Continuation::Response))
|
||||
|
|
@ -171,7 +171,7 @@ impl<H: PythonCallHooks> ChainHooks for HookAdapter<H> {
|
|||
result: PyResult<Py<PyAny>>,
|
||||
) -> PyResult<ChainStep<()>> {
|
||||
let Some(Continuation::Event(resume)) = self.continuation.take() else {
|
||||
return Err(missing_state());
|
||||
return Err(Error::MissingState.into());
|
||||
};
|
||||
let step = resume(&mut self.hooks, py, result)?;
|
||||
Ok(self.step(step, Continuation::Event))
|
||||
|
|
|
|||
|
|
@ -10,9 +10,7 @@ use pyo3::{
|
|||
types::PyDict,
|
||||
};
|
||||
|
||||
use crate::{
|
||||
HookStep, PythonCallEvent, PythonCallHooks, PythonOwned, PythonRuntime, missing_state,
|
||||
};
|
||||
use crate::{Error, HookStep, PythonCallEvent, PythonCallHooks, PythonOwned, PythonRuntime};
|
||||
|
||||
use super::adapter::{ChainHooks, ChainStep, HookAdapter};
|
||||
|
||||
|
|
@ -72,7 +70,7 @@ impl HookChain {
|
|||
py: Python<'_>,
|
||||
result: PyResult<Py<PyAny>>,
|
||||
) -> PyResult<HookStep<Self, Py<PyDict>>> {
|
||||
let (index, context) = self.arguments.take().ok_or_else(missing_state)?;
|
||||
let (index, context) = self.arguments.take().ok_or(Error::MissingState)?;
|
||||
let result = resume_unless_cancelled(py, result)?;
|
||||
let step = self.hooks[index].resume_arguments(py, result)?;
|
||||
self.arguments_step(py, index, step, context)
|
||||
|
|
@ -113,7 +111,7 @@ impl HookChain {
|
|||
py: Python<'_>,
|
||||
result: PyResult<Py<PyAny>>,
|
||||
) -> PyResult<HookStep<Self, Box<WireRequest>>> {
|
||||
let (index, context) = self.wire.take().ok_or_else(missing_state)?;
|
||||
let (index, context) = self.wire.take().ok_or(Error::MissingState)?;
|
||||
let result = resume_unless_cancelled(py, result)?;
|
||||
let step = self.hooks[index].resume_wire(py, result)?;
|
||||
self.wire_step(py, index, step, &context)
|
||||
|
|
@ -154,7 +152,7 @@ impl HookChain {
|
|||
py: Python<'_>,
|
||||
result: PyResult<Py<PyAny>>,
|
||||
) -> PyResult<HookStep<Self, Py<PyAny>>> {
|
||||
let (index, context) = self.response.take().ok_or_else(missing_state)?;
|
||||
let (index, context) = self.response.take().ok_or(Error::MissingState)?;
|
||||
let result = resume_unless_cancelled(py, result)?;
|
||||
let step = self.hooks[index].resume_response(py, result)?;
|
||||
self.response_step(py, index, step, context)
|
||||
|
|
@ -195,7 +193,7 @@ impl HookChain {
|
|||
py: Python<'_>,
|
||||
result: PyResult<Py<PyAny>>,
|
||||
) -> PyResult<HookStep<Self, ()>> {
|
||||
let (index, event) = self.event.take().ok_or_else(missing_state)?;
|
||||
let (index, event) = self.event.take().ok_or(Error::MissingState)?;
|
||||
let result = resume_unless_cancelled(py, result)?;
|
||||
let result = self.hooks[index].resume_event(py, result);
|
||||
let step = notification_result(py, is_terminal(&event), result)?;
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@
|
|||
|
||||
mod argument;
|
||||
mod binding;
|
||||
mod caller;
|
||||
mod conversion_cache;
|
||||
mod driver;
|
||||
mod error;
|
||||
|
|
@ -23,11 +24,11 @@ pub use argument::{effective_py_args, lookup, present};
|
|||
pub use binding::PythonBinding;
|
||||
pub use conversion_cache::{FromPythonCache, ToPythonCache};
|
||||
pub use driver::{CallOptions, run_call};
|
||||
pub use error::{InvokeError, missing_state};
|
||||
pub use error::{Error, InvokeError};
|
||||
pub use file_reader::{FileContent, PythonFileReader, py_bytes};
|
||||
pub use fork_gate::RuntimeAlreadyStarted;
|
||||
pub use gil::{PythonContext, attach_blocking, release_count, release_gil};
|
||||
pub use handle::{Execution, ExecutionBody, ExecutionStep, PythonLifecycle};
|
||||
pub use handle::{Execution, ExecutionBody, ExecutionStep, StepTypes, Suspension};
|
||||
pub use hooks::{HookChain, HookResume, HookStep, PythonCallEvent, PythonCallHooks, PythonRuntime};
|
||||
pub use marshal::{
|
||||
Pythonized, from_py, from_py_argument, json_loads, json_object_field, panic_to_pyerr, to_py,
|
||||
|
|
|
|||
|
|
@ -5,11 +5,10 @@ use futures_util::future::{AbortHandle, Abortable};
|
|||
use litellm_host::call::HostedCompletion;
|
||||
use litellm_host::machine::{HostFailure, Machine, MachineStep};
|
||||
use litellm_host::protocol::Protocol;
|
||||
use pyo3::exceptions::PyRuntimeError;
|
||||
use pyo3::prelude::*;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
use crate::missing_state;
|
||||
use crate::Error;
|
||||
use crate::runtime::{poll_async_value, run_async_value, run_sync_value};
|
||||
|
||||
type NativeResult<M> = Result<
|
||||
|
|
@ -65,7 +64,7 @@ where
|
|||
py: Python<'_>,
|
||||
interruption: Option<HostFailure<<M::Protocol as Protocol>::Error>>,
|
||||
) -> PyResult<NativePoll<NativeResult<M>>> {
|
||||
let state = Arc::clone(self.state.as_ref().ok_or_else(missing_state)?);
|
||||
let state = Arc::clone(self.state.as_ref().ok_or(Error::MissingState)?);
|
||||
let future = async move {
|
||||
let mut state = state.lock().await;
|
||||
let result = match interruption {
|
||||
|
|
@ -90,7 +89,7 @@ where
|
|||
run_async_value(py, async move {
|
||||
Abortable::new(future, registration)
|
||||
.await
|
||||
.map_err(|_| PyRuntimeError::new_err("native execution closed"))?
|
||||
.map_err(|_| Error::NativeExecutionClosed)?
|
||||
})?
|
||||
.unbind(),
|
||||
))
|
||||
|
|
@ -103,12 +102,12 @@ where
|
|||
pub(super) fn take_result(&self) -> PyResult<NativeResult<M>> {
|
||||
self.state
|
||||
.as_ref()
|
||||
.ok_or_else(missing_state)?
|
||||
.ok_or(Error::MissingState)?
|
||||
.try_lock()
|
||||
.map_err(|_| missing_state())?
|
||||
.map_err(|_| Error::MissingState)?
|
||||
.result
|
||||
.take()
|
||||
.ok_or_else(missing_state)
|
||||
.ok_or_else(|| Error::MissingState.into())
|
||||
.map(|result| {
|
||||
result.map(|step| match step {
|
||||
MachineStep::Suspended(op) => MachineStep::Suspended(op),
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ use std::task::{Context, Poll, Waker};
|
|||
use std::time::Duration;
|
||||
|
||||
use crate::fork_gate::{ForkGate, Refused, RuntimeAlreadyStarted};
|
||||
use crate::{Pythonized, panic_to_pyerr, release_gil};
|
||||
use crate::{Error, Pythonized, panic_to_pyerr, release_gil};
|
||||
use futures_util::FutureExt;
|
||||
use pyo3::exceptions::PyRuntimeError;
|
||||
use pyo3::prelude::*;
|
||||
|
|
@ -112,9 +112,7 @@ where
|
|||
F: Future<Output = PyResult<T>> + Send + 'static,
|
||||
{
|
||||
if Handle::try_current().is_ok() {
|
||||
return Err(PyRuntimeError::new_err(
|
||||
"synchronous native routes cannot run from a Tokio context; use the async route",
|
||||
));
|
||||
return Err(Error::SyncInsideTokio.into());
|
||||
}
|
||||
release_gil(py, move || runtime.block_on(wait_for_sync_result(future)))?
|
||||
}
|
||||
|
|
@ -131,9 +129,7 @@ where
|
|||
F: Future<Output = Result<T, E>> + Send + 'static,
|
||||
{
|
||||
if Handle::try_current().is_ok() {
|
||||
return Err(PyRuntimeError::new_err(
|
||||
"synchronous native routes cannot run from a Tokio context; use the async route",
|
||||
));
|
||||
return Err(Error::SyncInsideTokio.into());
|
||||
}
|
||||
|
||||
let result = release_gil(py, move || runtime.block_on(wait_for_sync_result(future)))?;
|
||||
|
|
|
|||
770
litellm-rust/crates/host-python/tests/driver.rs
Normal file
770
litellm-rust/crates/host-python/tests/driver.rs
Normal file
|
|
@ -0,0 +1,770 @@
|
|||
mod support;
|
||||
|
||||
use std::num::NonZeroUsize;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
|
||||
use litellm_host::{
|
||||
hooks::CallHooks,
|
||||
lifecycle::{CallEvent, ExecutionEvent, FailureOrigin},
|
||||
machine::CallMachine,
|
||||
observation::observation_channel,
|
||||
};
|
||||
use litellm_host_python::{CallOptions, HookChain, PythonOwned, PythonRuntime, run_call};
|
||||
use pyo3::{
|
||||
exceptions::{
|
||||
PyBaseException, PyException, PyLookupError, PyRuntimeError, PyStopIteration, PyTypeError,
|
||||
PyValueError, asyncio::CancelledError,
|
||||
},
|
||||
gc::{PyTraverseError, PyVisit},
|
||||
prelude::*,
|
||||
types::PyDict,
|
||||
};
|
||||
use rstest::rstest;
|
||||
use support::{
|
||||
HookScript, Interpreter, Log, MachineFactory, OpScript, Synthetic, binding, call_options,
|
||||
failing_machine, handed, hooks, interpreter, lifecycle_module, open_stream, read_all,
|
||||
run_composed, run_scripted, settled, success_machine,
|
||||
};
|
||||
|
||||
fn observed(capacity: usize) -> (CallOptions, impl FnMut() -> Vec<CallEvent>) {
|
||||
let (sender, mut receiver) = observation_channel(NonZeroUsize::new(capacity).unwrap());
|
||||
let options = CallOptions {
|
||||
asynchronous: false,
|
||||
steps: lifecycle_module,
|
||||
observers: Some(sender),
|
||||
};
|
||||
(options, move || {
|
||||
std::iter::from_fn(|| receiver.try_recv().ok()).collect()
|
||||
})
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::success(OpScript::AwaitAnswer)]
|
||||
#[case::failure(OpScript::AwaitFailure)]
|
||||
fn asynchronous_host_operations_resume_the_same_machine(
|
||||
#[from(interpreter)] _python: Interpreter,
|
||||
#[case] op: OpScript,
|
||||
) {
|
||||
Python::attach(|py| {
|
||||
let machine: MachineFactory = Box::new(|_| {
|
||||
CallMachine::<Synthetic>::new(None, |host| {
|
||||
Box::pin(async move { host.services.call(|reply| ("read", reply)).await })
|
||||
})
|
||||
});
|
||||
let (result, log) = run_scripted(py, machine, op, HookScript::Plain, true);
|
||||
match op {
|
||||
OpScript::AwaitAnswer => {
|
||||
assert_eq!(result.unwrap().extract::<String>(py).unwrap(), "awaited")
|
||||
}
|
||||
OpScript::AwaitFailure => {
|
||||
assert!(result.unwrap_err().is_instance_of::<PyLookupError>(py));
|
||||
assert_eq!(
|
||||
log.iter()
|
||||
.filter(|entry| entry.starts_with("failed:"))
|
||||
.count(),
|
||||
1
|
||||
);
|
||||
assert!(!log.iter().any(|entry| entry.starts_with("succeeded:")));
|
||||
}
|
||||
_ => unreachable!(),
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn composed_hooks_share_prepared_arguments_and_one_execution(
|
||||
#[from(interpreter)] _python: Interpreter,
|
||||
#[values(false, true)] asynchronous: bool,
|
||||
) {
|
||||
Python::attach(|py| {
|
||||
let log = Log::default();
|
||||
let chain = HookChain::new()
|
||||
.with(hooks(&log, HookScript::ObserveArguments))
|
||||
.with(hooks(&log, HookScript::RewriteArguments))
|
||||
.with(hooks(&log, HookScript::ReplaceResponse));
|
||||
let result = handed(
|
||||
py,
|
||||
run_call(
|
||||
py,
|
||||
move |_, _, request| Ok(success_machine()(request)),
|
||||
binding(&log, OpScript::Answer),
|
||||
chain,
|
||||
PyDict::new(py).unbind(),
|
||||
call_options(asynchronous),
|
||||
),
|
||||
asynchronous,
|
||||
);
|
||||
if asynchronous {
|
||||
assert!(log.entries().is_empty());
|
||||
}
|
||||
let response = settled(py, result, asynchronous).unwrap();
|
||||
assert_eq!(response.extract::<String>(py).unwrap(), "replaced");
|
||||
let entries = log.entries();
|
||||
assert!(entries.iter().any(|entry| entry == "adopted:hook"));
|
||||
assert_eq!(
|
||||
entries.iter().filter(|entry| *entry == "project").count(),
|
||||
1
|
||||
);
|
||||
assert_eq!(
|
||||
entries.iter().filter(|entry| *entry == "op:sign").count(),
|
||||
1
|
||||
);
|
||||
assert_eq!(
|
||||
entries
|
||||
.iter()
|
||||
.filter(|entry| *entry == "succeeded:replaced")
|
||||
.count(),
|
||||
3
|
||||
);
|
||||
assert!(!entries.iter().any(|entry| entry.starts_with("failed:")));
|
||||
});
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::draining(4, false)]
|
||||
#[case::full(1, false)]
|
||||
#[case::closed(4, true)]
|
||||
fn observation_delivery_preserves_callback_order_and_response(
|
||||
#[from(interpreter)] _python: Interpreter,
|
||||
#[case] capacity: usize,
|
||||
#[case] closed: bool,
|
||||
#[values(false, true)] asynchronous: bool,
|
||||
) {
|
||||
Python::attach(|py| {
|
||||
let (expected, expected_log) = run_scripted(
|
||||
py,
|
||||
success_machine(),
|
||||
OpScript::Answer,
|
||||
HookScript::ReplaceResponse,
|
||||
asynchronous,
|
||||
);
|
||||
let (sender, mut receiver) = observation_channel(NonZeroUsize::new(capacity).unwrap());
|
||||
if closed {
|
||||
receiver.close();
|
||||
}
|
||||
let log = Log::default();
|
||||
let (actual, actual_log) = run_composed(
|
||||
py,
|
||||
success_machine(),
|
||||
binding(&log, OpScript::Answer),
|
||||
hooks(&log, HookScript::ReplaceResponse),
|
||||
std::convert::identity,
|
||||
CallOptions {
|
||||
asynchronous,
|
||||
steps: lifecycle_module,
|
||||
observers: Some(sender.clone()),
|
||||
},
|
||||
);
|
||||
assert_eq!(
|
||||
actual.unwrap().extract::<String>(py).unwrap(),
|
||||
expected.unwrap().extract::<String>(py).unwrap()
|
||||
);
|
||||
assert_eq!(actual_log, expected_log);
|
||||
let events: Vec<_> = std::iter::from_fn(|| receiver.try_recv().ok()).collect();
|
||||
assert_eq!(events.len() as u64 + sender.dropped_events(), 3);
|
||||
if !closed && capacity >= 3 {
|
||||
let [
|
||||
CallEvent::Started { start_time },
|
||||
CallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }),
|
||||
CallEvent::Succeeded { timing, .. },
|
||||
] = events.as_slice()
|
||||
else {
|
||||
panic!("expected start, provider response and success: {events:?}");
|
||||
};
|
||||
assert_eq!(raw.body, "raw");
|
||||
assert_eq!(*start_time, timing.start_time);
|
||||
assert!(timing.end_time >= timing.start_time);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::preparation(HookScript::FailBegin, false, 1)]
|
||||
#[case::transformation(HookScript::FailAfterSuccess, false, 1)]
|
||||
#[case::terminal_cancellation(HookScript::CancelTerminal, true, 0)]
|
||||
#[case::terminal_failure(HookScript::FailTerminal, false, 0)]
|
||||
fn a_terminal_hook_failure_is_reported_once_and_never_replays_callbacks(
|
||||
#[from(interpreter)] _python: Interpreter,
|
||||
#[case] script: HookScript,
|
||||
#[case] cancelled: bool,
|
||||
#[case] failure_callbacks: usize,
|
||||
#[values(false, true)] asynchronous: bool,
|
||||
) {
|
||||
Python::attach(|py| {
|
||||
let (options, mut events) = observed(4);
|
||||
let log = Log::default();
|
||||
let (result, log) = run_composed(
|
||||
py,
|
||||
success_machine(),
|
||||
binding(&log, OpScript::Answer),
|
||||
hooks(&log, script),
|
||||
std::convert::identity,
|
||||
CallOptions {
|
||||
asynchronous,
|
||||
..options
|
||||
},
|
||||
);
|
||||
let error = result.unwrap_err();
|
||||
assert_eq!(error.is_instance_of::<CancelledError>(py), cancelled);
|
||||
let events = events();
|
||||
assert!(matches!(events.first(), Some(CallEvent::Started { .. })));
|
||||
assert_eq!(
|
||||
matches!(events.last(), Some(CallEvent::Cancelled { .. })),
|
||||
cancelled
|
||||
);
|
||||
assert_eq!(
|
||||
matches!(
|
||||
events.last(),
|
||||
Some(CallEvent::Failed {
|
||||
origin: FailureOrigin::Host,
|
||||
..
|
||||
})
|
||||
),
|
||||
!cancelled
|
||||
);
|
||||
let terminal = events.iter().filter(|event| {
|
||||
matches!(
|
||||
event,
|
||||
CallEvent::Failed { .. }
|
||||
| CallEvent::Succeeded { .. }
|
||||
| CallEvent::Cancelled { .. }
|
||||
)
|
||||
});
|
||||
assert_eq!(terminal.count(), 1);
|
||||
assert_eq!(
|
||||
log.iter()
|
||||
.filter(|entry| entry.starts_with("failed:"))
|
||||
.count(),
|
||||
failure_callbacks
|
||||
);
|
||||
assert!(
|
||||
!log.iter()
|
||||
.any(|entry| entry.starts_with("succeeded:") || entry == "cancelled")
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::preparation(HookScript::FailBegin, OpScript::Answer)]
|
||||
#[case::native_decode(HookScript::Plain, OpScript::RejectRequestNatively)]
|
||||
#[case::python_decode(HookScript::Plain, OpScript::RaiseRequestPython)]
|
||||
fn startup_failures_release_the_factory_without_constructing_a_machine(
|
||||
#[from(interpreter)] _python: Interpreter,
|
||||
#[case] script: HookScript,
|
||||
#[case] op: OpScript,
|
||||
#[values(false, true)] asynchronous: bool,
|
||||
) {
|
||||
struct Release(Arc<AtomicBool>);
|
||||
impl Drop for Release {
|
||||
fn drop(&mut self) {
|
||||
self.0.store(true, Ordering::SeqCst);
|
||||
}
|
||||
}
|
||||
Python::attach(|py| {
|
||||
let constructed = Arc::new(AtomicBool::new(false));
|
||||
let did_construct = constructed.clone();
|
||||
let released = Arc::new(AtomicBool::new(false));
|
||||
let release = Release(released.clone());
|
||||
let machine: MachineFactory = Box::new(move |request| {
|
||||
let _release = release;
|
||||
did_construct.store(true, Ordering::SeqCst);
|
||||
success_machine()(request)
|
||||
});
|
||||
let (result, log) = run_scripted(py, machine, op, script, asynchronous);
|
||||
assert!(result.is_err());
|
||||
assert!(!constructed.load(Ordering::SeqCst));
|
||||
assert!(released.load(Ordering::SeqCst));
|
||||
assert_eq!(
|
||||
log.iter()
|
||||
.filter(|event| event.starts_with("failed:"))
|
||||
.count(),
|
||||
1
|
||||
);
|
||||
assert_eq!(
|
||||
log.iter().filter(|event| *event == "adapter.close").count(),
|
||||
1
|
||||
);
|
||||
assert_eq!(log.iter().filter(|event| *event == "host.close").count(), 1);
|
||||
});
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn success_runs_every_step_in_order_and_returns_the_public_response(
|
||||
#[from(interpreter)] _python: Interpreter,
|
||||
#[values(false, true)] asynchronous: bool,
|
||||
) {
|
||||
Python::attach(|py| {
|
||||
let (result, log) = run_scripted(
|
||||
py,
|
||||
success_machine(),
|
||||
OpScript::Answer,
|
||||
HookScript::Plain,
|
||||
asynchronous,
|
||||
);
|
||||
assert_eq!(
|
||||
result.unwrap().extract::<String>(py).unwrap(),
|
||||
"project:1|sign|rewritten"
|
||||
);
|
||||
assert_eq!(
|
||||
log,
|
||||
[
|
||||
"started",
|
||||
"begin",
|
||||
"project",
|
||||
"op:sign",
|
||||
"before_provider_request",
|
||||
"response:raw",
|
||||
"complete",
|
||||
"after_success",
|
||||
"succeeded:project:1|sign|rewritten",
|
||||
"adapter.close",
|
||||
"host.close",
|
||||
]
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn explicitly_closing_a_stream_dispatches_success_for_delivered_chunks(
|
||||
#[from(interpreter)] _python: Interpreter,
|
||||
#[values(false, true)] asynchronous: bool,
|
||||
) {
|
||||
Python::attach(|py| {
|
||||
let log = Log::default();
|
||||
let (options, mut events) = observed(3);
|
||||
let stream = open_stream(
|
||||
py,
|
||||
hooks(&log, HookScript::Plain),
|
||||
CallOptions {
|
||||
asynchronous,
|
||||
..options
|
||||
},
|
||||
);
|
||||
if asynchronous {
|
||||
for method in ["__anext__", "aclose", "aclose"] {
|
||||
let stop = stream
|
||||
.call_method0(method)
|
||||
.unwrap()
|
||||
.call_method1("send", (py.None(),))
|
||||
.unwrap_err();
|
||||
assert!(stop.is_instance_of::<PyStopIteration>(py));
|
||||
}
|
||||
} else {
|
||||
assert_eq!(
|
||||
stream
|
||||
.call_method0("__next__")
|
||||
.unwrap()
|
||||
.extract::<String>()
|
||||
.unwrap(),
|
||||
"first"
|
||||
);
|
||||
stream.call_method0("close").unwrap();
|
||||
stream.call_method0("close").unwrap();
|
||||
}
|
||||
assert!(matches!(
|
||||
events().as_slice(),
|
||||
[CallEvent::Started { .. }, CallEvent::Succeeded { .. }]
|
||||
));
|
||||
assert_eq!(
|
||||
log.entries(),
|
||||
[
|
||||
"started",
|
||||
"begin",
|
||||
"opened",
|
||||
"delivered",
|
||||
"succeeded:None",
|
||||
"adapter.close"
|
||||
]
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn a_stream_carries_its_head_before_the_first_chunk(
|
||||
#[from(interpreter)] _python: Interpreter,
|
||||
#[values(false, true)] asynchronous: bool,
|
||||
) {
|
||||
Python::attach(|py| {
|
||||
let log = Log::default();
|
||||
let stream = open_stream(
|
||||
py,
|
||||
hooks(&log, HookScript::Plain),
|
||||
call_options(asynchronous),
|
||||
);
|
||||
let hidden: std::collections::HashMap<String, std::collections::HashMap<String, String>> =
|
||||
stream.getattr("_hidden_params").unwrap().extract().unwrap();
|
||||
assert_eq!(
|
||||
hidden["additional_headers"],
|
||||
std::collections::HashMap::from([("request-id".to_string(), "req_1".to_string())])
|
||||
);
|
||||
assert_eq!(log.entries(), ["started", "begin", "opened"]);
|
||||
assert_eq!(read_all(py, &stream, asynchronous), ["first", "second"]);
|
||||
});
|
||||
}
|
||||
|
||||
struct Classification {
|
||||
provider_fails: bool,
|
||||
op: OpScript,
|
||||
classifier_fails: bool,
|
||||
message: &'static str,
|
||||
native_context: Option<&'static str>,
|
||||
steps: &'static [&'static str],
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::provider_failure(Classification {
|
||||
provider_fails: true,
|
||||
op: OpScript::Answer,
|
||||
classifier_fails: false,
|
||||
message: "classified: provider exploded",
|
||||
native_context: None,
|
||||
steps: &["project", "classify:provider exploded"],
|
||||
})]
|
||||
#[case::host_op_rejection(Classification {
|
||||
provider_fails: false,
|
||||
op: OpScript::RejectNatively,
|
||||
classifier_fails: false,
|
||||
message: "classified: op rejected",
|
||||
native_context: None,
|
||||
steps: &["project", "op:sign", "classify:op rejected"],
|
||||
})]
|
||||
#[case::host_op_python_exception(Classification {
|
||||
provider_fails: false,
|
||||
op: OpScript::RaisePython,
|
||||
classifier_fails: false,
|
||||
message: "op failed",
|
||||
native_context: None,
|
||||
steps: &["project", "op:sign"],
|
||||
})]
|
||||
#[case::failing_classifier(Classification {
|
||||
provider_fails: true,
|
||||
op: OpScript::Answer,
|
||||
classifier_fails: true,
|
||||
message: "classifier failed",
|
||||
native_context: Some("provider exploded"),
|
||||
steps: &["project", "classify:provider exploded"],
|
||||
})]
|
||||
fn a_native_failure_is_classified_once_and_a_python_exception_is_raised_as_is(
|
||||
#[from(interpreter)] _python: Interpreter,
|
||||
#[case] case: Classification,
|
||||
#[values(false, true)] asynchronous: bool,
|
||||
) {
|
||||
Python::attach(|py| {
|
||||
let Classification {
|
||||
provider_fails,
|
||||
op,
|
||||
classifier_fails,
|
||||
message,
|
||||
native_context,
|
||||
steps,
|
||||
} = case;
|
||||
let log = Log::default();
|
||||
let binding = if classifier_fails {
|
||||
binding(&log, op).with_failing_classifier()
|
||||
} else {
|
||||
binding(&log, op)
|
||||
};
|
||||
let machine = if provider_fails {
|
||||
failing_machine()
|
||||
} else {
|
||||
success_machine()
|
||||
};
|
||||
let (result, log) = run_composed(
|
||||
py,
|
||||
machine,
|
||||
binding,
|
||||
hooks(&log, HookScript::Plain),
|
||||
std::convert::identity,
|
||||
call_options(asynchronous),
|
||||
);
|
||||
let error = result.unwrap_err();
|
||||
assert_eq!(error.value(py).to_string(), message);
|
||||
match native_context {
|
||||
Some(context) => {
|
||||
assert!(error.is_instance_of::<PyTypeError>(py));
|
||||
let chained = error.context(py).unwrap();
|
||||
assert!(chained.is_instance_of::<PyRuntimeError>(py));
|
||||
assert_eq!(chained.value(py).to_string(), context);
|
||||
}
|
||||
None => assert!(error.is_instance_of::<PyValueError>(py)),
|
||||
}
|
||||
let expected: Vec<String> = ["started", "begin"]
|
||||
.iter()
|
||||
.chain(steps)
|
||||
.map(|step| step.to_string())
|
||||
.chain([
|
||||
format!("failed:Call:{message}"),
|
||||
"adapter.close".into(),
|
||||
"host.close".into(),
|
||||
])
|
||||
.collect();
|
||||
assert_eq!(log, expected);
|
||||
});
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn begin_failures_are_host_failures_without_provider_mapping(
|
||||
#[from(interpreter)] _python: Interpreter,
|
||||
) {
|
||||
Python::attach(|py| {
|
||||
let (result, log) = run_scripted(
|
||||
py,
|
||||
success_machine(),
|
||||
OpScript::Answer,
|
||||
HookScript::FailBegin,
|
||||
false,
|
||||
);
|
||||
assert_eq!(result.unwrap_err().value(py).to_string(), "begin failed");
|
||||
assert_eq!(
|
||||
log,
|
||||
[
|
||||
"started",
|
||||
"begin",
|
||||
"failed:Host:begin failed",
|
||||
"adapter.close",
|
||||
"host.close"
|
||||
]
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
enum ArgumentPolicy {
|
||||
Inherit,
|
||||
Reject(Py<PyBaseException>),
|
||||
}
|
||||
|
||||
impl CallHooks<PythonRuntime> for ArgumentPolicy {
|
||||
fn arguments_prepared(&mut self, py: Python<'_>, arguments: &Py<PyDict>) -> PyResult<()> {
|
||||
match self {
|
||||
Self::Inherit => arguments.bind(py).set_item("api_key", "inherited"),
|
||||
Self::Reject(error) => Err(PyErr::from_value(error.bind(py).clone().into_any())),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl PythonOwned for ArgumentPolicy {
|
||||
fn close(&mut self, _: Python<'_>) {}
|
||||
|
||||
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
|
||||
match self {
|
||||
Self::Inherit => Ok(()),
|
||||
Self::Reject(error) => visit.call(error),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn resource_setup_uses_prepared_arguments_and_failures_are_terminal(
|
||||
#[from(interpreter)] _python: Interpreter,
|
||||
#[values(false, true)] asynchronous: bool,
|
||||
#[values(false, true)] fail_setup: bool,
|
||||
) {
|
||||
Python::attach(|py| {
|
||||
let log = Log::default();
|
||||
let setup_log = log.clone();
|
||||
let error = PyValueError::new_err("setup failed").into_value(py);
|
||||
let setup_error = error.clone_ref(py);
|
||||
let arguments = PyDict::new(py);
|
||||
arguments.set_item("api_key", "original").unwrap();
|
||||
let result = handed(
|
||||
py,
|
||||
run_call(
|
||||
py,
|
||||
move |py, prepared, request| {
|
||||
let source: String = prepared.get_item("prepared")?.unwrap().extract()?;
|
||||
let key: String = prepared.get_item("api_key")?.unwrap().extract()?;
|
||||
setup_log.push(format!("setup:{source}:{key}"));
|
||||
if fail_setup {
|
||||
return Err(PyErr::from_value(setup_error.into_bound(py).into_any()));
|
||||
}
|
||||
Ok(success_machine()(request))
|
||||
},
|
||||
binding(&log, OpScript::Answer),
|
||||
HookChain::new()
|
||||
.with(hooks(&log, HookScript::RewriteArguments))
|
||||
.with(ArgumentPolicy::Inherit),
|
||||
arguments.clone().unbind(),
|
||||
call_options(asynchronous),
|
||||
),
|
||||
asynchronous,
|
||||
);
|
||||
if asynchronous {
|
||||
assert!(log.entries().is_empty());
|
||||
}
|
||||
let settled = settled(py, result, asynchronous);
|
||||
assert_eq!(
|
||||
arguments
|
||||
.get_item("api_key")
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.extract::<String>()
|
||||
.unwrap(),
|
||||
"original"
|
||||
);
|
||||
assert_eq!(
|
||||
&log.entries()[..4],
|
||||
["started", "begin", "project", "setup:hook:inherited"]
|
||||
);
|
||||
if fail_setup {
|
||||
assert!(settled.unwrap_err().value(py).is(error.bind(py)));
|
||||
assert_eq!(
|
||||
log.entries(),
|
||||
[
|
||||
"started",
|
||||
"begin",
|
||||
"project",
|
||||
"setup:hook:inherited",
|
||||
"failed:Host:setup failed",
|
||||
"adapter.close",
|
||||
"host.close",
|
||||
]
|
||||
);
|
||||
} else {
|
||||
assert!(settled.is_ok());
|
||||
assert!(
|
||||
log.entries()
|
||||
.iter()
|
||||
.any(|entry| entry.starts_with("succeeded:"))
|
||||
);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn closing_an_unstarted_call_never_initializes_resources(
|
||||
#[from(interpreter)] _python: Interpreter,
|
||||
) {
|
||||
Python::attach(|py| {
|
||||
let log = Log::default();
|
||||
let setup_log = log.clone();
|
||||
let pending = handed(
|
||||
py,
|
||||
run_call(
|
||||
py,
|
||||
move |_, _, request| {
|
||||
setup_log.push("setup");
|
||||
Ok(success_machine()(request))
|
||||
},
|
||||
binding(&log, OpScript::Answer),
|
||||
hooks(&log, HookScript::Plain),
|
||||
PyDict::new(py).unbind(),
|
||||
call_options(true),
|
||||
),
|
||||
true,
|
||||
)
|
||||
.unwrap();
|
||||
assert!(log.entries().is_empty());
|
||||
pending.call_method0(py, "close").unwrap();
|
||||
drop(pending);
|
||||
assert_eq!(log.entries(), ["adapter.close", "host.close"]);
|
||||
});
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn an_argument_policy_rejection_is_the_callers_error_and_the_machine_never_starts(
|
||||
#[from(interpreter)] _python: Interpreter,
|
||||
#[values(false, true)] asynchronous: bool,
|
||||
) {
|
||||
Python::attach(|py| {
|
||||
let raised = PyValueError::new_err("over budget").into_value(py);
|
||||
let log = Log::default();
|
||||
let (result, log) = run_composed(
|
||||
py,
|
||||
success_machine(),
|
||||
binding(&log, OpScript::Answer),
|
||||
hooks(&log, HookScript::Plain),
|
||||
|hooks| {
|
||||
HookChain::new()
|
||||
.with(hooks)
|
||||
.with(ArgumentPolicy::Reject(raised.clone_ref(py)))
|
||||
},
|
||||
call_options(asynchronous),
|
||||
);
|
||||
assert!(result.unwrap_err().value(py).is(raised.bind(py)));
|
||||
assert_eq!(
|
||||
log,
|
||||
[
|
||||
"started",
|
||||
"begin",
|
||||
"failed:Host:over budget",
|
||||
"adapter.close",
|
||||
"host.close"
|
||||
]
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn the_host_projects_from_the_keyword_view_the_argument_policy_rewrote(
|
||||
#[from(interpreter)] _python: Interpreter,
|
||||
#[values(false, true)] asynchronous: bool,
|
||||
) {
|
||||
Python::attach(|py| {
|
||||
let log = Log::default();
|
||||
let (result, _) = run_composed(
|
||||
py,
|
||||
success_machine(),
|
||||
binding(&log, OpScript::Answer),
|
||||
hooks(&log, HookScript::Plain),
|
||||
|hooks| HookChain::new().with(hooks).with(ArgumentPolicy::Inherit),
|
||||
call_options(asynchronous),
|
||||
);
|
||||
assert_eq!(
|
||||
result.unwrap().extract::<String>(py).unwrap(),
|
||||
"project:2|sign|rewritten"
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn a_failure_while_finalizing_fails_the_call_instead_of_succeeding(
|
||||
#[from(interpreter)] _python: Interpreter,
|
||||
#[values(false, true)] asynchronous: bool,
|
||||
) {
|
||||
Python::attach(|py| {
|
||||
let (result, log) = run_scripted(
|
||||
py,
|
||||
success_machine(),
|
||||
OpScript::Answer,
|
||||
HookScript::FailAfterSuccess,
|
||||
asynchronous,
|
||||
);
|
||||
assert_eq!(
|
||||
result.unwrap_err().value(py).to_string(),
|
||||
"after_success failed"
|
||||
);
|
||||
assert_eq!(
|
||||
&log[log.len() - 4..],
|
||||
[
|
||||
"after_success",
|
||||
"failed:Host:after_success failed",
|
||||
"adapter.close",
|
||||
"host.close"
|
||||
]
|
||||
);
|
||||
assert!(!log.iter().any(|entry| entry.starts_with("succeeded")));
|
||||
});
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn cancellation_ends_the_call_without_terminal_dispatch(
|
||||
#[from(interpreter)] _python: Interpreter,
|
||||
#[values(false, true)] asynchronous: bool,
|
||||
) {
|
||||
Python::attach(|py| {
|
||||
let (result, log) = run_scripted(
|
||||
py,
|
||||
success_machine(),
|
||||
OpScript::CancelRequest,
|
||||
HookScript::Plain,
|
||||
asynchronous,
|
||||
);
|
||||
let error = result.unwrap_err();
|
||||
assert!(!error.is_instance_of::<PyException>(py));
|
||||
assert_eq!(
|
||||
log,
|
||||
["started", "begin", "project", "adapter.close", "host.close"]
|
||||
);
|
||||
});
|
||||
}
|
||||
343
litellm-rust/crates/host-python/tests/handle.rs
Normal file
343
litellm-rust/crates/host-python/tests/handle.rs
Normal file
|
|
@ -0,0 +1,343 @@
|
|||
mod support;
|
||||
|
||||
use litellm_host::{hooks::CallHooks, lifecycle::Timing};
|
||||
use litellm_host_python::{
|
||||
Execution, ExecutionBody, ExecutionStep, HookStep, PythonCallEvent, PythonOwned, PythonRuntime,
|
||||
Suspension, run_call,
|
||||
};
|
||||
use pyo3::{
|
||||
exceptions::PyBaseException,
|
||||
gc::{PyTraverseError, PyVisit},
|
||||
prelude::*,
|
||||
types::{PyCFunction, PyDict},
|
||||
};
|
||||
use rstest::rstest;
|
||||
use support::{
|
||||
Interpreter, Log, OpScript, binding, call_options, handed, install_lifecycle_module,
|
||||
interpreter, lifecycle_module, success_machine,
|
||||
};
|
||||
|
||||
struct RetainingHost {
|
||||
retained: Option<Py<PyAny>>,
|
||||
}
|
||||
|
||||
impl ExecutionBody for RetainingHost {
|
||||
fn resume(&mut self, _: Option<PyResult<Py<PyAny>>>) -> PyResult<ExecutionStep> {
|
||||
Python::attach(|py| Ok(ExecutionStep::Return(py.None())))
|
||||
}
|
||||
|
||||
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
|
||||
visit.call(&self.retained)
|
||||
}
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
fn retaining_coroutine(py: Python<'_>, retained: Py<PyAny>) -> PyResult<Py<Execution>> {
|
||||
Py::new(
|
||||
py,
|
||||
Execution::new(
|
||||
RetainingHost {
|
||||
retained: Some(retained),
|
||||
},
|
||||
lifecycle_module,
|
||||
),
|
||||
)
|
||||
}
|
||||
|
||||
struct AwaitBody(Option<Py<PyAny>>);
|
||||
|
||||
impl ExecutionBody for AwaitBody {
|
||||
fn resume(&mut self, result: Option<PyResult<Py<PyAny>>>) -> PyResult<ExecutionStep> {
|
||||
match self.0.take() {
|
||||
Some(awaitable) => Ok(ExecutionStep::Suspended(Suspension::Await(awaitable))),
|
||||
None => result
|
||||
.expect("selected await completed")
|
||||
.map(ExecutionStep::Return),
|
||||
}
|
||||
}
|
||||
|
||||
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
|
||||
visit.call(&self.0)
|
||||
}
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
fn await_execution(awaitable: Py<PyAny>) -> Execution {
|
||||
Execution::new(AwaitBody(Some(awaitable)), lifecycle_module)
|
||||
}
|
||||
|
||||
struct CallingBody(Py<PyAny>);
|
||||
|
||||
impl ExecutionBody for CallingBody {
|
||||
fn resume(&mut self, _: Option<PyResult<Py<PyAny>>>) -> PyResult<ExecutionStep> {
|
||||
Python::attach(|py| self.0.call0(py).map(ExecutionStep::Return))
|
||||
}
|
||||
|
||||
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
|
||||
visit.call(&self.0)
|
||||
}
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
fn calling_execution(callback: Py<PyAny>) -> Execution {
|
||||
Execution::new(CallingBody(callback), lifecycle_module)
|
||||
}
|
||||
|
||||
struct ErrorBody(Option<Py<PyBaseException>>);
|
||||
|
||||
impl ExecutionBody for ErrorBody {
|
||||
fn resume(&mut self, _: Option<PyResult<Py<PyAny>>>) -> PyResult<ExecutionStep> {
|
||||
Python::attach(|py| {
|
||||
Err(PyErr::from_value(
|
||||
self.0.take().unwrap().into_bound(py).into_any(),
|
||||
))
|
||||
})
|
||||
}
|
||||
|
||||
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
|
||||
visit.call(&self.0)
|
||||
}
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
fn error_execution(error: Bound<'_, PyBaseException>) -> Execution {
|
||||
Execution::new(ErrorBody(Some(error.unbind())), lifecycle_module)
|
||||
}
|
||||
|
||||
/// Where a paused call keeps the test's object: in its hooks, or selected as the call's
|
||||
/// response or exception, with the call paused in the hook that is told about it.
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
enum Holder {
|
||||
Hooks,
|
||||
Response,
|
||||
Exception,
|
||||
}
|
||||
|
||||
struct PausingHooks {
|
||||
retained: Option<Py<PyAny>>,
|
||||
pause: Option<Py<PyAny>>,
|
||||
holder: Holder,
|
||||
}
|
||||
|
||||
impl CallHooks<PythonRuntime> for PausingHooks {
|
||||
fn transform_response(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
response: Py<PyAny>,
|
||||
_: Timing,
|
||||
) -> PyResult<HookStep<Self, Py<PyAny>>> {
|
||||
match self.holder {
|
||||
Holder::Hooks => Ok(HookStep::Ready(response)),
|
||||
Holder::Response => Ok(HookStep::Ready(self.retained.take().unwrap())),
|
||||
Holder::Exception => Err(PyErr::from_value(
|
||||
self.retained.take().unwrap().into_bound(py),
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
fn on_event(
|
||||
&mut self,
|
||||
_: Python<'_>,
|
||||
event: PythonCallEvent<'_>,
|
||||
) -> PyResult<HookStep<Self, ()>> {
|
||||
let pauses_here = match self.holder {
|
||||
Holder::Hooks => true,
|
||||
Holder::Response => matches!(event, PythonCallEvent::Succeeded { .. }),
|
||||
Holder::Exception => matches!(event, PythonCallEvent::Failed { .. }),
|
||||
};
|
||||
Ok(match self.pause.take_if(|_| pauses_here) {
|
||||
Some(pause) => {
|
||||
HookStep::Await(pause, |_, _, answer| answer.map(|_| HookStep::Ready(())))
|
||||
}
|
||||
None => HookStep::Ready(()),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl PythonOwned for PausingHooks {
|
||||
fn close(&mut self, _: Python<'_>) {
|
||||
self.retained = None;
|
||||
self.pause = None;
|
||||
}
|
||||
|
||||
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
|
||||
visit.call(&self.retained)?;
|
||||
visit.call(&self.pause)
|
||||
}
|
||||
}
|
||||
|
||||
fn paused_call(
|
||||
py: Python<'_>,
|
||||
retained: Py<PyAny>,
|
||||
pause: Py<PyAny>,
|
||||
holder: Holder,
|
||||
) -> PyResult<Py<PyAny>> {
|
||||
let machine = success_machine();
|
||||
handed(
|
||||
py,
|
||||
run_call(
|
||||
py,
|
||||
move |_, _, request| Ok(machine(request)),
|
||||
binding(&Log::default(), OpScript::Answer),
|
||||
PausingHooks {
|
||||
retained: Some(retained),
|
||||
pause: Some(pause),
|
||||
holder,
|
||||
},
|
||||
PyDict::new(py).unbind(),
|
||||
call_options(true),
|
||||
),
|
||||
true,
|
||||
)
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn every_suspension_is_awaited_inline_on_the_callers_task(
|
||||
#[from(interpreter)] _python: Interpreter,
|
||||
) {
|
||||
Python::attach(|py| {
|
||||
let module = install_lifecycle_module(py);
|
||||
let locals = PyDict::new(py);
|
||||
locals
|
||||
.set_item("drive", module.getattr("drive").unwrap())
|
||||
.unwrap();
|
||||
locals
|
||||
.set_item(
|
||||
"await_execution",
|
||||
wrap_pyfunction!(await_execution, py).unwrap(),
|
||||
)
|
||||
.unwrap();
|
||||
locals
|
||||
.set_item(
|
||||
"calling_execution",
|
||||
wrap_pyfunction!(calling_execution, py).unwrap(),
|
||||
)
|
||||
.unwrap();
|
||||
let probe = std::ffi::CString::new(include_str!("lifecycle.py")).unwrap();
|
||||
py.run(&probe, Some(&locals), Some(&locals)).unwrap();
|
||||
});
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn retained_exception_frames_are_collectable(#[from(interpreter)] _python: Interpreter) {
|
||||
Python::attach(|py| {
|
||||
let locals = PyDict::new(py);
|
||||
locals
|
||||
.set_item(
|
||||
"error_execution",
|
||||
wrap_pyfunction!(error_execution, py).unwrap(),
|
||||
)
|
||||
.unwrap();
|
||||
py.run(
|
||||
pyo3::ffi::c_str!(
|
||||
r#"
|
||||
import gc
|
||||
import weakref
|
||||
|
||||
class Retained:
|
||||
pass
|
||||
|
||||
def cycle():
|
||||
retained = Retained()
|
||||
try:
|
||||
raise ValueError('retained traceback')
|
||||
except ValueError as error:
|
||||
retained.owner = error_execution(error)
|
||||
return weakref.ref(retained)
|
||||
|
||||
reference = cycle()
|
||||
gc.collect()
|
||||
assert reference() is None
|
||||
"#
|
||||
),
|
||||
Some(&locals),
|
||||
Some(&locals),
|
||||
)
|
||||
.unwrap();
|
||||
});
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn cycles_through_a_retained_execution_are_collectable(#[from(interpreter)] _python: Interpreter) {
|
||||
Python::attach(|py| {
|
||||
let locals = PyDict::new(py);
|
||||
locals
|
||||
.set_item(
|
||||
"retaining_coroutine",
|
||||
wrap_pyfunction!(retaining_coroutine, py).unwrap(),
|
||||
)
|
||||
.unwrap();
|
||||
py.run(
|
||||
pyo3::ffi::c_str!(
|
||||
r#"
|
||||
import gc
|
||||
import weakref
|
||||
|
||||
class Retained:
|
||||
pass
|
||||
|
||||
def cycle():
|
||||
retained = Retained()
|
||||
coroutine = retaining_coroutine(retained)
|
||||
retained.coroutine = coroutine
|
||||
return weakref.ref(retained)
|
||||
|
||||
retained_ref = cycle()
|
||||
gc.collect()
|
||||
assert retained_ref() is None
|
||||
"#
|
||||
),
|
||||
Some(&locals),
|
||||
Some(&locals),
|
||||
)
|
||||
.unwrap();
|
||||
});
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::held_by_the_hooks(Holder::Hooks)]
|
||||
#[case::selected_as_the_response(Holder::Response)]
|
||||
#[case::selected_as_the_exception(Holder::Exception)]
|
||||
fn cycles_through_a_call_suspended_in_a_hook_are_collectable(
|
||||
#[from(interpreter)] _python: Interpreter,
|
||||
#[case] holder: Holder,
|
||||
) {
|
||||
Python::attach(|py| {
|
||||
let locals = PyDict::new(py);
|
||||
let call = PyCFunction::new_closure(py, None, None, move |args, _| {
|
||||
let (retained, pause) = args.extract::<(Py<PyAny>, Py<PyAny>)>()?;
|
||||
paused_call(args.py(), retained, pause, holder)
|
||||
})
|
||||
.unwrap();
|
||||
locals.set_item("paused_call", call).unwrap();
|
||||
py.run(
|
||||
pyo3::ffi::c_str!(
|
||||
r#"
|
||||
import gc
|
||||
import weakref
|
||||
|
||||
class Retained(Exception):
|
||||
pass
|
||||
|
||||
class Pause:
|
||||
def __await__(self):
|
||||
yield
|
||||
|
||||
def cycle():
|
||||
retained = Retained()
|
||||
coroutine = paused_call(retained, Pause())
|
||||
coroutine.send(None)
|
||||
retained.coroutine = coroutine
|
||||
return weakref.ref(retained)
|
||||
|
||||
reference = cycle()
|
||||
gc.collect()
|
||||
assert reference() is None
|
||||
"#
|
||||
),
|
||||
Some(&locals),
|
||||
Some(&locals),
|
||||
)
|
||||
.unwrap();
|
||||
});
|
||||
}
|
||||
|
|
@ -504,7 +504,9 @@ impl litellm_host_python::ExecutionBody for Preparing {
|
|||
}
|
||||
HookStep::Await(awaitable, resume) => {
|
||||
self.resume = Some(resume);
|
||||
Ok(litellm_host_python::ExecutionStep::Await(awaitable))
|
||||
Ok(litellm_host_python::ExecutionStep::Suspended(
|
||||
litellm_host_python::Suspension::Await(awaitable),
|
||||
))
|
||||
}
|
||||
}
|
||||
})
|
||||
|
|
@ -518,9 +520,9 @@ impl litellm_host_python::ExecutionBody for Preparing {
|
|||
|
||||
fn lifecycle(py: Python<'_>) -> PyResult<Bound<'_, PyModule>> {
|
||||
let source =
|
||||
std::ffi::CString::new(include_str!("../../../../litellm/rust_bridge/lifecycle.py"))
|
||||
std::ffi::CString::new(include_str!("../../../../litellm/rust_bridge/execution.py"))
|
||||
.unwrap();
|
||||
PyModule::from_code(py, &source, c"lifecycle.py", c"hook_chain_test_lifecycle")
|
||||
PyModule::from_code(py, &source, c"execution.py", c"hook_chain_test_lifecycle")
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
|
|
@ -556,9 +558,16 @@ second = Hooks('b', invoke)
|
|||
},
|
||||
lifecycle,
|
||||
);
|
||||
locals
|
||||
.set_item("call", execution.into_coroutine(py).unwrap())
|
||||
let module = lifecycle(py).unwrap();
|
||||
let call = module
|
||||
.getattr("drive")
|
||||
.unwrap()
|
||||
.call1((
|
||||
Py::new(py, execution).unwrap(),
|
||||
module.getattr("AsyncStream").unwrap(),
|
||||
))
|
||||
.unwrap();
|
||||
locals.set_item("call", call).unwrap();
|
||||
py.run(
|
||||
c"
|
||||
async def exercise():
|
||||
|
|
|
|||
|
|
@ -36,11 +36,17 @@ async def exercise():
|
|||
releaser = asyncio.create_task(release())
|
||||
execution = await_execution(CustomAwaitable())
|
||||
try:
|
||||
execution.resume_value(None)
|
||||
execution.throw(ValueError("early"))
|
||||
except RuntimeError:
|
||||
pass
|
||||
else:
|
||||
raise AssertionError("resumed an unstarted execution")
|
||||
try:
|
||||
execution.send(1)
|
||||
except TypeError:
|
||||
pass
|
||||
else:
|
||||
raise AssertionError("sent a value into an unstarted execution")
|
||||
wrapped = drive(execution)
|
||||
try:
|
||||
wrapped.send(1)
|
||||
|
|
@ -129,7 +135,7 @@ async def exercise():
|
|||
|
||||
def reenter():
|
||||
try:
|
||||
active.start()
|
||||
active.send(None)
|
||||
except RuntimeError as error:
|
||||
observed.append(str(error))
|
||||
return result
|
||||
|
|
@ -151,7 +157,7 @@ async def exercise():
|
|||
execution = calling_execution(callback)
|
||||
callback.owner = execution
|
||||
if started:
|
||||
assert execution.start().value is result
|
||||
assert execution.send(None).value is result
|
||||
return weakref.ref(callback)
|
||||
|
||||
for started in (False, True):
|
||||
|
|
|
|||
684
litellm-rust/crates/host-python/tests/support/mod.rs
Normal file
684
litellm-rust/crates/host-python/tests/support/mod.rs
Normal file
|
|
@ -0,0 +1,684 @@
|
|||
#![allow(dead_code)] // each test binary compiles this module on its own and uses a different subset
|
||||
|
||||
use std::sync::{Arc, Mutex, MutexGuard, PoisonError};
|
||||
|
||||
use litellm_host::{
|
||||
call::{CallOutput, HostedMachine, hosted_call},
|
||||
hooks::CallHooks,
|
||||
interceptors::{Interceptors, RawResponse, RequestContext, WireRequest},
|
||||
lifecycle::{ExecutionEvent, Timing},
|
||||
machine::{CallMachine, MachineFault},
|
||||
protocol::{Protocol, Reply},
|
||||
};
|
||||
use litellm_host_python::{
|
||||
CallOptions, HookStep, InvokeError, PythonBinding, PythonCallEvent, PythonCallHooks,
|
||||
PythonHostCalls, PythonOwned, PythonRuntime, run_call,
|
||||
};
|
||||
use pyo3::{
|
||||
exceptions::{PyTypeError, PyValueError, asyncio::CancelledError},
|
||||
gc::{PyTraverseError, PyVisit},
|
||||
prelude::*,
|
||||
types::{PyDict, PyString},
|
||||
};
|
||||
use rstest::fixture;
|
||||
|
||||
static PYTHON_GLOBALS: Mutex<()> = Mutex::new(());
|
||||
|
||||
/// One initialized interpreter with the real lifecycle protocol installed, held for the
|
||||
/// length of a test so tests sharing the interpreter run one at a time.
|
||||
pub struct Interpreter(MutexGuard<'static, ()>);
|
||||
|
||||
#[fixture]
|
||||
pub fn interpreter() -> Interpreter {
|
||||
let guard = PYTHON_GLOBALS
|
||||
.lock()
|
||||
.unwrap_or_else(PoisonError::into_inner);
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
py.import("asyncio").unwrap();
|
||||
install_lifecycle_module(py);
|
||||
});
|
||||
Interpreter(guard)
|
||||
}
|
||||
|
||||
pub fn lifecycle_module(py: Python<'_>) -> PyResult<Bound<'_, PyModule>> {
|
||||
py.import("host_test_lifecycle")
|
||||
}
|
||||
|
||||
/// Hands a call's execution to its caller the way the product entrypoints do: an async call
|
||||
/// gets the coroutine from `drive`, a sync call the value or stream `run` settles to.
|
||||
pub fn handed(
|
||||
py: Python<'_>,
|
||||
execution: PyResult<Py<PyAny>>,
|
||||
asynchronous: bool,
|
||||
) -> PyResult<Py<PyAny>> {
|
||||
let execution = execution?;
|
||||
let module = lifecycle_module(py)?;
|
||||
let handed = if asynchronous {
|
||||
module.getattr("drive")?.call1((execution,))?
|
||||
} else {
|
||||
module
|
||||
.getattr("run")?
|
||||
.call1((execution, module.getattr("SyncStream")?))?
|
||||
};
|
||||
Ok(handed.unbind())
|
||||
}
|
||||
|
||||
pub fn call_options(asynchronous: bool) -> CallOptions {
|
||||
CallOptions::new(asynchronous, lifecycle_module)
|
||||
}
|
||||
|
||||
pub fn install_lifecycle_module(py: Python<'_>) -> Bound<'_, PyModule> {
|
||||
let source = std::ffi::CString::new(concat!(
|
||||
include_str!("../../../../../litellm/rust_bridge/execution.py"),
|
||||
"\n_drive = drive\n\n\nasync def drive(execution):\n return await _drive(execution, AsyncStream)\n",
|
||||
))
|
||||
.unwrap();
|
||||
PyModule::from_code(
|
||||
py,
|
||||
&source,
|
||||
pyo3::ffi::c_str!("execution.py"),
|
||||
pyo3::ffi::c_str!("host_test_lifecycle"),
|
||||
)
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct Error(pub String);
|
||||
|
||||
impl std::fmt::Display for Error {
|
||||
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
formatter.write_str(&self.0)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<MachineFault> for Error {
|
||||
fn from(fault: MachineFault) -> Self {
|
||||
Self(format!("{fault:?}"))
|
||||
}
|
||||
}
|
||||
|
||||
pub struct Synthetic;
|
||||
|
||||
impl Protocol for Synthetic {
|
||||
type Response = String;
|
||||
type Error = Error;
|
||||
type Request = String;
|
||||
type HostCall = (&'static str, Reply<String>);
|
||||
type Chunk = std::convert::Infallible;
|
||||
type StreamHead = std::convert::Infallible;
|
||||
}
|
||||
|
||||
pub fn wire() -> WireRequest {
|
||||
WireRequest {
|
||||
url: "https://example.invalid".into(),
|
||||
headers: Vec::new(),
|
||||
body: serde_json::json!({}),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn context() -> RequestContext {
|
||||
RequestContext {
|
||||
model: "model".into(),
|
||||
custom_llm_provider: "provider".into(),
|
||||
optional_params: serde_json::json!({}),
|
||||
secret_fields: Vec::new(),
|
||||
api_key: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Default)]
|
||||
pub struct Log(Arc<Mutex<Vec<String>>>);
|
||||
|
||||
impl Log {
|
||||
pub fn push(&self, entry: impl Into<String>) {
|
||||
self.0.lock().unwrap().push(entry.into());
|
||||
}
|
||||
|
||||
pub fn entries(&self) -> Vec<String> {
|
||||
self.0.lock().unwrap().clone()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
pub enum OpScript {
|
||||
Answer,
|
||||
RaisePython,
|
||||
RejectNatively,
|
||||
RejectRequestNatively,
|
||||
RaiseRequestPython,
|
||||
CancelRequest,
|
||||
AwaitAnswer,
|
||||
AwaitFailure,
|
||||
}
|
||||
|
||||
/// The fake route's public exception, kept as a value so a test sees what `classify`
|
||||
/// produced before the driver raises it.
|
||||
#[derive(Debug, PartialEq, Eq)]
|
||||
pub struct Classified(pub String);
|
||||
|
||||
impl From<Classified> for PyErr {
|
||||
fn from(classified: Classified) -> Self {
|
||||
PyValueError::new_err(format!("classified: {}", classified.0))
|
||||
}
|
||||
}
|
||||
|
||||
pub struct SyntheticBinding {
|
||||
log: Log,
|
||||
op: OpScript,
|
||||
classifier_fails: bool,
|
||||
pending_reply: Option<Reply<String>>,
|
||||
}
|
||||
|
||||
pub fn binding(log: &Log, op: OpScript) -> SyntheticBinding {
|
||||
SyntheticBinding {
|
||||
log: log.clone(),
|
||||
op,
|
||||
classifier_fails: false,
|
||||
pending_reply: None,
|
||||
}
|
||||
}
|
||||
|
||||
impl SyntheticBinding {
|
||||
pub fn with_failing_classifier(self) -> Self {
|
||||
Self {
|
||||
classifier_fails: true,
|
||||
..self
|
||||
}
|
||||
}
|
||||
|
||||
fn answer(&self, value: impl FnOnce() -> String) -> Result<String, InvokeError<Error>> {
|
||||
match self.op {
|
||||
OpScript::Answer | OpScript::AwaitAnswer | OpScript::AwaitFailure => Ok(value()),
|
||||
OpScript::RaisePython | OpScript::RaiseRequestPython => {
|
||||
Err(PyValueError::new_err("op failed").into())
|
||||
}
|
||||
OpScript::RejectNatively | OpScript::RejectRequestNatively => {
|
||||
Err(InvokeError::Native(Error("op rejected".into())))
|
||||
}
|
||||
OpScript::CancelRequest => Err(CancelledError::new_err(()).into()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl PythonBinding for SyntheticBinding {
|
||||
type Protocol = Synthetic;
|
||||
type Failure = Classified;
|
||||
|
||||
fn decode_request(
|
||||
&mut self,
|
||||
_: Python<'_>,
|
||||
arguments: &Bound<'_, PyDict>,
|
||||
) -> Result<String, InvokeError<Error>> {
|
||||
self.log.push("project");
|
||||
match self.op {
|
||||
OpScript::RejectRequestNatively
|
||||
| OpScript::RaiseRequestPython
|
||||
| OpScript::CancelRequest => self.answer(String::new),
|
||||
_ => Ok(format!("project:{}", arguments.len())),
|
||||
}
|
||||
}
|
||||
|
||||
fn encode_stream_head(
|
||||
&mut self,
|
||||
_: Python<'_>,
|
||||
head: std::convert::Infallible,
|
||||
) -> PyResult<Py<PyAny>> {
|
||||
match head {}
|
||||
}
|
||||
|
||||
fn encode_chunk(
|
||||
&mut self,
|
||||
_: Python<'_>,
|
||||
chunk: std::convert::Infallible,
|
||||
) -> PyResult<Py<PyAny>> {
|
||||
match chunk {}
|
||||
}
|
||||
|
||||
fn encode_response(&mut self, py: Python<'_>, response: String) -> PyResult<Py<PyAny>> {
|
||||
self.log.push("complete");
|
||||
Ok(PyString::new(py, &response).into_any().unbind())
|
||||
}
|
||||
|
||||
fn map_error(&self, _: Python<'_>, error: Error) -> PyResult<Classified> {
|
||||
self.log.push(format!("classify:{error}"));
|
||||
if self.classifier_fails {
|
||||
return Err(PyTypeError::new_err("classifier failed"));
|
||||
}
|
||||
Ok(Classified(error.0))
|
||||
}
|
||||
|
||||
fn host_error(error: &PyErr) -> Error {
|
||||
Error(error.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
impl PythonHostCalls<Synthetic> for SyntheticBinding {
|
||||
fn handle_host_call(
|
||||
&mut self,
|
||||
_: Python<'_>,
|
||||
(op, reply): (&'static str, Reply<String>),
|
||||
) -> Result<(), InvokeError<Error>> {
|
||||
self.log.push(format!("op:{op}"));
|
||||
self.answer(|| op.to_string())
|
||||
.map(|answer| reply.send(answer))
|
||||
}
|
||||
|
||||
fn begin_host_call(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
(op, reply): (&'static str, Reply<String>),
|
||||
) -> Result<Option<Py<PyAny>>, InvokeError<Error>> {
|
||||
if !matches!(self.op, OpScript::AwaitAnswer | OpScript::AwaitFailure) {
|
||||
return self.handle_host_call(py, (op, reply)).map(|()| None);
|
||||
}
|
||||
self.pending_reply = Some(reply);
|
||||
let module = PyModule::from_code(
|
||||
py,
|
||||
pyo3::ffi::c_str!(
|
||||
"async def answer(fail):\n if fail:\n raise LookupError('async host failed')\n return 'awaited'\n"
|
||||
),
|
||||
pyo3::ffi::c_str!("host_op.py"),
|
||||
pyo3::ffi::c_str!("host_op"),
|
||||
)?;
|
||||
Ok(Some(
|
||||
module
|
||||
.getattr("answer")?
|
||||
.call1((matches!(self.op, OpScript::AwaitFailure),))?
|
||||
.unbind(),
|
||||
))
|
||||
}
|
||||
|
||||
fn resume_host_call(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
result: PyResult<Py<PyAny>>,
|
||||
) -> Result<Option<Py<PyAny>>, InvokeError<Error>> {
|
||||
let answer = result?.extract::<String>(py)?;
|
||||
self.pending_reply.take().unwrap().send(answer);
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
|
||||
impl PythonOwned for SyntheticBinding {
|
||||
fn close(&mut self, _: Python<'_>) {
|
||||
self.log.push("host.close");
|
||||
}
|
||||
|
||||
fn traverse(&self, _: &PyVisit<'_>) -> Result<(), PyTraverseError> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
pub enum HookScript {
|
||||
Plain,
|
||||
RewriteArguments,
|
||||
ObserveArguments,
|
||||
FailBegin,
|
||||
ReplaceResponse,
|
||||
FailAfterSuccess,
|
||||
CancelTerminal,
|
||||
FailTerminal,
|
||||
}
|
||||
|
||||
pub struct SyntheticHooks {
|
||||
log: Log,
|
||||
script: HookScript,
|
||||
}
|
||||
|
||||
pub fn hooks(log: &Log, script: HookScript) -> SyntheticHooks {
|
||||
SyntheticHooks {
|
||||
log: log.clone(),
|
||||
script,
|
||||
}
|
||||
}
|
||||
|
||||
impl CallHooks<PythonRuntime> for SyntheticHooks {
|
||||
fn prepare_arguments(
|
||||
&mut self,
|
||||
_: Python<'_>,
|
||||
arguments: Py<PyDict>,
|
||||
_: f64,
|
||||
) -> PyResult<HookStep<Self, Py<PyDict>>> {
|
||||
self.log.push("begin");
|
||||
if matches!(self.script, HookScript::FailBegin) {
|
||||
return Err(PyValueError::new_err("begin failed"));
|
||||
}
|
||||
if matches!(self.script, HookScript::RewriteArguments) {
|
||||
let prepared = Python::attach(|py| -> PyResult<Py<PyDict>> {
|
||||
let prepared = arguments.bind(py).copy()?;
|
||||
prepared.set_item("prepared", "hook")?;
|
||||
prepared.set_item("api_key", "hook-key")?;
|
||||
Ok(prepared.unbind())
|
||||
})?;
|
||||
return Ok(HookStep::Ready(prepared));
|
||||
}
|
||||
Ok(HookStep::Ready(arguments))
|
||||
}
|
||||
|
||||
fn arguments_prepared(&mut self, py: Python<'_>, arguments: &Py<PyDict>) -> PyResult<()> {
|
||||
if matches!(self.script, HookScript::ObserveArguments) {
|
||||
let value: String = arguments
|
||||
.bind(py)
|
||||
.get_item("prepared")?
|
||||
.unwrap()
|
||||
.extract()?;
|
||||
self.log.push(format!("adopted:{value}"));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn before_provider_request(
|
||||
&mut self,
|
||||
_: Python<'_>,
|
||||
wire: Box<WireRequest>,
|
||||
_: &RequestContext,
|
||||
) -> PyResult<HookStep<Self, Box<WireRequest>>> {
|
||||
self.log.push("before_provider_request");
|
||||
Ok(HookStep::Ready(Box::new(WireRequest {
|
||||
url: "rewritten".into(),
|
||||
..*wire
|
||||
})))
|
||||
}
|
||||
|
||||
fn transform_response(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
response: Py<PyAny>,
|
||||
_: Timing,
|
||||
) -> PyResult<HookStep<Self, Py<PyAny>>> {
|
||||
self.log.push("after_success");
|
||||
match self.script {
|
||||
HookScript::ReplaceResponse => Ok(HookStep::Ready(
|
||||
"replaced".into_pyobject(py)?.into_any().unbind(),
|
||||
)),
|
||||
HookScript::FailAfterSuccess => Err(PyValueError::new_err("after_success failed")),
|
||||
_ => Ok(HookStep::Ready(response)),
|
||||
}
|
||||
}
|
||||
|
||||
fn on_event(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
event: PythonCallEvent<'_>,
|
||||
) -> PyResult<HookStep<Self, ()>> {
|
||||
let terminal = matches!(
|
||||
event,
|
||||
PythonCallEvent::Succeeded { .. } | PythonCallEvent::Failed { .. }
|
||||
);
|
||||
if terminal && matches!(self.script, HookScript::CancelTerminal) {
|
||||
return Err(CancelledError::new_err("callback cancelled"));
|
||||
}
|
||||
if terminal && matches!(self.script, HookScript::FailTerminal) {
|
||||
return Err(PyValueError::new_err("callback failed"));
|
||||
}
|
||||
self.log.push(match event {
|
||||
PythonCallEvent::Execution(ExecutionEvent::ResultReady { .. }) => "cache_hit".into(),
|
||||
PythonCallEvent::Started { .. } => "started".into(),
|
||||
PythonCallEvent::Cancelled { .. } => "cancelled".into(),
|
||||
PythonCallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }) => {
|
||||
format!("response:{}", raw.body)
|
||||
}
|
||||
PythonCallEvent::Succeeded { response, .. } => {
|
||||
format!("succeeded:{}", response.bind(py))
|
||||
}
|
||||
PythonCallEvent::Failed { origin, error, .. } => {
|
||||
format!("failed:{origin:?}:{}", error.value(py))
|
||||
}
|
||||
});
|
||||
Ok(HookStep::Ready(()))
|
||||
}
|
||||
|
||||
fn on_stream_open(&mut self, _: Python<'_>, _: &Py<PyAny>) -> PyResult<()> {
|
||||
self.log.push("opened");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn on_stream_chunk(&mut self, _: Python<'_>, _: &Py<PyAny>) -> PyResult<()> {
|
||||
self.log.push("delivered");
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl PythonOwned for SyntheticHooks {
|
||||
fn close(&mut self, _: Python<'_>) {
|
||||
self.log.push("adapter.close");
|
||||
}
|
||||
|
||||
fn traverse(&self, _: &PyVisit<'_>) -> Result<(), PyTraverseError> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
pub type MachineFactory = Box<dyn FnOnce(String) -> CallMachine<Synthetic> + Send + Sync>;
|
||||
|
||||
/// Answers to projection, to the route op and to `before_provider_request` all reach the
|
||||
/// response, so a driver that misroutes a reply changes what the call returns.
|
||||
pub fn success_machine() -> MachineFactory {
|
||||
Box::new(move |projected| {
|
||||
CallMachine::new(None, move |host| {
|
||||
Box::pin(async move {
|
||||
let signed = host.services.call(|reply| ("sign", reply)).await?;
|
||||
let wire = host
|
||||
.interceptors
|
||||
.before_provider_request(wire(), context())
|
||||
.await?;
|
||||
host.interceptors
|
||||
.after_provider_response(RawResponse { body: "raw".into() })
|
||||
.await?;
|
||||
Ok(format!("{projected}|{signed}|{}", wire.url))
|
||||
})
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
pub fn failing_machine() -> MachineFactory {
|
||||
Box::new(move |_| {
|
||||
CallMachine::new(None, |_| {
|
||||
Box::pin(async move { Err(Error("provider exploded".into())) })
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
/// Drives an asynchronous call's coroutine to its settled value from a sync test. A
|
||||
/// synchronous call is already settled.
|
||||
pub fn settled(
|
||||
py: Python<'_>,
|
||||
result: PyResult<Py<PyAny>>,
|
||||
asynchronous: bool,
|
||||
) -> PyResult<Py<PyAny>> {
|
||||
if !asynchronous {
|
||||
return result;
|
||||
}
|
||||
let coroutine = result?;
|
||||
let completed = coroutine
|
||||
.call_method1(py, "send", (py.None(),))
|
||||
.unwrap_err();
|
||||
if !completed.is_instance_of::<pyo3::exceptions::PyStopIteration>(py) {
|
||||
return Err(completed);
|
||||
}
|
||||
completed.value(py).getattr("value").map(Bound::unbind)
|
||||
}
|
||||
|
||||
pub fn run_scripted(
|
||||
py: Python<'_>,
|
||||
machine: MachineFactory,
|
||||
op: OpScript,
|
||||
script: HookScript,
|
||||
asynchronous: bool,
|
||||
) -> (PyResult<Py<PyAny>>, Vec<String>) {
|
||||
let log = Log::default();
|
||||
run_composed(
|
||||
py,
|
||||
machine,
|
||||
binding(&log, op),
|
||||
hooks(&log, script),
|
||||
std::convert::identity,
|
||||
call_options(asynchronous),
|
||||
)
|
||||
}
|
||||
|
||||
pub fn run_composed<L: PythonCallHooks + 'static>(
|
||||
py: Python<'_>,
|
||||
machine: MachineFactory,
|
||||
binding: SyntheticBinding,
|
||||
hooks: SyntheticHooks,
|
||||
compose: impl FnOnce(SyntheticHooks) -> L,
|
||||
options: CallOptions,
|
||||
) -> (PyResult<Py<PyAny>>, Vec<String>) {
|
||||
let asynchronous = options.asynchronous;
|
||||
let log = binding.log.clone();
|
||||
let arguments = PyDict::new(py);
|
||||
arguments.set_item("model", "m").unwrap();
|
||||
let result = handed(
|
||||
py,
|
||||
run_call(
|
||||
py,
|
||||
move |_, _, request| Ok(machine(request)),
|
||||
binding,
|
||||
compose(hooks),
|
||||
arguments.unbind(),
|
||||
options,
|
||||
),
|
||||
asynchronous,
|
||||
);
|
||||
(settled(py, result, asynchronous), log.entries())
|
||||
}
|
||||
|
||||
pub struct Streaming;
|
||||
|
||||
impl Protocol for Streaming {
|
||||
type Response = ();
|
||||
type Error = Error;
|
||||
type Request = ();
|
||||
type HostCall = std::convert::Infallible;
|
||||
type Chunk = &'static str;
|
||||
type StreamHead = Vec<(&'static str, &'static str)>;
|
||||
}
|
||||
|
||||
pub struct StreamingBinding;
|
||||
|
||||
impl PythonBinding for StreamingBinding {
|
||||
type Protocol = Streaming;
|
||||
type Failure = Classified;
|
||||
|
||||
fn decode_request(
|
||||
&mut self,
|
||||
_: Python<'_>,
|
||||
_: &Bound<'_, PyDict>,
|
||||
) -> Result<(), InvokeError<Error>> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn encode_stream_head(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
head: Vec<(&'static str, &'static str)>,
|
||||
) -> PyResult<Py<PyAny>> {
|
||||
let headers = PyDict::new(py);
|
||||
for (name, value) in head {
|
||||
headers.set_item(name, value)?;
|
||||
}
|
||||
let hidden = PyDict::new(py);
|
||||
hidden.set_item("additional_headers", headers)?;
|
||||
Ok(hidden.into_any().unbind())
|
||||
}
|
||||
|
||||
fn encode_chunk(&mut self, py: Python<'_>, chunk: &'static str) -> PyResult<Py<PyAny>> {
|
||||
Ok(PyString::new(py, chunk).into_any().unbind())
|
||||
}
|
||||
|
||||
fn encode_response(&mut self, py: Python<'_>, (): ()) -> PyResult<Py<PyAny>> {
|
||||
Ok(py.None())
|
||||
}
|
||||
|
||||
fn map_error(&self, _: Python<'_>, error: Error) -> PyResult<Classified> {
|
||||
Ok(Classified(error.0))
|
||||
}
|
||||
|
||||
fn host_error(error: &PyErr) -> Error {
|
||||
Error(error.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
impl PythonHostCalls<Streaming> for StreamingBinding {
|
||||
fn handle_host_call(
|
||||
&mut self,
|
||||
_: Python<'_>,
|
||||
op: std::convert::Infallible,
|
||||
) -> Result<(), InvokeError<Error>> {
|
||||
match op {}
|
||||
}
|
||||
}
|
||||
|
||||
impl PythonOwned for StreamingBinding {
|
||||
fn close(&mut self, _: Python<'_>) {}
|
||||
|
||||
fn traverse(&self, _: &PyVisit<'_>) -> Result<(), PyTraverseError> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
pub fn streaming_machine() -> impl FnOnce(()) -> HostedMachine<Streaming> + Send + Sync {
|
||||
move |()| {
|
||||
hosted_call((), None, |(), _, _, _observations| async {
|
||||
Ok(CallOutput::Stream {
|
||||
head: vec![("request-id", "req_1")],
|
||||
chunks: Box::pin(futures_util::stream::iter([Ok("first"), Ok("second")])),
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Starts a streaming call and returns the caller's stream object.
|
||||
pub fn open_stream(
|
||||
py: Python<'_>,
|
||||
hooks: SyntheticHooks,
|
||||
options: CallOptions,
|
||||
) -> Bound<'_, PyAny> {
|
||||
let asynchronous = options.asynchronous;
|
||||
let handed = handed(
|
||||
py,
|
||||
run_call(
|
||||
py,
|
||||
|_, _, request| Ok(streaming_machine()(request)),
|
||||
StreamingBinding,
|
||||
hooks,
|
||||
PyDict::new(py).unbind(),
|
||||
options,
|
||||
),
|
||||
asynchronous,
|
||||
)
|
||||
.unwrap();
|
||||
if asynchronous {
|
||||
let stop = handed.call_method1(py, "send", (py.None(),)).unwrap_err();
|
||||
stop.value(py).getattr("value").unwrap()
|
||||
} else {
|
||||
handed.into_bound(py)
|
||||
}
|
||||
}
|
||||
|
||||
/// Drives an `AsyncStream` or `SyncStream` to completion from a sync test.
|
||||
pub fn read_all(py: Python<'_>, stream: &Bound<'_, PyAny>, asynchronous: bool) -> Vec<String> {
|
||||
if !asynchronous {
|
||||
return stream
|
||||
.try_iter()
|
||||
.unwrap()
|
||||
.map(|chunk| chunk.unwrap().extract().unwrap())
|
||||
.collect();
|
||||
}
|
||||
std::iter::from_fn(|| {
|
||||
let stop = stream
|
||||
.call_method0("__anext__")
|
||||
.unwrap()
|
||||
.call_method1("send", (py.None(),))
|
||||
.unwrap_err();
|
||||
if stop.is_instance_of::<pyo3::exceptions::PyStopAsyncIteration>(py) {
|
||||
return None;
|
||||
}
|
||||
assert!(stop.is_instance_of::<pyo3::exceptions::PyStopIteration>(py));
|
||||
Some(stop.value(py).getattr("value").unwrap().extract().unwrap())
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
|
@ -42,10 +42,32 @@ pub enum ResultSource {
|
|||
Cache { key: String },
|
||||
}
|
||||
|
||||
/// What the call cost, settled once by the route. A host records it; it never re-derives it.
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub enum Cost {
|
||||
/// The amount the provider itself billed, taken from its response.
|
||||
Reported { amount: serde_json::Number },
|
||||
/// Served from the cache, so nothing was billed.
|
||||
CacheHit,
|
||||
/// No provider figure: the host estimates from usage and its price table.
|
||||
Deferred,
|
||||
}
|
||||
|
||||
impl Cost {
|
||||
pub fn settle(source: &ResultSource, reported: Option<serde_json::Number>) -> Self {
|
||||
match (source, reported) {
|
||||
(ResultSource::Cache { .. }, _) => Self::CacheHit,
|
||||
(ResultSource::Provider, Some(amount)) => Self::Reported { amount },
|
||||
(ResultSource::Provider, None) => Self::Deferred,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct ExecutionFacts {
|
||||
pub provider: ProviderIdentity,
|
||||
pub source: ResultSource,
|
||||
pub cost: Cost,
|
||||
}
|
||||
|
||||
pub trait Interceptors<E>: Send + Sync {
|
||||
|
|
|
|||
|
|
@ -138,7 +138,9 @@ pub fn has_bearer_auth(headers: &[(String, String)]) -> bool {
|
|||
}
|
||||
let value = value.trim();
|
||||
value.len() > 7
|
||||
&& value[..7].eq_ignore_ascii_case("bearer ")
|
||||
&& value
|
||||
.get(..7)
|
||||
.is_some_and(|prefix| prefix.eq_ignore_ascii_case("bearer "))
|
||||
&& !value[7..].trim().is_empty()
|
||||
})
|
||||
}
|
||||
|
|
@ -279,31 +281,19 @@ mod tests {
|
|||
assert!(!has_header(&headers, "authorization"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bearer_detection_requires_a_non_empty_token() {
|
||||
assert!(has_bearer_auth(&[(
|
||||
"Authorization".to_string(),
|
||||
"Bearer abc".to_string()
|
||||
)]));
|
||||
assert!(has_bearer_auth(&[(
|
||||
"authorization".to_string(),
|
||||
"bearer abc".to_string()
|
||||
)]));
|
||||
assert!(!has_bearer_auth(&[(
|
||||
"Authorization".to_string(),
|
||||
"Bearer ".to_string()
|
||||
)]));
|
||||
assert!(!has_bearer_auth(&[(
|
||||
"authorization".to_string(),
|
||||
String::new()
|
||||
)]));
|
||||
assert!(!has_bearer_auth(&[(
|
||||
"Authorization".to_string(),
|
||||
"Basic abc".to_string()
|
||||
)]));
|
||||
assert!(!has_bearer_auth(&[(
|
||||
"x-api-key".to_string(),
|
||||
"abc".to_string()
|
||||
)]));
|
||||
#[rstest::rstest]
|
||||
#[case::bearer("Authorization", "Bearer abc", true)]
|
||||
#[case::lowercase("authorization", "bearer abc", true)]
|
||||
#[case::blank_token("Authorization", "Bearer ", false)]
|
||||
#[case::empty("authorization", "", false)]
|
||||
#[case::basic("Authorization", "Basic abc", false)]
|
||||
#[case::different_header("x-api-key", "abc", false)]
|
||||
#[case::non_ascii("authorization", "invali\u{00e9}auth", false)]
|
||||
fn bearer_detection_requires_a_non_empty_token(
|
||||
#[case] name: &str,
|
||||
#[case] value: &str,
|
||||
#[case] expected: bool,
|
||||
) {
|
||||
assert_eq!(has_bearer_auth(&[(name.into(), value.into())]), expected);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -140,7 +140,7 @@ impl Default for HttpSettings {
|
|||
proxies: EnvironmentProxies::default(),
|
||||
connect_timeout: Duration::from_secs(10),
|
||||
tcp_keepalive: None,
|
||||
pool_idle_timeout: Duration::from_secs(120),
|
||||
pool_idle_timeout: Duration::from_secs(15),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -330,7 +330,7 @@ mod tests {
|
|||
&[("AIOHTTP_TCP_KEEPIDLE", "45")],
|
||||
))]);
|
||||
assert_eq!(settings.tcp_keepalive, None);
|
||||
assert_eq!(settings.pool_idle_timeout, Duration::from_secs(120));
|
||||
assert_eq!(settings.pool_idle_timeout, Duration::from_secs(15));
|
||||
}
|
||||
|
||||
fn proxy_flags(
|
||||
|
|
|
|||
|
|
@ -26,6 +26,7 @@ serde_json = { workspace = true, features = ["preserve_order"] }
|
|||
tracing.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
litellm-framer = { workspace = true, features = ["sse"] }
|
||||
litellm-cache-memory.workspace = true
|
||||
litellm-host-native.workspace = true
|
||||
litellm-http = { workspace = true, features = ["test-support"] }
|
||||
|
|
|
|||
|
|
@ -1,14 +1,21 @@
|
|||
use litellm_core_utils::get_llm_provider_logic::LlmProviders;
|
||||
use litellm_http::request::string_headers as shared_string_headers;
|
||||
pub(super) use litellm_http::request::truncate_error_body;
|
||||
use litellm_llms::tencent::messages::transformation::TENCENT_MESSAGES_CONFIG;
|
||||
use litellm_llms::{
|
||||
anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG,
|
||||
azure_ai::messages::transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG,
|
||||
base_llm::messages::transformation::BaseMessagesConfig,
|
||||
bedrock::messages::invoke_transformations::anthropic_claude3_transformation::BEDROCK_ANTHROPIC_MESSAGES_CONFIG,
|
||||
deepseek::messages::transformation::DEEPSEEK_ANTHROPIC_MESSAGES_CONFIG,
|
||||
edenai::messages::transformation::EDENAI_MESSAGES_CONFIG,
|
||||
minimax::messages::transformation::MINIMAX_MESSAGES_CONFIG,
|
||||
vertex_ai::messages::transformation::VERTEX_ANTHROPIC_MESSAGES_CONFIG,
|
||||
};
|
||||
use litellm_llms::{
|
||||
bedrock::claude_platform::messages::transformation::BEDROCK_CLAUDE_PLATFORM_MESSAGES_CONFIG,
|
||||
bedrock_mantle::messages::transformation::BEDROCK_MANTLE_MESSAGES_CONFIG,
|
||||
};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use super::Error;
|
||||
|
|
@ -34,14 +41,25 @@ impl MessagesProvider {
|
|||
/// Python's `get_provider_anthropic_messages_config`: Vertex AI serves only its Claude
|
||||
/// partner models on this route.
|
||||
pub(crate) fn messages_provider(provider: LlmProviders, model: &str) -> Option<MessagesProvider> {
|
||||
use litellm_llms::openrouter::messages::transformation::OPENROUTER_MESSAGES_CONFIG;
|
||||
|
||||
let config: &'static dyn BaseMessagesConfig = match provider {
|
||||
LlmProviders::Tencent => &TENCENT_MESSAGES_CONFIG,
|
||||
LlmProviders::Anthropic => &ANTHROPIC_MESSAGES_CONFIG,
|
||||
LlmProviders::AzureAi => &AZURE_ANTHROPIC_MESSAGES_CONFIG,
|
||||
LlmProviders::Bedrock if model.starts_with("claude_platform/") => {
|
||||
&BEDROCK_CLAUDE_PLATFORM_MESSAGES_CONFIG
|
||||
}
|
||||
LlmProviders::Bedrock if model.starts_with("mantle/") => &BEDROCK_MANTLE_MESSAGES_CONFIG,
|
||||
LlmProviders::Bedrock => &BEDROCK_ANTHROPIC_MESSAGES_CONFIG,
|
||||
LlmProviders::BedrockMantle => &BEDROCK_MANTLE_MESSAGES_CONFIG,
|
||||
LlmProviders::Deepseek => &DEEPSEEK_ANTHROPIC_MESSAGES_CONFIG,
|
||||
LlmProviders::Edenai => &EDENAI_MESSAGES_CONFIG,
|
||||
LlmProviders::Minimax => &MINIMAX_MESSAGES_CONFIG,
|
||||
LlmProviders::VertexAi if model.to_ascii_lowercase().contains("claude") => {
|
||||
&VERTEX_ANTHROPIC_MESSAGES_CONFIG
|
||||
}
|
||||
LlmProviders::Openrouter => &OPENROUTER_MESSAGES_CONFIG,
|
||||
_ => return None,
|
||||
};
|
||||
Some(MessagesProvider { provider, config })
|
||||
|
|
@ -64,10 +82,13 @@ mod tests {
|
|||
use litellm_core_utils::get_llm_provider_logic::LlmProviders;
|
||||
|
||||
#[rstest]
|
||||
#[case::tencent("tencent")]
|
||||
#[case::anthropic("anthropic")]
|
||||
#[case::azure_ai("azure_ai")]
|
||||
#[case::bedrock("bedrock")]
|
||||
#[case::deepseek("deepseek")]
|
||||
#[case::edenai("edenai")]
|
||||
#[case::minimax("minimax")]
|
||||
#[case::vertex_ai("vertex_ai")]
|
||||
fn provider_keeps_its_python_name(#[case] name: &str) {
|
||||
let provider =
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ use litellm_http::transport::Error as TransportError;
|
|||
use litellm_llms::base_llm::{
|
||||
auth::{Authenticated, resolve_auth},
|
||||
messages::{
|
||||
streaming::{ByteStream, StreamDecoder, encode_anthropic_sse},
|
||||
streaming::{self, ByteStream, StreamDecoder},
|
||||
transformation::BaseMessagesConfig,
|
||||
},
|
||||
};
|
||||
|
|
@ -25,7 +25,7 @@ use litellm_inference::{context::CallContext, outbound::outbound_request};
|
|||
pub(super) struct ProviderCall {
|
||||
pub identity: ProviderIdentity,
|
||||
pub wire: WireRequest,
|
||||
provider: super::common_utils::MessagesProvider,
|
||||
pub(super) provider: super::common_utils::MessagesProvider,
|
||||
signer: Option<litellm_auth_aws::SigV4Signer>,
|
||||
timeout: Option<Duration>,
|
||||
stream: bool,
|
||||
|
|
@ -245,7 +245,7 @@ fn decoded_chunks(
|
|||
let Some(event) = events.try_next().await? else {
|
||||
return Ok(None);
|
||||
};
|
||||
let chunk = encode_anthropic_sse(&event)?;
|
||||
let chunk = streaming::encode(&event)?;
|
||||
log_chunk(provider, "client_response", &chunk);
|
||||
Ok(Some((chunk, events)))
|
||||
})
|
||||
|
|
@ -271,7 +271,6 @@ fn log_chunk(provider: &str, stage: &str, data: &bytes::Bytes) {
|
|||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use litellm_llms::base_llm::messages::streaming::anthropic_sse_event_stream;
|
||||
use rstest::rstest;
|
||||
use wiremock::{Mock, MockServer, ResponseTemplate, matchers::any};
|
||||
|
||||
|
|
@ -300,7 +299,7 @@ mod tests {
|
|||
.await
|
||||
.unwrap();
|
||||
let MessagesCallResponse::Stream { mut chunks, .. } =
|
||||
streaming_response(response, Some(anthropic_sse_event_stream), "test")
|
||||
streaming_response(response, Some(streaming::event_stream), "test")
|
||||
else {
|
||||
panic!("a streaming response returns chunks");
|
||||
};
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ mod types;
|
|||
|
||||
use futures_util::FutureExt;
|
||||
use litellm_auth::AuthServices;
|
||||
use litellm_host::interceptors::{ExecutionFacts, Interceptors, ResultSource};
|
||||
use litellm_host::interceptors::{Cost, ExecutionFacts, Interceptors, ResultSource};
|
||||
|
||||
use litellm_inference::{caching::CallCache, context::CallContext};
|
||||
use litellm_secrets::source::SecretSource;
|
||||
|
|
@ -87,6 +87,7 @@ impl MessagesRoute {
|
|||
&request.wire,
|
||||
);
|
||||
let identity = request.identity.clone();
|
||||
let config = request.provider.config();
|
||||
let (output, source) = match cache.lookup().await {
|
||||
Some(hit) => hit,
|
||||
None => (
|
||||
|
|
@ -94,9 +95,14 @@ impl MessagesRoute {
|
|||
ResultSource::Provider,
|
||||
),
|
||||
};
|
||||
let reported = match &output {
|
||||
MessagesCallResponse::Complete(response) => config.reported_cost(response),
|
||||
MessagesCallResponse::Stream { .. } => None,
|
||||
};
|
||||
context
|
||||
.result_ready(ExecutionFacts {
|
||||
provider: identity,
|
||||
cost: Cost::settle(&source, reported),
|
||||
source: source.clone(),
|
||||
})
|
||||
.await?;
|
||||
|
|
|
|||
|
|
@ -76,6 +76,10 @@ fn prepare_provider_request(
|
|||
shaping,
|
||||
..
|
||||
} = call;
|
||||
let connection_params = litellm_router_types::LitellmParams {
|
||||
api_base: api_base.clone().or(litellm_params.api_base.clone()),
|
||||
..litellm_params
|
||||
};
|
||||
let config = provider.config();
|
||||
let env_lookup = |key: &str| secrets.get(key);
|
||||
|
||||
|
|
@ -95,25 +99,22 @@ fn prepare_provider_request(
|
|||
let forwarded = string_headers(Some(
|
||||
extra_headers.into_iter().flatten().chain(scoped).collect(),
|
||||
))?;
|
||||
let validated = config.validate_environment(
|
||||
let forwarded = config.request_headers(
|
||||
with_default_headers(forwarded, config.default_headers()),
|
||||
&transformed,
|
||||
);
|
||||
let environment = config.validate_environment(
|
||||
forwarded,
|
||||
api_key.as_deref(),
|
||||
&transformed.model,
|
||||
&litellm_params,
|
||||
&connection_params,
|
||||
&env_lookup,
|
||||
)?;
|
||||
let environment = ValidatedEnvironment {
|
||||
headers: config.request_headers(
|
||||
with_default_headers(validated.headers, config.default_headers()),
|
||||
&transformed,
|
||||
),
|
||||
auth: validated.auth,
|
||||
};
|
||||
|
||||
let url = config.get_complete_url(
|
||||
api_base.as_deref(),
|
||||
&transformed.model,
|
||||
&litellm_params,
|
||||
&connection_params,
|
||||
transformed.params.stream == Some(true),
|
||||
&env_lookup,
|
||||
)?;
|
||||
|
|
@ -475,4 +476,44 @@ mod tests {
|
|||
))
|
||||
);
|
||||
}
|
||||
#[rstest]
|
||||
#[case::mantle(
|
||||
"bedrock_mantle/claude-test",
|
||||
"https://bedrock-mantle.us-west-2.api.aws",
|
||||
"/anthropic/v1/messages",
|
||||
"bedrock"
|
||||
)]
|
||||
#[case::platform(
|
||||
"bedrock/claude_platform/claude-test",
|
||||
"https://aws-external-anthropic.us-west-2.api.aws",
|
||||
"/v1/messages",
|
||||
"aws-external-anthropic"
|
||||
)]
|
||||
fn explicit_endpoint_region_is_used_by_the_signer(
|
||||
shaping: MessagesShaping,
|
||||
#[case] model: &str,
|
||||
#[case] base: &str,
|
||||
#[case] path: &str,
|
||||
#[case] service: &str,
|
||||
) {
|
||||
let prepared = prepare_with_secrets(MessagesCall {
|
||||
body: body(json!({"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": "hi"}]})),
|
||||
api_key: None,
|
||||
api_base: Some(base.into()),
|
||||
custom_llm_provider: None,
|
||||
litellm_params: Default::default(),
|
||||
extra_headers: None,
|
||||
provider_specific_header: None,
|
||||
timeout: None,
|
||||
shaping,
|
||||
}, &|name: &str| match name {
|
||||
"AWS_REGION_NAME" => Some("us-east-1".into()),
|
||||
"ANTHROPIC_AWS_WORKSPACE_ID" => Some("workspace".into()),
|
||||
_ => None,
|
||||
}).unwrap();
|
||||
assert_eq!(prepared.url, format!("{base}{path}"));
|
||||
assert!(
|
||||
matches!(prepared.environment.auth, litellm_llms::base_llm::auth::AuthScheme::AwsSigV4 { region, service: actual, .. } if region == "us-west-2" && actual == service)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,323 @@
|
|||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use futures_util::TryStreamExt;
|
||||
use litellm_inference_messages::{Error, MessagesCall, MessagesCallResponse, MessagesShaping};
|
||||
use litellm_inference_testing::{RecordingSecrets, no_secrets};
|
||||
use litellm_llms_types::formats::messages::{MessagesRequest, MessagesResponse};
|
||||
use rstest::{fixture, rstest};
|
||||
use serde_json::{Value, json};
|
||||
use wiremock::ResponseTemplate;
|
||||
|
||||
#[path = "../support/mod.rs"]
|
||||
mod support;
|
||||
use support::*;
|
||||
|
||||
const MODEL: &str = "claude-test";
|
||||
|
||||
#[fixture]
|
||||
fn call() -> MessagesCall {
|
||||
MessagesCall {
|
||||
body: serde_json::from_value(json!({"model": MODEL, "max_tokens": 16, "messages": [{"role": "user", "content": "hi"}]})).unwrap(),
|
||||
api_key: None,
|
||||
api_base: None,
|
||||
custom_llm_provider: None,
|
||||
litellm_params: Default::default(),
|
||||
extra_headers: None,
|
||||
provider_specific_header: None,
|
||||
timeout: Some(Duration::from_secs(5)),
|
||||
shaping: MessagesShaping::default(),
|
||||
}
|
||||
}
|
||||
|
||||
fn with_fields(call: MessagesCall, fields: Value) -> MessagesCall {
|
||||
let body = serde_json::to_value(call.body).unwrap();
|
||||
MessagesCall {
|
||||
body: serde_json::from_value(Value::Object(
|
||||
body.as_object()
|
||||
.unwrap()
|
||||
.clone()
|
||||
.into_iter()
|
||||
.chain(fields.as_object().unwrap().clone())
|
||||
.collect(),
|
||||
))
|
||||
.unwrap(),
|
||||
..call
|
||||
}
|
||||
}
|
||||
|
||||
fn with_model(call: MessagesCall, model: &str) -> MessagesCall {
|
||||
MessagesCall {
|
||||
body: MessagesRequest {
|
||||
model: model.into(),
|
||||
..call.body
|
||||
},
|
||||
..call
|
||||
}
|
||||
}
|
||||
|
||||
fn message_response() -> ResponseTemplate {
|
||||
json_response(
|
||||
json!({"id": "msg_1", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "hi"}], "model": MODEL, "stop_reason": "end_turn", "usage": {"input_tokens": 1, "output_tokens": 2}}),
|
||||
)
|
||||
}
|
||||
|
||||
async fn run_with(
|
||||
secrets: Arc<RecordingSecrets>,
|
||||
call: MessagesCall,
|
||||
) -> Result<MessagesCallResponse, Error> {
|
||||
messages_route(secrets).execute(call, &(), None).await
|
||||
}
|
||||
|
||||
async fn run(call: MessagesCall) -> Result<MessagesCallResponse, Error> {
|
||||
run_with(Arc::new(RecordingSecrets::empty()), call).await
|
||||
}
|
||||
|
||||
async fn run_message(call: MessagesCall) -> MessagesResponse {
|
||||
match run(call).await.unwrap() {
|
||||
MessagesCallResponse::Complete(message) => *message,
|
||||
MessagesCallResponse::Stream { .. } => panic!("expected a completed message"),
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::platform(
|
||||
"bedrock/claude_platform/claude-test",
|
||||
"/v1/messages",
|
||||
"x-api-key",
|
||||
"key"
|
||||
)]
|
||||
#[case::mantle(
|
||||
"bedrock_mantle/anthropic.claude-test",
|
||||
"/anthropic/v1/messages",
|
||||
"authorization",
|
||||
"Bearer key"
|
||||
)]
|
||||
#[case::mantle_alias(
|
||||
"bedrock/mantle/anthropic.claude-test",
|
||||
"/anthropic/v1/messages",
|
||||
"authorization",
|
||||
"Bearer key"
|
||||
)]
|
||||
#[tokio::test]
|
||||
async fn prefixes_select_native_messages_and_strip_connection_fields(
|
||||
call: MessagesCall,
|
||||
#[case] model: &str,
|
||||
#[case] path: &str,
|
||||
#[case] header: &str,
|
||||
#[case] credential: &str,
|
||||
) {
|
||||
let upstream = upstream([message_response()]).await;
|
||||
let message = run_message(MessagesCall {
|
||||
api_key: Some("key".into()),
|
||||
api_base: Some(upstream.uri()),
|
||||
custom_llm_provider: None,
|
||||
..with_fields(
|
||||
call,
|
||||
json!({"model": model, "workspace_id": "request-workspace"}),
|
||||
)
|
||||
})
|
||||
.await;
|
||||
assert_eq!(message.model, MODEL);
|
||||
let request = only_request(&upstream).await;
|
||||
assert_eq!(request.url.path(), path);
|
||||
assert_eq!(request.header(header), Some(credential));
|
||||
assert_eq!(request.json()["model"], model.rsplit('/').next().unwrap());
|
||||
if path == "/v1/messages" {
|
||||
assert_eq!(
|
||||
request.header("anthropic-workspace-id"),
|
||||
Some("request-workspace")
|
||||
);
|
||||
assert_eq!(request.json().get("workspace_id"), None);
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::platform(
|
||||
"bedrock/claude_platform/claude-test",
|
||||
"ANTHROPIC_AWS_API_KEY",
|
||||
"ANTHROPIC_AWS_BASE_URL",
|
||||
"/v1/messages",
|
||||
"x-api-key",
|
||||
"managed"
|
||||
)]
|
||||
#[case::mantle(
|
||||
"bedrock_mantle/claude-test",
|
||||
"AWS_BEARER_TOKEN_BEDROCK",
|
||||
"BEDROCK_MANTLE_API_BASE",
|
||||
"/anthropic/v1/messages",
|
||||
"authorization",
|
||||
"Bearer managed"
|
||||
)]
|
||||
#[tokio::test]
|
||||
async fn managed_keys_bases_and_workspaces_reach_the_upstream(
|
||||
call: MessagesCall,
|
||||
#[case] model: &str,
|
||||
#[case] key_name: &str,
|
||||
#[case] base_name: &str,
|
||||
#[case] path: &str,
|
||||
#[case] header: &str,
|
||||
#[case] credential: &str,
|
||||
) {
|
||||
let upstream = upstream([message_response()]).await;
|
||||
let base = upstream.uri();
|
||||
let secrets = Arc::new(RecordingSecrets::new([
|
||||
(key_name, "managed"),
|
||||
(base_name, base.as_str()),
|
||||
("ANTHROPIC_AWS_WORKSPACE_ID", "managed-workspace"),
|
||||
]));
|
||||
let output = run_with(
|
||||
secrets,
|
||||
MessagesCall {
|
||||
custom_llm_provider: None,
|
||||
..with_model(call, model)
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(matches!(output, MessagesCallResponse::Complete(_)));
|
||||
let request = only_request(&upstream).await;
|
||||
assert_eq!(request.url.path(), path);
|
||||
assert_eq!(request.header(header), Some(credential));
|
||||
if path == "/v1/messages" {
|
||||
assert_eq!(
|
||||
request.header("anthropic-workspace-id"),
|
||||
Some("managed-workspace")
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::platform("bedrock/claude_platform/claude-test", "aws-external-anthropic")]
|
||||
#[case::mantle("bedrock_mantle/claude-test", "bedrock")]
|
||||
#[tokio::test]
|
||||
async fn managed_aws_credentials_sign_the_exact_native_request(
|
||||
call: MessagesCall,
|
||||
#[case] model: &str,
|
||||
#[case] service: &str,
|
||||
) {
|
||||
let upstream = upstream([message_response()]).await;
|
||||
let secrets = Arc::new(RecordingSecrets::new([
|
||||
("AWS_ACCESS_KEY_ID", "test-access"),
|
||||
("AWS_SECRET_ACCESS_KEY", "test-secret"),
|
||||
("AWS_SESSION_TOKEN", "test-session"),
|
||||
("AWS_REGION_NAME", "us-west-2"),
|
||||
("ANTHROPIC_AWS_WORKSPACE_ID", "managed-workspace"),
|
||||
]));
|
||||
run_with(
|
||||
secrets,
|
||||
MessagesCall {
|
||||
api_base: Some(upstream.uri()),
|
||||
custom_llm_provider: None,
|
||||
..with_model(call, model)
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let request = only_request(&upstream).await;
|
||||
let auth = request.header("authorization").unwrap();
|
||||
assert!(auth.starts_with("AWS4-HMAC-SHA256 Credential=test-access/"));
|
||||
assert!(auth.contains(&format!("/us-west-2/{service}/aws4_request")));
|
||||
assert_eq!(request.header("x-amz-security-token"), Some("test-session"));
|
||||
assert_eq!(request.header("x-api-key"), None);
|
||||
assert_eq!(request.json()["model"], "claude-test");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn private_mantle_endpoint_is_used_and_signed(call: MessagesCall) {
|
||||
let upstream = upstream([message_response()]).await;
|
||||
let params = litellm_router_types::LitellmParams {
|
||||
aws: litellm_auth::AwsParams {
|
||||
aws_bedrock_runtime_endpoint: Some(upstream.uri()),
|
||||
aws_access_key_id: Some("private-access".into()),
|
||||
aws_secret_access_key: Some("private-secret".into()),
|
||||
aws_region_name: Some("us-west-2".into()),
|
||||
..Default::default()
|
||||
},
|
||||
..Default::default()
|
||||
};
|
||||
run(MessagesCall {
|
||||
custom_llm_provider: None,
|
||||
litellm_params: params,
|
||||
..with_model(call, "bedrock/mantle/claude-test")
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
let request = only_request(&upstream).await;
|
||||
assert_eq!(request.url.path(), "/anthropic/v1/messages");
|
||||
assert!(
|
||||
request
|
||||
.header("authorization")
|
||||
.unwrap()
|
||||
.contains("/us-west-2/bedrock/aws4_request")
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::platform("bedrock/claude_platform/claude-test")]
|
||||
#[case::mantle("bedrock_mantle/claude-test")]
|
||||
#[tokio::test]
|
||||
async fn native_streams_relay_every_sse_byte(call: MessagesCall, #[case] model: &str) {
|
||||
let sse = "event: message_start\ndata: {\"type\":\"message_start\"}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n";
|
||||
let upstream =
|
||||
upstream([ResponseTemplate::new(200).set_body_raw(sse, "text/event-stream")]).await;
|
||||
let response = messages_route(no_secrets())
|
||||
.execute(
|
||||
MessagesCall {
|
||||
api_key: Some("key".into()),
|
||||
api_base: Some(upstream.uri()),
|
||||
custom_llm_provider: None,
|
||||
..with_fields(
|
||||
call,
|
||||
json!({"model": model, "stream": true, "workspace_id": "workspace"}),
|
||||
)
|
||||
},
|
||||
&(),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let MessagesCallResponse::Stream { chunks, .. } = response else {
|
||||
panic!("expected SSE");
|
||||
};
|
||||
assert_eq!(
|
||||
chunks.try_collect::<Vec<_>>().await.unwrap().concat(),
|
||||
sse.as_bytes()
|
||||
);
|
||||
assert_eq!(only_request(&upstream).await.json()["stream"], true);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::missing(false)]
|
||||
#[case::configured(true)]
|
||||
#[tokio::test]
|
||||
async fn deployment_workspace_is_required_and_sent_as_a_header(
|
||||
call: MessagesCall,
|
||||
#[case] configured: bool,
|
||||
) {
|
||||
let upstream = upstream([message_response()]).await;
|
||||
let params = if configured {
|
||||
serde_json::from_value(json!({"model": "bedrock/claude_platform/claude-test", "workspace_id": "deployment-workspace"})).unwrap()
|
||||
} else {
|
||||
Default::default()
|
||||
};
|
||||
let result = run(MessagesCall {
|
||||
api_key: Some("key".into()),
|
||||
api_base: Some(upstream.uri()),
|
||||
litellm_params: params,
|
||||
..with_model(call, "bedrock/claude_platform/claude-test")
|
||||
})
|
||||
.await;
|
||||
if !configured {
|
||||
assert!(matches!(result, Err(Error::Auth(_))));
|
||||
assert!(received(&upstream).await.is_empty());
|
||||
return;
|
||||
}
|
||||
assert!(matches!(result.unwrap(), MessagesCallResponse::Complete(_)));
|
||||
let sent = only_request(&upstream).await;
|
||||
assert_eq!(
|
||||
sent.header("anthropic-workspace-id"),
|
||||
Some("deployment-workspace")
|
||||
);
|
||||
assert_eq!(sent.json().get("workspace_id"), None);
|
||||
}
|
||||
|
|
@ -13,8 +13,8 @@ use litellm_cache_response::{
|
|||
CacheOptions, CacheScope, ResponseCache, ResponseCacheConfig, ResponseCacheService, ScopedCache,
|
||||
};
|
||||
use litellm_host::interceptors::{
|
||||
ExecutionFacts, Interceptors, ProviderIdentity, RawResponse, RequestContext, ResultSource,
|
||||
WireRequest,
|
||||
Cost, ExecutionFacts, Interceptors, ProviderIdentity, RawResponse, RequestContext,
|
||||
ResultSource, WireRequest,
|
||||
};
|
||||
use litellm_inference::{
|
||||
RouteError,
|
||||
|
|
@ -241,3 +241,55 @@ async fn messages_cache_identity_follows_resolved_configuration_and_request_call
|
|||
first.verify().await;
|
||||
second.verify().await;
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::reported(json!(0.37), Some(0.37))]
|
||||
#[case::free(json!(0), Some(0.0))]
|
||||
#[case::malformed(json!("invalid"), None)]
|
||||
#[case::missing(Value::Null, None)]
|
||||
#[tokio::test]
|
||||
async fn provider_reported_cost_is_delivered_once_and_never_charged_for_a_cache_hit(
|
||||
cache: Arc<dyn ResponseCacheService>,
|
||||
#[case] cost: Value,
|
||||
#[case] expected: Option<f64>,
|
||||
) {
|
||||
use litellm_inference_messages::{MessagesCall, MessagesCallResponse};
|
||||
let response = json!({"id": "msg_cost", "type": "message", "role": "assistant", "model": "native-test",
|
||||
"content": [{"type": "text", "text": "answer"}], "stop_reason": "end_turn", "stop_sequence": null,
|
||||
"usage": {"input_tokens": 3, "output_tokens": 2}, "cost": cost});
|
||||
let upstream = support::upstream([support::json_response(response)]).await;
|
||||
let hooks = ChangingHooks::default();
|
||||
let route = support::messages_route(litellm_inference_testing::no_secrets())
|
||||
.with_cache(ScopedCache::new(cache, CacheScope::Shared));
|
||||
let request = || {
|
||||
MessagesCall {
|
||||
body: serde_json::from_value(json!({"model": "native-test", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16})).unwrap(),
|
||||
api_key: Some("eden-key".into()), api_base: Some(upstream.uri()), custom_llm_provider: Some("edenai".into()),
|
||||
litellm_params: Default::default(), extra_headers: None, provider_specific_header: None,
|
||||
timeout: None, shaping: Default::default(),
|
||||
}
|
||||
};
|
||||
let first = route.execute(request(), &hooks, None).await.unwrap();
|
||||
let second = route.execute(request(), &hooks, None).await.unwrap();
|
||||
let (MessagesCallResponse::Complete(first), MessagesCallResponse::Complete(second)) =
|
||||
(first, second)
|
||||
else {
|
||||
panic!("expected Messages responses");
|
||||
};
|
||||
assert_eq!(first, second);
|
||||
assert_eq!(first.extra.get("cost"), Some(&cost));
|
||||
assert_eq!(support::received(&upstream).await.len(), 1);
|
||||
let facts = hooks.facts.lock().unwrap();
|
||||
let [provider, cached] = facts.as_slice() else {
|
||||
panic!("expected provider and cache results");
|
||||
};
|
||||
assert_eq!(provider.source, ResultSource::Provider);
|
||||
assert_eq!(
|
||||
provider.cost,
|
||||
expected.map_or(Cost::Deferred, |amount| Cost::Reported {
|
||||
amount: serde_json::Number::from_f64(amount).unwrap()
|
||||
})
|
||||
);
|
||||
assert!(matches!(cached.source, ResultSource::Cache { .. }));
|
||||
assert_eq!(cached.cost, Cost::CacheHit);
|
||||
}
|
||||
|
|
|
|||
8
litellm-rust/crates/inference-messages/tests/live.rs
Normal file
8
litellm-rust/crates/inference-messages/tests/live.rs
Normal file
|
|
@ -0,0 +1,8 @@
|
|||
#[path = "live/support.rs"]
|
||||
mod support;
|
||||
|
||||
#[path = "live/edenai.rs"]
|
||||
mod edenai;
|
||||
|
||||
#[path = "live/vertex_ai.rs"]
|
||||
mod vertex_ai;
|
||||
92
litellm-rust/crates/inference-messages/tests/live/edenai.rs
Normal file
92
litellm-rust/crates/inference-messages/tests/live/edenai.rs
Normal file
|
|
@ -0,0 +1,92 @@
|
|||
use litellm_host::interceptors::Cost;
|
||||
use litellm_inference_messages::MessagesRoute;
|
||||
use litellm_inference_testing::live::within_deadline;
|
||||
use rstest::rstest;
|
||||
use serde_json::{Number, Value, json};
|
||||
|
||||
use super::support::{LiveCall, assert_text, call, complete, request, route, stream_events};
|
||||
|
||||
const PROVIDER: &str = "edenai";
|
||||
|
||||
#[rstest]
|
||||
#[case::basic(json!({"messages": [{"role": "user", "content": "Say hello."}]}))]
|
||||
#[case::system_text(json!({
|
||||
"system": "You are a helpful assistant.",
|
||||
"messages": [{"role": "user", "content": "Say hello."}]
|
||||
}))]
|
||||
#[case::portable_cache_hint(json!({
|
||||
"system": [{"type": "text", "text": "You are a helpful assistant.", "cache_control": {"type": "ephemeral", "ttl": "1h"}}],
|
||||
"messages": [{"role": "user", "content": "Say hello."}]
|
||||
}))]
|
||||
#[ignore = "calls a real provider, requires credentials and a live model"]
|
||||
#[tokio::test]
|
||||
async fn completion_and_cost_handoff(route: MessagesRoute, #[case] payload: Value) {
|
||||
within_deadline(async {
|
||||
let host = LiveCall::new(PROVIDER);
|
||||
let message = complete(request(&route, call(PROVIDER, payload), &host).await);
|
||||
assert_text(&message);
|
||||
host.assert_provider_result();
|
||||
let expected_cost = message
|
||||
.extra
|
||||
.get("cost")
|
||||
.and_then(|value| match value {
|
||||
Value::Number(number) => number.as_f64(),
|
||||
Value::String(text) => text.parse::<f64>().ok(),
|
||||
_ => None,
|
||||
})
|
||||
.filter(|cost| *cost >= 0.0)
|
||||
.and_then(Number::from_f64)
|
||||
.map_or(Cost::Deferred, |amount| Cost::Reported { amount });
|
||||
assert_eq!(host.facts().cost, expected_cost);
|
||||
println!("{}", serde_json::to_string(&message).unwrap());
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[ignore = "calls a real provider, requires credentials and a live model"]
|
||||
#[tokio::test]
|
||||
async fn streaming(route: MessagesRoute) {
|
||||
within_deadline(async {
|
||||
let host = LiveCall::new(PROVIDER);
|
||||
let events = stream_events(
|
||||
request(
|
||||
&route,
|
||||
call(
|
||||
PROVIDER,
|
||||
json!({
|
||||
"stream": true, "messages": [{"role": "user", "content": "Say hello."}]
|
||||
}),
|
||||
),
|
||||
&host,
|
||||
)
|
||||
.await,
|
||||
)
|
||||
.await;
|
||||
assert!(events.iter().any(|event| {
|
||||
event["delta"]["text"]
|
||||
.as_str()
|
||||
.is_some_and(|text| !text.is_empty())
|
||||
}));
|
||||
host.assert_provider_result();
|
||||
assert_eq!(host.facts().cost, Cost::Deferred);
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::complete(false)]
|
||||
#[case::streaming(true)]
|
||||
#[ignore = "calls a real provider, requires credentials and a live model"]
|
||||
#[tokio::test]
|
||||
async fn tool_round_trip(route: MessagesRoute, #[case] stream: bool) {
|
||||
super::support::tool_round_trip(
|
||||
route,
|
||||
PROVIDER,
|
||||
stream,
|
||||
json!({
|
||||
"max_tokens": 512, "tool_choice": {"type": "auto"}
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
265
litellm-rust/crates/inference-messages/tests/live/support.rs
Normal file
265
litellm-rust/crates/inference-messages/tests/live/support.rs
Normal file
|
|
@ -0,0 +1,265 @@
|
|||
use std::{sync::OnceLock, time::Duration};
|
||||
|
||||
use futures_util::{FutureExt, TryStreamExt};
|
||||
use litellm_host::interceptors::{
|
||||
ExecutionFacts, Interceptors, RawResponse, RequestContext, ResultSource, WireRequest,
|
||||
};
|
||||
use litellm_inference_messages::{Error, MessagesCall, MessagesCallResponse, MessagesRoute};
|
||||
use litellm_inference_testing::live::{LiveResources, model, within_deadline};
|
||||
use litellm_llms_types::formats::messages::MessagesResponse;
|
||||
use serde_json::{Value, json};
|
||||
|
||||
pub(super) struct LiveCall {
|
||||
provider: &'static str,
|
||||
facts: OnceLock<ExecutionFacts>,
|
||||
}
|
||||
|
||||
impl LiveCall {
|
||||
pub(super) fn new(provider: &'static str) -> Self {
|
||||
Self {
|
||||
provider,
|
||||
facts: OnceLock::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn facts(&self) -> &ExecutionFacts {
|
||||
self.facts
|
||||
.get()
|
||||
.expect("the route must deliver execution facts")
|
||||
}
|
||||
|
||||
pub(super) fn assert_provider_result(&self) {
|
||||
let facts = self.facts();
|
||||
assert_eq!(facts.provider.provider, self.provider);
|
||||
assert_eq!(facts.source, ResultSource::Provider);
|
||||
println!(
|
||||
"{}",
|
||||
json!({"route": "messages", "provider": self.provider, "source": "provider"})
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
impl Interceptors<Error> for LiveCall {
|
||||
async fn result_ready(&self, facts: ExecutionFacts) -> Result<(), Error> {
|
||||
assert!(
|
||||
self.facts.set(facts).is_ok(),
|
||||
"result must be accepted exactly once"
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn before_provider_request(
|
||||
&self,
|
||||
wire: WireRequest,
|
||||
context: RequestContext,
|
||||
) -> Result<WireRequest, Error> {
|
||||
assert_eq!(context.custom_llm_provider, self.provider);
|
||||
println!(
|
||||
"{}",
|
||||
json!({"route": "messages", "provider": self.provider, "model": context.model})
|
||||
);
|
||||
Ok(wire)
|
||||
}
|
||||
|
||||
async fn after_provider_response(&self, _: RawResponse) -> Result<(), Error> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest::fixture]
|
||||
pub(super) fn route() -> MessagesRoute {
|
||||
let resources = LiveResources::default();
|
||||
MessagesRoute::new(resources.http, resources.auth, resources.secrets)
|
||||
}
|
||||
|
||||
pub(super) async fn request(
|
||||
route: &MessagesRoute,
|
||||
call: MessagesCall,
|
||||
host: &LiveCall,
|
||||
) -> MessagesCallResponse {
|
||||
route
|
||||
.execute(call, host, None)
|
||||
.boxed()
|
||||
.await
|
||||
.expect("live request failed")
|
||||
}
|
||||
|
||||
pub(super) fn call(provider: &str, payload: Value) -> MessagesCall {
|
||||
let Value::Object(fields) = payload else {
|
||||
panic!("live payload must be an object")
|
||||
};
|
||||
let body = serde_json::from_value(Value::Object(
|
||||
[("max_tokens".into(), json!(128))]
|
||||
.into_iter()
|
||||
.chain(fields)
|
||||
.chain([("model".into(), json!(model("messages", provider)))])
|
||||
.collect(),
|
||||
))
|
||||
.expect("live payload must be a valid Messages request");
|
||||
MessagesCall {
|
||||
body,
|
||||
api_key: None,
|
||||
api_base: None,
|
||||
custom_llm_provider: Some(provider.into()),
|
||||
litellm_params: Default::default(),
|
||||
extra_headers: None,
|
||||
provider_specific_header: None,
|
||||
timeout: Some(Duration::from_secs(60)),
|
||||
shaping: Default::default(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn complete(response: MessagesCallResponse) -> Box<MessagesResponse> {
|
||||
match response {
|
||||
MessagesCallResponse::Complete(message) => message,
|
||||
MessagesCallResponse::Stream { .. } => panic!("a complete request returned a stream"),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn assert_text(message: &MessagesResponse) {
|
||||
assert!(!message.id.is_empty());
|
||||
assert_eq!(message.role, "assistant");
|
||||
assert!(message.stop_reason.is_some());
|
||||
assert!(
|
||||
message
|
||||
.content
|
||||
.iter()
|
||||
.any(|block| block["text"].as_str().is_some_and(|text| !text.is_empty()))
|
||||
);
|
||||
assert!(
|
||||
message
|
||||
.usage
|
||||
.as_ref()
|
||||
.and_then(|usage| usage["output_tokens"].as_u64())
|
||||
.is_some_and(|count| count > 0)
|
||||
);
|
||||
}
|
||||
|
||||
pub(super) async fn stream_events(response: MessagesCallResponse) -> Vec<Value> {
|
||||
let MessagesCallResponse::Stream { head, chunks } = response else {
|
||||
panic!("stream request returned a complete response")
|
||||
};
|
||||
let content_types: Vec<_> = head
|
||||
.headers
|
||||
.iter()
|
||||
.filter(|(name, _)| name.eq_ignore_ascii_case("content-type"))
|
||||
.map(|(_, value)| value.split(';').next().unwrap().trim())
|
||||
.collect();
|
||||
assert_eq!(content_types, ["text/event-stream"]);
|
||||
let frames = litellm_framer::frames(chunks, litellm_framer::sse::SseCodec::default())
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.expect("live SSE framing failed");
|
||||
let data_frames = match frames.last() {
|
||||
Some(frame) if frame.data == "[DONE]" => &frames[..frames.len() - 1],
|
||||
_ => frames.as_slice(),
|
||||
};
|
||||
let events: Vec<Value> = data_frames
|
||||
.iter()
|
||||
.map(|frame| serde_json::from_str(&frame.data).expect("SSE data must be JSON"))
|
||||
.collect();
|
||||
assert_eq!(
|
||||
events.first().map(|event| event["type"].as_str()),
|
||||
Some(Some("message_start"))
|
||||
);
|
||||
assert_eq!(
|
||||
events.last().map(|event| event["type"].as_str()),
|
||||
Some(Some("message_stop"))
|
||||
);
|
||||
assert!(!events.iter().any(|event| event["type"] == "error"));
|
||||
assert!(events.iter().any(|event| {
|
||||
event["usage"]["output_tokens"]
|
||||
.as_u64()
|
||||
.is_some_and(|count| count > 0)
|
||||
}));
|
||||
println!("{}", json!({"events": events}));
|
||||
events
|
||||
}
|
||||
|
||||
fn prompt() -> Value {
|
||||
json!({"role": "user", "content": "Use the echo tool with value live-check, then tell me the result."})
|
||||
}
|
||||
|
||||
fn tools() -> Value {
|
||||
json!([{"name": "echo", "description": "Return a value", "input_schema": {
|
||||
"type": "object", "properties": {"value": {"type": "string", "enum": ["live-check"]}}, "required": ["value"]
|
||||
}}])
|
||||
}
|
||||
|
||||
fn fragments(events: &[Value], index: u64, field: &str) -> String {
|
||||
events
|
||||
.iter()
|
||||
.filter(|event| event["index"].as_u64() == Some(index))
|
||||
.filter_map(|event| event["delta"][field].as_str())
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn streamed_block(start: &Value, events: &[Value]) -> Value {
|
||||
let index = start["index"]
|
||||
.as_u64()
|
||||
.expect("content block must have an index");
|
||||
let block = &start["content_block"];
|
||||
match block["type"].as_str() {
|
||||
Some("tool_use") => {
|
||||
let partial = fragments(events, index, "partial_json");
|
||||
let input = match partial.is_empty() {
|
||||
true => block["input"].clone(),
|
||||
false => {
|
||||
serde_json::from_str(&partial).expect("streamed tool arguments must be JSON")
|
||||
}
|
||||
};
|
||||
json!({"type": "tool_use", "id": block["id"], "name": block["name"], "input": input})
|
||||
}
|
||||
Some("text") => {
|
||||
json!({"type": "text", "text": format!("{}{}", block["text"].as_str().unwrap_or_default(), fragments(events, index, "text"))})
|
||||
}
|
||||
Some("thinking") => {
|
||||
json!({"type": "thinking", "thinking": format!("{}{}", block["thinking"].as_str().unwrap_or_default(), fragments(events, index, "thinking")),
|
||||
"signature": format!("{}{}", block["signature"].as_str().unwrap_or_default(), fragments(events, index, "signature"))})
|
||||
}
|
||||
_ => block.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn tool_round_trip(
|
||||
route: MessagesRoute,
|
||||
provider: &'static str,
|
||||
stream: bool,
|
||||
tool_params: Value,
|
||||
) {
|
||||
within_deadline(async {
|
||||
let first_host = LiveCall::new(provider);
|
||||
let payload = json!({"messages": [prompt()], "tools": tools(), "stream": stream});
|
||||
let fields = payload.as_object().expect("tool payload must be an object").clone();
|
||||
let params = tool_params.as_object().expect("tool parameters must be an object").clone();
|
||||
let first = request(&route, call(provider, Value::Object(fields.into_iter().chain(params).collect())), &first_host).await;
|
||||
let content = match stream {
|
||||
false => {
|
||||
let message = complete(first);
|
||||
assert_eq!(message.stop_reason.as_deref(), Some("tool_use"));
|
||||
message.content
|
||||
}
|
||||
true => {
|
||||
let events = stream_events(first).await;
|
||||
assert!(events.iter().any(|event| event["delta"]["stop_reason"] == "tool_use"));
|
||||
events.iter().filter(|event| event["type"] == "content_block_start")
|
||||
.map(|event| streamed_block(event, &events)).collect()
|
||||
}
|
||||
};
|
||||
let tool = content.iter().find(|block| block["type"] == "tool_use").expect("requested tool must be returned");
|
||||
assert_eq!(tool["name"], "echo");
|
||||
assert_eq!(tool["input"]["value"], "live-check");
|
||||
let tool_id = tool["id"].as_str().filter(|id| !id.is_empty()).expect("tool must have an id");
|
||||
first_host.assert_provider_result();
|
||||
let second_host = LiveCall::new(provider);
|
||||
let second = complete(request(&route, call(provider, json!({
|
||||
"tools": tools(), "messages": [prompt(),
|
||||
{"role": "assistant", "content": content},
|
||||
{"role": "user", "content": [{"type": "tool_result", "tool_use_id": tool_id, "content": "live-check"}]}]
|
||||
})), &second_host).await);
|
||||
assert_text(&second);
|
||||
second_host.assert_provider_result();
|
||||
println!("{}", serde_json::to_string(&second).unwrap());
|
||||
})
|
||||
.await;
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue