mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
Merge origin/main into litellm_vector_store_files_team_model_hint
This commit is contained in:
commit
3fa0e6fbfd
2636 changed files with 280002 additions and 95204 deletions
7
.circleci/AGENTS.md
Normal file
7
.circleci/AGENTS.md
Normal file
|
|
@ -0,0 +1,7 @@
|
|||
# .circleci
|
||||
|
||||
The `integration` workflow on `main` runs from a scheduled pipeline in the CircleCI project settings, cron `9 0-8,14-23 * * *` (UTC) on `main` with no pipeline parameters. It has no `triggers:` block in this file because a workflow with one runs only on its schedule, so the API pipelines that set `run_migration_tests` or `routing_parity_base` would stop running it. If main stops getting integration runs, check that the schedule still exists and recreate it with these values
|
||||
|
||||
Jobs clone with the `checkout_over_https` command, never the built-in `checkout` step, because CircleCI falls back to the project's SSH checkout key without its GitHub App access, GitHub rejects that key, and the public repo needs no key. It fetches the job's commit with full history and no file contents (`--filter=blob:none`), like the built-in step did, so `scripts/path_filter.sh` can find a merge base
|
||||
|
||||
A step that checks out files from another commit fetches that commit with `--no-filter`, as `routing_parity` does, because the clone has no file contents and git would otherwise download each missing file in its own request, which took 210 to 896 seconds for the routing parity base
|
||||
|
|
@ -21,19 +21,30 @@ orbs:
|
|||
win: circleci/windows@5.0 # Add Windows orb
|
||||
|
||||
commands:
|
||||
checkout_migration_source:
|
||||
checkout_over_https:
|
||||
parameters:
|
||||
revision:
|
||||
type: string
|
||||
default: ""
|
||||
shell:
|
||||
type: string
|
||||
default: /bin/bash -eo pipefail
|
||||
steps:
|
||||
- run:
|
||||
name: Select the requested migration test revision
|
||||
name: Check out the public repository over HTTPS
|
||||
shell: << parameters.shell >>
|
||||
environment:
|
||||
MIGRATION_SOURCE_SHA: << pipeline.parameters.migration_source_sha >>
|
||||
REQUESTED_REVISION: << parameters.revision >>
|
||||
command: |
|
||||
revision="${MIGRATION_SOURCE_SHA:-$CIRCLE_SHA1}"
|
||||
revision="${REQUESTED_REVISION:-$CIRCLE_SHA1}"
|
||||
[[ "$revision" =~ ^[0-9a-f]{40}$ ]] || exit 1
|
||||
git init
|
||||
git remote add origin https://github.com/BerriAI/litellm.git
|
||||
git fetch --depth 1 origin "$revision"
|
||||
git checkout --detach FETCH_HEAD
|
||||
for attempt in 1 2 3; do
|
||||
git fetch --filter=blob:none origin "$revision" && git checkout --detach FETCH_HEAD && break
|
||||
[[ $attempt -lt 3 ]] || exit 1
|
||||
sleep 5
|
||||
done
|
||||
skip_if_unrelated_changes:
|
||||
parameters:
|
||||
category:
|
||||
|
|
@ -347,7 +358,7 @@ commands:
|
|||
- when:
|
||||
condition: << parameters.checkout >>
|
||||
steps:
|
||||
- checkout
|
||||
- checkout_over_https
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- install_rust
|
||||
|
|
@ -377,7 +388,8 @@ jobs:
|
|||
CARGO_HTTP_MULTIPLEXING: "false"
|
||||
CARGO_NET_RETRY: "5"
|
||||
steps:
|
||||
- checkout
|
||||
- checkout_over_https:
|
||||
shell: bash.exe
|
||||
- run:
|
||||
name: Skip job when no windows-release-relevant files changed
|
||||
shell: bash.exe
|
||||
|
|
@ -451,7 +463,7 @@ jobs:
|
|||
password: ${DOCKERHUB_PASSWORD}
|
||||
working_directory: ~/project
|
||||
steps:
|
||||
- checkout
|
||||
- checkout_over_https
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- install_rust
|
||||
|
|
@ -482,7 +494,7 @@ jobs:
|
|||
working_directory: ~/project
|
||||
|
||||
steps:
|
||||
- checkout
|
||||
- checkout_over_https
|
||||
- skip_if_unrelated_changes
|
||||
- attach_workspace:
|
||||
at: ~/project
|
||||
|
|
@ -559,7 +571,7 @@ jobs:
|
|||
resource_class: medium
|
||||
working_directory: ~/project
|
||||
steps:
|
||||
- checkout
|
||||
- checkout_over_https
|
||||
- run:
|
||||
name: Generate LiteLLM master key
|
||||
command: |
|
||||
|
|
@ -659,7 +671,7 @@ jobs:
|
|||
resource_class: large
|
||||
working_directory: ~/project
|
||||
steps:
|
||||
- checkout
|
||||
- checkout_over_https
|
||||
- run:
|
||||
name: Generate LiteLLM master key
|
||||
command: |
|
||||
|
|
@ -740,7 +752,7 @@ jobs:
|
|||
resource_class: large
|
||||
working_directory: ~/project
|
||||
steps:
|
||||
- checkout
|
||||
- checkout_over_https
|
||||
- skip_if_unrelated_changes:
|
||||
category: client
|
||||
- setup_google_dns
|
||||
|
|
@ -794,7 +806,7 @@ jobs:
|
|||
# the rest of the suite: nothing else performs a logout.
|
||||
PROXY_LOGOUT_URL: "https://www.example.com"
|
||||
steps:
|
||||
- checkout
|
||||
- checkout_over_https
|
||||
- run:
|
||||
name: Generate LiteLLM master key
|
||||
command: |
|
||||
|
|
@ -940,7 +952,7 @@ jobs:
|
|||
# and by migration.serverRootPath.config.ts, which refuses to run without it.
|
||||
SERVER_ROOT_PATH: "/litellm"
|
||||
steps:
|
||||
- checkout
|
||||
- checkout_over_https
|
||||
- run:
|
||||
name: Generate LiteLLM master key
|
||||
command: |
|
||||
|
|
@ -1066,11 +1078,12 @@ jobs:
|
|||
- when:
|
||||
condition: << parameters.migration_qualification >>
|
||||
steps:
|
||||
- checkout_migration_source
|
||||
- checkout_over_https:
|
||||
revision: << pipeline.parameters.migration_source_sha >>
|
||||
- unless:
|
||||
condition: << parameters.migration_qualification >>
|
||||
steps:
|
||||
- checkout
|
||||
- checkout_over_https
|
||||
- skip_if_unrelated_changes
|
||||
|
||||
- run:
|
||||
|
|
@ -1122,7 +1135,8 @@ jobs:
|
|||
MIGRATION_TEST_OUTPUT: /tmp/migration-results
|
||||
PYTHONPATH: tests/e2e
|
||||
steps:
|
||||
- checkout_migration_source
|
||||
- checkout_over_https:
|
||||
revision: << pipeline.parameters.migration_source_sha >>
|
||||
- run:
|
||||
name: Generate LiteLLM master key
|
||||
command: |
|
||||
|
|
@ -1197,7 +1211,7 @@ jobs:
|
|||
resource_class: medium
|
||||
working_directory: ~/project
|
||||
steps:
|
||||
- checkout
|
||||
- checkout_over_https
|
||||
- skip_if_unrelated_changes
|
||||
- attach_workspace:
|
||||
at: ~/project
|
||||
|
|
@ -1371,7 +1385,7 @@ jobs:
|
|||
environment:
|
||||
UV_PYTHON_DOWNLOADS: never
|
||||
steps:
|
||||
- checkout
|
||||
- checkout_over_https
|
||||
- install_uv
|
||||
- run:
|
||||
name: Install the locked shard verifier and coverage dependencies
|
||||
|
|
@ -1419,7 +1433,7 @@ jobs:
|
|||
ROUTING_PARITY_BASE: << pipeline.parameters.routing_parity_base >>
|
||||
command: |
|
||||
[[ "$ROUTING_PARITY_BASE" =~ ^[0-9a-f]{40}$ ]] || exit 1
|
||||
git fetch --depth 1 origin "$ROUTING_PARITY_BASE"
|
||||
git fetch --depth 1 --no-filter origin "$ROUTING_PARITY_BASE"
|
||||
git rm -r -f --quiet litellm enterprise litellm-proxy-extras
|
||||
git checkout "$ROUTING_PARITY_BASE" -- litellm enterprise litellm-proxy-extras
|
||||
git reset --quiet
|
||||
|
|
@ -1494,7 +1508,7 @@ jobs:
|
|||
environment:
|
||||
DATABASE_URL: postgresql://postgres:postgres@localhost:5432/litellm_test
|
||||
steps:
|
||||
- checkout
|
||||
- checkout_over_https
|
||||
- skip_if_unrelated_changes
|
||||
- setup_litellm_test_deps:
|
||||
checkout: false
|
||||
|
|
|
|||
|
|
@ -128,6 +128,8 @@ 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"
|
||||
|
||||
export UV_OFFLINE=true
|
||||
|
||||
if [ "$mode" = parity ]; then
|
||||
export INTEGRATION_WORKERS=1
|
||||
else
|
||||
|
|
|
|||
10
.github/PULL_REQUEST_TEMPLATE/general.md
vendored
10
.github/PULL_REQUEST_TEMPLATE/general.md
vendored
|
|
@ -4,9 +4,14 @@
|
|||
Drop every section you have nothing to put in, heading included: a bare "## Relevant issues" or
|
||||
"## Affected release" with nothing under it must not appear in the final description -->
|
||||
|
||||
<!-- Two or three plain sentences a teammate can read in 10 seconds, written the way you'd explain it in Slack
|
||||
## Handwritten 1-3 sentence description
|
||||
|
||||
<!-- The author writes this in their own words, and it can be left empty. AI agents never write here and leave the
|
||||
placeholder line below as is, and when updating an existing PR they keep whatever the author wrote. Keep this heading
|
||||
Say what the PR does for users and why, then anything that keeps it working over time (a sync job, a config someone owns)
|
||||
No headings, bullets, tables, or code names here. Everything else goes inside the Details block below -->
|
||||
No bullets, tables, or code names here. Everything else goes inside the Details block below -->
|
||||
|
||||
[insert handwritten description here]
|
||||
|
||||
<details>
|
||||
<summary>Details</summary>
|
||||
|
|
@ -88,6 +93,7 @@ If you're seeing a delay in your PR being merged, ping the LiteLLM Team on [Slac
|
|||
Otherwise proof must be fully e2e, no mocks, with real provider calls where applicable. Pytest alone is insufficient. The following format applies to live proof
|
||||
Show only the latest proof. Before uses the merge base and After the tested tip. Rerun affected QA after behavior changes. Verified neutral diffs may carry proof with the tested SHA, current tip, equivalence reason and focused checks, following AGENTS.md and /qa
|
||||
Use matching Before/After case headings with numbered commands and observed output. Put shared setup above them. For one case, omit case headings
|
||||
A reviewer must be able to rerun every case from the body alone, so the shared setup gives the proxy start command, the config.yaml, env var names (never values) and the client code or curl used, all exactly as run. Wrap any block longer than about 20 lines in a collapsed <details> with a one-line summary
|
||||
|
||||
### Before (<hash>)
|
||||
|
||||
|
|
|
|||
4
.github/workflows/compat-matrix-image.yml
vendored
4
.github/workflows/compat-matrix-image.yml
vendored
|
|
@ -4,7 +4,7 @@ on:
|
|||
pull_request:
|
||||
paths:
|
||||
- tests/e2e/claude_code/cron_vm/**
|
||||
- tests/e2e/claude_code/pr_gate_version_resolver.py
|
||||
- tests/e2e/cli_harness/npm_version_resolver.py
|
||||
- .github/workflows/compat-matrix-image.yml
|
||||
workflow_dispatch:
|
||||
|
||||
|
|
@ -36,7 +36,7 @@ jobs:
|
|||
whoami
|
||||
gh --version
|
||||
uv --version
|
||||
version="$(uv run --no-project --python 3.12 python /opt/litellm/tests/e2e/claude_code/pr_gate_version_resolver.py)"
|
||||
version="$(uv run --no-project --python 3.12 python /opt/litellm/tests/e2e/cli_harness/npm_version_resolver.py @anthropic-ai/claude-code)"
|
||||
/opt/litellm/tests/e2e/claude_code/cron_vm/install_claude_code.sh "${version}" /tmp/claude-cli
|
||||
/tmp/claude-cli/claude --version
|
||||
'
|
||||
|
|
|
|||
2
.github/workflows/create-rc-branch.yml
vendored
2
.github/workflows/create-rc-branch.yml
vendored
|
|
@ -2,7 +2,7 @@ name: Create RC Branch
|
|||
|
||||
on:
|
||||
schedule:
|
||||
- cron: "0 3 * * 5"
|
||||
- cron: "0 3 * * 1"
|
||||
timezone: "America/Los_Angeles"
|
||||
workflow_dispatch:
|
||||
|
||||
|
|
|
|||
30
.github/workflows/deploy-dev-on-main.yml
vendored
30
.github/workflows/deploy-dev-on-main.yml
vendored
|
|
@ -1,30 +0,0 @@
|
|||
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
|
||||
2
.github/workflows/image-scan.yml
vendored
2
.github/workflows/image-scan.yml
vendored
|
|
@ -53,7 +53,7 @@ jobs:
|
|||
if: >-
|
||||
github.event_name != 'pull_request' ||
|
||||
github.event.pull_request.head.repo.full_name == github.repository
|
||||
timeout-minutes: 30
|
||||
timeout-minutes: 60
|
||||
permissions:
|
||||
contents: read
|
||||
steps:
|
||||
|
|
|
|||
|
|
@ -49,9 +49,9 @@ jobs:
|
|||
id: compliance
|
||||
run: |
|
||||
uv run --no-sync pytest tests/unit/interactions/test_openapi_compliance.py \
|
||||
--interactions-spec-url=https://ai.google.dev/static/api/interactions.openapi.json \
|
||||
--interactions-spec-url=https://aistudio.google.com/api-reference/interactions.openapi.json \
|
||||
--tb=short -vv -p no:rerunfailures -p no:pytest-retry \
|
||||
--allow-hosts=127.0.0.1,::1,localhost,ai.google.dev
|
||||
--allow-hosts=127.0.0.1,::1,localhost,ai.google.dev,aistudio.google.com
|
||||
|
||||
tracking-issue:
|
||||
name: tracking-issue
|
||||
|
|
|
|||
3
.github/workflows/osv-scan.yml
vendored
3
.github/workflows/osv-scan.yml
vendored
|
|
@ -5,6 +5,9 @@ on:
|
|||
branches:
|
||||
- main
|
||||
- "litellm_**"
|
||||
- "rc/**"
|
||||
- "stable/**"
|
||||
- "lts/**"
|
||||
schedule:
|
||||
- cron: "23 6 * * *"
|
||||
workflow_dispatch:
|
||||
|
|
|
|||
4
.github/workflows/test-linting.yml
vendored
4
.github/workflows/test-linting.yml
vendored
|
|
@ -21,7 +21,7 @@ jobs:
|
|||
pull-requests: read
|
||||
actions: read
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 20
|
||||
timeout-minutes: 30
|
||||
services:
|
||||
redis:
|
||||
image: redis@sha256:e2debfb7956fa12c7ddc79d7e645c8cf26b30c99a6e9161ea9bf4171e1668a5f
|
||||
|
|
@ -208,7 +208,7 @@ jobs:
|
|||
exit 0
|
||||
fi
|
||||
retry() { "$@" || { sleep 15; "$@"; } || { sleep 30; "$@"; }; }
|
||||
CLAUDE_VERSION="$(retry uv run --no-sync python tests/e2e/claude_code/pr_gate_version_resolver.py)"
|
||||
CLAUDE_VERSION="$(retry uv run --no-sync python tests/e2e/cli_harness/npm_version_resolver.py @anthropic-ai/claude-code)"
|
||||
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 \
|
||||
|
|
|
|||
52
.github/workflows/test-unit.yml
vendored
52
.github/workflows/test-unit.yml
vendored
|
|
@ -5,6 +5,9 @@ on:
|
|||
branches:
|
||||
- main
|
||||
- "litellm_**"
|
||||
- "rc/**"
|
||||
- "stable/**"
|
||||
- "lts/**"
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
|
|
@ -42,7 +45,7 @@ jobs:
|
|||
id: rust-bridge-key
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
env:
|
||||
RUST_BRIDGE_CACHE_KEY: ${{ runner.os }}-rust-bridge-editable-${{ hashFiles('litellm-rust/Cargo.lock', 'litellm-rust/Cargo.toml', 'litellm-rust/crates/**', 'model_prices_and_context_window.json', 'pyproject.toml') }}
|
||||
RUST_BRIDGE_CACHE_KEY: ${{ runner.os }}-rust-bridge-editable-${{ hashFiles('litellm-rust/Cargo.lock', 'litellm-rust/Cargo.toml', 'litellm-rust/crates/**', 'model_prices_and_context_window.json', 'pyproject.toml', 'packaging/litellm-proxy/pyproject.toml') }}
|
||||
run: echo "key=$RUST_BRIDGE_CACHE_KEY" >> "$GITHUB_OUTPUT"
|
||||
|
||||
- name: Set up Python
|
||||
|
|
@ -95,6 +98,18 @@ jobs:
|
|||
mkdir -p rust-bridge-dist/litellm/rust_bridge
|
||||
cp litellm/rust_bridge/_native.abi3.so rust-bridge-dist/litellm/rust_bridge/_native.abi3.so
|
||||
|
||||
- name: Build the gateway wheel once for test shards
|
||||
if: steps.changes.outputs.decision != 'skip' && steps.rust-bridge-cache.outputs.cache-hit != 'true'
|
||||
timeout-minutes: 15
|
||||
working-directory: packaging/litellm-proxy
|
||||
env:
|
||||
RUSTC_WRAPPER: sccache
|
||||
SCCACHE_GHA_ENABLED: "true"
|
||||
SCCACHE_GHA_VERSION: litellm-ci-v1
|
||||
SCCACHE_GHA_RW_MODE: ${{ github.event_name == 'push' && github.ref == 'refs/heads/main' && 'READ_WRITE' || 'READ_ONLY' }}
|
||||
CARGO_INCREMENTAL: "0"
|
||||
run: uv run --no-project --with maturin==1.15.0 maturin build --profile dev --strip --out ../../rust-bridge-dist
|
||||
|
||||
- name: Upload Rust bridge artifact
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
timeout-minutes: 10
|
||||
|
|
@ -280,6 +295,7 @@ jobs:
|
|||
tests/unit/llms/e2b
|
||||
tests/unit/llms/edenai
|
||||
tests/unit/llms/elevenlabs
|
||||
tests/unit/llms/empiriolabs
|
||||
tests/unit/llms/exa_ai
|
||||
tests/unit/llms/fal_ai
|
||||
tests/unit/llms/fastcrw
|
||||
|
|
@ -415,6 +431,7 @@ jobs:
|
|||
|
||||
- shard: router
|
||||
test-path: >-
|
||||
tests/unit/router_rust
|
||||
tests/unit/test_router
|
||||
python-version: "3.12"
|
||||
workers: 4
|
||||
|
|
@ -441,6 +458,7 @@ jobs:
|
|||
tests/unit/ocr
|
||||
tests/unit/passthrough
|
||||
tests/unit/rag
|
||||
tests/unit/cookbook
|
||||
tests/unit/realtime_api
|
||||
tests/unit/rerank_api
|
||||
tests/unit/search
|
||||
|
|
@ -972,10 +990,10 @@ jobs:
|
|||
dist: loadscope
|
||||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
timeout-minutes: ${{ matrix.shard == 'db-and-spend' && 10 || 3 }}
|
||||
timeout-minutes: 3
|
||||
with:
|
||||
persist-credentials: false
|
||||
fetch-depth: ${{ matrix.shard == 'db-and-spend' && '0' || '1' }}
|
||||
fetch-depth: 1
|
||||
|
||||
- name: Detect relevant changes
|
||||
id: changes
|
||||
|
|
@ -1008,27 +1026,37 @@ jobs:
|
|||
with:
|
||||
artifact: ${{ needs.rust-bridge.outputs.artifact }}
|
||||
|
||||
- name: Install dependencies and check schema
|
||||
- name: Install dependencies
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
timeout-minutes: 8
|
||||
timeout-minutes: 5
|
||||
env:
|
||||
RUST_BRIDGE_ARTIFACT: ${{ needs.rust-bridge.outputs.artifact }}
|
||||
UNIT_SHARD: ${{ matrix.shard }}
|
||||
SCHEMA_BASE_REF: origin/${{ github.event.repository.default_branch }}
|
||||
run: |
|
||||
diff -u model_prices_and_context_window.json litellm/model_prices_and_context_window_backup.json
|
||||
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 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 extra_proxy --extra proxy-runtime --extra utils --no-install-project
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra extra_proxy --extra proxy-runtime --extra utils --no-install-project --no-install-package litellm-proxy
|
||||
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
|
||||
uv run --no-sync python -c "import importlib.metadata; import litellm.rust_bridge._native; print(importlib.metadata.version('litellm'))"
|
||||
fi
|
||||
uv run --no-sync python -c 'import os, sys; print(sys.version); assert f"{sys.version_info.major}.{sys.version_info.minor}" == os.environ["UV_PYTHON"]'
|
||||
if [ "$UNIT_SHARD" = "db-and-spend" ]; then
|
||||
uv run --no-sync python scripts/check_schema_compatibility.py --base-ref "$SCHEMA_BASE_REF"
|
||||
fi
|
||||
|
||||
- name: Check schema.prisma against the merge base
|
||||
if: steps.changes.outputs.decision != 'skip' && matrix.shard == 'db-and-spend' && github.event_name == 'pull_request'
|
||||
timeout-minutes: 3
|
||||
env:
|
||||
GH_TOKEN: ${{ github.token }}
|
||||
PR_NUMBER: ${{ github.event.pull_request.number }}
|
||||
EVENT_BASE_REF: ${{ github.event.pull_request.base.ref }}
|
||||
HEAD_SHA: ${{ github.event.pull_request.head.sha }}
|
||||
run: |
|
||||
base_ref=$(gh api "repos/${GITHUB_REPOSITORY}/branches/${EVENT_BASE_REF}" --jq .name) || base_ref=$(gh api "repos/${GITHUB_REPOSITORY}/pulls/${PR_NUMBER}" --jq .base.ref)
|
||||
merge_base=$(gh api "repos/${GITHUB_REPOSITORY}/compare/${base_ref}...${HEAD_SHA}?per_page=1" --jq '.merge_base_commit.sha')
|
||||
test -n "$merge_base"
|
||||
git fetch --no-tags --depth=1 origin "$merge_base"
|
||||
uv run --no-sync python scripts/check_schema_compatibility.py --base-commit "$merge_base"
|
||||
|
||||
- name: Cache Prisma binaries
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
|
|
@ -1470,7 +1498,7 @@ jobs:
|
|||
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 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 extra_proxy --extra proxy-runtime --extra utils --no-install-project
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra extra_proxy --extra proxy-runtime --extra utils --no-install-project --no-install-package litellm-proxy
|
||||
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
|
||||
|
|
|
|||
85
.github/workflows/weekly_load_anomaly.yml
vendored
85
.github/workflows/weekly_load_anomaly.yml
vendored
|
|
@ -1,85 +0,0 @@
|
|||
name: "Weekly Load Anomaly Check"
|
||||
|
||||
on:
|
||||
schedule:
|
||||
- cron: "0 12 * * 6"
|
||||
workflow_dispatch:
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
jobs:
|
||||
weekly-load-anomaly:
|
||||
if: github.event_name != 'schedule' || github.repository == 'BerriAI/litellm'
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 45
|
||||
services:
|
||||
postgres:
|
||||
image: postgres:16.6
|
||||
env:
|
||||
POSTGRES_USER: llmproxy
|
||||
POSTGRES_PASSWORD: dbpassword9090
|
||||
POSTGRES_DB: litellm
|
||||
ports:
|
||||
- 5432:5432
|
||||
options: >-
|
||||
--health-cmd "pg_isready -U llmproxy"
|
||||
--health-interval 5s
|
||||
--health-timeout 5s
|
||||
--health-retries 10
|
||||
env:
|
||||
DATABASE_URL: postgresql://llmproxy:dbpassword9090@localhost:5432/litellm
|
||||
LITELLM_MASTER_KEY: sk-weekly-anomaly-check
|
||||
ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
|
||||
AWS_BEARER_TOKEN_BEDROCK: ${{ secrets.AWS_BEARER_TOKEN_BEDROCK }}
|
||||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
with:
|
||||
python-version: "3.12"
|
||||
|
||||
- name: Set up uv
|
||||
uses: ./.github/actions/setup-uv-with-retries
|
||||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
- name: Cache the Rust build
|
||||
uses: ./.github/actions/cache-cargo-build
|
||||
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra proxy
|
||||
|
||||
- name: Cache Prisma binaries
|
||||
uses: ./.github/actions/cache-prisma-binaries
|
||||
|
||||
- name: Generate Prisma client
|
||||
run: |
|
||||
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
|
||||
|
||||
- name: Start the proxy
|
||||
run: |
|
||||
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
|
||||
fi
|
||||
sleep 2
|
||||
done
|
||||
echo "proxy never became live"
|
||||
tail -n 100 proxy.log
|
||||
exit 1
|
||||
|
||||
- name: Run the weekly session anomaly test
|
||||
env:
|
||||
E2E_WEEKLY_ANOMALY: "1"
|
||||
run: |
|
||||
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()
|
||||
run: tail -n 300 proxy.log
|
||||
|
|
@ -41,9 +41,11 @@ Same applies for filing bug reports and feature requests, with .github/ISSUE_TEM
|
|||
|
||||
If you're resolving a linear ticket, in the "## Linear ticket" section of the PR, say "Resolves LIT-1234", replacing "LIT-1234" with the actual ticket id that you're resolving. If you don't have the ticket id, don't make one up or search for it. Just drop the section
|
||||
|
||||
Run QA proxies with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --num_workers 2 --detailed_debug --reload --use_v2_migration_resolver 2>&1 | tee litellm.log`, and swap its `--config` for a copy of `litellm/proxy/dev_config.yaml` that points both `router_settings` (`redis_host`, `redis_port`) and `general_settings.coordination_redis` (`host`, `port`) at one dedicated Redis instance, with credentials when required, so the workers share router state and spend, budget, and rate-limit counters. Check that two worker PIDs are serving, and use the same worker count and Redis setup for the before and after proof. When the PR is about `--run_gunicorn`, `--run_hypercorn`, or `--run_granian`, which `--reload` does not support, drop `--reload`. When it is about behavior without Redis, drop Redis. When the proxy prints `"workers" flag is ignored when reloading is enabled`, that checkout or its uvicorn reloads only one worker, so drop the two-worker check for it. Say in the proof which part you dropped and why
|
||||
|
||||
Simple PRs may skip before/after e2e when the full diff is mechanical and focused checks completely prove it, such as #49030 restoring an existing helper name. Trace affected callers and run relevant tests/static checks at the tip. For a bug fix, establish that the regression catches the original defect. Size alone never qualifies. New behavior, API/provider semantics, auth, billing, persistence, concurrency or uncertain integration risk need live e2e. In Screenshots / Proof of Fix, record `E2E exception at <tip SHA>`, scope, rationale, commands and results. Tests count as proof only for this exception. Reassess changed diffs, retain CI/bot gates, and honor explicit e2e requests (2026-10-28, unnecessary live QA for #49030)
|
||||
|
||||
For other PRs, proof must show the end user's real flow before and after, with commands and output against a live proxy and real provider calls where applicable. Pytest alone is insufficient. For UI or interactive coding-tool flows, drive that surface and embed before/after screenshots, with URLs, clicks and input fields a reviewer can reproduce
|
||||
For other PRs, proof must show the end user's real flow before and after, with commands and output against a live proxy and real provider calls where applicable. The body alone must let a reviewer rerun it, with the proxy command, config, env var names and client code exactly as run, and long blocks collapsed in <details>. Pytest alone is insufficient. For UI or interactive coding-tool flows, drive that surface and embed before/after screenshots, with URLs, clicks and input fields a reviewer can reproduce
|
||||
|
||||
If you ever write any human-facing text (pull requests, issues, commit messages, discussion posts, github comments, release notes, docs, etc.), always follow these guidelines to sound less AI-y:
|
||||
- don't use emojis
|
||||
|
|
|
|||
|
|
@ -76,6 +76,7 @@ ENV UV_PROJECT_ENVIRONMENT=/app/.venv \
|
|||
COPY pyproject.toml uv.lock ./
|
||||
COPY enterprise/pyproject.toml enterprise/
|
||||
COPY litellm-proxy-extras/pyproject.toml litellm-proxy-extras/
|
||||
COPY packaging/litellm-proxy/pyproject.toml packaging/litellm-proxy/
|
||||
|
||||
# Install third-party dependencies (cached unless pyproject.toml/uv.lock change)
|
||||
RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-groups --no-editable \
|
||||
|
|
@ -101,7 +102,8 @@ COPY --from=ui-builder /ui/out/. litellm/proxy/_experimental/out/
|
|||
RUN sed -i 's/\r$//' docker/build_admin_ui.sh && chmod +x docker/build_admin_ui.sh && ./docker/build_admin_ui.sh
|
||||
|
||||
# Install project and workspace packages (fast - deps already cached)
|
||||
RUN uv sync --frozen --no-default-groups --no-editable \
|
||||
RUN uv build --package litellm-proxy --wheel --out-dir /wheels
|
||||
RUN uv sync --frozen --no-install-package litellm-proxy --no-default-groups --no-editable \
|
||||
--extra proxy \
|
||||
--extra proxy-runtime \
|
||||
--group admin-mcp \
|
||||
|
|
@ -111,6 +113,8 @@ RUN uv sync --frozen --no-default-groups --no-editable \
|
|||
--extra bedrock-realtime \
|
||||
--python python3.13
|
||||
|
||||
RUN uv pip install --python /app/.venv/bin/python --no-deps /wheels/litellm_proxy-*.whl
|
||||
|
||||
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
|
||||
npm_config_cache=/root/.npm \
|
||||
prisma generate --schema=./schema.prisma
|
||||
|
|
|
|||
|
|
@ -398,6 +398,7 @@ Set `LITELLM_PROXY_API_BASE` and `LITELLM_PROXY_API_KEY` and every model call th
|
|||
| [Deepseek (`deepseek`)](https://docs.litellm.ai/docs/providers/deepseek) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [Eden AI (`edenai`)](https://docs.litellm.ai/docs/providers/edenai) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | | | |
|
||||
| [ElevenLabs (`elevenlabs`)](https://docs.litellm.ai/docs/providers/elevenlabs) | ✅ | ✅ | ✅ | | | ✅ | ✅ | | | |
|
||||
| [EmpirioLabs AI (`empiriolabs`)](https://docs.litellm.ai/docs/providers/empiriolabs) | ✅ | | ✅ | | | | | | | |
|
||||
| [Empower (`empower`)](https://docs.litellm.ai/docs/providers/empower) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [eu/jev Decisions API (`eu_jev`)](https://docs.litellm.ai/docs/pass_through/eu_jev) | | | | | | | | | | |
|
||||
| [Fal AI (`fal_ai`)](https://docs.litellm.ai/docs/providers/fal_ai) | ✅ | ✅ | ✅ | | ✅ | | | | | |
|
||||
|
|
@ -444,6 +445,7 @@ Set `LITELLM_PROXY_API_BASE` and `LITELLM_PROXY_API_KEY` and every model call th
|
|||
| [OpenAI-like (`openai_like`)](https://docs.litellm.ai/docs/providers/openai_compatible) | | | | ✅ | | | | | | |
|
||||
| [OpenRouter (`openrouter`)](https://docs.litellm.ai/docs/providers/openrouter) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [OVHCloud AI Endpoints (`ovhcloud`)](https://docs.litellm.ai/docs/providers/ovhcloud) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [Parallel AI (`parallel_ai`)](https://docs.litellm.ai/docs/search/parallel_ai) | | | | | | | | | | |
|
||||
| [Perplexity AI (`perplexity`)](https://docs.litellm.ai/docs/providers/perplexity) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [Petals (`petals`)](https://docs.litellm.ai/docs/providers/petals) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [Pinstripes (`pinstripes`)](https://docs.litellm.ai/docs/providers/pinstripes) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
|
|
@ -696,7 +698,7 @@ For companies that need better security, user management and professional suppor
|
|||
[Talk to founders](https://enterprise.litellm.ai/demo)
|
||||
|
||||
This covers:
|
||||
- ✅ **Features under the [LiteLLM Commercial License](https://docs.litellm.ai/docs/proxy/enterprise):**
|
||||
- ✅ **Features under the [LiteLLM Commercial License](https://docs.litellm.ai/docs/enterprise):**
|
||||
- ✅ **Feature Prioritization**
|
||||
- ✅ **Custom Integrations**
|
||||
- ✅ **Professional Support - Dedicated discord + slack**
|
||||
|
|
|
|||
245
cookbook/gpt_realtime_translate.py
Normal file
245
cookbook/gpt_realtime_translate.py
Normal file
|
|
@ -0,0 +1,245 @@
|
|||
#!/usr/bin/env python3
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import wave
|
||||
from collections.abc import Iterator, Sequence
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from urllib.parse import urlencode, urlsplit, urlunsplit
|
||||
|
||||
import websockets
|
||||
from websockets.asyncio.client import ClientConnection
|
||||
|
||||
SAMPLE_RATE = 24_000
|
||||
CHANNELS = 1
|
||||
SAMPLE_WIDTH = 2
|
||||
CHUNK_DURATION_SECONDS = 0.1
|
||||
CHUNK_BYTES = int(SAMPLE_RATE * CHANNELS * SAMPLE_WIDTH * CHUNK_DURATION_SECONDS)
|
||||
RESPONSE_TIMEOUT_SECONDS = 30.0
|
||||
AUDIO_EVENT_TYPES = frozenset(
|
||||
{
|
||||
"session.output_audio.delta",
|
||||
"response.audio.delta",
|
||||
"response.output_audio.delta",
|
||||
}
|
||||
)
|
||||
TRANSCRIPT_EVENT_TYPES = frozenset(
|
||||
{
|
||||
"session.output_transcript.delta",
|
||||
"response.text.delta",
|
||||
"response.output_audio_transcript.delta",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Settings:
|
||||
input_wav: Path
|
||||
output_wav: Path
|
||||
base_url: str
|
||||
model: str
|
||||
target_language: str
|
||||
trailing_silence_seconds: float
|
||||
api_key: str
|
||||
|
||||
|
||||
def write_stdout(message: str = "", *, end: str = "\n", flush: bool = False) -> None:
|
||||
sys.stdout.write(f"{message}{end}")
|
||||
if flush:
|
||||
sys.stdout.flush()
|
||||
|
||||
|
||||
def write_stderr(message: str) -> None:
|
||||
sys.stderr.write(f"{message}\n")
|
||||
|
||||
|
||||
def parse_args(argv: Sequence[str] | None = None) -> Settings | str:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Stream a 24 kHz PCM16 WAV through gpt-realtime-translate and save the translated audio",
|
||||
)
|
||||
parser.add_argument("input_wav", type=Path)
|
||||
parser.add_argument("--output", type=Path, default=Path("translated.wav"))
|
||||
parser.add_argument("--base-url", default=os.getenv("LITELLM_BASE_URL", "http://localhost:4000"))
|
||||
parser.add_argument("--model", default=os.getenv("REALTIME_TRANSLATE_MODEL", "gpt-realtime-translate"))
|
||||
parser.add_argument("--target-language", default="fr")
|
||||
parser.add_argument("--trailing-silence", type=float, default=1.5)
|
||||
parsed = parser.parse_args(argv)
|
||||
api_key = os.getenv("LITELLM_API_KEY") or os.getenv("OPENAI_API_KEY")
|
||||
if not api_key:
|
||||
return "Set LITELLM_API_KEY or OPENAI_API_KEY before running the script"
|
||||
if parsed.trailing_silence < 0:
|
||||
return "--trailing-silence must be zero or greater"
|
||||
return Settings(
|
||||
input_wav=parsed.input_wav,
|
||||
output_wav=parsed.output,
|
||||
base_url=parsed.base_url,
|
||||
model=parsed.model,
|
||||
target_language=parsed.target_language,
|
||||
trailing_silence_seconds=parsed.trailing_silence,
|
||||
api_key=api_key,
|
||||
)
|
||||
|
||||
|
||||
def translation_url(base_url: str, model: str) -> str | None:
|
||||
parsed = urlsplit(base_url.rstrip("/"))
|
||||
scheme = {"http": "ws", "https": "wss", "ws": "ws", "wss": "wss"}.get(parsed.scheme)
|
||||
if not scheme or not parsed.netloc:
|
||||
return None
|
||||
base_path = parsed.path.rstrip("/")
|
||||
realtime_path = (
|
||||
f"{base_path}/realtime/translations" if base_path.endswith("/v1") else f"{base_path}/v1/realtime/translations"
|
||||
)
|
||||
return urlunsplit((scheme, parsed.netloc, realtime_path, urlencode({"model": model}), ""))
|
||||
|
||||
|
||||
def read_pcm16_wav(path: Path) -> bytes | str:
|
||||
try:
|
||||
with wave.open(str(path), "rb") as source:
|
||||
actual_format = (
|
||||
source.getnchannels(),
|
||||
source.getsampwidth(),
|
||||
source.getframerate(),
|
||||
source.getcomptype(),
|
||||
)
|
||||
expected_format = (CHANNELS, SAMPLE_WIDTH, SAMPLE_RATE, "NONE")
|
||||
if actual_format != expected_format:
|
||||
return (
|
||||
f"{path} must be mono, 16-bit PCM, 24 kHz WAV; received "
|
||||
f"channels={actual_format[0]}, sample_width={actual_format[1]}, "
|
||||
f"sample_rate={actual_format[2]}, compression={actual_format[3]}"
|
||||
)
|
||||
return source.readframes(source.getnframes())
|
||||
except (OSError, EOFError, wave.Error) as exc:
|
||||
return f"Unable to read {path}: {exc}"
|
||||
|
||||
|
||||
def audio_chunks(audio: bytes) -> Iterator[bytes]:
|
||||
return (audio[offset : offset + CHUNK_BYTES] for offset in range(0, len(audio), CHUNK_BYTES))
|
||||
|
||||
|
||||
def audio_message(audio: bytes) -> str:
|
||||
return json.dumps(
|
||||
{
|
||||
"type": "session.input_audio_buffer.append",
|
||||
"audio": base64.b64encode(audio).decode("ascii"),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
async def configure_session(connection: ClientConnection, target_language: str) -> str | None:
|
||||
await connection.send(
|
||||
json.dumps(
|
||||
{
|
||||
"type": "session.update",
|
||||
"session": {"audio": {"output": {"language": target_language}}},
|
||||
}
|
||||
)
|
||||
)
|
||||
while True:
|
||||
raw_event = await asyncio.wait_for(connection.recv(), timeout=20)
|
||||
event = json.loads(raw_event)
|
||||
event_type = event.get("type")
|
||||
if event_type == "session.created":
|
||||
write_stdout(f"Session: {event.get('session', {}).get('id', 'created')}")
|
||||
if event_type == "session.updated":
|
||||
return None
|
||||
if event_type == "error":
|
||||
return f"Realtime API error: {json.dumps(event.get('error', event), ensure_ascii=False)}"
|
||||
|
||||
|
||||
async def send_audio(
|
||||
connection: ClientConnection, pcm: bytes, trailing_silence_seconds: float, finished: asyncio.Event
|
||||
) -> None:
|
||||
silence = bytes(round(trailing_silence_seconds * SAMPLE_RATE * CHANNELS * SAMPLE_WIDTH))
|
||||
try:
|
||||
for chunk in audio_chunks(pcm + silence):
|
||||
await connection.send(audio_message(chunk))
|
||||
await asyncio.sleep(len(chunk) / (SAMPLE_RATE * CHANNELS * SAMPLE_WIDTH))
|
||||
await connection.send(json.dumps({"type": "session.close"}))
|
||||
finally:
|
||||
finished.set()
|
||||
|
||||
|
||||
async def receive_translation(
|
||||
connection: ClientConnection, output_path: Path, sender_finished: asyncio.Event
|
||||
) -> str | None:
|
||||
audio_received = asyncio.Event()
|
||||
try:
|
||||
with wave.open(str(output_path), "wb") as output:
|
||||
output.setnchannels(CHANNELS)
|
||||
output.setsampwidth(SAMPLE_WIDTH)
|
||||
output.setframerate(SAMPLE_RATE)
|
||||
write_stdout("Translation: ", end="", flush=True)
|
||||
while True:
|
||||
try:
|
||||
raw_event = await asyncio.wait_for(connection.recv(), timeout=RESPONSE_TIMEOUT_SECONDS)
|
||||
except TimeoutError:
|
||||
if not sender_finished.is_set():
|
||||
continue
|
||||
return (
|
||||
f"The translation stream went quiet for {RESPONSE_TIMEOUT_SECONDS:g} s without closing the session"
|
||||
)
|
||||
event = json.loads(raw_event)
|
||||
event_type = event.get("type")
|
||||
if event_type in AUDIO_EVENT_TYPES:
|
||||
output.writeframes(base64.b64decode(event.get("delta", ""), validate=True))
|
||||
audio_received.set()
|
||||
elif event_type in TRANSCRIPT_EVENT_TYPES:
|
||||
write_stdout(event.get("delta", event.get("text", "")), end="", flush=True)
|
||||
elif event_type == "error":
|
||||
return f"Realtime API error: {json.dumps(event.get('error', event), ensure_ascii=False)}"
|
||||
elif event_type == "session.closed":
|
||||
write_stdout()
|
||||
return None if audio_received.is_set() else "The translation stream ended without translated audio"
|
||||
except (OSError, wave.Error) as exc:
|
||||
return f"Unable to write {output_path}: {exc}"
|
||||
|
||||
|
||||
async def translate(settings: Settings, pcm: bytes) -> str | None:
|
||||
url = translation_url(settings.base_url, settings.model)
|
||||
if not url:
|
||||
return f"Invalid --base-url: {settings.base_url}"
|
||||
sender_finished = asyncio.Event()
|
||||
try:
|
||||
async with websockets.connect(
|
||||
url,
|
||||
additional_headers={"Authorization": f"Bearer {settings.api_key}"},
|
||||
proxy=None,
|
||||
open_timeout=20,
|
||||
close_timeout=5,
|
||||
) as connection:
|
||||
configuration_error = await configure_session(connection, settings.target_language)
|
||||
if configuration_error:
|
||||
return configuration_error
|
||||
async with asyncio.TaskGroup() as tasks:
|
||||
receiver = tasks.create_task(receive_translation(connection, settings.output_wav, sender_finished))
|
||||
tasks.create_task(send_audio(connection, pcm, settings.trailing_silence_seconds, sender_finished))
|
||||
return receiver.result()
|
||||
except Exception as exc:
|
||||
return f"Translation failed: {type(exc).__name__}: {exc}"
|
||||
|
||||
|
||||
def main(argv: Sequence[str] | None = None) -> int:
|
||||
settings = parse_args(argv)
|
||||
if isinstance(settings, str):
|
||||
write_stderr(settings)
|
||||
return 2
|
||||
pcm = read_pcm16_wav(settings.input_wav)
|
||||
if isinstance(pcm, str):
|
||||
write_stderr(pcm)
|
||||
return 2
|
||||
error = asyncio.run(translate(settings, pcm))
|
||||
if error:
|
||||
write_stderr(error)
|
||||
return 1
|
||||
write_stdout(f"Translated audio: {settings.output_wav}")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
60
cookbook/skill_scanner/README.md
Normal file
60
cookbook/skill_scanner/README.md
Normal file
|
|
@ -0,0 +1,60 @@
|
|||
# Manual skill security scans
|
||||
|
||||
This standalone adapter runs NVIDIA SkillSpector on a separate machine. It is not installed in the LiteLLM proxy and does not require a shared Helm or Compose deployment
|
||||
|
||||
Start it with Python 3.12 or newer and uv. The script pins SkillSpector and its HTTP dependencies
|
||||
|
||||
```sh
|
||||
export SKILL_SCANNER_API_KEY="$(openssl rand -hex 32)"
|
||||
uv run server.py
|
||||
```
|
||||
|
||||
For another host, place the loopback service behind an HTTPS reverse proxy that forwards the Authorization header and enforces a 65 MiB request limit. Run it under an unprivileged account with a writable temporary directory and no production files or credentials. The scanner host needs Git and network access to requested sources and optional model providers. Apply network egress restrictions suitable for those sources
|
||||
|
||||
Configure the same scanner secret in the gateway's environment, then add this to LiteLLM's configuration
|
||||
|
||||
```yaml
|
||||
general_settings:
|
||||
skill_scanner:
|
||||
endpoint: https://scanner.example.com/scan
|
||||
api_key_env: SKILL_SCANNER_API_KEY
|
||||
use_llm: false
|
||||
timeout_seconds: 180
|
||||
```
|
||||
|
||||
The endpoint must be reachable from the gateway. A customer's laptop works while it is reachable through their network, VPN or authenticated HTTPS tunnel. `localhost` refers to the gateway machine, not the admin's browser. HTTP is accepted only for loopback development
|
||||
|
||||
Open Skills, Submissions, and Review as a proxy admin. Click **Scan for threats**, inspect the report, then independently approve or reject the proposal. Scanning is optional and never changes approval permissions or blocks publication. The report shows findings, locations, remediation, risk score, scanner version and whether semantic analysis actually ran. Incomplete reports do not establish safety
|
||||
|
||||
Uploaded ZIPs are sent privately to the scanner with their exact stored bytes. Metadata-only edits reuse the matching stored archive. If that archive is missing or changed, submit a new ZIP before scanning. The displayed SHA-256 identifies that archive, and the scan examines its files rather than proposed catalog metadata. Remote URL reports describe the source at scan time and do not freeze its contents. GitHub branch/tag refs and explicit-ref subdirectories are supported. Other ref/subdirectory combinations require a ZIP so the adapter does not silently scan the wrong revision
|
||||
|
||||
The latest scan result is retained with the submission. Running scans have a deadline. If a gateway restarts or the scanner becomes unavailable, the attempt expires and an admin can retry. Retries replace the latest report. Approving or rejecting during a scan does not cancel it or change the decision
|
||||
|
||||
Set `use_llm: true` to request semantic analysis and configure the provider on the scanner machine using [SkillSpector's provider instructions](https://github.com/NVIDIA/skillspector#llm-analysis). Provider credentials stay on that machine. Static analysis can also call OSV for dependency checks
|
||||
|
||||
## HTTP contract
|
||||
|
||||
The gateway sends an authenticated `POST` to the configured endpoint. The body contains form fields `source` (JSON source descriptor), `use_llm` (boolean), and `expires_at` (timezone-aware ISO timestamp), plus `file` (ZIP) when bytes are available. The service returns HTTP 200 with this JSON shape
|
||||
|
||||
```json
|
||||
{
|
||||
"scanner_version": "2.12.0",
|
||||
"complete": true,
|
||||
"scan_mode": "static-only",
|
||||
"risk_score": 70,
|
||||
"findings": [{
|
||||
"rule_id": "example-rule",
|
||||
"severity": "HIGH",
|
||||
"message": "Potential data exfiltration",
|
||||
"file": "SKILL.md",
|
||||
"start_line": 8,
|
||||
"explanation": "Instructions request sending local files to an external host",
|
||||
"remediation": "Remove the external upload instruction"
|
||||
}],
|
||||
"limitations": []
|
||||
}
|
||||
```
|
||||
|
||||
Reports are limited to 1 MiB and 500 findings. Scanner uploads allow 64 MiB because the gateway can expand compressed skill uploads into stored compatibility ZIPs. The gateway rejects malformed reports, refuses redirects and verifies TLS. Only the configured scanner credential is sent. A timeout, HTTP error or incomplete report remains visible to the reviewer
|
||||
|
||||
The bundled adapter runs at most two child scans concurrently. Each child and its descendants are terminated at the gateway deadline, capped at 570 seconds. Keep the gateway and scanner clocks synchronized. Direct adapter clients may omit `expires_at` to use the cap. Child processes never execute submitted skill scripts. Raise the gateway deadline up to 600 seconds when needed for semantic analysis and configure the reverse proxy timeout to match
|
||||
172
cookbook/skill_scanner/server.py
Normal file
172
cookbook/skill_scanner/server.py
Normal file
|
|
@ -0,0 +1,172 @@
|
|||
# /// script
|
||||
# requires-python = ">=3.12"
|
||||
# dependencies = [
|
||||
# "fastapi==0.162.0", "uvicorn==0.54.0", "python-multipart==0.0.33",
|
||||
# "skillspector @ git+https://github.com/NVIDIA/skillspector.git@62c2fe8622f580ac292d4078c3b71f790194dbbd",
|
||||
# ]
|
||||
# ///
|
||||
import asyncio
|
||||
import hmac
|
||||
import json
|
||||
import os
|
||||
import signal
|
||||
import sys
|
||||
import tempfile
|
||||
from contextlib import suppress
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Annotated, Final
|
||||
from urllib.parse import quote, urlsplit
|
||||
|
||||
from fastapi import Depends, FastAPI, File, Form, HTTPException, UploadFile
|
||||
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
|
||||
from pydantic import AwareDatetime, TypeAdapter, ValidationError
|
||||
|
||||
MAX_UPLOAD: Final = 64 * 1024 * 1024
|
||||
MAX_REPORT: Final = 1024 * 1024
|
||||
SCAN_TIMEOUT: Final = 570
|
||||
SOURCE: Final = TypeAdapter(dict[str, str])
|
||||
SLOTS: Final = asyncio.Semaphore(2)
|
||||
app: Final = FastAPI(title="SkillSpector adapter", docs_url=None, redoc_url=None, openapi_url=None)
|
||||
|
||||
|
||||
def authenticate(credentials: Annotated[HTTPAuthorizationCredentials, Depends(HTTPBearer())]) -> None:
|
||||
token: Final = os.environ.get("SKILL_SCANNER_API_KEY", "")
|
||||
if not token or not hmac.compare_digest(credentials.credentials, token):
|
||||
raise HTTPException(401, "Invalid scanner credentials")
|
||||
|
||||
|
||||
def source_target(source: dict[str, str]) -> str:
|
||||
target: Final = (
|
||||
"https://github.com/" + source.get("repo", "") if source.get("source") == "github" else source.get("url", "")
|
||||
)
|
||||
parsed: Final = urlsplit(target)
|
||||
if parsed.scheme != "https" or not parsed.hostname or parsed.username or parsed.password or parsed.fragment:
|
||||
raise HTTPException(400, "Remote sources require an HTTPS URL without credentials or a fragment")
|
||||
ref: Final = source.get("ref")
|
||||
subdir: Final = source.get("path")
|
||||
if ref or subdir:
|
||||
if parsed.hostname != "github.com" or len(parsed.path.strip("/").split("/")) != 2 or not ref:
|
||||
raise HTTPException(400, "Upload a ZIP to scan this source's exact ref and subdirectory")
|
||||
suffix: Final = "/" + quote(subdir, safe="/") if subdir else ""
|
||||
return target.removesuffix(".git").rstrip("/") + "/tree/" + quote(ref, safe="/") + suffix
|
||||
return target
|
||||
|
||||
|
||||
async def worker(target: str, output: Path, use_llm: bool) -> None:
|
||||
from skillspector.mcp_server import run_scan
|
||||
|
||||
result: Final = await run_scan(target, use_llm=use_llm, allow_local_targets=not target.startswith("https://"))
|
||||
coverage: Final = result.get("analysis_completeness", {})
|
||||
findings: Final = result["findings"]
|
||||
complete: Final = (
|
||||
result.get("execution_successful") is True
|
||||
and coverage.get("is_complete") is True
|
||||
and (not use_llm or (result.get("llm_available") is True and result.get("llm_used") is True))
|
||||
and len(findings) <= 500
|
||||
)
|
||||
limitations: Final = tuple(str(value)[:2000] for value in coverage.get("limitations", ()))
|
||||
report: Final = {
|
||||
"scanner_version": result["version"],
|
||||
"complete": complete,
|
||||
"scan_mode": result["scan_mode"],
|
||||
"risk_score": result["risk_score"],
|
||||
"findings": [
|
||||
{
|
||||
"rule_id": finding["id"],
|
||||
"severity": finding["severity"],
|
||||
"message": finding.get("pattern") or finding.get("finding") or "Security finding",
|
||||
"file": finding["location"]["file"],
|
||||
"start_line": finding["location"].get("start_line"),
|
||||
"explanation": finding.get("explanation"),
|
||||
"remediation": finding.get("remediation"),
|
||||
}
|
||||
for finding in findings[:500]
|
||||
],
|
||||
"limitations": (
|
||||
*limitations,
|
||||
*(("Only the first 500 findings are displayed.",) if len(findings) > 500 else ()),
|
||||
),
|
||||
}
|
||||
encoded: Final = json.dumps(report).encode()
|
||||
if len(encoded) > MAX_REPORT:
|
||||
raise ValueError("Scanner report exceeded the response limit")
|
||||
output.write_bytes(encoded)
|
||||
|
||||
|
||||
async def run_worker(target: str, output: Path, use_llm: bool, expires_at: datetime | None = None) -> dict[str, object]:
|
||||
timeout: Final = (
|
||||
SCAN_TIMEOUT
|
||||
if expires_at is None
|
||||
else min(SCAN_TIMEOUT, (expires_at - datetime.now(timezone.utc)).total_seconds())
|
||||
)
|
||||
if timeout <= 0:
|
||||
raise HTTPException(504, "SkillSpector scan deadline expired")
|
||||
process: Final = await asyncio.create_subprocess_exec(
|
||||
sys.executable,
|
||||
str(Path(__file__).resolve()),
|
||||
"--scan",
|
||||
target,
|
||||
str(output),
|
||||
str(use_llm).lower(),
|
||||
stdout=asyncio.subprocess.DEVNULL,
|
||||
stderr=asyncio.subprocess.DEVNULL,
|
||||
start_new_session=True,
|
||||
env={
|
||||
**os.environ,
|
||||
"HOME": str(output.parent),
|
||||
"TMPDIR": str(output.parent),
|
||||
"GIT_CONFIG_GLOBAL": os.devnull,
|
||||
"GIT_CONFIG_NOSYSTEM": "1",
|
||||
"GIT_TERMINAL_PROMPT": "0",
|
||||
"GIT_ALLOW_PROTOCOL": "https",
|
||||
},
|
||||
)
|
||||
try:
|
||||
code: Final = await asyncio.wait_for(process.wait(), timeout)
|
||||
if code != 0 or not output.is_file() or output.stat().st_size > MAX_REPORT:
|
||||
raise HTTPException(502, "SkillSpector could not produce a report for this source")
|
||||
return TypeAdapter(dict[str, object]).validate_json(output.read_bytes())
|
||||
except TimeoutError:
|
||||
raise HTTPException(504, "SkillSpector scan timed out") from None
|
||||
finally:
|
||||
with suppress(ProcessLookupError):
|
||||
os.killpg(process.pid, signal.SIGKILL)
|
||||
await process.wait()
|
||||
|
||||
|
||||
@app.post("/scan", dependencies=[Depends(authenticate)])
|
||||
async def scan(
|
||||
source: Annotated[str, Form(max_length=16384)],
|
||||
use_llm: Annotated[bool, Form()] = False,
|
||||
expires_at: Annotated[AwareDatetime | None, Form()] = None,
|
||||
file: Annotated[UploadFile | None, File()] = None,
|
||||
) -> dict[str, object]:
|
||||
if SLOTS.locked():
|
||||
raise HTTPException(429, "Scanner is busy; try again shortly")
|
||||
try:
|
||||
parsed: Final = SOURCE.validate_json(source)
|
||||
except ValidationError:
|
||||
raise HTTPException(400, "Invalid source metadata") from None
|
||||
async with SLOTS:
|
||||
with tempfile.TemporaryDirectory(prefix="skill-scan-") as directory:
|
||||
root: Final = Path(directory)
|
||||
if file is not None:
|
||||
content: Final = await file.read(MAX_UPLOAD + 1)
|
||||
if len(content) > MAX_UPLOAD:
|
||||
raise HTTPException(413, "Upload exceeds 64 MiB")
|
||||
archive: Final = root / "skill.zip"
|
||||
archive.write_bytes(content)
|
||||
return await run_worker(str(archive), root / "report.json", use_llm, expires_at)
|
||||
return await run_worker(source_target(parsed), root / "report.json", use_llm, expires_at)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
if len(sys.argv) > 1 and sys.argv[1] == "--scan":
|
||||
asyncio.run(worker(sys.argv[2], Path(sys.argv[3]), sys.argv[4] == "true"))
|
||||
else:
|
||||
import uvicorn
|
||||
|
||||
if not os.environ.get("SKILL_SCANNER_API_KEY"):
|
||||
raise SystemExit("Set SKILL_SCANNER_API_KEY before starting the scanner")
|
||||
uvicorn.run(app, host="127.0.0.1", port=int(os.environ.get("PORT", "8099")))
|
||||
|
|
@ -74,6 +74,7 @@ ENV UV_PROJECT_ENVIRONMENT=/app/.venv \
|
|||
COPY pyproject.toml uv.lock ./
|
||||
COPY enterprise/pyproject.toml enterprise/
|
||||
COPY litellm-proxy-extras/pyproject.toml litellm-proxy-extras/
|
||||
COPY packaging/litellm-proxy/pyproject.toml packaging/litellm-proxy/
|
||||
|
||||
# Install third-party dependencies (cached unless pyproject.toml/uv.lock change)
|
||||
RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-groups --no-editable \
|
||||
|
|
@ -99,7 +100,8 @@ COPY --from=ui-builder /ui/out/. litellm/proxy/_experimental/out/
|
|||
RUN sed -i 's/\r$//' docker/build_admin_ui.sh && chmod +x docker/build_admin_ui.sh && ./docker/build_admin_ui.sh
|
||||
|
||||
# Install project and workspace packages (fast - deps already cached)
|
||||
RUN uv sync --frozen --no-default-groups --no-editable \
|
||||
RUN uv build --package litellm-proxy --wheel --out-dir /wheels
|
||||
RUN uv sync --frozen --no-install-package litellm-proxy --no-default-groups --no-editable \
|
||||
--extra proxy \
|
||||
--extra proxy-runtime \
|
||||
--group admin-mcp \
|
||||
|
|
@ -109,6 +111,8 @@ RUN uv sync --frozen --no-default-groups --no-editable \
|
|||
--extra bedrock-realtime \
|
||||
--python python3.13
|
||||
|
||||
RUN uv pip install --python /app/.venv/bin/python --no-deps /wheels/litellm_proxy-*.whl
|
||||
|
||||
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
|
||||
npm_config_cache=/root/.npm \
|
||||
prisma generate --schema=./schema.prisma
|
||||
|
|
|
|||
|
|
@ -3,9 +3,11 @@
|
|||
|
||||
import asyncio
|
||||
import base64
|
||||
import copy
|
||||
import json
|
||||
import time
|
||||
from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta
|
||||
from itertools import chain
|
||||
from types import MappingProxyType
|
||||
|
|
@ -34,7 +36,7 @@ from typing_extensions import ReadOnly, Unpack
|
|||
|
||||
import litellm
|
||||
from litellm import Router, verbose_logger
|
||||
from litellm._internal_context import with_service_target
|
||||
from litellm._internal_context import internal_sub_call, post_response_phase, with_service_target
|
||||
from litellm._uuid import uuid
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.constants import (
|
||||
|
|
@ -95,6 +97,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.streaming_iterator import spawn_logging_task
|
||||
from litellm.responses.utils import ResponsesAPIRequestUtils
|
||||
from litellm.types.llms.openai import ( # pyright: ignore[reportAttributeAccessIssue]
|
||||
AllMessageValues,
|
||||
|
|
@ -395,6 +398,134 @@ class _RouterFileCallKwargs(TypedDict, total=False):
|
|||
custom_llm_provider: ReadOnly[str | None]
|
||||
|
||||
|
||||
_MODEL_FILE_ID_MAPPING: Final = TypeAdapter(dict[str, str])
|
||||
_FILE_UPLOAD_ONLY_FIELDS: Final = frozenset({"file", "purpose", "expires_after", "extra_body"})
|
||||
_FILE_DELETE_RESERVED_FIELDS: Final = frozenset({"model", "file_id", "_litellm_internal_model_credentials"})
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _ProviderUpload:
|
||||
model_id: str
|
||||
provider_file_id: str
|
||||
upload_call_id: str | None = None
|
||||
|
||||
|
||||
def _provider_uploads(file_object: OpenAIFileObject) -> tuple[_ProviderUpload, ...]:
|
||||
file_hidden_params: Final = cast( # cast-ok: preserve mapping operations on dynamic file metadata
|
||||
dict[str, object], getattr(file_object, HIDDEN_PARAMS_ATTR)
|
||||
)
|
||||
model_file_id_mapping: Final = _MODEL_FILE_ID_MAPPING.validate_python(
|
||||
file_hidden_params.get("model_file_id_mapping") or {}
|
||||
)
|
||||
upload_call_id: Final = file_hidden_params.get("litellm_call_id")
|
||||
return tuple(
|
||||
_ProviderUpload(
|
||||
model_id=model_id,
|
||||
provider_file_id=provider_file_id,
|
||||
upload_call_id=upload_call_id if isinstance(upload_call_id, str) else None,
|
||||
)
|
||||
for model_id, provider_file_id in model_file_id_mapping.items()
|
||||
)
|
||||
|
||||
|
||||
def _keeps_one_upload_per_target(llm_router: Router, model_id: str) -> bool:
|
||||
credentials: Final = _provider_file_retrieve_credentials(llm_router=llm_router, model_id=model_id)
|
||||
return (credentials or {}).get("custom_llm_provider") == LlmProviders.LITELLM_PROXY.value
|
||||
|
||||
|
||||
def _is_kept_upload(uploads: Sequence[_ProviderUpload], index: int, per_target_model_ids: frozenset[str]) -> bool:
|
||||
upload: Final = uploads[index]
|
||||
other_uploads: Final = uploads[index + 1 :] if upload.model_id in per_target_model_ids else uploads[:index]
|
||||
return all(other.model_id != upload.model_id for other in other_uploads)
|
||||
|
||||
|
||||
def _upload_from_trailing_id_fields(decoded_unified_file_id: str) -> _ProviderUpload:
|
||||
trailing_fields: Final = decoded_unified_file_id.rpartition(";llm_output_file_id,")[2]
|
||||
provider_file_id, _, model_id = trailing_fields.partition(";llm_output_file_model_id,")
|
||||
return _ProviderUpload(model_id=model_id, provider_file_id=provider_file_id)
|
||||
|
||||
|
||||
def _carried_upload_missing_from(unified_file_id: str, model_file_ids: Mapping[str, str]) -> _ProviderUpload | None:
|
||||
decoded: Final = _is_base64_encoded_unified_file_id(unified_file_id)
|
||||
carried: Final = _upload_from_trailing_id_fields(decoded) if decoded else None
|
||||
if carried is None or model_file_ids.get(carried.model_id) in (None, carried.provider_file_id):
|
||||
return None
|
||||
return carried
|
||||
|
||||
|
||||
async def _delete_provider_file(
|
||||
llm_router: Router,
|
||||
model_id: str,
|
||||
provider_file_id: str,
|
||||
request_data: Mapping[str, object],
|
||||
) -> None:
|
||||
credentials: Final = llm_router.registry().credentials(model_id, form="with_provider")
|
||||
delete_data: Final = {
|
||||
**{k: v for k, v in request_data.items() if k not in _FILE_DELETE_RESERVED_FIELDS},
|
||||
**(
|
||||
{"_litellm_internal_model_credentials": MappingProxyType(dict(credentials))}
|
||||
if credentials is not None
|
||||
else {}
|
||||
),
|
||||
}
|
||||
await llm_router.afile_delete(
|
||||
model=model_id,
|
||||
file_id=provider_file_id,
|
||||
**cast(_RouterFileCallKwargs, delete_data),
|
||||
)
|
||||
|
||||
|
||||
def _with_own_nested_dicts(request_data: Mapping[str, object]) -> MappingProxyType[str, object]:
|
||||
return MappingProxyType({k: copy.copy(v) if isinstance(v, dict) else v for k, v in request_data.items()})
|
||||
|
||||
|
||||
async def _delete_redundant_upload(
|
||||
llm_router: Router,
|
||||
upload: _ProviderUpload,
|
||||
delete_request: Mapping[str, object],
|
||||
parent_call_id: str | None,
|
||||
) -> None:
|
||||
try:
|
||||
with internal_sub_call(parent_call_id):
|
||||
await _delete_provider_file(
|
||||
llm_router, upload.model_id, upload.provider_file_id, _with_own_nested_dicts(delete_request)
|
||||
)
|
||||
except Exception as error:
|
||||
verbose_logger.warning(
|
||||
f"Could not delete the redundant upload provider_file_id={sanitize_for_log(upload.provider_file_id)} "
|
||||
f"at deployment model_id={sanitize_for_log(upload.model_id)}: "
|
||||
f"{type(error).__name__} {sanitize_for_log(error)}"
|
||||
)
|
||||
|
||||
|
||||
async def _delete_redundant_uploads(
|
||||
llm_router: Router,
|
||||
uploads: Sequence[_ProviderUpload],
|
||||
delete_request: Mapping[str, object],
|
||||
) -> None:
|
||||
await asyncio.gather(
|
||||
*(_delete_redundant_upload(llm_router, upload, delete_request, upload.upload_call_id) for upload in uploads)
|
||||
)
|
||||
|
||||
|
||||
def _start_deleting_redundant_uploads(
|
||||
llm_router: Router,
|
||||
uploads: Sequence[_ProviderUpload],
|
||||
create_file_request: CreateFileRequest,
|
||||
) -> None:
|
||||
if not uploads:
|
||||
return
|
||||
delete_request: Final = MappingProxyType(
|
||||
{k: v for k, v in create_file_request.items() if k not in _FILE_UPLOAD_ONLY_FIELDS}
|
||||
)
|
||||
with post_response_phase():
|
||||
spawn_logging_task(
|
||||
asyncio.get_running_loop(),
|
||||
_delete_redundant_uploads(llm_router, uploads, delete_request),
|
||||
task_name="managed file repeat upload cleanup",
|
||||
)
|
||||
|
||||
|
||||
def _managed_database(prisma_client: _ManagedResourcePrismaClient) -> _ManagedResourceDatabase:
|
||||
return prisma_client.db
|
||||
|
||||
|
|
@ -2034,16 +2165,16 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
target_model_names_list=target_model_names_list,
|
||||
)
|
||||
|
||||
## STORE MODEL MAPPINGS IN DB
|
||||
model_mappings: Dict[str, str] = {}
|
||||
|
||||
for file_object in responses:
|
||||
file_hidden_params = cast( # cast-ok: preserve mapping operations on dynamic file metadata
|
||||
dict[str, object], getattr(file_object, HIDDEN_PARAMS_ATTR)
|
||||
)
|
||||
model_file_id_mapping = file_hidden_params.get("model_file_id_mapping")
|
||||
if model_file_id_mapping and isinstance(model_file_id_mapping, dict):
|
||||
model_mappings.update(model_file_id_mapping)
|
||||
uploads: Final = tuple(chain.from_iterable(map(_provider_uploads, responses)))
|
||||
per_target_model_ids: Final = frozenset(
|
||||
model_id
|
||||
for model_id in {upload.model_id for upload in uploads}
|
||||
if _keeps_one_upload_per_target(llm_router, model_id)
|
||||
)
|
||||
kept_uploads: Final = tuple(
|
||||
upload for index, upload in enumerate(uploads) if _is_kept_upload(uploads, index, per_target_model_ids)
|
||||
)
|
||||
model_mappings: Final = {upload.model_id: upload.provider_file_id for upload in kept_uploads}
|
||||
|
||||
await self.store_unified_file_id(
|
||||
file_id=response.id,
|
||||
|
|
@ -2052,6 +2183,15 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
model_mappings=model_mappings,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
_start_deleting_redundant_uploads(
|
||||
llm_router,
|
||||
tuple(
|
||||
upload
|
||||
for upload in uploads
|
||||
if upload not in kept_uploads and upload.model_id not in per_target_model_ids
|
||||
),
|
||||
create_file_request,
|
||||
)
|
||||
|
||||
# Emit Prometheus metrics for managed file creation
|
||||
prom_logger = self._get_prometheus_logger()
|
||||
|
|
@ -2692,23 +2832,13 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
specific_model_file_id_mapping: Final = model_file_id_mapping.get(file_id)
|
||||
if not specific_model_file_id_mapping:
|
||||
return
|
||||
filtered_data: Final = {
|
||||
k: v for k, v in data.items() if k not in ("model", "file_id", "_litellm_internal_model_credentials")
|
||||
}
|
||||
for model_id, model_file_id in specific_model_file_id_mapping.items():
|
||||
credentials = llm_router.registry().credentials(model_id, form="with_provider")
|
||||
delete_data = {
|
||||
**filtered_data,
|
||||
**(
|
||||
{"_litellm_internal_model_credentials": MappingProxyType(dict(credentials))}
|
||||
if credentials is not None
|
||||
else {}
|
||||
),
|
||||
}
|
||||
await llm_router.afile_delete(
|
||||
model=model_id,
|
||||
file_id=model_file_id,
|
||||
**cast(_RouterFileCallKwargs, delete_data),
|
||||
await _delete_provider_file(llm_router, model_id, model_file_id, data)
|
||||
unrecorded_upload: Final = _carried_upload_missing_from(file_id, specific_model_file_id_mapping)
|
||||
delete_call_id: Final = data.get("litellm_call_id")
|
||||
if unrecorded_upload is not None:
|
||||
await _delete_redundant_upload(
|
||||
llm_router, unrecorded_upload, data, delete_call_id if isinstance(delete_call_id, str) else None
|
||||
)
|
||||
|
||||
async def afile_content(
|
||||
|
|
|
|||
|
|
@ -1,9 +1,9 @@
|
|||
import enum
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
from pydantic import Field
|
||||
|
||||
from litellm.proxy._types import WebhookEvent
|
||||
from litellm.types.integrations.model_deprecation_email import MODEL_DEPRECATION_EMAIL_EVENT
|
||||
from litellm.types.llms.base import LiteLLMBaseModel
|
||||
|
||||
|
||||
|
|
@ -26,7 +26,7 @@ class SendKeyCreatedEmailEvent(WebhookEvent):
|
|||
|
||||
class SendKeyRotatedEmailEvent(WebhookEvent):
|
||||
virtual_key: str
|
||||
key_alias: Optional[str] = None
|
||||
key_alias: str | None = None
|
||||
"""
|
||||
The virtual key that was rotated
|
||||
this will be sk-123xxx, since we will be emailing this to the user to start using the new key
|
||||
|
|
@ -34,23 +34,33 @@ class SendKeyRotatedEmailEvent(WebhookEvent):
|
|||
|
||||
|
||||
class EmailEvent(str, enum.Enum):
|
||||
model_deprecations = MODEL_DEPRECATION_EMAIL_EVENT
|
||||
virtual_key_created = "Virtual Key Created"
|
||||
new_user_invitation = "New User Invitation"
|
||||
virtual_key_rotated = "Virtual Key Rotated"
|
||||
soft_budget_crossed = "Soft Budget Crossed"
|
||||
max_budget_alert = "Max Budget Alert"
|
||||
|
||||
|
||||
class EmailEventSettings(LiteLLMBaseModel):
|
||||
event: EmailEvent
|
||||
enabled: bool
|
||||
|
||||
|
||||
class EmailEventSettingsUpdateRequest(LiteLLMBaseModel):
|
||||
settings: List[EmailEventSettings]
|
||||
settings: list[EmailEventSettings]
|
||||
|
||||
|
||||
class EmailEventSettingsResponse(LiteLLMBaseModel):
|
||||
settings: List[EmailEventSettings]
|
||||
settings: list[EmailEventSettings]
|
||||
|
||||
|
||||
class DefaultEmailSettings(LiteLLMBaseModel):
|
||||
"""Default settings for email events"""
|
||||
settings: Dict[EmailEvent, bool] = Field(
|
||||
|
||||
settings: dict[EmailEvent, bool] = Field(
|
||||
default_factory=lambda: {
|
||||
EmailEvent.model_deprecations: False, # Explicit opt-in for owner retirement reminders.
|
||||
EmailEvent.virtual_key_created: True, # On by default
|
||||
EmailEvent.new_user_invitation: True, # On by default
|
||||
EmailEvent.virtual_key_rotated: True, # On by default
|
||||
|
|
@ -58,10 +68,12 @@ class DefaultEmailSettings(LiteLLMBaseModel):
|
|||
EmailEvent.max_budget_alert: True, # On by default
|
||||
}
|
||||
)
|
||||
def to_dict(self) -> Dict[str, bool]:
|
||||
|
||||
def to_dict(self) -> dict[str, bool]:
|
||||
"""Convert to dictionary with string keys for storage"""
|
||||
return {event.value: enabled for event, enabled in self.settings.items()}
|
||||
|
||||
@classmethod
|
||||
def get_defaults(cls) -> Dict[str, bool]:
|
||||
def get_defaults(cls) -> dict[str, bool]:
|
||||
"""Get the default settings as a dictionary with string keys"""
|
||||
return cls().to_dict()
|
||||
return cls().to_dict()
|
||||
|
|
|
|||
|
|
@ -41,10 +41,14 @@ COPY pyproject.toml README.md rust-toolchain.toml ./
|
|||
COPY .cargo ./.cargo
|
||||
COPY litellm-rust ./litellm-rust
|
||||
COPY model_prices_and_context_window.json ./
|
||||
COPY packaging/litellm-proxy/pyproject.toml packaging/litellm-proxy/
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
mkdir -p litellm/rust_bridge /wheels && \
|
||||
touch litellm/__init__.py litellm/rust_bridge/__init__.py && \
|
||||
uvx --python python3.13 --from maturin==1.15.0 --with puccinialin==0.1.14 \
|
||||
python -c "import maturin; maturin.build_wheel('/wheels')" && \
|
||||
cd packaging/litellm-proxy && \
|
||||
uvx --python python3.13 --from maturin==1.15.0 --with puccinialin==0.1.14 \
|
||||
python -c "import maturin; maturin.build_wheel('/wheels')"
|
||||
|
||||
|
|
@ -85,6 +89,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
|||
--mount=type=bind,source=uv.lock,target=uv.lock \
|
||||
--mount=type=bind,source=enterprise/pyproject.toml,target=enterprise/pyproject.toml \
|
||||
--mount=type=bind,source=litellm-proxy-extras/pyproject.toml,target=litellm-proxy-extras/pyproject.toml \
|
||||
--mount=type=bind,source=packaging/litellm-proxy/pyproject.toml,target=packaging/litellm-proxy/pyproject.toml \
|
||||
uv sync --frozen --no-install-project --no-install-workspace --no-default-groups --no-editable \
|
||||
--extra proxy \
|
||||
--extra proxy-runtime \
|
||||
|
|
@ -98,7 +103,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
|||
COPY . .
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv sync --frozen --no-install-project --no-default-groups --no-editable \
|
||||
uv sync --frozen --no-install-project --no-install-package litellm-proxy --no-default-groups --no-editable \
|
||||
--extra proxy \
|
||||
--extra proxy-runtime \
|
||||
--group admin-mcp \
|
||||
|
|
@ -108,7 +113,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
|||
--python python3.13
|
||||
|
||||
RUN --mount=type=bind,from=rust-native,source=/wheels,target=/wheels \
|
||||
uv pip install --python /app/.venv/bin/python --no-deps /wheels/litellm-*.whl
|
||||
uv pip install --python /app/.venv/bin/python --no-deps /wheels/litellm-*.whl /wheels/litellm_proxy-*.whl
|
||||
|
||||
# PYTHONPATH=/app makes the source tree shadow the installed package, so the
|
||||
# compiled Rust extension must live next to the source or it is never imported.
|
||||
|
|
|
|||
|
|
@ -18,7 +18,7 @@ type: application
|
|||
# This is the chart version. This version number should be incremented each time you make changes
|
||||
# to the chart and its templates, including the app version.
|
||||
# Versions are expected to follow Semantic Versioning (https://semver.org/)
|
||||
version: 1.1.3
|
||||
version: 1.1.4
|
||||
|
||||
# This is the version number of the application being deployed. This version number should be
|
||||
# incremented each time you make changes to the application. Versions are not expected to
|
||||
|
|
|
|||
|
|
@ -56,8 +56,8 @@ If `db.useStackgresOperator` is used (not yet implemented):
|
|||
|
||||
| `billingMetrics.enabled` | Enable enterprise billable-request metering. Requires an enterprise license. | `false` |
|
||||
| `billingMetrics.endpoint` | Collector that the billable-request counter is pushed to. | `https://telemetry.litellm.ai` |
|
||||
| `billingMetrics.secretName` | Name of an existing Secret holding the mTLS client certificate, under the keys `tls.crt` and `tls.key`. | `litellm-billing-metrics-mtls` |
|
||||
| `billingMetrics.caSecretName` | Name of an existing Secret holding a CA bundle under the key `ca.crt`. Only needed for a private or test collector whose server certificate is not on the public web PKI. | `""` |
|
||||
| `billingMetrics.tokenUrl` | Cognito OAuth2 client-credentials token endpoint. | `https://litellm-prod.auth.us-west-1.amazoncognito.com/oauth2/token` |
|
||||
| `billingMetrics.secretName` | Name of an existing Secret with `client-id` and `client-secret` keys. | `litellm-billing-metrics` |
|
||||
| `billingMetrics.exportIntervalMs` | How often the counter is pushed, in milliseconds. The proxy defaults to `60000` when unset. | `""` |
|
||||
|
||||
#### Example `proxy_config` ConfigMap from values (default):
|
||||
|
|
@ -109,10 +109,10 @@ type: Opaque
|
|||
|
||||
#### Enterprise billable-request metering
|
||||
|
||||
Enterprise licenses meter billable requests by pushing a counter to LiteLLM's collector over mutual TLS. The chart does not create the client certificate; it mounts one you already hold, read-only, so the private key is never exposed through the environment. Create the Secret under the name the chart expects, then turn the block on:
|
||||
Enterprise licenses meter billable requests by pushing a counter to LiteLLM's collector with Cognito bearer authentication. Create a Secret with `client-id` and `client-secret` keys, then enable the chart values:
|
||||
|
||||
```
|
||||
kubectl create secret tls litellm-billing-metrics-mtls --cert=client.crt --key=client.key
|
||||
kubectl create secret generic litellm-billing-metrics --from-literal=client-id=... --from-literal=client-secret=...
|
||||
```
|
||||
|
||||
```
|
||||
|
|
@ -120,7 +120,8 @@ billingMetrics:
|
|||
enabled: true
|
||||
```
|
||||
|
||||
Set `billingMetrics.caSecretName` only when the collector is a private or test one whose server certificate is not on the public web PKI; the production collector needs no CA override. The chart fails the render rather than deploying a proxy that silently never exports, so a missing `secretName` or an emptied `endpoint` surfaces at `helm install` time.
|
||||
The endpoint and token URL must use HTTPS. The chart requires both URLs and
|
||||
the Secret name when billing metrics are enabled
|
||||
|
||||
### Database Settings
|
||||
|
||||
|
|
|
|||
|
|
@ -50,53 +50,27 @@ app.kubernetes.io/name: {{ include "litellm.name" . }}
|
|||
app.kubernetes.io/instance: {{ .Release.Name }}
|
||||
{{- end }}
|
||||
|
||||
{{/*
|
||||
Enterprise billable-request metering. The client certificate identifies the
|
||||
deployment to LiteLLM's collector, so it is mounted read-only from an existing
|
||||
Secret rather than passed through the environment.
|
||||
*/}}
|
||||
{{- define "litellm.billingMetrics.certDir" -}}/etc/litellm/billing-mtls{{- end -}}
|
||||
{{- define "litellm.billingMetrics.caDir" -}}/etc/litellm/billing-mtls-ca{{- end -}}
|
||||
|
||||
{{- define "litellm.billingMetricsEnv" -}}
|
||||
- name: LITELLM_BILLING_METRICS_ENDPOINT
|
||||
value: {{ required "billingMetrics.endpoint is required when billingMetrics.enabled is true" .Values.billingMetrics.endpoint | quote }}
|
||||
- name: LITELLM_BILLING_METRICS_CLIENT_CERT
|
||||
value: {{ printf "%s/tls.crt" (include "litellm.billingMetrics.certDir" .) | quote }}
|
||||
- name: LITELLM_BILLING_METRICS_CLIENT_KEY
|
||||
value: {{ printf "%s/tls.key" (include "litellm.billingMetrics.certDir" .) | quote }}
|
||||
{{- if .Values.billingMetrics.caSecretName }}
|
||||
- name: LITELLM_BILLING_METRICS_CA_CERT
|
||||
value: {{ printf "%s/ca.crt" (include "litellm.billingMetrics.caDir" .) | quote }}
|
||||
{{- end }}
|
||||
- name: LITELLM_BILLING_METRICS_TOKEN_URL
|
||||
value: {{ required "billingMetrics.tokenUrl is required when billingMetrics.enabled is true" .Values.billingMetrics.tokenUrl | quote }}
|
||||
- name: LITELLM_BILLING_METRICS_CLIENT_ID
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: {{ required "billingMetrics.secretName is required when billingMetrics.enabled is true" .Values.billingMetrics.secretName }}
|
||||
key: client-id
|
||||
- name: LITELLM_BILLING_METRICS_CLIENT_SECRET
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: {{ required "billingMetrics.secretName is required when billingMetrics.enabled is true" .Values.billingMetrics.secretName }}
|
||||
key: client-secret
|
||||
{{- with .Values.billingMetrics.exportIntervalMs }}
|
||||
- name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS
|
||||
value: {{ . | quote }}
|
||||
{{- end }}
|
||||
{{- end -}}
|
||||
|
||||
{{- define "litellm.billingMetricsVolumes" -}}
|
||||
- name: billing-metrics-mtls
|
||||
secret:
|
||||
secretName: {{ required "billingMetrics.secretName is required when billingMetrics.enabled is true (an existing Secret with tls.crt and tls.key)" .Values.billingMetrics.secretName }}
|
||||
{{- if .Values.billingMetrics.caSecretName }}
|
||||
- name: billing-metrics-mtls-ca
|
||||
secret:
|
||||
secretName: {{ .Values.billingMetrics.caSecretName }}
|
||||
{{- end }}
|
||||
{{- end -}}
|
||||
|
||||
{{- define "litellm.billingMetricsVolumeMounts" -}}
|
||||
- name: billing-metrics-mtls
|
||||
mountPath: {{ include "litellm.billingMetrics.certDir" . }}
|
||||
readOnly: true
|
||||
{{- if .Values.billingMetrics.caSecretName }}
|
||||
- name: billing-metrics-mtls-ca
|
||||
mountPath: {{ include "litellm.billingMetrics.caDir" . }}
|
||||
readOnly: true
|
||||
{{- end }}
|
||||
{{- end -}}
|
||||
|
||||
{{/*
|
||||
Create the name of the service account to use
|
||||
*/}}
|
||||
|
|
|
|||
|
|
@ -148,9 +148,6 @@ spec:
|
|||
- name: npm
|
||||
mountPath: /.npm
|
||||
{{- end }}
|
||||
{{- if .Values.billingMetrics.enabled }}
|
||||
{{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }}
|
||||
{{- end }}
|
||||
{{- if include "litellm.collector.socketDir" . }}
|
||||
- name: collector-socket
|
||||
mountPath: {{ include "litellm.collector.socketDir" . }}
|
||||
|
|
@ -234,9 +231,6 @@ spec:
|
|||
items:
|
||||
- key: {{ .Values.proxyConfigMap.key | default "config.yaml" }}
|
||||
path: "config.yaml"
|
||||
{{- if .Values.billingMetrics.enabled }}
|
||||
{{- include "litellm.billingMetricsVolumes" . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- if include "litellm.collector.socketDir" . }}
|
||||
- name: collector-socket
|
||||
emptyDir:
|
||||
|
|
|
|||
|
|
@ -1,31 +1,21 @@
|
|||
suite: test billingMetrics wiring on the proxy deployment
|
||||
suite: test bearer-only billingMetrics wiring
|
||||
templates:
|
||||
- deployment.yaml
|
||||
- configmap-litellm.yaml
|
||||
- migrations-job.yaml
|
||||
tests:
|
||||
- it: is off by default, adding no env, volume, or mount
|
||||
- it: is disabled by default
|
||||
template: deployment.yaml
|
||||
asserts:
|
||||
- notContains:
|
||||
path: spec.template.spec.volumes
|
||||
content:
|
||||
name: billing-metrics-mtls
|
||||
secret:
|
||||
secretName: litellm-billing-metrics-mtls
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].volumeMounts
|
||||
content:
|
||||
name: billing-metrics-mtls
|
||||
mountPath: /etc/litellm/billing-mtls
|
||||
readOnly: true
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_BILLING_METRICS_ENDPOINT
|
||||
value: https://telemetry.litellm.ai
|
||||
- notContains:
|
||||
path: spec.template.spec.volumes
|
||||
content:
|
||||
name: billing-metrics
|
||||
|
||||
- it: renders the endpoint and the mounted cert paths when enabled
|
||||
- it: renders required URLs and Cognito secret references when enabled
|
||||
template: deployment.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
|
|
@ -39,119 +29,34 @@ tests:
|
|||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_BILLING_METRICS_CLIENT_CERT
|
||||
value: /etc/litellm/billing-mtls/tls.crt
|
||||
name: LITELLM_BILLING_METRICS_TOKEN_URL
|
||||
value: https://litellm-prod.auth.us-west-1.amazoncognito.com/oauth2/token
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_BILLING_METRICS_CLIENT_KEY
|
||||
value: /etc/litellm/billing-mtls/tls.key
|
||||
|
||||
# The conventional Secret name is the default, so enabling the block is enough.
|
||||
- it: mounts the default cert secret read-only alongside the config volume
|
||||
template: deployment.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.volumes
|
||||
content:
|
||||
name: billing-metrics-mtls
|
||||
secret:
|
||||
secretName: litellm-billing-metrics-mtls
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].volumeMounts
|
||||
content:
|
||||
name: billing-metrics-mtls
|
||||
mountPath: /etc/litellm/billing-mtls
|
||||
readOnly: true
|
||||
|
||||
- it: honours a secretName override
|
||||
template: deployment.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
secretName: my-billing-mtls
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.volumes
|
||||
content:
|
||||
name: billing-metrics-mtls
|
||||
secret:
|
||||
secretName: my-billing-mtls
|
||||
- notContains:
|
||||
path: spec.template.spec.volumes
|
||||
content:
|
||||
name: billing-metrics-mtls
|
||||
secret:
|
||||
secretName: litellm-billing-metrics-mtls
|
||||
|
||||
- it: honours an endpoint override
|
||||
template: deployment.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
endpoint: https://collector.internal:4318
|
||||
asserts:
|
||||
name: LITELLM_BILLING_METRICS_CLIENT_ID
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: litellm-billing-metrics
|
||||
key: client-id
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_BILLING_METRICS_ENDPOINT
|
||||
value: https://collector.internal:4318
|
||||
|
||||
# The production collector presents a public web-PKI certificate, so the CA
|
||||
# override must stay absent unless a private collector is configured.
|
||||
- it: omits the CA env, volume, and mount when no caSecretName is set
|
||||
template: deployment.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
asserts:
|
||||
name: LITELLM_BILLING_METRICS_CLIENT_SECRET
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: litellm-billing-metrics
|
||||
key: client-secret
|
||||
- notContains:
|
||||
path: spec.template.spec.volumes
|
||||
content:
|
||||
name: billing-metrics-mtls-ca
|
||||
secret:
|
||||
secretName: billing-ca
|
||||
name: billing-metrics
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].volumeMounts
|
||||
content:
|
||||
name: billing-metrics-mtls-ca
|
||||
mountPath: /etc/litellm/billing-mtls-ca
|
||||
readOnly: true
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_BILLING_METRICS_CA_CERT
|
||||
value: /etc/litellm/billing-mtls-ca/ca.crt
|
||||
name: billing-metrics
|
||||
|
||||
- it: mounts the CA secret when caSecretName is set
|
||||
template: deployment.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
caSecretName: billing-ca
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_BILLING_METRICS_CA_CERT
|
||||
value: /etc/litellm/billing-mtls-ca/ca.crt
|
||||
- contains:
|
||||
path: spec.template.spec.volumes
|
||||
content:
|
||||
name: billing-metrics-mtls-ca
|
||||
secret:
|
||||
secretName: billing-ca
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].volumeMounts
|
||||
content:
|
||||
name: billing-metrics-mtls-ca
|
||||
mountPath: /etc/litellm/billing-mtls-ca
|
||||
readOnly: true
|
||||
|
||||
- it: passes the export interval through only when set
|
||||
- it: renders the export interval only when set
|
||||
template: deployment.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
|
|
@ -164,129 +69,7 @@ tests:
|
|||
name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS
|
||||
value: "5000"
|
||||
|
||||
- it: omits the export interval when unset
|
||||
template: deployment.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
asserts:
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS
|
||||
value: "60000"
|
||||
|
||||
# Kubernetes resolves duplicate env names last-wins, so the chart-owned billing
|
||||
# entries must render after .Values.envVars or a user could silently redirect
|
||||
# the metering export. The three billing entries are the last ones emitted here
|
||||
# (migrationJob, which appends DISABLE_SCHEMA_UPDATE, is off for this case).
|
||||
- it: renders the billing endpoint after envVars so it cannot be shadowed
|
||||
template: deployment.yaml
|
||||
set:
|
||||
migrationJob:
|
||||
enabled: false
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
envVars:
|
||||
LITELLM_BILLING_METRICS_ENDPOINT: https://shadowed.example
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_BILLING_METRICS_ENDPOINT
|
||||
value: https://shadowed.example
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].env[-3]
|
||||
value:
|
||||
name: LITELLM_BILLING_METRICS_ENDPOINT
|
||||
value: https://telemetry.litellm.ai
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].env[-2].name
|
||||
value: LITELLM_BILLING_METRICS_CLIENT_CERT
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].env[-1].name
|
||||
value: LITELLM_BILLING_METRICS_CLIENT_KEY
|
||||
|
||||
- it: keeps user-supplied volumes and mounts alongside the billing secret
|
||||
template: deployment.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
volumes:
|
||||
- name: custom-callbacks
|
||||
configMap:
|
||||
name: my-callbacks
|
||||
volumeMounts:
|
||||
- name: custom-callbacks
|
||||
mountPath: /app/callbacks
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.volumes
|
||||
content:
|
||||
name: custom-callbacks
|
||||
configMap:
|
||||
name: my-callbacks
|
||||
- contains:
|
||||
path: spec.template.spec.volumes
|
||||
content:
|
||||
name: billing-metrics-mtls
|
||||
secret:
|
||||
secretName: litellm-billing-metrics-mtls
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].volumeMounts
|
||||
content:
|
||||
name: custom-callbacks
|
||||
mountPath: /app/callbacks
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].volumeMounts
|
||||
content:
|
||||
name: billing-metrics-mtls
|
||||
mountPath: /etc/litellm/billing-mtls
|
||||
readOnly: true
|
||||
|
||||
- it: still mounts the proxy config when enabled
|
||||
template: deployment.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].volumeMounts
|
||||
content:
|
||||
name: litellm-config
|
||||
mountPath: /etc/litellm/config.yaml
|
||||
subPath: config.yaml
|
||||
|
||||
# Only the proxy serves billable traffic. The migrations Job must never mount
|
||||
# the client certificate, and it renders its own env and volumes, so nothing
|
||||
# stops a future edit from wiring the billing include into it by mistake.
|
||||
- it: does not touch the migrations job when enabled
|
||||
template: migrations-job.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
asserts:
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_BILLING_METRICS_ENDPOINT
|
||||
value: https://telemetry.litellm.ai
|
||||
- notExists:
|
||||
path: spec.template.spec.containers[0].volumeMounts
|
||||
- notExists:
|
||||
path: spec.template.spec.volumes
|
||||
|
||||
- it: fails loudly when enabled with an emptied secretName
|
||||
template: deployment.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
secretName: ""
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: billingMetrics.secretName is required when billingMetrics.enabled is true (an existing Secret with tls.crt and tls.key)
|
||||
|
||||
- it: fails loudly when enabled without an endpoint
|
||||
- it: requires the endpoint when enabled
|
||||
template: deployment.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
|
|
@ -295,3 +78,23 @@ tests:
|
|||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: billingMetrics.endpoint is required when billingMetrics.enabled is true
|
||||
|
||||
- it: requires the token URL when enabled
|
||||
template: deployment.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
tokenUrl: ""
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: billingMetrics.tokenUrl is required when billingMetrics.enabled is true
|
||||
|
||||
- it: requires the secret name when enabled
|
||||
template: deployment.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
secretName: ""
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: billingMetrics.secretName is required when billingMetrics.enabled is true
|
||||
|
|
|
|||
|
|
@ -191,7 +191,7 @@ tests:
|
|||
metricsServer.port: 9090
|
||||
billingMetrics.enabled: true
|
||||
billingMetrics.endpoint: https://metering.example.com
|
||||
billingMetrics.secretName: billing-mtls
|
||||
billingMetrics.secretName: billing-metrics
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
|
|
@ -214,9 +214,9 @@ tests:
|
|||
name: LITELLM_BILLING_METRICS_ENDPOINT
|
||||
any: true
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[1].volumeMounts
|
||||
path: spec.template.spec.containers[1].env
|
||||
content:
|
||||
name: billing-metrics-mtls
|
||||
name: LITELLM_BILLING_METRICS_CLIENT_SECRET
|
||||
any: true
|
||||
|
||||
- it: should give the sidecar the same scratch mounts as the proxy on a read-only root
|
||||
|
|
|
|||
|
|
@ -154,19 +154,13 @@ masterkeySecretName: ""
|
|||
# if set, use this secret key for the master key; otherwise, use the default key
|
||||
masterkeySecretKey: ""
|
||||
|
||||
# Optional: enterprise billable-request metering. When enabled, the proxy counts
|
||||
# successful requests to inference, MCP, and A2A endpoints and pushes them to
|
||||
# LiteLLM's collector over mutual TLS. Requires an enterprise license.
|
||||
# The client certificate identifies the deployment, so it is mounted read-only
|
||||
# from an existing Secret and never passed through the environment.
|
||||
# Optional: enterprise billable-request metering. Requires an enterprise license.
|
||||
billingMetrics:
|
||||
enabled: false
|
||||
endpoint: https://telemetry.litellm.ai # collector to push the counter to
|
||||
secretName: litellm-billing-metrics-mtls # existing Secret holding tls.crt and tls.key
|
||||
# Only for private or test collectors whose server certificate is not on the
|
||||
# public web PKI. The production collector needs no CA override.
|
||||
caSecretName: "" # existing Secret holding ca.crt
|
||||
exportIntervalMs: "" # push cadence; the proxy defaults to 60000
|
||||
endpoint: https://telemetry.litellm.ai
|
||||
tokenUrl: https://litellm-prod.auth.us-west-1.amazoncognito.com/oauth2/token
|
||||
secretName: litellm-billing-metrics
|
||||
exportIntervalMs: ""
|
||||
|
||||
proxyConfigMap:
|
||||
# when true, creates a new configmap
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@ apiVersion: v2
|
|||
name: litellm
|
||||
description: LiteLLM componentized — gateway, UI backend, and UI as separate services
|
||||
type: application
|
||||
version: 0.1.0
|
||||
version: 0.1.1
|
||||
appVersion: "0.1.0"
|
||||
annotations:
|
||||
org.opencontainers.image.source: "https://github.com/BerriAI/litellm"
|
||||
|
|
|
|||
|
|
@ -22,6 +22,23 @@ backend:
|
|||
gatewayUrl: https://gateway.example.com
|
||||
```
|
||||
|
||||
## Enterprise billing metrics
|
||||
|
||||
Billing metrics use Cognito bearer authentication. Create the Secret with
|
||||
`client-id` and `client-secret` keys, then enable billing metrics:
|
||||
|
||||
```bash
|
||||
kubectl create secret generic litellm-billing-metrics --from-literal=client-id=... --from-literal=client-secret=...
|
||||
```
|
||||
|
||||
```yaml
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
```
|
||||
|
||||
The chart requires an HTTPS endpoint, an HTTPS token URL, and the Secret name
|
||||
when billing metrics are enabled
|
||||
|
||||
## Routes to another Service
|
||||
|
||||
For a shared hostname, `ingress.extraPaths` can route to a Service in the same namespace using `serviceName` and `servicePort`. The port can be a number or a named Service port. Omit `service` when using these fields
|
||||
|
|
|
|||
|
|
@ -58,56 +58,27 @@ app.kubernetes.io/managed-by: {{ .Release.Service }}
|
|||
helm.sh/chart: {{ printf "%s-%s" .Chart.Name .Chart.Version | replace "+" "_" }}
|
||||
{{- end -}}
|
||||
|
||||
{{/*
|
||||
Enterprise billable-request metering. Wired into gateway and backend, not the
|
||||
migrations job. The gateway serves nearly all billable traffic, but the backend
|
||||
keeps the named-server MCP transport (/{mcp_server_name}/mcp), which writes a
|
||||
SpendLogs row, so metering only the gateway would silently drop that traffic.
|
||||
The client certificate identifies the deployment to LiteLLM's collector, so it is
|
||||
mounted read-only from an existing Secret rather than passed through the
|
||||
environment.
|
||||
*/}}
|
||||
{{- define "litellm.billingMetrics.certDir" -}}/etc/litellm/billing-mtls{{- end -}}
|
||||
{{- define "litellm.billingMetrics.caDir" -}}/etc/litellm/billing-mtls-ca{{- end -}}
|
||||
|
||||
{{- define "litellm.billingMetricsEnv" -}}
|
||||
- name: LITELLM_BILLING_METRICS_ENDPOINT
|
||||
value: {{ required "billingMetrics.endpoint is required when billingMetrics.enabled is true" .Values.billingMetrics.endpoint | quote }}
|
||||
- name: LITELLM_BILLING_METRICS_CLIENT_CERT
|
||||
value: {{ printf "%s/tls.crt" (include "litellm.billingMetrics.certDir" .) | quote }}
|
||||
- name: LITELLM_BILLING_METRICS_CLIENT_KEY
|
||||
value: {{ printf "%s/tls.key" (include "litellm.billingMetrics.certDir" .) | quote }}
|
||||
{{- if .Values.billingMetrics.caSecretName }}
|
||||
- name: LITELLM_BILLING_METRICS_CA_CERT
|
||||
value: {{ printf "%s/ca.crt" (include "litellm.billingMetrics.caDir" .) | quote }}
|
||||
{{- end }}
|
||||
- name: LITELLM_BILLING_METRICS_TOKEN_URL
|
||||
value: {{ required "billingMetrics.tokenUrl is required when billingMetrics.enabled is true" .Values.billingMetrics.tokenUrl | quote }}
|
||||
- name: LITELLM_BILLING_METRICS_CLIENT_ID
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: {{ required "billingMetrics.secretName is required when billingMetrics.enabled is true" .Values.billingMetrics.secretName }}
|
||||
key: client-id
|
||||
- name: LITELLM_BILLING_METRICS_CLIENT_SECRET
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: {{ required "billingMetrics.secretName is required when billingMetrics.enabled is true" .Values.billingMetrics.secretName }}
|
||||
key: client-secret
|
||||
{{- with .Values.billingMetrics.exportIntervalMs }}
|
||||
- name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS
|
||||
value: {{ . | quote }}
|
||||
{{- end }}
|
||||
{{- end -}}
|
||||
|
||||
{{- define "litellm.billingMetricsVolumes" -}}
|
||||
- name: billing-metrics-mtls
|
||||
secret:
|
||||
secretName: {{ required "billingMetrics.secretName is required when billingMetrics.enabled is true (an existing Secret with tls.crt and tls.key)" .Values.billingMetrics.secretName }}
|
||||
{{- if .Values.billingMetrics.caSecretName }}
|
||||
- name: billing-metrics-mtls-ca
|
||||
secret:
|
||||
secretName: {{ .Values.billingMetrics.caSecretName }}
|
||||
{{- end }}
|
||||
{{- end -}}
|
||||
|
||||
{{- define "litellm.billingMetricsVolumeMounts" -}}
|
||||
- name: billing-metrics-mtls
|
||||
mountPath: {{ include "litellm.billingMetrics.certDir" . }}
|
||||
readOnly: true
|
||||
{{- if .Values.billingMetrics.caSecretName }}
|
||||
- name: billing-metrics-mtls-ca
|
||||
mountPath: {{ include "litellm.billingMetrics.caDir" . }}
|
||||
readOnly: true
|
||||
{{- end }}
|
||||
{{- end -}}
|
||||
|
||||
{{/*
|
||||
Per-component selector labels — used in both Service selectors and Deployment matchLabels.
|
||||
|
|
|
|||
|
|
@ -79,16 +79,13 @@ spec:
|
|||
{{- include "litellm.billingMetricsEnv" . | nindent 12 }}
|
||||
{{- end }}
|
||||
{{- include "litellm.envFrom" .Values.backend | nindent 10 }}
|
||||
{{- if or .Values.gateway.config.create .Values.backend.volumeMounts .Values.billingMetrics.enabled }}
|
||||
{{- if or .Values.gateway.config.create .Values.backend.volumeMounts }}
|
||||
volumeMounts:
|
||||
{{- if .Values.gateway.config.create }}
|
||||
- name: gateway-config
|
||||
mountPath: /app/config/config.yaml
|
||||
subPath: config.yaml
|
||||
{{- end }}
|
||||
{{- if .Values.billingMetrics.enabled }}
|
||||
{{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }}
|
||||
{{- end }}
|
||||
{{- with .Values.backend.volumeMounts }}
|
||||
{{- toYaml . | nindent 12 }}
|
||||
{{- end }}
|
||||
|
|
@ -114,16 +111,13 @@ spec:
|
|||
{{- with .Values.backend.extraContainers }}
|
||||
{{- tpl (toYaml .) $ | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- if or .Values.gateway.config.create .Values.backend.volumes .Values.billingMetrics.enabled }}
|
||||
{{- if or .Values.gateway.config.create .Values.backend.volumes }}
|
||||
volumes:
|
||||
{{- if .Values.gateway.config.create }}
|
||||
- name: gateway-config
|
||||
configMap:
|
||||
name: {{ include "litellm.gateway.fullname" . }}-config
|
||||
{{- end }}
|
||||
{{- if .Values.billingMetrics.enabled }}
|
||||
{{- include "litellm.billingMetricsVolumes" . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- with .Values.backend.volumes }}
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
|
|
|
|||
|
|
@ -86,7 +86,7 @@ spec:
|
|||
{{- include "litellm.gateway.collectorEnv" . | nindent 12 }}
|
||||
{{- end }}
|
||||
{{- include "litellm.envFrom" .Values.gateway | nindent 10 }}
|
||||
{{- if or .Values.gateway.config.create .Values.gateway.volumeMounts .Values.billingMetrics.enabled .Values.gateway.metricsServer.enabled (include "litellm.gateway.collectorSocketDir" .) }}
|
||||
{{- if or .Values.gateway.config.create .Values.gateway.volumeMounts .Values.gateway.metricsServer.enabled (include "litellm.gateway.collectorSocketDir" .) }}
|
||||
volumeMounts:
|
||||
{{- if .Values.gateway.config.create }}
|
||||
- name: gateway-config
|
||||
|
|
@ -101,9 +101,6 @@ spec:
|
|||
- name: collector-socket
|
||||
mountPath: {{ include "litellm.gateway.collectorSocketDir" . }}
|
||||
{{- end }}
|
||||
{{- if .Values.billingMetrics.enabled }}
|
||||
{{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }}
|
||||
{{- end }}
|
||||
{{- with .Values.gateway.volumeMounts }}
|
||||
{{- toYaml . | nindent 12 }}
|
||||
{{- end }}
|
||||
|
|
@ -206,7 +203,7 @@ spec:
|
|||
{{- with .Values.gateway.extraContainers }}
|
||||
{{- tpl (toYaml .) $ | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- if or .Values.gateway.config.create .Values.gateway.volumes .Values.billingMetrics.enabled .Values.gateway.metricsServer.enabled (include "litellm.gateway.collectorSocketDir" .) }}
|
||||
{{- if or .Values.gateway.config.create .Values.gateway.volumes .Values.gateway.metricsServer.enabled (include "litellm.gateway.collectorSocketDir" .) }}
|
||||
volumes:
|
||||
{{- if .Values.gateway.config.create }}
|
||||
- name: gateway-config
|
||||
|
|
@ -222,9 +219,6 @@ spec:
|
|||
emptyDir:
|
||||
sizeLimit: 1Mi
|
||||
{{- end }}
|
||||
{{- if .Values.billingMetrics.enabled }}
|
||||
{{- include "litellm.billingMetricsVolumes" . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- with .Values.gateway.volumes }}
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
|
|
|
|||
|
|
@ -1,39 +1,30 @@
|
|||
suite: test billingMetrics wiring on gateway and backend
|
||||
suite: test bearer-only billingMetrics wiring
|
||||
templates:
|
||||
- gateway/deployment.yaml
|
||||
- gateway/configmap.yaml
|
||||
- backend/deployment.yaml
|
||||
- backend/configmap.yaml
|
||||
- migrations-job.yaml
|
||||
values:
|
||||
- ./values/required.yaml
|
||||
tests:
|
||||
- it: is off by default, adding no env, volume, or mount
|
||||
- it: is disabled by default
|
||||
template: gateway/deployment.yaml
|
||||
asserts:
|
||||
- notContains:
|
||||
path: spec.template.spec.volumes
|
||||
content:
|
||||
name: billing-metrics-mtls
|
||||
secret:
|
||||
secretName: billing-mtls
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].volumeMounts
|
||||
value:
|
||||
- name: gateway-config
|
||||
mountPath: /app/config/config.yaml
|
||||
subPath: config.yaml
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_BILLING_METRICS_ENDPOINT
|
||||
value: https://telemetry.litellm.ai
|
||||
- notContains:
|
||||
path: spec.template.spec.volumes
|
||||
content:
|
||||
name: billing-metrics
|
||||
|
||||
- it: renders the endpoint and the mounted cert paths when enabled
|
||||
- it: wires gateway URLs and bearer credentials from Secret keys
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
secretName: billing-mtls
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
|
|
@ -43,87 +34,68 @@ tests:
|
|||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_BILLING_METRICS_CLIENT_CERT
|
||||
value: /etc/litellm/billing-mtls/tls.crt
|
||||
name: LITELLM_BILLING_METRICS_TOKEN_URL
|
||||
value: https://litellm-prod.auth.us-west-1.amazoncognito.com/oauth2/token
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_BILLING_METRICS_CLIENT_KEY
|
||||
value: /etc/litellm/billing-mtls/tls.key
|
||||
|
||||
- it: mounts the cert secret read-only alongside the config volume
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
secretName: billing-mtls
|
||||
asserts:
|
||||
name: LITELLM_BILLING_METRICS_CLIENT_ID
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: litellm-billing-metrics
|
||||
key: client-id
|
||||
- contains:
|
||||
path: spec.template.spec.volumes
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: billing-metrics-mtls
|
||||
secret:
|
||||
secretName: billing-mtls
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].volumeMounts
|
||||
content:
|
||||
name: billing-metrics-mtls
|
||||
mountPath: /etc/litellm/billing-mtls
|
||||
readOnly: true
|
||||
|
||||
# The production collector presents a public web-PKI certificate, so the CA
|
||||
# override must stay absent unless a private collector is configured.
|
||||
- it: omits the CA env, volume, and mount when no caSecretName is set
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
secretName: billing-mtls
|
||||
asserts:
|
||||
name: LITELLM_BILLING_METRICS_CLIENT_SECRET
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: litellm-billing-metrics
|
||||
key: client-secret
|
||||
- notContains:
|
||||
path: spec.template.spec.volumes
|
||||
content:
|
||||
name: billing-metrics-mtls-ca
|
||||
secret:
|
||||
secretName: billing-ca
|
||||
name: billing-metrics
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
path: spec.template.spec.containers[0].volumeMounts
|
||||
content:
|
||||
name: LITELLM_BILLING_METRICS_CA_CERT
|
||||
value: /etc/litellm/billing-mtls-ca/ca.crt
|
||||
name: billing-metrics
|
||||
|
||||
- it: mounts the CA secret when caSecretName is set
|
||||
template: gateway/deployment.yaml
|
||||
- it: wires billing metrics on the backend without volumes
|
||||
template: backend/deployment.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
secretName: billing-mtls
|
||||
caSecretName: billing-ca
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_BILLING_METRICS_CA_CERT
|
||||
value: /etc/litellm/billing-mtls-ca/ca.crt
|
||||
- contains:
|
||||
name: LITELLM_BILLING_METRICS_CLIENT_SECRET
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: litellm-billing-metrics
|
||||
key: client-secret
|
||||
- notContains:
|
||||
path: spec.template.spec.volumes
|
||||
content:
|
||||
name: billing-metrics-mtls-ca
|
||||
secret:
|
||||
secretName: billing-ca
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].volumeMounts
|
||||
content:
|
||||
name: billing-metrics-mtls-ca
|
||||
mountPath: /etc/litellm/billing-mtls-ca
|
||||
readOnly: true
|
||||
name: billing-metrics
|
||||
|
||||
- it: passes the export interval through only when set
|
||||
- it: does not add billing metrics to migrations
|
||||
template: migrations-job.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
asserts:
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_BILLING_METRICS_ENDPOINT
|
||||
|
||||
- it: renders an optional export interval
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
secretName: billing-mtls
|
||||
exportIntervalMs: 5000
|
||||
asserts:
|
||||
- contains:
|
||||
|
|
@ -132,102 +104,27 @@ tests:
|
|||
name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS
|
||||
value: "5000"
|
||||
|
||||
- it: keeps user-supplied gateway volumes alongside the billing secret
|
||||
- it: requires the endpoint when enabled
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
secretName: billing-mtls
|
||||
gateway.volumes:
|
||||
- name: custom-callbacks
|
||||
configMap:
|
||||
name: my-callbacks
|
||||
gateway.volumeMounts:
|
||||
- name: custom-callbacks
|
||||
mountPath: /app/callbacks
|
||||
endpoint: ""
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.volumes
|
||||
content:
|
||||
name: custom-callbacks
|
||||
configMap:
|
||||
name: my-callbacks
|
||||
- contains:
|
||||
path: spec.template.spec.volumes
|
||||
content:
|
||||
name: billing-metrics-mtls
|
||||
secret:
|
||||
secretName: billing-mtls
|
||||
- failedTemplate:
|
||||
errorMessage: billingMetrics.endpoint is required when billingMetrics.enabled is true
|
||||
|
||||
# The backend keeps the named-server MCP transport (/{mcp_server_name}/mcp),
|
||||
# which writes a SpendLogs row, so it must meter too or that traffic is lost.
|
||||
- it: meters the backend as well, since it serves the MCP transport
|
||||
template: backend/deployment.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
secretName: billing-mtls
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_BILLING_METRICS_ENDPOINT
|
||||
value: https://telemetry.litellm.ai
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].volumeMounts
|
||||
content:
|
||||
name: billing-metrics-mtls
|
||||
mountPath: /etc/litellm/billing-mtls
|
||||
readOnly: true
|
||||
- contains:
|
||||
path: spec.template.spec.volumes
|
||||
content:
|
||||
name: billing-metrics-mtls
|
||||
secret:
|
||||
secretName: billing-mtls
|
||||
|
||||
- it: leaves the backend alone when metering is off
|
||||
template: backend/deployment.yaml
|
||||
asserts:
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_BILLING_METRICS_ENDPOINT
|
||||
value: https://telemetry.litellm.ai
|
||||
|
||||
# The migrations job runs prisma and serves no traffic; it must never receive
|
||||
# the client key.
|
||||
- it: never mounts the billing cert on the migrations job
|
||||
template: migrations-job.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
secretName: billing-mtls
|
||||
asserts:
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_BILLING_METRICS_ENDPOINT
|
||||
value: https://telemetry.litellm.ai
|
||||
- isNull:
|
||||
path: spec.template.spec.volumes
|
||||
|
||||
# The conventional Secret name is the default, so enabling metering needs no
|
||||
# secretName at all; the guard below only fires on an explicitly blanked one.
|
||||
- it: uses the conventional secret name by default
|
||||
- it: requires the token URL when enabled
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
tokenUrl: ""
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.volumes
|
||||
content:
|
||||
name: billing-metrics-mtls
|
||||
secret:
|
||||
secretName: litellm-billing-metrics-mtls
|
||||
- failedTemplate:
|
||||
errorMessage: billingMetrics.tokenUrl is required when billingMetrics.enabled is true
|
||||
|
||||
- it: fails loudly when the secretName is explicitly blanked
|
||||
- it: requires the secret name when enabled
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
|
|
@ -235,15 +132,4 @@ tests:
|
|||
secretName: ""
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: billingMetrics.secretName is required when billingMetrics.enabled is true (an existing Secret with tls.crt and tls.key)
|
||||
|
||||
- it: fails loudly when enabled without an endpoint
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
endpoint: ""
|
||||
secretName: billing-mtls
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: billingMetrics.endpoint is required when billingMetrics.enabled is true
|
||||
errorMessage: billingMetrics.secretName is required when billingMetrics.enabled is true
|
||||
|
|
|
|||
|
|
@ -91,7 +91,7 @@ tests:
|
|||
app.kubernetes.io/name: litellm
|
||||
app.kubernetes.io/instance: RELEASE-NAME
|
||||
app.kubernetes.io/managed-by: Helm
|
||||
helm.sh/chart: litellm-0.1.0
|
||||
helm.sh/chart: litellm-0.1.1
|
||||
app.kubernetes.io/component: migrations
|
||||
|
||||
- it: renders pod-level and container-level securityContext in their own scopes
|
||||
|
|
|
|||
|
|
@ -180,23 +180,14 @@ masterKey:
|
|||
secretKey: master-key
|
||||
|
||||
# Optional: enterprise billable-request metering. When enabled, the gateway and
|
||||
# backend count successful requests to inference, MCP, and A2A endpoints and push
|
||||
# them to LiteLLM's collector over mutual TLS. Both components serve billable
|
||||
# routes: the backend keeps the named-server MCP transport. Requires an
|
||||
# enterprise license. The client certificate identifies the deployment, so it is
|
||||
# mounted read-only from an existing Secret and never passed through the env.
|
||||
# backend count successful requests to inference, MCP, and A2A endpoints. Requires
|
||||
# an enterprise license.
|
||||
billingMetrics:
|
||||
enabled: false
|
||||
endpoint: https://telemetry.litellm.ai # collector to push the counter to
|
||||
# An existing Secret holding the client certificate under tls.crt and its key
|
||||
# under tls.key, usually created from the onboarding artifact. The default is
|
||||
# the conventional name, so the common path is to create that Secret and set
|
||||
# enabled: true. Override only if yours is named differently.
|
||||
secretName: litellm-billing-metrics-mtls
|
||||
# Only for private or test collectors whose server certificate is not on the
|
||||
# public web PKI. The production collector needs no CA override.
|
||||
caSecretName: "" # existing Secret holding ca.crt
|
||||
exportIntervalMs: "" # push cadence; the proxy defaults to 60000
|
||||
endpoint: https://telemetry.litellm.ai
|
||||
tokenUrl: https://litellm-prod.auth.us-west-1.amazoncognito.com/oauth2/token
|
||||
secretName: litellm-billing-metrics
|
||||
exportIntervalMs: ""
|
||||
|
||||
# External Postgres connection.
|
||||
database:
|
||||
|
|
|
|||
|
|
@ -0,0 +1,19 @@
|
|||
SET LOCAL lock_timeout = '10s';
|
||||
CREATE TABLE IF NOT EXISTS "LiteLLM_SpendLogDelivery" (
|
||||
"delivery_id" TEXT PRIMARY KEY DEFAULT gen_random_uuid()::text,
|
||||
"request_id" TEXT NOT NULL,
|
||||
"state" "AccountingDeliveryState" NOT NULL DEFAULT 'pending',
|
||||
"attempts" INTEGER NOT NULL DEFAULT 0,
|
||||
"next_attempt_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
"applied_at" TIMESTAMP(3),
|
||||
"last_error" TEXT,
|
||||
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
"accounting_payload" JSONB NOT NULL,
|
||||
"is_batch" BOOLEAN NOT NULL
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogDelivery_state_next_attempt_at_idx"
|
||||
ON "LiteLLM_SpendLogDelivery" ("state", "next_attempt_at");
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogDelivery_state_applied_at_idx"
|
||||
ON "LiteLLM_SpendLogDelivery" ("state", "applied_at");
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogDelivery_request_id_idx"
|
||||
ON "LiteLLM_SpendLogDelivery" ("request_id");
|
||||
|
|
@ -0,0 +1,9 @@
|
|||
SET LOCAL lock_timeout = '10s';
|
||||
|
||||
CREATE TABLE IF NOT EXISTS "LiteLLM_ClaudeCodePluginUpload" (
|
||||
"plugin_name" TEXT NOT NULL,
|
||||
"content" BYTEA NOT NULL,
|
||||
CONSTRAINT "LiteLLM_ClaudeCodePluginUpload_pkey" PRIMARY KEY ("plugin_name"),
|
||||
CONSTRAINT "LiteLLM_ClaudeCodePluginUpload_plugin_name_fkey" FOREIGN KEY ("plugin_name")
|
||||
REFERENCES "LiteLLM_ClaudeCodePluginTable"("name") ON DELETE CASCADE ON UPDATE CASCADE
|
||||
);
|
||||
|
|
@ -0,0 +1,5 @@
|
|||
SET LOCAL lock_timeout = '10s';
|
||||
DROP TABLE IF EXISTS "LiteLLM_AccountingDelivery";
|
||||
DROP TABLE IF EXISTS "LiteLLM_AccountingEvent";
|
||||
DROP TYPE IF EXISTS "AccountingDeliveryTarget";
|
||||
DROP TYPE IF EXISTS "AccountingEventKind";
|
||||
|
|
@ -0,0 +1,16 @@
|
|||
SET LOCAL lock_timeout = '10s';
|
||||
CREATE TABLE IF NOT EXISTS "LiteLLM_ClaudeCodePluginVersion" (
|
||||
"id" TEXT NOT NULL,
|
||||
"plugin_id" TEXT NOT NULL,
|
||||
"version" TEXT,
|
||||
"description" TEXT,
|
||||
"manifest_json" TEXT,
|
||||
"published_at" TIMESTAMP(3),
|
||||
"published_by" TEXT,
|
||||
"archive_available" BOOLEAN NOT NULL DEFAULT false,
|
||||
"content" BYTEA,
|
||||
CONSTRAINT "LiteLLM_ClaudeCodePluginVersion_pkey" PRIMARY KEY ("id"),
|
||||
CONSTRAINT "LiteLLM_ClaudeCodePluginVersion_plugin_id_fkey" FOREIGN KEY ("plugin_id") REFERENCES "LiteLLM_ClaudeCodePluginTable"("id") ON DELETE CASCADE ON UPDATE CASCADE
|
||||
);
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_ClaudeCodePluginVersion_plugin_id_version_key" ON "LiteLLM_ClaudeCodePluginVersion"("plugin_id", "version");
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_ClaudeCodePluginVersion_plugin_id_published_at_idx" ON "LiteLLM_ClaudeCodePluginVersion"("plugin_id", "published_at");
|
||||
|
|
@ -0,0 +1,27 @@
|
|||
SET LOCAL lock_timeout = '10s';
|
||||
|
||||
CREATE TABLE IF NOT EXISTS "LiteLLM_SkillSubmission" (
|
||||
"id" TEXT NOT NULL,
|
||||
"plugin_name" TEXT NOT NULL,
|
||||
"operation" TEXT NOT NULL,
|
||||
"status" TEXT NOT NULL DEFAULT 'pending_review',
|
||||
"proposal_json" TEXT NOT NULL,
|
||||
"archive_content" BYTEA,
|
||||
"submitted_by" TEXT NOT NULL,
|
||||
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
"base_plugin_id" TEXT,
|
||||
"base_updated_at" TIMESTAMP(3),
|
||||
"reviewed_by" TEXT,
|
||||
"reviewed_at" TIMESTAMP(3),
|
||||
"rejection_reason" TEXT,
|
||||
CONSTRAINT "LiteLLM_SkillSubmission_pkey" PRIMARY KEY ("id")
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_SkillSubmission_submitted_by_created_at_id_idx"
|
||||
ON "LiteLLM_SkillSubmission"("submitted_by", "created_at" DESC, "id" DESC);
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_SkillSubmission_submitted_by_status_created_at_id_idx"
|
||||
ON "LiteLLM_SkillSubmission"("submitted_by", "status", "created_at" DESC, "id" DESC);
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_SkillSubmission_status_created_at_id_idx"
|
||||
ON "LiteLLM_SkillSubmission"("status", "created_at" DESC, "id" DESC);
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_SkillSubmission_created_at_id_idx"
|
||||
ON "LiteLLM_SkillSubmission"("created_at" DESC, "id" DESC);
|
||||
|
|
@ -0,0 +1,3 @@
|
|||
SET LOCAL lock_timeout = '10s';
|
||||
|
||||
ALTER TABLE "LiteLLM_SkillSubmission" ADD COLUMN IF NOT EXISTS "scan_json" TEXT;
|
||||
|
|
@ -0,0 +1,37 @@
|
|||
SET LOCAL lock_timeout = '10s';
|
||||
|
||||
DO $$
|
||||
BEGIN
|
||||
IF NOT EXISTS (
|
||||
SELECT 1 FROM pg_attribute
|
||||
WHERE attrelid = to_regclass('"LiteLLM_ClaudeCodePluginTable"')
|
||||
AND attname = 'anonymous_download_count' AND NOT attisdropped
|
||||
) THEN
|
||||
ALTER TABLE "LiteLLM_ClaudeCodePluginTable"
|
||||
ADD COLUMN IF NOT EXISTS "anonymous_download_count" BIGINT NOT NULL DEFAULT 0;
|
||||
END IF;
|
||||
|
||||
IF to_regclass('"LiteLLM_ClaudeCodePluginDownload"') IS NULL THEN
|
||||
CREATE TABLE IF NOT EXISTS "LiteLLM_ClaudeCodePluginDownload" (
|
||||
"id" BIGSERIAL NOT NULL,
|
||||
"plugin_id" TEXT NOT NULL,
|
||||
"downloaded_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
"version" TEXT,
|
||||
"digest" TEXT NOT NULL,
|
||||
"authenticated" BOOLEAN NOT NULL,
|
||||
"user_id" TEXT,
|
||||
"user_email" TEXT,
|
||||
"key_alias" TEXT,
|
||||
"team_id" TEXT,
|
||||
"team_alias" TEXT,
|
||||
CONSTRAINT "LiteLLM_ClaudeCodePluginDownload_pkey" PRIMARY KEY ("id"),
|
||||
CONSTRAINT "LiteLLM_ClaudeCodePluginDownload_plugin_id_fkey" FOREIGN KEY ("plugin_id")
|
||||
REFERENCES "LiteLLM_ClaudeCodePluginTable"("id") ON DELETE CASCADE ON UPDATE CASCADE
|
||||
);
|
||||
END IF;
|
||||
|
||||
IF to_regclass('"LiteLLM_ClaudeCodePluginDownload_plugin_id_id_idx"') IS NULL THEN
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_ClaudeCodePluginDownload_plugin_id_id_idx"
|
||||
ON "LiteLLM_ClaudeCodePluginDownload"("plugin_id", "id");
|
||||
END IF;
|
||||
END $$;
|
||||
|
|
@ -1615,9 +1615,77 @@ model LiteLLM_ClaudeCodePluginTable {
|
|||
updated_at DateTime? @default(now()) @updatedAt
|
||||
created_by String?
|
||||
|
||||
uploaded_archive LiteLLM_ClaudeCodePluginUpload?
|
||||
versions LiteLLM_ClaudeCodePluginVersion[]
|
||||
downloads LiteLLM_ClaudeCodePluginDownload[]
|
||||
anonymous_download_count BigInt @default(0)
|
||||
|
||||
@@map("LiteLLM_ClaudeCodePluginTable")
|
||||
}
|
||||
|
||||
model LiteLLM_ClaudeCodePluginVersion {
|
||||
id String @id @default(uuid())
|
||||
plugin_id String
|
||||
version String?
|
||||
description String?
|
||||
manifest_json String?
|
||||
published_at DateTime?
|
||||
published_by String?
|
||||
archive_available Boolean @default(false)
|
||||
content Bytes?
|
||||
plugin LiteLLM_ClaudeCodePluginTable @relation(fields: [plugin_id], references: [id], onDelete: Cascade, onUpdate: Cascade)
|
||||
|
||||
@@unique([plugin_id, version])
|
||||
@@index([plugin_id, published_at])
|
||||
}
|
||||
|
||||
|
||||
model LiteLLM_SkillSubmission {
|
||||
id String @id @default(uuid())
|
||||
plugin_name String
|
||||
operation String
|
||||
status String @default("pending_review")
|
||||
proposal_json String
|
||||
archive_content Bytes?
|
||||
submitted_by String
|
||||
created_at DateTime @default(now())
|
||||
base_plugin_id String?
|
||||
base_updated_at DateTime?
|
||||
reviewed_by String?
|
||||
reviewed_at DateTime?
|
||||
rejection_reason String?
|
||||
scan_json String?
|
||||
|
||||
@@index([submitted_by, created_at(sort: Desc), id(sort: Desc)])
|
||||
@@index([submitted_by, status, created_at(sort: Desc), id(sort: Desc)])
|
||||
@@index([status, created_at(sort: Desc), id(sort: Desc)])
|
||||
@@index([created_at(sort: Desc), id(sort: Desc)])
|
||||
}
|
||||
|
||||
model LiteLLM_ClaudeCodePluginDownload {
|
||||
id BigInt @id @default(autoincrement())
|
||||
plugin_id String
|
||||
downloaded_at DateTime @default(now())
|
||||
version String?
|
||||
digest String
|
||||
authenticated Boolean
|
||||
user_id String?
|
||||
user_email String?
|
||||
key_alias String?
|
||||
team_id String?
|
||||
team_alias String?
|
||||
plugin LiteLLM_ClaudeCodePluginTable @relation(fields: [plugin_id], references: [id], onDelete: Cascade, onUpdate: Cascade)
|
||||
|
||||
@@index([plugin_id, id])
|
||||
}
|
||||
|
||||
model LiteLLM_ClaudeCodePluginUpload {
|
||||
plugin_name String @id
|
||||
content Bytes
|
||||
plugin LiteLLM_ClaudeCodePluginTable @relation(fields: [plugin_name], references: [name], onDelete: Cascade, onUpdate: Cascade)
|
||||
}
|
||||
|
||||
|
||||
// User/team-scoped memory store with a GLOBAL unique key.
|
||||
// `value` is a string (typically markdown/text meant for LLM context);
|
||||
// `metadata` is an optional JSON envelope for structured tags without schema changes.
|
||||
|
|
@ -2062,74 +2130,27 @@ model LiteLLM_TelemetryReport {
|
|||
@@index([created_at])
|
||||
}
|
||||
|
||||
enum AccountingEventKind {
|
||||
inbound_usage
|
||||
outbound_usage
|
||||
cost
|
||||
}
|
||||
|
||||
enum AccountingDeliveryState {
|
||||
pending
|
||||
applied
|
||||
parked
|
||||
}
|
||||
|
||||
enum AccountingDeliveryTarget {
|
||||
LiteLLM_VerificationToken
|
||||
LiteLLM_UserTable
|
||||
LiteLLM_TeamTable
|
||||
LiteLLM_TeamMembership
|
||||
LiteLLM_OrganizationTable
|
||||
LiteLLM_OrganizationMembership
|
||||
LiteLLM_EndUserTable
|
||||
LiteLLM_AgentsTable
|
||||
LiteLLM_TagTable
|
||||
LiteLLM_ProjectTable
|
||||
LiteLLM_ModelAccessGroupBudgetTable
|
||||
LiteLLM_BudgetTable
|
||||
LiteLLM_BudgetWindowSpend
|
||||
LiteLLM_DailyUserSpend
|
||||
LiteLLM_DailyTeamSpend
|
||||
LiteLLM_DailyOrganizationSpend
|
||||
LiteLLM_DailyEndUserSpend
|
||||
LiteLLM_DailyAgentSpend
|
||||
LiteLLM_DailyTagSpend
|
||||
LiteLLM_DailyModelUsage
|
||||
LiteLLM_DailyToolSpend
|
||||
LiteLLM_AutoRouterDailySpend
|
||||
}
|
||||
model LiteLLM_SpendLogDelivery {
|
||||
delivery_id String @id @default(dbgenerated("(gen_random_uuid())::text"))
|
||||
request_id String
|
||||
state AccountingDeliveryState @default(pending)
|
||||
accounting_payload Json
|
||||
is_batch Boolean
|
||||
attempts Int @default(0)
|
||||
next_attempt_at DateTime @default(now())
|
||||
applied_at DateTime?
|
||||
last_error String?
|
||||
created_at DateTime @default(now())
|
||||
|
||||
// Durable accounting ledger. One row per usage or cost event; payload holds the
|
||||
// private event envelope versioned by schema_version. Usage rows retain request-scoped
|
||||
// pricing context; cost rows include resolved rates and the finalized pricing snapshot.
|
||||
model LiteLLM_AccountingEvent {
|
||||
event_id String @id
|
||||
kind AccountingEventKind
|
||||
source_event_id String?
|
||||
schema_version Int
|
||||
payload_hash String
|
||||
occurred_at DateTime
|
||||
recorded_at DateTime @default(now())
|
||||
payload Json
|
||||
|
||||
@@index([recorded_at])
|
||||
@@index([kind, source_event_id])
|
||||
}
|
||||
|
||||
// One row per (event, target table, version): pending work for the projector and,
|
||||
// once applied in the same transaction as the target write, the receipt that stops a retry.
|
||||
model LiteLLM_AccountingDelivery {
|
||||
event_id String
|
||||
target AccountingDeliveryTarget
|
||||
version Int
|
||||
state AccountingDeliveryState @default(pending)
|
||||
attempts Int @default(0)
|
||||
next_attempt_at DateTime @default(now())
|
||||
applied_at DateTime?
|
||||
last_error String?
|
||||
|
||||
@@id([event_id, target, version])
|
||||
@@index([target, state, next_attempt_at])
|
||||
@@index([state, next_attempt_at])
|
||||
@@index([state, applied_at])
|
||||
@@index([request_id])
|
||||
}
|
||||
|
||||
model LiteLLM_SubscriptionAccountTable {
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ metadata:
|
|||
|
||||
# Rust gateway
|
||||
|
||||
`crates/gateway` mounts endpoint routers and serves the listener. `crates/gateway-inference`, `gateway-mcp` and `gateway-ui` own their endpoints, `gateway-auth` is the inbound authentication middleware and `gateway-management` holds management logic that has no router yet. A crate's `AGENTS.md`, where it has one, says what belongs in it. Read the reference for the part you are touching
|
||||
`crates/gateway` mounts endpoint routers and serves the listener. `crates/gateway-inference`, `gateway-mcp` and `gateway-ui` own their endpoints, `gateway-auth` is the inbound authentication middleware and `gateway-management` holds management logic that has no router yet. `gateway-forward` is the reverse proxy that relays every request to the Python proxy and has no routes of its own. A crate's `AGENTS.md`, where it has one, says what belongs in it. Read the reference for the part you are touching
|
||||
|
||||
- [surfaces](references/surfaces.md)
|
||||
- Each API surface mirrors the envelope the Python proxy returns there, so the shapes differ on purpose
|
||||
|
|
|
|||
175
litellm-rust/Cargo.lock
generated
175
litellm-rust/Cargo.lock
generated
|
|
@ -547,7 +547,7 @@ dependencies = [
|
|||
"http 1.4.2",
|
||||
"http-body 0.4.6",
|
||||
"hyper 0.14.32",
|
||||
"hyper 1.10.1",
|
||||
"hyper 1.12.0",
|
||||
"hyper-rustls 0.24.2",
|
||||
"hyper-rustls 0.27.9",
|
||||
"hyper-util",
|
||||
|
|
@ -736,7 +736,7 @@ dependencies = [
|
|||
"http 1.4.2",
|
||||
"http-body 1.1.0",
|
||||
"http-body-util",
|
||||
"hyper 1.10.1",
|
||||
"hyper 1.12.0",
|
||||
"hyper-util",
|
||||
"itoa",
|
||||
"matchit",
|
||||
|
|
@ -1000,7 +1000,7 @@ dependencies = [
|
|||
"home",
|
||||
"http 1.4.2",
|
||||
"http-body-util",
|
||||
"hyper 1.10.1",
|
||||
"hyper 1.12.0",
|
||||
"hyper-named-pipe",
|
||||
"hyper-rustls 0.27.9",
|
||||
"hyper-util",
|
||||
|
|
@ -2255,7 +2255,7 @@ dependencies = [
|
|||
"chrono",
|
||||
"http 1.4.2",
|
||||
"http-body-util",
|
||||
"hyper 1.10.1",
|
||||
"hyper 1.12.0",
|
||||
"hyper-rustls 0.27.9",
|
||||
"hyper-util",
|
||||
"ring",
|
||||
|
|
@ -2391,7 +2391,7 @@ dependencies = [
|
|||
"http 1.4.2",
|
||||
"http-body 1.1.0",
|
||||
"http-body-util",
|
||||
"hyper 1.10.1",
|
||||
"hyper 1.12.0",
|
||||
"lazy_static",
|
||||
"opentelemetry",
|
||||
"opentelemetry-semantic-conventions",
|
||||
|
|
@ -2833,9 +2833,9 @@ dependencies = [
|
|||
|
||||
[[package]]
|
||||
name = "hyper"
|
||||
version = "1.10.1"
|
||||
version = "1.12.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "55281c53a1894c864990125767da440a4e630446785086f52523b20033b74498"
|
||||
checksum = "2c3e324da4c95177d6291d4c8730197c0d1822f8a9766814a4a44fa5ab797c9c"
|
||||
dependencies = [
|
||||
"atomic-waker",
|
||||
"bytes",
|
||||
|
|
@ -2860,7 +2860,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
|||
checksum = "fab3637d6b04a8037af8a266fdf6cf92ea957e8c53981a2bf6136572531025bf"
|
||||
dependencies = [
|
||||
"hex",
|
||||
"hyper 1.10.1",
|
||||
"hyper 1.12.0",
|
||||
"hyper-util",
|
||||
"pin-project-lite",
|
||||
"tokio",
|
||||
|
|
@ -2889,7 +2889,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
|||
checksum = "33ca68d021ef39cf6463ab54c1d0f5daf03377b70561305bb89a8f83aab66e0f"
|
||||
dependencies = [
|
||||
"http 1.4.2",
|
||||
"hyper 1.10.1",
|
||||
"hyper 1.12.0",
|
||||
"hyper-util",
|
||||
"rustls 0.23.42",
|
||||
"rustls-native-certs",
|
||||
|
|
@ -2904,7 +2904,7 @@ version = "0.5.2"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2b90d566bffbce6a75bd8b09a05aa8c2cb1fabb6cb348f8840c9e4c90a0d83b0"
|
||||
dependencies = [
|
||||
"hyper 1.10.1",
|
||||
"hyper 1.12.0",
|
||||
"hyper-util",
|
||||
"pin-project-lite",
|
||||
"tokio",
|
||||
|
|
@ -2923,7 +2923,7 @@ dependencies = [
|
|||
"futures-util",
|
||||
"http 1.4.2",
|
||||
"http-body 1.1.0",
|
||||
"hyper 1.10.1",
|
||||
"hyper 1.12.0",
|
||||
"ipnet",
|
||||
"libc",
|
||||
"percent-encoding",
|
||||
|
|
@ -2942,7 +2942,7 @@ checksum = "986c5ce3b994526b3cd75578e62554abd09f0899d6206de48b3e96ab34ccc8c7"
|
|||
dependencies = [
|
||||
"hex",
|
||||
"http-body-util",
|
||||
"hyper 1.10.1",
|
||||
"hyper 1.12.0",
|
||||
"hyper-util",
|
||||
"pin-project-lite",
|
||||
"tokio",
|
||||
|
|
@ -3389,15 +3389,19 @@ version = "0.1.0"
|
|||
dependencies = [
|
||||
"aws-config",
|
||||
"aws-credential-types",
|
||||
"aws-runtime",
|
||||
"aws-sdk-sts",
|
||||
"aws-sigv4",
|
||||
"aws-smithy-runtime-api",
|
||||
"aws-smithy-types",
|
||||
"aws-types",
|
||||
"bytes",
|
||||
"futures-util",
|
||||
"litellm-auth-types",
|
||||
"litellm-error",
|
||||
"litellm-http",
|
||||
"litellm-secret-value",
|
||||
"litellm-secrets",
|
||||
"moka",
|
||||
"reqwest",
|
||||
"rstest",
|
||||
|
|
@ -3406,6 +3410,7 @@ dependencies = [
|
|||
"sha2 0.10.9",
|
||||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
"url",
|
||||
"veil",
|
||||
]
|
||||
|
||||
|
|
@ -3418,6 +3423,7 @@ dependencies = [
|
|||
"litellm-auth-types",
|
||||
"litellm-error",
|
||||
"litellm-secret-value",
|
||||
"litellm-secrets",
|
||||
"moka",
|
||||
"rstest",
|
||||
"serde_json",
|
||||
|
|
@ -3431,6 +3437,7 @@ dependencies = [
|
|||
name = "litellm-auth-gcp"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"axum",
|
||||
"base64 0.22.1",
|
||||
"gcp_auth",
|
||||
"google-cloud-auth",
|
||||
|
|
@ -3440,10 +3447,12 @@ dependencies = [
|
|||
"litellm-secret-value",
|
||||
"moka",
|
||||
"rstest",
|
||||
"rustls 0.23.42",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"sha2 0.10.9",
|
||||
"tempfile",
|
||||
"time",
|
||||
"tokio",
|
||||
"veil",
|
||||
]
|
||||
|
|
@ -3569,37 +3578,6 @@ dependencies = [
|
|||
"uuid",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-cache-redis"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"litellm-cache",
|
||||
"litellm-cache-testing",
|
||||
"litellm-error",
|
||||
"r2d2",
|
||||
"redis",
|
||||
"redis-test",
|
||||
"rstest",
|
||||
"serde_json",
|
||||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-cache-redis-semantic"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"litellm-cache",
|
||||
"litellm-cache-redis",
|
||||
"litellm-cache-testing",
|
||||
"redis",
|
||||
"redis-test",
|
||||
"rstest",
|
||||
"serde_json",
|
||||
"sha2 0.10.9",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-cache-response"
|
||||
version = "0.1.0"
|
||||
|
|
@ -3607,11 +3585,8 @@ dependencies = [
|
|||
"litellm-cache",
|
||||
"litellm-cache-gcs",
|
||||
"litellm-cache-memory",
|
||||
"litellm-cache-redis",
|
||||
"litellm-http",
|
||||
"py_literal",
|
||||
"redis",
|
||||
"redis-test",
|
||||
"rstest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
|
|
@ -3650,24 +3625,6 @@ dependencies = [
|
|||
"litellm-cache",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-cache-valkey-semantic"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"litellm-cache",
|
||||
"litellm-cache-redis",
|
||||
"litellm-cache-testing",
|
||||
"litellm-error",
|
||||
"redis",
|
||||
"redis-test",
|
||||
"rstest",
|
||||
"serde_json",
|
||||
"sha2 0.10.9",
|
||||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
"uuid",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-callbacks-legacy-python"
|
||||
version = "0.1.0"
|
||||
|
|
@ -3812,6 +3769,7 @@ dependencies = [
|
|||
"litellm-config",
|
||||
"litellm-error",
|
||||
"litellm-gateway-auth",
|
||||
"litellm-gateway-forward",
|
||||
"litellm-gateway-inference",
|
||||
"litellm-gateway-mcp",
|
||||
"litellm-gateway-ui",
|
||||
|
|
@ -3858,6 +3816,29 @@ dependencies = [
|
|||
"veil",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-gateway-forward"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"axum",
|
||||
"bytes",
|
||||
"envy",
|
||||
"futures-util",
|
||||
"http-body-util",
|
||||
"httparse",
|
||||
"hyper 1.12.0",
|
||||
"hyper-util",
|
||||
"proptest",
|
||||
"rstest",
|
||||
"serde",
|
||||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
"tokio-tungstenite",
|
||||
"tokio-util",
|
||||
"tower-service",
|
||||
"tracing",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-gateway-inference"
|
||||
version = "0.1.0"
|
||||
|
|
@ -4331,6 +4312,7 @@ dependencies = [
|
|||
name = "litellm-python-bridge"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"async-trait",
|
||||
"aws-sdk-secretsmanager",
|
||||
"bytes",
|
||||
"criterion",
|
||||
|
|
@ -4339,8 +4321,6 @@ dependencies = [
|
|||
"litellm-auth",
|
||||
"litellm-auth-aws",
|
||||
"litellm-cache",
|
||||
"litellm-cache-memory",
|
||||
"litellm-cache-redis",
|
||||
"litellm-cache-response",
|
||||
"litellm-callbacks-legacy-python",
|
||||
"litellm-core-utils",
|
||||
|
|
@ -4357,6 +4337,7 @@ dependencies = [
|
|||
"litellm-llms",
|
||||
"litellm-llms-types",
|
||||
"litellm-model-catalog",
|
||||
"litellm-router-core",
|
||||
"litellm-router-types",
|
||||
"litellm-secret-value",
|
||||
"litellm-secrets",
|
||||
|
|
@ -4369,7 +4350,6 @@ dependencies = [
|
|||
"prost",
|
||||
"pyo3",
|
||||
"pyo3-async-runtimes",
|
||||
"redis",
|
||||
"reqwest",
|
||||
"rstest",
|
||||
"serde",
|
||||
|
|
@ -4439,6 +4419,24 @@ dependencies = [
|
|||
[[package]]
|
||||
name = "litellm-router-core"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"arc-swap",
|
||||
"async-trait",
|
||||
"jiff",
|
||||
"litellm-cache",
|
||||
"litellm-router-types",
|
||||
"parking_lot",
|
||||
"rand 0.8.7",
|
||||
"rstest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"serde_with",
|
||||
"sha2 0.10.9",
|
||||
"strum",
|
||||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
"tracing",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-router-types"
|
||||
|
|
@ -5782,17 +5780,6 @@ version = "6.0.0"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf"
|
||||
|
||||
[[package]]
|
||||
name = "r2d2"
|
||||
version = "0.8.10"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "51de85fb3fb6524929c8a2eb85e6b6d363de4e8c48f9e2c2eac4944abc181c93"
|
||||
dependencies = [
|
||||
"log",
|
||||
"parking_lot",
|
||||
"scheduled-thread-pool",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rand"
|
||||
version = "0.8.7"
|
||||
|
|
@ -5966,18 +5953,6 @@ dependencies = [
|
|||
"xxhash-rust",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "redis-test"
|
||||
version = "1.0.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "804d36862e4323b69f96440cbb13c9894fc90176abdeaf91264e21d5d77f6aca"
|
||||
dependencies = [
|
||||
"rand 0.9.5",
|
||||
"redis",
|
||||
"socket2 0.6.5",
|
||||
"tempfile",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "redox_syscall"
|
||||
version = "0.5.18"
|
||||
|
|
@ -6080,7 +6055,7 @@ dependencies = [
|
|||
"http 1.4.2",
|
||||
"http-body 1.1.0",
|
||||
"http-body-util",
|
||||
"hyper 1.10.1",
|
||||
"hyper 1.12.0",
|
||||
"hyper-rustls 0.27.9",
|
||||
"hyper-util",
|
||||
"js-sys",
|
||||
|
|
@ -6454,15 +6429,6 @@ dependencies = [
|
|||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "scheduled-thread-pool"
|
||||
version = "0.2.7"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3cbc66816425a074528352f5789333ecff06ca41b36b0b0efdfbb29edc391a19"
|
||||
dependencies = [
|
||||
"parking_lot",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "schemars"
|
||||
version = "0.9.0"
|
||||
|
|
@ -7570,6 +7536,7 @@ dependencies = [
|
|||
"bytes",
|
||||
"futures-core",
|
||||
"futures-sink",
|
||||
"futures-util",
|
||||
"pin-project-lite",
|
||||
"tokio",
|
||||
]
|
||||
|
|
@ -7649,7 +7616,7 @@ dependencies = [
|
|||
"http 1.4.2",
|
||||
"http-body 1.1.0",
|
||||
"http-body-util",
|
||||
"hyper 1.10.1",
|
||||
"hyper 1.12.0",
|
||||
"hyper-timeout",
|
||||
"hyper-util",
|
||||
"percent-encoding",
|
||||
|
|
@ -8225,9 +8192,9 @@ dependencies = [
|
|||
|
||||
[[package]]
|
||||
name = "want"
|
||||
version = "0.3.1"
|
||||
version = "0.3.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "bfa7760aed19e106de2c7c0b581b509f2f25d3dacaf737cb82ac61bc6d760b0e"
|
||||
checksum = "ec4cdd0dd910afe868b7ef477227d8d538b46b3075031afee8a9f2acb0a2ed0b"
|
||||
dependencies = [
|
||||
"try-lock",
|
||||
]
|
||||
|
|
@ -8609,7 +8576,7 @@ dependencies = [
|
|||
"futures",
|
||||
"http 1.4.2",
|
||||
"http-body-util",
|
||||
"hyper 1.10.1",
|
||||
"hyper 1.12.0",
|
||||
"hyper-util",
|
||||
"log",
|
||||
"once_cell",
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ litellm-config = { path = "crates/config" }
|
|||
litellm-error = { path = "crates/error" }
|
||||
litellm-router = { path = "crates/router" }
|
||||
litellm-router-types = { path = "crates/router-types" }
|
||||
litellm-router-core = { path = "crates/router-core" }
|
||||
litellm-tracing = { path = "crates/tracing" }
|
||||
litellm-spend-clickhouse = { path = "crates/spend-clickhouse" }
|
||||
litellm-storage-clickhouse = { path = "crates/storage-clickhouse" }
|
||||
|
|
@ -25,6 +26,7 @@ litellm-inference-ocr = { path = "crates/inference-ocr" }
|
|||
litellm-inference-testing = { path = "crates/inference-testing" }
|
||||
litellm-gateway-mcp = { path = "crates/gateway-mcp" }
|
||||
litellm-gateway = { path = "crates/gateway" }
|
||||
litellm-gateway-forward = { path = "crates/gateway-forward" }
|
||||
litellm-gateway-inference = { path = "crates/gateway-inference" }
|
||||
litellm-gateway-auth = { path = "crates/gateway-auth" }
|
||||
litellm-gateway-management = { path = "crates/gateway-management" }
|
||||
|
|
@ -57,11 +59,9 @@ litellm-db-testing = { path = "crates/db-testing" }
|
|||
litellm-cache = { path = "crates/cache" }
|
||||
litellm-cache-azure-blob = { path = "crates/cache-azure-blob" }
|
||||
litellm-cache-memory = { path = "crates/cache-memory" }
|
||||
litellm-cache-redis = { path = "crates/cache-redis" }
|
||||
litellm-cache-s3 = { path = "crates/cache-s3" }
|
||||
litellm-cache-gcs = { path = "crates/cache-gcs" }
|
||||
litellm-cache-disk = { path = "crates/cache-disk" }
|
||||
litellm-cache-redis-semantic = { path = "crates/cache-redis-semantic" }
|
||||
litellm-cache-response = { path = "crates/cache-response" }
|
||||
litellm-cache-qdrant-semantic = { path = "crates/cache-qdrant-semantic" }
|
||||
litellm-cache-testing = { path = "crates/cache-testing" }
|
||||
|
|
@ -78,6 +78,8 @@ litellm-cost = { path = "crates/cost" }
|
|||
litellm-model-catalog = { path = "crates/model-catalog" }
|
||||
litellm-model-catalog-refresh = { path = "crates/model-catalog-refresh" }
|
||||
|
||||
arc-swap = "1.9.2"
|
||||
async-trait = "0.1.91"
|
||||
askama = { version = "0.16.1", default-features = false, features = ["derive", "std"] }
|
||||
tracing = "0.1"
|
||||
axum = { version = "0.8.9", default-features = false, features = ["http1", "tokio", "multipart"] }
|
||||
|
|
@ -97,7 +99,6 @@ macro_rules_attribute = "0.2.3"
|
|||
schemars = "1"
|
||||
reqwest = { version = "0.13", default-features = false, features = ["json", "multipart", "rustls", "http2", "stream", "query"] }
|
||||
redis = "1.8.0"
|
||||
redis-test = "1.0.4"
|
||||
testcontainers-modules = "0.15.0"
|
||||
qdrant-client = { version = "1.19.0", default-features = false }
|
||||
uuid = { version = "1", features = ["v4"] }
|
||||
|
|
|
|||
|
|
@ -11,14 +11,16 @@ litellm-error.workspace = true
|
|||
litellm-http.workspace = true
|
||||
reqwest.workspace = true
|
||||
|
||||
moka = { workspace = true, features = ["sync"] }
|
||||
moka.workspace = true
|
||||
litellm-secret-value.workspace = true
|
||||
serde.workspace = true
|
||||
serde_json = { workspace = true, features = ["preserve_order"] }
|
||||
sha2.workspace = true
|
||||
thiserror.workspace = true
|
||||
url.workspace = true
|
||||
veil.workspace = true
|
||||
|
||||
aws-config = { version = "1.9.0", default-features = false, features = ["rustls", "rt-tokio"] }
|
||||
aws-config = { version = "1.9.0", default-features = false, features = ["rustls", "rt-tokio", "credentials-process"] }
|
||||
aws-credential-types = { version = "1.3.0", features = ["hardcoded-credentials"] }
|
||||
aws-sdk-sts = { version = "1.108.0", default-features = false, features = ["rustls", "rt-tokio"] }
|
||||
aws-sigv4 = "1.5.1"
|
||||
|
|
@ -26,6 +28,9 @@ aws-types = "1.4.0"
|
|||
aws-smithy-runtime-api = "1.13.0"
|
||||
|
||||
[dev-dependencies]
|
||||
litellm-secrets.workspace = true
|
||||
aws-smithy-types = "1.6.4"
|
||||
aws-runtime = { version = "1.9.4", default-features = false }
|
||||
rstest.workspace = true
|
||||
litellm-http = { workspace = true, features = ["test-support"] }
|
||||
bytes.workspace = true
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -11,8 +11,8 @@ pub const AWS_SESSION_NAME: &str = AwsParams::SESSION_NAME.env[0];
|
|||
pub const AWS_PROFILE_NAME: &str = AwsParams::PROFILE_NAME.env[0];
|
||||
pub const AWS_ROLE_NAME: &str = AwsParams::ROLE_NAME.env[0];
|
||||
pub const AWS_WEB_IDENTITY_TOKEN: &str = AwsParams::WEB_IDENTITY_TOKEN.env[0];
|
||||
pub const AWS_ROLE_ARN: &str = "AWS_ROLE_ARN";
|
||||
pub const AWS_WEB_IDENTITY_TOKEN_FILE: &str = "AWS_WEB_IDENTITY_TOKEN_FILE";
|
||||
pub const AWS_ROLE_ARN: &str = AwsParams::ROLE_ARN_ENV;
|
||||
pub const AWS_WEB_IDENTITY_TOKEN_FILE: &str = AwsParams::WEB_IDENTITY_TOKEN_FILE_ENV;
|
||||
pub const AWS_STS_ENDPOINT: &str = AwsParams::STS_ENDPOINT.env[0];
|
||||
pub const AWS_EXTERNAL_ID: &str = AwsParams::EXTERNAL_ID.env[0];
|
||||
pub const AWS_BEARER_TOKEN_BEDROCK: &str = "AWS_BEARER_TOKEN_BEDROCK";
|
||||
|
|
|
|||
|
|
@ -2,6 +2,8 @@ use litellm_error::Source;
|
|||
|
||||
#[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)]
|
||||
pub enum Error {
|
||||
#[error("AWS credential provider returned expired credentials")]
|
||||
ExpiredCredentials,
|
||||
#[error("could not load AWS {flow} credentials")]
|
||||
Credentials {
|
||||
flow: CredentialFlow,
|
||||
|
|
|
|||
|
|
@ -75,7 +75,8 @@ impl RequestSigner for SigV4Signer {
|
|||
)
|
||||
.map_err(|error| match error {
|
||||
Error::NonAsciiHeader { .. } => litellm_http::Error::unsignable_header(error),
|
||||
Error::Credentials { .. }
|
||||
Error::ExpiredCredentials
|
||||
| Error::Credentials { .. }
|
||||
| Error::MissingWebIdentityCredentials
|
||||
| Error::WebIdentityExpiration(_)
|
||||
| Error::Signing { .. } => litellm_http::Error::signature(error),
|
||||
|
|
|
|||
|
|
@ -8,17 +8,19 @@ repository.workspace = true
|
|||
[dependencies]
|
||||
litellm-secret-value.workspace = true
|
||||
litellm-auth-types.workspace = true
|
||||
litellm-error.workspace = true
|
||||
|
||||
moka.workspace = true
|
||||
serde_json.workspace = true
|
||||
sha2.workspace = true
|
||||
strum.workspace = true
|
||||
tokio.workspace = true
|
||||
url.workspace = true
|
||||
|
||||
azure_core = "1.0.0"
|
||||
azure_identity = { version = "1.0.0", features = ["tokio"] }
|
||||
|
||||
[dev-dependencies]
|
||||
litellm-error.workspace = true
|
||||
litellm-secrets.workspace = true
|
||||
rstest.workspace = true
|
||||
tokio.workspace = true
|
||||
tokio = { workspace = true, features = ["test-util"] }
|
||||
|
|
|
|||
|
|
@ -3,5 +3,5 @@ mod native;
|
|||
mod resolve;
|
||||
mod types;
|
||||
|
||||
pub use resolve::{AzureAuthService, SECRET_NAMES};
|
||||
pub use resolve::{AzureAuthService, AzureIdentity, SECRET_NAMES};
|
||||
pub use types::{AzureAuthInputs, ConfigValue};
|
||||
|
|
|
|||
|
|
@ -59,6 +59,7 @@ pub(crate) enum NativeAzureRequest {
|
|||
pub(crate) struct ValidatedAzureRequest {
|
||||
request: NativeAzureRequest,
|
||||
credential_source: InputSource,
|
||||
probe_imds: bool,
|
||||
}
|
||||
|
||||
impl ValidatedAzureRequest {
|
||||
|
|
@ -68,9 +69,18 @@ impl ValidatedAzureRequest {
|
|||
Ok(Self {
|
||||
request,
|
||||
credential_source,
|
||||
probe_imds: false,
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn with_imds_probe(self, enabled: bool) -> Self {
|
||||
Self {
|
||||
probe_imds: enabled
|
||||
&& matches!(self.request, NativeAzureRequest::ManagedIdentity { .. }),
|
||||
..self
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn credential_source(&self) -> InputSource {
|
||||
self.credential_source
|
||||
}
|
||||
|
|
@ -87,9 +97,34 @@ impl ValidatedAzureRequest {
|
|||
}
|
||||
}
|
||||
|
||||
pub(crate) enum CredentialAttempt {
|
||||
Resolved(ResolvedCredential),
|
||||
Unavailable(Error),
|
||||
Failed(Error),
|
||||
}
|
||||
|
||||
impl CredentialAttempt {
|
||||
pub(crate) fn into_result(self) -> Result<ResolvedCredential, Error> {
|
||||
match self {
|
||||
Self::Resolved(credential) => Ok(credential),
|
||||
Self::Unavailable(error) | Self::Failed(error) => Err(error),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) struct NativeAzureTokenAcquirer {
|
||||
cache: AzureCredentialProviderCache,
|
||||
transport: Option<azure_core::http::Transport>,
|
||||
imds_available: tokio::sync::Mutex<ImdsAvailability>,
|
||||
}
|
||||
|
||||
enum ImdsAvailability {
|
||||
Unknown,
|
||||
Available,
|
||||
Unavailable {
|
||||
retry_at: tokio::time::Instant,
|
||||
error: Error,
|
||||
},
|
||||
}
|
||||
|
||||
impl Default for NativeAzureTokenAcquirer {
|
||||
|
|
@ -103,6 +138,7 @@ impl NativeAzureTokenAcquirer {
|
|||
Self {
|
||||
cache: AzureCredentialProviderCache::new(cache_capacity),
|
||||
transport: None,
|
||||
imds_available: tokio::sync::Mutex::new(ImdsAvailability::Unknown),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -114,10 +150,51 @@ impl NativeAzureTokenAcquirer {
|
|||
Self {
|
||||
cache: AzureCredentialProviderCache::new(cache_capacity),
|
||||
transport: Some(transport),
|
||||
imds_available: tokio::sync::Mutex::new(ImdsAvailability::Unknown),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn acquire(
|
||||
pub(crate) async fn acquire(&self, request: ValidatedAzureRequest) -> CredentialAttempt {
|
||||
let discover_imds = request.probe_imds;
|
||||
if discover_imds && let Err(error) = self.discover_imds().await {
|
||||
return CredentialAttempt::Unavailable(error);
|
||||
}
|
||||
match self.acquire_token(request).await {
|
||||
Ok(credential) => CredentialAttempt::Resolved(credential),
|
||||
Err(error) => {
|
||||
let missing_identity = discover_imds && unassigned_identity(&error);
|
||||
if missing_identity {
|
||||
CredentialAttempt::Unavailable(error)
|
||||
} else {
|
||||
CredentialAttempt::Failed(error)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn discover_imds(&self) -> Result<(), Error> {
|
||||
let mut availability = self.imds_available.lock().await;
|
||||
match &*availability {
|
||||
ImdsAvailability::Available => return Ok(()),
|
||||
ImdsAvailability::Unavailable { retry_at, error }
|
||||
if tokio::time::Instant::now() < *retry_at =>
|
||||
{
|
||||
return Err(error.clone());
|
||||
}
|
||||
ImdsAvailability::Unknown | ImdsAvailability::Unavailable { .. } => {}
|
||||
}
|
||||
let result = probe_imds(self.transport.clone().unwrap_or_default()).await;
|
||||
*availability = match &result {
|
||||
Ok(()) => ImdsAvailability::Available,
|
||||
Err(error) => ImdsAvailability::Unavailable {
|
||||
retry_at: tokio::time::Instant::now() + Duration::from_secs(30),
|
||||
error: error.clone(),
|
||||
},
|
||||
};
|
||||
result
|
||||
}
|
||||
|
||||
async fn acquire_token(
|
||||
&self,
|
||||
request: ValidatedAzureRequest,
|
||||
) -> Result<ResolvedCredential, Error> {
|
||||
|
|
@ -148,6 +225,54 @@ impl NativeAzureTokenAcquirer {
|
|||
}
|
||||
}
|
||||
|
||||
fn unassigned_identity(error: &(dyn std::error::Error + 'static)) -> bool {
|
||||
std::iter::successors(Some(error), |cause| cause.source()).any(|cause| {
|
||||
cause
|
||||
.downcast_ref::<azure_core::Error>()
|
||||
.or_else(|| {
|
||||
cause
|
||||
.downcast_ref::<litellm_error::Source>()?
|
||||
.downcast_ref::<azure_core::Error>()
|
||||
})
|
||||
.is_some_and(|error| {
|
||||
error.http_status() == Some(azure_core::http::StatusCode::BadRequest)
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
async fn probe_imds(transport: azure_core::http::Transport) -> Result<(), Error> {
|
||||
use azure_core::http::{Context, Method, Request, Url};
|
||||
|
||||
let endpoint =
|
||||
Url::parse("http://169.254.169.254/metadata/identity/oauth2/token?api-version=2019-08-01")
|
||||
.map_err(|error| {
|
||||
Error::InvalidConfiguration(ErrorDetail::failed("Azure IMDS probe endpoint", error))
|
||||
})?;
|
||||
let mut request = Request::new(endpoint, Method::Get);
|
||||
let response = tokio::time::timeout(
|
||||
Duration::from_secs(1),
|
||||
transport.send(&Context::new(), &mut request),
|
||||
)
|
||||
.await
|
||||
.map_err(|error| {
|
||||
Error::CredentialAcquisition(ErrorDetail::failed("Azure IMDS availability probe", error))
|
||||
})?
|
||||
.map_err(|error| {
|
||||
Error::CredentialAcquisition(ErrorDetail::failed("Azure IMDS availability probe", error))
|
||||
})?;
|
||||
// Headerless IMDS discovery expects 400, checked 2026-11-01
|
||||
// https://learn.microsoft.com/en-us/azure/virtual-machines/instance-metadata-service#errors-and-debugging
|
||||
if response.status() != azure_core::http::StatusCode::BadRequest {
|
||||
return Err(Error::CredentialAcquisition(
|
||||
ErrorDetail::UnexpectedStatus {
|
||||
operation: "Azure IMDS availability probe",
|
||||
status: u16::from(response.status()),
|
||||
},
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
impl NativeAzureRequest {
|
||||
fn scope(&self) -> &str {
|
||||
match self {
|
||||
|
|
@ -493,6 +618,7 @@ impl ClientAssertion for StaticAssertion {
|
|||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::Duration;
|
||||
|
||||
use azure_core::http::headers::Headers;
|
||||
use azure_core::http::{AsyncRawResponse, HttpClient, Request, StatusCode, Transport};
|
||||
|
|
@ -569,6 +695,296 @@ mod tests {
|
|||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct AvailableImdsClient {
|
||||
probe_status: StatusCode,
|
||||
probe_delay: Duration,
|
||||
recover_after: Option<usize>,
|
||||
response_status: Option<StatusCode>,
|
||||
probes: std::sync::atomic::AtomicUsize,
|
||||
}
|
||||
|
||||
impl Default for AvailableImdsClient {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
probe_status: StatusCode::BadRequest,
|
||||
probe_delay: Duration::ZERO,
|
||||
recover_after: None,
|
||||
response_status: None,
|
||||
probes: std::sync::atomic::AtomicUsize::new(0),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl HttpClient for AvailableImdsClient {
|
||||
fn execute_request<'life0, 'life1, 'async_trait>(
|
||||
&'life0 self,
|
||||
request: &'life1 Request,
|
||||
) -> std::pin::Pin<
|
||||
Box<dyn std::future::Future<Output = Result<AsyncRawResponse>> + Send + 'async_trait>,
|
||||
>
|
||||
where
|
||||
'life0: 'async_trait,
|
||||
'life1: 'async_trait,
|
||||
Self: 'async_trait,
|
||||
{
|
||||
Box::pin(async move {
|
||||
if request
|
||||
.headers()
|
||||
.get_optional_str(&azure_core::http::headers::HeaderName::from_static(
|
||||
"metadata",
|
||||
))
|
||||
.is_none()
|
||||
{
|
||||
let probe = self
|
||||
.probes
|
||||
.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
|
||||
if self.recover_after.is_some_and(|failures| probe >= failures) {
|
||||
return Ok(AsyncRawResponse::from_bytes(
|
||||
StatusCode::BadRequest,
|
||||
Headers::new(),
|
||||
"metadata header required",
|
||||
));
|
||||
}
|
||||
tokio::time::sleep(self.probe_delay).await;
|
||||
return Ok(AsyncRawResponse::from_bytes(
|
||||
self.probe_status,
|
||||
Headers::new(),
|
||||
"metadata header required",
|
||||
));
|
||||
}
|
||||
if let Some(status) = self.response_status {
|
||||
return Ok(AsyncRawResponse::from_bytes(
|
||||
status,
|
||||
Headers::new(),
|
||||
"identity endpoint failure",
|
||||
));
|
||||
}
|
||||
tokio::time::sleep(std::time::Duration::from_secs(2)).await;
|
||||
let expires = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_secs()
|
||||
+ 3600;
|
||||
Ok(AsyncRawResponse::from_bytes(
|
||||
StatusCode::Ok,
|
||||
Headers::new(),
|
||||
serde_json::json!({
|
||||
"access_token": "managed-token", "expires_on": expires.to_string(),
|
||||
"token_type": "Bearer", "resource": "https://service.test",
|
||||
})
|
||||
.to_string(),
|
||||
))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[tokio::test(start_paused = true)]
|
||||
async fn available_imds_retains_full_sdk_token_acquisition() {
|
||||
let acquirer = NativeAzureTokenAcquirer::with_transport(
|
||||
4,
|
||||
Transport::new(Arc::new(AvailableImdsClient {
|
||||
probe_status: StatusCode::BadRequest,
|
||||
response_status: None,
|
||||
..Default::default()
|
||||
})),
|
||||
);
|
||||
let request = ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity {
|
||||
client_id: None,
|
||||
scope: deployment("https://service.test/.default".into()),
|
||||
selection_source: InputSource::Deployment,
|
||||
})
|
||||
.unwrap()
|
||||
.with_imds_probe(true);
|
||||
let credential = acquirer.acquire(request).await.into_result().unwrap();
|
||||
assert_eq!(credential.secret().expose(), "managed-token");
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::another_metadata_service(StatusCode::NotFound)]
|
||||
#[case::token_required(StatusCode::Unauthorized)]
|
||||
#[case::proxy_rejection(StatusCode::Forbidden)]
|
||||
#[case::unavailable(StatusCode::ServiceUnavailable)]
|
||||
#[tokio::test(start_paused = true)]
|
||||
async fn unrelated_metadata_responses_do_not_start_managed_identity_acquisition(
|
||||
#[case] status: StatusCode,
|
||||
) {
|
||||
let client = Arc::new(AvailableImdsClient {
|
||||
probe_status: status,
|
||||
..Default::default()
|
||||
});
|
||||
let acquirer = NativeAzureTokenAcquirer::with_transport(4, Transport::new(client.clone()));
|
||||
let request = ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity {
|
||||
client_id: None,
|
||||
scope: deployment("https://service.test/.default".into()),
|
||||
selection_source: InputSource::Deployment,
|
||||
})
|
||||
.unwrap()
|
||||
.with_imds_probe(true);
|
||||
assert!(matches!(
|
||||
acquirer.acquire(request.clone()).await,
|
||||
super::CredentialAttempt::Unavailable(_)
|
||||
));
|
||||
assert!(matches!(
|
||||
acquirer.acquire(request).await,
|
||||
super::CredentialAttempt::Unavailable(_)
|
||||
));
|
||||
assert_eq!(client.probes.load(std::sync::atomic::Ordering::SeqCst), 1);
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[tokio::test(start_paused = true)]
|
||||
async fn an_imds_timeout_is_not_repeated_on_the_next_request() {
|
||||
let client = Arc::new(AvailableImdsClient {
|
||||
probe_delay: Duration::from_secs(10),
|
||||
..Default::default()
|
||||
});
|
||||
let acquirer = NativeAzureTokenAcquirer::with_transport(4, Transport::new(client.clone()));
|
||||
let request = ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity {
|
||||
client_id: None,
|
||||
scope: deployment("https://service.test/.default".into()),
|
||||
selection_source: InputSource::Deployment,
|
||||
})
|
||||
.unwrap()
|
||||
.with_imds_probe(true);
|
||||
assert!(matches!(
|
||||
acquirer.acquire(request.clone()).await,
|
||||
super::CredentialAttempt::Unavailable(_)
|
||||
));
|
||||
let second = tokio::time::timeout(Duration::from_millis(10), acquirer.acquire(request))
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(matches!(second, super::CredentialAttempt::Unavailable(_)));
|
||||
assert_eq!(client.probes.load(std::sync::atomic::Ordering::SeqCst), 1);
|
||||
}
|
||||
|
||||
#[rstest::fixture]
|
||||
fn imds_request() -> ValidatedAzureRequest {
|
||||
ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity {
|
||||
client_id: None,
|
||||
scope: deployment("https://service.test/.default".into()),
|
||||
selection_source: InputSource::Deployment,
|
||||
})
|
||||
.unwrap()
|
||||
.with_imds_probe(true)
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::timeout(StatusCode::BadRequest, Duration::from_secs(10))]
|
||||
#[case::service_unavailable(StatusCode::ServiceUnavailable, Duration::ZERO)]
|
||||
#[tokio::test(start_paused = true)]
|
||||
async fn failed_imds_discovery_recovers_without_restarting_the_acquirer(
|
||||
imds_request: ValidatedAzureRequest,
|
||||
#[case] status: StatusCode,
|
||||
#[case] delay: Duration,
|
||||
) {
|
||||
let client = Arc::new(AvailableImdsClient {
|
||||
probe_status: status,
|
||||
probe_delay: delay,
|
||||
recover_after: Some(1),
|
||||
..Default::default()
|
||||
});
|
||||
let acquirer = NativeAzureTokenAcquirer::with_transport(4, Transport::new(client.clone()));
|
||||
let (first, second) = tokio::join!(
|
||||
acquirer.acquire(imds_request.clone()),
|
||||
acquirer.acquire(imds_request.clone()),
|
||||
);
|
||||
assert!(matches!(first, super::CredentialAttempt::Unavailable(_)));
|
||||
assert!(matches!(second, super::CredentialAttempt::Unavailable(_)));
|
||||
assert_eq!(client.probes.load(std::sync::atomic::Ordering::SeqCst), 1);
|
||||
|
||||
tokio::time::advance(Duration::from_secs(30)).await;
|
||||
let (first, second) = tokio::join!(
|
||||
acquirer.acquire(imds_request.clone()),
|
||||
acquirer.acquire(imds_request.clone()),
|
||||
);
|
||||
assert_eq!(
|
||||
first.into_result().unwrap().secret().expose(),
|
||||
"managed-token"
|
||||
);
|
||||
assert_eq!(
|
||||
second.into_result().unwrap().secret().expose(),
|
||||
"managed-token"
|
||||
);
|
||||
assert_eq!(client.probes.load(std::sync::atomic::Ordering::SeqCst), 2);
|
||||
|
||||
tokio::time::advance(Duration::from_secs(60)).await;
|
||||
assert_eq!(
|
||||
acquirer
|
||||
.acquire(imds_request)
|
||||
.await
|
||||
.into_result()
|
||||
.unwrap()
|
||||
.secret()
|
||||
.expose(),
|
||||
"managed-token",
|
||||
);
|
||||
assert_eq!(client.probes.load(std::sync::atomic::Ordering::SeqCst), 2);
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[tokio::test(start_paused = true)]
|
||||
async fn a_persistent_imds_outage_reuses_each_failed_probe(
|
||||
imds_request: ValidatedAzureRequest,
|
||||
) {
|
||||
let client = Arc::new(AvailableImdsClient {
|
||||
probe_delay: Duration::from_secs(10),
|
||||
..Default::default()
|
||||
});
|
||||
let acquirer = NativeAzureTokenAcquirer::with_transport(4, Transport::new(client.clone()));
|
||||
assert!(matches!(
|
||||
acquirer.acquire(imds_request.clone()).await,
|
||||
super::CredentialAttempt::Unavailable(_),
|
||||
));
|
||||
tokio::time::advance(Duration::from_secs(30)).await;
|
||||
let (first, second) = tokio::join!(
|
||||
acquirer.acquire(imds_request.clone()),
|
||||
acquirer.acquire(imds_request),
|
||||
);
|
||||
assert!(matches!(first, super::CredentialAttempt::Unavailable(_)));
|
||||
assert!(matches!(second, super::CredentialAttempt::Unavailable(_)));
|
||||
assert_eq!(client.probes.load(std::sync::atomic::Ordering::SeqCst), 2);
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::default_identity_not_assigned(true, StatusCode::BadRequest, true)]
|
||||
#[case::explicit_identity_not_assigned(false, StatusCode::BadRequest, false)]
|
||||
#[case::default_identity_forbidden(true, StatusCode::Forbidden, false)]
|
||||
#[case::default_identity_service_failure(true, StatusCode::InternalServerError, false)]
|
||||
#[tokio::test(start_paused = true)]
|
||||
async fn only_unavailable_default_identities_allow_fallback(
|
||||
#[case] discovery: bool,
|
||||
#[case] status: StatusCode,
|
||||
#[case] unavailable: bool,
|
||||
) {
|
||||
let acquirer = NativeAzureTokenAcquirer::with_transport(
|
||||
4,
|
||||
Transport::new(Arc::new(AvailableImdsClient {
|
||||
probe_status: StatusCode::BadRequest,
|
||||
response_status: Some(status),
|
||||
..Default::default()
|
||||
})),
|
||||
);
|
||||
let request = ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity {
|
||||
client_id: None,
|
||||
scope: deployment("https://service.test/.default".into()),
|
||||
selection_source: InputSource::Deployment,
|
||||
})
|
||||
.unwrap()
|
||||
.with_imds_probe(discovery);
|
||||
let outcome = acquirer.acquire(request).await;
|
||||
let is_unavailable = matches!(outcome, super::CredentialAttempt::Unavailable(_));
|
||||
let error = outcome.into_result().unwrap_err();
|
||||
assert_eq!(
|
||||
is_unavailable,
|
||||
unavailable,
|
||||
"{}",
|
||||
litellm_error::Report::new(&error)
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[tokio::test]
|
||||
async fn client_secret_uses_sdk_protocol_and_reuses_cached_credential() {
|
||||
let transport = Arc::new(RecordingTokenClient::default());
|
||||
|
|
@ -582,8 +998,12 @@ mod tests {
|
|||
"https://login.test",
|
||||
);
|
||||
|
||||
let first = acquirer.acquire(request.clone()).await.unwrap();
|
||||
let second = acquirer.acquire(request).await.unwrap();
|
||||
let first = acquirer
|
||||
.acquire(request.clone())
|
||||
.await
|
||||
.into_result()
|
||||
.unwrap();
|
||||
let second = acquirer.acquire(request).await.into_result().unwrap();
|
||||
|
||||
assert_eq!(first.secret().expose(), "native-token");
|
||||
assert_eq!(second.secret().expose(), "native-token");
|
||||
|
|
@ -599,6 +1019,7 @@ mod tests {
|
|||
);
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[tokio::test]
|
||||
async fn credential_provider_cache_isolates_every_client_secret_identity_field() {
|
||||
let transport = Arc::new(RecordingTokenClient::default());
|
||||
|
|
@ -645,10 +1066,10 @@ mod tests {
|
|||
),
|
||||
];
|
||||
|
||||
acquirer.acquire(base.clone()).await.unwrap();
|
||||
acquirer.acquire(base).await.unwrap();
|
||||
acquirer.acquire(base.clone()).await.into_result().unwrap();
|
||||
acquirer.acquire(base).await.into_result().unwrap();
|
||||
for request in variants.into_iter().skip(1) {
|
||||
acquirer.acquire(request).await.unwrap();
|
||||
acquirer.acquire(request).await.into_result().unwrap();
|
||||
}
|
||||
|
||||
assert_eq!(transport.requests.lock().unwrap().len(), 6);
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
use crate::native::CredentialAttempt;
|
||||
use litellm_auth_types::{
|
||||
CredentialFileRef, CredentialLookup, CredentialRef, InputSource, ResolvedCredential, Sourced,
|
||||
TokenProviderHandle,
|
||||
|
|
@ -12,25 +13,31 @@ use std::sync::Arc;
|
|||
use super::native::{NativeAzureRequest, NativeAzureTokenAcquirer, ValidatedAzureRequest};
|
||||
use super::types::{AzureAuthInputs, AzureCredentialType, ConfigValue, DEFAULT_AZURE_SCOPE};
|
||||
|
||||
const AZURE_AD_TOKEN_ENV: &str = "AZURE_AD_TOKEN";
|
||||
const AZURE_TENANT_ID_ENV: &str = "AZURE_TENANT_ID";
|
||||
const AZURE_CLIENT_ID_ENV: &str = "AZURE_CLIENT_ID";
|
||||
const AZURE_CLIENT_SECRET_ENV: &str = "AZURE_CLIENT_SECRET";
|
||||
const AZURE_SCOPE_ENV: &str = "AZURE_SCOPE";
|
||||
const AZURE_AUTHORITY_HOST_ENV: &str = "AZURE_AUTHORITY_HOST";
|
||||
const AZURE_CREDENTIAL_ENV: &str = "AZURE_CREDENTIAL";
|
||||
const AZURE_FEDERATED_TOKEN_FILE_ENV: &str = "AZURE_FEDERATED_TOKEN_FILE";
|
||||
#[derive(strum::IntoStaticStr, strum::VariantArray, strum::VariantNames)]
|
||||
enum AzureSecretName {
|
||||
#[strum(serialize = "AZURE_AD_TOKEN")]
|
||||
AdToken,
|
||||
#[strum(serialize = "AZURE_TENANT_ID")]
|
||||
TenantId,
|
||||
#[strum(serialize = "AZURE_CLIENT_ID")]
|
||||
ClientId,
|
||||
#[strum(serialize = "AZURE_CLIENT_SECRET")]
|
||||
ClientSecret,
|
||||
#[strum(serialize = "AZURE_SCOPE")]
|
||||
Scope,
|
||||
#[strum(serialize = "AZURE_AUTHORITY_HOST")]
|
||||
AuthorityHost,
|
||||
#[strum(serialize = "AZURE_CREDENTIAL")]
|
||||
Credential,
|
||||
#[strum(serialize = "AZURE_FEDERATED_TOKEN_FILE")]
|
||||
FederatedTokenFile,
|
||||
#[strum(serialize = "IDENTITY_ENDPOINT")]
|
||||
IdentityEndpoint,
|
||||
#[strum(serialize = "MSI_ENDPOINT")]
|
||||
MsiEndpoint,
|
||||
}
|
||||
|
||||
pub const SECRET_NAMES: &[&str] = &[
|
||||
AZURE_AD_TOKEN_ENV,
|
||||
AZURE_TENANT_ID_ENV,
|
||||
AZURE_CLIENT_ID_ENV,
|
||||
AZURE_CLIENT_SECRET_ENV,
|
||||
AZURE_SCOPE_ENV,
|
||||
AZURE_AUTHORITY_HOST_ENV,
|
||||
AZURE_CREDENTIAL_ENV,
|
||||
AZURE_FEDERATED_TOKEN_FILE_ENV,
|
||||
];
|
||||
pub const SECRET_NAMES: &[&str] = <AzureSecretName as strum::VariantNames>::VARIANTS;
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub(crate) enum AzureCredentialPlan {
|
||||
|
|
@ -57,14 +64,14 @@ trait AzureTokenAcquirer: Send + Sync {
|
|||
fn acquire(
|
||||
&self,
|
||||
request: ValidatedAzureRequest,
|
||||
) -> Pin<Box<dyn Future<Output = Result<ResolvedCredential, Error>> + Send + '_>>;
|
||||
) -> Pin<Box<dyn Future<Output = CredentialAttempt> + Send + '_>>;
|
||||
}
|
||||
|
||||
impl AzureTokenAcquirer for NativeAzureTokenAcquirer {
|
||||
fn acquire(
|
||||
&self,
|
||||
request: ValidatedAzureRequest,
|
||||
) -> Pin<Box<dyn Future<Output = Result<ResolvedCredential, Error>> + Send + '_>> {
|
||||
) -> Pin<Box<dyn Future<Output = CredentialAttempt> + Send + '_>> {
|
||||
Box::pin(NativeAzureTokenAcquirer::acquire(self, request))
|
||||
}
|
||||
}
|
||||
|
|
@ -83,12 +90,36 @@ impl AzureAuthService {
|
|||
Self { native }
|
||||
}
|
||||
|
||||
pub async fn identity_token(
|
||||
&self,
|
||||
identity: &AzureIdentity,
|
||||
scope: &str,
|
||||
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
|
||||
) -> Result<Option<ResolvedCredential>, Error> {
|
||||
let plan = identity_plan(identity, scope, env_lookup)?;
|
||||
let inputs = AzureAuthInputs::default_credential_for_scope(scope);
|
||||
Ok(self
|
||||
.run_plan(plan, &inputs, env_lookup)
|
||||
.await?
|
||||
.map(Sourced::into_value))
|
||||
}
|
||||
|
||||
pub async fn get_azure_ad_token(
|
||||
&self,
|
||||
inputs: &AzureAuthInputs,
|
||||
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
|
||||
) -> Result<Option<Sourced<ResolvedCredential>>, Error> {
|
||||
match select_auth_plan(inputs, env_lookup)? {
|
||||
self.run_plan(select_auth_plan(inputs, env_lookup)?, inputs, env_lookup)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn run_plan(
|
||||
&self,
|
||||
plan: AzureCredentialPlan,
|
||||
inputs: &AzureAuthInputs,
|
||||
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
|
||||
) -> Result<Option<Sourced<ResolvedCredential>>, Error> {
|
||||
match plan {
|
||||
AzureCredentialPlan::Supplied(credential) => Ok(Some(credential)),
|
||||
AzureCredentialPlan::Caller(caller) => {
|
||||
let credential = caller.acquire().await?;
|
||||
|
|
@ -123,6 +154,7 @@ impl AzureAuthService {
|
|||
self.native
|
||||
.acquire(request)
|
||||
.await
|
||||
.into_result()
|
||||
.map(|credential| Sourced::new(credential, source))
|
||||
.map(Some)
|
||||
}
|
||||
|
|
@ -131,6 +163,7 @@ impl AzureAuthService {
|
|||
self.native
|
||||
.acquire(request)
|
||||
.await
|
||||
.into_result()
|
||||
.map(|credential| Some(Sourced::new(credential, source)))
|
||||
}
|
||||
AzureCredentialPlan::Chain(requests) => {
|
||||
|
|
@ -138,8 +171,16 @@ impl AzureAuthService {
|
|||
for request in requests {
|
||||
let source = request.credential_source();
|
||||
match self.native.acquire(request).await {
|
||||
Ok(credential) => return Ok(Some(Sourced::new(credential, source))),
|
||||
Err(error) => failures.push(error),
|
||||
CredentialAttempt::Resolved(credential) => {
|
||||
return Ok(Some(Sourced::new(credential, source)));
|
||||
}
|
||||
CredentialAttempt::Unavailable(error) => failures.push(error),
|
||||
CredentialAttempt::Failed(error) => {
|
||||
return Err(Error::CredentialChain {
|
||||
last: Box::new(error),
|
||||
earlier: failures,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
match failures.pop() {
|
||||
|
|
@ -155,39 +196,161 @@ impl AzureAuthService {
|
|||
}
|
||||
}
|
||||
|
||||
/// The identity azure-identity builds for one explicit credential class, without LiteLLM's provider settings
|
||||
/// (`AZURE_AD_TOKEN`, `AZURE_CREDENTIAL`, `AZURE_SCOPE`) that `get_azure_ad_token` also reads.
|
||||
#[derive(Clone, Debug)]
|
||||
pub enum AzureIdentity {
|
||||
ClientSecret {
|
||||
tenant_id: String,
|
||||
client_id: String,
|
||||
client_secret: SecretValue,
|
||||
},
|
||||
ManagedIdentity {
|
||||
client_id: String,
|
||||
},
|
||||
Default,
|
||||
}
|
||||
|
||||
fn identity_plan(
|
||||
identity: &AzureIdentity,
|
||||
scope: &str,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> Result<AzureCredentialPlan, Error> {
|
||||
let deployment = |value: &str| Sourced::new(value.to_owned(), InputSource::Deployment);
|
||||
let environment = |name: &str| {
|
||||
env_lookup(name)
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
.map(|value| Sourced::new(value, InputSource::Environment))
|
||||
};
|
||||
let scope = deployment(scope);
|
||||
let authority = environment(AzureSecretName::AuthorityHost.into());
|
||||
match identity {
|
||||
AzureIdentity::ClientSecret {
|
||||
tenant_id,
|
||||
client_id,
|
||||
client_secret,
|
||||
} => Ok(AzureCredentialPlan::Native(ValidatedAzureRequest::new(
|
||||
NativeAzureRequest::ClientSecret {
|
||||
tenant_id: deployment(tenant_id),
|
||||
client_id: deployment(client_id),
|
||||
client_secret: Sourced::new(client_secret.clone(), InputSource::Deployment),
|
||||
scope,
|
||||
authority,
|
||||
},
|
||||
)?)),
|
||||
AzureIdentity::ManagedIdentity { client_id } => Ok(AzureCredentialPlan::Native(
|
||||
ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity {
|
||||
client_id: Some(deployment(client_id)),
|
||||
scope,
|
||||
selection_source: InputSource::Deployment,
|
||||
})?,
|
||||
)),
|
||||
AzureIdentity::Default => {
|
||||
let tenant_id = environment(AzureSecretName::TenantId.into());
|
||||
let client_id = environment(AzureSecretName::ClientId.into());
|
||||
let from_environment = match (
|
||||
tenant_id.clone(),
|
||||
client_id.clone(),
|
||||
environment(AzureSecretName::ClientSecret.into()),
|
||||
) {
|
||||
(Some(tenant_id), Some(client_id), Some(client_secret)) => {
|
||||
Some(NativeAzureRequest::ClientSecret {
|
||||
tenant_id,
|
||||
client_id,
|
||||
client_secret: client_secret.map(SecretValue::new),
|
||||
scope: scope.clone(),
|
||||
authority: authority.clone(),
|
||||
})
|
||||
}
|
||||
_ => None,
|
||||
};
|
||||
let probe_imds = env_lookup(AzureSecretName::IdentityEndpoint.into()).is_none()
|
||||
&& env_lookup(AzureSecretName::MsiEndpoint.into()).is_none();
|
||||
let AzureCredentialPlan::Chain(chain) = select_native_plan(
|
||||
Some(Sourced::new(
|
||||
AzureCredentialType::DefaultAzureCredential,
|
||||
InputSource::Deployment,
|
||||
)),
|
||||
tenant_id,
|
||||
client_id,
|
||||
environment(AzureSecretName::FederatedTokenFile.into()),
|
||||
scope,
|
||||
authority,
|
||||
InputSource::Environment,
|
||||
)?
|
||||
else {
|
||||
unreachable!("DefaultAzureCredential always plans a chain")
|
||||
};
|
||||
let environment_request = from_environment
|
||||
.map(ValidatedAzureRequest::new)
|
||||
.transpose()?;
|
||||
Ok(AzureCredentialPlan::Chain(
|
||||
environment_request
|
||||
.into_iter()
|
||||
.chain(chain)
|
||||
.map(|request| request.with_imds_probe(probe_imds))
|
||||
.collect(),
|
||||
))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn select_auth_plan(
|
||||
inputs: &AzureAuthInputs,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> Result<AzureCredentialPlan, Error> {
|
||||
let token = configured_secret(&inputs.azure_ad_token, AZURE_AD_TOKEN_ENV, env_lookup);
|
||||
let tenant_id = configured_string(&inputs.tenant_id, AZURE_TENANT_ID_ENV, env_lookup);
|
||||
let client_id = configured_string(&inputs.client_id, AZURE_CLIENT_ID_ENV, env_lookup);
|
||||
let client_secret =
|
||||
configured_secret(&inputs.client_secret, AZURE_CLIENT_SECRET_ENV, env_lookup);
|
||||
let scope = configured_string(&inputs.azure_scope, AZURE_SCOPE_ENV, env_lookup)
|
||||
.unwrap_or_else(|| Sourced::new(DEFAULT_AZURE_SCOPE.to_string(), InputSource::Environment));
|
||||
let authority = configured_string(
|
||||
&inputs.azure_authority_host,
|
||||
AZURE_AUTHORITY_HOST_ENV,
|
||||
let token = configured_secret(
|
||||
&inputs.azure_ad_token,
|
||||
AzureSecretName::AdToken.into(),
|
||||
env_lookup,
|
||||
);
|
||||
let selector = configured_string(&inputs.azure_credential, AZURE_CREDENTIAL_ENV, env_lookup)
|
||||
.map(|value| {
|
||||
value
|
||||
.value()
|
||||
.parse::<AzureCredentialType>()
|
||||
.map(|selector| Sourced::new(selector, value.source()))
|
||||
})
|
||||
.transpose()
|
||||
.map_err(|error| {
|
||||
Error::InvalidConfiguration(ErrorDetail::failed(
|
||||
"Azure credential selector parsing",
|
||||
error,
|
||||
))
|
||||
})?;
|
||||
let tenant_id = configured_string(
|
||||
&inputs.tenant_id,
|
||||
AzureSecretName::TenantId.into(),
|
||||
env_lookup,
|
||||
);
|
||||
let client_id = configured_string(
|
||||
&inputs.client_id,
|
||||
AzureSecretName::ClientId.into(),
|
||||
env_lookup,
|
||||
);
|
||||
let client_secret = configured_secret(
|
||||
&inputs.client_secret,
|
||||
AzureSecretName::ClientSecret.into(),
|
||||
env_lookup,
|
||||
);
|
||||
let scope = configured_string(
|
||||
&inputs.azure_scope,
|
||||
AzureSecretName::Scope.into(),
|
||||
env_lookup,
|
||||
)
|
||||
.unwrap_or_else(|| Sourced::new(DEFAULT_AZURE_SCOPE.to_string(), InputSource::Environment));
|
||||
let authority = configured_string(
|
||||
&inputs.azure_authority_host,
|
||||
AzureSecretName::AuthorityHost.into(),
|
||||
env_lookup,
|
||||
);
|
||||
let selector = configured_string(
|
||||
&inputs.azure_credential,
|
||||
AzureSecretName::Credential.into(),
|
||||
env_lookup,
|
||||
)
|
||||
.map(|value| {
|
||||
value
|
||||
.value()
|
||||
.parse::<AzureCredentialType>()
|
||||
.map(|selector| Sourced::new(selector, value.source()))
|
||||
})
|
||||
.transpose()
|
||||
.map_err(|error| {
|
||||
Error::InvalidConfiguration(ErrorDetail::failed(
|
||||
"Azure credential selector parsing",
|
||||
error,
|
||||
))
|
||||
})?;
|
||||
let federated_token_file = configured_string(
|
||||
&inputs.federated_token_file,
|
||||
AZURE_FEDERATED_TOKEN_FILE_ENV,
|
||||
AzureSecretName::FederatedTokenFile.into(),
|
||||
env_lookup,
|
||||
);
|
||||
|
||||
|
|
@ -237,7 +400,12 @@ pub(crate) fn select_auth_plan(
|
|||
return Ok(AzureCredentialPlan::Missing);
|
||||
}
|
||||
|
||||
select_native_plan(
|
||||
let probe_imds = selector
|
||||
.as_ref()
|
||||
.is_none_or(|selector| *selector.value() == AzureCredentialType::DefaultAzureCredential)
|
||||
&& env_lookup(AzureSecretName::IdentityEndpoint.into()).is_none()
|
||||
&& env_lookup(AzureSecretName::MsiEndpoint.into()).is_none();
|
||||
match select_native_plan(
|
||||
selector,
|
||||
tenant_id,
|
||||
client_id,
|
||||
|
|
@ -245,7 +413,15 @@ pub(crate) fn select_auth_plan(
|
|||
scope,
|
||||
authority,
|
||||
inputs.enable_azure_ad_token_refresh.source(),
|
||||
)
|
||||
)? {
|
||||
AzureCredentialPlan::Chain(requests) => Ok(AzureCredentialPlan::Chain(
|
||||
requests
|
||||
.into_iter()
|
||||
.map(|request| request.with_imds_probe(probe_imds))
|
||||
.collect(),
|
||||
)),
|
||||
plan => Ok(plan),
|
||||
}
|
||||
}
|
||||
|
||||
fn select_native_plan(
|
||||
|
|
@ -509,7 +685,10 @@ mod tests {
|
|||
#[derive(Debug)]
|
||||
struct FileResolver;
|
||||
|
||||
use crate::native::CredentialAttempt;
|
||||
|
||||
struct ChainAcquirer {
|
||||
terminal_on: Option<&'static str>,
|
||||
requests: Mutex<Vec<&'static str>>,
|
||||
succeed_on: Option<&'static str>,
|
||||
}
|
||||
|
|
@ -518,20 +697,25 @@ mod tests {
|
|||
fn acquire(
|
||||
&self,
|
||||
request: ValidatedAzureRequest,
|
||||
) -> std::pin::Pin<Box<dyn Future<Output = Result<ResolvedCredential, Error>> + Send + '_>>
|
||||
{
|
||||
) -> std::pin::Pin<Box<dyn Future<Output = CredentialAttempt> + Send + '_>> {
|
||||
let kind = request.kind();
|
||||
self.requests.lock().unwrap().push(kind);
|
||||
Box::pin(async move {
|
||||
if self.succeed_on == Some(kind) {
|
||||
Ok(ResolvedCredential::AccessToken {
|
||||
CredentialAttempt::Resolved(ResolvedCredential::AccessToken {
|
||||
token: SecretValue::new("chain-token"),
|
||||
expires_on: None,
|
||||
})
|
||||
} else {
|
||||
Err(Error::CredentialAcquisition(
|
||||
litellm_auth_types::ErrorDetail::Unresolved { subject: kind },
|
||||
))
|
||||
let error =
|
||||
Error::CredentialAcquisition(litellm_auth_types::ErrorDetail::Unresolved {
|
||||
subject: kind,
|
||||
});
|
||||
if self.terminal_on == Some(kind) {
|
||||
CredentialAttempt::Failed(error)
|
||||
} else {
|
||||
CredentialAttempt::Unavailable(error)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
|
@ -690,9 +874,210 @@ mod tests {
|
|||
assert_eq!(resolved, Some(SecretValue::new("rotated-assertion")));
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct UnavailableImdsClient;
|
||||
|
||||
impl azure_core::http::HttpClient for UnavailableImdsClient {
|
||||
fn execute_request<'life0, 'life1, 'async_trait>(
|
||||
&'life0 self,
|
||||
_request: &'life1 azure_core::http::Request,
|
||||
) -> std::pin::Pin<
|
||||
Box<
|
||||
dyn Future<Output = azure_core::Result<azure_core::http::AsyncRawResponse>>
|
||||
+ Send
|
||||
+ 'async_trait,
|
||||
>,
|
||||
>
|
||||
where
|
||||
'life0: 'async_trait,
|
||||
'life1: 'async_trait,
|
||||
Self: 'async_trait,
|
||||
{
|
||||
Box::pin(std::future::pending())
|
||||
}
|
||||
}
|
||||
|
||||
struct OfflineChainAcquirer(super::NativeAzureTokenAcquirer);
|
||||
|
||||
impl AzureTokenAcquirer for OfflineChainAcquirer {
|
||||
fn acquire(
|
||||
&self,
|
||||
request: ValidatedAzureRequest,
|
||||
) -> std::pin::Pin<Box<dyn Future<Output = CredentialAttempt> + Send + '_>> {
|
||||
Box::pin(async move {
|
||||
if request.kind() == "developer-tools" {
|
||||
return CredentialAttempt::Resolved(ResolvedCredential::AccessToken {
|
||||
token: SecretValue::new("cli-token"),
|
||||
expires_on: None,
|
||||
});
|
||||
}
|
||||
self.0.acquire(request).await
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::default_off_azure("DefaultAzureCredential", None, true)]
|
||||
#[case::explicit_managed_identity("ManagedIdentityCredential", None, false)]
|
||||
#[case::app_service("DefaultAzureCredential", Some("http://app-service.test"), false)]
|
||||
#[tokio::test(start_paused = true)]
|
||||
async fn only_default_imds_discovery_is_bounded(
|
||||
#[case] selector: &str,
|
||||
#[case] identity_endpoint: Option<&str>,
|
||||
#[case] expect_cli: bool,
|
||||
) {
|
||||
let native = super::NativeAzureTokenAcquirer::with_transport(
|
||||
4,
|
||||
azure_core::http::Transport::new(Arc::new(UnavailableImdsClient)),
|
||||
);
|
||||
let service = AzureAuthService::with_acquirer(Arc::new(OfflineChainAcquirer(native)));
|
||||
let inputs = AzureAuthInputs {
|
||||
azure_credential: crate::types::ConfigValue::Value(Sourced::new(
|
||||
selector.into(),
|
||||
InputSource::Deployment,
|
||||
)),
|
||||
enable_azure_ad_token_refresh: Sourced::new(true, InputSource::Deployment),
|
||||
..Default::default()
|
||||
};
|
||||
let environment = |name: &str| {
|
||||
(name == "IDENTITY_ENDPOINT")
|
||||
.then_some(identity_endpoint)
|
||||
.flatten()
|
||||
.map(str::to_string)
|
||||
};
|
||||
let outcome = tokio::time::timeout(
|
||||
std::time::Duration::from_secs(2),
|
||||
service.get_azure_ad_token(&inputs, &environment),
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(outcome.is_ok(), expect_cli);
|
||||
if expect_cli {
|
||||
assert_eq!(
|
||||
outcome.unwrap().unwrap().unwrap().value().secret().expose(),
|
||||
"cli-token"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::app_service(Some("IDENTITY_ENDPOINT"), false)]
|
||||
#[case::legacy_app_service(Some("MSI_ENDPOINT"), false)]
|
||||
#[case::off_azure(None, true)]
|
||||
#[tokio::test(start_paused = true)]
|
||||
async fn default_chain_uses_endpoints_from_the_secret_snapshot(
|
||||
#[case] endpoint: Option<&str>,
|
||||
#[case] expect_cli: bool,
|
||||
#[values(false, true)] consumer_identity: bool,
|
||||
) {
|
||||
use litellm_secrets::source::SecretSource;
|
||||
|
||||
struct EndpointSecrets<'a>(Option<&'a str>);
|
||||
|
||||
impl SecretSource for EndpointSecrets<'_> {
|
||||
fn get_secret_str<'a>(
|
||||
&'a self,
|
||||
name: &'a str,
|
||||
) -> std::pin::Pin<
|
||||
Box<
|
||||
dyn Future<Output = Result<Option<SecretValue>, litellm_secrets::Error>>
|
||||
+ Send
|
||||
+ 'a,
|
||||
>,
|
||||
> {
|
||||
Box::pin(async move {
|
||||
Ok((self.0 == Some(name)).then(|| SecretValue::new("http://app-service.test")))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
let snapshot = EndpointSecrets(endpoint)
|
||||
.resolve(SECRET_NAMES)
|
||||
.await
|
||||
.unwrap();
|
||||
let native = super::NativeAzureTokenAcquirer::with_transport(
|
||||
4,
|
||||
azure_core::http::Transport::new(Arc::new(UnavailableImdsClient)),
|
||||
);
|
||||
let service = AzureAuthService::with_acquirer(Arc::new(OfflineChainAcquirer(native)));
|
||||
let lookup = |name: &str| snapshot.get(name);
|
||||
let scope = "https://service.test/.default";
|
||||
let outcome = tokio::time::timeout(std::time::Duration::from_secs(2), async {
|
||||
if consumer_identity {
|
||||
service
|
||||
.identity_token(&super::AzureIdentity::Default, scope, &lookup)
|
||||
.await
|
||||
} else {
|
||||
service
|
||||
.get_azure_ad_token(
|
||||
&AzureAuthInputs::default_credential_for_scope(scope),
|
||||
&lookup,
|
||||
)
|
||||
.await
|
||||
.map(|token| token.map(Sourced::into_value))
|
||||
}
|
||||
})
|
||||
.await;
|
||||
|
||||
assert_eq!(outcome.is_ok(), expect_cli);
|
||||
if expect_cli {
|
||||
assert_eq!(
|
||||
outcome.unwrap().unwrap().unwrap().secret().expose(),
|
||||
"cli-token"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::litellm_provider_settings_are_ignored(
|
||||
&[("AZURE_AD_TOKEN", "provider-token"), ("AZURE_CREDENTIAL", "ClientSecretCredential")],
|
||||
&["managed-identity", "developer-tools"]
|
||||
)]
|
||||
#[case::environment_client_secret_comes_first(
|
||||
&[("AZURE_TENANT_ID", "tenant"), ("AZURE_CLIENT_ID", "client"), ("AZURE_CLIENT_SECRET", "secret")],
|
||||
&["client-secret", "managed-identity", "developer-tools"]
|
||||
)]
|
||||
#[case::aks_workload_identity(
|
||||
&[("AZURE_TENANT_ID", "tenant"), ("AZURE_CLIENT_ID", "client"), ("AZURE_FEDERATED_TOKEN_FILE", "/var/run/token")],
|
||||
&["workload-identity", "managed-identity", "developer-tools"]
|
||||
)]
|
||||
#[tokio::test]
|
||||
async fn the_default_identity_follows_azure_identity_defaults(
|
||||
#[case] environment: &[(&str, &str)],
|
||||
#[case] expected: &[&str],
|
||||
) {
|
||||
let acquirer = Arc::new(ChainAcquirer {
|
||||
terminal_on: None,
|
||||
requests: Mutex::new(Vec::new()),
|
||||
succeed_on: Some("developer-tools"),
|
||||
});
|
||||
let service = AzureAuthService::with_acquirer(acquirer.clone());
|
||||
let lookup = |name: &str| {
|
||||
environment
|
||||
.iter()
|
||||
.find(|(key, _)| *key == name)
|
||||
.map(|(_, value)| (*value).to_owned())
|
||||
};
|
||||
|
||||
let credential = service
|
||||
.identity_token(
|
||||
&super::AzureIdentity::Default,
|
||||
"https://redis.azure.com/.default",
|
||||
&lookup,
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(credential.secret().expose(), "chain-token");
|
||||
assert_eq!(*acquirer.requests.lock().unwrap(), expected);
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[tokio::test]
|
||||
async fn default_chain_uses_declared_order_and_stops_after_success() {
|
||||
let acquirer = Arc::new(ChainAcquirer {
|
||||
terminal_on: None,
|
||||
requests: Mutex::new(Vec::new()),
|
||||
succeed_on: Some("developer-tools"),
|
||||
});
|
||||
|
|
@ -715,9 +1100,13 @@ mod tests {
|
|||
);
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::all_unavailable(false)]
|
||||
#[case::terminal_failure_after_unavailable(true)]
|
||||
#[tokio::test]
|
||||
async fn chain_reports_each_acquisition_failure() {
|
||||
async fn chain_reports_each_acquisition_failure(#[case] terminal: bool) {
|
||||
let acquirer = Arc::new(ChainAcquirer {
|
||||
terminal_on: terminal.then_some("developer-tools"),
|
||||
requests: Mutex::new(Vec::new()),
|
||||
succeed_on: None,
|
||||
});
|
||||
|
|
@ -744,6 +1133,27 @@ mod tests {
|
|||
);
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[tokio::test]
|
||||
async fn a_configured_identity_failure_does_not_fall_back_to_a_developer_identity() {
|
||||
let acquirer = Arc::new(ChainAcquirer {
|
||||
terminal_on: Some("managed-identity"),
|
||||
requests: Mutex::new(Vec::new()),
|
||||
succeed_on: Some("developer-tools"),
|
||||
});
|
||||
let service = AzureAuthService::with_acquirer(acquirer.clone());
|
||||
let inputs = AzureAuthInputs {
|
||||
enable_azure_ad_token_refresh: Sourced::new(true, InputSource::Deployment),
|
||||
..Default::default()
|
||||
};
|
||||
let error = service
|
||||
.get_azure_ad_token(&inputs, &|_| None)
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert_eq!(*acquirer.requests.lock().unwrap(), ["managed-identity"]);
|
||||
assert!(matches!(error, Error::CredentialChain { earlier, .. } if earlier.is_empty()));
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct CallerToken(&'static str);
|
||||
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ license.workspace = true
|
|||
repository.workspace = true
|
||||
|
||||
[features]
|
||||
google-sdk = ["dep:google-cloud-auth", "dep:http"]
|
||||
google-sdk = ["dep:http"]
|
||||
|
||||
[dependencies]
|
||||
litellm-secret-value.workspace = true
|
||||
|
|
@ -21,10 +21,14 @@ tokio.workspace = true
|
|||
veil.workspace = true
|
||||
|
||||
gcp_auth = "0.12.7"
|
||||
google-cloud-auth = { workspace = true, optional = true }
|
||||
google-cloud-auth.workspace = true
|
||||
http = { workspace = true, optional = true }
|
||||
|
||||
[dev-dependencies]
|
||||
time = { workspace = true, features = ["formatting"] }
|
||||
axum = { workspace = true, features = ["json"] }
|
||||
litellm-error.workspace = true
|
||||
rstest.workspace = true
|
||||
rustls.workspace = true
|
||||
tempfile.workspace = true
|
||||
tokio = { workspace = true, features = ["test-util"] }
|
||||
|
|
|
|||
|
|
@ -1,20 +1,11 @@
|
|||
use std::{future::Future, pin::Pin, sync::Arc};
|
||||
use std::sync::Arc;
|
||||
|
||||
use gcp_auth::{CustomServiceAccount, TokenProvider};
|
||||
use litellm_auth_types::{CredentialPlacement, Error, ErrorDetail, http::apply_credential};
|
||||
use moka::future::Cache;
|
||||
use serde_json::Value;
|
||||
use sha2::{Digest, Sha256};
|
||||
use litellm_auth_types::{CredentialPlacement, Error, http::apply_credential};
|
||||
|
||||
use crate::{
|
||||
config::{
|
||||
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,
|
||||
VERTEXAI_API_KEY_ENV,
|
||||
},
|
||||
config::{VertexConfig, credential_source, get_vertex_ai_project, non_empty_env},
|
||||
constants::{VERTEX_AI_API_KEY_ENV, VERTEXAI_API_KEY_ENV},
|
||||
identity::{GcpAuth, GcpProviderLoader, GcpTokenSource},
|
||||
};
|
||||
|
||||
pub struct VertexEnvironment {
|
||||
|
|
@ -22,28 +13,15 @@ pub struct VertexEnvironment {
|
|||
pub project_id: String,
|
||||
}
|
||||
|
||||
struct VertexAccessToken {
|
||||
token: String,
|
||||
project_id: String,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
#[derive(Clone, Default)]
|
||||
pub struct VertexAuth {
|
||||
providers: Cache<CredentialCacheKey, Arc<dyn VertexTokenSource>>,
|
||||
loader: Arc<dyn VertexProviderLoader>,
|
||||
}
|
||||
|
||||
impl Default for VertexAuth {
|
||||
fn default() -> Self {
|
||||
Self::new(Arc::new(GcpProviderLoader))
|
||||
}
|
||||
identity: GcpAuth,
|
||||
}
|
||||
|
||||
impl VertexAuth {
|
||||
pub fn new(loader: Arc<dyn VertexProviderLoader>) -> Self {
|
||||
pub fn new(loader: Arc<dyn GcpProviderLoader>) -> Self {
|
||||
Self {
|
||||
providers: Cache::builder().max_capacity(64).build(),
|
||||
loader,
|
||||
identity: GcpAuth::new(loader),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -52,7 +30,10 @@ impl VertexAuth {
|
|||
config: &VertexConfig,
|
||||
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
|
||||
) -> Result<String, Error> {
|
||||
self.load_provider(config, env_lookup).await?.token().await
|
||||
self.identity
|
||||
.access_token(credential_source(config, env_lookup))
|
||||
.await
|
||||
.map(|token| token.expose().to_owned())
|
||||
}
|
||||
|
||||
pub async fn validate_environment(
|
||||
|
|
@ -77,10 +58,18 @@ impl VertexAuth {
|
|||
(true, _) => None,
|
||||
(false, Some(token)) => Some(token),
|
||||
(false, None) => {
|
||||
let access = self.get_access_token(config, env_lookup).await?;
|
||||
let provider = self.load_provider(config, env_lookup).await?;
|
||||
let (token, project_id) = match project_id {
|
||||
Some(project_id) => (provider.token().await?, project_id),
|
||||
None => tokio::try_join!(provider.token(), provider.project_id())?,
|
||||
};
|
||||
return Ok(VertexEnvironment {
|
||||
headers: apply_credential(headers, &access.token, CredentialPlacement::Bearer)?,
|
||||
project_id: project_id.unwrap_or(access.project_id),
|
||||
headers: apply_credential(
|
||||
headers,
|
||||
token.expose(),
|
||||
CredentialPlacement::Bearer,
|
||||
)?,
|
||||
project_id,
|
||||
});
|
||||
}
|
||||
};
|
||||
|
|
@ -104,248 +93,171 @@ impl VertexAuth {
|
|||
})
|
||||
}
|
||||
|
||||
async fn get_access_token(
|
||||
&self,
|
||||
config: &VertexConfig,
|
||||
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
|
||||
) -> Result<VertexAccessToken, Error> {
|
||||
let provider = self.load_provider(config, env_lookup).await?;
|
||||
let (token, project_id) = tokio::try_join!(provider.token(), provider.project_id())?;
|
||||
Ok(VertexAccessToken { token, project_id })
|
||||
}
|
||||
|
||||
async fn load_provider(
|
||||
&self,
|
||||
config: &VertexConfig,
|
||||
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
|
||||
) -> Result<Arc<dyn VertexTokenSource>, Error> {
|
||||
let source = credential_source(config, env_lookup);
|
||||
let key = source.cache_key();
|
||||
self.providers
|
||||
.try_get_with(key, self.loader.load(source))
|
||||
) -> Result<Arc<dyn GcpTokenSource>, Error> {
|
||||
self.identity
|
||||
.load_provider(credential_source(config, env_lookup))
|
||||
.await
|
||||
.map_err(|error| (*error).clone())
|
||||
}
|
||||
}
|
||||
|
||||
pub trait VertexTokenSource: Send + Sync {
|
||||
fn project_id(&self) -> VertexAuthFuture<'_, String>;
|
||||
fn token(&self) -> VertexAuthFuture<'_, String>;
|
||||
}
|
||||
|
||||
pub trait VertexProviderLoader: Send + Sync {
|
||||
fn load(&self, source: CredentialSource) -> VertexAuthFuture<'_, Arc<dyn VertexTokenSource>>;
|
||||
}
|
||||
|
||||
pub type VertexAuthFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, Error>> + Send + 'a>>;
|
||||
|
||||
struct GcpTokenSource(Arc<dyn TokenProvider>);
|
||||
|
||||
impl VertexTokenSource for GcpTokenSource {
|
||||
fn project_id(&self) -> VertexAuthFuture<'_, String> {
|
||||
Box::pin(async move {
|
||||
self.0
|
||||
.project_id()
|
||||
.await
|
||||
.map(|project| project.to_string())
|
||||
.map_err(auth_acquisition_error)
|
||||
})
|
||||
}
|
||||
|
||||
fn token(&self) -> VertexAuthFuture<'_, String> {
|
||||
Box::pin(async move {
|
||||
self.0
|
||||
.token(&[CLOUD_PLATFORM_SCOPE])
|
||||
.await
|
||||
.map(|token| token.as_str().to_string())
|
||||
.map_err(auth_acquisition_error)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
struct GcpProviderLoader;
|
||||
|
||||
impl VertexProviderLoader for GcpProviderLoader {
|
||||
fn load(&self, source: CredentialSource) -> VertexAuthFuture<'_, Arc<dyn VertexTokenSource>> {
|
||||
Box::pin(async move {
|
||||
let provider: Arc<dyn TokenProvider> = match source {
|
||||
CredentialSource::Inline(configured) => Arc::new(
|
||||
CustomServiceAccount::from_json(validate_request_credentials(
|
||||
configured.expose(),
|
||||
)?)
|
||||
.map_err(auth_acquisition_error)?,
|
||||
),
|
||||
CredentialSource::Trusted(configured) => {
|
||||
Arc::new(load_trusted_service_account(configured.expose())?)
|
||||
}
|
||||
CredentialSource::ApplicationCredentials(path) => {
|
||||
Arc::new(CustomServiceAccount::from_file(path).map_err(auth_acquisition_error)?)
|
||||
}
|
||||
CredentialSource::Adc => {
|
||||
gcp_auth::provider().await.map_err(auth_acquisition_error)?
|
||||
}
|
||||
};
|
||||
Ok(Arc::new(GcpTokenSource(provider)) as Arc<dyn VertexTokenSource>)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn validate_request_credentials(configured: &str) -> Result<&str, Error> {
|
||||
let token_uri = serde_json::from_str::<Value>(configured)
|
||||
.ok()
|
||||
.and_then(|credentials| {
|
||||
credentials
|
||||
.get("token_uri")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::to_string)
|
||||
});
|
||||
if token_uri.as_deref() != Some(GOOGLE_OAUTH_TOKEN_ENDPOINT) {
|
||||
return Err(Error::InvalidConfiguration(ErrorDetail::InvalidType {
|
||||
field: "request-controlled vertex_credentials token_uri".into(),
|
||||
expected: "the canonical Google OAuth token endpoint",
|
||||
}));
|
||||
}
|
||||
Ok(configured)
|
||||
}
|
||||
|
||||
impl CredentialSource {
|
||||
fn cache_key(&self) -> CredentialCacheKey {
|
||||
match self {
|
||||
Self::Inline(configured) => {
|
||||
CredentialCacheKey::Inline(Sha256::digest(configured.expose()).into())
|
||||
}
|
||||
Self::Trusted(configured) => {
|
||||
CredentialCacheKey::Trusted(Sha256::digest(configured.expose()).into())
|
||||
}
|
||||
Self::ApplicationCredentials(path) => {
|
||||
CredentialCacheKey::ApplicationCredentials(path.clone())
|
||||
}
|
||||
Self::Adc => CredentialCacheKey::Adc,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Hash, PartialEq, Eq)]
|
||||
enum CredentialCacheKey {
|
||||
Inline([u8; 32]),
|
||||
Trusted([u8; 32]),
|
||||
ApplicationCredentials(String),
|
||||
Adc,
|
||||
}
|
||||
|
||||
fn load_trusted_service_account(configured: &str) -> Result<CustomServiceAccount, Error> {
|
||||
if is_base64_encoded_json(configured) {
|
||||
return Err(Error::InvalidConfiguration(
|
||||
ErrorDetail::Base64EncodedJson {
|
||||
field: "vertex_credentials",
|
||||
},
|
||||
));
|
||||
}
|
||||
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",
|
||||
error,
|
||||
))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::{
|
||||
collections::{BTreeMap, BTreeSet},
|
||||
sync::{
|
||||
Arc,
|
||||
atomic::{AtomicUsize, Ordering},
|
||||
},
|
||||
};
|
||||
|
||||
use crate::{
|
||||
CredentialSource, GcpAuthFuture, GcpProviderLoader, GcpTokenSource, VertexAuth,
|
||||
VertexConfig, get_vertex_ai_location, secret_names,
|
||||
};
|
||||
use litellm_secret_value::SecretValue;
|
||||
use rstest::rstest;
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use super::*;
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::canonical_endpoint(r#"{"token_uri":"https://oauth2.googleapis.com/token"}"#, true)]
|
||||
#[case::noncanonical_endpoint(r#"{"token_uri":"http://127.0.0.1/token"}"#, false)]
|
||||
#[case::missing_endpoint("{}", false)]
|
||||
fn request_credentials_require_canonical_token_endpoint(
|
||||
#[case] credentials: &str,
|
||||
#[case] accepted: bool,
|
||||
) {
|
||||
assert_eq!(validate_request_credentials(credentials).is_ok(), accepted);
|
||||
struct FakeProvider {
|
||||
project: Option<&'static str>,
|
||||
calls: Arc<AtomicUsize>,
|
||||
}
|
||||
|
||||
#[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");
|
||||
};
|
||||
let report = litellm_error::Report::new(&error).to_string();
|
||||
assert!(report.contains(expected), "{report}");
|
||||
}
|
||||
impl GcpTokenSource for FakeProvider {
|
||||
fn project_id(&self) -> GcpAuthFuture<'_, String> {
|
||||
self.calls.fetch_add(1, Ordering::SeqCst);
|
||||
Box::pin(async move {
|
||||
self.project.map(str::to_owned).ok_or(
|
||||
litellm_auth_types::Error::CredentialAcquisition(
|
||||
litellm_auth_types::ErrorDetail::Unresolved {
|
||||
subject: "credential project",
|
||||
},
|
||||
),
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
#[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) => litellm_error::Report::new(&error).to_string(),
|
||||
fn token(&self) -> GcpAuthFuture<'_, SecretValue> {
|
||||
self.calls.fetch_add(1, Ordering::SeqCst);
|
||||
Box::pin(async { Ok(SecretValue::new("adc-token")) })
|
||||
}
|
||||
}
|
||||
|
||||
#[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
|
||||
);
|
||||
struct FakeLoader {
|
||||
loads: Arc<AtomicUsize>,
|
||||
provider: Arc<dyn GcpTokenSource>,
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
fn inline_and_trusted_credentials_never_share_a_cache_entry() {
|
||||
assert_ne!(
|
||||
CredentialSource::Inline(SecretValue::new("same-value")).cache_key(),
|
||||
CredentialSource::Trusted(SecretValue::new("same-value")).cache_key()
|
||||
);
|
||||
impl GcpProviderLoader for FakeLoader {
|
||||
fn load(&self, _source: CredentialSource) -> GcpAuthFuture<'_, Arc<dyn GcpTokenSource>> {
|
||||
let loads = self.loads.clone();
|
||||
let provider = self.provider.clone();
|
||||
Box::pin(async move {
|
||||
loads.fetch_add(1, Ordering::SeqCst);
|
||||
Ok(provider)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn config(value: Value) -> VertexConfig {
|
||||
VertexConfig::from_sourced_optional_params(value.as_object().unwrap(), &BTreeMap::new())
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
fn auth(calls: Arc<AtomicUsize>, loads: Arc<AtomicUsize>) -> VertexAuth {
|
||||
let provider: Arc<dyn GcpTokenSource> = Arc::new(FakeProvider {
|
||||
calls,
|
||||
project: Some("adc-project"),
|
||||
});
|
||||
VertexAuth::new(Arc::new(FakeLoader { loads, provider }))
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn secret_names_cover_environment_reads() {
|
||||
let seen = Arc::new(std::sync::Mutex::new(BTreeSet::<String>::new()));
|
||||
let recorded = seen.clone();
|
||||
let env = |name: &str| {
|
||||
recorded.lock().unwrap().insert(name.to_string());
|
||||
None
|
||||
};
|
||||
let auth = auth(Arc::new(AtomicUsize::new(0)), Arc::new(AtomicUsize::new(0)));
|
||||
auth.validate_environment(Vec::new(), None, &VertexConfig::default(), &env)
|
||||
.await
|
||||
.unwrap();
|
||||
get_vertex_ai_location(&VertexConfig::default(), &env);
|
||||
let allowed = secret_names()
|
||||
.iter()
|
||||
.map(|name| (*name).to_owned())
|
||||
.collect::<BTreeSet<_>>();
|
||||
assert!(seen.lock().unwrap().is_subset(&allowed));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn explicit_token_and_header_do_not_acquire_adc() {
|
||||
let loads = Arc::new(AtomicUsize::new(0));
|
||||
let auth = auth(Arc::new(AtomicUsize::new(0)), loads.clone());
|
||||
let configured = config(json!({"vertex_project":"project-1"}));
|
||||
let explicit = auth
|
||||
.validate_environment(Vec::new(), Some("access-token"), &configured, &|_| None)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(explicit.headers[0].1, "Bearer access-token");
|
||||
let existing = auth
|
||||
.validate_environment(
|
||||
vec![("authorization".into(), "Bearer existing".into())],
|
||||
None,
|
||||
&configured,
|
||||
&|_| None,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(existing.headers[0].1, "Bearer existing");
|
||||
assert_eq!(loads.load(Ordering::SeqCst), 0);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn provider_is_reused_across_authentication_calls() {
|
||||
let calls = Arc::new(AtomicUsize::new(0));
|
||||
let loads = Arc::new(AtomicUsize::new(0));
|
||||
let auth = auth(calls.clone(), loads.clone());
|
||||
let config = VertexConfig::default();
|
||||
let environment = |_: &str| None;
|
||||
let (first, second) = tokio::try_join!(
|
||||
auth.validate_environment(Vec::new(), None, &config, &environment),
|
||||
auth.validate_environment(Vec::new(), None, &config, &environment),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(first.project_id, "adc-project");
|
||||
assert_eq!(first.headers[0].1, "Bearer adc-token");
|
||||
assert_eq!(first.project_id, second.project_id);
|
||||
assert_eq!(first.headers, second.headers);
|
||||
assert_eq!(loads.load(Ordering::SeqCst), 1);
|
||||
assert_eq!(calls.load(Ordering::SeqCst), 4);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn an_explicit_project_does_not_require_project_metadata_in_the_credentials() {
|
||||
let auth = VertexAuth::new(Arc::new(FakeLoader {
|
||||
loads: Arc::new(AtomicUsize::new(0)),
|
||||
provider: Arc::new(FakeProvider {
|
||||
calls: Arc::new(AtomicUsize::new(0)),
|
||||
project: None,
|
||||
}),
|
||||
}));
|
||||
let configured = VertexConfig::new(None, Some("configured-project".into()), None);
|
||||
let environment = auth
|
||||
.validate_environment(Vec::new(), None, &configured, &|_| None)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(environment.project_id, "configured-project");
|
||||
assert_eq!(environment.headers[0].1, "Bearer adc-token");
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -146,6 +146,20 @@ pub enum CredentialSource {
|
|||
Adc,
|
||||
}
|
||||
|
||||
impl CredentialSource {
|
||||
pub fn from_vertex_environment(env_lookup: &dyn Fn(&str) -> Option<String>) -> Self {
|
||||
VertexParams::CREDENTIALS
|
||||
.resolve(&|_| None, env_lookup)
|
||||
.map(|configured| Self::Trusted(SecretValue::new(configured)))
|
||||
.unwrap_or_else(|| Self::application_default(env_lookup))
|
||||
}
|
||||
|
||||
pub fn application_default(env_lookup: &dyn Fn(&str) -> Option<String>) -> Self {
|
||||
non_empty_env(env_lookup, GOOGLE_APPLICATION_CREDENTIALS_ENV)
|
||||
.map_or(Self::Adc, Self::ApplicationCredentials)
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn credential_source(
|
||||
config: &VertexConfig,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
|
|
@ -158,12 +172,7 @@ pub(crate) fn credential_source(
|
|||
}
|
||||
};
|
||||
}
|
||||
if let Some(configured) = VertexParams::CREDENTIALS.resolve(&|_| None, env_lookup) {
|
||||
return CredentialSource::Trusted(SecretValue::new(configured));
|
||||
}
|
||||
non_empty_env(env_lookup, GOOGLE_APPLICATION_CREDENTIALS_ENV)
|
||||
.map(CredentialSource::ApplicationCredentials)
|
||||
.unwrap_or(CredentialSource::Adc)
|
||||
CredentialSource::from_vertex_environment(env_lookup)
|
||||
}
|
||||
|
||||
fn optional_credentials(
|
||||
|
|
@ -319,6 +328,24 @@ mod tests {
|
|||
));
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
fn application_default_ignores_vertex_credentials() {
|
||||
let env = |name: &str| match name {
|
||||
VERTEXAI_CREDENTIALS_ENV => Some("environment-json".to_string()),
|
||||
GOOGLE_APPLICATION_CREDENTIALS_ENV => Some("adc.json".to_string()),
|
||||
_ => None,
|
||||
};
|
||||
assert!(
|
||||
matches!(CredentialSource::application_default(&env), CredentialSource::ApplicationCredentials(path) if path == "adc.json")
|
||||
);
|
||||
assert!(matches!(
|
||||
CredentialSource::application_default(
|
||||
&|name| (name == VERTEXAI_CREDENTIALS_ENV).then(|| "environment-json".into())
|
||||
),
|
||||
CredentialSource::Adc
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
fn params_become_a_deployment_sourced_config() {
|
||||
let params = VertexParams {
|
||||
|
|
|
|||
445
litellm-rust/crates/auth-gcp/src/identity.rs
Normal file
445
litellm-rust/crates/auth-gcp/src/identity.rs
Normal file
|
|
@ -0,0 +1,445 @@
|
|||
use crate::{
|
||||
config::{CredentialSource, is_base64_encoded_json, names_a_credentials_file},
|
||||
constants::{CLOUD_PLATFORM_SCOPE, GOOGLE_OAUTH_TOKEN_ENDPOINT},
|
||||
};
|
||||
use gcp_auth::{CustomServiceAccount, TokenProvider};
|
||||
use litellm_auth_types::{Error, ErrorDetail};
|
||||
use litellm_secret_value::SecretValue;
|
||||
use moka::future::Cache;
|
||||
use serde_json::Value;
|
||||
use sha2::{Digest, Sha256};
|
||||
use std::{future::Future, pin::Pin, sync::Arc};
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct GcpAuth {
|
||||
providers: Cache<CredentialCacheKey, Arc<dyn GcpTokenSource>>,
|
||||
loader: Arc<dyn GcpProviderLoader>,
|
||||
}
|
||||
|
||||
impl Default for GcpAuth {
|
||||
fn default() -> Self {
|
||||
Self::new(Arc::new(NativeProviderLoader))
|
||||
}
|
||||
}
|
||||
|
||||
impl GcpAuth {
|
||||
pub fn new(loader: Arc<dyn GcpProviderLoader>) -> Self {
|
||||
Self {
|
||||
providers: Cache::builder().max_capacity(64).build(),
|
||||
loader,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn access_token(&self, source: CredentialSource) -> Result<SecretValue, Error> {
|
||||
self.load_provider(source).await?.token().await
|
||||
}
|
||||
|
||||
pub(crate) async fn load_provider(
|
||||
&self,
|
||||
source: CredentialSource,
|
||||
) -> Result<Arc<dyn GcpTokenSource>, Error> {
|
||||
self.providers
|
||||
.try_get_with(source.cache_key(), self.loader.load(source))
|
||||
.await
|
||||
.map_err(|error| (*error).clone())
|
||||
}
|
||||
}
|
||||
|
||||
pub trait GcpTokenSource: Send + Sync {
|
||||
fn project_id(&self) -> GcpAuthFuture<'_, String>;
|
||||
fn token(&self) -> GcpAuthFuture<'_, SecretValue>;
|
||||
}
|
||||
|
||||
pub trait GcpProviderLoader: Send + Sync {
|
||||
fn load(&self, source: CredentialSource) -> GcpAuthFuture<'_, Arc<dyn GcpTokenSource>>;
|
||||
}
|
||||
|
||||
pub type GcpAuthFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, Error>> + Send + 'a>>;
|
||||
|
||||
struct ServiceAccountTokenSource(Arc<dyn TokenProvider>);
|
||||
|
||||
impl GcpTokenSource for ServiceAccountTokenSource {
|
||||
fn project_id(&self) -> GcpAuthFuture<'_, String> {
|
||||
Box::pin(async move {
|
||||
self.0
|
||||
.project_id()
|
||||
.await
|
||||
.map(|project| project.to_string())
|
||||
.map_err(auth_acquisition_error)
|
||||
})
|
||||
}
|
||||
|
||||
fn token(&self) -> GcpAuthFuture<'_, SecretValue> {
|
||||
Box::pin(async move {
|
||||
self.0
|
||||
.token(&[CLOUD_PLATFORM_SCOPE])
|
||||
.await
|
||||
.map(|token| SecretValue::new(token.as_str()))
|
||||
.map_err(auth_acquisition_error)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
struct SdkTokenSource {
|
||||
credentials: google_cloud_auth::credentials::AccessTokenCredentials,
|
||||
project_id: Option<String>,
|
||||
}
|
||||
|
||||
impl GcpTokenSource for SdkTokenSource {
|
||||
fn project_id(&self) -> GcpAuthFuture<'_, String> {
|
||||
Box::pin(async move {
|
||||
self.project_id
|
||||
.clone()
|
||||
.ok_or(Error::CredentialAcquisition(ErrorDetail::Unresolved {
|
||||
subject: "GCP resource project",
|
||||
}))
|
||||
})
|
||||
}
|
||||
|
||||
fn token(&self) -> GcpAuthFuture<'_, SecretValue> {
|
||||
Box::pin(async move {
|
||||
self.credentials
|
||||
.access_token()
|
||||
.await
|
||||
.map(|token| SecretValue::new(token.token))
|
||||
.map_err(auth_acquisition_error)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
async fn load_application_credentials(path: &str) -> Result<Arc<dyn GcpTokenSource>, Error> {
|
||||
let contents = tokio::fs::read_to_string(path)
|
||||
.await
|
||||
.map_err(auth_acquisition_error)?;
|
||||
load_credentials_json(&contents)
|
||||
}
|
||||
|
||||
fn load_credentials_json(contents: &str) -> Result<Arc<dyn GcpTokenSource>, Error> {
|
||||
let configured: Value = serde_json::from_str(contents).map_err(auth_acquisition_error)?;
|
||||
let kind = configured.get("type").and_then(Value::as_str);
|
||||
let project_id = configured
|
||||
.get("project_id")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::to_string);
|
||||
let credentials = match kind {
|
||||
Some("authorized_user") => {
|
||||
google_cloud_auth::credentials::user_account::Builder::new(configured)
|
||||
.with_scopes([CLOUD_PLATFORM_SCOPE])
|
||||
.build_access_token_credentials()
|
||||
}
|
||||
Some("impersonated_service_account") => {
|
||||
google_cloud_auth::credentials::impersonated::Builder::new(configured)
|
||||
.with_scopes([CLOUD_PLATFORM_SCOPE])
|
||||
.build_access_token_credentials()
|
||||
}
|
||||
Some("external_account") => {
|
||||
google_cloud_auth::credentials::external_account::Builder::new(configured)
|
||||
.with_scopes([CLOUD_PLATFORM_SCOPE])
|
||||
.build_access_token_credentials()
|
||||
}
|
||||
Some("service_account") | None => {
|
||||
let provider =
|
||||
CustomServiceAccount::from_json(contents).map_err(auth_acquisition_error)?;
|
||||
return Ok(Arc::new(ServiceAccountTokenSource(Arc::new(provider))));
|
||||
}
|
||||
Some(_) => {
|
||||
return Err(Error::InvalidConfiguration(ErrorDetail::InvalidType {
|
||||
field: "Google application credentials type".into(),
|
||||
expected: "a supported Google credential type",
|
||||
}));
|
||||
}
|
||||
}
|
||||
.map_err(auth_acquisition_error)?;
|
||||
Ok(Arc::new(SdkTokenSource {
|
||||
credentials,
|
||||
project_id,
|
||||
}))
|
||||
}
|
||||
|
||||
struct NativeProviderLoader;
|
||||
|
||||
impl GcpProviderLoader for NativeProviderLoader {
|
||||
fn load(&self, source: CredentialSource) -> GcpAuthFuture<'_, Arc<dyn GcpTokenSource>> {
|
||||
Box::pin(async move {
|
||||
let provider: Arc<dyn TokenProvider> = match source {
|
||||
CredentialSource::Inline(configured) => Arc::new(
|
||||
CustomServiceAccount::from_json(validate_request_credentials(
|
||||
configured.expose(),
|
||||
)?)
|
||||
.map_err(auth_acquisition_error)?,
|
||||
),
|
||||
CredentialSource::Trusted(configured) => {
|
||||
return load_trusted_credentials(configured.expose()).await;
|
||||
}
|
||||
CredentialSource::ApplicationCredentials(path) => {
|
||||
return load_application_credentials(&path).await;
|
||||
}
|
||||
CredentialSource::Adc => {
|
||||
gcp_auth::provider().await.map_err(auth_acquisition_error)?
|
||||
}
|
||||
};
|
||||
Ok(Arc::new(ServiceAccountTokenSource(provider)) as Arc<dyn GcpTokenSource>)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn validate_request_credentials(configured: &str) -> Result<&str, Error> {
|
||||
let token_uri = serde_json::from_str::<Value>(configured)
|
||||
.ok()
|
||||
.and_then(|credentials| {
|
||||
credentials
|
||||
.get("token_uri")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::to_string)
|
||||
});
|
||||
if token_uri.as_deref() != Some(GOOGLE_OAUTH_TOKEN_ENDPOINT) {
|
||||
return Err(Error::InvalidConfiguration(ErrorDetail::InvalidType {
|
||||
field: "request-controlled vertex_credentials token_uri".into(),
|
||||
expected: "the canonical Google OAuth token endpoint",
|
||||
}));
|
||||
}
|
||||
Ok(configured)
|
||||
}
|
||||
|
||||
impl CredentialSource {
|
||||
fn cache_key(&self) -> CredentialCacheKey {
|
||||
match self {
|
||||
Self::Inline(configured) => {
|
||||
CredentialCacheKey::Inline(Sha256::digest(configured.expose()).into())
|
||||
}
|
||||
Self::Trusted(configured) => {
|
||||
CredentialCacheKey::Trusted(Sha256::digest(configured.expose()).into())
|
||||
}
|
||||
Self::ApplicationCredentials(path) => {
|
||||
CredentialCacheKey::ApplicationCredentials(path.clone())
|
||||
}
|
||||
Self::Adc => CredentialCacheKey::Adc,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Hash, PartialEq, Eq)]
|
||||
enum CredentialCacheKey {
|
||||
Inline([u8; 32]),
|
||||
Trusted([u8; 32]),
|
||||
ApplicationCredentials(String),
|
||||
Adc,
|
||||
}
|
||||
|
||||
async fn load_trusted_credentials(configured: &str) -> Result<Arc<dyn GcpTokenSource>, Error> {
|
||||
if is_base64_encoded_json(configured) {
|
||||
return Err(Error::InvalidConfiguration(
|
||||
ErrorDetail::Base64EncodedJson {
|
||||
field: "vertex_credentials",
|
||||
},
|
||||
));
|
||||
}
|
||||
if names_a_credentials_file(configured) {
|
||||
load_application_credentials(configured).await
|
||||
} else {
|
||||
load_credentials_json(configured)
|
||||
}
|
||||
}
|
||||
|
||||
fn auth_acquisition_error(error: impl std::error::Error + Send + Sync + 'static) -> Error {
|
||||
Error::CredentialAcquisition(litellm_auth_types::ErrorDetail::failed(
|
||||
"Google credentials",
|
||||
error,
|
||||
))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use litellm_secret_value::SecretValue;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::authorized_user("authorized_user", 2)]
|
||||
#[case::external_account("external_account", 2)]
|
||||
#[case::impersonated_service_account("impersonated_service_account", 4)]
|
||||
#[tokio::test]
|
||||
async fn explicit_application_credentials_refresh_expired_tokens(
|
||||
#[case] kind: &str,
|
||||
#[case] expected_requests: usize,
|
||||
#[values(false, true)] inline: bool,
|
||||
) {
|
||||
use axum::{Json, Router, routing::post};
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
|
||||
let _ = rustls::crypto::ring::default_provider().install_default();
|
||||
let requests = Arc::new(AtomicUsize::new(0));
|
||||
let recorded = requests.clone();
|
||||
let impersonations = requests.clone();
|
||||
let app = Router::new().route("/token", post(move || {
|
||||
let recorded = recorded.clone();
|
||||
async move {
|
||||
Json(serde_json::json!({
|
||||
"access_token": format!("token-{}", recorded.fetch_add(1, Ordering::SeqCst)),
|
||||
"expires_in": 3600,
|
||||
"token_type": "Bearer",
|
||||
"issued_token_type": "urn:ietf:params:oauth:token-type:access_token",
|
||||
}))
|
||||
}
|
||||
}));
|
||||
let app = app.route("/v1/projects/-/serviceAccounts/test-principal:generateAccessToken", post(move || {
|
||||
let recorded = impersonations.clone();
|
||||
async move {
|
||||
let expiration = (time::OffsetDateTime::now_utc() + time::Duration::hours(1)).format(&time::format_description::well_known::Rfc3339).unwrap();
|
||||
Json(serde_json::json!({
|
||||
"accessToken": format!("impersonated-{}", recorded.fetch_add(1, Ordering::SeqCst)),
|
||||
"expireTime": expiration,
|
||||
}))
|
||||
}
|
||||
}));
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let address = listener.local_addr().unwrap();
|
||||
let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
|
||||
let file = tempfile::NamedTempFile::new().unwrap();
|
||||
let subject_token = tempfile::NamedTempFile::new().unwrap();
|
||||
std::fs::write(subject_token.path(), "subject-token").unwrap();
|
||||
let configured = if kind == "external_account" {
|
||||
serde_json::json!({
|
||||
"type": "external_account", "audience": "//iam.googleapis.com/projects/test/locations/global/workloadIdentityPools/test/providers/test",
|
||||
"subject_token_type": "urn:ietf:params:oauth:token-type:jwt", "token_url": format!("http://{address}/token"),
|
||||
"credential_source": { "file": subject_token.path().to_str().unwrap() }, "quota_project_id": "quota-project"
|
||||
})
|
||||
} else {
|
||||
let source = serde_json::json!({
|
||||
"type": "authorized_user", "client_id": "user-client", "client_secret": "user-secret",
|
||||
"refresh_token": "user-refresh", "quota_project_id": "quota-project", "token_uri": format!("http://{address}/token")
|
||||
});
|
||||
if kind == "impersonated_service_account" {
|
||||
serde_json::json!({
|
||||
"type": kind, "source_credentials": source, "quota_project_id": "quota-project",
|
||||
"service_account_impersonation_url": format!("http://{address}/v1/projects/-/serviceAccounts/test-principal:generateAccessToken")
|
||||
})
|
||||
} else {
|
||||
source
|
||||
}
|
||||
};
|
||||
std::fs::write(file.path(), configured.to_string()).unwrap();
|
||||
let source = if inline {
|
||||
CredentialSource::Trusted(SecretValue::new(configured.to_string()))
|
||||
} else {
|
||||
CredentialSource::ApplicationCredentials(file.path().to_str().unwrap().into())
|
||||
};
|
||||
let provider = NativeProviderLoader.load(source).await.unwrap();
|
||||
|
||||
assert!(matches!(
|
||||
provider.project_id().await,
|
||||
Err(Error::CredentialAcquisition(ErrorDetail::Unresolved {
|
||||
subject: "GCP resource project"
|
||||
}))
|
||||
));
|
||||
let original = provider.token().await.unwrap();
|
||||
tokio::time::pause();
|
||||
tokio::time::advance(std::time::Duration::from_secs(3601)).await;
|
||||
tokio::time::resume();
|
||||
let refreshed = provider.token().await.unwrap();
|
||||
assert_ne!(original, refreshed);
|
||||
assert_eq!(requests.load(Ordering::SeqCst), expected_requests);
|
||||
server.abort();
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::canonical_endpoint(r#"{"token_uri":"https://oauth2.googleapis.com/token"}"#, true)]
|
||||
#[case::noncanonical_endpoint(r#"{"token_uri":"http://127.0.0.1/token"}"#, false)]
|
||||
#[case::missing_endpoint("{}", false)]
|
||||
fn request_credentials_require_canonical_token_endpoint(
|
||||
#[case] credentials: &str,
|
||||
#[case] accepted: bool,
|
||||
) {
|
||||
assert_eq!(validate_request_credentials(credentials).is_ok(), accepted);
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::missing_file("/nonexistent/vertex-key.json", "Google credentials")]
|
||||
#[case::malformed_json(" {\"type\": \"service_account\",", "Google credentials")]
|
||||
#[case::quoted_json("'{\"type\": \"service_account\"}'", "Google credentials")]
|
||||
#[tokio::test]
|
||||
async fn trusted_credentials_are_read_as_a_file_unless_they_are_json(
|
||||
#[case] configured: &str,
|
||||
#[case] expected: &str,
|
||||
) {
|
||||
let Err(error) = NativeProviderLoader
|
||||
.load(CredentialSource::Trusted(SecretValue::new(configured)))
|
||||
.await
|
||||
else {
|
||||
panic!("{configured:?} should not load");
|
||||
};
|
||||
let report = litellm_error::Report::new(&error).to_string();
|
||||
assert!(report.contains(expected), "{report}");
|
||||
}
|
||||
|
||||
#[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 NativeProviderLoader
|
||||
.load(CredentialSource::Trusted(SecretValue::new(configured)))
|
||||
.await
|
||||
{
|
||||
Ok(_) => panic!("{configured:?} should not load"),
|
||||
Err(error) => litellm_error::Report::new(&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]
|
||||
#[tokio::test]
|
||||
async fn unknown_application_credential_types_are_rejected_without_service_account_fallback() {
|
||||
let file = tempfile::NamedTempFile::new().unwrap();
|
||||
std::fs::write(
|
||||
file.path(),
|
||||
r#"{"type":"unsupported","private_key":"credential-value"}"#,
|
||||
)
|
||||
.unwrap();
|
||||
let result = GcpAuth::default()
|
||||
.access_token(CredentialSource::ApplicationCredentials(
|
||||
file.path().to_str().unwrap().into(),
|
||||
))
|
||||
.await;
|
||||
let error = result.unwrap_err();
|
||||
assert!(matches!(error, Error::InvalidConfiguration(_)));
|
||||
assert!(
|
||||
!litellm_error::Report::new(&error)
|
||||
.to_string()
|
||||
.contains("credential-value")
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
fn inline_and_trusted_credentials_never_share_a_cache_entry() {
|
||||
assert_ne!(
|
||||
CredentialSource::Inline(SecretValue::new("same-value")).cache_key(),
|
||||
CredentialSource::Trusted(SecretValue::new("same-value")).cache_key()
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
@ -1,16 +1,16 @@
|
|||
mod auth;
|
||||
mod config;
|
||||
pub mod constants;
|
||||
mod identity;
|
||||
#[cfg(feature = "google-sdk")]
|
||||
mod sdk;
|
||||
|
||||
pub use auth::{
|
||||
VertexAuth, VertexAuthFuture, VertexEnvironment, VertexProviderLoader, VertexTokenSource,
|
||||
};
|
||||
pub use auth::{VertexAuth, VertexEnvironment};
|
||||
pub use config::{
|
||||
CredentialSource, VertexConfig, get_vertex_ai_location, get_vertex_ai_project,
|
||||
get_vertex_ai_project_from_credentials,
|
||||
};
|
||||
pub use constants::secret_names;
|
||||
pub use identity::{GcpAuth, GcpAuthFuture, GcpProviderLoader, GcpTokenSource};
|
||||
#[cfg(feature = "google-sdk")]
|
||||
pub use sdk::GoogleCredentials;
|
||||
|
|
|
|||
|
|
@ -1,120 +0,0 @@
|
|||
use std::{
|
||||
collections::{BTreeMap, BTreeSet},
|
||||
sync::{
|
||||
Arc,
|
||||
atomic::{AtomicUsize, Ordering},
|
||||
},
|
||||
};
|
||||
|
||||
use litellm_auth_gcp::{
|
||||
CredentialSource, VertexAuth, VertexAuthFuture, VertexConfig, VertexProviderLoader,
|
||||
VertexTokenSource, get_vertex_ai_location, secret_names,
|
||||
};
|
||||
use rstest::rstest;
|
||||
use serde_json::{Value, json};
|
||||
|
||||
struct FakeProvider {
|
||||
calls: Arc<AtomicUsize>,
|
||||
}
|
||||
|
||||
impl VertexTokenSource for FakeProvider {
|
||||
fn project_id(&self) -> VertexAuthFuture<'_, String> {
|
||||
self.calls.fetch_add(1, Ordering::SeqCst);
|
||||
Box::pin(async { Ok("adc-project".into()) })
|
||||
}
|
||||
|
||||
fn token(&self) -> VertexAuthFuture<'_, String> {
|
||||
self.calls.fetch_add(1, Ordering::SeqCst);
|
||||
Box::pin(async { Ok("adc-token".into()) })
|
||||
}
|
||||
}
|
||||
|
||||
struct FakeLoader {
|
||||
loads: Arc<AtomicUsize>,
|
||||
provider: Arc<dyn VertexTokenSource>,
|
||||
}
|
||||
|
||||
impl VertexProviderLoader for FakeLoader {
|
||||
fn load(&self, _source: CredentialSource) -> VertexAuthFuture<'_, Arc<dyn VertexTokenSource>> {
|
||||
let loads = self.loads.clone();
|
||||
let provider = self.provider.clone();
|
||||
Box::pin(async move {
|
||||
loads.fetch_add(1, Ordering::SeqCst);
|
||||
Ok(provider)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn config(value: Value) -> VertexConfig {
|
||||
VertexConfig::from_sourced_optional_params(value.as_object().unwrap(), &BTreeMap::new())
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
fn auth(calls: Arc<AtomicUsize>, loads: Arc<AtomicUsize>) -> VertexAuth {
|
||||
let provider: Arc<dyn VertexTokenSource> = Arc::new(FakeProvider { calls });
|
||||
VertexAuth::new(Arc::new(FakeLoader { loads, provider }))
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn secret_names_cover_environment_reads() {
|
||||
let seen = Arc::new(std::sync::Mutex::new(BTreeSet::<String>::new()));
|
||||
let recorded = seen.clone();
|
||||
let env = |name: &str| {
|
||||
recorded.lock().unwrap().insert(name.to_string());
|
||||
None
|
||||
};
|
||||
let auth = auth(Arc::new(AtomicUsize::new(0)), Arc::new(AtomicUsize::new(0)));
|
||||
auth.validate_environment(Vec::new(), None, &VertexConfig::default(), &env)
|
||||
.await
|
||||
.unwrap();
|
||||
get_vertex_ai_location(&VertexConfig::default(), &env);
|
||||
assert!(
|
||||
seen.lock()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.all(|name| secret_names().contains(&name.as_str()))
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn explicit_token_and_header_do_not_acquire_adc() {
|
||||
let loads = Arc::new(AtomicUsize::new(0));
|
||||
let auth = auth(Arc::new(AtomicUsize::new(0)), loads.clone());
|
||||
let configured = config(json!({"vertex_project":"project-1"}));
|
||||
let explicit = auth
|
||||
.validate_environment(Vec::new(), Some("access-token"), &configured, &|_| None)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(explicit.headers[0].1, "Bearer access-token");
|
||||
let existing = auth
|
||||
.validate_environment(
|
||||
vec![("authorization".into(), "Bearer existing".into())],
|
||||
None,
|
||||
&configured,
|
||||
&|_| None,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(existing.headers[0].1, "Bearer existing");
|
||||
assert_eq!(loads.load(Ordering::SeqCst), 0);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn provider_is_reused_across_authentication_calls() {
|
||||
let calls = Arc::new(AtomicUsize::new(0));
|
||||
let loads = Arc::new(AtomicUsize::new(0));
|
||||
let auth = auth(calls.clone(), loads.clone());
|
||||
for _ in 0..2 {
|
||||
let environment = auth
|
||||
.validate_environment(Vec::new(), None, &VertexConfig::default(), &|_| None)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(environment.project_id, "adc-project");
|
||||
assert_eq!(environment.headers[0].1, "Bearer adc-token");
|
||||
}
|
||||
assert_eq!(loads.load(Ordering::SeqCst), 1);
|
||||
assert_eq!(calls.load(Ordering::SeqCst), 4);
|
||||
}
|
||||
|
|
@ -92,6 +92,11 @@ impl Error {
|
|||
|
||||
#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
|
||||
pub enum ErrorDetail {
|
||||
#[error("{operation} returned HTTP {status}")]
|
||||
UnexpectedStatus {
|
||||
operation: &'static str,
|
||||
status: u16,
|
||||
},
|
||||
#[error("{field} must be {expected}")]
|
||||
InvalidType {
|
||||
field: String,
|
||||
|
|
|
|||
|
|
@ -53,6 +53,9 @@ pub struct AwsParams {
|
|||
}
|
||||
|
||||
impl AwsParams {
|
||||
pub const ROLE_ARN_ENV: &str = "AWS_ROLE_ARN";
|
||||
pub const WEB_IDENTITY_TOKEN_FILE_ENV: &str = "AWS_WEB_IDENTITY_TOKEN_FILE";
|
||||
|
||||
pub const ACCESS_KEY_ID: ParamSpec = spec(
|
||||
"access key id",
|
||||
ParamKind::Credential,
|
||||
|
|
@ -149,6 +152,10 @@ impl AwsParams {
|
|||
AwsParams::SPECS
|
||||
.iter()
|
||||
.flat_map(|spec| spec.env.iter().copied())
|
||||
.chain([
|
||||
AwsParams::ROLE_ARN_ENV,
|
||||
AwsParams::WEB_IDENTITY_TOKEN_FILE_ENV,
|
||||
])
|
||||
.fold(Vec::new(), |mut names, name| {
|
||||
if !names.contains(&name) {
|
||||
names.push(name);
|
||||
|
|
|
|||
|
|
@ -3,21 +3,13 @@ use rstest::rstest;
|
|||
use serde_json::{Map, Value};
|
||||
|
||||
#[rstest]
|
||||
fn fields_name_every_param_once_and_in_declaration_order() {
|
||||
fn every_named_param_roundtrips_without_loss() {
|
||||
let filled: Map<String, Value> = AwsParams::fields()
|
||||
.map(|name| (name.to_string(), Value::from(format!("value-of-{name}"))))
|
||||
.collect();
|
||||
let typed = AwsParams::from_optional_params(&filled);
|
||||
let serialized = serde_json::to_value(&typed).unwrap();
|
||||
assert_eq!(
|
||||
serialized
|
||||
.as_object()
|
||||
.unwrap()
|
||||
.keys()
|
||||
.map(String::as_str)
|
||||
.collect::<Vec<_>>(),
|
||||
AwsParams::fields().collect::<Vec<_>>()
|
||||
);
|
||||
assert_eq!(filled.len(), AwsParams::fields().count());
|
||||
assert_eq!(serialized, Value::Object(filled.clone()));
|
||||
assert_eq!(
|
||||
serde_json::from_value::<AwsParams>(Value::Object(filled)).unwrap(),
|
||||
|
|
|
|||
|
|
@ -8,7 +8,7 @@ use litellm_cache::{
|
|||
use litellm_http::Client;
|
||||
use percent_encoding::{AsciiSet, NON_ALPHANUMERIC, percent_encode};
|
||||
|
||||
use crate::{GcpTokenSource, TokenSource};
|
||||
use crate::{GcsTokenSource, TokenSource};
|
||||
|
||||
pub const DEFAULT_ENDPOINT: &str = "https://storage.googleapis.com";
|
||||
|
||||
|
|
@ -62,7 +62,7 @@ pub struct GcsCache<S: CacheCodec> {
|
|||
|
||||
impl<S: CacheCodec> GcsCache<S> {
|
||||
pub fn new(config: GcsConfig, client: Client, codec: S) -> Self {
|
||||
let token = Arc::new(GcpTokenSource::new(config.path_service_account.clone()));
|
||||
let token = Arc::new(GcsTokenSource::new(config.path_service_account.clone()));
|
||||
Self::with_token_source(config, client, codec, token)
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -2,4 +2,4 @@ mod cache;
|
|||
mod token;
|
||||
|
||||
pub use cache::{DEFAULT_ENDPOINT, GcsCache, GcsConfig, key_prefix};
|
||||
pub use token::{GcpTokenSource, StaticTokenSource, TokenSource};
|
||||
pub use token::{GcsTokenSource, StaticTokenSource, TokenSource};
|
||||
|
|
|
|||
|
|
@ -1,7 +1,6 @@
|
|||
use std::{future::Future, pin::Pin};
|
||||
|
||||
use litellm_auth_gcp::{VertexAuth, VertexConfig};
|
||||
use litellm_auth_types::{InputSource, Sourced};
|
||||
use litellm_auth_gcp::{CredentialSource, GcpAuth};
|
||||
use litellm_secret_value::SecretValue;
|
||||
|
||||
use litellm_cache::Error;
|
||||
|
|
@ -10,28 +9,30 @@ pub trait TokenSource: Send + Sync + 'static {
|
|||
fn bearer_token(&self) -> Pin<Box<dyn Future<Output = Result<String, Error>> + Send + '_>>;
|
||||
}
|
||||
|
||||
pub struct GcpTokenSource {
|
||||
auth: VertexAuth,
|
||||
config: VertexConfig,
|
||||
pub struct GcsTokenSource {
|
||||
auth: GcpAuth,
|
||||
credentials: Option<String>,
|
||||
}
|
||||
|
||||
impl GcpTokenSource {
|
||||
impl GcsTokenSource {
|
||||
pub fn new(path_service_account: Option<String>) -> Self {
|
||||
let credentials = path_service_account
|
||||
.map(|path| Sourced::new(SecretValue::new(path), InputSource::Deployment));
|
||||
Self {
|
||||
auth: VertexAuth::default(),
|
||||
config: VertexConfig::new(credentials, None, None),
|
||||
auth: GcpAuth::default(),
|
||||
credentials: path_service_account.filter(|path| !path.trim().is_empty()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl TokenSource for GcpTokenSource {
|
||||
impl TokenSource for GcsTokenSource {
|
||||
fn bearer_token(&self) -> Pin<Box<dyn Future<Output = Result<String, Error>> + Send + '_>> {
|
||||
Box::pin(async move {
|
||||
self.auth
|
||||
.access_token(&self.config, &|name| std::env::var(name).ok())
|
||||
.access_token(match &self.credentials {
|
||||
Some(path) => CredentialSource::Trusted(SecretValue::new(path.clone())),
|
||||
None => CredentialSource::application_default(&|name| std::env::var(name).ok()),
|
||||
})
|
||||
.await
|
||||
.map(|token| token.expose().to_owned())
|
||||
.map_err(Error::unavailable)
|
||||
})
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,19 +0,0 @@
|
|||
[package]
|
||||
name = "litellm-cache-redis-semantic"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
litellm-cache.workspace = true
|
||||
litellm-cache-redis.workspace = true
|
||||
redis = { workspace = true, features = ["tls-rustls"] }
|
||||
sha2.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
litellm-cache-testing.workspace = true
|
||||
redis-test.workspace = true
|
||||
rstest.workspace = true
|
||||
serde_json.workspace = true
|
||||
tokio.workspace = true
|
||||
|
|
@ -1,404 +0,0 @@
|
|||
use std::{
|
||||
sync::Arc,
|
||||
time::{Duration, SystemTime, UNIX_EPOCH},
|
||||
};
|
||||
|
||||
use litellm_cache::{
|
||||
BaseCache, CacheCodec, Error, SemanticCacheContext,
|
||||
semantic::{Embedder, SemanticCache, SemanticLookup, prompt_from_context},
|
||||
};
|
||||
use litellm_cache_redis::{
|
||||
RedisTopology,
|
||||
connection::{ConnectionRef, Connections},
|
||||
};
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
use crate::{
|
||||
RedisSemanticConfig,
|
||||
index::{CACHE_KEY_FIELD, Index, VECTOR_FIELD},
|
||||
reply::{bytes_field, first_document, number_field, string_field},
|
||||
};
|
||||
|
||||
struct Inner {
|
||||
index: Index,
|
||||
distance_threshold: f64,
|
||||
clock: fn() -> f64,
|
||||
}
|
||||
|
||||
impl Inner {
|
||||
fn new(config: RedisSemanticConfig, clock: fn() -> f64) -> Self {
|
||||
Self {
|
||||
index: Index::new(config.index_name),
|
||||
distance_threshold: 1.0 - f64::from(config.similarity_threshold),
|
||||
clock,
|
||||
}
|
||||
}
|
||||
|
||||
fn store(
|
||||
&self,
|
||||
connection: &mut ConnectionRef<'_>,
|
||||
tag: &str,
|
||||
response: Vec<u8>,
|
||||
prompt: &str,
|
||||
vector: &[f32],
|
||||
ttl: Option<Duration>,
|
||||
) -> Result<(), Error> {
|
||||
let index = self.index.ensure(connection, vector.len())?;
|
||||
let entry_id = entry_id(prompt, tag);
|
||||
let hash_key = format!("{index}:{entry_id}");
|
||||
redis::cmd("HSET")
|
||||
.arg(&hash_key)
|
||||
.arg("entry_id")
|
||||
.arg(&entry_id)
|
||||
.arg("prompt")
|
||||
.arg(prompt)
|
||||
.arg("response")
|
||||
.arg(response)
|
||||
.arg(VECTOR_FIELD)
|
||||
.arg(vector_buffer(vector))
|
||||
.arg("inserted_at")
|
||||
.arg(format!("{}", (self.clock)()))
|
||||
.arg("updated_at")
|
||||
.arg(format!("{}", (self.clock)()))
|
||||
.arg(CACHE_KEY_FIELD)
|
||||
.arg(tag)
|
||||
.query::<()>(connection)
|
||||
.map_err(Error::unavailable)?;
|
||||
if let Some(ttl) = ttl {
|
||||
redis::cmd("EXPIRE")
|
||||
.arg(&hash_key)
|
||||
.arg(ttl_seconds(ttl))
|
||||
.query::<()>(connection)
|
||||
.map_err(Error::unavailable)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn lookup(
|
||||
&self,
|
||||
connection: &mut ConnectionRef<'_>,
|
||||
tag: &str,
|
||||
vector: &[f32],
|
||||
) -> Result<SemanticLookup<Vec<u8>>, Error> {
|
||||
let index = self.index.ensure(connection, vector.len())?;
|
||||
let query = format!(
|
||||
"(@{CACHE_KEY_FIELD}:{{{}}})=>[KNN 1 @{VECTOR_FIELD} $vector AS vector_distance]",
|
||||
escape_tag(tag)
|
||||
);
|
||||
let result = redis::cmd("FT.SEARCH")
|
||||
.arg(&index)
|
||||
.arg(query)
|
||||
.arg("RETURN")
|
||||
.arg(8)
|
||||
.arg("entry_id")
|
||||
.arg("prompt")
|
||||
.arg("response")
|
||||
.arg("inserted_at")
|
||||
.arg("updated_at")
|
||||
.arg("metadata")
|
||||
.arg(CACHE_KEY_FIELD)
|
||||
.arg("vector_distance")
|
||||
.arg("SORTBY")
|
||||
.arg("vector_distance")
|
||||
.arg("ASC")
|
||||
.arg("DIALECT")
|
||||
.arg(2)
|
||||
.arg("LIMIT")
|
||||
.arg(0)
|
||||
.arg(1)
|
||||
.arg("PARAMS")
|
||||
.arg(2)
|
||||
.arg("vector")
|
||||
.arg(vector_buffer(vector))
|
||||
.query::<redis::Value>(connection)
|
||||
.map_err(Error::unavailable)?;
|
||||
let Some(fields) = first_document(&result) else {
|
||||
return Ok(SemanticLookup::miss(Some(0.0)));
|
||||
};
|
||||
if string_field(fields, CACHE_KEY_FIELD).as_deref() != Some(tag) {
|
||||
return Ok(SemanticLookup::miss(Some(0.0)));
|
||||
}
|
||||
// redisvl's range query only returns entries within the distance threshold, so a
|
||||
// farther hit reads as no result.
|
||||
let Some(distance) = number_field(fields, "vector_distance")
|
||||
.filter(|distance| *distance <= self.distance_threshold)
|
||||
else {
|
||||
return Ok(SemanticLookup::miss(Some(0.0)));
|
||||
};
|
||||
let Some(response) = bytes_field(fields, "response") else {
|
||||
return Ok(SemanticLookup::miss(Some(0.0)));
|
||||
};
|
||||
Ok(SemanticLookup {
|
||||
value: Some(response),
|
||||
similarity: Some(1.0 - distance),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// `RedisSemanticCache`: a redisvl-compatible semantic index on Redis Stack. Values go through
|
||||
/// the injected codec, so the response layer decides what a cached entry is.
|
||||
pub struct RedisSemanticCache<E, S, C = redis::Connection> {
|
||||
connections: Arc<Connections<C>>,
|
||||
embedder: E,
|
||||
codec: S,
|
||||
inner: Arc<Inner>,
|
||||
}
|
||||
|
||||
impl<E: Embedder, S: CacheCodec> RedisSemanticCache<E, S> {
|
||||
pub fn new(
|
||||
url: &str,
|
||||
embedder: E,
|
||||
codec: S,
|
||||
config: RedisSemanticConfig,
|
||||
) -> Result<Self, Error> {
|
||||
Ok(Self {
|
||||
connections: Arc::new(Connections::open(url, &RedisTopology::Standalone)?),
|
||||
embedder,
|
||||
codec,
|
||||
inner: Arc::new(Inner::new(config, timestamp)),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl<E, S, C> RedisSemanticCache<E, S, C>
|
||||
where
|
||||
E: Embedder,
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
pub fn with_connection(
|
||||
connection: C,
|
||||
embedder: E,
|
||||
codec: S,
|
||||
config: RedisSemanticConfig,
|
||||
) -> Self {
|
||||
Self {
|
||||
connections: Arc::new(Connections::fixed(connection)),
|
||||
embedder,
|
||||
codec,
|
||||
inner: Arc::new(Inner::new(config, timestamp)),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_clock(self, clock: fn() -> f64) -> Self {
|
||||
let config = RedisSemanticConfig {
|
||||
index_name: self.index_name().to_owned(),
|
||||
similarity_threshold: self.similarity_threshold(),
|
||||
};
|
||||
Self {
|
||||
inner: Arc::new(Inner::new(config, clock)),
|
||||
..self
|
||||
}
|
||||
}
|
||||
|
||||
pub fn embedder(&self) -> &E {
|
||||
&self.embedder
|
||||
}
|
||||
|
||||
pub fn index_name(&self) -> &str {
|
||||
self.inner.index.name()
|
||||
}
|
||||
|
||||
pub fn similarity_threshold(&self) -> f32 {
|
||||
(1.0 - self.inner.distance_threshold) as f32
|
||||
}
|
||||
|
||||
fn tag<'a>(key: &'a str, context: &'a SemanticCacheContext) -> &'a str {
|
||||
context.scope.as_deref().unwrap_or(key)
|
||||
}
|
||||
|
||||
fn decode(&self, lookup: SemanticLookup<Vec<u8>>) -> Result<SemanticLookup<S::Value>, Error> {
|
||||
Ok(SemanticLookup {
|
||||
value: lookup
|
||||
.value
|
||||
.map(|bytes| self.codec.decode(&bytes))
|
||||
.transpose()?,
|
||||
similarity: lookup.similarity,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl<E, S, C> BaseCache for RedisSemanticCache<E, S, C>
|
||||
where
|
||||
E: Embedder,
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
type Value = S::Value;
|
||||
type Context = SemanticCacheContext;
|
||||
|
||||
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
|
||||
context.ttl
|
||||
}
|
||||
|
||||
fn set_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
value: Self::Value,
|
||||
context: &Self::Context,
|
||||
) -> Result<(), Error> {
|
||||
let Some(prompt) = prompt_from_context(context) else {
|
||||
return Ok(());
|
||||
};
|
||||
let response = self.codec.encode(&value)?;
|
||||
let vector = self.embedder.embed(&prompt, context.metadata.as_ref())?;
|
||||
let tag = Self::tag(key, context);
|
||||
self.connections.execute(|connection| {
|
||||
self.inner
|
||||
.store(connection, tag, response, &prompt, &vector, context.ttl)
|
||||
})
|
||||
}
|
||||
|
||||
fn get_cache(&self, key: &str, context: &Self::Context) -> Result<Option<Self::Value>, Error> {
|
||||
self.get_cache_with_similarity(key, context)
|
||||
.map(|lookup| lookup.value)
|
||||
}
|
||||
|
||||
async fn async_set_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
value: Self::Value,
|
||||
context: Self::Context,
|
||||
) -> Result<(), Error> {
|
||||
let Some(prompt) = prompt_from_context(&context) else {
|
||||
return Ok(());
|
||||
};
|
||||
let response = self.codec.encode(&value)?;
|
||||
let vector = self
|
||||
.embedder
|
||||
.async_embed(&prompt, context.metadata.as_ref())
|
||||
.await?;
|
||||
let tag = Self::tag(key, &context).to_owned();
|
||||
let inner = Arc::clone(&self.inner);
|
||||
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
|
||||
inner.store(connection, &tag, response, &prompt, &vector, context.ttl)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
async fn async_get_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
context: &Self::Context,
|
||||
) -> Result<Option<Self::Value>, Error> {
|
||||
self.async_get_cache_with_similarity(key, context)
|
||||
.await
|
||||
.map(|lookup| lookup.value)
|
||||
}
|
||||
}
|
||||
|
||||
/// Python stamps a similarity of `0.0` when there is no prompt or no hit in the key's scope.
|
||||
impl<E, S, C> SemanticCache for RedisSemanticCache<E, S, C>
|
||||
where
|
||||
E: Embedder,
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
fn get_cache_with_similarity(
|
||||
&self,
|
||||
key: &str,
|
||||
context: &Self::Context,
|
||||
) -> Result<SemanticLookup<Self::Value>, Error> {
|
||||
let Some(prompt) = prompt_from_context(context) else {
|
||||
return Ok(SemanticLookup::miss(Some(0.0)));
|
||||
};
|
||||
let vector = self.embedder.embed(&prompt, context.metadata.as_ref())?;
|
||||
let tag = Self::tag(key, context);
|
||||
let lookup = self
|
||||
.connections
|
||||
.execute(|connection| self.inner.lookup(connection, tag, &vector))?;
|
||||
self.decode(lookup)
|
||||
}
|
||||
|
||||
async fn async_get_cache_with_similarity(
|
||||
&self,
|
||||
key: &str,
|
||||
context: &Self::Context,
|
||||
) -> Result<SemanticLookup<Self::Value>, Error> {
|
||||
let Some(prompt) = prompt_from_context(context) else {
|
||||
return Ok(SemanticLookup::miss(Some(0.0)));
|
||||
};
|
||||
let vector = self
|
||||
.embedder
|
||||
.async_embed(&prompt, context.metadata.as_ref())
|
||||
.await?;
|
||||
let tag = Self::tag(key, context).to_owned();
|
||||
let inner = Arc::clone(&self.inner);
|
||||
let lookup = Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
|
||||
inner.lookup(connection, &tag, &vector)
|
||||
})
|
||||
.await?;
|
||||
self.decode(lookup)
|
||||
}
|
||||
}
|
||||
|
||||
fn timestamp() -> f64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|duration| duration.as_secs_f64())
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
fn entry_id(prompt: &str, tag: &str) -> String {
|
||||
let mut digest = Sha256::new();
|
||||
digest.update(prompt.as_bytes());
|
||||
digest.update(CACHE_KEY_FIELD.as_bytes());
|
||||
digest.update(tag.as_bytes());
|
||||
format!("{:x}", digest.finalize())
|
||||
}
|
||||
|
||||
fn vector_buffer(vector: &[f32]) -> Vec<u8> {
|
||||
vector
|
||||
.iter()
|
||||
.flat_map(|component| component.to_le_bytes())
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn escape_tag(value: &str) -> String {
|
||||
let mut escaped = String::with_capacity(value.len());
|
||||
for ch in value.chars() {
|
||||
if matches!(
|
||||
ch,
|
||||
',' | '.'
|
||||
| '<'
|
||||
| '>'
|
||||
| '{'
|
||||
| '}'
|
||||
| '['
|
||||
| ']'
|
||||
| '\\'
|
||||
| '"'
|
||||
| '\''
|
||||
| ':'
|
||||
| ';'
|
||||
| '!'
|
||||
| '@'
|
||||
| '#'
|
||||
| '$'
|
||||
| '%'
|
||||
| '^'
|
||||
| '&'
|
||||
| '*'
|
||||
| '('
|
||||
| ')'
|
||||
| '-'
|
||||
| '+'
|
||||
| '='
|
||||
| '~'
|
||||
| '|'
|
||||
| '/'
|
||||
| ' '
|
||||
| '?'
|
||||
) {
|
||||
escaped.push('\\');
|
||||
}
|
||||
escaped.push(ch);
|
||||
}
|
||||
escaped
|
||||
}
|
||||
|
||||
fn ttl_seconds(ttl: Duration) -> u64 {
|
||||
ttl.as_secs()
|
||||
.saturating_add(u64::from(ttl.subsec_nanos() > 0))
|
||||
.max(1)
|
||||
}
|
||||
|
|
@ -1,8 +0,0 @@
|
|||
/// `RedisSemanticCache.DEFAULT_REDIS_INDEX_NAME`.
|
||||
pub const DEFAULT_INDEX_NAME: &str = "litellm_semantic_cache_index";
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct RedisSemanticConfig {
|
||||
pub index_name: String,
|
||||
pub similarity_threshold: f32,
|
||||
}
|
||||
|
|
@ -1,209 +0,0 @@
|
|||
use std::sync::OnceLock;
|
||||
|
||||
use litellm_cache::Error;
|
||||
use litellm_cache_redis::connection::ConnectionRef;
|
||||
|
||||
use crate::reply::{number_value, string_value};
|
||||
|
||||
pub(crate) const CACHE_KEY_FIELD: &str = "litellm_cache_key";
|
||||
pub(crate) const VECTOR_FIELD: &str = "prompt_vector";
|
||||
|
||||
/// The redisvl `SemanticCache` index, resolved once per cache: the configured name when its
|
||||
/// schema fits, else `<name>_isolated`, recreated when that one is stale too.
|
||||
pub(crate) struct Index {
|
||||
name: String,
|
||||
resolved: OnceLock<String>,
|
||||
}
|
||||
|
||||
impl Index {
|
||||
pub(crate) fn new(name: String) -> Self {
|
||||
Self {
|
||||
name,
|
||||
resolved: OnceLock::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn name(&self) -> &str {
|
||||
&self.name
|
||||
}
|
||||
|
||||
pub(crate) fn ensure(
|
||||
&self,
|
||||
connection: &mut ConnectionRef<'_>,
|
||||
dims: usize,
|
||||
) -> Result<String, Error> {
|
||||
if let Some(name) = self.resolved.get() {
|
||||
return Ok(name.clone());
|
||||
}
|
||||
let name = match index_compatible(connection, &self.name, dims)? {
|
||||
Some(true) => self.name.clone(),
|
||||
Some(false) => self.isolated(connection, dims)?,
|
||||
None => match create_index(connection, &self.name, dims) {
|
||||
Ok(()) => self.name.clone(),
|
||||
Err(error) => match index_compatible(connection, &self.name, dims)? {
|
||||
Some(true) => self.name.clone(),
|
||||
Some(false) => self.isolated(connection, dims)?,
|
||||
None => return Err(error),
|
||||
},
|
||||
},
|
||||
};
|
||||
let _ = self.resolved.set(name.clone());
|
||||
Ok(name)
|
||||
}
|
||||
|
||||
fn isolated(&self, connection: &mut ConnectionRef<'_>, dims: usize) -> Result<String, Error> {
|
||||
let name = format!("{}_isolated", self.name);
|
||||
match index_compatible(connection, &name, dims)? {
|
||||
Some(true) => Ok(name),
|
||||
Some(false) => {
|
||||
redis::cmd("FT.DROPINDEX")
|
||||
.arg(&name)
|
||||
.query::<()>(connection)
|
||||
.map_err(Error::unavailable)?;
|
||||
create_index(connection, &name, dims)?;
|
||||
Ok(name)
|
||||
}
|
||||
None => {
|
||||
create_index(connection, &name, dims)?;
|
||||
Ok(name)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn create_index(connection: &mut ConnectionRef<'_>, name: &str, dims: usize) -> Result<(), Error> {
|
||||
redis::cmd("FT.CREATE")
|
||||
.arg(name)
|
||||
.arg("ON")
|
||||
.arg("HASH")
|
||||
.arg("PREFIX")
|
||||
.arg(1)
|
||||
.arg(name)
|
||||
.arg("SCORE")
|
||||
.arg(1.0)
|
||||
.arg("SCHEMA")
|
||||
.arg("prompt")
|
||||
.arg("TEXT")
|
||||
.arg("WEIGHT")
|
||||
.arg(1)
|
||||
.arg("response")
|
||||
.arg("TEXT")
|
||||
.arg("WEIGHT")
|
||||
.arg(1)
|
||||
.arg("inserted_at")
|
||||
.arg("NUMERIC")
|
||||
.arg("updated_at")
|
||||
.arg("NUMERIC")
|
||||
.arg(VECTOR_FIELD)
|
||||
.arg("VECTOR")
|
||||
.arg("FLAT")
|
||||
.arg(6)
|
||||
.arg("TYPE")
|
||||
.arg("FLOAT32")
|
||||
.arg("DIM")
|
||||
.arg(dims)
|
||||
.arg("DISTANCE_METRIC")
|
||||
.arg("COSINE")
|
||||
.arg(CACHE_KEY_FIELD)
|
||||
.arg("TAG")
|
||||
.arg("SEPARATOR")
|
||||
.arg(",")
|
||||
.query::<()>(connection)
|
||||
.map_err(Error::unavailable)
|
||||
}
|
||||
|
||||
fn index_compatible(
|
||||
connection: &mut ConnectionRef<'_>,
|
||||
name: &str,
|
||||
dims: usize,
|
||||
) -> Result<Option<bool>, Error> {
|
||||
let info = match redis::cmd("FT.INFO")
|
||||
.arg(name)
|
||||
.query::<redis::Value>(connection)
|
||||
{
|
||||
Ok(info) => info,
|
||||
Err(error) if unknown_index(&error) => return Ok(None),
|
||||
Err(error) => return Err(Error::unavailable(error)),
|
||||
};
|
||||
Ok(Some(schema_compatible(&info, dims)))
|
||||
}
|
||||
|
||||
fn unknown_index(error: &redis::RedisError) -> bool {
|
||||
error
|
||||
.code()
|
||||
.is_some_and(|code| code.eq_ignore_ascii_case("unknown"))
|
||||
&& error
|
||||
.detail()
|
||||
.is_some_and(|detail| detail.to_ascii_lowercase().starts_with("index"))
|
||||
}
|
||||
|
||||
struct Attribute {
|
||||
name: Option<String>,
|
||||
field_type: Option<String>,
|
||||
dim: Option<f64>,
|
||||
data_type: Option<String>,
|
||||
distance_metric: Option<String>,
|
||||
}
|
||||
|
||||
fn attribute(value: &redis::Value) -> Option<Attribute> {
|
||||
let redis::Value::Array(pairs) = value else {
|
||||
return None;
|
||||
};
|
||||
let mut attribute = Attribute {
|
||||
name: None,
|
||||
field_type: None,
|
||||
dim: None,
|
||||
data_type: None,
|
||||
distance_metric: None,
|
||||
};
|
||||
for pair in pairs.as_chunks::<2>().0 {
|
||||
match string_value(&pair[0]).as_deref() {
|
||||
Some("identifier") => attribute.name = string_value(&pair[1]),
|
||||
Some("type") => attribute.field_type = string_value(&pair[1]),
|
||||
Some("dim") => attribute.dim = number_value(&pair[1]),
|
||||
Some("data_type") => attribute.data_type = string_value(&pair[1]),
|
||||
Some("distance_metric") => attribute.distance_metric = string_value(&pair[1]),
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
Some(attribute)
|
||||
}
|
||||
|
||||
fn schema_compatible(info: &redis::Value, dims: usize) -> bool {
|
||||
let redis::Value::Array(entries) = info else {
|
||||
return false;
|
||||
};
|
||||
let attributes = entries
|
||||
.as_chunks::<2>()
|
||||
.0
|
||||
.iter()
|
||||
.find(|pair| string_value(&pair[0]).as_deref() == Some("attributes"))
|
||||
.map(|pair| &pair[1]);
|
||||
let Some(redis::Value::Array(attributes)) = attributes else {
|
||||
return false;
|
||||
};
|
||||
let fields = attributes.iter().filter_map(attribute).collect::<Vec<_>>();
|
||||
let has_field = |name: &str, field_type: &str| {
|
||||
fields.iter().any(|field| {
|
||||
field.name.as_deref() == Some(name) && field.field_type.as_deref() == Some(field_type)
|
||||
})
|
||||
};
|
||||
has_field("prompt", "TEXT")
|
||||
&& has_field("response", "TEXT")
|
||||
&& has_field("inserted_at", "NUMERIC")
|
||||
&& has_field("updated_at", "NUMERIC")
|
||||
&& has_field(CACHE_KEY_FIELD, "TAG")
|
||||
&& fields.iter().any(|field| {
|
||||
field.name.as_deref() == Some(VECTOR_FIELD)
|
||||
&& field.field_type.as_deref() == Some("VECTOR")
|
||||
&& field.dim == Some(dims as f64)
|
||||
&& field
|
||||
.data_type
|
||||
.as_deref()
|
||||
.is_some_and(|data| data.eq_ignore_ascii_case("float32"))
|
||||
&& field
|
||||
.distance_metric
|
||||
.as_deref()
|
||||
.is_some_and(|metric| metric.eq_ignore_ascii_case("cosine"))
|
||||
})
|
||||
}
|
||||
|
|
@ -1,7 +0,0 @@
|
|||
mod cache;
|
||||
mod config;
|
||||
mod index;
|
||||
mod reply;
|
||||
|
||||
pub use cache::RedisSemanticCache;
|
||||
pub use config::{DEFAULT_INDEX_NAME, RedisSemanticConfig};
|
||||
|
|
@ -1,57 +0,0 @@
|
|||
pub(crate) fn string_value(value: &redis::Value) -> Option<String> {
|
||||
match value {
|
||||
redis::Value::BulkString(bytes) => String::from_utf8(bytes.clone()).ok(),
|
||||
redis::Value::SimpleString(text) => Some(text.clone()),
|
||||
redis::Value::VerbatimString { text, .. } => Some(text.clone()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn number_value(value: &redis::Value) -> Option<f64> {
|
||||
match value {
|
||||
redis::Value::Int(number) => Some(*number as f64),
|
||||
redis::Value::Double(number) => Some(*number),
|
||||
_ => string_value(value).and_then(|text| text.parse().ok()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn first_document(result: &redis::Value) -> Option<&[redis::Value]> {
|
||||
let redis::Value::Array(items) = result else {
|
||||
return None;
|
||||
};
|
||||
let [count, _document_id, fields, ..] = items.as_slice() else {
|
||||
return None;
|
||||
};
|
||||
if !matches!(count, redis::Value::Int(count) if *count > 0) {
|
||||
return None;
|
||||
}
|
||||
match fields {
|
||||
redis::Value::Array(fields) => Some(fields.as_slice()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn field_value<'a>(fields: &'a [redis::Value], name: &str) -> Option<&'a redis::Value> {
|
||||
fields
|
||||
.as_chunks::<2>()
|
||||
.0
|
||||
.iter()
|
||||
.find(|pair| string_value(&pair[0]).as_deref() == Some(name))
|
||||
.map(|pair| &pair[1])
|
||||
}
|
||||
|
||||
pub(crate) fn string_field(fields: &[redis::Value], name: &str) -> Option<String> {
|
||||
field_value(fields, name).and_then(string_value)
|
||||
}
|
||||
|
||||
pub(crate) fn number_field(fields: &[redis::Value], name: &str) -> Option<f64> {
|
||||
field_value(fields, name).and_then(number_value)
|
||||
}
|
||||
|
||||
pub(crate) fn bytes_field(fields: &[redis::Value], name: &str) -> Option<Vec<u8>> {
|
||||
match field_value(fields, name)? {
|
||||
redis::Value::BulkString(bytes) => Some(bytes.clone()),
|
||||
redis::Value::SimpleString(text) => Some(text.clone().into_bytes()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
|
@ -1,778 +0,0 @@
|
|||
mod support;
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use litellm_cache::{
|
||||
BaseCache, Error, JsonCodec, SemanticCacheContext,
|
||||
semantic::{SemanticCache, SemanticLookup},
|
||||
};
|
||||
use litellm_cache_redis_semantic::{DEFAULT_INDEX_NAME, RedisSemanticCache, RedisSemanticConfig};
|
||||
use redis_test::{MockCmd, MockRedisConnection};
|
||||
use rstest::{fixture, rstest};
|
||||
use serde_json::{Value, json};
|
||||
use sha2::{Digest, Sha256};
|
||||
use support::FakeEmbedder;
|
||||
|
||||
const INDEX: &str = DEFAULT_INDEX_NAME;
|
||||
const PROMPT: &str = "hello prompt";
|
||||
const CLOCK: fn() -> f64 = || 1700000000.5;
|
||||
const VECTOR: [f32; 3] = [0.1, 0.2, 0.3];
|
||||
|
||||
type MockCache = RedisSemanticCache<FakeEmbedder, JsonCodec<Value>, MockRedisConnection>;
|
||||
|
||||
#[fixture]
|
||||
fn config() -> RedisSemanticConfig {
|
||||
RedisSemanticConfig {
|
||||
index_name: INDEX.into(),
|
||||
similarity_threshold: 0.9,
|
||||
}
|
||||
}
|
||||
|
||||
#[fixture]
|
||||
fn entry() -> Value {
|
||||
json!({"timestamp": 1.0, "response": {"answer": "yes"}})
|
||||
}
|
||||
|
||||
#[fixture]
|
||||
fn context() -> SemanticCacheContext {
|
||||
messages_context(vec![json!({"role": "user", "content": PROMPT})])
|
||||
}
|
||||
|
||||
fn messages_context(messages: Vec<Value>) -> SemanticCacheContext {
|
||||
SemanticCacheContext {
|
||||
messages: Some(Value::Array(messages)),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
fn cache(commands: Vec<MockCmd>, embedder: FakeEmbedder) -> MockCache {
|
||||
RedisSemanticCache::with_connection(
|
||||
MockRedisConnection::new(commands).assert_all_commands_consumed(),
|
||||
embedder,
|
||||
JsonCodec::new(),
|
||||
config(),
|
||||
)
|
||||
.with_clock(CLOCK)
|
||||
}
|
||||
|
||||
fn encoded(value: &Value) -> Vec<u8> {
|
||||
serde_json::to_vec(value).unwrap()
|
||||
}
|
||||
|
||||
fn vector_bytes(vector: &[f32]) -> Vec<u8> {
|
||||
vector
|
||||
.iter()
|
||||
.flat_map(|component| component.to_le_bytes())
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn entry_id(prompt: &str, tag: &str) -> String {
|
||||
let mut digest = Sha256::new();
|
||||
digest.update(prompt.as_bytes());
|
||||
digest.update(b"litellm_cache_key");
|
||||
digest.update(tag.as_bytes());
|
||||
format!("{:x}", digest.finalize())
|
||||
}
|
||||
|
||||
fn s(value: &str) -> redis::Value {
|
||||
redis::Value::BulkString(value.as_bytes().to_vec())
|
||||
}
|
||||
|
||||
fn unknown_index_error() -> redis::RedisError {
|
||||
support::server_error("Unknown index name")
|
||||
}
|
||||
|
||||
fn info_missing(index: &str) -> MockCmd {
|
||||
MockCmd::new(
|
||||
redis::cmd("FT.INFO").arg(index),
|
||||
Err::<redis::Value, _>(unknown_index_error()),
|
||||
)
|
||||
}
|
||||
|
||||
fn info(index: &str, value: redis::Value) -> MockCmd {
|
||||
MockCmd::new(redis::cmd("FT.INFO").arg(index), Ok(value))
|
||||
}
|
||||
|
||||
fn attribute(name: &str, field_type: &str, extra: Vec<redis::Value>) -> redis::Value {
|
||||
let mut parts = vec![
|
||||
s("identifier"),
|
||||
s(name),
|
||||
s("attribute"),
|
||||
s(name),
|
||||
s("type"),
|
||||
s(field_type),
|
||||
];
|
||||
parts.extend(extra);
|
||||
redis::Value::Array(parts)
|
||||
}
|
||||
|
||||
fn index_info(attributes: Vec<redis::Value>) -> redis::Value {
|
||||
redis::Value::Array(vec![
|
||||
s("index_name"),
|
||||
s(INDEX),
|
||||
s("attributes"),
|
||||
redis::Value::Array(attributes),
|
||||
])
|
||||
}
|
||||
|
||||
fn vector_attribute_with(dims: i64, data_type: &str, distance_metric: &str) -> redis::Value {
|
||||
attribute(
|
||||
"prompt_vector",
|
||||
"VECTOR",
|
||||
vec![
|
||||
s("algorithm"),
|
||||
s("FLAT"),
|
||||
s("data_type"),
|
||||
s(data_type),
|
||||
s("dim"),
|
||||
redis::Value::Int(dims),
|
||||
s("distance_metric"),
|
||||
s(distance_metric),
|
||||
],
|
||||
)
|
||||
}
|
||||
|
||||
fn info_with_vector(vector: redis::Value) -> redis::Value {
|
||||
index_info(vec![
|
||||
attribute("prompt", "TEXT", vec![]),
|
||||
attribute("response", "TEXT", vec![]),
|
||||
attribute("inserted_at", "NUMERIC", vec![]),
|
||||
attribute("updated_at", "NUMERIC", vec![]),
|
||||
vector,
|
||||
attribute("litellm_cache_key", "TAG", vec![]),
|
||||
])
|
||||
}
|
||||
|
||||
fn compatible_info(dims: i64) -> redis::Value {
|
||||
info_with_vector(vector_attribute_with(dims, "FLOAT32", "COSINE"))
|
||||
}
|
||||
|
||||
fn unscoped_info(dims: i64) -> redis::Value {
|
||||
index_info(vec![
|
||||
attribute("prompt", "TEXT", vec![]),
|
||||
attribute("response", "TEXT", vec![]),
|
||||
attribute("inserted_at", "NUMERIC", vec![]),
|
||||
attribute("updated_at", "NUMERIC", vec![]),
|
||||
vector_attribute_with(dims, "FLOAT32", "COSINE"),
|
||||
])
|
||||
}
|
||||
|
||||
fn create_index(name: &str, dims: usize) -> MockCmd {
|
||||
MockCmd::new(create_index_command(name, dims), Ok("OK"))
|
||||
}
|
||||
|
||||
fn create_index_command(name: &str, dims: usize) -> redis::Cmd {
|
||||
let mut command = redis::cmd("FT.CREATE");
|
||||
command
|
||||
.arg(name)
|
||||
.arg("ON")
|
||||
.arg("HASH")
|
||||
.arg("PREFIX")
|
||||
.arg(1)
|
||||
.arg(name)
|
||||
.arg("SCORE")
|
||||
.arg(1.0)
|
||||
.arg("SCHEMA")
|
||||
.arg("prompt")
|
||||
.arg("TEXT")
|
||||
.arg("WEIGHT")
|
||||
.arg(1)
|
||||
.arg("response")
|
||||
.arg("TEXT")
|
||||
.arg("WEIGHT")
|
||||
.arg(1)
|
||||
.arg("inserted_at")
|
||||
.arg("NUMERIC")
|
||||
.arg("updated_at")
|
||||
.arg("NUMERIC")
|
||||
.arg("prompt_vector")
|
||||
.arg("VECTOR")
|
||||
.arg("FLAT")
|
||||
.arg(6)
|
||||
.arg("TYPE")
|
||||
.arg("FLOAT32")
|
||||
.arg("DIM")
|
||||
.arg(dims)
|
||||
.arg("DISTANCE_METRIC")
|
||||
.arg("COSINE")
|
||||
.arg("litellm_cache_key")
|
||||
.arg("TAG")
|
||||
.arg("SEPARATOR")
|
||||
.arg(",");
|
||||
command
|
||||
}
|
||||
|
||||
fn hset(index: &str, prompt: &str, tag: &str, vector: &[f32], value: &Value) -> MockCmd {
|
||||
MockCmd::new(
|
||||
redis::cmd("HSET")
|
||||
.arg(format!("{index}:{}", entry_id(prompt, tag)))
|
||||
.arg("entry_id")
|
||||
.arg(entry_id(prompt, tag))
|
||||
.arg("prompt")
|
||||
.arg(prompt)
|
||||
.arg("response")
|
||||
.arg(encoded(value))
|
||||
.arg("prompt_vector")
|
||||
.arg(vector_bytes(vector))
|
||||
.arg("inserted_at")
|
||||
.arg("1700000000.5")
|
||||
.arg("updated_at")
|
||||
.arg("1700000000.5")
|
||||
.arg("litellm_cache_key")
|
||||
.arg(tag),
|
||||
Ok(7),
|
||||
)
|
||||
}
|
||||
|
||||
fn search(
|
||||
index: &str,
|
||||
tag: &str,
|
||||
vector: &[f32],
|
||||
reply: redis::RedisResult<redis::Value>,
|
||||
) -> MockCmd {
|
||||
let mut command = redis::cmd("FT.SEARCH");
|
||||
command
|
||||
.arg(index)
|
||||
.arg(format!(
|
||||
"(@litellm_cache_key:{{{tag}}})=>[KNN 1 @prompt_vector $vector AS vector_distance]"
|
||||
))
|
||||
.arg("RETURN")
|
||||
.arg(8)
|
||||
.arg("entry_id")
|
||||
.arg("prompt")
|
||||
.arg("response")
|
||||
.arg("inserted_at")
|
||||
.arg("updated_at")
|
||||
.arg("metadata")
|
||||
.arg("litellm_cache_key")
|
||||
.arg("vector_distance")
|
||||
.arg("SORTBY")
|
||||
.arg("vector_distance")
|
||||
.arg("ASC")
|
||||
.arg("DIALECT")
|
||||
.arg(2)
|
||||
.arg("LIMIT")
|
||||
.arg(0)
|
||||
.arg(1)
|
||||
.arg("PARAMS")
|
||||
.arg(2)
|
||||
.arg("vector")
|
||||
.arg(vector_bytes(vector));
|
||||
MockCmd::new(command, reply)
|
||||
}
|
||||
|
||||
fn hit(tag: &str, distance: &str, response: Vec<u8>) -> redis::Value {
|
||||
redis::Value::Array(vec![
|
||||
redis::Value::Int(1),
|
||||
s("litellm_semantic_cache_index:stored-id"),
|
||||
redis::Value::Array(vec![
|
||||
s("entry_id"),
|
||||
s("stored-id"),
|
||||
s("prompt"),
|
||||
s(PROMPT),
|
||||
s("response"),
|
||||
redis::Value::BulkString(response),
|
||||
s("inserted_at"),
|
||||
s("1700000000.5"),
|
||||
s("updated_at"),
|
||||
s("1700000000.5"),
|
||||
s("litellm_cache_key"),
|
||||
s(tag),
|
||||
s("vector_distance"),
|
||||
s(distance),
|
||||
]),
|
||||
])
|
||||
}
|
||||
|
||||
fn empty_result() -> redis::Value {
|
||||
redis::Value::Array(vec![redis::Value::Int(0)])
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::creates_index_and_expires(false, Some(Duration::from_secs(5)), Some(5))]
|
||||
#[case::existing_index_without_ttl(true, None, None)]
|
||||
#[case::fractional_ttl_rounds_up(true, Some(Duration::from_millis(1500)), Some(2))]
|
||||
fn store_writes_the_redisvl_hash(
|
||||
#[case] index_exists: bool,
|
||||
#[case] ttl: Option<Duration>,
|
||||
#[case] expire: Option<u64>,
|
||||
entry: Value,
|
||||
context: SemanticCacheContext,
|
||||
) {
|
||||
let hash_key = format!("{INDEX}:{}", entry_id(PROMPT, "key1"));
|
||||
let mut commands = if index_exists {
|
||||
vec![info(INDEX, compatible_info(3))]
|
||||
} else {
|
||||
vec![info_missing(INDEX), create_index(INDEX, 3)]
|
||||
};
|
||||
commands.push(hset(INDEX, PROMPT, "key1", &VECTOR, &entry));
|
||||
commands.extend(
|
||||
expire.map(|seconds| MockCmd::new(redis::cmd("EXPIRE").arg(&hash_key).arg(seconds), Ok(1))),
|
||||
);
|
||||
let cache = cache(commands, FakeEmbedder::new(&[]));
|
||||
|
||||
cache
|
||||
.set_cache("key1", entry, &SemanticCacheContext { ttl, ..context })
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::below_distance_threshold("key1", "0.05", None, Ok(Some(entry())))]
|
||||
#[case::above_distance_threshold("key1", "0.5", None, Ok(None))]
|
||||
#[case::other_cache_key("other", "0.05", None, Ok(None))]
|
||||
#[case::malformed_response("key1", "0.05", Some(b"not json!".as_slice()), Err(Error::InvalidEntry(None)))]
|
||||
fn lookup_applies_threshold_scope_and_codec(
|
||||
#[case] stored_tag: &str,
|
||||
#[case] distance: &str,
|
||||
#[case] response: Option<&[u8]>,
|
||||
#[case] expected: Result<Option<Value>, Error>,
|
||||
entry: Value,
|
||||
context: SemanticCacheContext,
|
||||
) {
|
||||
let response = response.map_or_else(|| encoded(&entry), <[u8]>::to_vec);
|
||||
let cache = cache(
|
||||
vec![
|
||||
info(INDEX, compatible_info(3)),
|
||||
search(
|
||||
INDEX,
|
||||
"key1",
|
||||
&VECTOR,
|
||||
Ok(hit(stored_tag, distance, response)),
|
||||
),
|
||||
],
|
||||
FakeEmbedder::new(&[]),
|
||||
);
|
||||
|
||||
let kind = |result: Result<Option<Value>, Error>| {
|
||||
result.map_err(|error| std::mem::discriminant(&error))
|
||||
};
|
||||
assert_eq!(kind(cache.get_cache("key1", &context)), kind(expected));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::hit(context(), Some(hit("key1", "0.05", encoded(&entry()))), Some(entry()), Some(1.0 - 0.05))]
|
||||
#[case::beyond_distance_threshold(context(), Some(hit("key1", "0.5", encoded(&entry()))), None, Some(0.0))]
|
||||
#[case::no_results(context(), Some(empty_result()), None, Some(0.0))]
|
||||
#[case::other_cache_key(context(), Some(hit("other", "0.05", encoded(&entry()))), None, Some(0.0))]
|
||||
#[case::no_prompt(SemanticCacheContext::default(), None, None, Some(0.0))]
|
||||
#[tokio::test]
|
||||
async fn lookup_reports_python_semantic_similarity(
|
||||
#[case] context: SemanticCacheContext,
|
||||
#[case] reply: Option<redis::Value>,
|
||||
#[case] value: Option<Value>,
|
||||
#[case] similarity: Option<f64>,
|
||||
#[values(false, true)] use_async: bool,
|
||||
) {
|
||||
let commands = reply.map_or_else(Vec::new, |reply| {
|
||||
vec![
|
||||
info(INDEX, compatible_info(3)),
|
||||
search(INDEX, "key1", &VECTOR, Ok(reply)),
|
||||
]
|
||||
});
|
||||
let cache = cache(commands, FakeEmbedder::new(&[]));
|
||||
|
||||
let lookup = if use_async {
|
||||
cache
|
||||
.async_get_cache_with_similarity("key1", &context)
|
||||
.await
|
||||
} else {
|
||||
cache.get_cache_with_similarity("key1", &context)
|
||||
};
|
||||
|
||||
assert_eq!(lookup, Ok(SemanticLookup { value, similarity }));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn missing_prompt_is_a_noop_that_never_embeds(entry: Value) {
|
||||
let embedder = FakeEmbedder::new(&[]);
|
||||
let calls = embedder.calls.clone();
|
||||
let cache = cache(Vec::new(), embedder);
|
||||
let context = SemanticCacheContext::default();
|
||||
|
||||
cache.set_cache("key1", entry.clone(), &context).unwrap();
|
||||
assert_eq!(cache.get_cache("key1", &context).unwrap(), None);
|
||||
cache
|
||||
.async_set_cache("key1", entry, context.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(cache.async_get_cache("key1", &context).await.unwrap(), None);
|
||||
assert!(calls.lock().unwrap().is_empty());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn scope_overrides_key_as_filter_tag(entry: Value, context: SemanticCacheContext) {
|
||||
let cache = cache(
|
||||
vec![
|
||||
info(INDEX, compatible_info(3)),
|
||||
hset(INDEX, PROMPT, "scope-a", &VECTOR, &entry),
|
||||
search(
|
||||
INDEX,
|
||||
"scope\\-a",
|
||||
&VECTOR,
|
||||
Ok(hit("scope-a", "0.05", encoded(&entry))),
|
||||
),
|
||||
],
|
||||
FakeEmbedder::new(&[]),
|
||||
);
|
||||
let context = SemanticCacheContext {
|
||||
scope: Some("scope-a".into()),
|
||||
..context
|
||||
};
|
||||
|
||||
cache.set_cache("key1", entry.clone(), &context).unwrap();
|
||||
assert_eq!(cache.get_cache("key1", &context).unwrap(), Some(entry));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::unscoped_schema(unscoped_info(3))]
|
||||
#[case::wrong_distance_metric(info_with_vector(vector_attribute_with(3, "FLOAT32", "L2")))]
|
||||
#[case::wrong_data_type(info_with_vector(vector_attribute_with(3, "FLOAT64", "COSINE")))]
|
||||
fn incompatible_schema_falls_back_to_isolated_index(
|
||||
#[case] base_info: redis::Value,
|
||||
entry: Value,
|
||||
context: SemanticCacheContext,
|
||||
) {
|
||||
let isolated = format!("{INDEX}_isolated");
|
||||
let cache = cache(
|
||||
vec![
|
||||
info(INDEX, base_info),
|
||||
info_missing(&isolated),
|
||||
create_index(&isolated, 3),
|
||||
hset(&isolated, PROMPT, "key1", &VECTOR, &entry),
|
||||
],
|
||||
FakeEmbedder::new(&[]),
|
||||
);
|
||||
|
||||
cache.set_cache("key1", entry, &context).unwrap();
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn create_index_race_rechecks_schema_and_stores(entry: Value, context: SemanticCacheContext) {
|
||||
let cache = cache(
|
||||
vec![
|
||||
info_missing(INDEX),
|
||||
MockCmd::new(
|
||||
create_index_command(INDEX, 3),
|
||||
Err::<&str, _>(support::server_error("Index already exists")),
|
||||
),
|
||||
info(INDEX, compatible_info(3)),
|
||||
hset(INDEX, PROMPT, "key1", &VECTOR, &entry),
|
||||
],
|
||||
FakeEmbedder::new(&[]),
|
||||
);
|
||||
|
||||
cache.set_cache("key1", entry, &context).unwrap();
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::punctuation_and_spaces("a:b, c|d", "a\\:b\\,\\ c\\|d")]
|
||||
#[case::braces_and_dots("{x}.y", "\\{x\\}\\.y")]
|
||||
#[case::plain("key1", "key1")]
|
||||
fn tag_special_characters_are_escaped_in_search_filter(
|
||||
#[case] tag: &str,
|
||||
#[case] escaped: &str,
|
||||
context: SemanticCacheContext,
|
||||
) {
|
||||
let cache = cache(
|
||||
vec![
|
||||
info(INDEX, compatible_info(3)),
|
||||
search(INDEX, escaped, &VECTOR, Ok(empty_result())),
|
||||
],
|
||||
FakeEmbedder::new(&[]),
|
||||
);
|
||||
|
||||
assert_eq!(cache.get_cache(tag, &context).unwrap(), None);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::content_parts(
|
||||
SemanticCacheContext {
|
||||
messages: Some(json!([
|
||||
{"role": "user", "content": [{"type": "text", "text": "first"}, {"type": "text", "text": "second"}]},
|
||||
{"role": "assistant", "content": "reply"},
|
||||
])),
|
||||
..Default::default()
|
||||
},
|
||||
"firstsecondreply"
|
||||
)]
|
||||
#[case::responses_string_input(
|
||||
SemanticCacheContext { input: Some(json!(" plain input ")), ..Default::default() },
|
||||
"plain input"
|
||||
)]
|
||||
#[case::responses_nested_input(
|
||||
SemanticCacheContext {
|
||||
input: Some(json!([{"content": [{"type": "input_text", "text": "nested"}]}, "tail"])),
|
||||
..Default::default()
|
||||
},
|
||||
"nested\ntail"
|
||||
)]
|
||||
#[case::responses_output_text(
|
||||
SemanticCacheContext { input: Some(json!({"output_text": " result text "})), ..Default::default() },
|
||||
"result text"
|
||||
)]
|
||||
#[case::search_results(
|
||||
messages_context(vec![json!({
|
||||
"role": "user",
|
||||
"content": "question",
|
||||
"search_results": [{"source": "src", "title": "t", "content": [{"text": "found"}], "citations": {"a": 1}}],
|
||||
})]),
|
||||
"questionsrctfound{\"a\":1}"
|
||||
)]
|
||||
#[case::empty_messages_fall_back_to_input(
|
||||
SemanticCacheContext { messages: Some(json!([])), input: Some(json!("fallback")), ..Default::default() },
|
||||
"fallback"
|
||||
)]
|
||||
fn prompt_extraction_matches_python_message_and_input_shapes(
|
||||
#[case] context: SemanticCacheContext,
|
||||
#[case] prompt: &str,
|
||||
) {
|
||||
let embedder = FakeEmbedder::new(&[]);
|
||||
let calls = embedder.calls.clone();
|
||||
let cache = cache(
|
||||
vec![
|
||||
info(INDEX, compatible_info(3)),
|
||||
search(INDEX, "key1", &VECTOR, Ok(empty_result())),
|
||||
],
|
||||
embedder,
|
||||
);
|
||||
|
||||
cache.get_cache("key1", &context).unwrap();
|
||||
|
||||
assert_eq!(*calls.lock().unwrap(), vec![(prompt.to_owned(), None)]);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case(None)]
|
||||
#[case(Some(Duration::from_secs(9)))]
|
||||
fn ttl_passes_through_context_only(#[case] ttl: Option<Duration>) {
|
||||
let cache = cache(Vec::new(), FakeEmbedder::new(&[]));
|
||||
|
||||
assert_eq!(
|
||||
cache.get_ttl(&SemanticCacheContext {
|
||||
ttl,
|
||||
..Default::default()
|
||||
}),
|
||||
ttl
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn async_paths_embed_with_metadata_then_run_blocking_redis_work(
|
||||
entry: Value,
|
||||
context: SemanticCacheContext,
|
||||
) {
|
||||
let embedder = FakeEmbedder::new(&[]);
|
||||
let calls = embedder.calls.clone();
|
||||
let cache = cache(
|
||||
vec![
|
||||
info(INDEX, compatible_info(3)),
|
||||
hset(INDEX, PROMPT, "key1", &VECTOR, &entry),
|
||||
search(
|
||||
INDEX,
|
||||
"key1",
|
||||
&VECTOR,
|
||||
Ok(hit("key1", "0.05", encoded(&entry))),
|
||||
),
|
||||
],
|
||||
embedder,
|
||||
);
|
||||
let context = SemanticCacheContext {
|
||||
metadata: Some(json!({"tenant": "team"})),
|
||||
..context
|
||||
};
|
||||
|
||||
cache
|
||||
.async_set_cache("key1", entry.clone(), context.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
cache.async_get_cache("key1", &context).await.unwrap(),
|
||||
Some(entry)
|
||||
);
|
||||
assert_eq!(
|
||||
*calls.lock().unwrap(),
|
||||
vec![(PROMPT.to_owned(), context.metadata.clone()); 2]
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn accessors_report_the_config(config: RedisSemanticConfig) {
|
||||
let cache = cache(Vec::new(), FakeEmbedder::new(&[]));
|
||||
|
||||
assert_eq!(cache.index_name(), config.index_name);
|
||||
assert!((cache.similarity_threshold() - config.similarity_threshold).abs() < 1e-6);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn shared_base_index_across_dimensions_replaces_the_isolated_index(entry: Value) {
|
||||
// Pins parity with Python's `_isolated` + overwrite=True flow.
|
||||
let prompt = "shared prompt";
|
||||
let isolated = format!("{INDEX}_isolated");
|
||||
let context = || messages_context(vec![json!({"role": "user", "content": prompt})]);
|
||||
|
||||
let vector_a = vec![0.1f32; 8];
|
||||
let worker_a = cache(
|
||||
vec![
|
||||
info_missing(INDEX),
|
||||
create_index(INDEX, 8),
|
||||
hset(INDEX, prompt, "key1", &vector_a, &entry),
|
||||
],
|
||||
FakeEmbedder::new(&[(prompt, &vector_a)]),
|
||||
);
|
||||
worker_a
|
||||
.set_cache("key1", entry.clone(), &context())
|
||||
.unwrap();
|
||||
|
||||
let vector_b = vec![0.2f32; 4];
|
||||
let worker_b = cache(
|
||||
vec![
|
||||
info(INDEX, compatible_info(8)),
|
||||
info_missing(&isolated),
|
||||
create_index(&isolated, 4),
|
||||
hset(&isolated, prompt, "key1", &vector_b, &entry),
|
||||
search(
|
||||
&isolated,
|
||||
"key1",
|
||||
&vector_b,
|
||||
Ok(hit("key1", "0.0", encoded(&entry))),
|
||||
),
|
||||
search(
|
||||
&isolated,
|
||||
"key1",
|
||||
&vector_b,
|
||||
Err(redis::RedisError::from((
|
||||
redis::ErrorKind::Extension,
|
||||
"Vector dimension mismatch",
|
||||
))),
|
||||
),
|
||||
],
|
||||
FakeEmbedder::new(&[(prompt, &vector_b)]),
|
||||
);
|
||||
worker_b
|
||||
.set_cache("key1", entry.clone(), &context())
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
worker_b.get_cache("key1", &context()).unwrap(),
|
||||
Some(entry.clone())
|
||||
);
|
||||
|
||||
let vector_c = vec![0.3f32; 16];
|
||||
let worker_c = cache(
|
||||
vec![
|
||||
info(INDEX, compatible_info(8)),
|
||||
info(&isolated, compatible_info(4)),
|
||||
MockCmd::new(redis::cmd("FT.DROPINDEX").arg(&isolated), Ok("OK")),
|
||||
create_index(&isolated, 16),
|
||||
hset(&isolated, prompt, "key1", &vector_c, &entry),
|
||||
],
|
||||
FakeEmbedder::new(&[(prompt, &vector_c)]),
|
||||
);
|
||||
worker_c
|
||||
.set_cache("key1", entry.clone(), &context())
|
||||
.unwrap();
|
||||
|
||||
assert!(matches!(
|
||||
worker_b.get_cache("key1", &context()),
|
||||
Err(Error::Unavailable(Some(_)))
|
||||
));
|
||||
}
|
||||
|
||||
#[fixture]
|
||||
fn redis_stack_url() -> Option<String> {
|
||||
std::env::var("LITELLM_REDIS_STACK_URL").ok()
|
||||
}
|
||||
|
||||
fn live_cache(
|
||||
url: &str,
|
||||
index_name: &str,
|
||||
prompt: &str,
|
||||
vector: Vec<f32>,
|
||||
) -> RedisSemanticCache<FakeEmbedder, JsonCodec<Value>> {
|
||||
RedisSemanticCache::new(
|
||||
url,
|
||||
FakeEmbedder::new(&[(prompt, vector.as_slice())]),
|
||||
JsonCodec::<Value>::new(),
|
||||
RedisSemanticConfig {
|
||||
index_name: index_name.to_owned(),
|
||||
similarity_threshold: 0.9,
|
||||
},
|
||||
)
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn live_shared_index_is_replaced_across_dimensions(redis_stack_url: Option<String>, entry: Value) {
|
||||
let Some(url) = redis_stack_url else {
|
||||
return;
|
||||
};
|
||||
// Pins parity with Python's `_isolated` + overwrite=True flow.
|
||||
let base = format!("rust_semantic_shared_{}", std::process::id());
|
||||
let isolated = format!("{base}_isolated");
|
||||
let prompt = "shared live prompt";
|
||||
let context = || messages_context(vec![json!({"role": "user", "content": prompt})]);
|
||||
|
||||
let worker_a = live_cache(&url, &base, prompt, vec![0.1f32; 8]);
|
||||
worker_a
|
||||
.set_cache("key1", entry.clone(), &context())
|
||||
.unwrap();
|
||||
|
||||
let worker_b = live_cache(&url, &base, prompt, vec![0.2f32; 4]);
|
||||
worker_b
|
||||
.set_cache("key1", entry.clone(), &context())
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
worker_b.get_cache("key1", &context()).unwrap(),
|
||||
Some(entry.clone())
|
||||
);
|
||||
|
||||
let worker_c = live_cache(&url, &base, prompt, vec![0.3f32; 16]);
|
||||
worker_c
|
||||
.set_cache("key1", entry.clone(), &context())
|
||||
.unwrap();
|
||||
|
||||
assert!(matches!(
|
||||
worker_b.get_cache("key1", &context()),
|
||||
Err(Error::Unavailable(Some(_)))
|
||||
));
|
||||
|
||||
let mut connection = redis::Client::open(url).unwrap().get_connection().unwrap();
|
||||
for index in [&base, &isolated] {
|
||||
let _: Result<(), _> = redis::cmd("FT.DROPINDEX")
|
||||
.arg(index)
|
||||
.arg("DD")
|
||||
.query(&mut connection);
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn live_store_lookup_and_ttl_against_redis_stack(redis_stack_url: Option<String>, entry: Value) {
|
||||
let Some(url) = redis_stack_url else {
|
||||
return;
|
||||
};
|
||||
let prompt = "rust semantic cache live prompt";
|
||||
let index_name = format!("rust_semantic_test_{}", std::process::id());
|
||||
let cache = live_cache(&url, &index_name, prompt, vec![0.1, 0.2, 0.3, 0.4]);
|
||||
let context = SemanticCacheContext {
|
||||
ttl: Some(Duration::from_secs(120)),
|
||||
..messages_context(vec![json!({"role": "user", "content": prompt})])
|
||||
};
|
||||
|
||||
cache
|
||||
.set_cache("live-key", entry.clone(), &context)
|
||||
.unwrap();
|
||||
assert_eq!(cache.get_cache("live-key", &context).unwrap(), Some(entry));
|
||||
assert_eq!(cache.get_cache("other-key", &context).unwrap(), None);
|
||||
|
||||
let mut connection = redis::Client::open(url).unwrap().get_connection().unwrap();
|
||||
let ttl: i64 = redis::Commands::ttl(
|
||||
&mut connection,
|
||||
format!("{index_name}:{}", entry_id(prompt, "live-key")),
|
||||
)
|
||||
.unwrap();
|
||||
assert!(
|
||||
ttl > 0,
|
||||
"expected stored hash to carry an expiry, got {ttl}"
|
||||
);
|
||||
}
|
||||
|
|
@ -1,63 +0,0 @@
|
|||
mod support;
|
||||
|
||||
use litellm_cache::{JsonCodec, SemanticCacheContext, semantic::PreparedEmbedding};
|
||||
use litellm_cache_redis_semantic::{DEFAULT_INDEX_NAME, RedisSemanticCache, RedisSemanticConfig};
|
||||
use litellm_cache_testing as contract;
|
||||
use rstest::{fixture, rstest};
|
||||
use serde_json::{Value, json};
|
||||
use support::FakeSearch;
|
||||
|
||||
type Cache = RedisSemanticCache<PreparedEmbedding, JsonCodec<Value>, FakeSearch>;
|
||||
|
||||
const PREFIX: &str = "contract:";
|
||||
|
||||
#[fixture]
|
||||
fn cache() -> Cache {
|
||||
RedisSemanticCache::with_connection(
|
||||
FakeSearch::default(),
|
||||
PreparedEmbedding(vec![0.6, 0.8]),
|
||||
JsonCodec::new(),
|
||||
RedisSemanticConfig {
|
||||
index_name: DEFAULT_INDEX_NAME.into(),
|
||||
similarity_threshold: 0.9,
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
#[fixture]
|
||||
fn context() -> SemanticCacheContext {
|
||||
SemanticCacheContext {
|
||||
messages: Some(json!([{"role": "user", "content": "contract prompt"}])),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn hit_and_miss(cache: Cache, context: SemanticCacheContext) {
|
||||
contract::hit_and_miss(&cache, context, PREFIX, json!({"answer": 42})).await;
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn sync_async_equivalence(cache: Cache, context: SemanticCacheContext) {
|
||||
contract::sync_async_equivalence(&cache, context, PREFIX, json!("first"), json!([2])).await;
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn overwrite_replaces(cache: Cache, context: SemanticCacheContext) {
|
||||
contract::overwrite_replaces(&cache, context, PREFIX, json!(1), json!({"b": 2})).await;
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn pipeline_writes_every_entry(cache: Cache, context: SemanticCacheContext) {
|
||||
contract::pipeline_writes_every_entry(
|
||||
&cache,
|
||||
context,
|
||||
PREFIX,
|
||||
vec![json!("a"), json!(2), json!({"c": true})],
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
|
@ -1,302 +0,0 @@
|
|||
#![allow(dead_code)]
|
||||
|
||||
use std::{
|
||||
collections::{BTreeMap, HashMap},
|
||||
sync::{Arc, Mutex},
|
||||
};
|
||||
|
||||
use litellm_cache::{Error, semantic::Embedder};
|
||||
use serde_json::Value;
|
||||
|
||||
pub type EmbedCalls = Arc<Mutex<Vec<(String, Option<Value>)>>>;
|
||||
|
||||
/// Embeds known prompts to fixed vectors, anything else to `[0.1, 0.2, 0.3]`, and records every
|
||||
/// prompt with its metadata.
|
||||
pub struct FakeEmbedder {
|
||||
vectors: HashMap<String, Vec<f32>>,
|
||||
pub calls: EmbedCalls,
|
||||
}
|
||||
|
||||
impl FakeEmbedder {
|
||||
pub fn new(vectors: &[(&str, &[f32])]) -> Self {
|
||||
Self {
|
||||
vectors: vectors
|
||||
.iter()
|
||||
.map(|(prompt, vector)| ((*prompt).to_owned(), vector.to_vec()))
|
||||
.collect(),
|
||||
calls: EmbedCalls::default(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Embedder for FakeEmbedder {
|
||||
fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result<Vec<f32>, Error> {
|
||||
self.calls
|
||||
.lock()
|
||||
.unwrap()
|
||||
.push((prompt.to_owned(), metadata.cloned()));
|
||||
Ok(self
|
||||
.vectors
|
||||
.get(prompt)
|
||||
.cloned()
|
||||
.unwrap_or_else(|| vec![0.1, 0.2, 0.3]))
|
||||
}
|
||||
|
||||
async fn async_embed(&self, prompt: &str, metadata: Option<&Value>) -> Result<Vec<f32>, Error> {
|
||||
self.embed(prompt, metadata)
|
||||
}
|
||||
}
|
||||
|
||||
struct FakeIndex {
|
||||
prefix: Vec<u8>,
|
||||
dims: usize,
|
||||
vector_field: String,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct SearchState {
|
||||
indexes: HashMap<String, FakeIndex>,
|
||||
hashes: BTreeMap<Vec<u8>, BTreeMap<String, Vec<u8>>>,
|
||||
}
|
||||
|
||||
/// An in-memory Redis Stack speaking the `FT.*`, `HSET` and `EXPIRE` subset the semantic cache
|
||||
/// sends, with exact cosine KNN over the hashes under an index prefix.
|
||||
#[derive(Clone, Default)]
|
||||
pub struct FakeSearch {
|
||||
state: Arc<Mutex<SearchState>>,
|
||||
}
|
||||
|
||||
impl FakeSearch {
|
||||
fn run(&self, args: Vec<Vec<u8>>) -> redis::RedisResult<redis::Value> {
|
||||
let mut state = self.state.lock().unwrap();
|
||||
let text = |index: usize| String::from_utf8_lossy(&args[index]).into_owned();
|
||||
match text(0).to_uppercase().as_str() {
|
||||
"FT.CREATE" => {
|
||||
let name = text(1);
|
||||
if state.indexes.contains_key(&name) {
|
||||
return Err(server_error("Index already exists"));
|
||||
}
|
||||
let position = |token: &str| args.iter().position(|arg| arg == token.as_bytes());
|
||||
let prefix = args[position("PREFIX").unwrap() + 2].clone();
|
||||
let dims = text(position("DIM").unwrap() + 1).parse().unwrap();
|
||||
let vector_field = text(position("VECTOR").unwrap() - 1);
|
||||
state.indexes.insert(
|
||||
name,
|
||||
FakeIndex {
|
||||
prefix,
|
||||
dims,
|
||||
vector_field,
|
||||
},
|
||||
);
|
||||
Ok(redis::Value::Okay)
|
||||
}
|
||||
"FT.INFO" => {
|
||||
let index = state
|
||||
.indexes
|
||||
.get(&text(1))
|
||||
.ok_or_else(|| server_error("Unknown index name"))?;
|
||||
Ok(index_info(index))
|
||||
}
|
||||
"FT.DROPINDEX" => {
|
||||
state.indexes.remove(&text(1));
|
||||
Ok(redis::Value::Okay)
|
||||
}
|
||||
"HSET" => {
|
||||
let hash = state.hashes.entry(args[1].clone()).or_default();
|
||||
for pair in args[2..].chunks(2) {
|
||||
hash.insert(
|
||||
String::from_utf8_lossy(&pair[0]).into_owned(),
|
||||
pair[1].clone(),
|
||||
);
|
||||
}
|
||||
Ok(redis::Value::Int(((args.len() - 2) / 2) as i64))
|
||||
}
|
||||
"EXPIRE" => Ok(redis::Value::Int(i64::from(
|
||||
state.hashes.contains_key(&args[1]),
|
||||
))),
|
||||
"FT.SEARCH" => {
|
||||
let index = state
|
||||
.indexes
|
||||
.get(&text(1))
|
||||
.ok_or_else(|| server_error("no such index"))?;
|
||||
let query = text(2);
|
||||
let tag = query_tag(&query);
|
||||
let params = args.iter().position(|arg| arg == b"PARAMS").unwrap();
|
||||
let vector = floats(&args[params + 3]);
|
||||
let best = state
|
||||
.hashes
|
||||
.iter()
|
||||
.filter(|(key, _)| key.starts_with(&index.prefix))
|
||||
.filter(|(_, fields)| {
|
||||
fields.get("litellm_cache_key").map(Vec::as_slice) == Some(tag.as_bytes())
|
||||
})
|
||||
.filter_map(|(key, fields)| {
|
||||
let stored = floats(fields.get(&index.vector_field)?);
|
||||
(stored.len() == index.dims)
|
||||
.then(|| (key, fields, 1.0 - cosine(&vector, &stored)))
|
||||
})
|
||||
.min_by(|left, right| left.2.total_cmp(&right.2));
|
||||
let Some((key, fields, distance)) = best else {
|
||||
return Ok(redis::Value::Array(vec![redis::Value::Int(0)]));
|
||||
};
|
||||
let mut reply = fields
|
||||
.iter()
|
||||
.filter(|(name, _)| **name != index.vector_field)
|
||||
.flat_map(|(name, value)| [bulk(name.as_bytes()), bulk(value)])
|
||||
.collect::<Vec<_>>();
|
||||
reply.extend([
|
||||
bulk(b"vector_distance"),
|
||||
bulk(distance.to_string().as_bytes()),
|
||||
]);
|
||||
Ok(redis::Value::Array(vec![
|
||||
redis::Value::Int(1),
|
||||
bulk(key),
|
||||
redis::Value::Array(reply),
|
||||
]))
|
||||
}
|
||||
"PING" => Ok(redis::Value::SimpleString("PONG".into())),
|
||||
_ => Err(server_error("unsupported command")),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl redis::ConnectionLike for FakeSearch {
|
||||
fn req_packed_command(&mut self, command: &[u8]) -> redis::RedisResult<redis::Value> {
|
||||
let mut commands = parse_commands(command);
|
||||
self.run(commands.remove(0))
|
||||
}
|
||||
|
||||
fn req_packed_commands(
|
||||
&mut self,
|
||||
commands: &[u8],
|
||||
offset: usize,
|
||||
count: usize,
|
||||
) -> redis::RedisResult<Vec<redis::Value>> {
|
||||
let replies = parse_commands(commands)
|
||||
.into_iter()
|
||||
.map(|args| self.run(args))
|
||||
.collect::<redis::RedisResult<Vec<_>>>()?;
|
||||
Ok(replies.into_iter().skip(offset).take(count).collect())
|
||||
}
|
||||
|
||||
fn get_db(&self) -> i64 {
|
||||
0
|
||||
}
|
||||
|
||||
fn check_connection(&mut self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn is_open(&self) -> bool {
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
pub fn server_error(reply: &str) -> redis::RedisError {
|
||||
redis::parse_redis_value(format!("-{reply}\r\n").as_bytes())
|
||||
.expect("error reply")
|
||||
.extract_error()
|
||||
.expect_err("server error")
|
||||
}
|
||||
|
||||
fn bulk(bytes: &[u8]) -> redis::Value {
|
||||
redis::Value::BulkString(bytes.to_vec())
|
||||
}
|
||||
|
||||
fn index_info(index: &FakeIndex) -> redis::Value {
|
||||
let attribute = |name: &str, field_type: &str| {
|
||||
redis::Value::Array(vec![
|
||||
bulk(b"identifier"),
|
||||
bulk(name.as_bytes()),
|
||||
bulk(b"type"),
|
||||
bulk(field_type.as_bytes()),
|
||||
])
|
||||
};
|
||||
redis::Value::Array(vec![
|
||||
bulk(b"attributes"),
|
||||
redis::Value::Array(vec![
|
||||
attribute("prompt", "TEXT"),
|
||||
attribute("response", "TEXT"),
|
||||
attribute("inserted_at", "NUMERIC"),
|
||||
attribute("updated_at", "NUMERIC"),
|
||||
attribute("litellm_cache_key", "TAG"),
|
||||
redis::Value::Array(vec![
|
||||
bulk(b"identifier"),
|
||||
bulk(index.vector_field.as_bytes()),
|
||||
bulk(b"type"),
|
||||
bulk(b"VECTOR"),
|
||||
bulk(b"dim"),
|
||||
redis::Value::Int(index.dims as i64),
|
||||
bulk(b"data_type"),
|
||||
bulk(b"FLOAT32"),
|
||||
bulk(b"distance_metric"),
|
||||
bulk(b"COSINE"),
|
||||
]),
|
||||
]),
|
||||
])
|
||||
}
|
||||
|
||||
/// The tag inside `@litellm_cache_key:{...}`, with query escapes removed.
|
||||
fn query_tag(query: &str) -> String {
|
||||
let start = query.find("@litellm_cache_key:{").unwrap() + "@litellm_cache_key:{".len();
|
||||
let mut tag = String::new();
|
||||
let mut characters = query[start..].chars();
|
||||
while let Some(character) = characters.next() {
|
||||
match character {
|
||||
'\\' => tag.extend(characters.next()),
|
||||
'}' => break,
|
||||
character => tag.push(character),
|
||||
}
|
||||
}
|
||||
tag
|
||||
}
|
||||
|
||||
fn floats(bytes: &[u8]) -> Vec<f32> {
|
||||
bytes
|
||||
.as_chunks::<4>()
|
||||
.0
|
||||
.iter()
|
||||
.map(|chunk| f32::from_le_bytes(*chunk))
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn cosine(left: &[f32], right: &[f32]) -> f64 {
|
||||
let dot = left
|
||||
.iter()
|
||||
.zip(right)
|
||||
.map(|(left, right)| f64::from(*left) * f64::from(*right))
|
||||
.sum::<f64>();
|
||||
let norm = |vector: &[f32]| {
|
||||
vector
|
||||
.iter()
|
||||
.map(|value| f64::from(*value).powi(2))
|
||||
.sum::<f64>()
|
||||
.sqrt()
|
||||
};
|
||||
dot / (norm(left) * norm(right))
|
||||
}
|
||||
|
||||
/// Splits a packed RESP request into each command's arguments.
|
||||
fn parse_commands(mut bytes: &[u8]) -> Vec<Vec<Vec<u8>>> {
|
||||
let line = |bytes: &mut &[u8]| {
|
||||
let end = bytes
|
||||
.windows(2)
|
||||
.position(|window| window == b"\r\n")
|
||||
.unwrap();
|
||||
let text = String::from_utf8(bytes[1..end].to_vec()).unwrap();
|
||||
*bytes = &bytes[end + 2..];
|
||||
text.parse::<usize>().unwrap()
|
||||
};
|
||||
let mut commands = Vec::new();
|
||||
while !bytes.is_empty() {
|
||||
let count = line(&mut bytes);
|
||||
let mut args = Vec::with_capacity(count);
|
||||
for _ in 0..count {
|
||||
let length = line(&mut bytes);
|
||||
args.push(bytes[..length].to_vec());
|
||||
bytes = &bytes[length + 2..];
|
||||
}
|
||||
commands.push(args);
|
||||
}
|
||||
commands
|
||||
}
|
||||
|
|
@ -1,20 +0,0 @@
|
|||
[package]
|
||||
name = "litellm-cache-redis"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
litellm-cache.workspace = true
|
||||
litellm-error.workspace = true
|
||||
redis = { workspace = true, features = ["cluster", "tls-rustls"] }
|
||||
r2d2 = "0.8.10"
|
||||
thiserror.workspace = true
|
||||
tokio.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
litellm-cache-testing.workspace = true
|
||||
redis-test.workspace = true
|
||||
rstest.workspace = true
|
||||
serde_json.workspace = true
|
||||
|
|
@ -1,192 +0,0 @@
|
|||
use std::{
|
||||
sync::{Arc, OnceLock},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use litellm_cache::{BatchEntry, CacheCodec, Error};
|
||||
|
||||
use crate::{
|
||||
connection::{ConnectionRef, Connections},
|
||||
topology::RedisTopology,
|
||||
};
|
||||
|
||||
const DEFAULT_TTL: Duration = Duration::from_secs(600);
|
||||
|
||||
pub struct RedisCache<S, C = redis::Connection> {
|
||||
pub(crate) connections: Arc<Connections<C>>,
|
||||
pub(crate) default_ttl: Duration,
|
||||
pub(crate) codec: S,
|
||||
pub(crate) namespace: Option<String>,
|
||||
pub(crate) topology: RedisTopology,
|
||||
/// The server's major version, read from `INFO` once, like Python's `redis_version`.
|
||||
pub(crate) major_version: Arc<OnceLock<u32>>,
|
||||
}
|
||||
|
||||
impl<S: CacheCodec> RedisCache<S> {
|
||||
pub fn new(url: &str, default_ttl: Option<Duration>, codec: S) -> Result<Self, Error> {
|
||||
Self::connect(url, &RedisTopology::Standalone, default_ttl, codec)
|
||||
}
|
||||
|
||||
pub fn connect(
|
||||
url: &str,
|
||||
topology: &RedisTopology,
|
||||
default_ttl: Option<Duration>,
|
||||
codec: S,
|
||||
) -> Result<Self, Error> {
|
||||
let connections = Connections::open(url, topology)?;
|
||||
Ok(Self {
|
||||
connections: Arc::new(connections),
|
||||
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
|
||||
codec,
|
||||
namespace: None,
|
||||
topology: topology.clone(),
|
||||
major_version: Arc::default(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, C> RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
pub fn with_connection(connection: C, default_ttl: Option<Duration>, codec: S) -> Self {
|
||||
Self {
|
||||
connections: Arc::new(Connections::fixed(connection)),
|
||||
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
|
||||
codec,
|
||||
namespace: None,
|
||||
topology: RedisTopology::Standalone,
|
||||
major_version: Arc::default(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_namespace(self, namespace: Option<String>) -> Self {
|
||||
Self {
|
||||
namespace: namespace.filter(|value| !value.is_empty()),
|
||||
..self
|
||||
}
|
||||
}
|
||||
|
||||
pub fn namespace(&self) -> Option<&str> {
|
||||
self.namespace.as_deref()
|
||||
}
|
||||
|
||||
pub fn topology(&self) -> &RedisTopology {
|
||||
&self.topology
|
||||
}
|
||||
|
||||
pub(crate) fn namespaced_key(&self, key: &str) -> String {
|
||||
namespaced_key(self.namespace.as_deref(), key)
|
||||
}
|
||||
|
||||
pub(crate) fn namespaced_keys(&self, keys: &[String]) -> Vec<String> {
|
||||
keys.iter().map(|key| self.namespaced_key(key)).collect()
|
||||
}
|
||||
|
||||
/// Whole seconds for `ttl`, falling back to the default TTL like Python's `get_ttl`.
|
||||
pub(crate) fn ttl_or_default(&self, ttl: Option<Duration>) -> u64 {
|
||||
ttl_seconds(ttl.unwrap_or(self.default_ttl))
|
||||
}
|
||||
|
||||
pub(crate) fn execute<T>(
|
||||
&self,
|
||||
operation: impl FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error>,
|
||||
) -> Result<T, Error> {
|
||||
self.connections.execute(operation)
|
||||
}
|
||||
|
||||
/// `_parse_redis_major_version`: the major version from `INFO`, or
|
||||
/// `DEFAULT_REDIS_MAJOR_VERSION` when `INFO` fails or its version does not parse. The first
|
||||
/// answer is kept, as Python reads `redis_version` once at construction.
|
||||
pub(crate) async fn major_version(&self) -> u32 {
|
||||
if let Some(version) = self.major_version.get() {
|
||||
return *version;
|
||||
}
|
||||
let info = self
|
||||
.run(|connection| connection.node_text(&redis::cmd("INFO")))
|
||||
.await;
|
||||
let version = info
|
||||
.ok()
|
||||
.and_then(|info| parse_major_version(&info))
|
||||
.unwrap_or_else(default_major_version);
|
||||
*self.major_version.get_or_init(|| version)
|
||||
}
|
||||
|
||||
pub(crate) async fn run<T, F>(&self, operation: F) -> Result<T, Error>
|
||||
where
|
||||
T: Send + 'static,
|
||||
F: FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error> + Send + 'static,
|
||||
{
|
||||
Connections::run_blocking(Arc::clone(&self.connections), operation).await
|
||||
}
|
||||
|
||||
pub(crate) fn namespaced_pattern(&self) -> Result<String, Error> {
|
||||
let namespace = self.namespace.as_ref().ok_or(Error::UnscopedFlush)?;
|
||||
let escaped: String = namespace
|
||||
.chars()
|
||||
.flat_map(|ch| {
|
||||
if matches!(ch, '*' | '?' | '[' | ']' | '\\') {
|
||||
vec!['\\', ch]
|
||||
} else {
|
||||
vec![ch]
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
Ok(format!("{escaped}:*"))
|
||||
}
|
||||
|
||||
pub(crate) fn decode_response(&self, value: redis::Value) -> Result<Option<S::Value>, Error> {
|
||||
match value {
|
||||
redis::Value::Nil => Ok(None),
|
||||
redis::Value::BulkString(bytes) => self.codec.decode(&bytes).map(Some),
|
||||
redis::Value::SimpleString(text) => self.codec.decode(text.as_bytes()).map(Some),
|
||||
_ => Err(Error::InvalidEntry(None)),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn decode_batch_response(
|
||||
&self,
|
||||
value: redis::Value,
|
||||
) -> Result<BatchEntry<S::Value>, Error> {
|
||||
match self.decode_response(value) {
|
||||
Ok(Some(value)) => Ok(BatchEntry::Hit(value)),
|
||||
Ok(None) => Ok(BatchEntry::Miss),
|
||||
Err(Error::InvalidEntry(_)) => Ok(BatchEntry::Invalid),
|
||||
Err(error) => Err(error),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn namespaced_key(namespace: Option<&str>, key: &str) -> String {
|
||||
match namespace {
|
||||
Some(namespace) if !key.starts_with(&format!("{namespace}:")) => {
|
||||
format!("{namespace}:{key}")
|
||||
}
|
||||
_ => key.into(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn ttl_seconds(ttl: Duration) -> u64 {
|
||||
ttl.as_secs()
|
||||
.saturating_add(u64::from(ttl.subsec_nanos() > 0))
|
||||
.max(1)
|
||||
}
|
||||
|
||||
fn parse_major_version(info: &str) -> Option<u32> {
|
||||
let version = info
|
||||
.lines()
|
||||
.find_map(|line| line.trim().strip_prefix("redis_version:"))?
|
||||
.trim();
|
||||
match version.split_once('.') {
|
||||
Some((major, _)) => major.parse().ok(),
|
||||
None => version.parse::<f64>().ok().map(|major| major as u32),
|
||||
}
|
||||
}
|
||||
|
||||
fn default_major_version() -> u32 {
|
||||
std::env::var("DEFAULT_REDIS_MAJOR_VERSION")
|
||||
.ok()
|
||||
.and_then(|value| value.parse().ok())
|
||||
.unwrap_or(7)
|
||||
}
|
||||
|
|
@ -1,109 +0,0 @@
|
|||
use litellm_cache::{CacheCodec, ClaimCache, Error, ExactCacheContext};
|
||||
use redis::Commands;
|
||||
|
||||
use crate::{cache::RedisCache, connection::ConnectionRef};
|
||||
|
||||
const CLAIM_SCRIPT: &str = concat!(
|
||||
"local current = redis.call('GET', KEYS[1]); ",
|
||||
"if ARGV[1] == '' then if current ~= false and current ~= '' then return 0; end; ",
|
||||
"elseif current ~= ARGV[1] then return 0; end; ",
|
||||
"if ARGV[3] ~= '' then redis.call('SET', KEYS[1], ARGV[3], 'EX', ARGV[2]); ",
|
||||
"elseif ARGV[4] == '1' then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return 1"
|
||||
);
|
||||
const CLAIM_ATTEMPTS: usize = 8;
|
||||
|
||||
fn stored_bytes(value: redis::Value) -> Result<Option<Vec<u8>>, Error> {
|
||||
match value {
|
||||
redis::Value::Nil => Ok(None),
|
||||
redis::Value::BulkString(bytes) => Ok(Some(bytes)),
|
||||
redis::Value::SimpleString(text) => Ok(Some(text.into_bytes())),
|
||||
_ => Err(Error::InvalidEntry(None)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Eligibility is decided on decoded values, so a pin written by another encoder (Python's
|
||||
/// `json.dumps` spacing or key order) still matches. The write is a compare-and-set on the
|
||||
/// bytes that decision was made on, retried when another claimant wins the race.
|
||||
fn claim<S: CacheCodec>(
|
||||
connection: &mut ConnectionRef<'_>,
|
||||
codec: &S,
|
||||
key: &str,
|
||||
candidate: S::Value,
|
||||
eligible: &[S::Value],
|
||||
ttl: u64,
|
||||
) -> Result<S::Value, Error>
|
||||
where
|
||||
S::Value: PartialEq,
|
||||
{
|
||||
let payload = codec.encode(&candidate)?;
|
||||
if payload.is_empty() {
|
||||
return Err(Error::InvalidEntry(None));
|
||||
}
|
||||
for _ in 0..CLAIM_ATTEMPTS {
|
||||
let current = stored_bytes(
|
||||
connection
|
||||
.get::<_, redis::Value>(key)
|
||||
.map_err(Error::unavailable)?,
|
||||
)?
|
||||
.filter(|bytes| !bytes.is_empty());
|
||||
let existing = current
|
||||
.as_deref()
|
||||
.and_then(|bytes| codec.decode(bytes).ok())
|
||||
.filter(|existing| eligible.is_empty() || eligible.contains(existing));
|
||||
let refresh = existing
|
||||
.as_ref()
|
||||
.is_some_and(|existing| !eligible.is_empty() || *existing == candidate);
|
||||
let write: &[u8] = if existing.is_some() { b"" } else { &payload };
|
||||
let applied = redis::cmd("EVAL")
|
||||
.arg(CLAIM_SCRIPT)
|
||||
.arg(1)
|
||||
.arg(key)
|
||||
.arg(current.as_deref().unwrap_or_default())
|
||||
.arg(ttl)
|
||||
.arg(write)
|
||||
.arg(u8::from(refresh))
|
||||
.query::<bool>(connection)
|
||||
.map_err(Error::unavailable)?;
|
||||
if applied {
|
||||
return Ok(existing.unwrap_or(candidate));
|
||||
}
|
||||
}
|
||||
Err(Error::unavailable(ClaimContended))
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
#[error("claim kept losing the race after {CLAIM_ATTEMPTS} attempts")]
|
||||
struct ClaimContended;
|
||||
|
||||
impl<S, C> ClaimCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec + Clone + 'static,
|
||||
S::Value: PartialEq,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
fn claim_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
candidate: S::Value,
|
||||
eligible: &[S::Value],
|
||||
context: ExactCacheContext,
|
||||
) -> Result<S::Value, Error> {
|
||||
let key = self.namespaced_key(key);
|
||||
let ttl = self.ttl_or_default(context.ttl);
|
||||
self.execute(|connection| claim(connection, &self.codec, &key, candidate, eligible, ttl))
|
||||
}
|
||||
|
||||
async fn async_claim_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
candidate: S::Value,
|
||||
eligible: Vec<S::Value>,
|
||||
context: ExactCacheContext,
|
||||
) -> Result<S::Value, Error> {
|
||||
let key = self.namespaced_key(key);
|
||||
let ttl = self.ttl_or_default(context.ttl);
|
||||
let codec = self.codec.clone();
|
||||
self.run(move |connection| claim(connection, &codec, &key, candidate, &eligible, ttl))
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
|
@ -1,513 +0,0 @@
|
|||
use std::{
|
||||
sync::{Arc, Mutex},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use litellm_cache::Error;
|
||||
use redis::{
|
||||
ConnectionAddr, ConnectionInfo, IntoConnectionInfo,
|
||||
cluster::{
|
||||
ClusterClient, ClusterClientBuilder, ClusterConnection, ClusterPipeline, NodeAddress,
|
||||
},
|
||||
cluster_routing::{
|
||||
MultipleNodeRoutingInfo, ResponsePolicy, RoutingInfo, SingleNodeRoutingInfo,
|
||||
},
|
||||
};
|
||||
|
||||
use crate::topology::{RedisNode, RedisTopology};
|
||||
|
||||
pub(crate) const REDIS_TIMEOUT: Duration = Duration::from_secs(5);
|
||||
const REDIS_POOL_SIZE: u32 = 16;
|
||||
|
||||
#[allow(private_interfaces)]
|
||||
pub enum Connections<C> {
|
||||
Pool(r2d2::Pool<ConnectionManager>),
|
||||
Cluster(r2d2::Pool<ClusterConnectionManager>),
|
||||
Fixed(Mutex<C>),
|
||||
}
|
||||
|
||||
impl<C> Connections<C>
|
||||
where
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
pub fn execute<T>(
|
||||
&self,
|
||||
operation: impl FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error>,
|
||||
) -> Result<T, Error> {
|
||||
match self {
|
||||
Self::Pool(pool) => {
|
||||
let mut pooled = pool.get().map_err(Error::unavailable)?;
|
||||
let result = operation(&mut ConnectionRef::Node(&mut pooled.connection));
|
||||
pooled.failed = matches!(result, Err(Error::Unavailable(_)));
|
||||
result
|
||||
}
|
||||
Self::Cluster(pool) => {
|
||||
let mut pooled = pool.get().map_err(Error::unavailable)?;
|
||||
let result = operation(&mut ConnectionRef::Cluster(&mut pooled.connection));
|
||||
pooled.failed = matches!(result, Err(Error::Unavailable(_)));
|
||||
result
|
||||
}
|
||||
Self::Fixed(connection) => {
|
||||
let mut connection = connection.lock()?;
|
||||
operation(&mut ConnectionRef::Node(&mut *connection))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn run_blocking<T, F>(connections: Arc<Self>, operation: F) -> Result<T, Error>
|
||||
where
|
||||
T: Send + 'static,
|
||||
F: FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error> + Send + 'static,
|
||||
{
|
||||
tokio::task::spawn_blocking(move || connections.execute(operation))
|
||||
.await
|
||||
.map_err(Error::unavailable)?
|
||||
}
|
||||
|
||||
pub fn fixed(connection: C) -> Self {
|
||||
Self::Fixed(Mutex::new(connection))
|
||||
}
|
||||
|
||||
pub fn open(url: &str, topology: &RedisTopology) -> Result<Self, Error> {
|
||||
match topology {
|
||||
RedisTopology::Standalone => Ok(Self::Pool(pool(ConnectionManager::open(url)?)?)),
|
||||
RedisTopology::Cluster { startup_nodes } => Ok(Self::Cluster(pool(
|
||||
ClusterConnectionManager::open(url, startup_nodes)?,
|
||||
)?)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Closes every idle pooled connection; the next operation opens a fresh one. Connections
|
||||
/// checked out right now return to the pool, and a caller-owned connection stays open.
|
||||
pub fn disconnect(&self) {
|
||||
match self {
|
||||
Self::Pool(pool) => close_idle(pool),
|
||||
Self::Cluster(pool) => close_idle(pool),
|
||||
Self::Fixed(_) => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn pool<M: r2d2::ManageConnection>(manager: M) -> Result<r2d2::Pool<M>, Error> {
|
||||
r2d2::Pool::builder()
|
||||
.max_size(REDIS_POOL_SIZE)
|
||||
.min_idle(Some(0))
|
||||
.connection_timeout(REDIS_TIMEOUT)
|
||||
.test_on_check_out(false)
|
||||
.build(manager)
|
||||
.map_err(Error::unavailable)
|
||||
}
|
||||
|
||||
fn close_idle<M, T>(pool: &r2d2::Pool<M>)
|
||||
where
|
||||
M: r2d2::ManageConnection<Connection = PooledConnection<T>>,
|
||||
{
|
||||
let mut idle = Vec::new();
|
||||
while let Some(mut connection) = pool.try_get() {
|
||||
connection.failed = true;
|
||||
idle.push(connection);
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) struct PooledConnection<C> {
|
||||
connection: C,
|
||||
failed: bool,
|
||||
}
|
||||
|
||||
/// Pools connections without a checkout PING, which would double every operation's round trips.
|
||||
/// A timed-out command leaves its reply on the socket while redis still reports the connection
|
||||
/// open, so any connection whose operation failed is discarded instead of being reused.
|
||||
pub(crate) struct ConnectionManager(redis::Client);
|
||||
|
||||
impl ConnectionManager {
|
||||
fn open(url: &str) -> Result<Self, Error> {
|
||||
redis::Client::open(url)
|
||||
.map(Self)
|
||||
.map_err(Error::unavailable)
|
||||
}
|
||||
}
|
||||
|
||||
impl r2d2::ManageConnection for ConnectionManager {
|
||||
type Connection = PooledConnection<redis::Connection>;
|
||||
type Error = redis::RedisError;
|
||||
|
||||
fn connect(&self) -> Result<Self::Connection, redis::RedisError> {
|
||||
let connection = self.0.get_connection()?;
|
||||
connection.set_read_timeout(Some(REDIS_TIMEOUT))?;
|
||||
connection.set_write_timeout(Some(REDIS_TIMEOUT))?;
|
||||
Ok(PooledConnection {
|
||||
connection,
|
||||
failed: false,
|
||||
})
|
||||
}
|
||||
|
||||
fn is_valid(&self, connection: &mut Self::Connection) -> Result<(), redis::RedisError> {
|
||||
redis::cmd("PING").query::<String>(&mut connection.connection)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn has_broken(&self, connection: &mut Self::Connection) -> bool {
|
||||
connection.failed || !redis::ConnectionLike::is_open(&connection.connection)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
enum TopologyFault {
|
||||
#[error("Redis Cluster needs at least one startup node")]
|
||||
NoStartupNodes,
|
||||
#[error("Redis Cluster nodes need a TCP address")]
|
||||
NonTcpAddress,
|
||||
#[error("Redis Cluster answered a fan-out command without a per-node map")]
|
||||
FanOutReply,
|
||||
#[error("scan target does not match the connection topology")]
|
||||
ScanTarget,
|
||||
}
|
||||
|
||||
pub(crate) struct ClusterConnectionManager(ClusterClient);
|
||||
|
||||
impl ClusterConnectionManager {
|
||||
fn open(url: &str, startup_nodes: &[RedisNode]) -> Result<Self, Error> {
|
||||
if startup_nodes.is_empty() {
|
||||
return Err(Error::unavailable(TopologyFault::NoStartupNodes));
|
||||
}
|
||||
let info = url.into_connection_info().map_err(Error::unavailable)?;
|
||||
let nodes = startup_nodes
|
||||
.iter()
|
||||
.map(|node| node_info(&info, node))
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
ClusterClientBuilder::new(nodes)
|
||||
.connection_timeout(REDIS_TIMEOUT)
|
||||
.response_timeout(REDIS_TIMEOUT)
|
||||
.build()
|
||||
.map(Self)
|
||||
.map_err(Error::unavailable)
|
||||
}
|
||||
}
|
||||
|
||||
fn node_info(info: &ConnectionInfo, node: &RedisNode) -> Result<ConnectionInfo, Error> {
|
||||
let addr = match info.addr() {
|
||||
ConnectionAddr::Tcp(..) => ConnectionAddr::Tcp(node.host.clone(), node.port),
|
||||
ConnectionAddr::TcpTls {
|
||||
insecure,
|
||||
tls_params,
|
||||
..
|
||||
} => ConnectionAddr::TcpTls {
|
||||
host: node.host.clone(),
|
||||
port: node.port,
|
||||
insecure: *insecure,
|
||||
tls_params: tls_params.clone(),
|
||||
},
|
||||
_ => return Err(Error::unavailable(TopologyFault::NonTcpAddress)),
|
||||
};
|
||||
Ok(info.clone().set_addr(addr))
|
||||
}
|
||||
|
||||
impl r2d2::ManageConnection for ClusterConnectionManager {
|
||||
type Connection = PooledConnection<ClusterConnection>;
|
||||
type Error = redis::RedisError;
|
||||
|
||||
fn connect(&self) -> Result<Self::Connection, redis::RedisError> {
|
||||
let connection = self.0.get_connection()?;
|
||||
connection.set_read_timeout(Some(REDIS_TIMEOUT))?;
|
||||
connection.set_write_timeout(Some(REDIS_TIMEOUT))?;
|
||||
Ok(PooledConnection {
|
||||
connection,
|
||||
failed: false,
|
||||
})
|
||||
}
|
||||
|
||||
fn is_valid(&self, connection: &mut Self::Connection) -> Result<(), redis::RedisError> {
|
||||
redis::cmd("PING").query::<String>(&mut connection.connection)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn has_broken(&self, connection: &mut Self::Connection) -> bool {
|
||||
connection.failed || !redis::ConnectionLike::is_open(&connection.connection)
|
||||
}
|
||||
}
|
||||
|
||||
pub enum ConnectionRef<'a> {
|
||||
Node(&'a mut dyn redis::ConnectionLike),
|
||||
Cluster(&'a mut ClusterConnection),
|
||||
}
|
||||
|
||||
impl redis::ConnectionLike for ConnectionRef<'_> {
|
||||
fn req_packed_command(&mut self, cmd: &[u8]) -> redis::RedisResult<redis::Value> {
|
||||
match self {
|
||||
Self::Node(connection) => connection.req_packed_command(cmd),
|
||||
Self::Cluster(connection) => connection.req_packed_command(cmd),
|
||||
}
|
||||
}
|
||||
|
||||
fn req_packed_commands(
|
||||
&mut self,
|
||||
cmd: &[u8],
|
||||
offset: usize,
|
||||
count: usize,
|
||||
) -> redis::RedisResult<Vec<redis::Value>> {
|
||||
match self {
|
||||
Self::Node(connection) => connection.req_packed_commands(cmd, offset, count),
|
||||
Self::Cluster(connection) => connection.req_packed_commands(cmd, offset, count),
|
||||
}
|
||||
}
|
||||
|
||||
fn get_db(&self) -> i64 {
|
||||
match self {
|
||||
Self::Node(connection) => connection.get_db(),
|
||||
Self::Cluster(connection) => redis::ConnectionLike::get_db(*connection),
|
||||
}
|
||||
}
|
||||
|
||||
fn supports_pipelining(&self) -> bool {
|
||||
match self {
|
||||
Self::Node(connection) => connection.supports_pipelining(),
|
||||
Self::Cluster(connection) => redis::ConnectionLike::supports_pipelining(*connection),
|
||||
}
|
||||
}
|
||||
|
||||
fn check_connection(&mut self) -> bool {
|
||||
match self {
|
||||
Self::Node(connection) => connection.check_connection(),
|
||||
Self::Cluster(connection) => connection.check_connection(),
|
||||
}
|
||||
}
|
||||
|
||||
fn is_open(&self) -> bool {
|
||||
match self {
|
||||
Self::Node(connection) => connection.is_open(),
|
||||
Self::Cluster(connection) => redis::ConnectionLike::is_open(*connection),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ConnectionRef<'_> {
|
||||
/// Runs `pipeline` and decodes its non-ignored replies as `T`. A cluster connection refuses
|
||||
/// `Pipeline::query`, so there a transaction goes to its keys' slot as one MULTI/EXEC and
|
||||
/// anything else is split per node by `ClusterPipeline`; either way the raw replies are
|
||||
/// handed back to `pipeline` to decode.
|
||||
pub(crate) fn query_pipeline<T: redis::FromRedisValue>(
|
||||
&mut self,
|
||||
pipeline: &redis::Pipeline,
|
||||
) -> Result<T, Error> {
|
||||
match self {
|
||||
Self::Node(connection) => pipeline.query(*connection),
|
||||
Self::Cluster(connection) if pipeline.is_transaction() => {
|
||||
redis::ConnectionLike::req_packed_commands(
|
||||
*connection,
|
||||
&pipeline.get_packed_pipeline(),
|
||||
pipeline.len() + 1,
|
||||
1,
|
||||
)
|
||||
.and_then(|replies| pipeline.query(&mut Replies(Some(replies))))
|
||||
}
|
||||
Self::Cluster(connection) => {
|
||||
let mut cluster = ClusterPipeline::with_capacity(pipeline.len());
|
||||
for command in pipeline.cmd_iter() {
|
||||
cluster.add_command(command.clone());
|
||||
}
|
||||
cluster
|
||||
.query(connection)
|
||||
.and_then(|replies| pipeline.query(&mut Replies(Some(replies))))
|
||||
}
|
||||
}
|
||||
.map_err(Error::unavailable)
|
||||
}
|
||||
|
||||
pub(crate) fn scan(
|
||||
&mut self,
|
||||
pattern: &str,
|
||||
count: usize,
|
||||
mut visit: impl FnMut(&mut Self, Vec<String>) -> Result<bool, Error>,
|
||||
) -> Result<(), Error> {
|
||||
let pages = match self {
|
||||
Self::Node(connection) => {
|
||||
let page = scan_command(0, pattern, count)
|
||||
.query::<ScanPage>(*connection)
|
||||
.map_err(Error::unavailable)?;
|
||||
vec![(None, page)]
|
||||
}
|
||||
Self::Cluster(connection) => connection
|
||||
.route_command(
|
||||
&scan_command(0, pattern, count),
|
||||
RoutingInfo::MultiNode((
|
||||
MultipleNodeRoutingInfo::AllMasters,
|
||||
Some(ResponsePolicy::Special),
|
||||
)),
|
||||
)
|
||||
.map_err(Error::unavailable)
|
||||
.and_then(primary_pages)?
|
||||
.into_iter()
|
||||
.map(|(node, page)| (Some(node), page))
|
||||
.collect(),
|
||||
};
|
||||
for (node, (mut cursor, mut keys)) in pages {
|
||||
loop {
|
||||
if !visit(self, keys)? {
|
||||
return Ok(());
|
||||
}
|
||||
if cursor == 0 {
|
||||
break;
|
||||
}
|
||||
(cursor, keys) = self.scan_page(node.as_ref(), cursor, pattern, count)?;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) fn ping(&mut self) -> Result<bool, redis::RedisError> {
|
||||
let command = redis::cmd("PING");
|
||||
match self {
|
||||
Self::Node(connection) => command
|
||||
.query::<String>(*connection)
|
||||
.map(|response| response == "PONG"),
|
||||
Self::Cluster(connection) => connection
|
||||
.route_command(
|
||||
&command,
|
||||
RoutingInfo::MultiNode((
|
||||
MultipleNodeRoutingInfo::AllNodes,
|
||||
Some(ResponsePolicy::AllSucceeded),
|
||||
)),
|
||||
)
|
||||
.map(|_| true),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn node_text(&mut self, command: &redis::Cmd) -> Result<String, Error> {
|
||||
match self {
|
||||
Self::Node(connection) => command.query(*connection).map_err(Error::unavailable),
|
||||
Self::Cluster(connection) => {
|
||||
let value = connection
|
||||
.route_command(
|
||||
command,
|
||||
RoutingInfo::MultiNode((
|
||||
MultipleNodeRoutingInfo::AllNodes,
|
||||
Some(ResponsePolicy::Special),
|
||||
)),
|
||||
)
|
||||
.map_err(Error::unavailable)?;
|
||||
let redis::Value::Map(entries) = value else {
|
||||
return Err(Error::unavailable(TopologyFault::FanOutReply));
|
||||
};
|
||||
let mut replies = entries
|
||||
.into_iter()
|
||||
.map(|(node, reply)| {
|
||||
Ok((
|
||||
redis::from_redis_value::<String>(node).map_err(Error::unavailable)?,
|
||||
redis::from_redis_value::<String>(reply).map_err(Error::unavailable)?,
|
||||
))
|
||||
})
|
||||
.collect::<Result<Vec<(String, String)>, Error>>()?;
|
||||
replies.sort();
|
||||
Ok(replies
|
||||
.into_iter()
|
||||
.map(|(_, reply)| reply)
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n"))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn flushall(&mut self) -> Result<(), Error> {
|
||||
let command = redis::cmd("FLUSHALL");
|
||||
match self {
|
||||
Self::Node(connection) => command.query(*connection).map_err(Error::unavailable),
|
||||
Self::Cluster(connection) => connection
|
||||
.route_command(
|
||||
&command,
|
||||
RoutingInfo::MultiNode((
|
||||
MultipleNodeRoutingInfo::AllMasters,
|
||||
Some(ResponsePolicy::AllSucceeded),
|
||||
)),
|
||||
)
|
||||
.map(|_| ())
|
||||
.map_err(Error::unavailable),
|
||||
}
|
||||
}
|
||||
|
||||
fn scan_page(
|
||||
&mut self,
|
||||
node: Option<&NodeAddress>,
|
||||
cursor: u64,
|
||||
pattern: &str,
|
||||
count: usize,
|
||||
) -> Result<ScanPage, Error> {
|
||||
let command = scan_command(cursor, pattern, count);
|
||||
match (self, node) {
|
||||
(Self::Node(connection), None) => {
|
||||
command.query(*connection).map_err(Error::unavailable)
|
||||
}
|
||||
(Self::Cluster(connection), Some(node)) => connection
|
||||
.route_command(
|
||||
&command,
|
||||
RoutingInfo::SingleNode(SingleNodeRoutingInfo::ByAddress {
|
||||
host: node.host().to_string(),
|
||||
port: node.port(),
|
||||
}),
|
||||
)
|
||||
.map_err(Error::unavailable)
|
||||
.and_then(|value| redis::from_redis_value(value).map_err(Error::unavailable)),
|
||||
_ => Err(Error::unavailable(TopologyFault::ScanTarget)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
type ScanPage = (u64, Vec<String>);
|
||||
|
||||
fn primary_pages(value: redis::Value) -> Result<Vec<(NodeAddress, ScanPage)>, Error> {
|
||||
let redis::Value::Map(entries) = value else {
|
||||
return Err(Error::unavailable(TopologyFault::FanOutReply));
|
||||
};
|
||||
entries
|
||||
.into_iter()
|
||||
.map(|(node, page)| {
|
||||
let node = redis::from_redis_value::<String>(node).map_err(Error::unavailable)?;
|
||||
let node = NodeAddress::try_from(node.as_str()).map_err(Error::unavailable)?;
|
||||
let page = redis::from_redis_value::<ScanPage>(page).map_err(Error::unavailable)?;
|
||||
Ok((node, page))
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn scan_command(cursor: u64, pattern: &str, count: usize) -> redis::Cmd {
|
||||
let mut command = redis::cmd("SCAN");
|
||||
command
|
||||
.cursor_arg(cursor)
|
||||
.arg("MATCH")
|
||||
.arg(pattern)
|
||||
.arg("COUNT")
|
||||
.arg(count);
|
||||
command
|
||||
}
|
||||
|
||||
/// Hands already received pipeline replies to `Pipeline::query`, so it applies its own
|
||||
/// ignore and error handling to replies a cluster pipeline gathered from several nodes.
|
||||
struct Replies(Option<Vec<redis::Value>>);
|
||||
|
||||
impl redis::ConnectionLike for Replies {
|
||||
fn req_packed_command(&mut self, _: &[u8]) -> redis::RedisResult<redis::Value> {
|
||||
Err((redis::ErrorKind::Client, "replies hold a pipeline only").into())
|
||||
}
|
||||
|
||||
fn req_packed_commands(
|
||||
&mut self,
|
||||
_: &[u8],
|
||||
_: usize,
|
||||
_: usize,
|
||||
) -> redis::RedisResult<Vec<redis::Value>> {
|
||||
self.0
|
||||
.take()
|
||||
.ok_or_else(|| (redis::ErrorKind::Client, "replies were already read").into())
|
||||
}
|
||||
|
||||
fn get_db(&self) -> i64 {
|
||||
0
|
||||
}
|
||||
|
||||
fn check_connection(&mut self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn is_open(&self) -> bool {
|
||||
true
|
||||
}
|
||||
}
|
||||
|
|
@ -1,205 +0,0 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use litellm_cache::{
|
||||
BoundedCounterCache, CacheCodec, CountReadCache, CounterCache, Error, ExactCacheContext,
|
||||
IncrementOperation,
|
||||
};
|
||||
|
||||
use crate::{
|
||||
cache::{RedisCache, ttl_seconds},
|
||||
connection::ConnectionRef,
|
||||
store::mget,
|
||||
};
|
||||
|
||||
const INCREMENT_SCRIPT: &str = concat!(
|
||||
"local value = redis.call('INCRBYFLOAT', KEYS[1], ARGV[1]); ",
|
||||
"if redis.call('TTL', KEYS[1]) == -1 then ",
|
||||
"redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return value"
|
||||
);
|
||||
const INCREMENT_WITH_FLOOR_SCRIPT: &str = concat!(
|
||||
"local count = redis.call('INCRBY', KEYS[1], ARGV[1]); ",
|
||||
"if count < 0 then count = redis.call('INCRBY', KEYS[1], -count); end; ",
|
||||
"if redis.call('TTL', KEYS[1]) < 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; ",
|
||||
"return count"
|
||||
);
|
||||
const SET_MAX_SCRIPT: &str = concat!(
|
||||
"local current = redis.call('GET', KEYS[1]); ",
|
||||
"if current == false or tonumber(current) < tonumber(ARGV[1]) then ",
|
||||
"redis.call('SET', KEYS[1], ARGV[1]); ",
|
||||
"if tonumber(ARGV[2]) > 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; ",
|
||||
"return ARGV[1]; end; return current"
|
||||
);
|
||||
|
||||
impl<S, C> CounterCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
fn increment_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
amount: f64,
|
||||
context: ExactCacheContext,
|
||||
) -> Result<f64, Error> {
|
||||
let key = self.namespaced_key(key);
|
||||
let ttl = self.ttl_or_default(context.ttl);
|
||||
self.execute(|connection| increment(connection, key, amount, ttl, false))
|
||||
}
|
||||
|
||||
/// Python `_incrbyfloat_with_ttl`: without `refresh_ttl` the TTL is set only on a key that
|
||||
/// has none, in one atomic script; with it, every increment re-arms the TTL.
|
||||
async fn async_increment(
|
||||
&self,
|
||||
key: &str,
|
||||
amount: f64,
|
||||
context: ExactCacheContext,
|
||||
refresh_ttl: bool,
|
||||
) -> Result<f64, Error> {
|
||||
let key = self.namespaced_key(key);
|
||||
let ttl = self.ttl_or_default(context.ttl);
|
||||
self.run(move |connection| increment(connection, key, amount, ttl, refresh_ttl))
|
||||
.await
|
||||
}
|
||||
|
||||
async fn async_increment_pipeline(
|
||||
&self,
|
||||
operations: Vec<IncrementOperation>,
|
||||
) -> Result<Vec<f64>, Error> {
|
||||
if operations.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let mut pipeline = redis::pipe();
|
||||
for operation in operations {
|
||||
let key = self.namespaced_key(&operation.key);
|
||||
pipeline.cmd("INCRBYFLOAT").arg(&key).arg(operation.amount);
|
||||
if let Some(ttl) = operation.ttl {
|
||||
pipeline
|
||||
.cmd("EXPIRE")
|
||||
.arg(key)
|
||||
.arg(ttl_seconds(ttl))
|
||||
.ignore();
|
||||
}
|
||||
}
|
||||
self.run(move |connection| connection.query_pipeline(&pipeline))
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
fn increment(
|
||||
connection: &mut ConnectionRef<'_>,
|
||||
key: String,
|
||||
amount: f64,
|
||||
ttl: u64,
|
||||
refresh_ttl: bool,
|
||||
) -> Result<f64, Error> {
|
||||
if !refresh_ttl {
|
||||
return redis::cmd("EVAL")
|
||||
.arg(INCREMENT_SCRIPT)
|
||||
.arg(1)
|
||||
.arg(key)
|
||||
.arg(amount)
|
||||
.arg(ttl)
|
||||
.query(connection)
|
||||
.map_err(Error::unavailable);
|
||||
}
|
||||
connection
|
||||
.query_pipeline(
|
||||
redis::pipe()
|
||||
.cmd("INCRBYFLOAT")
|
||||
.arg(&key)
|
||||
.arg(amount)
|
||||
.cmd("EXPIRE")
|
||||
.arg(&key)
|
||||
.arg(ttl)
|
||||
.ignore(),
|
||||
)
|
||||
.map(|(value,)| value)
|
||||
}
|
||||
|
||||
impl<S, C> CountReadCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
fn batch_get_counts(&self, keys: &[String]) -> Result<Vec<Option<i64>>, Error> {
|
||||
let keys = self.namespaced_keys(keys);
|
||||
self.execute(|connection| mget(connection, keys))?
|
||||
.into_iter()
|
||||
.map(count)
|
||||
.collect()
|
||||
}
|
||||
|
||||
async fn async_batch_get_counts(&self, keys: Vec<String>) -> Result<Vec<Option<i64>>, Error> {
|
||||
let keys = self.namespaced_keys(&keys);
|
||||
self.run(move |connection| mget(connection, keys))
|
||||
.await?
|
||||
.into_iter()
|
||||
.map(count)
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
fn count(value: redis::Value) -> Result<Option<i64>, Error> {
|
||||
redis::from_redis_value(value).map_err(Error::invalid_entry)
|
||||
}
|
||||
|
||||
impl<S, C> BoundedCounterCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
fn increment_with_floor(&self, key: &str, amount: i64, ttl: Duration) -> Result<i64, Error> {
|
||||
let key = self.namespaced_key(key);
|
||||
let ttl = ttl_seconds(ttl);
|
||||
self.execute(|connection| increment_with_floor(connection, key, amount, ttl))
|
||||
}
|
||||
|
||||
async fn async_increment_with_floor(
|
||||
&self,
|
||||
key: &str,
|
||||
amount: i64,
|
||||
ttl: Duration,
|
||||
) -> Result<i64, Error> {
|
||||
let key = self.namespaced_key(key);
|
||||
let ttl = ttl_seconds(ttl);
|
||||
self.run(move |connection| increment_with_floor(connection, key, amount, ttl))
|
||||
.await
|
||||
}
|
||||
|
||||
async fn async_set_max(
|
||||
&self,
|
||||
key: &str,
|
||||
value: f64,
|
||||
ttl: Option<Duration>,
|
||||
) -> Result<f64, Error> {
|
||||
let key = self.namespaced_key(key);
|
||||
let ttl = self.ttl_or_default(ttl);
|
||||
self.run(move |connection| {
|
||||
redis::cmd("EVAL")
|
||||
.arg(SET_MAX_SCRIPT)
|
||||
.arg(1)
|
||||
.arg(key)
|
||||
.arg(value)
|
||||
.arg(ttl)
|
||||
.query(connection)
|
||||
.map_err(Error::unavailable)
|
||||
})
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
fn increment_with_floor(
|
||||
connection: &mut ConnectionRef<'_>,
|
||||
key: String,
|
||||
amount: i64,
|
||||
ttl: u64,
|
||||
) -> Result<i64, Error> {
|
||||
redis::cmd("EVAL")
|
||||
.arg(INCREMENT_WITH_FLOOR_SCRIPT)
|
||||
.arg(1)
|
||||
.arg(key)
|
||||
.arg(amount)
|
||||
.arg(ttl)
|
||||
.query(connection)
|
||||
.map_err(Error::unavailable)
|
||||
}
|
||||
|
|
@ -1,63 +0,0 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use litellm_cache::{CacheCodec, Error, RefreshTtlCache, ScanCache, TtlCache};
|
||||
|
||||
use crate::cache::RedisCache;
|
||||
|
||||
impl<S, C> TtlCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
async fn async_get_ttl(&self, key: &str) -> Result<Option<Duration>, Error> {
|
||||
let key = self.namespaced_key(key);
|
||||
let ttl = self
|
||||
.run(move |connection| {
|
||||
redis::cmd("TTL")
|
||||
.arg(key)
|
||||
.query::<i64>(connection)
|
||||
.map_err(Error::unavailable)
|
||||
})
|
||||
.await?;
|
||||
Ok(u64::try_from(ttl).ok().map(Duration::from_secs))
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, C> RefreshTtlCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
async fn async_refresh_ttl(&self, key: &str, ttl: Option<Duration>) -> Result<bool, Error> {
|
||||
let key = self.namespaced_key(key);
|
||||
let ttl = self.ttl_or_default(ttl);
|
||||
self.run(move |connection| {
|
||||
redis::cmd("EXPIRE")
|
||||
.arg(key)
|
||||
.arg(ttl)
|
||||
.query(connection)
|
||||
.map_err(Error::unavailable)
|
||||
})
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, C> ScanCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
async fn async_scan_iter(&self, pattern: &str, count: usize) -> Result<Vec<String>, Error> {
|
||||
let pattern = format!("{}*", self.namespaced_key(pattern));
|
||||
self.run(move |connection| {
|
||||
let mut matches = Vec::new();
|
||||
connection.scan(&pattern, count, |_, keys| {
|
||||
matches.extend(keys);
|
||||
Ok(matches.len() < count)
|
||||
})?;
|
||||
matches.truncate(count);
|
||||
Ok(matches)
|
||||
})
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
|
@ -1,15 +0,0 @@
|
|||
mod cache;
|
||||
mod claim;
|
||||
pub mod connection;
|
||||
mod counter;
|
||||
mod keys;
|
||||
mod lifecycle;
|
||||
mod queue;
|
||||
mod script;
|
||||
mod store;
|
||||
mod topology;
|
||||
|
||||
pub use cache::RedisCache;
|
||||
pub use queue::{RedisLpopOperation, RedisLpopResult, RedisRpushOperation};
|
||||
pub use script::{RedisArg, RedisScript};
|
||||
pub use topology::{RedisNode, RedisTopology};
|
||||
|
|
@ -1,86 +0,0 @@
|
|||
use litellm_cache::{
|
||||
CacheCodec, CacheConnectionResult, CacheConnectionStatus, ClientInfoCache, ConnectionCache,
|
||||
DisconnectCache, Error, PingCache,
|
||||
};
|
||||
|
||||
use crate::{cache::RedisCache, topology::RedisTopology};
|
||||
|
||||
impl<S, C> PingCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
fn sync_ping(&self) -> Result<bool, Error> {
|
||||
self.execute(|connection| connection.ping().map_err(Error::unavailable))
|
||||
}
|
||||
|
||||
async fn ping(&self) -> Result<bool, Error> {
|
||||
self.run(|connection| connection.ping().map_err(Error::unavailable))
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, C> ConnectionCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
/// Python `RedisCache.test_connection`, or `RedisClusterCache.test_connection` for a
|
||||
/// cluster topology, which differs only in its messages.
|
||||
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
|
||||
let label = match self.topology {
|
||||
RedisTopology::Standalone => "Redis",
|
||||
RedisTopology::Cluster { .. } => "Redis Cluster",
|
||||
};
|
||||
let ping = self
|
||||
.run(|connection| Ok(connection.ping().map_err(Error::unavailable)))
|
||||
.await
|
||||
.and_then(|ping| ping)
|
||||
.map_err(|error| litellm_error::Report::new(&error).to_string());
|
||||
Ok(match ping {
|
||||
Ok(true) => CacheConnectionResult {
|
||||
status: CacheConnectionStatus::Success,
|
||||
message: format!("{label} connection test successful"),
|
||||
error: None,
|
||||
},
|
||||
Ok(false) => CacheConnectionResult {
|
||||
status: CacheConnectionStatus::Failed,
|
||||
message: format!("{label} ping returned False"),
|
||||
error: None,
|
||||
},
|
||||
Err(error) => CacheConnectionResult {
|
||||
status: CacheConnectionStatus::Failed,
|
||||
message: format!("{label} connection failed: {error}"),
|
||||
error: Some(error),
|
||||
},
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, C> DisconnectCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
async fn disconnect(&self) -> Result<(), Error> {
|
||||
self.connections.disconnect();
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, C> ClientInfoCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
type ClientList = String;
|
||||
type Info = String;
|
||||
|
||||
fn client_list(&self) -> Result<Self::ClientList, Error> {
|
||||
self.execute(|connection| connection.node_text(redis::cmd("CLIENT").arg("LIST")))
|
||||
}
|
||||
|
||||
fn info(&self) -> Result<Self::Info, Error> {
|
||||
self.execute(|connection| connection.node_text(&redis::cmd("INFO")))
|
||||
}
|
||||
}
|
||||
|
|
@ -1,226 +0,0 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use litellm_cache::{CacheCodec, Error, PopOperation, PushOperation, QueueCache, SetCache};
|
||||
|
||||
use crate::{cache::RedisCache, script::RedisArg};
|
||||
|
||||
pub type RedisRpushOperation = PushOperation<RedisArg>;
|
||||
pub type RedisLpopOperation = PopOperation;
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub enum RedisLpopResult {
|
||||
Missing,
|
||||
Value(Vec<u8>),
|
||||
Values(Vec<Vec<u8>>),
|
||||
}
|
||||
|
||||
impl<S, C> SetCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
type SetValue = RedisArg;
|
||||
type SetResult = usize;
|
||||
|
||||
async fn async_set_cache_sadd(
|
||||
&self,
|
||||
key: &str,
|
||||
values: Vec<Self::SetValue>,
|
||||
ttl: Option<Duration>,
|
||||
) -> Result<Self::SetResult, Error> {
|
||||
if values.is_empty() {
|
||||
return Err(Error::InvalidEntry(None));
|
||||
}
|
||||
let key = self.namespaced_key(key);
|
||||
let ttl = self.ttl_or_default(ttl);
|
||||
let mut pipeline = redis::pipe();
|
||||
pipeline
|
||||
.cmd("SADD")
|
||||
.arg(&key)
|
||||
.arg(values)
|
||||
.cmd("EXPIRE")
|
||||
.arg(&key)
|
||||
.arg(ttl)
|
||||
.ignore();
|
||||
self.run(move |connection| connection.query_pipeline(&pipeline))
|
||||
.await
|
||||
.map(|(added,)| added)
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, C> QueueCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
type QueueValue = RedisArg;
|
||||
type PopResult = RedisLpopResult;
|
||||
|
||||
async fn async_rpush(&self, key: &str, values: Vec<Self::QueueValue>) -> Result<usize, Error> {
|
||||
if values.is_empty() {
|
||||
return Err(Error::InvalidEntry(None));
|
||||
}
|
||||
let key = self.namespaced_key(key);
|
||||
self.run(move |connection| {
|
||||
redis::cmd("RPUSH")
|
||||
.arg(key)
|
||||
.arg(values)
|
||||
.query(connection)
|
||||
.map_err(Error::unavailable)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
async fn async_rpush_and_trim(
|
||||
&self,
|
||||
key: &str,
|
||||
values: Vec<Self::QueueValue>,
|
||||
max_len: usize,
|
||||
) -> Result<usize, Error> {
|
||||
if values.is_empty() {
|
||||
return Err(Error::InvalidEntry(None));
|
||||
}
|
||||
let key = self.namespaced_key(key);
|
||||
let start = i64::try_from(max_len).map_or(i64::MIN, |max_len| -max_len);
|
||||
let mut pipeline = redis::pipe();
|
||||
pipeline
|
||||
.atomic()
|
||||
.cmd("RPUSH")
|
||||
.arg(&key)
|
||||
.arg(values)
|
||||
.cmd("LTRIM")
|
||||
.arg(&key)
|
||||
.arg(start)
|
||||
.arg(-1)
|
||||
.ignore();
|
||||
self.run(move |connection| connection.query_pipeline(&pipeline))
|
||||
.await
|
||||
.map(|(length,)| length)
|
||||
}
|
||||
|
||||
async fn async_rpush_pipeline(
|
||||
&self,
|
||||
operations: Vec<RedisRpushOperation>,
|
||||
) -> Result<Vec<usize>, Error> {
|
||||
if operations.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let mut pipeline = redis::pipe();
|
||||
for operation in operations {
|
||||
if operation.values.is_empty() {
|
||||
return Err(Error::InvalidEntry(None));
|
||||
}
|
||||
pipeline
|
||||
.cmd("RPUSH")
|
||||
.arg(self.namespaced_key(&operation.key))
|
||||
.arg(operation.values);
|
||||
}
|
||||
self.run(move |connection| connection.query_pipeline(&pipeline))
|
||||
.await
|
||||
}
|
||||
|
||||
async fn async_lpop(&self, key: &str, count: Option<usize>) -> Result<Self::PopResult, Error> {
|
||||
if let Some(count) = count
|
||||
&& self.major_version().await < 7
|
||||
{
|
||||
return self.lpop_one_at_a_time(key, count).await;
|
||||
}
|
||||
let command = lpop(self.namespaced_key(key), count);
|
||||
let value = self
|
||||
.run(move |connection| {
|
||||
command
|
||||
.query::<redis::Value>(connection)
|
||||
.map_err(Error::unavailable)
|
||||
})
|
||||
.await?;
|
||||
lpop_result(value, count.is_some())
|
||||
}
|
||||
|
||||
async fn async_lpop_pipeline(
|
||||
&self,
|
||||
operations: Vec<RedisLpopOperation>,
|
||||
) -> Result<Vec<Self::PopResult>, Error> {
|
||||
if operations.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
if operations.iter().any(|operation| operation.count.is_some())
|
||||
&& self.major_version().await < 7
|
||||
{
|
||||
let mut results = Vec::with_capacity(operations.len());
|
||||
for operation in &operations {
|
||||
results.push(self.async_lpop(&operation.key, operation.count).await?);
|
||||
}
|
||||
return Ok(results);
|
||||
}
|
||||
let multiple = operations
|
||||
.iter()
|
||||
.map(|operation| operation.count.is_some())
|
||||
.collect::<Vec<_>>();
|
||||
let mut pipeline = redis::pipe();
|
||||
for operation in operations {
|
||||
pipeline.add_command(lpop(self.namespaced_key(&operation.key), operation.count));
|
||||
}
|
||||
self.run(move |connection| connection.query_pipeline::<Vec<redis::Value>>(&pipeline))
|
||||
.await?
|
||||
.into_iter()
|
||||
.zip(multiple)
|
||||
.map(|(value, multiple)| lpop_result(value, multiple))
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
fn lpop(key: String, count: Option<usize>) -> redis::Cmd {
|
||||
let mut command = redis::cmd("LPOP");
|
||||
command.arg(key);
|
||||
if let Some(count) = count {
|
||||
command.arg(count);
|
||||
}
|
||||
command
|
||||
}
|
||||
|
||||
fn redis_bytes(value: redis::Value) -> Result<Vec<u8>, Error> {
|
||||
match value {
|
||||
redis::Value::BulkString(bytes) => Ok(bytes),
|
||||
redis::Value::SimpleString(text) => Ok(text.into_bytes()),
|
||||
_ => Err(Error::InvalidEntry(None)),
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, C> RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
/// `handle_lpop_count_for_older_redis_versions`: `count` single-`LPOP` pipelines, keeping
|
||||
/// only the values actually popped.
|
||||
async fn lpop_one_at_a_time(&self, key: &str, count: usize) -> Result<RedisLpopResult, Error> {
|
||||
let key = self.namespaced_key(key);
|
||||
let mut values = Vec::new();
|
||||
for _ in 0..count {
|
||||
let mut pipeline = redis::pipe();
|
||||
pipeline.add_command(lpop(key.clone(), None));
|
||||
let replies = self
|
||||
.run(move |connection| connection.query_pipeline::<Vec<redis::Value>>(&pipeline))
|
||||
.await?;
|
||||
for reply in replies {
|
||||
if reply != redis::Value::Nil {
|
||||
values.push(redis_bytes(reply)?);
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(RedisLpopResult::Values(values))
|
||||
}
|
||||
}
|
||||
|
||||
fn lpop_result(value: redis::Value, multiple: bool) -> Result<RedisLpopResult, Error> {
|
||||
match value {
|
||||
redis::Value::Nil => Ok(RedisLpopResult::Missing),
|
||||
redis::Value::Array(values) if multiple => values
|
||||
.into_iter()
|
||||
.map(redis_bytes)
|
||||
.collect::<Result<Vec<_>, _>>()
|
||||
.map(RedisLpopResult::Values),
|
||||
value if !multiple => redis_bytes(value).map(RedisLpopResult::Value),
|
||||
_ => Err(Error::InvalidEntry(None)),
|
||||
}
|
||||
}
|
||||
|
|
@ -1,136 +0,0 @@
|
|||
use std::sync::Arc;
|
||||
|
||||
use litellm_cache::{CacheCodec, CacheScript, Error, ScriptCache};
|
||||
|
||||
use crate::{
|
||||
cache::{RedisCache, namespaced_key},
|
||||
connection::Connections,
|
||||
};
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub enum RedisArg {
|
||||
Bytes(Vec<u8>),
|
||||
Integer(i64),
|
||||
Float(f64),
|
||||
}
|
||||
|
||||
impl From<&str> for RedisArg {
|
||||
fn from(value: &str) -> Self {
|
||||
Self::Bytes(value.as_bytes().to_vec())
|
||||
}
|
||||
}
|
||||
|
||||
impl From<String> for RedisArg {
|
||||
fn from(value: String) -> Self {
|
||||
Self::Bytes(value.into_bytes())
|
||||
}
|
||||
}
|
||||
|
||||
impl From<Vec<u8>> for RedisArg {
|
||||
fn from(value: Vec<u8>) -> Self {
|
||||
Self::Bytes(value)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<i64> for RedisArg {
|
||||
fn from(value: i64) -> Self {
|
||||
Self::Integer(value)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<f64> for RedisArg {
|
||||
fn from(value: f64) -> Self {
|
||||
Self::Float(value)
|
||||
}
|
||||
}
|
||||
|
||||
impl redis::ToRedisArgs for RedisArg {
|
||||
fn write_redis_args<W>(&self, out: &mut W)
|
||||
where
|
||||
W: ?Sized + redis::RedisWrite,
|
||||
{
|
||||
match self {
|
||||
Self::Bytes(value) => value.write_redis_args(out),
|
||||
Self::Integer(value) => value.write_redis_args(out),
|
||||
Self::Float(value) => value.write_redis_args(out),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct RedisScript<C> {
|
||||
connections: Arc<Connections<C>>,
|
||||
namespace: Option<String>,
|
||||
source: String,
|
||||
}
|
||||
|
||||
impl<C> CacheScript for RedisScript<C>
|
||||
where
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
type Argument = RedisArg;
|
||||
type Output = redis::Value;
|
||||
|
||||
async fn invoke(
|
||||
&self,
|
||||
keys: Vec<String>,
|
||||
arguments: Vec<Self::Argument>,
|
||||
) -> Result<Self::Output, Error> {
|
||||
let keys = keys
|
||||
.into_iter()
|
||||
.map(|key| namespaced_key(self.namespace.as_deref(), &key))
|
||||
.collect::<Vec<_>>();
|
||||
let source = self.source.clone();
|
||||
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
|
||||
eval(connection, &source, keys, arguments)
|
||||
})
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, C> RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
pub async fn async_eval(
|
||||
&self,
|
||||
script: String,
|
||||
keys: Vec<String>,
|
||||
arguments: Vec<RedisArg>,
|
||||
) -> Result<redis::Value, Error> {
|
||||
let keys = self.namespaced_keys(&keys);
|
||||
self.run(move |connection| eval(connection, &script, keys, arguments))
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
fn eval(
|
||||
connection: &mut impl redis::ConnectionLike,
|
||||
script: &str,
|
||||
keys: Vec<String>,
|
||||
arguments: Vec<RedisArg>,
|
||||
) -> Result<redis::Value, Error> {
|
||||
redis::cmd("EVAL")
|
||||
.arg(script)
|
||||
.arg(keys.len())
|
||||
.arg(keys)
|
||||
.arg(arguments)
|
||||
.query(connection)
|
||||
.map_err(Error::unavailable)
|
||||
}
|
||||
|
||||
impl<S, C> ScriptCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
type Script = RedisScript<C>;
|
||||
|
||||
fn async_register_script(&self, source: String) -> Self::Script {
|
||||
RedisScript {
|
||||
connections: Arc::clone(&self.connections),
|
||||
namespace: self.namespace.clone(),
|
||||
source,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -1,232 +0,0 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use litellm_cache::{
|
||||
BaseCache, BatchCache, BatchEntry, BulkDeleteCache, CacheCodec, DeleteCache, Error,
|
||||
ExactCacheContext, FlushAllCache, FlushCache, TtlPipelineCache,
|
||||
};
|
||||
use redis::Commands;
|
||||
|
||||
use crate::{cache::RedisCache, connection::ConnectionRef};
|
||||
|
||||
impl<S, C> BaseCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
type Value = S::Value;
|
||||
type Context = ExactCacheContext;
|
||||
|
||||
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
|
||||
context.ttl.or(Some(self.default_ttl))
|
||||
}
|
||||
|
||||
fn set_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
value: Self::Value,
|
||||
context: &ExactCacheContext,
|
||||
) -> Result<(), Error> {
|
||||
let payload = self.codec.encode(&value)?;
|
||||
let ttl = self.ttl_or_default(context.ttl);
|
||||
let key = self.namespaced_key(key);
|
||||
self.execute(|connection| {
|
||||
connection
|
||||
.set_ex::<_, _, ()>(key, payload, ttl)
|
||||
.map_err(Error::unavailable)
|
||||
})
|
||||
}
|
||||
|
||||
fn get_cache(&self, key: &str, _: &ExactCacheContext) -> Result<Option<Self::Value>, Error> {
|
||||
let key = self.namespaced_key(key);
|
||||
let value = self.execute(|connection| {
|
||||
connection
|
||||
.get::<_, redis::Value>(key)
|
||||
.map_err(Error::unavailable)
|
||||
})?;
|
||||
self.decode_response(value)
|
||||
}
|
||||
|
||||
async fn async_set_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
value: Self::Value,
|
||||
context: ExactCacheContext,
|
||||
) -> Result<(), Error> {
|
||||
let payload = self.codec.encode(&value)?;
|
||||
let key = self.namespaced_key(key);
|
||||
let ttl = self.ttl_or_default(context.ttl);
|
||||
self.run(move |connection| {
|
||||
connection
|
||||
.set_ex::<_, _, ()>(key, payload, ttl)
|
||||
.map_err(Error::unavailable)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
async fn async_get_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
_: &ExactCacheContext,
|
||||
) -> Result<Option<Self::Value>, Error> {
|
||||
let key = self.namespaced_key(key);
|
||||
let value = self
|
||||
.run(move |connection| {
|
||||
connection
|
||||
.get::<_, redis::Value>(key)
|
||||
.map_err(Error::unavailable)
|
||||
})
|
||||
.await?;
|
||||
self.decode_response(value)
|
||||
}
|
||||
|
||||
async fn async_set_cache_pipeline(
|
||||
&self,
|
||||
cache_list: Vec<(String, Self::Value)>,
|
||||
context: ExactCacheContext,
|
||||
) -> Result<(), Error> {
|
||||
self.async_set_cache_pipeline_with_ttls(
|
||||
cache_list
|
||||
.into_iter()
|
||||
.map(|(key, value)| (key, value, context.ttl))
|
||||
.collect(),
|
||||
)
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, C> TtlPipelineCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
async fn async_set_cache_pipeline_with_ttls(
|
||||
&self,
|
||||
entries: Vec<(String, Self::Value, Option<Duration>)>,
|
||||
) -> Result<(), Error> {
|
||||
if entries.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
let mut pipeline = redis::pipe();
|
||||
for (key, value, ttl) in entries {
|
||||
pipeline
|
||||
.cmd("SETEX")
|
||||
.arg(self.namespaced_key(&key))
|
||||
.arg(self.ttl_or_default(ttl))
|
||||
.arg(self.codec.encode(&value)?)
|
||||
.ignore();
|
||||
}
|
||||
self.run(move |connection| connection.query_pipeline(&pipeline))
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, C> BatchCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
fn batch_get_cache(
|
||||
&self,
|
||||
keys: &[String],
|
||||
_: &ExactCacheContext,
|
||||
) -> Result<Vec<BatchEntry<Self::Value>>, Error> {
|
||||
let keys = self.namespaced_keys(keys);
|
||||
self.execute(|connection| mget(connection, keys))?
|
||||
.into_iter()
|
||||
.map(|value| self.decode_batch_response(value))
|
||||
.collect()
|
||||
}
|
||||
|
||||
async fn async_batch_get_cache(
|
||||
&self,
|
||||
keys: Vec<String>,
|
||||
_: ExactCacheContext,
|
||||
) -> Result<Vec<BatchEntry<Self::Value>>, Error> {
|
||||
let keys = self.namespaced_keys(&keys);
|
||||
self.run(move |connection| mget(connection, keys))
|
||||
.await?
|
||||
.into_iter()
|
||||
.map(|value| self.decode_batch_response(value))
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn mget(
|
||||
connection: &mut ConnectionRef<'_>,
|
||||
keys: Vec<String>,
|
||||
) -> Result<Vec<redis::Value>, Error> {
|
||||
redis::cmd("MGET")
|
||||
.arg(keys)
|
||||
.query(connection)
|
||||
.map_err(Error::unavailable)
|
||||
}
|
||||
|
||||
impl<S, C> DeleteCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
fn delete_cache(&self, key: &str) -> Result<(), Error> {
|
||||
let key = self.namespaced_key(key);
|
||||
self.execute(|connection| connection.del::<_, ()>(key).map_err(Error::unavailable))
|
||||
}
|
||||
|
||||
async fn async_delete_cache(&self, key: &str) -> Result<(), Error> {
|
||||
let key = self.namespaced_key(key);
|
||||
self.run(move |connection| connection.del::<_, ()>(key).map_err(Error::unavailable))
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, C> BulkDeleteCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
async fn delete_cache_keys(&self, keys: Vec<String>) -> Result<usize, Error> {
|
||||
if keys.is_empty() {
|
||||
return Ok(0);
|
||||
}
|
||||
let keys = self.namespaced_keys(&keys);
|
||||
self.run(move |connection| connection.del(keys).map_err(Error::unavailable))
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, C> FlushCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
fn flush_cache(&self) -> Result<(), Error> {
|
||||
let pattern = self.namespaced_pattern()?;
|
||||
self.execute(|connection| flush_matching(connection, &pattern))
|
||||
}
|
||||
|
||||
async fn async_flush_cache(&self) -> Result<(), Error> {
|
||||
let pattern = self.namespaced_pattern()?;
|
||||
self.run(move |connection| flush_matching(connection, &pattern))
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, C> FlushAllCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
fn flushall(&self) -> Result<(), Error> {
|
||||
self.execute(|connection| connection.flushall())
|
||||
}
|
||||
}
|
||||
|
||||
fn flush_matching(connection: &mut ConnectionRef<'_>, pattern: &str) -> Result<(), Error> {
|
||||
connection.scan(pattern, 1000, |connection, keys| {
|
||||
if !keys.is_empty() {
|
||||
connection
|
||||
.del::<_, usize>(keys)
|
||||
.map_err(Error::unavailable)?;
|
||||
}
|
||||
Ok(true)
|
||||
})
|
||||
}
|
||||
|
|
@ -1,14 +0,0 @@
|
|||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct RedisNode {
|
||||
pub host: String,
|
||||
pub port: u16,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default, PartialEq, Eq)]
|
||||
pub enum RedisTopology {
|
||||
#[default]
|
||||
Standalone,
|
||||
Cluster {
|
||||
startup_nodes: Vec<RedisNode>,
|
||||
},
|
||||
}
|
||||
File diff suppressed because it is too large
Load diff
|
|
@ -1,563 +0,0 @@
|
|||
//! Tests against a real Redis Cluster. Set `LITELLM_TEST_REDIS_CLUSTER_NODES` to a comma
|
||||
//! separated `host:port` list (for example `127.0.0.1:7000,127.0.0.1:7001`) to run them.
|
||||
|
||||
mod support;
|
||||
|
||||
use std::{collections::HashSet, time::Duration};
|
||||
|
||||
use litellm_cache::{
|
||||
BaseCache, BatchCache, BatchEntry, BoundedCounterCache, BulkDeleteCache, CacheConnectionStatus,
|
||||
CacheScript, ClaimCache, ClientInfoCache, ConnectionCache, CounterCache, DeleteCache,
|
||||
DisconnectCache, Error, ExactCacheContext, FlushCache, IncrementOperation, JsonCodec,
|
||||
PingCache, QueueCache, RefreshTtlCache, ScanCache, ScriptCache, SetCache, TtlCache,
|
||||
TtlPipelineCache,
|
||||
};
|
||||
use litellm_cache_redis::{
|
||||
RedisArg, RedisCache, RedisLpopOperation, RedisLpopResult, RedisNode, RedisRpushOperation,
|
||||
RedisTopology,
|
||||
};
|
||||
use redis::cluster_routing::Slot;
|
||||
use rstest::{fixture, rstest};
|
||||
use serde_json::json;
|
||||
use support::{JsonCache, cluster_cache, cluster_url};
|
||||
|
||||
type Counter = RedisCache<JsonCodec<f64>>;
|
||||
|
||||
#[fixture]
|
||||
fn cache(#[default("cache")] label: &str) -> Option<JsonCache> {
|
||||
cluster_cache(label, Duration::from_secs(120), JsonCodec::new())
|
||||
}
|
||||
|
||||
#[fixture]
|
||||
fn counter(#[default("counter")] label: &str) -> Option<Counter> {
|
||||
cluster_cache(label, Duration::from_secs(60), JsonCodec::new())
|
||||
}
|
||||
|
||||
#[fixture]
|
||||
fn context() -> ExactCacheContext {
|
||||
ExactCacheContext::default()
|
||||
}
|
||||
|
||||
fn multi_slot_keys(count: usize) -> Vec<String> {
|
||||
let keys: Vec<String> = (0..count).map(|index| format!("key-{index}")).collect();
|
||||
let slots: HashSet<Slot> = keys.iter().map(Slot::for_key).collect();
|
||||
assert!(slots.len() > 1, "keys must span multiple slots");
|
||||
keys
|
||||
}
|
||||
|
||||
fn seconds(seconds: u64) -> Option<Duration> {
|
||||
Some(Duration::from_secs(seconds))
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::no_startup_nodes("redis://127.0.0.1:7000", Vec::new())]
|
||||
#[case::unix_socket_url(
|
||||
"redis+unix:///tmp/redis.sock",
|
||||
vec![RedisNode { host: "127.0.0.1".into(), port: 7000 }]
|
||||
)]
|
||||
fn constructor_rejects_unusable_cluster_configs(
|
||||
#[case] url: &str,
|
||||
#[case] startup_nodes: Vec<RedisNode>,
|
||||
) {
|
||||
let error = JsonCache::connect(
|
||||
url,
|
||||
&RedisTopology::Cluster { startup_nodes },
|
||||
None,
|
||||
JsonCodec::new(),
|
||||
)
|
||||
.err();
|
||||
assert!(matches!(error, Some(Error::Unavailable(_))));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn single_key_operations_round_trip_with_ttl_rounding(
|
||||
#[with("single")] cache: Option<JsonCache>,
|
||||
) {
|
||||
let Some(cache) = cache else { return };
|
||||
let context = ExactCacheContext {
|
||||
ttl: Some(Duration::from_millis(1500)),
|
||||
};
|
||||
let keys = multi_slot_keys(12);
|
||||
for (index, key) in keys.iter().enumerate() {
|
||||
cache
|
||||
.set_cache(key, json!({ "index": index }), &context)
|
||||
.unwrap();
|
||||
}
|
||||
for (index, key) in keys.iter().enumerate() {
|
||||
assert_eq!(
|
||||
cache.get_cache(key, &context).unwrap(),
|
||||
Some(json!({ "index": index }))
|
||||
);
|
||||
}
|
||||
assert_eq!(cache.async_get_ttl(&keys[0]).await.unwrap(), seconds(2));
|
||||
assert!(
|
||||
cache
|
||||
.async_refresh_ttl(&keys[0], seconds(40))
|
||||
.await
|
||||
.unwrap()
|
||||
);
|
||||
assert_eq!(cache.async_get_ttl(&keys[0]).await.unwrap(), seconds(40));
|
||||
cache.delete_cache(&keys[0]).unwrap();
|
||||
assert_eq!(cache.get_cache(&keys[0], &context).unwrap(), None);
|
||||
assert!(!cache.async_refresh_ttl(&keys[0], None).await.unwrap());
|
||||
assert!(cache.sync_ping().unwrap());
|
||||
cache.async_flush_cache().await.unwrap();
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn batch_reads_span_slots_and_preserve_order_with_malformed_entries(
|
||||
#[with("batch")] cache: Option<JsonCache>,
|
||||
context: ExactCacheContext,
|
||||
) {
|
||||
let Some(cache) = cache else { return };
|
||||
let keys = multi_slot_keys(40);
|
||||
for (index, key) in keys.iter().enumerate() {
|
||||
if index % 5 == 0 {
|
||||
continue;
|
||||
}
|
||||
cache
|
||||
.async_set_cache(key, json!(index), context.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
let mut raw = redis::cluster::ClusterClient::new(vec![cluster_url()])
|
||||
.unwrap()
|
||||
.get_connection()
|
||||
.unwrap();
|
||||
let malformed = format!("{}:{}", cache.namespace().unwrap(), keys[1]);
|
||||
redis::cmd("SET")
|
||||
.arg(&malformed)
|
||||
.arg("not json")
|
||||
.exec(&mut raw)
|
||||
.unwrap();
|
||||
|
||||
let entries = cache
|
||||
.async_batch_get_cache(keys.clone(), context.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(entries.len(), keys.len());
|
||||
for (index, entry) in entries.iter().enumerate() {
|
||||
let expected = if index == 1 {
|
||||
BatchEntry::Invalid
|
||||
} else if index % 5 == 0 {
|
||||
BatchEntry::Miss
|
||||
} else {
|
||||
BatchEntry::Hit(json!(index))
|
||||
};
|
||||
assert_eq!(*entry, expected, "entry {index}");
|
||||
}
|
||||
assert_eq!(cache.batch_get_cache(&keys, &context).unwrap(), entries);
|
||||
|
||||
cache.delete_cache_keys(keys.clone()).await.unwrap();
|
||||
let entries = cache.async_batch_get_cache(keys, context).await.unwrap();
|
||||
assert!(entries.iter().all(|entry| *entry == BatchEntry::Miss));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn pipelines_group_by_slot_and_return_results_in_submission_order(
|
||||
#[with("pipeline")] cache: Option<JsonCache>,
|
||||
counter: Option<Counter>,
|
||||
context: ExactCacheContext,
|
||||
) {
|
||||
let (Some(cache), Some(counter)) = (cache, counter) else {
|
||||
return;
|
||||
};
|
||||
let keys = multi_slot_keys(30);
|
||||
cache
|
||||
.async_set_cache_pipeline(
|
||||
keys.iter()
|
||||
.enumerate()
|
||||
.map(|(index, key)| (key.clone(), json!(index)))
|
||||
.collect(),
|
||||
context.clone(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let hits = cache
|
||||
.async_batch_get_cache(keys.clone(), context.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(
|
||||
hits.iter()
|
||||
.enumerate()
|
||||
.all(|(index, entry)| *entry == BatchEntry::Hit(json!(index)))
|
||||
);
|
||||
|
||||
cache
|
||||
.async_set_cache_pipeline_with_ttls(
|
||||
keys.iter()
|
||||
.enumerate()
|
||||
.map(|(index, key)| (key.clone(), json!(index), seconds(index as u64 + 10)))
|
||||
.collect(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
for (index, key) in keys.iter().enumerate() {
|
||||
assert_eq!(
|
||||
cache.async_get_ttl(key).await.unwrap(),
|
||||
seconds(index as u64 + 10),
|
||||
"{key}"
|
||||
);
|
||||
}
|
||||
|
||||
let queues: Vec<String> = keys.iter().map(|key| format!("queue:{key}")).collect();
|
||||
let pushed = cache
|
||||
.async_rpush_pipeline(
|
||||
queues
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(index, key)| RedisRpushOperation {
|
||||
key: key.clone(),
|
||||
values: (0..=index)
|
||||
.map(|value| RedisArg::Integer(value as i64))
|
||||
.collect(),
|
||||
})
|
||||
.collect(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(pushed, (1..=keys.len()).collect::<Vec<_>>());
|
||||
let popped = cache
|
||||
.async_lpop_pipeline(
|
||||
queues
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(index, key)| RedisLpopOperation {
|
||||
key: key.clone(),
|
||||
count: (index % 2 == 0).then_some(2),
|
||||
})
|
||||
.collect(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
for (index, result) in popped.into_iter().enumerate() {
|
||||
match result {
|
||||
RedisLpopResult::Value(value) => {
|
||||
assert_eq!(index % 2, 1, "queue {index}");
|
||||
assert_eq!(value, b"0");
|
||||
}
|
||||
RedisLpopResult::Values(values) => {
|
||||
assert_eq!(index % 2, 0, "queue {index}");
|
||||
let expected: Vec<Vec<u8>> = (0..=index)
|
||||
.take(2)
|
||||
.map(|value| value.to_string().into_bytes())
|
||||
.collect();
|
||||
assert_eq!(values, expected);
|
||||
}
|
||||
other => panic!("queue {index}: {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
let counters: Vec<String> = keys.iter().map(|key| format!("counter:{key}")).collect();
|
||||
let totals = counter
|
||||
.async_increment_pipeline(
|
||||
counters
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(index, key)| IncrementOperation {
|
||||
key: key.clone(),
|
||||
amount: index as f64 + 0.5,
|
||||
ttl: (index % 3 == 0).then_some(Duration::from_secs(30)),
|
||||
})
|
||||
.collect(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let expected: Vec<f64> = (0..keys.len()).map(|index| index as f64 + 0.5).collect();
|
||||
assert_eq!(totals, expected);
|
||||
assert_eq!(
|
||||
counter.async_get_ttl(&counters[0]).await.unwrap(),
|
||||
seconds(30)
|
||||
);
|
||||
assert_eq!(counter.async_get_ttl(&counters[1]).await.unwrap(), None);
|
||||
counter.async_flush_cache().await.unwrap();
|
||||
cache.async_flush_cache().await.unwrap();
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn rpush_and_trim_is_one_transaction_on_the_key_slot(
|
||||
#[with("trim")] cache: Option<JsonCache>,
|
||||
) {
|
||||
let Some(cache) = cache else { return };
|
||||
let values = |values: &[&str]| values.iter().map(|value| RedisArg::from(*value)).collect();
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_rpush_and_trim("buf", values(&["a", "b"]), 3)
|
||||
.await
|
||||
.unwrap(),
|
||||
2
|
||||
);
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_rpush_and_trim("buf", values(&["c", "d"]), 3)
|
||||
.await
|
||||
.unwrap(),
|
||||
4
|
||||
);
|
||||
assert_eq!(
|
||||
cache.async_lpop("buf", Some(10)).await.unwrap(),
|
||||
RedisLpopResult::Values(vec![b"b".to_vec(), b"c".to_vec(), b"d".to_vec()])
|
||||
);
|
||||
cache.async_flush_cache().await.unwrap();
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn scan_and_scoped_flush_cover_every_primary(
|
||||
#[with("flush")] cache: Option<JsonCache>,
|
||||
#[from(cache)]
|
||||
#[with("other")]
|
||||
other: Option<JsonCache>,
|
||||
context: ExactCacheContext,
|
||||
) {
|
||||
let (Some(cache), Some(other)) = (cache, other) else {
|
||||
return;
|
||||
};
|
||||
let keys = multi_slot_keys(60);
|
||||
for key in &keys {
|
||||
cache
|
||||
.async_set_cache(key, json!(true), context.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
other
|
||||
.async_set_cache(key, json!(true), context.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
let mut scanned = cache.async_scan_iter("key-", 1000).await.unwrap();
|
||||
scanned.sort();
|
||||
let mut expected: Vec<String> = keys
|
||||
.iter()
|
||||
.map(|key| format!("{}:{key}", cache.namespace().unwrap()))
|
||||
.collect();
|
||||
expected.sort();
|
||||
assert_eq!(scanned, expected);
|
||||
assert_eq!(cache.async_scan_iter("key-", 7).await.unwrap().len(), 7);
|
||||
|
||||
cache.flush_cache().unwrap();
|
||||
let flushed = cache
|
||||
.async_batch_get_cache(keys.clone(), context.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(flushed.iter().all(|entry| *entry == BatchEntry::Miss));
|
||||
let kept = other.async_batch_get_cache(keys, context).await.unwrap();
|
||||
assert!(
|
||||
kept.iter()
|
||||
.all(|entry| *entry == BatchEntry::Hit(json!(true)))
|
||||
);
|
||||
other.async_flush_cache().await.unwrap();
|
||||
}
|
||||
|
||||
fn ping_calls_per_node(startup: &redis::Client) -> Vec<(String, u64)> {
|
||||
let mut connection = startup.get_connection().unwrap();
|
||||
let nodes: String = redis::cmd("CLUSTER")
|
||||
.arg("NODES")
|
||||
.query(&mut connection)
|
||||
.unwrap();
|
||||
let mut counts: Vec<(String, u64)> = nodes
|
||||
.lines()
|
||||
.map(|line| {
|
||||
let address = line.split_whitespace().nth(1).unwrap();
|
||||
let address = address.split('@').next().unwrap();
|
||||
let mut node = redis::Client::open(format!("redis://{address}"))
|
||||
.unwrap()
|
||||
.get_connection()
|
||||
.unwrap();
|
||||
let stats: String = redis::cmd("INFO")
|
||||
.arg("commandstats")
|
||||
.query(&mut node)
|
||||
.unwrap();
|
||||
let calls = stats
|
||||
.lines()
|
||||
.find_map(|stat| stat.strip_prefix("cmdstat_ping:calls="))
|
||||
.and_then(|rest| rest.split(',').next())
|
||||
.map_or(0, |calls| calls.parse().unwrap());
|
||||
(address.to_string(), calls)
|
||||
})
|
||||
.collect();
|
||||
counts.sort();
|
||||
counts
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn ping_reaches_every_node(#[with("ping")] cache: Option<JsonCache>) {
|
||||
let Some(cache) = cache else { return };
|
||||
let startup = redis::Client::open(cluster_url()).unwrap();
|
||||
let before = ping_calls_per_node(&startup);
|
||||
assert!(before.len() >= 2, "{before:?}");
|
||||
assert!(cache.ping().await.unwrap());
|
||||
let after = ping_calls_per_node(&startup);
|
||||
for ((node, calls_before), (_, calls_after)) in before.iter().zip(&after) {
|
||||
assert!(calls_after > calls_before, "{node} was not pinged");
|
||||
}
|
||||
assert!(cache.sync_ping().unwrap());
|
||||
let result = cache.test_connection().await.unwrap();
|
||||
assert_eq!(result.status, CacheConnectionStatus::Success);
|
||||
assert_eq!(result.message, "Redis Cluster connection test successful");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn disconnect_closes_idle_connections_and_reconnects_on_demand(
|
||||
#[with("disconnect")] cache: Option<JsonCache>,
|
||||
) {
|
||||
let Some(cache) = cache else { return };
|
||||
assert!(cache.ping().await.unwrap());
|
||||
cache.disconnect().await.unwrap();
|
||||
assert!(cache.ping().await.unwrap());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::keep_existing_ttl(false)]
|
||||
#[case::refresh_ttl(true)]
|
||||
#[tokio::test]
|
||||
async fn increments_refresh_the_ttl_only_when_asked(
|
||||
counter: Option<Counter>,
|
||||
#[case] refresh_ttl: bool,
|
||||
) {
|
||||
let Some(counter) = counter else { return };
|
||||
let context = ExactCacheContext { ttl: seconds(60) };
|
||||
counter
|
||||
.async_set_cache("spend", 0.0, ExactCacheContext { ttl: seconds(600) })
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
counter
|
||||
.async_increment("spend", 1.5, context.clone(), refresh_ttl)
|
||||
.await
|
||||
.unwrap(),
|
||||
1.5
|
||||
);
|
||||
assert_eq!(
|
||||
counter
|
||||
.async_increment("spend", 2.0, context, refresh_ttl)
|
||||
.await
|
||||
.unwrap(),
|
||||
3.5
|
||||
);
|
||||
let ttl = counter.async_get_ttl("spend").await.unwrap().unwrap();
|
||||
assert_eq!(ttl <= Duration::from_secs(60), refresh_ttl, "{ttl:?}");
|
||||
counter.async_flush_cache().await.unwrap();
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn counters_claims_scripts_and_sets_work_on_the_cluster(
|
||||
counter: Option<Counter>,
|
||||
#[with("claim")] cache: Option<JsonCache>,
|
||||
context: ExactCacheContext,
|
||||
) {
|
||||
let (Some(counter), Some(cache)) = (counter, cache) else {
|
||||
return;
|
||||
};
|
||||
assert_eq!(
|
||||
counter
|
||||
.increment_cache("spend", 1.5, context.clone())
|
||||
.unwrap(),
|
||||
1.5
|
||||
);
|
||||
assert_eq!(
|
||||
counter
|
||||
.async_increment("spend", 2.0, context.clone(), false)
|
||||
.await
|
||||
.unwrap(),
|
||||
3.5
|
||||
);
|
||||
assert_eq!(
|
||||
counter
|
||||
.increment_with_floor("budget", -3, Duration::from_secs(30))
|
||||
.unwrap(),
|
||||
0
|
||||
);
|
||||
assert_eq!(
|
||||
counter
|
||||
.async_increment_with_floor("budget", 7, Duration::from_secs(30))
|
||||
.await
|
||||
.unwrap(),
|
||||
7
|
||||
);
|
||||
assert_eq!(counter.async_set_max("peak", 4.0, None).await.unwrap(), 4.0);
|
||||
assert_eq!(counter.async_set_max("peak", 2.0, None).await.unwrap(), 4.0);
|
||||
counter.flush_cache().unwrap();
|
||||
|
||||
let owner = json!("owner-a");
|
||||
let rival = json!("owner-b");
|
||||
assert_eq!(
|
||||
cache
|
||||
.claim_cache("lock", owner.clone(), &[], context.clone())
|
||||
.unwrap(),
|
||||
owner
|
||||
);
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_claim_cache("lock", rival.clone(), vec![owner.clone()], context.clone())
|
||||
.await
|
||||
.unwrap(),
|
||||
owner
|
||||
);
|
||||
assert_eq!(
|
||||
cache
|
||||
.claim_cache("lock", rival.clone(), &[], context.clone())
|
||||
.unwrap(),
|
||||
owner
|
||||
);
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_claim_cache("lock", rival.clone(), vec![rival.clone()], context.clone())
|
||||
.await
|
||||
.unwrap(),
|
||||
rival
|
||||
);
|
||||
|
||||
let script = cache
|
||||
.async_register_script("return redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2])".into());
|
||||
let reply = script
|
||||
.invoke(
|
||||
vec!["scripted".into()],
|
||||
vec![RedisArg::Bytes(b"payload".to_vec()), RedisArg::Integer(5)],
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(reply, redis::Value::Okay);
|
||||
assert_eq!(cache.async_get_ttl("scripted").await.unwrap(), seconds(5));
|
||||
let evaluated = cache
|
||||
.async_eval(
|
||||
"return redis.call('GET', KEYS[1])".into(),
|
||||
vec!["scripted".into()],
|
||||
Vec::new(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(evaluated, redis::Value::BulkString(b"payload".to_vec()));
|
||||
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_set_cache_sadd(
|
||||
"members",
|
||||
vec![
|
||||
RedisArg::Bytes(b"a".to_vec()),
|
||||
RedisArg::Bytes(b"b".to_vec())
|
||||
],
|
||||
seconds(9),
|
||||
)
|
||||
.await
|
||||
.unwrap(),
|
||||
2
|
||||
);
|
||||
assert_eq!(cache.async_get_ttl("members").await.unwrap(), seconds(9));
|
||||
|
||||
let result = cache.test_connection().await.unwrap();
|
||||
assert_eq!(result.status, CacheConnectionStatus::Success);
|
||||
assert!(cache.ping().await.unwrap());
|
||||
let info = cache.info().unwrap();
|
||||
assert!(info.matches("redis_version").count() > 1, "{info}");
|
||||
assert!(cache.client_list().unwrap().contains("id="));
|
||||
cache.async_flush_cache().await.unwrap();
|
||||
assert_eq!(cache.async_get_ttl("members").await.unwrap(), None);
|
||||
assert_eq!(cache.get_cache("lock", &context).unwrap(), None);
|
||||
}
|
||||
|
|
@ -1,95 +0,0 @@
|
|||
//! The shared cache contracts, run against the in-process fake connection and, when
|
||||
//! `LITELLM_TEST_REDIS_CLUSTER_NODES` is set, against a live Redis Cluster.
|
||||
|
||||
mod support;
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use litellm_cache::{ExactCacheContext, JsonCodec};
|
||||
use litellm_cache_testing as contract;
|
||||
use rstest::rstest;
|
||||
use serde_json::json;
|
||||
use support::{JsonCache, cluster_cache, fake_cache};
|
||||
|
||||
const PREFIX: &str = "contract:";
|
||||
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
enum Contract {
|
||||
HitAndMiss,
|
||||
SyncAsyncEquivalence,
|
||||
OverwriteReplaces,
|
||||
PipelineWritesEveryEntry,
|
||||
BatchPreservesOrder,
|
||||
DeleteRemovesKey,
|
||||
FlushClears,
|
||||
CounterAccumulates,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
enum Server {
|
||||
Fake,
|
||||
Cluster,
|
||||
}
|
||||
|
||||
async fn check<C>(contract: Contract, cache: &JsonCache<C>)
|
||||
where
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
let context = ExactCacheContext::default();
|
||||
match contract {
|
||||
Contract::HitAndMiss => {
|
||||
contract::hit_and_miss(cache, context, PREFIX, json!({"answer": 42})).await
|
||||
}
|
||||
Contract::SyncAsyncEquivalence => {
|
||||
contract::sync_async_equivalence(cache, context, PREFIX, json!("first"), json!([2]))
|
||||
.await
|
||||
}
|
||||
Contract::OverwriteReplaces => {
|
||||
contract::overwrite_replaces(cache, context, PREFIX, json!(1), json!({"b": 2})).await
|
||||
}
|
||||
Contract::PipelineWritesEveryEntry => {
|
||||
contract::pipeline_writes_every_entry(
|
||||
cache,
|
||||
context,
|
||||
PREFIX,
|
||||
vec![json!("a"), json!(2), json!({"c": true})],
|
||||
)
|
||||
.await
|
||||
}
|
||||
Contract::BatchPreservesOrder => {
|
||||
contract::batch_preserves_order(cache, context, PREFIX, json!("first"), json!(2)).await
|
||||
}
|
||||
Contract::DeleteRemovesKey => {
|
||||
contract::delete_removes_key(cache, context, PREFIX, json!("value")).await
|
||||
}
|
||||
Contract::FlushClears => {
|
||||
contract::flush_clears(cache, context, PREFIX, json!("value")).await
|
||||
}
|
||||
Contract::CounterAccumulates => contract::counter_accumulates(cache, context, PREFIX).await,
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::hit_and_miss(Contract::HitAndMiss)]
|
||||
#[case::sync_async_equivalence(Contract::SyncAsyncEquivalence)]
|
||||
#[case::overwrite_replaces(Contract::OverwriteReplaces)]
|
||||
#[case::pipeline_writes_every_entry(Contract::PipelineWritesEveryEntry)]
|
||||
#[case::batch_preserves_order(Contract::BatchPreservesOrder)]
|
||||
#[case::delete_removes_key(Contract::DeleteRemovesKey)]
|
||||
#[case::flush_clears(Contract::FlushClears)]
|
||||
#[case::counter_accumulates(Contract::CounterAccumulates)]
|
||||
#[tokio::test]
|
||||
async fn redis_satisfies_the_cache_contract(
|
||||
#[case] contract: Contract,
|
||||
#[values(Server::Fake, Server::Cluster)] server: Server,
|
||||
) {
|
||||
match server {
|
||||
Server::Fake => check(contract, &fake_cache("contract")).await,
|
||||
Server::Cluster => {
|
||||
let label = format!("{contract:?}");
|
||||
if let Some(cache) = cluster_cache(&label, Duration::from_secs(120), JsonCodec::new()) {
|
||||
check(contract, &cache).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -1,231 +0,0 @@
|
|||
#![allow(dead_code)]
|
||||
|
||||
use std::{
|
||||
collections::BTreeMap,
|
||||
time::{Duration, SystemTime, UNIX_EPOCH},
|
||||
};
|
||||
|
||||
use litellm_cache::{CacheCodec, Error, JsonCodec};
|
||||
use litellm_cache_redis::{RedisCache, RedisNode, RedisTopology};
|
||||
|
||||
/// Encodes a byte behind a tag, so a value written with another tag decodes as invalid.
|
||||
pub struct TaggedByteCodec(pub u8);
|
||||
|
||||
impl CacheCodec for TaggedByteCodec {
|
||||
type Value = u8;
|
||||
|
||||
fn encode(&self, value: &u8) -> Result<Vec<u8>, Error> {
|
||||
if *value > 127 {
|
||||
return Err(Error::InvalidEntry(None));
|
||||
}
|
||||
Ok(vec![self.0, *value])
|
||||
}
|
||||
|
||||
fn decode(&self, bytes: &[u8]) -> Result<u8, Error> {
|
||||
match bytes {
|
||||
[tag, value] if *tag == self.0 => Ok(*value),
|
||||
_ => Err(Error::InvalidEntry(None)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A stateful in-process stand-in for a Redis server that understands the string commands the
|
||||
/// shared contracts exercise, so they run without a live server. TTLs are accepted and ignored.
|
||||
#[derive(Default)]
|
||||
pub struct FakeRedis {
|
||||
strings: BTreeMap<Vec<u8>, Vec<u8>>,
|
||||
}
|
||||
|
||||
impl FakeRedis {
|
||||
fn run(&mut self, command: Vec<Vec<u8>>) -> redis::RedisResult<redis::Value> {
|
||||
let name = String::from_utf8_lossy(&command[0]).to_ascii_uppercase();
|
||||
let args = &command[1..];
|
||||
Ok(match name.as_str() {
|
||||
"PING" => redis::Value::SimpleString("PONG".into()),
|
||||
"SET" | "SETEX" => {
|
||||
let value = if name == "SET" { &args[1] } else { &args[2] };
|
||||
self.strings.insert(args[0].clone(), value.clone());
|
||||
redis::Value::Okay
|
||||
}
|
||||
"GET" => self.get(&args[0]),
|
||||
"MGET" => redis::Value::Array(args.iter().map(|key| self.get(key)).collect()),
|
||||
"DEL" => {
|
||||
let removed = args
|
||||
.iter()
|
||||
.filter(|key| self.strings.remove(*key).is_some())
|
||||
.count();
|
||||
redis::Value::Int(removed as i64)
|
||||
}
|
||||
"SCAN" => {
|
||||
let pattern = &args[2];
|
||||
let keys = self
|
||||
.strings
|
||||
.keys()
|
||||
.filter(|key| glob(pattern, key))
|
||||
.map(|key| redis::Value::BulkString(key.clone()))
|
||||
.collect();
|
||||
redis::Value::Array(vec![
|
||||
redis::Value::BulkString(b"0".to_vec()),
|
||||
redis::Value::Array(keys),
|
||||
])
|
||||
}
|
||||
"EVAL" if args[0].windows(11).any(|window| window == b"INCRBYFLOAT") => {
|
||||
self.increment_by_float(&args[2], &args[3])
|
||||
}
|
||||
"INCRBYFLOAT" => self.increment_by_float(&args[0], &args[1]),
|
||||
_ => {
|
||||
return Err(redis::RedisError::from((
|
||||
redis::ErrorKind::Client,
|
||||
"unsupported command",
|
||||
name,
|
||||
)));
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
fn get(&self, key: &[u8]) -> redis::Value {
|
||||
self.strings.get(key).map_or(redis::Value::Nil, |value| {
|
||||
redis::Value::BulkString(value.clone())
|
||||
})
|
||||
}
|
||||
|
||||
fn increment_by_float(&mut self, key: &[u8], amount: &[u8]) -> redis::Value {
|
||||
let current = self
|
||||
.strings
|
||||
.get(key)
|
||||
.map_or(0.0, |value| parse_float(value));
|
||||
let total = format!("{}", current + parse_float(amount));
|
||||
self.strings
|
||||
.insert(key.to_vec(), total.clone().into_bytes());
|
||||
redis::Value::BulkString(total.into_bytes())
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_float(bytes: &[u8]) -> f64 {
|
||||
std::str::from_utf8(bytes).unwrap().parse().unwrap()
|
||||
}
|
||||
|
||||
/// Redis `MATCH` globbing for `*`, `?` and backslash escapes.
|
||||
fn glob(pattern: &[u8], key: &[u8]) -> bool {
|
||||
match pattern.split_first() {
|
||||
None => key.is_empty(),
|
||||
Some((b'*', rest)) => (0..=key.len()).any(|skip| glob(rest, &key[skip..])),
|
||||
Some((b'?', rest)) => !key.is_empty() && glob(rest, &key[1..]),
|
||||
Some((b'\\', [escaped, rest @ ..])) => {
|
||||
key.first() == Some(escaped) && glob(rest, &key[1..])
|
||||
}
|
||||
Some((literal, rest)) => key.first() == Some(literal) && glob(rest, &key[1..]),
|
||||
}
|
||||
}
|
||||
|
||||
/// Splits RESP request bytes into the commands they carry.
|
||||
fn commands(mut bytes: &[u8]) -> Vec<Vec<Vec<u8>>> {
|
||||
fn line<'a>(bytes: &mut &'a [u8]) -> &'a [u8] {
|
||||
let end = bytes
|
||||
.windows(2)
|
||||
.position(|window| window == b"\r\n")
|
||||
.unwrap();
|
||||
let (line, rest) = bytes.split_at(end);
|
||||
*bytes = &rest[2..];
|
||||
line
|
||||
}
|
||||
fn length(line: &[u8]) -> usize {
|
||||
std::str::from_utf8(&line[1..]).unwrap().parse().unwrap()
|
||||
}
|
||||
let mut commands = Vec::new();
|
||||
while !bytes.is_empty() {
|
||||
let count = length(line(&mut bytes));
|
||||
let command = (0..count)
|
||||
.map(|_| {
|
||||
let size = length(line(&mut bytes));
|
||||
let (argument, rest) = bytes.split_at(size);
|
||||
bytes = &rest[2..];
|
||||
argument.to_vec()
|
||||
})
|
||||
.collect();
|
||||
commands.push(command);
|
||||
}
|
||||
commands
|
||||
}
|
||||
|
||||
impl redis::ConnectionLike for FakeRedis {
|
||||
fn req_packed_command(&mut self, cmd: &[u8]) -> redis::RedisResult<redis::Value> {
|
||||
let command = commands(cmd).into_iter().next().unwrap();
|
||||
self.run(command)
|
||||
}
|
||||
|
||||
fn req_packed_commands(
|
||||
&mut self,
|
||||
cmd: &[u8],
|
||||
offset: usize,
|
||||
count: usize,
|
||||
) -> redis::RedisResult<Vec<redis::Value>> {
|
||||
let replies = commands(cmd)
|
||||
.into_iter()
|
||||
.map(|command| self.run(command))
|
||||
.collect::<redis::RedisResult<Vec<_>>>()?;
|
||||
Ok(replies.into_iter().skip(offset).take(count).collect())
|
||||
}
|
||||
|
||||
fn get_db(&self) -> i64 {
|
||||
0
|
||||
}
|
||||
|
||||
fn check_connection(&mut self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn is_open(&self) -> bool {
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
pub type JsonCache<C = redis::Connection> = RedisCache<JsonCodec<serde_json::Value>, C>;
|
||||
|
||||
pub fn fake_cache(namespace: &str) -> JsonCache<FakeRedis> {
|
||||
RedisCache::with_connection(FakeRedis::default(), None, JsonCodec::new())
|
||||
.with_namespace(Some(namespace.into()))
|
||||
}
|
||||
|
||||
/// Startup nodes from `LITELLM_TEST_REDIS_CLUSTER_NODES` (`host:port,host:port`); tests that
|
||||
/// need a live cluster skip when it is unset.
|
||||
pub fn cluster_topology() -> Option<RedisTopology> {
|
||||
let nodes = std::env::var("LITELLM_TEST_REDIS_CLUSTER_NODES").ok()?;
|
||||
let startup_nodes = nodes
|
||||
.split(',')
|
||||
.map(|node| {
|
||||
let (host, port) = node.trim().rsplit_once(':').expect("host:port");
|
||||
RedisNode {
|
||||
host: host.to_string(),
|
||||
port: port.parse().expect("port"),
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
Some(RedisTopology::Cluster { startup_nodes })
|
||||
}
|
||||
|
||||
pub fn cluster_url() -> String {
|
||||
std::env::var("LITELLM_TEST_REDIS_CLUSTER_URL")
|
||||
.unwrap_or_else(|_| "redis://127.0.0.1:7000".into())
|
||||
}
|
||||
|
||||
pub fn unique_namespace(label: &str) -> String {
|
||||
let nanos = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_nanos();
|
||||
format!("cluster-test:{label}:{nanos}")
|
||||
}
|
||||
|
||||
pub fn cluster_cache<S: CacheCodec>(
|
||||
label: &str,
|
||||
default_ttl: Duration,
|
||||
codec: S,
|
||||
) -> Option<RedisCache<S>> {
|
||||
let topology = cluster_topology()?;
|
||||
Some(
|
||||
RedisCache::connect(&cluster_url(), &topology, Some(default_ttl), codec)
|
||||
.expect("cluster connection")
|
||||
.with_namespace(Some(unique_namespace(label))),
|
||||
)
|
||||
}
|
||||
|
|
@ -16,9 +16,6 @@ sha2.workspace = true
|
|||
litellm-cache-gcs.workspace = true
|
||||
litellm-http = { workspace = true, features = ["test-support"] }
|
||||
litellm-cache-memory.workspace = true
|
||||
litellm-cache-redis.workspace = true
|
||||
redis.workspace = true
|
||||
redis-test.workspace = true
|
||||
rstest.workspace = true
|
||||
tokio.workspace = true
|
||||
wiremock.workspace = true
|
||||
|
|
|
|||
|
|
@ -2,79 +2,11 @@ mod support;
|
|||
|
||||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use litellm_cache::CacheConnectionStatus;
|
||||
use litellm_cache_memory::InMemoryCache;
|
||||
use litellm_cache_response::{
|
||||
CacheEntry, ConnectionProbe, ExactResponseCache, ResponseCache, ResponseCacheRequest,
|
||||
};
|
||||
use redis_test::MockCmd;
|
||||
use litellm_cache_response::{CacheEntry, ExactResponseCache, ResponseCache, ResponseCacheRequest};
|
||||
use rstest::rstest;
|
||||
use serde_json::json;
|
||||
use support::{keyed, memory, redis, request};
|
||||
|
||||
#[rstest]
|
||||
#[case::reachable(
|
||||
Ok("PONG"),
|
||||
CacheConnectionStatus::Success,
|
||||
"Redis connection test successful",
|
||||
false
|
||||
)]
|
||||
#[case::unexpected_reply(
|
||||
Ok("NOPE"),
|
||||
CacheConnectionStatus::Failed,
|
||||
"Redis ping returned False",
|
||||
false
|
||||
)]
|
||||
#[case::connection_refused(
|
||||
Err(redis::RedisError::from((redis::ErrorKind::Io, "connection refused"))),
|
||||
CacheConnectionStatus::Failed,
|
||||
"Redis connection failed:",
|
||||
true
|
||||
)]
|
||||
#[tokio::test]
|
||||
async fn connection_backends_are_reachable_as_a_probe(
|
||||
#[case] reply: redis::RedisResult<&'static str>,
|
||||
#[case] status: CacheConnectionStatus,
|
||||
#[case] message: &str,
|
||||
#[case] has_error: bool,
|
||||
) {
|
||||
let probe: Arc<dyn ConnectionProbe> =
|
||||
Arc::new(redis(vec![MockCmd::new(redis::cmd("PING"), reply)], None));
|
||||
|
||||
let result = probe.test_connection().await.unwrap();
|
||||
assert_eq!(result.status, status);
|
||||
assert!(result.message.starts_with(message), "{}", result.message);
|
||||
assert_eq!(result.error.is_some(), has_error);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn one_service_serves_both_the_exact_cache_and_its_probe(request: ResponseCacheRequest) {
|
||||
let service = Arc::new(redis(
|
||||
vec![
|
||||
MockCmd::new(redis::cmd("PING"), Ok("PONG")),
|
||||
MockCmd::new(
|
||||
redis::cmd("GET").arg("tenant:key"),
|
||||
Ok(br#"{"timestamp":100.0,"response":{"ok":true}}"#.to_vec()),
|
||||
),
|
||||
],
|
||||
Some("tenant"),
|
||||
));
|
||||
let probe: Arc<dyn ConnectionProbe> = service.clone();
|
||||
let exact: Arc<dyn ExactResponseCache> = service;
|
||||
|
||||
assert_eq!(
|
||||
probe.test_connection().await.unwrap().status,
|
||||
CacheConnectionStatus::Success
|
||||
);
|
||||
assert_eq!(
|
||||
exact
|
||||
.async_lookup(&request, Duration::from_secs(100))
|
||||
.await
|
||||
.unwrap(),
|
||||
Some(json!({"ok": true}))
|
||||
);
|
||||
}
|
||||
use support::{keyed, memory, request};
|
||||
|
||||
/// The in-memory backend has no `test_connection`, as in Python, and still serves every response
|
||||
/// operation.
|
||||
|
|
|
|||
|
|
@ -17,10 +17,9 @@ use litellm_cache_response::{
|
|||
CacheControls, CacheEntry, CacheKeyField, CacheKeyInput, ResponseCache, ResponseCacheRequest,
|
||||
WriteBuffer, cache_key,
|
||||
};
|
||||
use redis_test::MockCmd;
|
||||
use rstest::{fixture, rstest};
|
||||
use serde_json::{Value, json};
|
||||
use support::{keyed, memory, redis, request};
|
||||
use support::{keyed, memory, request};
|
||||
|
||||
type Memory = Arc<ResponseCache<InMemoryCache<CacheEntry>>>;
|
||||
|
||||
|
|
@ -331,79 +330,6 @@ async fn directives_skip_io_and_keep_reads_and_writes_independent(
|
|||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::python_sync_literal(
|
||||
br#"{'timestamp': 100.0, 'response': '{"ok": true, "text": "cached"}'}"#.as_slice()
|
||||
)]
|
||||
#[case::python_async_json(br#"{"timestamp":100.0,"response":{"ok":true,"text":"cached"}}"#.as_slice())]
|
||||
#[tokio::test]
|
||||
async fn redis_consumer_reads_python_sync_and_async_envelopes(
|
||||
request: ResponseCacheRequest,
|
||||
#[case] stored: &[u8],
|
||||
#[values(false, true)] asynchronous: bool,
|
||||
) {
|
||||
let cache = redis(
|
||||
vec![MockCmd::new(
|
||||
redis::cmd("GET").arg("tenant:key"),
|
||||
Ok(stored.to_vec()),
|
||||
)],
|
||||
Some("tenant"),
|
||||
);
|
||||
let now = Duration::from_secs(101);
|
||||
let found = if asynchronous {
|
||||
cache.async_lookup(&request, now).await.unwrap()
|
||||
} else {
|
||||
cache.lookup(&request, now).unwrap()
|
||||
};
|
||||
assert_eq!(found, Some(json!({"ok": true, "text": "cached"})));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::object(
|
||||
json!({"ok": true, "text": "cached"}),
|
||||
br#"{"timestamp":100.0,"response":{"ok":true,"text":"cached"}}"#.as_slice()
|
||||
)]
|
||||
#[case::array(json!([1, 2]), br#"{"timestamp":100.0,"response":"[1,2]"}"#.as_slice())]
|
||||
#[tokio::test]
|
||||
async fn redis_consumer_writes_python_compatible_json(
|
||||
request: ResponseCacheRequest,
|
||||
#[case] response: Value,
|
||||
#[case] wire: &[u8],
|
||||
) {
|
||||
let cache = redis(
|
||||
vec![MockCmd::new(
|
||||
redis::cmd("SETEX").arg("tenant:key").arg(600).arg(wire),
|
||||
Ok("OK"),
|
||||
)],
|
||||
Some("tenant"),
|
||||
);
|
||||
cache
|
||||
.async_store(&request, response, Duration::from_secs(100))
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn invalid_entries_are_misses_and_disabled_reads_do_not_touch_redis(
|
||||
mut request: ResponseCacheRequest,
|
||||
) {
|
||||
let cache = redis(
|
||||
vec![MockCmd::new(
|
||||
redis::cmd("GET").arg("tenant:key"),
|
||||
Ok(b"invalid".to_vec()),
|
||||
)],
|
||||
None,
|
||||
);
|
||||
request.controls.no_cache = true;
|
||||
assert_eq!(cache.lookup(&request, Duration::ZERO).unwrap(), None);
|
||||
request.controls.no_cache = false;
|
||||
assert_eq!(
|
||||
cache.async_lookup(&request, Duration::ZERO).await.unwrap(),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn captured_service_keeps_the_selected_backend_for_background_writes(
|
||||
|
|
@ -541,27 +467,6 @@ async fn batch_lookup_reports_partial_hits_and_batch_store_populates_misses(
|
|||
assert_eq!(memory.lookup(&requests[2], now).unwrap(), None);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn batch_lookup_with_no_readable_request_skips_the_backend(
|
||||
#[values(false, true)] asynchronous: bool,
|
||||
) {
|
||||
let cache = redis(Vec::new(), None);
|
||||
let mut request = keyed("key");
|
||||
request.controls.no_cache = true;
|
||||
let requests = [request.clone(), request];
|
||||
let partial = if asynchronous {
|
||||
cache
|
||||
.async_lookup_batch(&requests, Duration::ZERO)
|
||||
.await
|
||||
.unwrap()
|
||||
} else {
|
||||
cache.lookup_batch(&requests, Duration::ZERO).unwrap()
|
||||
};
|
||||
assert_eq!(partial.values, vec![None, None]);
|
||||
assert_eq!(partial.missing_indices, vec![0, 1]);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn deferred_entries_keep_the_time_they_were_produced(
|
||||
|
|
|
|||
|
|
@ -1,15 +1,9 @@
|
|||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use litellm_cache_memory::InMemoryCache;
|
||||
use litellm_cache_redis::RedisCache;
|
||||
use litellm_cache_response::{
|
||||
CacheEntry, CacheKeyInput, ResponseCache, ResponseCacheCodec, ResponseCacheRequest,
|
||||
};
|
||||
use redis_test::{MockCmd, MockRedisConnection};
|
||||
use litellm_cache_response::{CacheEntry, CacheKeyInput, ResponseCache, ResponseCacheRequest};
|
||||
use rstest::fixture;
|
||||
|
||||
pub type MockedRedis = RedisCache<ResponseCacheCodec, MockRedisConnection>;
|
||||
|
||||
#[fixture]
|
||||
pub fn memory() -> Arc<ResponseCache<InMemoryCache<CacheEntry>>> {
|
||||
Arc::new(ResponseCache::new(Arc::new(InMemoryCache::new(
|
||||
|
|
@ -29,12 +23,3 @@ pub fn keyed(key: &str) -> ResponseCacheRequest {
|
|||
..Default::default()
|
||||
})
|
||||
}
|
||||
|
||||
/// A Redis response cache that must receive exactly `commands`, in order.
|
||||
pub fn redis(commands: Vec<MockCmd>, namespace: Option<&str>) -> ResponseCache<MockedRedis> {
|
||||
let connection = MockRedisConnection::new(commands).assert_all_commands_consumed();
|
||||
ResponseCache::new(Arc::new(
|
||||
RedisCache::with_connection(connection, None, ResponseCacheCodec)
|
||||
.with_namespace(namespace.map(str::to_owned)),
|
||||
))
|
||||
}
|
||||
|
|
|
|||
|
|
@ -36,9 +36,12 @@ impl ProvideCredentials for S3Credentials {
|
|||
self.config.secret_access_key.clone(),
|
||||
) {
|
||||
return Ok(AwsCredentials::new(
|
||||
access_key_id,
|
||||
secret_access_key,
|
||||
self.config.session_token.clone(),
|
||||
access_key_id.expose(),
|
||||
secret_access_key.expose(),
|
||||
self.config
|
||||
.session_token
|
||||
.as_ref()
|
||||
.map(|value| value.expose().to_owned()),
|
||||
None,
|
||||
"litellm-s3-cache",
|
||||
));
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue