scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-26 01:12:21 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 99

chore: merge main into litellm_durable_background_interaction_settlement

Browse source
This commit is contained in:
mateo-berri 2026-09-24 12:38:32 -07:00
commit 8b451c6cf9
1376 changed files with 145227 additions and 35780 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

128
.circleci/config.yml
View file

@ -12,6 +12,9 @@ parameters:
migration_source_sha:
type: string
default: ""
routing_parity_base:
type: string
default: ""
orbs:
codecov: codecov/codecov@4.0.1
node: circleci/node@5.1.0 # Add this line to declare the node orb
@ -176,6 +179,9 @@ commands:
image:
type: string
default: postgres:14@sha256:6a70deda415ec296f977890e11aba04a0db9f632a362e3fce45e845e3db74f26
server_args:
type: string
default: ""
steps:
- run:
name: Start PostgreSQL
@ -186,7 +192,7 @@ commands:
-e POSTGRES_PASSWORD=postgres \
-e POSTGRES_DB=<< parameters.db_name >> \
-p 5432:5432 \
<< parameters.image >>
<< parameters.image >> << parameters.server_args >>
- wait_for_service:
url: tcp://localhost:5432
timeout: "60"
@ -3108,6 +3114,10 @@ jobs:
parameters:
suite:
type: string
mode:
type: enum
enum: [standard, replica]
default: standard
machine:
image: ubuntu-2204:2024.04.1
resource_class: large
@ -3142,18 +3152,19 @@ jobs:
command: cd ui/litellm-dashboard && NEXT_TELEMETRY_DISABLED=1 npm run build
- start_postgres:
image: postgres:16@sha256:e17e86066e5ef83e0952a9347f5c792b7ece00972e2aa787a6986f471b3dd3d5
server_args: "-c shared_preload_libraries=pg_stat_statements -c pg_stat_statements.track=all -c pg_stat_statements.max=20000"
- start_redis
- run:
name: Run owned integration contracts
command: bash .circleci/scripts/run_integration.sh << parameters.suite >>
command: bash .circleci/scripts/run_integration.sh << parameters.suite >> << parameters.mode >>
no_output_timeout: 15m
- run:
name: Stop owned database and Redis
when: always
command: |
mkdir -p test-results/integration-<< parameters.suite >>
docker logs postgres-db > test-results/integration-<< parameters.suite >>/postgres.log 2>&1 || true
docker logs redis-cache > test-results/integration-<< parameters.suite >>/redis.log 2>&1 || true
mkdir -p test-results/services-<< parameters.suite >>-<< parameters.mode >>
docker logs postgres-db > test-results/services-<< parameters.suite >>-<< parameters.mode >>/postgres.log 2>&1 || true
docker logs redis-cache > test-results/services-<< parameters.suite >>-<< parameters.mode >>/redis.log 2>&1 || true
docker rm -f postgres-db redis-cache
test -z "$(docker ps -aq --filter name=postgres-db --filter name=redis-cache)"
- store_test_results:
@ -3161,6 +3172,76 @@ jobs:
- store_artifacts:
path: test-results
routing_parity:
parameters:
suite:
type: string
machine:
image: ubuntu-2204:2024.04.1
resource_class: large
working_directory: ~/project
steps:
- setup_litellm_test_deps
- run:
name: Check out base product code
environment:
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 rm -r -f --quiet litellm enterprise litellm-proxy-extras
git checkout "$ROUTING_PARITY_BASE" -- litellm enterprise litellm-proxy-extras
git reset --quiet
test -f litellm/rust_bridge/_native.abi3.so
- start_postgres:
image: postgres:16@sha256:e17e86066e5ef83e0952a9347f5c792b7ece00972e2aa787a6986f471b3dd3d5
server_args: "-c shared_preload_libraries=pg_stat_statements -c pg_stat_statements.track=all -c pg_stat_statements.max=20000"
- start_redis
- run:
name: Run base side
command: bash .circleci/scripts/run_integration.sh << parameters.suite >> parity base
no_output_timeout: 15m
- run:
name: Stop base database and Redis
when: always
command: |
mkdir -p test-results/services-<< parameters.suite >>-parity-base
docker logs postgres-db > test-results/services-<< parameters.suite >>-parity-base/postgres.log 2>&1 || true
docker logs redis-cache > test-results/services-<< parameters.suite >>-parity-base/redis.log 2>&1 || true
docker rm -f postgres-db redis-cache
test -z "$(docker ps -aq --filter name=postgres-db --filter name=redis-cache)"
- run:
name: Check out head product code
command: |
git rm -r -f --quiet litellm enterprise litellm-proxy-extras
git checkout "$CIRCLE_SHA1" -- litellm enterprise litellm-proxy-extras
git reset --quiet
test -f litellm/rust_bridge/_native.abi3.so
- start_postgres:
image: postgres:16@sha256:e17e86066e5ef83e0952a9347f5c792b7ece00972e2aa787a6986f471b3dd3d5
server_args: "-c shared_preload_libraries=pg_stat_statements -c pg_stat_statements.track=all -c pg_stat_statements.max=20000"
- start_redis
- run:
name: Run head side
command: bash .circleci/scripts/run_integration.sh << parameters.suite >> parity head
no_output_timeout: 15m
- run:
name: Stop head database and Redis
when: always
command: |
mkdir -p test-results/services-<< parameters.suite >>-parity-head
docker logs postgres-db > test-results/services-<< parameters.suite >>-parity-head/postgres.log 2>&1 || true
docker logs redis-cache > test-results/services-<< parameters.suite >>-parity-head/redis.log 2>&1 || true
docker rm -f postgres-db redis-cache
test -z "$(docker ps -aq --filter name=postgres-db --filter name=redis-cache)"
- run:
name: Compare routing parity
command: PYTHONPATH="$PWD/tests" .venv/bin/python -m integration._support.routing check test-results/parity-<< parameters.suite >>/base test-results/parity-<< parameters.suite >>/head
- store_test_results:
path: test-results
- store_artifacts:
path: test-results
unit:
machine:
image: ubuntu-2204:2024.04.1
@ -3222,23 +3303,52 @@ workflows:
cron: "17 0,6,12,18 * * *"
filters:
branches:
only: litellm_internal_staging
only: main
jobs: *migration_jobs
routing_parity:
when:
not:
equal: ["", << pipeline.parameters.routing_parity_base >>]
jobs:
- routing_parity:
name: routing-parity-<< matrix.suite >>
matrix:
parameters:
suite: [management, accounting, database, providers, extensions, cost, mcp]
integration:
unless: << pipeline.parameters.run_migration_tests >>
unless:
or:
- << pipeline.parameters.run_migration_tests >>
- not:
equal: ["", << pipeline.parameters.routing_parity_base >>]
jobs:
- integration_contracts:
name: integration-<< matrix.suite >>
matrix:
parameters:
suite: [management, accounting, database, providers, extensions, sdk, cost, browser]
suite: [management, accounting, database, providers, extensions, mcp, sdk, cost, browser]
filters:
branches:
only:
- main
- /litellm_.*/
- integration_contracts:
name: integration-<< matrix.suite >>-replica
matrix:
parameters:
suite: [management, database]
mode: [replica]
filters:
branches:
only:
- main
- /litellm_.*/
build_and_test:
unless: << pipeline.parameters.run_migration_tests >>
unless:
or:
- << pipeline.parameters.run_migration_tests >>
- not:
equal: ["", << pipeline.parameters.routing_parity_base >>]
jobs:
- using_litellm_on_windows:
filters: &main_branches

2
.circleci/scripts/path_filter.sh
View file

@ -11,7 +11,7 @@ run_full() {
[ -n "${CIRCLE_PULL_REQUEST:-}" ] || run_full "not a pull request"
candidate_bases="main litellm_internal_staging litellm_oss_staging"
candidate_bases="${PATH_FILTER_BASE_BRANCH:-main}"
merge_base=""
for base in $candidate_bases; do
git fetch --quiet origin "$base" 2>/dev/null || continue

52
.circleci/scripts/prepare_replica_roles.py Normal file
View file

@ -0,0 +1,52 @@
from __future__ import annotations
import os
from typing import Final
from urllib.parse import urlsplit, urlunsplit
import psycopg
DATABASE_URL: Final = os.environ["DATABASE_URL"]
def postgres_url() -> str:
parsed: Final = urlsplit(DATABASE_URL)
return urlunsplit(parsed._replace(path="/postgres"))
def main() -> None:
with psycopg.connect(postgres_url(), autocommit=True) as admin:
admin.execute("CREATE EXTENSION IF NOT EXISTS pg_stat_statements")
admin.execute("CREATE ROLE litellm_writer LOGIN PASSWORD 'litellm-writer' NOSUPERUSER")
admin.execute("CREATE ROLE litellm_reader LOGIN PASSWORD 'litellm-reader' NOSUPERUSER NOINHERIT")
admin.execute("ALTER ROLE litellm_reader SET default_transaction_read_only = on")
admin.execute("ALTER DATABASE circle_test OWNER TO litellm_writer")
admin.execute("GRANT CONNECT ON DATABASE circle_test TO litellm_reader")
with psycopg.connect(DATABASE_URL, autocommit=True) as admin:
admin.execute("GRANT USAGE ON SCHEMA public TO litellm_reader")
admin.execute(
"ALTER DEFAULT PRIVILEGES FOR ROLE litellm_writer IN SCHEMA public GRANT SELECT ON TABLES TO litellm_reader"
)
admin.execute("GRANT SELECT ON ALL TABLES IN SCHEMA public TO litellm_reader")
parsed: Final = urlsplit(DATABASE_URL)
reader_url: Final = urlunsplit(
parsed._replace(netloc=f"litellm_reader:litellm-reader@{parsed.hostname}:{parsed.port}")
)
writer_url: Final = urlunsplit(
parsed._replace(netloc=f"litellm_writer:litellm-writer@{parsed.hostname}:{parsed.port}")
)
with psycopg.connect(reader_url, autocommit=True) as reader:
assert reader.execute("SHOW transaction_read_only").fetchone() == ("on",)
try:
reader.execute("CREATE TABLE integration_readonly_probe (id int)")
except psycopg.errors.ReadOnlySqlTransaction:
pass
else:
raise AssertionError("litellm_reader executed a write statement")
with psycopg.connect(writer_url, autocommit=True) as writer:
assert writer.execute("SELECT current_user").fetchone() == ("litellm_writer",)
if __name__ == "__main__":
main()

70
.circleci/scripts/run_integration.sh
View file

@ -7,9 +7,16 @@ if [ "${GITHUB_ACTIONS:-}" = true ]; then
fi
suite="${1:?integration suite required}"
results="test-results/integration-${suite}"
mode="${2:-standard}"
side="${3:-}"
if [ "$mode" = replica ]; then
results="test-results/integration-${suite}-replica"
elif [ "$mode" = parity ]; then
results="test-results/parity-${suite}/${side:?parity side required}"
else
results="test-results/integration-${suite}"
fi
mkdir -p "$results"
shard_timeout=11m
integration_identity="$(.venv/bin/python -c 'import uuid; print(uuid.uuid4().hex)')"
upstream_pid=""
proxy_pid=""
@ -81,6 +88,18 @@ export INTEGRATION_ORDER_SEED="$INTEGRATION_SEED"
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma > "$results/prisma-generate.log" 2>&1
export INTEGRATION_PROXY_DATABASE_URL=""
export INTEGRATION_PROXY_READ_REPLICA_URL=""
export INTEGRATION_ROUTING=""
if [ "$mode" = replica ] || [ "$mode" = parity ]; then
.venv/bin/python .circleci/scripts/prepare_replica_roles.py > "$results/prepare-replica-roles.log" 2>&1
export INTEGRATION_PROXY_DATABASE_URL="postgresql://litellm_writer:litellm-writer@127.0.0.1:5432/circle_test"
export INTEGRATION_PROXY_READ_REPLICA_URL="postgresql://litellm_reader:litellm-reader@127.0.0.1:5432/circle_test"
fi
if [ "$mode" = parity ]; then
export INTEGRATION_ROUTING=capture
fi
sudo iptables -N integration_only
guard_created=true
sudo iptables -A integration_only -o lo -j ACCEPT
@ -112,6 +131,15 @@ upstream_pid=$!
if [ "$suite" = cost ]; then
export INTEGRATION_WORKERS=8
fi
if [ "$suite" = mcp ]; then
export INTEGRATION_WORKERS=4 INTEGRATION_COVERAGE=1
fi
coverage_data="$PWD/$results/coverage/data"
proxy_command=(.venv/bin/python -m integration._support.proxy)
if [ "${INTEGRATION_COVERAGE:-0}" = 1 ]; then
mkdir -p "$(dirname "$coverage_data")"
proxy_command=(.venv/bin/python -m coverage run --rcfile=tests/integration/mcp_coverage.toml -m integration._support.proxy)
fi
start_proxy() {
local port="$1"
local log_name="$2"
@ -129,12 +157,17 @@ start_proxy() {
else
cost_map_env=("LITELLM_LOCAL_MODEL_COST_MAP=True")
fi
local -a database_env=("DATABASE_URL=${INTEGRATION_PROXY_DATABASE_URL:-$DATABASE_URL}")
if [ -n "$INTEGRATION_PROXY_READ_REPLICA_URL" ]; then
database_env+=("DATABASE_URL_READ_REPLICA=$INTEGRATION_PROXY_READ_REPLICA_URL")
fi
setsid env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" INTEGRATION_RUN_ID="$integration_identity" \
DATABASE_URL="$DATABASE_URL" REDIS_HOST="$REDIS_HOST" REDIS_PORT="$REDIS_PORT" \
"${database_env[@]}" REDIS_HOST="$REDIS_HOST" REDIS_PORT="$REDIS_PORT" \
INTEGRATION_UPSTREAM_URL="$INTEGRATION_UPSTREAM_URL" \
LITELLM_MASTER_KEY="$LITELLM_MASTER_KEY" LITELLM_SALT_KEY="$LITELLM_SALT_KEY" LITELLM_UI_PATH="$LITELLM_UI_PATH" PROXY_BASE_URL="http://127.0.0.1:$port" \
LITELLM_MODE=PRODUCTION STORE_MODEL_IN_DB=True "${cost_map_env[@]}" \
AWS_EC2_METADATA_DISABLED=true DO_NOT_TRACK=1 \
.venv/bin/python -m integration._support.proxy --config tests/integration/proxy_config.yaml \
AWS_EC2_METADATA_DISABLED=true DO_NOT_TRACK=1 COVERAGE_FILE="$coverage_data" \
"${proxy_command[@]}" --config tests/integration/proxy_config.yaml \
--host 127.0.0.1 --port "$port" --num_workers 1 --telemetry False \
--use_prisma_db_push --enforce_prisma_migration_check \
> "$results/$log_name" 2>&1 &
@ -146,7 +179,7 @@ proxy_pid="$launched_pid"
curl --noproxy '*' -sSf -X POST "$INTEGRATION_PROXY_URL/config/update" \
-H "Authorization: Bearer $LITELLM_MASTER_KEY" -H 'Content-Type: application/json' \
-d '{"router_settings": {"num_retries": 0}}' > "$results/seed-router-settings.json"
if [ "$suite" = management ]; then
if [ "$suite" = management ] || [ "$suite" = mcp ]; then
export INTEGRATION_PEER_URL=http://127.0.0.1:4001
start_proxy 4001 peer.log
peer_pid="$launched_pid"
@ -176,7 +209,7 @@ if [ "$suite" = browser ]; then
exit 0
fi
timeout --signal=TERM --kill-after=20s "$shard_timeout" env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" \
env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" \
INTEGRATION_RUN_ID="$integration_identity" \
DATABASE_URL="$DATABASE_URL" REDIS_HOST="$REDIS_HOST" REDIS_PORT="$REDIS_PORT" \
INTEGRATION_PROXY_URL="$INTEGRATION_PROXY_URL" INTEGRATION_PEER_URL="$INTEGRATION_PEER_URL" \
@ -186,4 +219,27 @@ timeout --signal=TERM --kill-after=20s "$shard_timeout" env -i PATH="$PATH" HOME
INTEGRATION_SEED="$INTEGRATION_SEED" \
INTEGRATION_ORDER_SEED="$INTEGRATION_ORDER_SEED" \
LITELLM_LOCAL_MODEL_COST_MAP=True AWS_EC2_METADATA_DISABLED=true DO_NOT_TRACK=1 \
INTEGRATION_PROXY_DATABASE_URL="$INTEGRATION_PROXY_DATABASE_URL" \
INTEGRATION_PROXY_READ_REPLICA_URL="$INTEGRATION_PROXY_READ_REPLICA_URL" \
INTEGRATION_ROUTING="$INTEGRATION_ROUTING" \
.venv/bin/python tests/integration/run.py "$suite" --results "$results"
if [ "${INTEGRATION_COVERAGE:-0}" = 1 ]; then
for covered_pid in "$proxy_pid" "$peer_pid"; do
[ -n "$covered_pid" ] || continue
kill -TERM -- "-$covered_pid"
for _ in {1..300}; do
kill -0 "$covered_pid" 2>/dev/null || break
sleep 0.1
done
wait "$covered_pid" 2>/dev/null || true
done
proxy_pid=""
peer_pid=""
COVERAGE_FILE="$coverage_data" .venv/bin/python -m coverage combine --rcfile=tests/integration/mcp_coverage.toml
COVERAGE_FILE="$coverage_data" .venv/bin/python -m coverage report --rcfile=tests/integration/mcp_coverage.toml \
> "$results/coverage/coverage.txt"
COVERAGE_FILE="$coverage_data" .venv/bin/python -m coverage html --rcfile=tests/integration/mcp_coverage.toml \
-d "$results/coverage/html"
tail -n 1 "$results/coverage/coverage.txt"
fi

4
.circleci/scripts/verify_integration_browser.py
View file

@ -31,8 +31,8 @@ def main() -> None:
result: Final = json.loads(Path(sys.argv[1]).read_text())
assert not result.get("errors"), result.get("errors")
expected: Final = json.loads(
(Path(__file__).resolve().parents[2] / "tests/integration/contracts.json").read_text()
)["browser"]
(Path(__file__).resolve().parents[2] / "tests/e2e/ui/tests/integrationCritical/expected.json").read_text()
)
assert expected and result["stats"]["expected"] == len(expected)
assert all(result["stats"][name] == 0 for name in ("unexpected", "flaky", "skipped"))

292
.circleci/tests.yml Normal file
View file

@ -0,0 +1,292 @@
version: 2.1
commands:
wait_for_service:
parameters:
url:
type: string
timeout:
type: string
default: "60"
steps:
- run:
name: "Wait for << parameters.url >>"
command: |
TIMEOUT=<< parameters.timeout >>
URL="<< parameters.url >>"
ELAPSED=0
echo "Waiting up to ${TIMEOUT}s for ${URL} ..."
if echo "$URL" | grep -q '^tcp://'; then
HOST=$(echo "$URL" | sed 's|tcp://||' | cut -d: -f1)
PORT=$(echo "$URL" | sed 's|tcp://||' | cut -d: -f2)
while ! bash -c "echo > /dev/tcp/$HOST/$PORT" 2>/dev/null; do
sleep 2; ELAPSED=$((ELAPSED+2))
if [ "$ELAPSED" -ge "$TIMEOUT" ]; then echo "Timed out"; exit 1; fi
done
else
while ! curl -sf --max-time 5 "$URL" > /dev/null 2>&1; do
sleep 2; ELAPSED=$((ELAPSED+2))
if [ "$ELAPSED" -ge "$TIMEOUT" ]; then echo "Timed out"; exit 1; fi
done
fi
echo "Service ready after ${ELAPSED}s"
install_uv:
steps:
- run:
name: Install uv (pinned 0.10.9)
command: |
curl -LsSf -o /tmp/uv-install.sh https://astral.sh/uv/0.10.9/install.sh
echo "7fc46e39cb97290b57169c0c813a17970585ac519139f19006453c99b5f2f45f /tmp/uv-install.sh" | sha256sum -c -
env UV_NO_MODIFY_PATH=1 sh /tmp/uv-install.sh
rm -f /tmp/uv-install.sh
echo 'export PATH="$HOME/.local/bin:$PATH"' >> "$BASH_ENV"
export PATH="$HOME/.local/bin:$PATH"
install_rust:
steps:
- run:
name: Install Rust (rustup 1.28.2, toolchain 1.98.0)
command: |
case "$(uname -m)" in
x86_64)
RUSTUP_TRIPLE=x86_64-unknown-linux-gnu
RUSTUP_SHA256=20a06e644b0d9bd2fbdbfd52d42540bdde820ea7df86e92e533c073da0cdd43c
;;
aarch64)
RUSTUP_TRIPLE=aarch64-unknown-linux-gnu
RUSTUP_SHA256=e3853c5a252fca15252d07cb23a1bdd9377a8c6f3efa01531109281ae47f841c
;;
*)
echo "install_rust: unsupported architecture $(uname -m)" >&2
exit 1
;;
esac
curl -sSLf -o /tmp/rustup-init \
"https://static.rust-lang.org/rustup/archive/1.28.2/${RUSTUP_TRIPLE}/rustup-init"
echo "${RUSTUP_SHA256} /tmp/rustup-init" | sha256sum -c -
chmod +x /tmp/rustup-init
/tmp/rustup-init -y --no-modify-path --profile minimal --default-toolchain 1.98.0
rm -f /tmp/rustup-init
echo 'export PATH="$HOME/.cargo/bin:$PATH"' >> "$BASH_ENV"
export PATH="$HOME/.cargo/bin:$PATH"
rustc --version
cargo --version
install_codecov_cli:
steps:
- run:
name: Install Codecov CLI (pinned v11.3.1)
command: |
curl -sSLf -o /tmp/codecov https://cli.codecov.io/v11.3.1/linux/codecov
curl -sSLf -o /tmp/codecov.SHA256SUM https://cli.codecov.io/v11.3.1/linux/codecov.SHA256SUM
[ "$(cat /tmp/codecov.SHA256SUM)" = "ca1d64196d2d34771084afe76ea657d581bf628e31d993ff8e52ea09cc88a56d codecov" ]
(cd /tmp && sha256sum -c codecov.SHA256SUM)
chmod +x /tmp/codecov
mkdir -p "$HOME/.local/bin"
mv /tmp/codecov "$HOME/.local/bin/codecov"
setup_litellm_enterprise_pip:
steps:
- run:
name: "Install local version of litellm-enterprise"
command: |
uv run --no-sync python -c "import litellm_enterprise; print('litellm-enterprise OK:', litellm_enterprise.__file__)"
setup_test_deps:
steps:
- checkout
- install_uv
- install_rust
- restore_cache:
keys:
- v3-integration-uv-cache-{{ checksum "uv.lock" }}
- run:
name: Install Dependencies
command: |
uv sync --frozen --all-groups --all-extras --python 3.12
- setup_litellm_enterprise_pip
- save_cache:
paths:
- ~/.cache/uv
key: v3-integration-uv-cache-{{ checksum "uv.lock" }}
- run:
name: Generate Prisma client
command: uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
skip_unless_relevant:
parameters:
category:
type: string
default: backend
base_ref:
type: string
default: ""
pull_request_url:
type: string
default: ""
steps:
- run:
name: "Skip job when no << parameters.category >>-relevant files changed"
command: |
export CIRCLE_PULL_REQUEST="${CIRCLE_PULL_REQUEST:-<< parameters.pull_request_url >>}"
export PATH_FILTER_BASE_BRANCH="<< parameters.base_ref >>"
[ -n "$PATH_FILTER_BASE_BRANCH" ] || unset PATH_FILTER_BASE_BRANCH
bash .circleci/scripts/path_filter.sh << parameters.category >>
start_postgres:
parameters:
db_name:
type: string
default: circle_test
image:
type: string
default: postgres:14@sha256:6a70deda415ec296f977890e11aba04a0db9f632a362e3fce45e845e3db74f26
steps:
- run:
name: Start PostgreSQL
command: |
docker run -d \
--name postgres-db \
-e POSTGRES_USER=postgres \
-e POSTGRES_PASSWORD=postgres \
-e POSTGRES_DB=<< parameters.db_name >> \
-p 5432:5432 \
<< parameters.image >>
- wait_for_service:
url: tcp://localhost:5432
timeout: "60"
start_redis:
steps:
- run:
name: Start Redis
command: |
docker run -d \
--name redis-cache \
-p 6379:6379 \
redis:7-alpine@sha256:7aec734b2bb298a1d769fd8729f13b8514a41bf90fcdd1f38ec52267fbaa8ee6
- wait_for_service:
url: tcp://localhost:6379
timeout: "60"
jobs:
unit:
parameters:
tests_path:
type: string
default: tests/unit
flag:
type: string
default: unit
shards:
type: integer
default: 6
base_ref:
type: string
default: ""
pull_request_url:
type: string
default: ""
machine:
image: ubuntu-2204:2024.04.1
resource_class: large
working_directory: ~/project
parallelism: << parameters.shards >>
environment:
LITELLM_LOCAL_MODEL_COST_MAP: "True"
steps:
- setup_test_deps
- skip_unless_relevant:
base_ref: << parameters.base_ref >>
pull_request_url: << parameters.pull_request_url >>
- run:
name: "Run << parameters.tests_path >> shard"
no_output_timeout: 20m
command: |
mkdir -p test-results/<< parameters.flag >>
mapfile -t files < <(find << parameters.tests_path >> -name 'test_*.py' | sort | circleci tests split --split-by=timings --timings-type=filename)
if [ "${#files[@]}" -eq 0 ]; then echo "shard ${CIRCLE_NODE_INDEX} received no << parameters.tests_path >> files; nothing to run"; exit 0; fi
set +e
uv run --no-sync pytest "${files[@]}" -p no:rerunfailures -p no:pytest-retry --timeout=90 -n 4 --dist=loadscope --tb=short --durations=20 -o junit_family=xunit1 --junitxml=test-results/<< parameters.flag >>/junit.xml --cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml:coverage.xml --cov-config=pyproject.toml
status=$?
set -e
if [ "$status" -eq 5 ]; then echo "pytest collected no tests from the shard; passing"; exit 0; fi
exit "$status"
- install_codecov_cli
- run:
name: Upload coverage
when: always
command: |
[ -f coverage.xml ] || { echo "no coverage.xml produced; skipping upload"; exit 0; }
codecov upload-process --disable-search -f coverage.xml -F << parameters.flag >> -C "$CIRCLE_SHA1" -n "<< parameters.flag >>-${CIRCLE_NODE_INDEX}-${CIRCLE_BUILD_NUM}" --git-service github
- store_test_results:
path: test-results
- store_artifacts:
path: test-results
- store_artifacts:
path: coverage.xml
documentation:
machine:
image: ubuntu-2204:2024.04.1
resource_class: large
working_directory: ~/project
steps:
- setup_test_deps
- run:
name: Checkout litellm-docs
command: rm -rf docs/my-website && git clone --depth 1 https://github.com/BerriAI/litellm-docs.git docs/my-website
- run:
name: Run documentation validation
command: |
uv run --no-sync python ./tests/documentation_tests/test_env_keys.py
uv run --no-sync python ./tests/documentation_tests/test_router_settings.py
uv run --no-sync python ./tests/documentation_tests/test_api_docs.py
uv run --no-sync python ./tests/documentation_tests/test_circular_imports.py
integration:
parameters:
suite:
type: string
base_ref:
type: string
default: ""
pull_request_url:
type: string
default: ""
machine:
image: ubuntu-2204:2024.04.1
resource_class: large
working_directory: ~/project
steps:
- setup_test_deps
- skip_unless_relevant:
base_ref: << parameters.base_ref >>
pull_request_url: << parameters.pull_request_url >>
- start_postgres:
image: postgres:16@sha256:e17e86066e5ef83e0952a9347f5c792b7ece00972e2aa787a6986f471b3dd3d5
- start_redis
- run:
name: Run owned integration contracts
command: bash .circleci/scripts/run_integration.sh << parameters.suite >>
no_output_timeout: 15m
- run:
name: Stop owned database and Redis
when: always
command: |
mkdir -p test-results/integration-<< parameters.suite >>
docker logs postgres-db > test-results/integration-<< parameters.suite >>/postgres.log 2>&1 || true
docker logs redis-cache > test-results/integration-<< parameters.suite >>/redis.log 2>&1 || true
docker rm -f postgres-db redis-cache
test -z "$(docker ps -aq --filter name=postgres-db --filter name=redis-cache)"
- store_test_results:
path: test-results
- store_artifacts:
path: test-results
workflows:
tests:
when: (pipeline.event.name == "push" and pipeline.git.branch == "main") or pipeline.event.name == "api" or (pipeline.event.name == "pull_request" and (pipeline.event.github.pull_request.base.ref == "main" or pipeline.event.github.pull_request.base.ref starts-with "litellm_"))
jobs:
- unit:
base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >>
pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >>
- documentation
- integration:
name: integration-<< matrix.suite >>
matrix:
parameters:
suite: [sdk]
base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >>
pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >>

6
.githooks/pre-push
View file

@ -8,7 +8,6 @@
#
# Protected branches (always allowed):
# - main
# - litellm_internal_staging
# - dependabot/*
# - gh-readonly-queue/*
#
@ -22,7 +21,7 @@ ZERO_OID_SHA256="000000000000000000000000000000000000000000000000000000000000000
ALLOWED_TYPES="feature|bugfix|hotfix|release|chore"
BRANCH_PATTERN="^(${ALLOWED_TYPES})/.+"
PROTECTED_NAMES="main litellm_internal_staging"
PROTECTED_NAMES="main"
PROTECTED_PREFIXES="dependabot/ gh-readonly-queue/"
is_protected() {
@ -78,8 +77,7 @@ if [ -n "$invalid" ]; then
chore/bump-deps
hotfix/auth-bypass
Protected (always allowed): main, litellm_internal_staging,
dependabot/*, gh-readonly-queue/*.
Protected (always allowed): main, dependabot/*, gh-readonly-queue/*.
See https://conventional-branch.github.io/

13
.github/ci-coverage-allowlist.yml vendored
View file

@ -10,12 +10,13 @@ test_paths:
paths:
- tests/rust-python-harness
- reason: >-
What is left of the caching suite in tests/local_testing that runs nowhere. Every job that
globs that directory either deselects it (local_testing_part1 and part2 carry `-k "... and
not caching and not cache"`) or keeps only another keyword (langfuse, router, assistants),
and no job names these files the way redis_caching_unit_tests names test_dual_cache.py.
The gap was eight files and 118 tests when measured 2026-08-20; the five keyless ones now
run in the caching-local shard, leaving these three. Measured 2026-08-21 with no provider
Live-provider caching cases in tests/local_testing that remain outside CI. Jobs that
glob that directory either deselect them (local_testing_part1 and part2 carry `-k "... and
not caching and not cache"`) or keep only another keyword (langfuse, router, assistants).
Separately, test-redis-compat.yml selects two IAM cluster authentication tests in
test_caching.py by node ID. It does not run that file's other tests.
The gap was eight files and 118 tests when measured 2026-08-20; the five keyless files now
run in the caching-local shard, leaving live cases in these three. Measured 2026-08-21 with no provider
credentials and no Redis: test_caching.py needs both (37 of 65 fail without them),
test_disk_cache_unit_tests.py needs OPENAI_API_KEY for 2 of its 4, and
test_gcs_cache_unit_tests.py needs GCS credentials for all 4. They want the keyless/live

14
.github/e2e-stack/secrets_to_env.py vendored
View file

@ -9,6 +9,7 @@ from pydantic import TypeAdapter, ValidationError
secrets_adapter: Final[TypeAdapter[dict[str, str]]] = TypeAdapter(dict[str, str])
ENV_NAME: Final = re.compile(r"[A-Za-z_][A-Za-z0-9_]*")
MIN_MASKED_LENGTH: Final = 8
ACTIONS_RUNNER_FLAG: Final = "GITHUB_ACTIONS"
def main() -> int:
@ -30,10 +31,15 @@ def main() -> int:
f"these names or values cannot be represented in both bash and dotenv: {' '.join(sorted(unusable))}\n"
)
return 1
for value in secrets.values():
if len(value) >= MIN_MASKED_LENGTH:
_ = sys.stdout.write(f"::add-mask::{value.replace('%', '%25')}\n")
sys.stdout.flush()
if os.environ.get(ACTIONS_RUNNER_FLAG) == "true":
_ = sys.stdout.write(
"".join(
f"::add-mask::{value.replace('%', '%25')}\n"
for value in secrets.values()
if len(value) >= MIN_MASKED_LENGTH
)
)
sys.stdout.flush()
lines: Final = tuple(f"{key}='{value}'" for key, value in secrets.items() if value)
try:
with os.fdopen(os.open(env_path, os.O_WRONLY | os.O_APPEND | os.O_CREAT | os.O_NOFOLLOW, 0o600), "w") as handle:

2
.github/e2e-stack/select_tests.py vendored
View file

@ -10,6 +10,8 @@ UNSUPPORTED: Final = re.compile(
r"|^tests/e2e/batches/test_managed_files_enforcement_e2e\.py$"
r"|^tests/e2e/guardrails/test_presidio_masking_e2e\.py$"
r"|^tests/e2e/logging/test_otel_v2_langfuse_generation_output_e2e\.py$"
r"|^tests/e2e/logging/test_langsmith_batch_serialization_e2e\.py$"
r"|^tests/e2e/secret_manager/"
)
HARNESS: Final = re.compile(
r"^tests/e2e/[A-Za-z0-9_.-]+\.(py|ini)$"

29
.github/e2e-stack/up.sh vendored
View file

@ -24,6 +24,7 @@ DATABASE_USER="${E2E_DATABASE_USER:-litellm}"
DATABASE_PASSWORD="${E2E_DATABASE_PASSWORD:-dbpassword9090}"
DATABASE_NAME="${E2E_DATABASE_NAME:-litellm}"
JAEGER_OTLP_PORT="${E2E_JAEGER_OTLP_PORT:-4318}"
JAEGER_OTLP_TLS_PORT="${E2E_JAEGER_OTLP_TLS_PORT:-4319}"
JAEGER_QUERY_PORT="${E2E_JAEGER_QUERY_PORT:-16686}"
KEYCLOAK_PORT="${E2E_KEYCLOAK_PORT:-8081}"
@ -122,7 +123,7 @@ SERVER_ENV=(
"CONFIG_FILE_PATH=${CONFIG_PATH}"
"STORE_MODEL_IN_DB=True"
"OTEL_EXPORTER_OTLP_PROTOCOL=http/protobuf"
"OTEL_EXPORTER_OTLP_ENDPOINT=http://127.0.0.1:${JAEGER_OTLP_PORT}"
"OTEL_EXPORTER_OTLP_ENDPOINT=https://127.0.0.1:${JAEGER_OTLP_TLS_PORT}"
"SSL_CERT_FILE=${CERTS_DIR}/ca-bundle.pem"
"PYTHONPATH=${REPO_ROOT}"
"JWT_PUBLIC_KEY_URL=http://127.0.0.1:${KEYCLOAK_PORT}/realms/litellm-e2e/protocol/openid-connect/certs"
@ -147,16 +148,12 @@ start_server() {
echo $! > "${PIDS_DIR}/${name}.pid"
}
start_server backend uv run --no-sync uvicorn backend.main:app --host 0.0.0.0 --port "${BACKEND_PORT}"
start_server gateway-1 uv run --no-sync uvicorn gateway.main:app --workers 1 --host 0.0.0.0 --port "${GATEWAY_PORT_1}"
start_server gateway-2 uv run --no-sync uvicorn gateway.main:app --workers 1 --host 0.0.0.0 --port "${GATEWAY_PORT_2}"
if [[ "$(uname)" == "Linux" ]]; then
NGINX_UPSTREAM_HOST=127.0.0.1
NGINX_DOCKER_ARGS=(--network host)
else
NGINX_UPSTREAM_HOST=host.docker.internal
NGINX_DOCKER_ARGS=(-p "${LB_PORT}:${LB_PORT}")
NGINX_DOCKER_ARGS=(-p "${LB_PORT}:${LB_PORT}" -p "${JAEGER_OTLP_TLS_PORT}:${JAEGER_OTLP_TLS_PORT}")
fi
cat > "${STACK_DIR}/nginx.conf" <<EOF
@ -186,12 +183,29 @@ http {
proxy_send_timeout 600s;
}
}
server {
listen ${JAEGER_OTLP_TLS_PORT} ssl;
ssl_certificate /certs/server.crt;
ssl_certificate_key /certs/server.key;
client_max_body_size 100m;
location / {
proxy_pass http://${NGINX_UPSTREAM_HOST}:${JAEGER_OTLP_PORT};
}
}
}
EOF
docker rm -f e2e-nginx >/dev/null 2>&1 || true
docker run -d --name e2e-nginx "${NGINX_DOCKER_ARGS[@]}" \
-v "${STACK_DIR}/nginx.conf:/etc/nginx/nginx.conf:ro" "${NGINX_IMAGE}" >/dev/null
-v "${STACK_DIR}/nginx.conf:/etc/nginx/nginx.conf:ro" \
-v "${CERTS_DIR}:/certs:ro" "${NGINX_IMAGE}" >/dev/null
wait_for "Jaeger OTLP TLS listener" \
"curl -sS --cacert ${CERTS_DIR}/ca.crt https://127.0.0.1:${JAEGER_OTLP_TLS_PORT}/ -o /dev/null -w '%{http_code}' | grep -qE '^[2345]'"
start_server backend uv run --no-sync uvicorn backend.main:app --host 0.0.0.0 --port "${BACKEND_PORT}"
start_server gateway-1 uv run --no-sync uvicorn gateway.main:app --workers 1 --host 0.0.0.0 --port "${GATEWAY_PORT_1}"
start_server gateway-2 uv run --no-sync uvicorn gateway.main:app --workers 1 --host 0.0.0.0 --port "${GATEWAY_PORT_2}"
wait_for "backend" "curl -fs http://127.0.0.1:${BACKEND_PORT}/health/liveliness >/dev/null" 300
wait_for "gateway-1" "curl -fs http://127.0.0.1:${GATEWAY_PORT_1}/health/liveliness >/dev/null" 300
@ -206,6 +220,7 @@ LITELLM_MASTER_KEY=${MASTER_KEY}
REDIS_HOST=127.0.0.1
REDIS_PORT=${REDIS_PORT}
E2E_OTEL_QUERY_URL=http://127.0.0.1:${JAEGER_QUERY_PORT}
E2E_OTEL_EXPORTER_ENDPOINT=https://127.0.0.1:${JAEGER_OTLP_TLS_PORT}
E2E_KEYCLOAK_URL=http://127.0.0.1:${KEYCLOAK_PORT}
E2E_KEYCLOAK_ADMIN_USER=admin
E2E_KEYCLOAK_ADMIN_PASSWORD=e2e-ephemeral-idp-not-a-secret

15
.github/merge-smoke-tests.json vendored Normal file
View file

@ -0,0 +1,15 @@
{
"cases": {
"CHAT-JSON": "tests/test_litellm/llms/openai/test_openai.py::test_acompletion_returns_json_reply_over_injected_transport",
"CHAT-TEXT-STREAM": "tests/test_litellm/llms/openai/test_openai.py::test_acompletion_streams_text_deltas_over_injected_transport",
"CHAT-TOOL-STREAM": "tests/test_litellm/llms/openai/test_openai.py::test_acompletion_streams_tool_call_arguments_over_injected_transport",
"MODEL-ALLOW": "tests/test_litellm/proxy/auth/test_auth_checks.py::test_can_object_call_model_allows_listed_model_for_key",
"MODEL-DENY": "tests/test_litellm/proxy/auth/test_auth_checks.py::test_can_object_call_model_denials_return_forbidden[key-key_model_access_denied]",
"COST-EXPLICIT": "tests/test_litellm/test_cost_calculator.py::test_completion_cost_charges_explicit_per_token_rates_over_registered_ones",
"COST-ZERO": "tests/test_litellm/test_cost_calculator.py::test_completion_cost_is_zero_when_explicit_rates_are_zero",
"LOG-CONTENT-ON": "tests/test_litellm/litellm_core_utils/test_litellm_logging.py::test_standard_logging_payload_keeps_message_content_when_message_logging_is_on",
"LOG-CONTENT-OFF": "tests/test_litellm/litellm_core_utils/test_litellm_logging.py::test_standard_logging_payload_redacts_message_content_when_message_logging_is_off",
"CALLBACK-SUCCESS": "tests/test_litellm/litellm_core_utils/test_litellm_logging.py::test_async_success_handler_delivers_standard_logging_payload_to_custom_logger",
"CALLBACK-FAILURE": "tests/test_litellm/litellm_core_utils/test_litellm_logging.py::test_async_failure_handler_delivers_failure_payload_to_custom_logger"
}
}

13
.github/pull_request_template.md vendored
View file

@ -1,6 +1,8 @@
<!-- The whole description's target audience is humans, not AI agents: write it in plain, simple,
everyday engineering language, extremely parsable and readable at a glance. This goes double for
the TLDR, User Flow, and Caveats sections -->
the TLDR, User Flow, and Caveats sections
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 -->
## TLDR
@ -21,6 +23,7 @@ How it solves it:
<!-- Two ordered lists, Before and After, walking the same end user through the same task, written strictly from that user's seat
Read the linked issue, ticket, or customer thread first so the flow reflects the real application and the routes its users actually hit; don't invent a generic scenario
Lead each list with one plain sentence saying where the flow fails (Before) or succeeds (After), then number the steps
Keep it tight: aim for 3 to 5 steps per list, one line each, roughly 20 words max, and never pad a shorter flow with filler steps to hit the count. Cover the one path the PR changes and fold variants (case, other field, second endpoint) into a clause on the step they belong to rather than their own steps. The example below is the target length
Every step is something the user does or observes: the HTTP method and full URL they hit, what they sent, and what visibly came back (status code, error text, the shape of an ID). UI steps name the page URL and what is on screen
No LiteLLM internals: never name functions, files, DB tables, config classes, hooks, callbacks, or code paths. "The upload hands back an ID that looks like OpenAI's own `file-abc123` instead of the scrambled one the gateway returned" is right, "no managed-file row was registered" is wrong
Keep the two lists step-for-step identical until they diverge, so the changed step is obvious
@ -45,15 +48,15 @@ After: the same request comes back with real token counts, so the dashboard show
## Relevant issues
<!-- e.g., "Fixes #000" -->
<!-- e.g., "Fixes #000". Drop the section if there is none -->
## Affected release
<!-- Only for a fix to a regression in a released or rc version (perf, memory, crash, or behavior): name the version it regressed in, e.g. "regression in v1.100.0" or "since v1.101.0-rc.1", and add the `backport-stable` label so the fix is cherry-picked onto the rc line before the stable is tagged. Leave the section blank otherwise -->
<!-- Only for a fix to a regression in a released or rc version (perf, memory, crash, or behavior): name the version it regressed in, e.g. "regression in v1.100.0" or "since v1.101.0-rc.1", and add the `backport-stable` label so the fix is cherry-picked onto the rc line before the stable is tagged. Drop the section otherwise -->
## Linear ticket
<!-- if you are an internal contributor, add "Resolves " followed by the Linear ticket e.g., "Resolves LIT-1234" to link the Linear ticket to the GitHub PR. If you don't have one, leave the section blank rather than guessing -->
<!-- if you are an internal contributor, add "Resolves " followed by the Linear ticket e.g., "Resolves LIT-1234" to link the Linear ticket to the GitHub PR. If you don't have one, drop the section rather than guessing -->
## Pre-Submission checklist
@ -134,7 +137,7 @@ If you're seeing a delay in your PR being merged, ping the LiteLLM Team on [Slac
human reader
If you assumed something instead of testing it, e.g. "only reproduces with X on" or "no
user-observable behavior difference", list it here too with what breaks if it is wrong
Leave this section empty if there are none -->
Drop this section if there are none -->
## QA runbook

58
.github/scripts/assert_ci_coverage.py vendored
View file

@ -235,9 +235,7 @@ class Slice:
return True # a `-k` this parser cannot model is assumed to claim everything
if any(term.lower() in relative_path.lower() for term in self.excluded):
return False
return not self.required or any(
term.lower() in name.lower() for term in self.required for name in inner_names
)
return not self.required or any(term.lower() in name.lower() for term in self.required for name in inner_names)
def _strings(node: object) -> Iterable[str]:
@ -307,9 +305,7 @@ def _matchable_names(relative_path: str) -> frozenset[str]:
except (OSError, SyntaxError):
return frozenset({relative_path})
return frozenset({relative_path}) | frozenset(
node.name
for node in ast.walk(tree)
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef))
node.name for node in ast.walk(tree) if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef))
)
@ -331,9 +327,7 @@ def _deselected_everywhere(allowlist: Allowlist) -> tuple[Finding, ...]:
slices: Final = _slices()
named_by_workflow: Final = _workflow_named_tokens()
globbed: Final = tuple(
path
for path in _test_files()
if any(_token_covers(glob, path) for slice_ in slices for glob in slice_.globs)
path for path in _test_files() if any(_token_covers(glob, path) for slice_ in slices for glob in slice_.globs)
)
return tuple(
Finding(
@ -363,11 +357,7 @@ def _shard_children(root: str, repo_root: pathlib.Path = REPO_ROOT) -> tuple[str
child.relative_to(repo_root).as_posix()
for child in (repo_root / root).iterdir()
if not child.name.startswith(".")
and (
_holds_tests(child)
if child.is_dir()
else child.name.startswith("test_") and child.suffix == ".py"
)
and (_holds_tests(child) if child.is_dir() else child.name.startswith("test_") and child.suffix == ".py")
)
)
@ -499,13 +489,32 @@ def _check_shards() -> int:
return 0
def _integration_groups(runner: pathlib.Path) -> dict[str, tuple[str, ...]]:
module: Final = ast.parse(runner.read_text())
literal: Final = next(
node.value
for node in module.body
if isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name) and node.target.id == "GROUPS"
)
mapping: Final = literal.args[0] if isinstance(literal, ast.Call) else literal
return {group: tuple(folders) for group, folders in ast.literal_eval(mapping).items()}
def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozenset[str], tuple[Finding, ...]]:
manifest: Final = repo_root / "tests/integration/contracts.json"
if not manifest.exists():
runner: Final = repo_root / "tests/integration/run.py"
if not runner.exists():
return frozenset(), ()
entries: Final = json.loads(manifest.read_text())
paths: Final = frozenset(node.split("::", 1)[0] for node in entries["tests"])
browser_paths: Final = frozenset(node.split("::", 1)[0] for node in entries.get("browser", {}))
groups: Final = _integration_groups(runner)
integration_root: Final = repo_root / "tests/integration"
paths: Final = frozenset(
str(path.relative_to(repo_root))
for folders in groups.values()
for folder in folders
for path in (integration_root / folder).glob("test_*.py")
)
browser_manifest: Final = repo_root / "tests/e2e/ui/tests/integrationCritical/expected.json"
browser_nodes: Final = json.loads(browser_manifest.read_text()) if browser_manifest.exists() else ()
browser_paths: Final = frozenset(node.split("::", 1)[0] for node in browser_nodes)
circle_path: Final = repo_root / ".circleci/config.yml"
circle: Final = yaml.safe_load(circle_path.read_text()) if circle_path.exists() else {}
steps: Final = circle.get("jobs", {}).get("integration_contracts", {}).get("steps", ())
@ -526,15 +535,14 @@ def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozens
)
required: Final = (frozenset({"browser"}) if browser_paths else frozenset()) | frozenset(
group
for group, folders in entries["groups"].items()
for group, folders in groups.items()
if any(any(path.startswith(f"tests/integration/{folder}/") for folder in folders) for path in paths)
)
ungrouped: Final = frozenset(
path
for path in paths
if sum(
any(path.startswith(f"tests/integration/{folder}/") for folder in folders)
for folders in entries["groups"].values()
any(path.startswith(f"tests/integration/{folder}/") for folder in folders) for folders in groups.values()
)
!= 1
)
@ -547,10 +555,6 @@ def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozens
Finding(path, "integration contract is also selected by GitHub Actions")
for path in paths
if any(_token_covers(token, path) for token in gha_tokens)
) + tuple(
Finding(path, "canonical integration test file is missing")
for path in paths
if not (repo_root / path).is_file()
)
browser_commands: Final = tuple(
scalar.value
@ -592,7 +596,7 @@ def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozens
) + tuple(Finding(path, "canonical node must have exactly one integration group") for path in sorted(ungrouped))
if not paths or not invoked or not scheduled:
return frozenset(), findings + (
Finding(str(manifest.relative_to(repo_root)), "dedicated CircleCI runner is missing"),
Finding(str(runner.relative_to(repo_root)), "dedicated CircleCI runner is missing"),
)
return paths | browser_paths, findings + group_findings + browser_findings + exclusion_findings

493
.github/scripts/run_merge_smoke.py vendored Normal file
View file

@ -0,0 +1,493 @@
#!/usr/bin/env python3
"""Merge smoke harness: bounded checks run inside a loopback-only Linux network namespace."""
# ruff: noqa: T201 # CLI harness: stdout/stderr lines are the reported result
from __future__ import annotations
import argparse
import contextlib
import http.client
import json
import os
import secrets
import signal
import socket
import subprocess
import sys
import time
from collections import Counter
from collections.abc import Sequence
from dataclasses import dataclass, field
from pathlib import Path
from types import MappingProxyType
from typing import Final, NoReturn, TextIO, cast
import pytest
EXPECTED_CASES: Final = (
"CHAT-JSON",
"CHAT-TEXT-STREAM",
"CHAT-TOOL-STREAM",
"MODEL-ALLOW",
"MODEL-DENY",
"COST-EXPLICIT",
"COST-ZERO",
"LOG-CONTENT-ON",
"LOG-CONTENT-OFF",
"CALLBACK-SUCCESS",
"CALLBACK-FAILURE",
)
@dataclass(frozen=True, slots=True)
class CheckResult:
ok: bool
detail: str = ""
@dataclass(slots=True)
class _Args:
command: str = ""
no_child: bool = False
expect: str = ""
litellm_bin: str | None = None
lite_bin: str | None = None
diagnostics_dir: str = ""
ready_deadline: float = 120.0
shutdown_deadline: float = 20.0
poll_interval: float = 0.5
manifest: str = ""
rootdir: str | None = None
def fail(reason: str) -> NoReturn:
print(f"merge-smoke: FAIL {reason}", file=sys.stderr)
sys.exit(1)
def ok(step: str) -> None:
print(f"merge-smoke: OK {step}")
def tail(path: Path, lines: int = 20) -> str:
try:
return "\n".join(path.read_text(errors="replace").splitlines()[-lines:])
except OSError as exc:
return f"<cannot read {path}: {exc}>"
def cmd_verify_isolation(args: _Args) -> int:
if os.geteuid() == 0:
fail("verify-isolation must run unprivileged (geteuid()==0)")
try:
socket.create_connection(("192.0.2.1", 9), timeout=3)
except OSError as exc:
print(f"external connect blocked as expected: errno={exc.errno} {exc}")
else:
fail("external TCP connect to 192.0.2.1:9 succeeded; namespace is not isolated")
listener: Final = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
listener.bind(("127.0.0.1", 0))
listener.listen(1)
port: Final = cast(int, listener.getsockname()[1])
client: Final = socket.create_connection(("127.0.0.1", port), timeout=5)
accepted: Final = listener.accept()
accepted[0].close()
client.close()
listener.close()
print(f"loopback connect ok on 127.0.0.1:{port}")
if not args.no_child:
proc: Final = subprocess.run(
[sys.executable, str(Path(__file__).resolve()), "verify-isolation", "--no-child"],
timeout=30,
capture_output=True,
text=True,
)
if proc.returncode != 0:
fail(f"child process did not inherit isolation: {proc.stderr.strip()}")
print("child process inherits isolation")
ok("verify-isolation")
return 0
def cmd_interpreter(args: _Args) -> int:
print(sys.version)
print(sys.executable)
actual: Final = f"{sys.version_info.major}.{sys.version_info.minor}"
if actual != args.expect:
fail(f"interpreter is {actual}, expected {args.expect}")
ok(f"interpreter {actual}")
return 0
def _run_cli(argv: Sequence[str], label: str) -> CheckResult:
try:
proc: Final = subprocess.run(list(argv), timeout=120, capture_output=True, text=True)
except subprocess.TimeoutExpired:
return CheckResult(ok=False, detail=f"{label} timed out after 120s")
sys.stdout.write(proc.stdout)
sys.stderr.write(proc.stderr)
if proc.returncode != 0:
return CheckResult(ok=False, detail=f"{label} exited {proc.returncode}")
return CheckResult(ok=True)
def cmd_cli(args: _Args) -> int:
venv_bin: Final = Path(sys.executable).parent
litellm_bin: Final = Path(args.litellm_bin) if args.litellm_bin else venv_bin / "litellm"
lite_bin: Final = Path(args.lite_bin) if args.lite_bin else venv_bin / "lite"
commands: Final = (
("import litellm", [sys.executable, "-c", "import litellm"]),
("litellm --version", [str(litellm_bin), "--version"]),
("lite version", [str(lite_bin), "version"]),
)
for label, argv in commands:
result = _run_cli(argv, label)
if not result.ok:
fail(result.detail)
ok(label)
return 0
def _free_port() -> int:
sock: Final = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock.bind(("127.0.0.1", 0))
port: Final = cast(int, sock.getsockname()[1])
sock.close()
return port
_CONFIG_TEMPLATE: Final = """model_list:
- model_name: smoke-model
litellm_params:
model: openai/smoke-model
api_base: http://127.0.0.1:9/v1
api_key: synthetic-key
general_settings:
master_key: os.environ/LITELLM_MASTER_KEY
"""
def _listen_inode(port: int) -> str | None:
target: Final = f"{port:04X}"
for table in ("/proc/net/tcp", "/proc/net/tcp6"):
try:
rows = Path(table).read_text().splitlines()[1:]
except OSError:
continue
for row in rows:
cols = row.split()
if len(cols) > 9 and cols[3] == "0A" and cols[1].rsplit(":", 1)[-1] == target:
return cols[9]
return None
def _ancestors(pid: int) -> frozenset[int]:
chain: Final[set[int]] = set()
pending: Final[list[int]] = [pid]
while pending:
current = pending.pop()
if current <= 0 or current in chain:
continue
chain.add(current)
try:
stat = Path(f"/proc/{current}/stat").read_text()
except OSError:
continue
pending.append(int(stat.rpartition(")")[2].split()[1]))
return frozenset(chain)
def _socket_owner_pid(inode: str) -> int | None:
for proc_dir in Path("/proc").iterdir():
if not proc_dir.name.isdigit():
continue
fd_dir = proc_dir / "fd"
try:
for fd in fd_dir.iterdir():
try:
if os.readlink(fd) == f"socket:[{inode}]":
return int(proc_dir.name)
except OSError:
continue
except OSError:
continue
return None
def _verify_port_owner(port: int, proc: subprocess.Popen[bytes]) -> CheckResult:
inode: Final = _listen_inode(port)
if inode is None:
return CheckResult(ok=False, detail=f"no LISTEN socket found for port {port} in /proc/net/tcp")
owner: Final = _socket_owner_pid(inode)
if owner is None:
return CheckResult(ok=False, detail=f"no process owns the listen socket inode {inode} for port {port}")
if owner != proc.pid and proc.pid not in _ancestors(owner):
return CheckResult(
ok=False, detail=f"port {port} owned by pid {owner} outside the launched process group {proc.pid}"
)
if proc.poll() is not None:
return CheckResult(ok=False, detail=f"proxy exited with code {proc.returncode} after readiness")
return CheckResult(ok=True)
def cmd_proxy_startup(args: _Args) -> int:
diagnostics: Final = Path(args.diagnostics_dir)
diagnostics.mkdir(parents=True, exist_ok=True)
venv_bin: Final = Path(sys.executable).parent
litellm_bin: Final = Path(args.litellm_bin) if args.litellm_bin else venv_bin / "litellm"
port: Final = _free_port()
master_key: Final = "sk-smoke-" + secrets.token_hex(16)
config_path: Final = diagnostics / "config.yaml"
config_path.write_text(_CONFIG_TEMPLATE)
log_path: Final = diagnostics / "proxy.log"
result_path: Final = diagnostics / "result.json"
outcome: Final[dict[str, object]] = {
"port": port,
"time_to_ready_s": None,
"shutdown_s": None,
"readiness": None,
"outcome": "failed",
}
log_file: Final = log_path.open("w")
env: Final = {
**os.environ,
"LITELLM_MASTER_KEY": master_key,
"LITELLM_LOCAL_MODEL_COST_MAP": "True",
}
started: Final = time.monotonic()
proc: Final = subprocess.Popen(
[str(litellm_bin), "--config", str(config_path), "--host", "127.0.0.1", "--port", str(port)],
stdout=log_file,
stderr=subprocess.STDOUT,
start_new_session=True,
env=env,
)
body: str | None = None
last_status: int | None = None
while time.monotonic() - started < args.ready_deadline:
if proc.poll() is not None:
log_file.close()
result_path.write_text(json.dumps(outcome))
fail(f"proxy exited early with code {proc.returncode}\n{tail(log_path)}")
try:
conn = http.client.HTTPConnection("127.0.0.1", port, timeout=5)
conn.request("GET", "/health/readiness")
resp = conn.getresponse()
last_status = resp.status
candidate = resp.read().decode()
conn.close()
except (http.client.HTTPException, ConnectionError, OSError):
time.sleep(args.poll_interval)
continue
if last_status == 200:
body = candidate
break
time.sleep(args.poll_interval)
outcome["time_to_ready_s"] = round(time.monotonic() - started, 3)
if body is None:
_terminate(proc, log_file)
result_path.write_text(json.dumps(outcome))
detail = f"last status {last_status}" if last_status is not None else "no response"
fail(f"readiness not reached within {args.ready_deadline}s ({detail})\n{tail(log_path)}")
outcome["readiness"] = body
try:
readiness = cast(object, json.loads(body))
except json.JSONDecodeError:
readiness = None
if readiness != {"status": "healthy", "db": "Not connected"}:
_terminate(proc, log_file)
result_path.write_text(json.dumps(outcome))
fail(f"unexpected readiness body: {body}")
owner_check: Final = _verify_port_owner(port, proc)
if not owner_check.ok:
_terminate(proc, log_file)
result_path.write_text(json.dumps(outcome))
fail(owner_check.detail)
shutdown_started: Final = time.monotonic()
os.killpg(proc.pid, signal.SIGTERM)
try:
proc.wait(timeout=args.shutdown_deadline)
except subprocess.TimeoutExpired:
os.killpg(proc.pid, signal.SIGKILL)
proc.wait(timeout=10)
outcome["shutdown_s"] = round(time.monotonic() - shutdown_started, 3)
log_file.close()
result_path.write_text(json.dumps(outcome))
fail(f"forced kill after {args.shutdown_deadline}s\n{tail(log_path)}")
outcome["shutdown_s"] = round(time.monotonic() - shutdown_started, 3)
try:
os.killpg(proc.pid, 0)
except ProcessLookupError:
pass
else:
os.killpg(proc.pid, signal.SIGKILL)
log_file.close()
result_path.write_text(json.dumps(outcome))
fail("process group survived SIGTERM")
log_file.close()
outcome["outcome"] = "ok"
result_path.write_text(json.dumps(outcome))
ok(f"proxy-startup ready={outcome['time_to_ready_s']}s shutdown={outcome['shutdown_s']}s")
return 0
def _terminate(proc: subprocess.Popen[bytes], log_file: TextIO) -> None:
with contextlib.suppress(ProcessLookupError):
os.killpg(proc.pid, signal.SIGTERM)
try:
proc.wait(timeout=10)
except subprocess.TimeoutExpired:
with contextlib.suppress(ProcessLookupError):
os.killpg(proc.pid, signal.SIGKILL)
with contextlib.suppress(subprocess.TimeoutExpired):
proc.wait(timeout=10)
log_file.close()
def _load_manifest(path: Path) -> MappingProxyType[str, str]:
def no_duplicates(pairs: list[tuple[object, object]]) -> dict[object, object]:
seen: dict[object, object] = {}
for key, value in pairs:
if key in seen:
raise ValueError(f"duplicate key in manifest: {key}")
seen[key] = value
return seen
raw_value: object = cast(object, json.loads(path.read_text(), object_pairs_hook=no_duplicates))
if not isinstance(raw_value, dict):
raise ValueError("manifest must be an object")
loaded: Final = cast(dict[object, object], raw_value)
cases_value: object = loaded.get("cases")
if not isinstance(cases_value, dict):
raise ValueError("manifest must be an object with a 'cases' object")
cases_any: Final = cast(dict[object, object], cases_value)
cases: Final = {k: v for k, v in cases_any.items() if isinstance(k, str) and isinstance(v, str)}
if len(cases) != len(cases_any):
raise ValueError("manifest 'cases' must map string ids to string node ids")
return MappingProxyType(cases)
@dataclass(slots=True, eq=False)
class _Recorder:
collect_failed: list[str] = field(default_factory=list)
collected: tuple[str, ...] = ()
reports: dict[str, list[tuple[str, str, bool]]] = field(default_factory=dict)
def pytest_collectreport(self, report: pytest.CollectReport) -> None:
if report.failed:
self.collect_failed.append(report.nodeid)
def pytest_collection_finish(self, session: pytest.Session) -> None:
self.collected = tuple(item.nodeid for item in session.items)
def pytest_runtest_logreport(self, report: pytest.TestReport) -> None:
self.reports.setdefault(report.nodeid, []).append((report.when, report.outcome, hasattr(report, "wasxfail")))
def cmd_pytest(args: _Args) -> int:
try:
cases: Final = _load_manifest(Path(args.manifest))
except (OSError, ValueError, json.JSONDecodeError) as exc:
fail(f"manifest invalid: {exc}")
if tuple(cases) != EXPECTED_CASES:
fail(f"manifest case ids must be exactly {list(EXPECTED_CASES)} in order, got {list(cases)}")
node_ids: Final = tuple(cases.values())
if len(set(node_ids)) != len(node_ids):
fail("manifest node ids are not unique")
argv: Final = [
*node_ids,
"-p",
"no:cacheprovider",
"-p",
"no:xdist",
"-p",
"no:rerunfailures",
"-p",
"no:randomly",
"-rA",
"-q",
*(["--rootdir", args.rootdir] if args.rootdir else []),
]
recorder: Final = _Recorder()
code: Final = pytest.main(argv, plugins=[recorder])
name_of: Final = MappingProxyType({node_id: case_id for case_id, node_id in cases.items()})
problems: Final[list[str]] = []
if code != 0:
problems.append(f"pytest exit code {code}")
for failed_id in recorder.collect_failed:
problems.append(f"collection failed: {name_of.get(failed_id, failed_id)}")
expected: Final = Counter(node_ids)
collected: Final = Counter(recorder.collected)
for node_id in expected - collected:
problems.append(f"missing case {name_of[node_id]} ({node_id})")
for node_id in collected - expected:
problems.append(f"unexpected test collected: {node_id}")
for node_id, count in collected.items():
if count > 1:
problems.append(f"duplicated test id: {node_id}")
if len(recorder.collected) != len(EXPECTED_CASES):
problems.append(f"collected {len(recorder.collected)} tests, expected {len(EXPECTED_CASES)}")
rows: Final[list[tuple[str, bool]]] = []
for case_id, node_id in cases.items():
reports = recorder.reports.get(node_id, [])
case_ok = (
bool(reports)
and all(outcome == "passed" and not wasxfail for _, outcome, wasxfail in reports)
and {when for when, _, _ in reports} >= {"setup", "call", "teardown"}
)
rows.append((case_id, case_ok))
if not reports:
problems.append(f"{case_id} ({node_id}) produced no runtest reports")
continue
for when, outcome, wasxfail in reports:
if outcome != "passed":
problems.append(f"{case_id} ({node_id}) {when} outcome={outcome}")
if wasxfail:
problems.append(f"{case_id} ({node_id}) {when} was xfail/xpass")
missing_phases = {"setup", "call", "teardown"} - {when for when, _, _ in reports}
for phase in sorted(missing_phases):
problems.append(f"{case_id} ({node_id}) missing {phase} report")
for case_id, passed in rows:
print(f"{case_id} {'PASS' if passed else 'FAIL'} {cases[case_id]}")
if problems:
for problem in problems:
print(f"merge-smoke: {problem}", file=sys.stderr)
fail("pytest verdict failed")
ok("pytest 11 cases")
return 0
def main() -> int:
parser: Final = argparse.ArgumentParser(description=__doc__)
subs: Final = parser.add_subparsers(dest="command", required=True)
p_iso: Final = subs.add_parser("verify-isolation")
p_iso.add_argument("--no-child", action="store_true")
p_interp: Final = subs.add_parser("interpreter")
p_interp.add_argument("--expect", required=True)
p_cli: Final = subs.add_parser("cli")
p_cli.add_argument("--litellm-bin", default=None)
p_cli.add_argument("--lite-bin", default=None)
p_proxy: Final = subs.add_parser("proxy-startup")
p_proxy.add_argument("--diagnostics-dir", required=True)
p_proxy.add_argument("--litellm-bin", default=None)
p_proxy.add_argument("--ready-deadline", type=float, default=120)
p_proxy.add_argument("--shutdown-deadline", type=float, default=20)
p_proxy.add_argument("--poll-interval", type=float, default=0.5)
p_test: Final = subs.add_parser("pytest")
p_test.add_argument("--manifest", required=True)
p_test.add_argument("--rootdir", default=None)
args: Final = parser.parse_args(namespace=_Args())
handlers: Final = {
"verify-isolation": cmd_verify_isolation,
"interpreter": cmd_interpreter,
"cli": cmd_cli,
"proxy-startup": cmd_proxy_startup,
"pytest": cmd_pytest,
}
return handlers[args.command](args)
if __name__ == "__main__":
sys.exit(main())

2
.github/scripts/verify_linux_native_wheel.py vendored
View file

@ -214,7 +214,7 @@ def main(
native_module: Final = load_native_module(native_path)
native_module_loads: Final = native_module is not None
panic_test_hook_absent: Final = native_module is not None and not hasattr(native_module, "_panic_for_test")
native_size_limit: Final = 35_000_000
native_size_limit: Final = 40_000_000
native_size_within_limit: Final = native_member.file_size <= native_size_limit
validations: Final = (
(f"Python tag is {EXPECTED_PYTHON_TAG}", python_tag == EXPECTED_PYTHON_TAG),

32
.github/workflows/_test-unit-base.yml vendored
View file

@ -4,7 +4,13 @@ on:
workflow_call:
inputs:
test-path:
description: "Pytest path(s) to run"
description: >-
Space-separated pytest paths to run. A path that no longer exists is
dropped with a warning instead of being passed to pytest, because one
missing path makes pytest-xdist collect nothing and report exit 5, which
the step treats as a drained shard. Options are passed through as
written, so use the `--flag=value` form: a bare `--ignore path` would
have its path existence-checked like any other token.
required: true
type: string
workers:
@ -165,14 +171,22 @@ jobs:
DIST: ${{ inputs.dist }}
COVERAGE_CORE: sysmon
run: |
found_path=false
for path in ${TEST_PATH}; do
if [ -e "${path%%::*}" ]; then
found_path=true
break
fi
pytest_args=()
existing_paths=0
for token in ${TEST_PATH:?}; do
case "${token}" in
-*) pytest_args+=("${token}") ;;
*)
if [ -e "${token%%::*}" ]; then
pytest_args+=("${token}")
existing_paths=$((existing_paths + 1))
else
echo "::warning::${token} does not exist; drop it from this shard's test-path"
fi
;;
esac
done
if [ "$found_path" = false ]; then
if [ "${existing_paths}" -eq 0 ]; then
echo "No path in TEST_PATH exists (${TEST_PATH}); nothing to run"
exit 0
fi
@ -181,7 +195,7 @@ jobs:
xdist_args=(-n "${WORKERS}" --dist="${DIST}")
fi
set +e
uv run --no-sync pytest ${TEST_PATH:?} \
uv run --no-sync pytest "${pytest_args[@]}" \
--tb=short -vv \
--maxfail="${MAX_FAILURES}" \
"${xdist_args[@]}" \

2
.github/workflows/check-ui-api-types.yml vendored
View file

@ -4,8 +4,6 @@ on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
permissions:

3
.github/workflows/ci-coverage.yml vendored
View file

@ -4,13 +4,10 @@ on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
push:
branches:
- main
- litellm_internal_staging
permissions:
contents: read

8
.github/workflows/codeql.yml vendored
View file

@ -67,12 +67,18 @@ jobs:
# further up the stack are modified. The suppression is scoped to this one
# file/rule pair via SARIF post-filtering so every other callsite of
# py/weak-sensitive-data-hashing in the repository continues to be analyzed.
- name: Filter SARIF (OCI sha256)
# The same query fires on the HIBP k-anonymity lookup in
# litellm/proxy/auth/password_policy.py, where the password's SHA-1 is only
# a lookup key into the haveibeenpwned range API (the protocol mandates
# SHA-1) and the digest itself never leaves the proxy beyond its first 5
# characters.
- name: Filter SARIF (OCI sha256, HIBP sha1)
if: matrix.language == 'python'
uses: advanced-security/filter-sarif@2da736ff05ef065cb2894ac6892e47b5eac2c3c0 # v1.1
with:
patterns: |
-litellm/llms/oci/common_utils.py:py/weak-sensitive-data-hashing
-litellm/proxy/auth/password_policy.py:py/weak-sensitive-data-hashing
input: sarif-results/python.sarif
output: sarif-results/python.sarif

2
.github/workflows/codspeed.yml vendored
View file

@ -4,7 +4,6 @@ on:
push:
branches:
- main
- litellm_internal_staging
paths:
- "litellm/**"
- "tests/benchmarks/**"
@ -17,7 +16,6 @@ on:
pull_request:
branches:
- main
- litellm_internal_staging
paths:
- "litellm/**"
- "tests/benchmarks/**"

33
.github/workflows/compat-matrix-image.yml vendored Normal file
View file

@ -0,0 +1,33 @@
name: Compat Matrix Image
on:
pull_request:
paths:
- tests/e2e/claude_code/cron_vm/**
- .github/workflows/compat-matrix-image.yml
workflow_dispatch:
permissions: {}
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
cancel-in-progress: true
jobs:
compat-matrix-image:
name: compat-matrix-image
runs-on: ubuntu-latest
timeout-minutes: 15
permissions:
contents: read
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Build the Render cron image
run: docker build -f tests/e2e/claude_code/cron_vm/Dockerfile -t compat-matrix:${{ github.sha }} tests/e2e
- name: Run the pinned binaries as the cron user
run: |
docker run --rm compat-matrix:${{ github.sha }} bash -c 'set -e; whoami; claude --version; gh --version; uv --version'

2
.github/workflows/cost-map-guard.yml vendored
View file

@ -4,8 +4,6 @@ on: # zizmor: ignore[dangerous-triggers] runs the base branch's code only; the P
pull_request_target:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
permissions:

186
.github/workflows/create-release.yml vendored
View file

@ -1,186 +0,0 @@
name: Create Release
on:
workflow_dispatch:
inputs:
tag:
description: "Release tag (e.g. 1.84.0, 1.84.0rc1, 1.84.0.dev42, 1.84.0-dev.2, 1.84.0.post1; legacy v1.83.10-stable still accepted)"
required: true
type: string
commit_hash:
description: "Full 40-char commit SHA to target"
required: true
type: string
permissions: {}
jobs:
release:
name: Create Release
runs-on: ubuntu-latest
permissions:
contents: write
steps:
- name: Validate inputs
env:
TAG: ${{ inputs.tag }}
COMMIT_HASH: ${{ inputs.commit_hash }}
run: |
if ! echo "${COMMIT_HASH}" | grep -qE '^[0-9a-f]{40}$'; then
echo "::error::commit_hash must be a full 40-character commit SHA"
exit 1
fi
if ! echo "${TAG}" | grep -qE '^v?[0-9]+\.[0-9]+\.[0-9]+'; then
echo "::error::tag must start with X.Y.Z (optional leading v), e.g. 1.84.0, 1.84.0rc1, 1.84.0.dev42, or v1.83.10-stable"
exit 1
fi
- name: Create release
env:
TAG: ${{ inputs.tag }}
COMMIT_HASH: ${{ inputs.commit_hash }}
uses: actions/github-script@60a0d83039c74a4aee543508d2ffcb1c3799cdea # v7.0.1
with:
script: |
const tag = process.env.TAG;
const commitHash = process.env.COMMIT_HASH;
// Mark RC / dev / nightly / alpha / beta tags as GitHub pre-releases.
// Accept both PEP 440 (`.dev`) and SemVer (`-dev`) separators so tags
// like `1.84.0.dev2` and `1.84.0-dev.2` are both detected.
// PEP 440 post-releases (e.g. `1.84.0.post1`) and legacy `-stable[.patch.N]`
// are stable maintenance releases, not pre-releases.
const isPrerelease = /(?:rc|nightly|alpha|beta|[-.]dev)/i.test(tag);
// A stable release should only claim the repo "latest" badge when its
// version is >= the current latest. Otherwise a backport (e.g. 1.84.6)
// would steal "latest" from a newer line (e.g. 1.88.1).
const versionKey = (rawTag) => {
const m = String(rawTag).match(/^v?(\d+)\.(\d+)\.(\d+)/);
if (!m) return null;
const maintenance = String(rawTag).match(/(?:\.post|\.patch\.)(\d+)/i);
return [Number(m[1]), Number(m[2]), Number(m[3]), maintenance ? Number(maintenance[1]) : 0];
};
const isAtLeast = (a, b) => {
for (let i = 0; i < a.length; i++) {
if (a[i] !== b[i]) return a[i] > b[i];
}
return true;
};
const cosignSection = [
`## Verify Docker Image Signature`,
``,
`All LiteLLM Docker images are signed with [cosign](https://docs.sigstore.dev/cosign/overview/). Every release is signed with the same key introduced in [commit \`0112e53\`](https://github.com/BerriAI/litellm/commit/0112e53046018d726492c814b3644b7d376029d0).`,
``,
`**Verify using the pinned commit hash (recommended):**`,
``,
`A commit hash is cryptographically immutable, so this is the strongest way to ensure you are using the original signing key:`,
``,
'```bash',
`cosign verify \\`,
` --key https://raw.githubusercontent.com/BerriAI/litellm/0112e53046018d726492c814b3644b7d376029d0/cosign.pub \\`,
` ghcr.io/berriai/litellm:${tag}`,
'```',
``,
`**Verify using the release tag (convenience):**`,
``,
`Tags are protected in this repository and resolve to the same key. This option is easier to read but relies on tag protection rules:`,
``,
'```bash',
`cosign verify \\`,
` --key https://raw.githubusercontent.com/BerriAI/litellm/${tag}/cosign.pub \\`,
` ghcr.io/berriai/litellm:${tag}`,
'```',
``,
`Expected output:`,
``,
'```',
`The following checks were performed on each of these signatures:`,
` - The cosign claims were validated`,
` - The signatures were verified against the specified public key`,
'```',
``,
`---`,
``,
].join('\n');
try {
let makeLatest = "false";
const newVersion = versionKey(tag);
if (!isPrerelease && newVersion) {
let latestVersion = null;
try {
const latest = await github.rest.repos.getLatestRelease({
owner: context.repo.owner,
repo: context.repo.repo,
});
latestVersion = versionKey(latest.data.tag_name);
} catch (error) {
if (error.status !== 404) throw error;
}
makeLatest = (!latestVersion || isAtLeast(newVersion, latestVersion)) ? "true" : "false";
}
try {
await github.rest.git.createRef({
owner: context.repo.owner,
repo: context.repo.repo,
ref: `refs/tags/${tag}`,
sha: commitHash,
});
} catch (error) {
if (error.status !== 422) throw error;
const existing = await github.rest.git.getRef({
owner: context.repo.owner,
repo: context.repo.repo,
ref: `tags/${tag}`,
});
if (existing.data.object.sha !== commitHash) {
throw new Error(`Tag ${tag} already exists at ${existing.data.object.sha}, expected ${commitHash}`);
}
}
const response = await github.rest.repos.createRelease({
draft: true,
generate_release_notes: true,
name: tag,
owner: context.repo.owner,
prerelease: isPrerelease,
repo: context.repo.repo,
tag_name: tag,
});
const updatedBody = cosignSection + (response.data.body ?? '');
await github.rest.repos.updateRelease({
owner: context.repo.owner,
repo: context.repo.repo,
release_id: response.data.id,
tag_name: tag,
body: updatedBody,
draft: false,
});
if (!isPrerelease) {
await github.rest.repos.updateRelease({
owner: context.repo.owner,
repo: context.repo.repo,
release_id: response.data.id,
tag_name: tag,
make_latest: makeLatest,
});
}
} catch (error) {
core.setFailed(error.message);
}
create-branch:
name: Create Release Branch
needs: release
permissions:
contents: write
uses: ./.github/workflows/create-release-branch.yml
with:
tag: ${{ inputs.tag }}
commit_hash: ${{ inputs.commit_hash }}

49
.github/workflows/create_daily_staging_branch.yml vendored
View file

@ -1,49 +0,0 @@
name: Create Daily Staging Branch
on:
schedule:
- cron: "0 0,12 * * *" # Runs every 12 hours at midnight and noon UTC
workflow_dispatch: # Allow manual trigger
jobs:
create-staging-branch:
if: github.repository == 'BerriAI/litellm'
runs-on: ubuntu-latest
permissions:
contents: write
steps:
- name: Create daily staging branch
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
run: |
BRANCH_NAME="litellm_oss_staging_$(date +'%m_%d_%Y')"
echo "Creating branch: $BRANCH_NAME"
if gh api "repos/${{ github.repository }}/git/ref/heads/$BRANCH_NAME" --silent 2>/dev/null; then
echo "Branch $BRANCH_NAME already exists. Skipping creation."
exit 0
fi
MAIN_SHA=$(gh api "repos/${{ github.repository }}/git/ref/heads/main" --jq '.object.sha')
gh api "repos/${{ github.repository }}/git/refs" -f ref="refs/heads/$BRANCH_NAME" -f sha="$MAIN_SHA" --silent
echo "Successfully created branch: $BRANCH_NAME at $MAIN_SHA"
create-internal-dev-branch:
if: github.repository == 'BerriAI/litellm'
runs-on: ubuntu-latest
permissions:
contents: write
steps:
- name: Create internal dev branch
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
run: |
BRANCH_NAME="litellm_internal_dev_$(date +'%m_%d_%Y')"
echo "Creating branch: $BRANCH_NAME"
if gh api "repos/${{ github.repository }}/git/ref/heads/$BRANCH_NAME" --silent 2>/dev/null; then
echo "Branch $BRANCH_NAME already exists. Skipping creation."
exit 0
fi
MAIN_SHA=$(gh api "repos/${{ github.repository }}/git/ref/heads/main" --jq '.object.sha')
gh api "repos/${{ github.repository }}/git/refs" -f ref="refs/heads/$BRANCH_NAME" -f sha="$MAIN_SHA" --silent
echo "Successfully created branch: $BRANCH_NAME at $MAIN_SHA"

2
.github/workflows/guard-fork-dependencies.yml vendored
View file

@ -4,8 +4,6 @@ on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
paths:
- "uv.lock"

1
.github/workflows/image-scan.yml vendored
View file

@ -4,7 +4,6 @@ on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_branch
- "litellm_**"
paths:

20
.github/workflows/issue_fixed_comment.yml vendored
View file

@ -6,8 +6,12 @@ on:
workflow_dispatch:
inputs:
issue_number:
description: "Closed issue number to comment on manually."
required: true
description: "Closed issue number to comment on and close the superseded pull requests of. Ignored by a sweep."
required: false
sweep:
description: "Close every open pull request whose linked issues were all fixed on the default branch. Reads every open pull request, so run it at most once an hour."
type: boolean
default: false
pull_request:
paths:
- .github/workflows/issue_fixed_comment.yml
@ -39,16 +43,17 @@ jobs:
with:
bun-version: "1.4.0"
- name: Test the closer lookup, the release placement and the comment
- name: Test the closer lookup, the release placement, the comment and the superseded pull request close
run: bun test scripts/comment-fixed-issue.test.ts
comment-fixed-issue:
if: github.event_name != 'pull_request' && github.repository == 'BerriAI/litellm'
runs-on: ubuntu-latest
timeout-minutes: 5
timeout-minutes: 15
permissions:
contents: read
issues: write
pull-requests: write
steps:
- name: Checkout scripts
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
@ -59,13 +64,16 @@ jobs:
- name: Setup Bun
uses: oven-sh/setup-bun@0c5077e51419868618aeaa5fe8019c62421857d6 # v2.2.0
with:
# Exact version, never latest: the next step holds an issues: write token
# Exact version, never latest: the next step holds issues: write and pull-requests: write tokens
bun-version: "1.4.0"
- name: Name the release that carries the fix
- name: Name the release that carries the fix and close the pull requests it supersedes
shell: bash
run: bun run scripts/comment-fixed-issue.ts | tee -a "${GITHUB_STEP_SUMMARY}"
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
ISSUE_NUMBER: ${{ github.event.issue.number || github.event.inputs.issue_number }}
SWEEP: ${{ github.event.inputs.sweep }}
DEFAULT_BRANCH: ${{ github.event.repository.default_branch }}
DRY_RUN: ${{ vars.ISSUE_FIXED_COMMENT_ENABLED != 'true' }}
CLOSE_PRS_DRY_RUN: ${{ vars.ISSUE_FIXED_CLOSE_PRS_ENABLED != 'true' }}

2
.github/workflows/osv-scan.yml vendored
View file

@ -4,8 +4,6 @@ on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
schedule:
- cron: "23 6 * * *"

3
.github/workflows/publish-basedpyright-base-counts.yml vendored
View file

@ -1,6 +1,6 @@
name: Publish basedpyright base counts
# Every commit on main or litellm_internal_staging can become a future merge-base.
# Every commit on main can become a future merge-base.
# Publishing its per-rule basedpyright counts as an artifact lets
# scripts/type_check_gate.py download them in seconds instead of paying a
# 60-110s second basedpyright pass on every fresh worktree or moved merge-base.
@ -11,7 +11,6 @@ on:
push:
branches:
- main
- litellm_internal_staging
workflow_dispatch:
inputs:
ref:

6
.github/workflows/test-code-quality.yml vendored
View file

@ -4,13 +4,10 @@ on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
push:
branches:
- main
- litellm_internal_staging
permissions:
contents: read
@ -83,6 +80,9 @@ jobs:
- name: test_e2e_changed_gate
run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_e2e_changed_gate.py tests/code_coverage_tests/test_e2e_idp_stack.py
- name: Check merge smoke harness
run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_merge_smoke.py
- name: router_code_coverage
run: uv run --no-sync python ./tests/code_coverage_tests/router_code_coverage.py

2
.github/workflows/test-linting.yml vendored
View file

@ -4,8 +4,6 @@ on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
permissions:

2
.github/workflows/test-litellm-ui-build.yml vendored
View file

@ -7,8 +7,6 @@ on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
concurrency:

2
.github/workflows/test-litellm-ui-lint.yml vendored
View file

@ -6,8 +6,6 @@ on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
concurrency:

3
.github/workflows/test-litellm-ui-unit.yml vendored
View file

@ -7,13 +7,10 @@ on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
push:
branches:
- main
- litellm_internal_staging
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}

2
.github/workflows/test-mcp-dependency-resolution.yml vendored
View file

@ -4,8 +4,6 @@ on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
permissions:

95
.github/workflows/test-merge-smoke.yml vendored Normal file
View file

@ -0,0 +1,95 @@
name: Merge smoke checks
on:
pull_request:
branches: [main, litellm_internal_staging, litellm_oss_staging, "litellm_**"]
workflow_dispatch:
permissions:
contents: read
concurrency:
group: merge-smoke-${{ github.event.pull_request.number || github.run_id }}
cancel-in-progress: true
jobs:
dashboard-build:
name: Dashboard build
runs-on: ubuntu-24.04
timeout-minutes: 30
steps:
- name: Checkout
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Build the dashboard stage
run: docker build --target ui-builder -f Dockerfile .
core-checks:
name: Core checks (Python ${{ matrix.python-version }})
runs-on: ubuntu-24.04
timeout-minutes: 30
strategy:
fail-fast: false
matrix:
python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
env:
LITELLM_LOCAL_MODEL_COST_MAP: "True"
steps:
- name: Checkout
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: ${{ matrix.python-version }}
- name: Set up uv
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Install dependencies
run: .github/scripts/uv_sync_with_retries.sh --frozen --extra proxy --extra cli --group dev --group proxy-dev --python ${{ matrix.python-version }}
- name: Create the loopback-only network namespace
run: |
sudo ip netns add smoke
sudo ip netns exec smoke ip link set lo up
cat > "${RUNNER_TEMP}/in-netns" <<'WRAP'
#!/usr/bin/env bash
set -euo pipefail
exec sudo --preserve-env=LITELLM_LOCAL_MODEL_COST_MAP ip netns exec smoke setpriv --reuid "$(id -u)" --regid "$(id -g)" --init-groups -- env HOME="${HOME}" PATH="${PATH}" "$@"
WRAP
chmod +x "${RUNNER_TEMP}/in-netns"
echo "IN_NETNS=${RUNNER_TEMP}/in-netns" >> "${GITHUB_ENV}"
- name: Verify namespace isolation
run: $IN_NETNS .venv/bin/python .github/scripts/run_merge_smoke.py verify-isolation
- name: Verify interpreter version
run: $IN_NETNS .venv/bin/python .github/scripts/run_merge_smoke.py interpreter --expect ${{ matrix.python-version }}
- name: Import and CLI checks
run: $IN_NETNS .venv/bin/python .github/scripts/run_merge_smoke.py cli
- name: Proxy startup check
run: $IN_NETNS .venv/bin/python .github/scripts/run_merge_smoke.py proxy-startup --diagnostics-dir "${RUNNER_TEMP}/smoke-diagnostics"
- name: Run curated smoke cases
run: $IN_NETNS .venv/bin/python .github/scripts/run_merge_smoke.py pytest --manifest .github/merge-smoke-tests.json
- name: Upload smoke diagnostics
if: always()
uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1
with:
name: merge-smoke-diagnostics-py${{ matrix.python-version }}
path: ${{ runner.temp }}/smoke-diagnostics
if-no-files-found: ignore
- name: Remove the network namespace
if: always()
run: sudo ip netns delete smoke

3
.github/workflows/test-postgres.yml vendored
View file

@ -4,13 +4,10 @@ on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
push:
branches:
- main
- litellm_internal_staging
workflow_dispatch:
permissions:

31
.github/workflows/test-redis-compat.yml vendored
View file

@ -4,13 +4,12 @@ on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
paths:
- "litellm/_redis.py"
- "litellm/_redis_credential_provider.py"
- "tests/test_litellm/test_redis.py"
- "tests/local_testing/test_caching.py"
- "tests/test_litellm/caching/test_redis_connection_pool.py"
- ".github/workflows/test-redis-compat.yml"
- "pyproject.toml"
@ -28,6 +27,9 @@ jobs:
name: "redis-py ${{ matrix.redis-version }}"
runs-on: ubuntu-latest
timeout-minutes: 15
permissions:
contents: read
id-token: write
strategy:
fail-fast: false
@ -57,7 +59,7 @@ jobs:
- name: Install dependencies
run: |
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra extra_proxy --extra semantic-router
- name: Pin redis-py to the matrix version
env:
@ -66,12 +68,33 @@ jobs:
uv pip install "redis==${REDIS_VERSION:?}"
uv run --no-sync python -c "import redis; assert redis.__version__ == '${REDIS_VERSION:?}', redis.__version__; print('redis-py', redis.__version__)"
- name: Build Redis for cluster authentication tests
run: |
curl --fail --location --retry 3 https://download.redis.io/releases/redis-7.2.16.tar.gz -o "$RUNNER_TEMP/redis-7.2.16.tar.gz"
echo "960a8ec15e34ff40e57ff16837b26b33bd81f2da6d24497bb63de532a323a18e $RUNNER_TEMP/redis-7.2.16.tar.gz" | sha256sum --check
tar -xzf "$RUNNER_TEMP/redis-7.2.16.tar.gz" -C "$RUNNER_TEMP"
make -C "$RUNNER_TEMP/redis-7.2.16" -j2 MALLOC=libc OPTIMIZATION=-O1 redis-server
echo "$RUNNER_TEMP/redis-7.2.16/src" >> "$GITHUB_PATH"
- name: Run redis unit tests
run: |
redis-server --version
uv run --no-sync pytest \
tests/test_litellm/test_redis.py \
tests/test_litellm/caching/test_redis_connection_pool.py \
tests/local_testing/test_caching.py::test_sync_cluster_authenticates_with_azure_credentials \
tests/local_testing/test_caching.py::test_sync_cluster_authenticates_with_gcp_credentials \
--tb=short -vv \
--reruns 2 \
--reruns-delay 1 \
--durations=20
--durations=20 \
--cov=./litellm --cov-report=xml:coverage-redis.xml
- name: Upload Redis coverage
if: matrix.redis-version == '5.3.1'
uses: codecov/codecov-action@75cd11691c0faa626561e295848008c8a7dddffe # v5.5.4
with:
use_oidc: true
files: coverage-redis.xml
flags: redis-compat
fail_ci_if_error: false

12
.github/workflows/test-rust.yml vendored
View file

@ -29,8 +29,6 @@ on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
paths:
- "litellm-rust/**"
@ -105,6 +103,16 @@ jobs:
with:
python-version: "3.12"
- uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Install Python dependencies for the bridge tests
working-directory: .
run: |
uv sync --frozen --no-install-project
echo "PYTHONPATH=$PWD/.venv/lib/$(ls .venv/lib)/site-packages" >> "$GITHUB_ENV"
- run: rustup toolchain install --no-self-update
- uses: taiki-e/install-action@d438492cf8a250514fa2d34b30bc3c0dc37c65ff # v2.87.8

2
.github/workflows/test-semgrep.yml vendored
View file

@ -4,8 +4,6 @@ on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
permissions:

2
.github/workflows/test-terraform-modules.yml vendored
View file

@ -9,8 +9,6 @@ on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
paths:
- "terraform/litellm/aws/**"

2
.github/workflows/test-terraform-provider.yml vendored
View file

@ -8,8 +8,6 @@ on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
paths:
- "terraform/provider/**"

3
.github/workflows/test-unit-documentation.yml vendored
View file

@ -4,13 +4,10 @@ on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
push:
branches:
- main
- litellm_internal_staging
permissions:
contents: read

3
.github/workflows/test-unit-proxy-db.yml vendored
View file

@ -4,13 +4,10 @@ on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
push:
branches:
- main
- litellm_internal_staging
permissions:
contents: read

12
.github/workflows/test-unit.yml vendored
View file

@ -4,13 +4,10 @@ on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
push:
branches:
- main
- litellm_internal_staging
workflow_dispatch:
permissions:
@ -107,26 +104,19 @@ jobs:
tests/test_litellm/batches
tests/test_litellm/secret_managers
tests/test_litellm/a2a_protocol
tests/test_litellm/anthropic_interface
tests/test_litellm/chat_completions
tests/test_litellm/completion_extras
tests/test_litellm/compression
tests/test_litellm/containers
tests/test_litellm/endpoints
tests/test_litellm/models
tests/test_litellm/repositories
tests/test_litellm/images
tests/test_litellm/interactions
tests/test_litellm/messages
tests/test_litellm/embeddings
tests/test_litellm/ocr
tests/test_litellm/passthrough
tests/test_litellm/rag
tests/test_litellm/realtime_api
tests/test_litellm/rerank_api
tests/test_litellm/rust_bridge
tests/test_litellm/sandbox
tests/test_litellm/skills
tests/test_litellm/test_router
tests/test_litellm/vector_stores
tests/test_litellm/videos
tests/test_litellm/test_*.py

2
.github/workflows/test-vscode-extension.yml vendored
View file

@ -6,8 +6,6 @@ on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
paths:
- "vscode-extension/**"

4
.github/workflows/zizmor.yml vendored
View file

@ -2,12 +2,10 @@ name: GitHub Actions Security Analysis
on:
push:
branches: [main, litellm_internal_staging]
branches: [main]
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
concurrency:

4
AGENTS.md
View file

@ -33,11 +33,11 @@ End-to-end tests belong in `tests/e2e/` and must follow the harness conventions
When creating PRs, target the repository's current default branch for both internal and external / OSS contributions. Check it with `python3 scripts/default_branch.py --branch` instead of assuming a branch name or relying on cached `origin/HEAD`
When writing a PR body, treat the comments and imperative instructions inside .github/pull_request_template.md as rules to follow, not just layout. Agent harnesses may strip HTML comments from copies of that file injected into context, so read .github/pull_request_template.md from disk before writing a PR body to make sure you see every comment rule
When writing a PR body, treat the comments and imperative instructions inside .github/pull_request_template.md as rules to follow, not just layout. Agent harnesses may strip HTML comments from copies of that file injected into context, so read .github/pull_request_template.md from disk before writing a PR body to make sure you see every comment rule. A section you have nothing to put in (Relevant issues, Affected release, Linear ticket, Caveats, QA runbook, and so on) is removed entirely, heading included, never left as an empty title
Same applies for filing bug reports and feature requests, with .github/ISSUE_TEMPLATE/bug_report.yml and .github/ISSUE_TEMPLATE/feature_request.yml, respectively
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 leave the section blank
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
Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We prefer curl'ing a live proxy instance running on localhost:4000 (I like to run it with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --reload --use_v2_migration_resolver 2>&1 | tee litellm.log`; the Admin UI dev server is `npm run dev` in `ui/litellm-dashboard`, served on port 3000) and showing both the command run and the output. Also, it should hit real LLM provider APIs, not mocks, and cost real $$$ because that is the most realistic test. The proof of fix should be exactly what the end user / customer would see / do. The run logs in PR #27703 is a prime example of how to do it (not a huge fan of using a python test script that future me and the team will have no visibility into; I prefer just curl commands or a short list of bash commands (e.g., using `for`)). If it's a UI thing, or the main use case runs through a headful agentic coding tool like Claude Code or Codex, drive that surface yourself and embed your own before and after screenshots of it in the PR (the Admin UI page, or what the coding tool shows), next to an ordered list of the URLs to go to (e.g., http://localhost:4000/ui/?page=logs), where to click, and what fields to fill out so a reviewer can reproduce it

4
Dockerfile
View file

@ -1,10 +1,10 @@
# syntax=docker/dockerfile:1.7
# Base image for building
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d
# Runtime image
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
# Pinned by digest like the other base images; bump explicitly on Node upgrades.
ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43

4
backend/Dockerfile
View file

@ -1,5 +1,5 @@
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
FROM $UV_IMAGE AS uvbin

1
backend/routes/allowlist.py
View file

@ -26,6 +26,7 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = (
"/v2/login",
"/v3/login",
"/logout",
"/session/logout",
"/token",
"/onboarding/",
"/audit",

4
docker/Dockerfile.database
View file

@ -1,10 +1,10 @@
# syntax=docker/dockerfile:1.7
# Base image for building
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d
# Runtime image
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
# Pinned by digest like the other base images; bump explicitly on Node upgrades.
ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43

4
docker/Dockerfile.non_root
View file

@ -1,8 +1,8 @@
# syntax=docker/dockerfile:1.7
# Base images
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d
ARG PROXY_EXTRAS_SOURCE=published
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
# Pinned by digest like the other base images; bump explicitly on Node upgrades.

4
enterprise/pyproject.toml
View file

@ -1,6 +1,6 @@
[project]
name = "litellm-enterprise"
version = "0.1.69"
version = "0.1.70"
description = "Package for LiteLLM Enterprise features"
readme = "README.md"
requires-python = ">=3.9"
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.1.69"
version = "0.1.70"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-enterprise==",

4
gateway/Dockerfile
View file

@ -1,5 +1,5 @@
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
# Checksum from https://www.pgbouncer.org/downloads/ (the Wolfi repo only carries 1.24.x)
ARG PGBOUNCER_VERSION=1.25.2

1
gateway/routes/allowlist.py
View file

@ -131,6 +131,7 @@ GATEWAY_EXACT_PATHS: frozenset[str] = frozenset(
"/redoc",
"/test",
"/debug/memory/summary",
"/api/event_logging/batch",
}
)

2
helm/litellm/templates/ingress.yaml
View file

@ -61,7 +61,7 @@
"/v1/fine-tuning" "/fine-tuning" "/v1/responses" "/responses" "/v1/threads" "/threads"
"/v1/assistants" "/assistants" "/v1/vector_stores" "/vector_stores" "/v1/indexes"
"/v1/models" "/models" "/openai" "/engines"
"/v1/messages" "/messages" "/v1/skills" "/v1/a2a" "/a2a"
"/v1/messages" "/messages" "/v1/skills" "/v1/a2a" "/a2a" "/api/event_logging"
"/v1/rerank" "/v2/rerank" "/rerank" "/v1/ocr" "/ocr" "/v1/rag" "/rag"
"/v1/video" "/v1/videos" "/video" "/videos" "/v1/search" "/search"
"/v1/containers" "/containers" "/v1/evals" "/v1/memory" "/queue/chat"

4
litellm-proxy-extras/pyproject.toml
View file

@ -1,6 +1,6 @@
[project]
name = "litellm-proxy-extras"
version = "0.4.100"
version = "0.4.101"
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
readme = "README.md"
requires-python = ">=3.9"
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.4.100"
version = "0.4.101"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-proxy-extras==",

311
litellm-rust/Cargo.lock generated
View file

@ -88,6 +88,45 @@ version = "1.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "03918c3dbd7701a85c6b9887732e2921175f26c350b4563841d0958c21d57e6d"
[[package]]
name = "asn1-rs"
version = "0.7.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b7f43a50ac4fdca5df8e885c21b835997f0a1cdee65494a6847694a98652d9d8"
dependencies = [
"asn1-rs-derive",
"asn1-rs-impl",
"displaydoc",
"nom",
"num-traits",
"rusticata-macros",
"thiserror 2.0.19",
"time",
]
[[package]]
name = "asn1-rs-derive"
version = "0.6.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3109e49b1e4909e9db6515a30c633684d68cdeaa252f215214cb4fa1a5bfee2c"
dependencies = [
"proc-macro2",
"quote",
"syn 2.0.119",
"synstructure 0.13.2",
]
[[package]]
name = "asn1-rs-impl"
version = "0.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7b18050c2cd6fe86c3a76584ef5e0baf286d038cda203eb6223df2cc413565f7"
dependencies = [
"proc-macro2",
"quote",
"syn 2.0.119",
]
[[package]]
name = "assert-json-diff"
version = "2.0.2"
@ -810,7 +849,7 @@ version = "0.8.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "08807e080ed7f9d5433fa9b275196cfc35414f66a0c79d864dc51a0d825231a3"
dependencies = [
"bit-vec",
"bit-vec 0.8.0",
]
[[package]]
@ -819,6 +858,15 @@ version = "0.8.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5e764a1d40d510daf35e07be9eb06e75770908c27d411ee6c92109c9840eaaf7"
[[package]]
name = "bit-vec"
version = "0.9.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b71798fca2c1fe1086445a7258a4bc81e6e49dcd24c8d0dd9a1e57395b603f51"
dependencies = [
"serde",
]
[[package]]
name = "bitflags"
version = "1.3.2"
@ -1120,6 +1168,15 @@ version = "0.4.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "338089f42c427b86394a5ee60ff321da23a5c89c9d89514c829687b26359fcff"
[[package]]
name = "crc32c"
version = "0.6.8"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3a47af21622d091a8f0fb295b88bc886ac74efcc613efc19f5d0b21de5c89e47"
dependencies = [
"rustc_version",
]
[[package]]
name = "crc32fast"
version = "1.5.1"
@ -1378,6 +1435,20 @@ dependencies = [
"thiserror 2.0.19",
]
[[package]]
name = "der-parser"
version = "10.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "07da5016415d5a3c4dd39b11ed26f915f52fc4e0dc197d87908bc916e51bc1a6"
dependencies = [
"asn1-rs",
"displaydoc",
"nom",
"num-bigint 0.4.8",
"num-traits",
"rusticata-macros",
]
[[package]]
name = "deranged"
version = "0.5.8"
@ -2589,21 +2660,6 @@ dependencies = [
"wasm-bindgen",
]
[[package]]
name = "jsonwebtoken"
version = "11.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e75fe14a82d81e5f5af639997db37d8b96045938a7ac6ab18cdbe1c7467e05e1"
dependencies = [
"base64 0.22.1",
"getrandom 0.2.17",
"js-sys",
"serde",
"serde_json",
"signature",
"zeroize",
]
[[package]]
name = "lazy_static"
version = "1.5.0"
@ -2727,9 +2783,13 @@ dependencies = [
"litellm-auth-types",
"litellm-cache",
"litellm-cache-response",
"litellm-cache-testing",
"reqwest 0.12.28",
"rstest",
"serde_json",
"tokio",
"url",
"wiremock",
]
[[package]]
@ -2737,6 +2797,7 @@ name = "litellm-cache-disk"
version = "0.1.0"
dependencies = [
"litellm-cache",
"litellm-cache-testing",
"py_literal",
"rand 0.8.7",
"rstest",
@ -2755,8 +2816,10 @@ dependencies = [
"litellm-auth-gcp",
"litellm-auth-types",
"litellm-cache",
"litellm-cache-testing",
"percent-encoding",
"reqwest 0.12.28",
"rstest",
"serde_json",
"tokio",
"wiremock",
@ -2767,8 +2830,8 @@ name = "litellm-cache-memory"
version = "0.1.0"
dependencies = [
"litellm-cache",
"litellm-cache-testing",
"rstest",
"serde_json",
"tokio",
]
@ -2776,9 +2839,10 @@ dependencies = [
name = "litellm-cache-qdrant-semantic"
version = "0.1.0"
dependencies = [
"futures-executor",
"futures-util",
"litellm-cache",
"litellm-cache-response",
"litellm-cache-testing",
"qdrant-client",
"reqwest 0.12.28",
"rstest",
@ -2797,9 +2861,11 @@ name = "litellm-cache-redis"
version = "0.1.0"
dependencies = [
"litellm-cache",
"litellm-cache-testing",
"r2d2",
"redis",
"redis-test",
"rstest",
"serde_json",
"tokio",
]
@ -2810,10 +2876,10 @@ version = "0.1.0"
dependencies = [
"litellm-cache",
"litellm-cache-redis",
"litellm-cache-response",
"r2d2",
"litellm-cache-testing",
"redis",
"redis-test",
"rstest",
"serde_json",
"sha2 0.10.9",
"tokio",
@ -2829,6 +2895,7 @@ dependencies = [
"py_literal",
"redis",
"redis-test",
"rstest",
"serde",
"serde_json",
"sha2 0.10.9",
@ -2841,22 +2908,35 @@ version = "0.1.0"
dependencies = [
"aws-credential-types",
"aws-sdk-s3",
"aws-smithy-runtime-api",
"aws-smithy-types",
"aws-types",
"futures-util",
"http 1.4.2",
"litellm-auth-aws",
"litellm-cache",
"litellm-cache-testing",
"reqwest 0.12.28",
"rstest",
"serde_json",
"tokio",
"wiremock",
]
[[package]]
name = "litellm-cache-testing"
version = "0.1.0"
dependencies = [
"litellm-cache",
]
[[package]]
name = "litellm-cache-valkey-semantic"
version = "0.1.0"
dependencies = [
"litellm-cache",
"litellm-cache-redis",
"litellm-cache-response",
"litellm-cache-testing",
"redis",
"redis-test",
"rstest",
@ -2922,6 +3002,7 @@ name = "litellm-core-utils"
version = "0.1.0"
dependencies = [
"fancy-regex 0.19.2",
"litellm-tracing",
"litellm-types",
"rstest",
"serde",
@ -2932,6 +3013,14 @@ dependencies = [
"url",
]
[[package]]
name = "litellm-cost"
version = "0.1.0"
dependencies = [
"criterion",
"proptest",
]
[[package]]
name = "litellm-framing"
version = "0.1.0"
@ -3022,6 +3111,20 @@ dependencies = [
"url",
]
[[package]]
name = "litellm-model-catalog"
version = "0.1.0"
dependencies = [
"criterion",
"indexmap 2.14.0",
"litellm-model-catalog",
"rstest",
"schemars 1.2.2",
"serde",
"serde_json",
"thiserror 2.0.19",
]
[[package]]
name = "litellm-python-bridge"
version = "0.1.0"
@ -3029,6 +3132,7 @@ dependencies = [
"aws-sdk-secretsmanager",
"bytes",
"criterion",
"fancy-regex 0.19.2",
"futures-util",
"litellm-auth",
"litellm-auth-aws",
@ -3047,6 +3151,7 @@ dependencies = [
"litellm-callbacks-legacy-python",
"litellm-core",
"litellm-core-utils",
"litellm-host",
"litellm-host-python",
"litellm-http",
"litellm-llms",
@ -3054,6 +3159,7 @@ dependencies = [
"litellm-secrets-aws",
"litellm-secrets-types",
"litellm-token-counter",
"litellm-tracing",
"litellm-types",
"pyo3",
"pyo3-async-runtimes",
@ -3065,22 +3171,40 @@ dependencies = [
"serde_json",
"serde_with",
"sha2 0.10.9",
"thiserror 2.0.19",
"tokio",
"tokio-tungstenite",
"url",
"veil",
"wiremock",
]
[[package]]
name = "litellm-python-compat"
version = "0.1.0"
dependencies = [
"criterion",
"hex",
"num-bigint 0.4.8",
"num-traits",
"rstest",
"serde",
"serde-pickle",
"serde_json",
"thiserror 2.0.19",
]
[[package]]
name = "litellm-secrets"
version = "0.1.0"
dependencies = [
"aws-sdk-kms",
"base64 0.22.1",
"futures-util",
"google-cloud-auth",
"google-cloud-kms-v1",
"jsonwebtoken",
"litellm-core-utils",
"litellm-python-compat",
"litellm-secrets-aws",
"litellm-secrets-azure",
"litellm-secrets-cyberark",
@ -3110,11 +3234,12 @@ dependencies = [
"litellm-auth-aws",
"litellm-core-utils",
"litellm-secrets-types",
"litellm-tracing",
"rstest",
"serde_json",
"tempfile",
"thiserror 2.0.19",
"tokio",
"tracing",
"veil",
"wiremock",
]
@ -3146,15 +3271,17 @@ dependencies = [
"base64 0.22.1",
"litellm-core-utils",
"litellm-secrets-types",
"litellm-tracing",
"moka",
"percent-encoding",
"rcgen",
"reqwest 0.12.28",
"rstest",
"serde",
"serde_json",
"tempfile",
"thiserror 2.0.19",
"tokio",
"tracing",
"veil",
"wiremock",
]
@ -3164,6 +3291,7 @@ name = "litellm-secrets-google"
version = "0.1.0"
dependencies = [
"base64 0.22.1",
"crc32c",
"google-cloud-auth",
"google-cloud-gax",
"google-cloud-kms-v1",
@ -3189,7 +3317,6 @@ version = "0.1.0"
dependencies = [
"litellm-core-utils",
"litellm-secrets-types",
"moka",
"rstest",
"rustify",
"rustify_derive",
@ -3208,6 +3335,7 @@ name = "litellm-secrets-types"
version = "0.1.0"
dependencies = [
"litellm-auth-types",
"moka",
"rstest",
"serde",
"serde_json",
@ -3269,10 +3397,24 @@ dependencies = [
"tiktoken-rs",
]
[[package]]
name = "litellm-tracing"
version = "0.1.0"
dependencies = [
"fancy-regex 0.19.2",
"percent-encoding",
"rstest",
"serde_json",
"tokio",
"tracing",
"tracing-subscriber",
]
[[package]]
name = "litellm-types"
version = "0.1.0"
dependencies = [
"rstest",
"serde",
"serde_json",
]
@ -3509,6 +3651,15 @@ dependencies = [
"libc",
]
[[package]]
name = "oid-registry"
version = "0.8.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "12f40cff3dde1b6087cc5d5f5d4d65712f34016a03ed60e9c08dcc392736b5b7"
dependencies = [
"asn1-rs",
]
[[package]]
name = "once_cell"
version = "1.21.4"
@ -3642,6 +3793,16 @@ version = "0.2.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2ee67f1008b1ba2321834326597b8e186293b049a023cdef258527550b9935b4"
[[package]]
name = "pem"
version = "4.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d354a98a3d1251555de99e8fdd8afda05573c31b82f59063a7b0a29b5527f120"
dependencies = [
"base64 0.23.1",
"serde_core",
]
[[package]]
name = "percent-encoding"
version = "2.3.2"
@ -3820,7 +3981,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4b45fcc2344c680f5025fe57779faef368840d0bd1f42f216291f0dc4ace4744"
dependencies = [
"bit-set",
"bit-vec",
"bit-vec 0.8.0",
"bitflags 2.13.1",
"num-traits",
"rand 0.9.5",
@ -4209,6 +4370,20 @@ dependencies = [
"crossbeam-utils",
]
[[package]]
name = "rcgen"
version = "0.14.10"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8774e05a7d0de114588e6a28fe7e71694b82614ed569d86d8b389dfbc98b8ad8"
dependencies = [
"pem",
"ring",
"rustls-pki-types",
"time",
"x509-parser",
"yasna",
]
[[package]]
name = "redis"
version = "1.7.0"
@ -4491,6 +4666,15 @@ dependencies = [
"semver",
]
[[package]]
name = "rusticata-macros"
version = "4.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "faf0c4a6ece9950b9abdb62b1cfcf2a68b3b67a10ba445b3bb85be2a293d0632"
dependencies = [
"nom",
]
[[package]]
name = "rustify"
version = "0.7.0"
@ -4708,10 +4892,23 @@ checksum = "687274d293b6cdc6e73e0fee520bf2049650090d7164f87672d212a3c530cf4a"
dependencies = [
"dyn-clone",
"ref-cast",
"schemars_derive",
"serde",
"serde_json",
]
[[package]]
name = "schemars_derive"
version = "1.2.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d98c67716b46af2f0b8cf752abc930f6f9aecfbf671ecfb531db8a31dbe4e2ba"
dependencies = [
"proc-macro2",
"quote",
"serde_derive_internals",
"syn 3.0.0",
]
[[package]]
name = "scopeguard"
version = "1.2.0"
@ -4800,6 +4997,17 @@ dependencies = [
"syn 3.0.0",
]
[[package]]
name = "serde_derive_internals"
version = "0.30.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f852137cce035d6a4df67ccce505ff6b3e9fd3a10e3e52b24dc71e650bb1a9bd"
dependencies = [
"proc-macro2",
"quote",
"syn 3.0.0",
]
[[package]]
name = "serde_json"
version = "1.0.150"
@ -4943,15 +5151,6 @@ dependencies = [
"libc",
]
[[package]]
name = "signature"
version = "2.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "77549399552de45a898a580c1b41d445bf730df867cc44e6c0233bbc4b8329de"
dependencies = [
"rand_core 0.6.4",
]
[[package]]
name = "simd-adler32"
version = "0.3.10"
@ -6253,6 +6452,24 @@ version = "0.6.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1ffae5123b2d3fc086436f8834ae3ab053a283cfac8fe0a0b8eaae044768a4c4"
[[package]]
name = "x509-parser"
version = "0.18.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d43b0f71ce057da06bc0851b23ee24f3f86190b07203dd8f567d0b706a185202"
dependencies = [
"asn1-rs",
"data-encoding",
"der-parser",
"lazy_static",
"nom",
"oid-registry",
"ring",
"rusticata-macros",
"thiserror 2.0.19",
"time",
]
[[package]]
name = "xmlparser"
version = "0.13.6"
@ -6265,6 +6482,16 @@ version = "0.8.18"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "aee1b19627c7c60102ab80d3a9cbe18de90bfe03bfa6c3715447681f0e8c8af6"
[[package]]
name = "yasna"
version = "0.6.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b5f6765e852b9b4dc8e2a76843e4d64d1cea8e79bcde0b6901aea8e7c7f08282"
dependencies = [
"bit-vec 0.9.1",
"time",
]
[[package]]
name = "yoke"
version = "0.8.3"
@ -6334,20 +6561,6 @@ name = "zeroize"
version = "1.9.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e13c156562582aa81c60cb29407084cdb54c4164760106ab78e6c5b0858cf64e"
dependencies = [
"zeroize_derive",
]
[[package]]
name = "zeroize_derive"
version = "1.5.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3c50655cbb0fe3fc43170059e702f1ce5e19b84cec58dc87b037a09935c2f328"
dependencies = [
"proc-macro2",
"quote",
"syn 2.0.119",
]
[[package]]
name = "zerotrie"

3
litellm-rust/Cargo.toml
View file

@ -9,6 +9,8 @@ license = "MIT"
repository = "https://github.com/BerriAI/litellm"
[workspace.dependencies]
litellm-tracing = { path = "crates/tracing" }
tracing = "0.1"
litellm-core = { path = "crates/core" }
litellm-host = { path = "crates/host" }
litellm-callbacks-legacy-python = { path = "crates/callbacks-legacy-python" }
@ -39,6 +41,7 @@ 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" }
litellm-token-counter = { path = "crates/token-counter" }
litellm-token-counter-fast = { path = "crates/token-counter-fast" }
litellm-token-counter-huggingface = { path = "crates/token-counter-huggingface" }

5
litellm-rust/crates/cache-azure-blob/Cargo.toml
View file

@ -14,9 +14,14 @@ async-trait = "0.1"
azure_core = "1.1.0"
azure_storage_blob = "1.1.0"
futures-util.workspace = true
reqwest.workspace = true
tokio.workspace = true
url.workspace = true
[dev-dependencies]
litellm-cache-response.workspace = true
litellm-cache-testing.workspace = true
rstest.workspace = true
serde_json.workspace = true
tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }
wiremock = "0.6.5"

48
litellm-rust/crates/cache-azure-blob/src/cache.rs
View file

@ -3,7 +3,7 @@ use std::{sync::Arc, time::Duration};
use azure_core::{
credentials::TokenCredential,
error::ErrorKind,
http::{ClientOptions, RequestContent},
http::{ClientOptions, RequestContent, Transport},
};
use azure_storage_blob::{
BlobContainerClient, BlobContainerClientOptions,
@ -11,13 +11,12 @@ use azure_storage_blob::{
};
use futures_util::{TryStreamExt, future::try_join_all};
use litellm_cache::{
BaseCache, BatchCache, CacheCodec, CacheConnectionResult, CacheConnectionStatus, Error,
ExactCacheContext, FlushCache,
BaseCache, BatchCache, CacheCodec, DisconnectCache, Error, ExactCacheContext, FlushCache,
};
use tokio::runtime::Handle;
use url::Url;
use crate::credential::AzureBlobCredential;
use crate::{credential::AzureBlobCredential, transport::ReqwestTransport};
pub struct AzureBlobCache<C> {
container: BlobContainerClient,
@ -28,9 +27,11 @@ pub struct AzureBlobCache<C> {
}
impl<C: CacheCodec> AzureBlobCache<C> {
/// `http` is the host's pooled client; the SDK sends every request through it.
pub async fn connect(
account_url: &str,
container: &str,
http: reqwest::Client,
codec: C,
runtime: Handle,
) -> Result<Self, Error> {
@ -38,7 +39,10 @@ impl<C: CacheCodec> AzureBlobCache<C> {
account_url,
container,
Some(Arc::new(AzureBlobCredential::default())),
ClientOptions::default(),
ClientOptions {
transport: Some(Transport::new(Arc::new(ReqwestTransport(http)))),
..ClientOptions::default()
},
codec,
runtime,
)
@ -152,7 +156,11 @@ impl<C: CacheCodec> AzureBlobCache<C> {
}
fn block_on<T>(&self, future: impl Future<Output = T>) -> T {
self.runtime.block_on(future)
if Handle::try_current().is_ok() {
tokio::task::block_in_place(|| self.runtime.block_on(future))
} else {
self.runtime.block_on(future)
}
}
}
@ -217,25 +225,6 @@ impl<C: CacheCodec> BaseCache for AzureBlobCache<C> {
.await
.map(drop)
}
async fn disconnect(&self) -> Result<(), Error> {
Ok(())
}
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
Ok(match self.container.get_properties(None).await {
Ok(_) => CacheConnectionResult {
status: CacheConnectionStatus::Success,
message: "Azure Blob cache connection test successful".into(),
error: None,
},
Err(error) => CacheConnectionResult {
status: CacheConnectionStatus::Failed,
message: format!("Azure Blob connection failed: {error}"),
error: Some(error.to_string()),
},
})
}
}
impl<C: CacheCodec> BatchCache for AzureBlobCache<C> {}
@ -250,5 +239,10 @@ impl<C: CacheCodec> FlushCache for AzureBlobCache<C> {
}
}
#[cfg(test)]
mod tests;
impl<C: CacheCodec> DisconnectCache for AzureBlobCache<C> {
/// Python closes its two SDK clients; the Rust clients hold no connection of their own
/// (the pooled transport belongs to the host), so there is nothing to release.
async fn disconnect(&self) -> Result<(), Error> {
Ok(())
}
}

746
litellm-rust/crates/cache-azure-blob/src/cache/tests.rs vendored
View file

@ -1,746 +0,0 @@
use std::{
collections::BTreeMap,
sync::{Arc, Mutex},
time::Duration,
};
use azure_core::http::{
AsyncRawResponse, Body, ClientOptions, HttpClient, Method, Request, StatusCode, Transport,
headers::{HeaderName, Headers},
};
use litellm_cache::{
BaseCache, BatchCache, BatchEntry, CacheConnectionStatus, Error, ExactCacheContext, FlushCache,
};
use litellm_cache_response::{
CacheEntry, CacheKeyField, CacheKeyInput, ResponseCache, ResponseCacheCodec,
ResponseCacheRequest, cache_key,
};
use serde_json::json;
use tokio::runtime::Runtime;
use super::AzureBlobCache;
const ACCOUNT_URL: &str = "https://example.blob.core.windows.net";
const CONTAINER: &str = "litellm-cache";
const IF_NONE_MATCH: HeaderName = HeaderName::from_static("if-none-match");
const ERROR_CODE: HeaderName = HeaderName::from_static("x-ms-error-code");
#[derive(Clone, Debug, PartialEq, Eq)]
struct RecordedRequest {
method: Method,
path: String,
query: String,
if_none_match: Option<String>,
}
#[derive(Default)]
struct FakeState {
container_exists: bool,
blobs: BTreeMap<String, Vec<u8>>,
requests: Vec<RecordedRequest>,
failing: bool,
precondition_conflicts: bool,
}
#[derive(Clone, Default)]
struct FakeBlobService {
state: Arc<Mutex<FakeState>>,
}
impl std::fmt::Debug for FakeBlobService {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str("FakeBlobService")
}
}
impl FakeBlobService {
fn with_existing_container() -> Self {
let service = Self::default();
service.state.lock().unwrap().container_exists = true;
service
}
fn blob(&self, name: &str) -> Option<Vec<u8>> {
self.state.lock().unwrap().blobs.get(name).cloned()
}
fn blob_names(&self) -> Vec<String> {
self.state.lock().unwrap().blobs.keys().cloned().collect()
}
fn seed_blob(&self, name: &str, bytes: &[u8]) {
self.state
.lock()
.unwrap()
.blobs
.insert(name.to_string(), bytes.to_vec());
}
fn set_failing(&self, failing: bool) {
self.state.lock().unwrap().failing = failing;
}
fn set_precondition_conflicts(&self, enabled: bool) {
self.state.lock().unwrap().precondition_conflicts = enabled;
}
fn requests(&self) -> Vec<RecordedRequest> {
self.state.lock().unwrap().requests.clone()
}
fn container_exists(&self) -> bool {
self.state.lock().unwrap().container_exists
}
fn respond(status: StatusCode, error_code: Option<&str>, body: Vec<u8>) -> AsyncRawResponse {
let mut headers = Headers::new();
if let Some(code) = error_code {
headers.insert(ERROR_CODE, code.to_string());
}
AsyncRawResponse::from_bytes(status, headers, body)
}
fn list_body(state: &FakeState) -> Vec<u8> {
let mut xml = String::from(
r#"<?xml version="1.0" encoding="utf-8"?><EnumerationResults ServiceEndpoint="https://example.blob.core.windows.net/" ContainerName="litellm-cache"><Blobs>"#,
);
for name in state.blobs.keys() {
xml.push_str(&format!(
"<Blob><Name>{name}</Name><Properties><BlobType>BlockBlob</BlobType></Properties></Blob>"
));
}
xml.push_str("</Blobs><NextMarker /></EnumerationResults>");
xml.into_bytes()
}
}
#[async_trait::async_trait]
impl HttpClient for FakeBlobService {
async fn execute_request(&self, request: &Request) -> azure_core::Result<AsyncRawResponse> {
let mut state = self.state.lock().unwrap();
let path = request.url().path().to_string();
let query = request.url().query().unwrap_or_default().to_string();
let if_none_match = request
.headers()
.get_optional_str(&IF_NONE_MATCH)
.map(str::to_owned);
state.requests.push(RecordedRequest {
method: request.method(),
path: path.clone(),
query: query.clone(),
if_none_match: if_none_match.clone(),
});
if state.failing {
return Ok(Self::respond(
StatusCode::Forbidden,
Some("AuthorizationFailure"),
Vec::new(),
));
}
let container_path = format!("/{CONTAINER}");
let blob_name = path
.strip_prefix(&format!("{container_path}/"))
.map(str::to_owned);
let is_container = path == container_path && query.contains("restype=container");
let response = match (request.method(), is_container, blob_name) {
(Method::Put, true, None) if state.container_exists => Self::respond(
StatusCode::Conflict,
Some("ContainerAlreadyExists"),
Vec::new(),
),
(Method::Put, true, None) => {
state.container_exists = true;
Self::respond(StatusCode::Created, None, Vec::new())
}
(Method::Get, true, None) if query.contains("comp=list") => {
Self::respond(StatusCode::Ok, None, Self::list_body(&state))
}
(Method::Get, true, None) if state.container_exists => {
Self::respond(StatusCode::Ok, None, Vec::new())
}
(Method::Get, true, None) => {
Self::respond(StatusCode::NotFound, Some("ContainerNotFound"), Vec::new())
}
(Method::Put, false, Some(name)) => {
if if_none_match.as_deref() == Some("*") && state.blobs.contains_key(&name) {
if state.precondition_conflicts {
Self::respond(
StatusCode::PreconditionFailed,
Some("ConditionNotMet"),
Vec::new(),
)
} else {
Self::respond(StatusCode::Conflict, Some("BlobAlreadyExists"), Vec::new())
}
} else {
let bytes = match request.body() {
Body::Bytes(bytes) => bytes.to_vec(),
Body::SeekableStream(_) => panic!("unexpected streaming upload"),
};
state.blobs.insert(name, bytes);
Self::respond(StatusCode::Created, None, Vec::new())
}
}
(Method::Get, false, Some(name)) => match state.blobs.get(&name) {
Some(bytes) => Self::respond(StatusCode::Ok, None, bytes.clone()),
None => Self::respond(StatusCode::NotFound, Some("BlobNotFound"), Vec::new()),
},
(Method::Delete, false, Some(name)) => match state.blobs.remove(&name) {
Some(_) => Self::respond(StatusCode::Accepted, None, Vec::new()),
None => Self::respond(StatusCode::NotFound, Some("BlobNotFound"), Vec::new()),
},
(method, _, _) => panic!("unexpected request {method:?} {path}?{query}"),
};
Ok(response)
}
}
struct Fixture {
runtime: Runtime,
service: FakeBlobService,
cache: Arc<AzureBlobCache<ResponseCacheCodec>>,
}
impl Fixture {
fn new(service: FakeBlobService) -> Self {
let runtime = Runtime::new().unwrap();
let cache = runtime
.block_on(Self::connect(&service, runtime.handle().clone()))
.unwrap();
Self {
runtime,
service,
cache: Arc::new(cache),
}
}
async fn connect(
service: &FakeBlobService,
handle: tokio::runtime::Handle,
) -> Result<AzureBlobCache<ResponseCacheCodec>, Error> {
AzureBlobCache::connect_with_options(
ACCOUNT_URL,
CONTAINER,
None,
ClientOptions {
transport: Some(Transport::new(Arc::new(service.clone()))),
..ClientOptions::default()
},
ResponseCacheCodec,
handle,
)
.await
}
fn response_cache(&self) -> ResponseCache<AzureBlobCache<ResponseCacheCodec>> {
ResponseCache::new(self.cache.clone())
}
fn stored_json(&self, key: &str) -> serde_json::Value {
serde_json::from_slice(&self.service.blob(key).expect("blob should exist")).unwrap()
}
}
fn request(model: &str) -> ResponseCacheRequest {
ResponseCacheRequest::new(CacheKeyInput {
fields: vec![CacheKeyField {
name: "model".into(),
value: Some(model.into()),
api_parameter: true,
internal_parameter: false,
}],
preset: None,
namespace: None,
include_provider_parameters: false,
})
}
fn now() -> Duration {
Duration::from_secs(1_700_000_000)
}
fn entry(value: serde_json::Value) -> CacheEntry {
CacheEntry {
timestamp: Some(1_700_000_000.5),
response: value,
}
}
fn no_ttl() -> ExactCacheContext {
ExactCacheContext::default()
}
fn with_ttl(seconds: u64) -> ExactCacheContext {
ExactCacheContext {
ttl: Some(Duration::from_secs(seconds)),
}
}
#[test]
fn connect_creates_the_container_once() {
let fixture = Fixture::new(FakeBlobService::default());
assert!(fixture.service.container_exists());
assert_eq!(
fixture.service.requests(),
vec![RecordedRequest {
method: Method::Put,
path: format!("/{CONTAINER}"),
query: "restype=container".into(),
if_none_match: None,
}]
);
assert_eq!(fixture.cache.account_url(), ACCOUNT_URL);
assert_eq!(fixture.cache.container_name(), CONTAINER);
}
#[test]
fn connect_accepts_an_existing_container() {
let fixture = Fixture::new(FakeBlobService::with_existing_container());
assert!(fixture.service.container_exists());
assert_eq!(fixture.service.requests().len(), 1);
}
#[test]
fn connect_accepts_account_urls_with_trailing_slash() {
let runtime = Runtime::new().unwrap();
let service = FakeBlobService::default();
let cache = runtime
.block_on(AzureBlobCache::connect_with_options(
"https://example.blob.core.windows.net/",
CONTAINER,
None,
ClientOptions {
transport: Some(Transport::new(Arc::new(service.clone()))),
..ClientOptions::default()
},
ResponseCacheCodec,
runtime.handle().clone(),
))
.unwrap();
assert_eq!(service.requests()[0].path, format!("/{CONTAINER}"));
assert_eq!(cache.account_url(), "https://example.blob.core.windows.net");
}
#[test]
fn connect_keeps_account_url_query_parameters_on_the_container_path() {
let runtime = Runtime::new().unwrap();
let service = FakeBlobService::default();
runtime
.block_on(AzureBlobCache::connect_with_options(
"https://example.blob.core.windows.net/?sv=2024-01-01&sig=abc",
CONTAINER,
None,
ClientOptions {
transport: Some(Transport::new(Arc::new(service.clone()))),
..ClientOptions::default()
},
ResponseCacheCodec,
runtime.handle().clone(),
))
.unwrap();
let create = &service.requests()[0];
assert_eq!(create.path, format!("/{CONTAINER}"));
assert!(create.query.contains("sig=abc"));
}
#[test]
fn connect_surfaces_service_failures() {
let runtime = Runtime::new().unwrap();
let service = FakeBlobService::default();
service.set_failing(true);
let result = runtime.block_on(Fixture::connect(&service, runtime.handle().clone()));
assert!(matches!(result, Err(Error::Unavailable)));
}
#[test]
fn sync_set_and_get_round_trip_python_json_shape() {
let fixture = Fixture::new(FakeBlobService::default());
let value = entry(json!({"choices": [{"message": {"content": "héllo 🌍"}}]}));
fixture
.cache
.set_cache("key-1", value.clone(), &no_ttl())
.unwrap();
assert_eq!(
fixture.stored_json("key-1"),
json!({
"timestamp": 1_700_000_000.5,
"response": {"choices": [{"message": {"content": "héllo 🌍"}}]}
})
);
assert_eq!(
fixture.cache.get_cache("key-1", &no_ttl()).unwrap(),
Some(value)
);
}
#[test]
fn sync_set_does_not_overwrite_an_existing_blob() {
let fixture = Fixture::new(FakeBlobService::default());
fixture
.cache
.set_cache("key", entry(json!({"v": "first"})), &no_ttl())
.unwrap();
fixture
.cache
.set_cache("key", entry(json!({"v": "second"})), &no_ttl())
.unwrap();
assert_eq!(
fixture.stored_json("key")["response"],
json!({"v": "first"})
);
let uploads: Vec<_> = fixture
.service
.requests()
.into_iter()
.filter(|request| request.method == Method::Put && request.path.ends_with("/key"))
.collect();
assert_eq!(uploads.len(), 2);
assert!(
uploads
.iter()
.all(|request| request.if_none_match.as_deref() == Some("*"))
);
}
#[test]
fn sync_set_treats_a_precondition_conflict_as_an_existing_blob() {
let fixture = Fixture::new(FakeBlobService::default());
fixture.service.set_precondition_conflicts(true);
fixture
.cache
.set_cache("key", entry(json!({"v": "first"})), &no_ttl())
.unwrap();
fixture
.cache
.set_cache("key", entry(json!({"v": "second"})), &no_ttl())
.unwrap();
assert_eq!(
fixture.stored_json("key")["response"],
json!({"v": "first"})
);
}
#[test]
fn async_set_overwrites_an_existing_blob() {
let fixture = Fixture::new(FakeBlobService::default());
fixture.runtime.block_on(async {
fixture
.cache
.async_set_cache("key", entry(json!({"v": "first"})), no_ttl())
.await
.unwrap();
fixture
.cache
.async_set_cache("key", entry(json!({"v": "second"})), no_ttl())
.await
.unwrap();
assert_eq!(
fixture
.cache
.async_get_cache("key", &no_ttl())
.await
.unwrap(),
Some(entry(json!({"v": "second"})))
);
});
assert_eq!(
fixture.stored_json("key")["response"],
json!({"v": "second"})
);
assert!(
fixture
.service
.requests()
.iter()
.filter(|request| request.method == Method::Put && request.path.ends_with("/key"))
.all(|request| request.if_none_match.is_none())
);
}
#[test]
fn missing_blobs_are_misses() {
let fixture = Fixture::new(FakeBlobService::default());
assert_eq!(fixture.cache.get_cache("absent", &no_ttl()).unwrap(), None);
assert_eq!(
fixture
.runtime
.block_on(fixture.cache.async_get_cache("absent", &no_ttl()))
.unwrap(),
None
);
}
#[test]
fn ttl_is_ignored_and_entries_never_expire() {
let fixture = Fixture::new(FakeBlobService::default());
assert_eq!(fixture.cache.get_ttl(&with_ttl(1)), None);
assert_eq!(fixture.cache.get_ttl(&no_ttl()), None);
fixture
.cache
.set_cache("key", entry(json!("value")), &with_ttl(1))
.unwrap();
std::thread::sleep(Duration::from_millis(1100));
assert_eq!(
fixture.cache.get_cache("key", &with_ttl(1)).unwrap(),
Some(entry(json!("value")))
);
assert!(
fixture
.service
.requests()
.iter()
.all(|request| !request.query.contains("expiry"))
);
}
#[test]
fn malformed_blobs_are_invalid_entries_and_response_cache_misses() {
let fixture = Fixture::new(FakeBlobService::default());
fixture.service.seed_blob("broken-json", b"{not json");
fixture
.service
.seed_blob("broken-utf8", &[0xff, 0xfe, 0x22]);
fixture
.service
.seed_blob("wrong-shape", br#"{"timestamp": "yesterday"}"#);
for key in ["broken-json", "broken-utf8", "wrong-shape"] {
assert!(matches!(
fixture.cache.get_cache(key, &no_ttl()),
Err(Error::InvalidEntry)
));
}
let response_cache = fixture.response_cache();
let broken = request("broken");
fixture
.service
.seed_blob(&cache_key(&broken.key), b"{not json");
assert_eq!(response_cache.lookup(&broken, now()).unwrap(), None);
assert_eq!(
fixture
.runtime
.block_on(response_cache.async_lookup(&broken, now()))
.unwrap(),
None
);
}
#[test]
fn batch_get_preserves_order_and_marks_misses_and_invalid_entries() {
let fixture = Fixture::new(FakeBlobService::default());
fixture
.cache
.set_cache("a", entry(json!("A")), &no_ttl())
.unwrap();
fixture
.cache
.set_cache("c", entry(json!("C")), &no_ttl())
.unwrap();
fixture.service.seed_blob("bad", b"nope");
let keys = ["c", "missing", "a", "bad"].map(String::from);
let sync = fixture.cache.batch_get_cache(&keys, &no_ttl()).unwrap();
assert_eq!(
sync,
vec![
BatchEntry::Hit(entry(json!("C"))),
BatchEntry::Miss,
BatchEntry::Hit(entry(json!("A"))),
BatchEntry::Invalid,
]
);
let asynchronous = fixture
.runtime
.block_on(fixture.cache.async_batch_get_cache(keys.to_vec(), no_ttl()))
.unwrap();
assert_eq!(asynchronous, sync);
let response_cache = fixture.response_cache();
let requests = [request("hit"), request("missing"), request("bad")];
response_cache
.store(&requests[0], json!("HIT"), now())
.unwrap();
fixture
.service
.seed_blob(&cache_key(&requests[2].key), b"nope");
let hits = response_cache.lookup_batch(&requests, now()).unwrap();
assert_eq!(hits.values, vec![Some(json!("HIT")), None, None]);
assert_eq!(hits.missing_indices, vec![1, 2]);
let async_hits = fixture
.runtime
.block_on(response_cache.async_lookup_batch(&requests, now()))
.unwrap();
assert_eq!(async_hits.values, hits.values);
}
#[test]
fn async_pipeline_writes_every_entry_with_overwrite() {
let fixture = Fixture::new(FakeBlobService::default());
fixture.service.seed_blob("k2", b"stale");
fixture
.runtime
.block_on(fixture.cache.async_set_cache_pipeline(
vec![
("k1".into(), entry(json!({"n": 1}))),
("k2".into(), entry(json!({"n": 2}))),
("k3".into(), entry(json!({"n": 3}))),
],
with_ttl(30),
))
.unwrap();
assert_eq!(fixture.service.blob_names(), ["k1", "k2", "k3"]);
assert_eq!(fixture.stored_json("k2")["response"], json!({"n": 2}));
}
#[test]
fn flush_deletes_every_blob_in_the_container() {
let fixture = Fixture::new(FakeBlobService::default());
for key in ["x", "y", "z"] {
fixture
.cache
.set_cache(key, entry(json!(key)), &no_ttl())
.unwrap();
}
fixture.cache.flush_cache().unwrap();
assert!(fixture.service.blob_names().is_empty());
assert!(fixture.service.container_exists());
fixture
.cache
.set_cache("again", entry(json!(1)), &no_ttl())
.unwrap();
fixture
.runtime
.block_on(fixture.cache.async_flush_cache())
.unwrap();
assert!(fixture.service.blob_names().is_empty());
}
#[test]
fn service_failures_map_to_unavailable() {
let fixture = Fixture::new(FakeBlobService::default());
fixture.service.set_failing(true);
assert!(matches!(
fixture.cache.get_cache("key", &no_ttl()),
Err(Error::Unavailable)
));
assert!(matches!(
fixture.cache.set_cache("key", entry(json!(1)), &no_ttl()),
Err(Error::Unavailable)
));
assert!(matches!(
fixture.cache.flush_cache(),
Err(Error::Unavailable)
));
assert!(matches!(
fixture.runtime.block_on(
fixture
.cache
.async_set_cache_pipeline(vec![("k".into(), entry(json!(1)))], no_ttl())
),
Err(Error::Unavailable)
));
}
#[test]
fn test_connection_reports_container_reachability() {
let fixture = Fixture::new(FakeBlobService::default());
let ok = fixture
.runtime
.block_on(fixture.cache.test_connection())
.unwrap();
assert_eq!(ok.status, CacheConnectionStatus::Success);
assert!(ok.error.is_none());
fixture.service.set_failing(true);
let failed = fixture
.runtime
.block_on(fixture.cache.test_connection())
.unwrap();
assert_eq!(failed.status, CacheConnectionStatus::Failed);
assert!(failed.error.is_some());
}
#[test]
fn disconnect_is_idempotent_and_keeps_data() {
let fixture = Fixture::new(FakeBlobService::default());
fixture
.cache
.set_cache("key", entry(json!(1)), &no_ttl())
.unwrap();
fixture.runtime.block_on(async {
fixture.cache.disconnect().await.unwrap();
fixture.cache.disconnect().await.unwrap();
});
assert_eq!(
fixture.cache.get_cache("key", &no_ttl()).unwrap(),
Some(entry(json!(1)))
);
}
#[test]
fn response_cache_stores_and_reads_through_the_backend() {
let fixture = Fixture::new(FakeBlobService::default());
let response_cache = fixture.response_cache();
let mut request = request("gpt");
request.context = with_ttl(60);
let response = json!({"id": "chatcmpl-1"});
response_cache
.store(&request, response.clone(), now())
.unwrap();
assert_eq!(
fixture.stored_json(&cache_key(&request.key)),
json!({"timestamp": 1_700_000_000.0, "response": {"id": "chatcmpl-1"}})
);
assert_eq!(
response_cache
.lookup(&request, now() + Duration::from_secs(3600))
.unwrap(),
Some(response.clone())
);
assert_eq!(
fixture
.runtime
.block_on(response_cache.async_lookup(&request, now() + Duration::from_secs(3600)))
.unwrap(),
Some(response.clone())
);
fixture.runtime.block_on(async {
response_cache
.async_store(&request, json!("replaced"), now())
.await
.unwrap();
assert_eq!(
response_cache.async_lookup(&request, now()).await.unwrap(),
Some(json!("replaced"))
);
response_cache.async_flush().await.unwrap();
assert_eq!(
response_cache.async_lookup(&request, now()).await.unwrap(),
None
);
});
}
#[test]
fn non_object_responses_are_written_serialized_like_python() {
let fixture = Fixture::new(FakeBlobService::default());
fixture
.cache
.set_cache("s", entry(json!("plain")), &no_ttl())
.unwrap();
assert_eq!(
fixture.stored_json("s"),
json!({"timestamp": 1_700_000_000.5, "response": "\"plain\""})
);
assert_eq!(
fixture.cache.get_cache("s", &no_ttl()).unwrap(),
Some(entry(json!("plain")))
);
}

2
litellm-rust/crates/cache-azure-blob/src/lib.rs
View file

@ -1,5 +1,7 @@
mod cache;
mod credential;
mod transport;
pub use cache::AzureBlobCache;
pub use credential::AzureBlobCredential;
pub use transport::ReqwestTransport;

49
litellm-rust/crates/cache-azure-blob/src/transport.rs Normal file
View file

@ -0,0 +1,49 @@
use azure_core::{
error::ErrorKind,
http::{
AsyncRawResponse, Body, HttpClient, Request,
headers::{HeaderName, HeaderValue, Headers},
},
};
use futures_util::TryStreamExt;
#[derive(Debug)]
pub struct ReqwestTransport(pub reqwest::Client);
#[async_trait::async_trait]
impl HttpClient for ReqwestTransport {
async fn execute_request(&self, request: &Request) -> azure_core::Result<AsyncRawResponse> {
let method = reqwest::Method::from_bytes(request.method().as_ref().as_bytes())
.map_err(|error| azure_core::Error::new(ErrorKind::Other, error))?;
let mut outgoing = self.0.request(method, request.url().as_str());
for (name, value) in request.headers().iter() {
outgoing = outgoing.header(name.as_str(), value.as_str());
}
let outgoing = match request.body().clone() {
Body::Bytes(bytes) => outgoing.body(bytes),
Body::SeekableStream(stream) => outgoing.body(reqwest::Body::wrap_stream(stream)),
};
let response = outgoing.send().await.map_err(|error| {
let kind = if error.is_connect() {
ErrorKind::Connection
} else {
ErrorKind::Io
};
azure_core::Error::new(kind, error)
})?;
let status = response.status().as_u16().into();
let mut headers = Headers::new();
for (name, value) in response.headers() {
if let Ok(value) = value.to_str() {
headers.insert(
HeaderName::from(name.as_str().to_owned()),
HeaderValue::from(value.to_owned()),
);
}
}
let body = response
.bytes_stream()
.map_err(|error| azure_core::Error::new(ErrorKind::Io, error));
Ok(AsyncRawResponse::new(status, headers, Box::pin(body)))
}
}

494
litellm-rust/crates/cache-azure-blob/tests/cache.rs Normal file
View file

@ -0,0 +1,494 @@
mod support;
use std::{sync::Arc, time::Duration};
use azure_core::http::Method;
use litellm_cache::{
BaseCache, BatchCache, BatchEntry, DisconnectCache, Error, ExactCacheContext, FlushCache,
};
use litellm_cache_azure_blob::AzureBlobCache;
use litellm_cache_response::{
CacheEntry, CacheKeyField, CacheKeyInput, ResponseCache, ResponseCacheCodec,
ResponseCacheRequest, cache_key,
};
use rstest::{fixture, rstest};
use serde_json::json;
use support::{ACCOUNT_URL, CONTAINER, FakeBlobService, RecordedRequest};
use tokio::runtime::Runtime;
type Fixture = support::Fixture<ResponseCacheCodec>;
#[fixture]
fn fixture() -> Fixture {
Fixture::new(FakeBlobService::default(), ResponseCacheCodec)
}
fn response_cache(fixture: &Fixture) -> ResponseCache<AzureBlobCache<ResponseCacheCodec>> {
ResponseCache::new(fixture.cache.clone())
}
fn request(model: &str) -> ResponseCacheRequest {
ResponseCacheRequest::new(CacheKeyInput {
fields: vec![CacheKeyField {
name: "model".into(),
value: Some(model.into()),
api_parameter: true,
internal_parameter: false,
}],
preset: None,
namespace: None,
include_provider_parameters: false,
})
}
fn now() -> Duration {
Duration::from_secs(1_700_000_000)
}
fn entry(value: serde_json::Value) -> CacheEntry {
CacheEntry {
timestamp: Some(1_700_000_000.5),
response: value,
}
}
fn no_ttl() -> ExactCacheContext {
ExactCacheContext::default()
}
fn with_ttl(seconds: u64) -> ExactCacheContext {
ExactCacheContext {
ttl: Some(Duration::from_secs(seconds)),
}
}
fn connect_to(account_url: &str) -> (FakeBlobService, AzureBlobCache<ResponseCacheCodec>) {
let runtime = Runtime::new().unwrap();
let service = FakeBlobService::default();
let cache = runtime
.block_on(support::connect(
&service,
account_url,
ResponseCacheCodec,
runtime.handle().clone(),
))
.unwrap();
(service, cache)
}
#[rstest]
fn connect_creates_the_container_once(fixture: Fixture) {
assert!(fixture.service.container_exists());
assert_eq!(
fixture.service.requests(),
vec![RecordedRequest {
method: Method::Put,
path: format!("/{CONTAINER}"),
query: "restype=container".into(),
if_none_match: None,
}]
);
assert_eq!(fixture.cache.account_url(), ACCOUNT_URL);
assert_eq!(fixture.cache.container_name(), CONTAINER);
}
#[rstest]
fn connect_accepts_an_existing_container() {
let fixture = Fixture::new(
FakeBlobService::with_existing_container(),
ResponseCacheCodec,
);
assert!(fixture.service.container_exists());
assert_eq!(fixture.service.requests().len(), 1);
}
#[rstest]
fn connect_accepts_account_urls_with_trailing_slash() {
let (service, cache) = connect_to("https://example.blob.core.windows.net/");
assert_eq!(service.requests()[0].path, format!("/{CONTAINER}"));
assert_eq!(cache.account_url(), "https://example.blob.core.windows.net");
}
#[rstest]
fn connect_keeps_account_url_query_parameters_on_the_container_path() {
let (service, _) = connect_to("https://example.blob.core.windows.net/?sv=2024-01-01&sig=abc");
let create = &service.requests()[0];
assert_eq!(create.path, format!("/{CONTAINER}"));
assert!(create.query.contains("sig=abc"));
}
#[rstest]
fn connect_surfaces_service_failures() {
let runtime = Runtime::new().unwrap();
let service = FakeBlobService::default();
service.set_failing(true);
let result = runtime.block_on(support::connect(
&service,
ACCOUNT_URL,
ResponseCacheCodec,
runtime.handle().clone(),
));
assert!(matches!(result, Err(Error::Unavailable)));
}
#[rstest]
fn sync_set_and_get_round_trip_python_json_shape(fixture: Fixture) {
let value = entry(json!({"choices": [{"message": {"content": "héllo 🌍"}}]}));
fixture
.cache
.set_cache("key-1", value.clone(), &no_ttl())
.unwrap();
assert_eq!(
fixture.stored_json("key-1"),
json!({
"timestamp": 1_700_000_000.5,
"response": {"choices": [{"message": {"content": "héllo 🌍"}}]}
})
);
assert_eq!(
fixture.cache.get_cache("key-1", &no_ttl()).unwrap(),
Some(value)
);
}
#[rstest]
#[case::blob_already_exists(false)]
#[case::precondition_conflict(true)]
fn sync_set_does_not_overwrite_an_existing_blob(fixture: Fixture, #[case] precondition: bool) {
fixture.service.set_precondition_conflicts(precondition);
fixture
.cache
.set_cache("key", entry(json!({"v": "first"})), &no_ttl())
.unwrap();
fixture
.cache
.set_cache("key", entry(json!({"v": "second"})), &no_ttl())
.unwrap();
assert_eq!(
fixture.stored_json("key")["response"],
json!({"v": "first"})
);
let uploads: Vec<_> = fixture
.service
.requests()
.into_iter()
.filter(|request| request.method == Method::Put && request.path.ends_with("/key"))
.collect();
assert_eq!(uploads.len(), 2);
assert!(
uploads
.iter()
.all(|request| request.if_none_match.as_deref() == Some("*"))
);
}
#[rstest]
fn async_set_overwrites_an_existing_blob(fixture: Fixture) {
fixture.runtime.block_on(async {
fixture
.cache
.async_set_cache("key", entry(json!({"v": "first"})), no_ttl())
.await
.unwrap();
fixture
.cache
.async_set_cache("key", entry(json!({"v": "second"})), no_ttl())
.await
.unwrap();
assert_eq!(
fixture
.cache
.async_get_cache("key", &no_ttl())
.await
.unwrap(),
Some(entry(json!({"v": "second"})))
);
});
assert_eq!(
fixture.stored_json("key")["response"],
json!({"v": "second"})
);
assert!(
fixture
.service
.requests()
.iter()
.filter(|request| request.method == Method::Put && request.path.ends_with("/key"))
.all(|request| request.if_none_match.is_none())
);
}
#[rstest]
fn missing_blobs_are_misses(fixture: Fixture) {
assert_eq!(fixture.cache.get_cache("absent", &no_ttl()).unwrap(), None);
assert_eq!(
fixture
.runtime
.block_on(fixture.cache.async_get_cache("absent", &no_ttl()))
.unwrap(),
None
);
}
#[rstest]
fn ttl_is_ignored_and_entries_never_expire(fixture: Fixture) {
assert_eq!(fixture.cache.get_ttl(&with_ttl(1)), None);
assert_eq!(fixture.cache.get_ttl(&no_ttl()), None);
fixture
.cache
.set_cache("key", entry(json!("value")), &with_ttl(1))
.unwrap();
std::thread::sleep(Duration::from_millis(1100));
assert_eq!(
fixture.cache.get_cache("key", &with_ttl(1)).unwrap(),
Some(entry(json!("value")))
);
assert!(
fixture
.service
.requests()
.iter()
.all(|request| !request.query.contains("expiry"))
);
}
#[rstest]
#[case::broken_json("broken-json", b"{not json".as_slice())]
#[case::broken_utf8("broken-utf8", &[0xff, 0xfe, 0x22])]
#[case::wrong_shape("wrong-shape", br#"{"timestamp": "yesterday"}"#.as_slice())]
fn malformed_blobs_are_invalid_entries(fixture: Fixture, #[case] key: &str, #[case] bytes: &[u8]) {
fixture.service.seed_blob(key, bytes);
assert!(matches!(
fixture.cache.get_cache(key, &no_ttl()),
Err(Error::InvalidEntry)
));
}
#[rstest]
fn malformed_blobs_are_response_cache_misses(fixture: Fixture) {
let response_cache = response_cache(&fixture);
let broken = request("broken");
fixture
.service
.seed_blob(&cache_key(&broken.key), b"{not json");
assert_eq!(response_cache.lookup(&broken, now()).unwrap(), None);
assert_eq!(
fixture
.runtime
.block_on(response_cache.async_lookup(&broken, now()))
.unwrap(),
None
);
}
#[rstest]
fn batch_get_preserves_order_and_marks_misses_and_invalid_entries(fixture: Fixture) {
fixture
.cache
.set_cache("a", entry(json!("A")), &no_ttl())
.unwrap();
fixture
.cache
.set_cache("c", entry(json!("C")), &no_ttl())
.unwrap();
fixture.service.seed_blob("bad", b"nope");
let keys = ["c", "missing", "a", "bad"].map(String::from);
let sync = fixture.cache.batch_get_cache(&keys, &no_ttl()).unwrap();
assert_eq!(
sync,
vec![
BatchEntry::Hit(entry(json!("C"))),
BatchEntry::Miss,
BatchEntry::Hit(entry(json!("A"))),
BatchEntry::Invalid,
]
);
let asynchronous = fixture
.runtime
.block_on(fixture.cache.async_batch_get_cache(keys.to_vec(), no_ttl()))
.unwrap();
assert_eq!(asynchronous, sync);
let response_cache = response_cache(&fixture);
let requests = [request("hit"), request("missing"), request("bad")];
response_cache
.store(&requests[0], json!("HIT"), now())
.unwrap();
fixture
.service
.seed_blob(&cache_key(&requests[2].key), b"nope");
let hits = response_cache.lookup_batch(&requests, now()).unwrap();
assert_eq!(hits.values, vec![Some(json!("HIT")), None, None]);
assert_eq!(hits.missing_indices, vec![1, 2]);
let async_hits = fixture
.runtime
.block_on(response_cache.async_lookup_batch(&requests, now()))
.unwrap();
assert_eq!(async_hits.values, hits.values);
}
#[rstest]
fn async_pipeline_writes_every_entry_with_overwrite(fixture: Fixture) {
fixture.service.seed_blob("k2", b"stale");
fixture
.runtime
.block_on(fixture.cache.async_set_cache_pipeline(
vec![
("k1".into(), entry(json!({"n": 1}))),
("k2".into(), entry(json!({"n": 2}))),
("k3".into(), entry(json!({"n": 3}))),
],
with_ttl(30),
))
.unwrap();
assert_eq!(fixture.service.blob_names(), ["k1", "k2", "k3"]);
assert_eq!(fixture.stored_json("k2")["response"], json!({"n": 2}));
}
#[rstest]
fn flush_deletes_every_blob_in_the_container(fixture: Fixture) {
for key in ["x", "y", "z"] {
fixture
.cache
.set_cache(key, entry(json!(key)), &no_ttl())
.unwrap();
}
fixture.cache.flush_cache().unwrap();
assert!(fixture.service.blob_names().is_empty());
assert!(fixture.service.container_exists());
fixture
.cache
.set_cache("again", entry(json!(1)), &no_ttl())
.unwrap();
fixture
.runtime
.block_on(fixture.cache.async_flush_cache())
.unwrap();
assert!(fixture.service.blob_names().is_empty());
}
#[rstest]
fn service_failures_map_to_unavailable(fixture: Fixture) {
fixture.service.set_failing(true);
assert!(matches!(
fixture.cache.get_cache("key", &no_ttl()),
Err(Error::Unavailable)
));
assert!(matches!(
fixture.cache.set_cache("key", entry(json!(1)), &no_ttl()),
Err(Error::Unavailable)
));
assert!(matches!(
fixture.cache.flush_cache(),
Err(Error::Unavailable)
));
assert!(matches!(
fixture.runtime.block_on(
fixture
.cache
.async_set_cache_pipeline(vec![("k".into(), entry(json!(1)))], no_ttl())
),
Err(Error::Unavailable)
));
}
#[rstest]
fn disconnect_is_idempotent_and_keeps_data(fixture: Fixture) {
fixture
.cache
.set_cache("key", entry(json!(1)), &no_ttl())
.unwrap();
fixture.runtime.block_on(async {
fixture.cache.disconnect().await.unwrap();
fixture.cache.disconnect().await.unwrap();
});
assert_eq!(
fixture.cache.get_cache("key", &no_ttl()).unwrap(),
Some(entry(json!(1)))
);
}
#[rstest]
fn response_cache_stores_and_reads_through_the_backend(fixture: Fixture) {
let response_cache = response_cache(&fixture);
let mut request = request("gpt");
request.context = with_ttl(60);
let response = json!({"id": "chatcmpl-1"});
response_cache
.store(&request, response.clone(), now())
.unwrap();
assert_eq!(
fixture.stored_json(&cache_key(&request.key)),
json!({"timestamp": 1_700_000_000.0, "response": {"id": "chatcmpl-1"}})
);
assert_eq!(
response_cache
.lookup(&request, now() + Duration::from_secs(3600))
.unwrap(),
Some(response.clone())
);
assert_eq!(
fixture
.runtime
.block_on(response_cache.async_lookup(&request, now() + Duration::from_secs(3600)))
.unwrap(),
Some(response.clone())
);
fixture.runtime.block_on(async {
response_cache
.async_store(&request, json!("replaced"), now())
.await
.unwrap();
assert_eq!(
response_cache.async_lookup(&request, now()).await.unwrap(),
Some(json!("replaced"))
);
response_cache.async_flush().await.unwrap();
assert_eq!(
response_cache.async_lookup(&request, now()).await.unwrap(),
None
);
});
}
#[rstest]
fn non_object_responses_are_written_serialized_like_python(fixture: Fixture) {
fixture
.cache
.set_cache("s", entry(json!("plain")), &no_ttl())
.unwrap();
assert_eq!(
fixture.stored_json("s"),
json!({"timestamp": 1_700_000_000.5, "response": "\"plain\""})
);
assert_eq!(
fixture.cache.get_cache("s", &no_ttl()).unwrap(),
Some(entry(json!("plain")))
);
}
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn sync_methods_block_inside_a_multi_thread_runtime() {
let service = FakeBlobService::default();
let cache = support::connect(
&service,
ACCOUNT_URL,
ResponseCacheCodec,
tokio::runtime::Handle::current(),
)
.await
.map(Arc::new)
.unwrap();
cache.set_cache("key", entry(json!(1)), &no_ttl()).unwrap();
assert_eq!(
cache.get_cache("key", &no_ttl()).unwrap(),
Some(entry(json!(1)))
);
}

81
litellm-rust/crates/cache-azure-blob/tests/contract.rs Normal file
View file

@ -0,0 +1,81 @@
mod support;
use litellm_cache::{ExactCacheContext, JsonCodec};
use litellm_cache_azure_blob::AzureBlobCache;
use litellm_cache_testing as contract;
use rstest::{fixture, rstest};
use serde_json::{Value, json};
use support::{ACCOUNT_URL, FakeBlobService};
use tokio::runtime::Handle;
#[fixture]
async fn azure() -> AzureBlobCache<JsonCodec<Value>> {
support::connect(
&FakeBlobService::default(),
ACCOUNT_URL,
JsonCodec::new(),
Handle::current(),
)
.await
.unwrap()
}
#[fixture]
fn context() -> ExactCacheContext {
ExactCacheContext::default()
}
const PREFIX: &str = "contract:";
// `overwrite_replaces` does not apply: sync `set_cache` never overwrites a blob, as in Python.
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn hit_and_miss(
#[future(awt)] azure: AzureBlobCache<JsonCodec<Value>>,
context: ExactCacheContext,
) {
contract::hit_and_miss(&azure, context, PREFIX, json!({"answer": 42})).await;
}
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn sync_async_equivalence(
#[future(awt)] azure: AzureBlobCache<JsonCodec<Value>>,
context: ExactCacheContext,
) {
contract::sync_async_equivalence(&azure, context, PREFIX, json!("first"), json!([2])).await;
}
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn pipeline_writes_every_entry(
#[future(awt)] azure: AzureBlobCache<JsonCodec<Value>>,
context: ExactCacheContext,
) {
contract::pipeline_writes_every_entry(
&azure,
context,
PREFIX,
vec![json!("a"), json!(2), json!({"c": true})],
)
.await;
}
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn batch_preserves_order(
#[future(awt)] azure: AzureBlobCache<JsonCodec<Value>>,
context: ExactCacheContext,
) {
contract::batch_preserves_order(&azure, context, PREFIX, json!("first"), json!(2)).await;
}
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn flush_clears(
#[future(awt)] azure: AzureBlobCache<JsonCodec<Value>>,
context: ExactCacheContext,
) {
contract::flush_clears(&azure, context, PREFIX, json!("value")).await;
}

239
litellm-rust/crates/cache-azure-blob/tests/support/mod.rs Normal file
View file

@ -0,0 +1,239 @@
#![allow(dead_code)]
use std::{
collections::BTreeMap,
sync::{Arc, Mutex},
};
use azure_core::http::{
AsyncRawResponse, Body, ClientOptions, HttpClient, Method, Request, StatusCode, Transport,
headers::{HeaderName, Headers},
};
use litellm_cache::{CacheCodec, Error};
use litellm_cache_azure_blob::AzureBlobCache;
use tokio::runtime::{Handle, Runtime};
pub const ACCOUNT_URL: &str = "https://example.blob.core.windows.net";
pub const CONTAINER: &str = "litellm-cache";
const IF_NONE_MATCH: HeaderName = HeaderName::from_static("if-none-match");
const ERROR_CODE: HeaderName = HeaderName::from_static("x-ms-error-code");
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct RecordedRequest {
pub method: Method,
pub path: String,
pub query: String,
pub if_none_match: Option<String>,
}
#[derive(Default)]
struct FakeState {
container_exists: bool,
blobs: BTreeMap<String, Vec<u8>>,
requests: Vec<RecordedRequest>,
failing: bool,
precondition_conflicts: bool,
}
#[derive(Clone, Default)]
pub struct FakeBlobService {
state: Arc<Mutex<FakeState>>,
}
impl std::fmt::Debug for FakeBlobService {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str("FakeBlobService")
}
}
impl FakeBlobService {
pub fn with_existing_container() -> Self {
let service = Self::default();
service.state.lock().unwrap().container_exists = true;
service
}
pub fn blob(&self, name: &str) -> Option<Vec<u8>> {
self.state.lock().unwrap().blobs.get(name).cloned()
}
pub fn blob_names(&self) -> Vec<String> {
self.state.lock().unwrap().blobs.keys().cloned().collect()
}
pub fn seed_blob(&self, name: &str, bytes: &[u8]) {
self.state
.lock()
.unwrap()
.blobs
.insert(name.to_string(), bytes.to_vec());
}
pub fn set_failing(&self, failing: bool) {
self.state.lock().unwrap().failing = failing;
}
pub fn set_precondition_conflicts(&self, enabled: bool) {
self.state.lock().unwrap().precondition_conflicts = enabled;
}
pub fn requests(&self) -> Vec<RecordedRequest> {
self.state.lock().unwrap().requests.clone()
}
pub fn container_exists(&self) -> bool {
self.state.lock().unwrap().container_exists
}
fn respond(status: StatusCode, error_code: Option<&str>, body: Vec<u8>) -> AsyncRawResponse {
let mut headers = Headers::new();
if let Some(code) = error_code {
headers.insert(ERROR_CODE, code.to_string());
}
AsyncRawResponse::from_bytes(status, headers, body)
}
fn list_body(state: &FakeState) -> Vec<u8> {
let mut xml = String::from(
r#"<?xml version="1.0" encoding="utf-8"?><EnumerationResults ServiceEndpoint="https://example.blob.core.windows.net/" ContainerName="litellm-cache"><Blobs>"#,
);
for name in state.blobs.keys() {
xml.push_str(&format!(
"<Blob><Name>{name}</Name><Properties><BlobType>BlockBlob</BlobType></Properties></Blob>"
));
}
xml.push_str("</Blobs><NextMarker /></EnumerationResults>");
xml.into_bytes()
}
}
#[async_trait::async_trait]
impl HttpClient for FakeBlobService {
async fn execute_request(&self, request: &Request) -> azure_core::Result<AsyncRawResponse> {
let mut state = self.state.lock().unwrap();
let path = request.url().path().to_string();
let query = request.url().query().unwrap_or_default().to_string();
let if_none_match = request
.headers()
.get_optional_str(&IF_NONE_MATCH)
.map(str::to_owned);
state.requests.push(RecordedRequest {
method: request.method(),
path: path.clone(),
query: query.clone(),
if_none_match: if_none_match.clone(),
});
if state.failing {
return Ok(Self::respond(
StatusCode::Forbidden,
Some("AuthorizationFailure"),
Vec::new(),
));
}
let container_path = format!("/{CONTAINER}");
let blob_name = path
.strip_prefix(&format!("{container_path}/"))
.map(str::to_owned);
let is_container = path == container_path && query.contains("restype=container");
let response = match (request.method(), is_container, blob_name) {
(Method::Put, true, None) if state.container_exists => Self::respond(
StatusCode::Conflict,
Some("ContainerAlreadyExists"),
Vec::new(),
),
(Method::Put, true, None) => {
state.container_exists = true;
Self::respond(StatusCode::Created, None, Vec::new())
}
(Method::Get, true, None) if query.contains("comp=list") => {
Self::respond(StatusCode::Ok, None, Self::list_body(&state))
}
(Method::Get, true, None) if state.container_exists => {
Self::respond(StatusCode::Ok, None, Vec::new())
}
(Method::Get, true, None) => {
Self::respond(StatusCode::NotFound, Some("ContainerNotFound"), Vec::new())
}
(Method::Put, false, Some(name)) => {
if if_none_match.as_deref() == Some("*") && state.blobs.contains_key(&name) {
if state.precondition_conflicts {
Self::respond(
StatusCode::PreconditionFailed,
Some("ConditionNotMet"),
Vec::new(),
)
} else {
Self::respond(StatusCode::Conflict, Some("BlobAlreadyExists"), Vec::new())
}
} else {
let bytes = match request.body() {
Body::Bytes(bytes) => bytes.to_vec(),
Body::SeekableStream(_) => panic!("unexpected streaming upload"),
};
state.blobs.insert(name, bytes);
Self::respond(StatusCode::Created, None, Vec::new())
}
}
(Method::Get, false, Some(name)) => match state.blobs.get(&name) {
Some(bytes) => Self::respond(StatusCode::Ok, None, bytes.clone()),
None => Self::respond(StatusCode::NotFound, Some("BlobNotFound"), Vec::new()),
},
(Method::Delete, false, Some(name)) => match state.blobs.remove(&name) {
Some(_) => Self::respond(StatusCode::Accepted, None, Vec::new()),
None => Self::respond(StatusCode::NotFound, Some("BlobNotFound"), Vec::new()),
},
(method, _, _) => panic!("unexpected request {method:?} {path}?{query}"),
};
Ok(response)
}
}
pub async fn connect<C: CacheCodec>(
service: &FakeBlobService,
account_url: &str,
codec: C,
handle: Handle,
) -> Result<AzureBlobCache<C>, Error> {
AzureBlobCache::connect_with_options(
account_url,
CONTAINER,
None,
ClientOptions {
transport: Some(Transport::new(Arc::new(service.clone()))),
..ClientOptions::default()
},
codec,
handle,
)
.await
}
/// A cache on its own fake service and runtime, so sync methods run outside any runtime.
pub struct Fixture<C> {
pub runtime: Runtime,
pub service: FakeBlobService,
pub cache: Arc<AzureBlobCache<C>>,
}
impl<C: CacheCodec> Fixture<C> {
pub fn new(service: FakeBlobService, codec: C) -> Self {
let runtime = Runtime::new().unwrap();
let cache = runtime
.block_on(connect(
&service,
ACCOUNT_URL,
codec,
runtime.handle().clone(),
))
.unwrap();
Self {
runtime,
service,
cache: Arc::new(cache),
}
}
pub fn stored_json(&self, key: &str) -> serde_json::Value {
serde_json::from_slice(&self.service.blob(key).expect("blob should exist")).unwrap()
}
}

90
litellm-rust/crates/cache-azure-blob/tests/transport.rs Normal file
View file

@ -0,0 +1,90 @@
use std::sync::Arc;
use azure_core::http::{ClientOptions, Transport};
use litellm_cache::{BaseCache, ExactCacheContext, JsonCodec};
use litellm_cache_azure_blob::{AzureBlobCache, ReqwestTransport};
use rstest::{fixture, rstest};
use serde_json::{Value, json};
use tokio::runtime::Handle;
use wiremock::{
Mock, MockServer, ResponseTemplate,
matchers::{body_json, header, method, path, query_param},
};
#[fixture]
async fn server() -> MockServer {
let server = MockServer::start().await;
Mock::given(method("PUT"))
.and(path("/litellm-cache"))
.and(query_param("restype", "container"))
.respond_with(ResponseTemplate::new(201))
.expect(1)
.mount(&server)
.await;
server
}
async fn connect(server: &MockServer) -> AzureBlobCache<JsonCodec<Value>> {
AzureBlobCache::connect_with_options(
&server.uri(),
"litellm-cache",
None,
ClientOptions {
transport: Some(Transport::new(Arc::new(ReqwestTransport(
reqwest::Client::new(),
)))),
..ClientOptions::default()
},
JsonCodec::new(),
Handle::current(),
)
.await
.unwrap()
}
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn uploads_go_through_the_host_client(#[future(awt)] server: MockServer) {
Mock::given(method("PUT"))
.and(path("/litellm-cache/key"))
.and(header("if-none-match", "*"))
.and(body_json(json!({"answer": 1})))
.respond_with(ResponseTemplate::new(201))
.expect(1)
.mount(&server)
.await;
connect(&server)
.await
.set_cache("key", json!({"answer": 1}), &ExactCacheContext::default())
.unwrap();
}
#[rstest]
#[case::hit(
ResponseTemplate::new(200).set_body_json(json!({"answer": 2})),
Some(json!({"answer": 2}))
)]
#[case::blob_not_found(
ResponseTemplate::new(404).insert_header("x-ms-error-code", "BlobNotFound"),
None
)]
#[tokio::test(flavor = "multi_thread")]
async fn downloads_map_the_host_client_response(
#[future(awt)] server: MockServer,
#[case] response: ResponseTemplate,
#[case] expected: Option<Value>,
) {
Mock::given(method("GET"))
.and(path("/litellm-cache/key"))
.respond_with(response)
.mount(&server)
.await;
assert_eq!(
connect(&server)
.await
.async_get_cache("key", &ExactCacheContext::default())
.await
.unwrap(),
expected
);
}

1
litellm-rust/crates/cache-disk/Cargo.toml
View file

@ -15,5 +15,6 @@ serde_json.workspace = true
tokio.workspace = true
[dev-dependencies]
litellm-cache-testing.workspace = true
rstest.workspace = true
tempfile = "3.27.0"

38
litellm-rust/crates/cache-disk/src/cache.rs
View file

@ -5,8 +5,8 @@ use std::{
};
use litellm_cache::{
BaseCache, BatchCache, BatchEntry, CacheCodec, CacheConnectionResult, CacheConnectionStatus,
CounterCache, DeleteCache, Error, ExactCacheContext, FlushCache,
BaseCache, BatchCache, BatchEntry, CacheCodec, CounterCache, DeleteCache, DisconnectCache,
Error, ExactCacheContext, FlushCache,
};
use crate::{DiskStore, DiskcacheSqliteStore, PythonDiskCacheAdapter, StoredValue, ValueAdapter};
@ -150,29 +150,6 @@ impl<S: CacheCodec, D: DiskStore, A: ValueAdapter> BaseCache for DiskCache<S, D,
})
.await
}
async fn disconnect(&self) -> Result<(), Error> {
Ok(())
}
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
let result = Self::run_blocking(Arc::clone(&self.store), |store| {
store.probe().map(|_| CacheConnectionResult {
status: CacheConnectionStatus::Success,
message: "Disk cache connection test successful".into(),
error: None,
})
})
.await;
Ok(match result {
Ok(result) => result,
Err(error) => CacheConnectionResult {
status: CacheConnectionStatus::Failed,
message: format!("Disk cache connection failed: {error}"),
error: Some(error.to_string()),
},
})
}
}
impl<S: CacheCodec, D: DiskStore, A: ValueAdapter> BatchCache for DiskCache<S, D, A> {
@ -241,9 +218,13 @@ impl<S: CacheCodec, D: DiskStore, A: ValueAdapter> FlushCache for DiskCache<S, D
}
}
impl<S: CacheCodec<Value = f64>, D: DiskStore, A: ValueAdapter> CounterCache
for DiskCache<S, D, A>
{
impl<S: CacheCodec, D: DiskStore, A: ValueAdapter> DisconnectCache for DiskCache<S, D, A> {
async fn disconnect(&self) -> Result<(), Error> {
Ok(())
}
}
impl<S: CacheCodec, D: DiskStore, A: ValueAdapter> CounterCache for DiskCache<S, D, A> {
fn increment_cache(
&self,
key: &str,
@ -264,6 +245,7 @@ impl<S: CacheCodec<Value = f64>, D: DiskStore, A: ValueAdapter> CounterCache
key: &str,
amount: f64,
context: ExactCacheContext,
_refresh_ttl: bool,
) -> Result<f64, Error> {
let key = key.to_string();
let adapter = Arc::clone(&self.adapter);

12
litellm-rust/crates/cache-disk/src/sqlite.rs
View file

@ -544,18 +544,6 @@ impl DiskStore for DiskcacheSqliteStore {
}
}
}
fn probe(&self) -> Result<(), Error> {
let connection = self.connection.lock().map_err(|_| Error::Unavailable)?;
connection
.query_row(
"SELECT value FROM Settings WHERE key = 'count'",
[],
|row| row.get::<_, i64>(0),
)
.map(|_| ())
.map_err(|_| Error::Unavailable)
}
}
fn default_settings() -> HashMap<String, Value> {

1
litellm-rust/crates/cache-disk/src/store.rs
View file

@ -29,5 +29,4 @@ pub trait DiskStore: Send + Sync + 'static {
now: f64,
apply: &mut dyn FnMut(Option<StoredValue>) -> Result<(StoredValue, Option<f64>), Error>,
) -> Result<(), Error>;
fn probe(&self) -> Result<(), Error>;
}

103
litellm-rust/crates/cache-disk/tests/cache.rs
View file

@ -7,8 +7,8 @@ use std::{
};
use litellm_cache::{
BaseCache, BatchCache, BatchEntry, CacheCodec, CounterCache, DeleteCache, ExactCacheContext,
FlushCache, JsonCodec,
BaseCache, BatchCache, BatchEntry, CacheCodec, CounterCache, DeleteCache, DisconnectCache,
ExactCacheContext, FlushCache, JsonCodec,
};
use litellm_cache_disk::{DiskCache, DiskStore, DiskcacheSqliteStore, StoredValue, ValueAdapter};
use rstest::{fixture, rstest};
@ -395,7 +395,7 @@ fn delete_flush_and_spilled_file_replacement_clean_up_storage(sandbox: Sandbox)
#[rstest]
#[tokio::test]
async fn async_operations_connection_and_delete_match_sync_operations(sandbox: Sandbox) {
async fn async_operations_disconnect_and_delete_match_sync_operations(sandbox: Sandbox) {
let cache = sandbox.cache::<Value>();
let context = ExactCacheContext {
ttl: Some(Duration::from_secs(60)),
@ -424,8 +424,101 @@ async fn async_operations_connection_and_delete_match_sync_operations(sandbox: S
);
cache.async_delete_cache("a").await.unwrap();
cache.async_flush_cache().await.unwrap();
cache.disconnect().await.unwrap();
}
#[derive(Clone, Copy, Debug)]
enum Increment {
Sync,
Async { refresh_ttl: bool },
}
impl Increment {
async fn apply(
self,
cache: &DiskCache<JsonCodec<Value>>,
key: &str,
amount: f64,
context: ExactCacheContext,
) -> f64 {
match self {
Self::Sync => cache.increment_cache(key, amount, context).unwrap(),
Self::Async { refresh_ttl } => cache
.async_increment(key, amount, context, refresh_ttl)
.await
.unwrap(),
}
}
}
#[rstest]
#[case::sync_missing(Increment::Sync, None, 3.0, 3.0)]
#[case::sync_existing_int(Increment::Sync, Some(json!(7)), 5.0, 12.0)]
#[case::sync_non_int(Increment::Sync, Some(json!("not-a-number")), 4.0, 4.0)]
#[case::async_missing(Increment::Async { refresh_ttl: false }, None, 2.0, 2.0)]
#[case::async_existing_int(Increment::Async { refresh_ttl: false }, Some(json!(10)), 5.0, 15.0)]
#[case::async_non_int(Increment::Async { refresh_ttl: false }, Some(json!("corrupt")), 9.0, 9.0)]
#[case::async_refresh_ttl_is_ignored(Increment::Async { refresh_ttl: true }, Some(json!(1)), 1.0, 2.0)]
#[tokio::test]
async fn increments_read_back_through_get_cache(
sandbox: Sandbox,
#[case] increment: Increment,
#[case] initial: Option<Value>,
#[case] amount: f64,
#[case] expected: f64,
) {
let cache = sandbox.cache::<Value>();
let context = ExactCacheContext::default();
if let Some(initial) = initial {
cache
.async_set_cache("counter", initial, context.clone())
.await
.unwrap();
}
assert_eq!(
cache.test_connection().await.unwrap().status,
litellm_cache::CacheConnectionStatus::Success
increment
.apply(&cache, "counter", amount, context.clone())
.await,
expected
);
assert_eq!(
cache.get_cache("counter", &context).unwrap(),
Some(json!(expected as i64))
);
}
#[rstest]
#[case::without_refresh(false)]
#[case::with_refresh(true)]
#[tokio::test]
async fn async_increment_rewrites_ttl_on_every_write(sandbox: Sandbox, #[case] refresh_ttl: bool) {
let cache = sandbox.cache::<Value>();
let expiry = || {
sandbox
.db()
.query_row(
"SELECT expire_time IS NOT NULL FROM Cache WHERE key = 'counter'",
[],
|row| row.get::<_, bool>(0),
)
.unwrap()
};
let ttl = ExactCacheContext {
ttl: Some(Duration::from_secs(60)),
};
cache
.async_increment("counter", 1.0, ttl.clone(), refresh_ttl)
.await
.unwrap();
assert!(expiry());
cache
.async_increment("counter", 1.0, ExactCacheContext::default(), refresh_ttl)
.await
.unwrap();
assert!(!expiry());
cache
.async_increment("counter", 1.0, ttl, refresh_ttl)
.await
.unwrap();
assert!(expiry());
}

82
litellm-rust/crates/cache-disk/tests/contract.rs Normal file
View file

@ -0,0 +1,82 @@
use litellm_cache::{ExactCacheContext, JsonCodec};
use litellm_cache_disk::DiskCache;
use litellm_cache_testing as contract;
use rstest::{fixture, rstest};
use serde_json::{Value, json};
use tempfile::TempDir;
struct Disk {
cache: DiskCache<JsonCodec<Value>>,
_directory: TempDir,
}
#[fixture]
fn disk() -> Disk {
let directory = tempfile::tempdir().unwrap();
Disk {
cache: DiskCache::open(directory.path(), JsonCodec::new()).unwrap(),
_directory: directory,
}
}
#[fixture]
fn context() -> ExactCacheContext {
ExactCacheContext::default()
}
const PREFIX: &str = "contract:";
#[rstest]
#[tokio::test]
async fn hit_and_miss(disk: Disk, context: ExactCacheContext) {
contract::hit_and_miss(&disk.cache, context, PREFIX, json!({"answer": 42})).await;
}
#[rstest]
#[tokio::test]
async fn sync_async_equivalence(disk: Disk, context: ExactCacheContext) {
contract::sync_async_equivalence(&disk.cache, context, PREFIX, json!("first"), json!([2]))
.await;
}
#[rstest]
#[tokio::test]
async fn overwrite_replaces(disk: Disk, context: ExactCacheContext) {
contract::overwrite_replaces(&disk.cache, context, PREFIX, json!(1), json!({"b": 2})).await;
}
#[rstest]
#[tokio::test]
async fn pipeline_writes_every_entry(disk: Disk, context: ExactCacheContext) {
contract::pipeline_writes_every_entry(
&disk.cache,
context,
PREFIX,
vec![json!("a"), json!(2), json!({"c": true})],
)
.await;
}
#[rstest]
#[tokio::test]
async fn batch_preserves_order(disk: Disk, context: ExactCacheContext) {
contract::batch_preserves_order(&disk.cache, context, PREFIX, json!("first"), json!(2)).await;
}
#[rstest]
#[tokio::test]
async fn delete_removes_key(disk: Disk, context: ExactCacheContext) {
contract::delete_removes_key(&disk.cache, context, PREFIX, json!("value")).await;
}
#[rstest]
#[tokio::test]
async fn flush_clears(disk: Disk, context: ExactCacheContext) {
contract::flush_clears(&disk.cache, context, PREFIX, json!("value")).await;
}
#[rstest]
#[tokio::test]
async fn counter_accumulates(disk: Disk, context: ExactCacheContext) {
contract::counter_accumulates(&disk.cache, context, PREFIX).await;
}

2
litellm-rust/crates/cache-gcs/Cargo.toml
View file

@ -15,6 +15,8 @@ reqwest.workspace = true
tokio.workspace = true
[dev-dependencies]
litellm-cache-testing.workspace = true
rstest.workspace = true
serde_json.workspace = true
tokio.workspace = true
wiremock = "0.6.5"

36
litellm-rust/crates/cache-gcs/src/cache.rs
View file

@ -2,7 +2,7 @@ use std::{future::Future, sync::Arc, time::Duration};
use futures_util::future::try_join_all;
use litellm_cache::{
BaseCache, BatchCache, BatchEntry, CacheCodec, CacheConnectionResult, Error, ExactCacheContext,
BaseCache, BatchCache, BatchEntry, CacheCodec, DisconnectCache, Error, ExactCacheContext,
FlushCache,
};
use percent_encoding::{AsciiSet, NON_ALPHANUMERIC, percent_encode};
@ -53,25 +53,25 @@ pub struct GcsCache<S: CacheCodec> {
}
impl<S: CacheCodec> GcsCache<S> {
pub fn new(config: GcsConfig, codec: S) -> Result<Self, Error> {
pub fn new(config: GcsConfig, client: Client, codec: S) -> Self {
let token = Arc::new(GcpTokenSource::new(config.path_service_account.clone()));
Self::with_token_source(config, codec, token)
Self::with_token_source(config, client, codec, token)
}
pub fn with_token_source(
config: GcsConfig,
client: Client,
codec: S,
token: Arc<dyn TokenSource>,
) -> Result<Self, Error> {
let client = Client::builder().build().map_err(|_| Error::Unavailable)?;
) -> Self {
let key_prefix = key_prefix(config.gcs_path.as_deref());
Ok(Self {
Self {
config,
key_prefix,
client,
token,
codec,
})
}
}
pub fn bucket_name(&self) -> &str {
@ -154,26 +154,26 @@ impl<S: CacheCodec> GcsCache<S> {
F: Future<Output = Result<T, Error>> + Send,
T: Send,
{
let run = || {
let run = |future: F| {
tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.map_err(|_| Error::Unavailable)
.and_then(|runtime| runtime.block_on(future))
};
if let Ok(handle) = tokio::runtime::Handle::try_current() {
if handle.runtime_flavor() == tokio::runtime::RuntimeFlavor::MultiThread {
return tokio::task::block_in_place(run);
match tokio::runtime::Handle::try_current() {
Ok(handle) if handle.runtime_flavor() == tokio::runtime::RuntimeFlavor::MultiThread => {
tokio::task::block_in_place(|| handle.block_on(future))
}
return std::thread::scope(|scope| {
Ok(_) => std::thread::scope(|scope| {
scope
.spawn(run)
.spawn(|| run(future))
.join()
.map_err(|_| Error::Unavailable)
.and_then(|result| result)
});
}),
Err(_) => run(future),
}
run()
}
}
@ -222,14 +222,12 @@ impl<S: CacheCodec> BaseCache for GcsCache<S> {
.await
.map(|_| ())
}
}
impl<S: CacheCodec> DisconnectCache for GcsCache<S> {
async fn disconnect(&self) -> Result<(), Error> {
Ok(())
}
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
Err(Error::UnsupportedOperation)
}
}
impl<S: CacheCodec> BatchCache for GcsCache<S> {

351
litellm-rust/crates/cache-gcs/tests/cache.rs
View file

@ -1,37 +1,32 @@
use std::{sync::Arc, time::Duration};
mod support;
use std::{future::Future, pin::Pin, sync::Arc, time::Duration};
use litellm_cache::{
BaseCache, BatchCache, BatchEntry, CacheContext, Error, ExactCacheContext, FlushCache,
JsonCodec,
BaseCache, BatchCache, BatchEntry, CacheContext, DisconnectCache, Error, ExactCacheContext,
FlushCache,
};
use litellm_cache_gcs::{GcsCache, GcsConfig, StaticTokenSource, TokenSource, key_prefix};
use serde_json::json;
use litellm_cache_gcs::{GcsCache, GcsConfig, TokenSource, key_prefix};
use rstest::{fixture, rstest};
use serde_json::{Value, json};
use support::FakeBucket;
use wiremock::{
Mock, MockServer, ResponseTemplate,
matchers::{body_bytes, header, method, path, query_param},
};
fn config(server: &MockServer, gcs_path: Option<&str>) -> GcsConfig {
GcsConfig {
bucket_name: "bucket".into(),
gcs_path: gcs_path.map(str::to_string),
path_service_account: None,
endpoint: server.uri(),
}
#[fixture]
async fn server() -> MockServer {
MockServer::start().await
}
fn cache(server: &MockServer, gcs_path: Option<&str>) -> GcsCache<JsonCodec<serde_json::Value>> {
GcsCache::with_token_source(
config(server, gcs_path),
JsonCodec::new(),
Arc::new(StaticTokenSource("tok".into())),
)
.unwrap()
fn context() -> ExactCacheContext {
ExactCacheContext::default()
}
#[rstest]
#[tokio::test]
async fn set_writes_encoded_object_and_headers() {
let server = MockServer::start().await;
async fn set_writes_encoded_object_and_headers(#[future(awt)] server: MockServer) {
Mock::given(method("POST"))
.and(path("/upload/storage/v1/b/bucket/o"))
.and(query_param("uploadType", "media"))
@ -42,12 +37,8 @@ async fn set_writes_encoded_object_and_headers() {
.expect(1)
.mount(&server)
.await;
cache(&server, Some("cache/"))
.set_cache(
"team:a b/c",
json!({"value": "entry"}),
&ExactCacheContext::default(),
)
support::cache(&server, Some("cache/"))
.set_cache("team:a b/c", json!({"value": "entry"}), &context())
.unwrap();
let requests = server.received_requests().await.unwrap();
assert_eq!(requests.len(), 1);
@ -57,103 +48,108 @@ async fn set_writes_encoded_object_and_headers() {
);
}
#[rstest]
#[case::hit(
"hit",
ResponseTemplate::new(200).set_body_json(json!({"value": "entry"})),
Ok(Some(json!({"value": "entry"})))
)]
#[case::missing("missing", ResponseTemplate::new(404), Ok(None))]
#[case::server_error("server-error", ResponseTemplate::new(500), Err(Error::Unavailable))]
#[case::invalid(
"invalid",
ResponseTemplate::new(200).set_body_string("not json"),
Err(Error::InvalidEntry)
)]
#[tokio::test]
async fn get_maps_statuses_and_decode_failures() {
let server = MockServer::start().await;
async fn get_maps_statuses_and_decode_failures(
#[future(awt)] server: MockServer,
#[case] key: &str,
#[case] response: ResponseTemplate,
#[case] expected: Result<Option<Value>, Error>,
) {
Mock::given(method("GET"))
.and(path("/storage/v1/b/bucket/o/hit"))
.and(path(format!("/storage/v1/b/bucket/o/{key}")))
.and(query_param("alt", "media"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"value": "entry"})))
.respond_with(response)
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/storage/v1/b/bucket/o/missing"))
.respond_with(ResponseTemplate::new(404))
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/storage/v1/b/bucket/o/server-error"))
.respond_with(ResponseTemplate::new(500))
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/storage/v1/b/bucket/o/invalid"))
.respond_with(ResponseTemplate::new(200).set_body_string("not json"))
.mount(&server)
.await;
let cache = cache(&server, None);
assert_eq!(
cache
.get_cache("hit", &ExactCacheContext::default())
.unwrap(),
Some(json!({"value": "entry"}))
);
assert_eq!(
cache
.get_cache("missing", &ExactCacheContext::default())
.unwrap(),
None
);
assert_eq!(
cache
.get_cache("server-error", &ExactCacheContext::default())
.unwrap_err(),
Error::Unavailable
);
assert_eq!(
cache
.get_cache("invalid", &ExactCacheContext::default())
.unwrap_err(),
Error::InvalidEntry
);
let cache = support::cache(&server, None);
assert_eq!(cache.get_cache(key, &context()), expected);
assert_eq!(cache.async_get_cache(key, &context()).await, expected);
}
#[test]
fn key_prefix_normalizes_paths() {
assert_eq!(key_prefix(None), "");
assert_eq!(key_prefix(Some("a/b/")), "a/b/");
assert_eq!(key_prefix(Some("a/b")), "a/b/");
assert_eq!(key_prefix(Some("")), "");
#[rstest]
#[case::none(None, "")]
#[case::trailing_slash(Some("a/b/"), "a/b/")]
#[case::no_trailing_slash(Some("a/b"), "a/b/")]
#[case::empty(Some(""), "")]
fn key_prefix_normalizes_paths(#[case] gcs_path: Option<&str>, #[case] expected: &str) {
assert_eq!(key_prefix(gcs_path), expected);
}
#[rstest]
#[tokio::test]
async fn object_names_use_python_quote_encoding() {
let server = MockServer::start().await;
async fn cache_exposes_its_configuration(#[future(awt)] server: MockServer) {
let cache = GcsCache::new(
GcsConfig {
path_service_account: Some("/secrets/sa.json".into()),
..support::config(&server, Some("folder"))
},
reqwest::Client::new(),
litellm_cache::JsonCodec::<Value>::new(),
);
assert_eq!(cache.bucket_name(), "bucket");
assert_eq!(cache.key_prefix(), "folder/");
assert_eq!(cache.path_service_account(), Some("/secrets/sa.json"));
assert_eq!(cache.object_name("k"), "folder/k");
}
#[rstest]
#[case::punctuation("a~b-c_d.e/f g%h", "uploadType=media&name=p%2Fa~b-c_d.e%2Ff%20g%25h")]
#[case::utf8("ключ", "uploadType=media&name=p%2F%D0%BA%D0%BB%D1%8E%D1%87")]
#[tokio::test]
async fn object_names_use_python_quote_encoding(
#[future(awt)] server: MockServer,
#[case] key: &str,
#[case] query: &str,
) {
Mock::given(method("POST"))
.and(path("/upload/storage/v1/b/bucket/o"))
.and(query_param("uploadType", "media"))
.respond_with(ResponseTemplate::new(200))
.expect(1)
.mount(&server)
.await;
support::cache(&server, Some("p/"))
.async_set_cache(key, json!({"value": key}), context())
.await
.unwrap();
let requests = server.received_requests().await.unwrap();
assert_eq!(requests[0].url.query(), Some(query));
}
#[rstest]
#[tokio::test]
async fn object_names_are_encoded_in_the_download_path(#[future(awt)] server: MockServer) {
Mock::given(method("GET"))
.and(path("/storage/v1/b/bucket/o/p%2Fa%3Ab%20c"))
.and(query_param("alt", "media"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!(1)))
.expect(2)
.mount(&server)
.await;
let cache = cache(&server, Some("p/"));
cache
.set_cache(
"a~b-c_d.e/f g%h",
json!({"value": "punctuation"}),
&ExactCacheContext::default(),
)
.unwrap();
cache
.set_cache(
"ключ",
json!({"value": "utf8"}),
&ExactCacheContext::default(),
)
.unwrap();
let requests = server.received_requests().await.unwrap();
let queries: Vec<_> = requests
.iter()
.filter_map(|request| request.url.query())
.collect();
assert!(queries.contains(&"uploadType=media&name=p%2Fa~b-c_d.e%2Ff%20g%25h"));
assert!(queries.contains(&"uploadType=media&name=p%2F%D0%BA%D0%BB%D1%8E%D1%87"));
let cache = support::cache(&server, Some("p"));
assert_eq!(cache.get_cache("a:b c", &context()), Ok(Some(json!(1))));
assert_eq!(
cache.async_get_cache("a:b c", &context()).await,
Ok(Some(json!(1)))
);
}
#[rstest]
#[tokio::test]
async fn ignores_ttl_and_writes_pipeline_concurrently() {
let server = MockServer::start().await;
async fn ignores_ttl_and_writes_pipeline_concurrently(#[future(awt)] server: MockServer) {
for key in ["one", "two", "three"] {
Mock::given(method("POST"))
.and(path("/upload/storage/v1/b/bucket/o"))
@ -164,12 +160,10 @@ async fn ignores_ttl_and_writes_pipeline_concurrently() {
.mount(&server)
.await;
}
let cache = cache(&server, None);
assert_eq!(cache.get_ttl(&ExactCacheContext::default()), None);
assert_eq!(
cache.get_ttl(&ExactCacheContext::default().with_ttl(Some(Duration::from_secs(5)))),
None
);
let cache = support::cache(&server, None);
let with_ttl = context().with_ttl(Some(Duration::from_secs(5)));
assert_eq!(cache.get_ttl(&context()), None);
assert_eq!(cache.get_ttl(&with_ttl), None);
cache
.async_set_cache_pipeline(
vec![
@ -177,15 +171,15 @@ async fn ignores_ttl_and_writes_pipeline_concurrently() {
("two".into(), json!({"key": "two"})),
("three".into(), json!({"key": "three"})),
],
ExactCacheContext::default().with_ttl(Some(Duration::from_secs(5))),
with_ttl,
)
.await
.unwrap();
}
#[rstest]
#[tokio::test]
async fn async_batch_get_preserves_hits_misses_and_invalid_entries() {
let server = MockServer::start().await;
async fn batch_get_preserves_hits_misses_and_invalid_entries(#[future(awt)] server: MockServer) {
Mock::given(method("GET"))
.and(path("/storage/v1/b/bucket/o/hit"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"value": "entry"})))
@ -201,122 +195,83 @@ async fn async_batch_get_preserves_hits_misses_and_invalid_entries() {
.respond_with(ResponseTemplate::new(200).set_body_string("not json"))
.mount(&server)
.await;
let cache = support::cache(&server, None);
let keys = vec!["hit".to_string(), "missing".into(), "invalid".into()];
let expected = vec![
BatchEntry::Hit(json!({"value": "entry"})),
BatchEntry::Miss,
BatchEntry::Invalid,
];
assert_eq!(cache.batch_get_cache(&keys, &context()).unwrap(), expected);
assert_eq!(
cache(&server, None)
.async_batch_get_cache(
vec!["hit".into(), "missing".into(), "invalid".into()],
ExactCacheContext::default(),
)
.await
.unwrap(),
vec![
BatchEntry::Hit(json!({"value": "entry"})),
BatchEntry::Miss,
BatchEntry::Invalid,
]
cache.async_batch_get_cache(keys, context()).await.unwrap(),
expected
);
}
#[rstest]
#[tokio::test]
async fn lifecycle_operations_are_noops_and_connection_test_is_unsupported() {
let server = MockServer::start().await;
let cache = cache(&server, None);
async fn flush_and_disconnect_are_noops_like_python(#[future(awt)] server: MockServer) {
let cache = support::cache(&server, None);
assert_eq!(cache.flush_cache(), Ok(()));
assert_eq!(cache.async_flush_cache().await, Ok(()));
assert_eq!(cache.disconnect().await, Ok(()));
assert_eq!(
cache.test_connection().await,
Err(Error::UnsupportedOperation)
);
assert!(server.received_requests().await.unwrap().is_empty());
}
#[test]
fn round_trip(cache: &support::JsonGcsCache) -> Result<Option<Value>, Error> {
cache.set_cache("key", json!({"value": "entry"}), &context())?;
cache.get_cache("key", &context())
}
#[rstest]
fn sync_operations_work_without_an_active_runtime() {
let runtime = tokio::runtime::Builder::new_multi_thread()
.enable_all()
.build()
.unwrap();
let server = runtime.block_on(MockServer::start());
runtime.block_on(
Mock::given(method("POST"))
.and(path("/upload/storage/v1/b/bucket/o"))
.respond_with(ResponseTemplate::new(200))
.mount(&server),
);
runtime.block_on(
Mock::given(method("GET"))
.and(path("/storage/v1/b/bucket/o/key"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"value": "entry"})))
.mount(&server),
);
let cache = cache(&server, None);
cache
.set_cache(
"key",
json!({"value": "entry"}),
&ExactCacheContext::default(),
)
.unwrap();
assert_eq!(
cache
.get_cache("key", &ExactCacheContext::default())
.unwrap(),
Some(json!({"value": "entry"}))
);
let server = runtime.block_on(FakeBucket::serve());
let cache = support::cache(&server, None);
assert_eq!(round_trip(&cache), Ok(Some(json!({"value": "entry"}))));
}
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn sync_operations_work_inside_a_multi_thread_runtime() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/upload/storage/v1/b/bucket/o"))
.respond_with(ResponseTemplate::new(200))
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/storage/v1/b/bucket/o/key"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"value": "entry"})))
.mount(&server)
.await;
let cache = cache(&server, None);
cache
.set_cache(
"key",
json!({"value": "entry"}),
&ExactCacheContext::default(),
)
.unwrap();
assert_eq!(
cache
.get_cache("key", &ExactCacheContext::default())
.unwrap(),
Some(json!({"value": "entry"}))
);
let server = FakeBucket::serve().await;
let cache = support::cache(&server, None);
assert_eq!(round_trip(&cache), Ok(Some(json!({"value": "entry"}))));
}
#[rstest]
#[tokio::test]
async fn sync_operations_work_inside_a_current_thread_runtime() {
let server = FakeBucket::serve().await;
let cache = support::cache(&server, None);
assert_eq!(round_trip(&cache), Ok(Some(json!({"value": "entry"}))));
}
struct FailingTokenSource;
impl TokenSource for FailingTokenSource {
fn bearer_token(
&self,
) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<String, Error>> + Send + '_>>
{
fn bearer_token(&self) -> Pin<Box<dyn Future<Output = Result<String, Error>> + Send + '_>> {
Box::pin(async { Err(Error::Unavailable) })
}
}
#[rstest]
#[tokio::test]
async fn token_source_failure_skips_http() {
let server = MockServer::start().await;
let cache = GcsCache::with_token_source(
config(&server, None),
JsonCodec::<serde_json::Value>::new(),
Arc::new(FailingTokenSource),
)
.unwrap();
async fn token_source_failure_skips_http(#[future(awt)] server: MockServer) {
let cache = support::cache_with_token(&server, None, Arc::new(FailingTokenSource));
assert_eq!(
cache.get_cache("key", &context()).unwrap_err(),
Error::Unavailable
);
assert_eq!(
cache
.get_cache("key", &ExactCacheContext::default())
.async_set_cache("key", json!(1), context())
.await
.unwrap_err(),
Error::Unavailable
);

65
litellm-rust/crates/cache-gcs/tests/contract.rs Normal file
View file

@ -0,0 +1,65 @@
mod support;
use litellm_cache::ExactCacheContext;
use litellm_cache_testing as contract;
use rstest::{fixture, rstest};
use serde_json::json;
use support::{FakeBucket, JsonGcsCache};
use wiremock::MockServer;
struct Gcs {
cache: JsonGcsCache,
_server: MockServer,
}
#[fixture]
async fn gcs() -> Gcs {
let server = FakeBucket::serve().await;
Gcs {
cache: support::cache(&server, Some("contract")),
_server: server,
}
}
#[fixture]
fn context() -> ExactCacheContext {
ExactCacheContext::default()
}
const PREFIX: &str = "contract:";
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn hit_and_miss(#[future(awt)] gcs: Gcs, context: ExactCacheContext) {
contract::hit_and_miss(&gcs.cache, context, PREFIX, json!({"answer": 42})).await;
}
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn sync_async_equivalence(#[future(awt)] gcs: Gcs, context: ExactCacheContext) {
contract::sync_async_equivalence(&gcs.cache, context, PREFIX, json!("first"), json!([2])).await;
}
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn overwrite_replaces(#[future(awt)] gcs: Gcs, context: ExactCacheContext) {
contract::overwrite_replaces(&gcs.cache, context, PREFIX, json!(1), json!({"b": 2})).await;
}
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn pipeline_writes_every_entry(#[future(awt)] gcs: Gcs, context: ExactCacheContext) {
contract::pipeline_writes_every_entry(
&gcs.cache,
context,
PREFIX,
vec![json!("a"), json!(2), json!({"c": true})],
)
.await;
}
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn batch_preserves_order(#[future(awt)] gcs: Gcs, context: ExactCacheContext) {
contract::batch_preserves_order(&gcs.cache, context, PREFIX, json!("first"), json!(2)).await;
}

87
litellm-rust/crates/cache-gcs/tests/support/mod.rs Normal file
View file

@ -0,0 +1,87 @@
#![allow(dead_code)]
use std::{
collections::HashMap,
sync::{Arc, Mutex},
};
use litellm_cache::JsonCodec;
use litellm_cache_gcs::{GcsCache, GcsConfig, StaticTokenSource, TokenSource};
use percent_encoding::percent_decode_str;
use serde_json::Value;
use wiremock::{Mock, MockServer, Request, Respond, ResponseTemplate, http::Method, matchers::any};
pub type JsonGcsCache = GcsCache<JsonCodec<Value>>;
pub fn config(server: &MockServer, gcs_path: Option<&str>) -> GcsConfig {
GcsConfig {
bucket_name: "bucket".into(),
gcs_path: gcs_path.map(str::to_string),
path_service_account: None,
endpoint: server.uri(),
}
}
pub fn cache_with_token(
server: &MockServer,
gcs_path: Option<&str>,
token: Arc<dyn TokenSource>,
) -> JsonGcsCache {
GcsCache::with_token_source(
config(server, gcs_path),
reqwest::Client::new(),
JsonCodec::new(),
token,
)
}
pub fn cache(server: &MockServer, gcs_path: Option<&str>) -> JsonGcsCache {
cache_with_token(server, gcs_path, Arc::new(StaticTokenSource("tok".into())))
}
/// An in-memory bucket speaking the JSON API's media upload and `alt=media` download.
#[derive(Clone, Default)]
pub struct FakeBucket {
objects: Arc<Mutex<HashMap<String, Vec<u8>>>>,
}
impl FakeBucket {
pub async fn serve() -> MockServer {
let server = MockServer::start().await;
Mock::given(any())
.respond_with(Self::default())
.mount(&server)
.await;
server
}
}
impl Respond for FakeBucket {
fn respond(&self, request: &Request) -> ResponseTemplate {
let mut objects = self.objects.lock().unwrap();
match request.method {
Method::POST => {
let name = request
.url
.query_pairs()
.find_map(|(key, value)| (key == "name").then(|| value.into_owned()))
.expect("uploads carry the object name");
objects.insert(name, request.body.clone());
ResponseTemplate::new(200)
}
Method::GET => {
let encoded = request
.url
.path()
.strip_prefix("/storage/v1/b/bucket/o/")
.expect("downloads address an object");
let name = percent_decode_str(encoded).decode_utf8().unwrap();
match objects.get(name.as_ref()) {
Some(body) => ResponseTemplate::new(200).set_body_bytes(body.clone()),
None => ResponseTemplate::new(404),
}
}
_ => ResponseTemplate::new(405),
}
}
}

2
litellm-rust/crates/cache-memory/Cargo.toml
View file

@ -9,6 +9,6 @@ repository.workspace = true
litellm-cache.workspace = true
[dev-dependencies]
serde_json.workspace = true
litellm-cache-testing.workspace = true
rstest.workspace = true
tokio.workspace = true

197
litellm-rust/crates/cache-memory/src/cache.rs
View file

@ -7,8 +7,8 @@ use std::{
};
use litellm_cache::{
BaseCache, BatchCache, CacheConnectionResult, CacheConnectionStatus, ClaimCache, CounterCache,
DeleteCache, Error, ExactCacheContext, FlushCache, IncrementOperation, SetCache, TtlCache,
BaseCache, BatchCache, ClaimCache, CounterCache, DeleteCache, DisconnectCache, Error,
ExactCacheContext, FlushCache, SetCache, TtlCache,
};
const DEFAULT_MAX_SIZE_IN_MEMORY: usize = 200;
@ -75,7 +75,9 @@ impl<V: Clone> InMemoryCache<V> {
expiration_heap: BinaryHeap::new(),
}),
max_size_in_memory: max_size_in_memory.unwrap_or(DEFAULT_MAX_SIZE_IN_MEMORY),
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
default_ttl: default_ttl
.filter(|ttl| !ttl.is_zero())
.unwrap_or(DEFAULT_TTL),
max_entry_bytes,
measure_value,
now: Arc::new(now),
@ -91,21 +93,9 @@ impl<V: Clone> InMemoryCache<V> {
if self.max_size_in_memory == 0 {
return Ok(CacheWrite::Disabled);
}
if let (Some(limit), Some(measure)) = (self.max_entry_bytes, &self.measure_value)
&& measure(&value)? > limit
{
return Ok(CacheWrite::TooLarge);
}
let now = (self.now)();
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
let key = key.into();
Self::evict(&mut state, self.max_size_in_memory, now, &key);
let expiration = state.expirations.get(&key).copied();
if expiration.is_none_or(|expiration| expiration < now) {
Self::set_expiration(&mut state, &key, now + ttl.unwrap_or(self.default_ttl));
}
state.values.insert(key, value);
Ok(CacheWrite::Stored)
self.store(&mut state, key.into(), value, ttl, now)
}
pub fn get_cache(&self, key: &str) -> Result<Option<V>, Error> {
@ -121,6 +111,70 @@ impl<V: Clone> InMemoryCache<V> {
Ok(state.values.get(key).cloned())
}
/// `check_value_size`: whether `value` fits `max_entry_bytes`. Always `true` without a
/// limit and a measure, since typed values have no generic size.
pub fn check_value_size(&self, value: &V) -> Result<bool, Error> {
match (self.max_entry_bytes, &self.measure_value) {
(Some(limit), Some(measure)) => Ok(measure(value)? <= limit),
_ => Ok(true),
}
}
/// `evict_cache`: drops expired entries, then the earliest-expiring ones until a new key
/// fits.
pub fn evict_cache(&self) -> Result<(), Error> {
let now = (self.now)();
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
Self::evict(&mut state, self.max_size_in_memory, now, None);
Ok(())
}
/// `evict_element_if_expired`: `true` when `key` had expired and was removed.
pub fn evict_element_if_expired(&self, key: &str) -> Result<bool, Error> {
let now = (self.now)();
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
let expired = state
.expirations
.get(key)
.is_some_and(|expiration| *expiration < now);
if expired {
Self::remove(&mut state, key);
}
Ok(expired)
}
/// `allow_ttl_override`: a write may set the TTL when the key has none or it has passed.
pub fn allow_ttl_override(&self, key: &str) -> Result<bool, Error> {
let now = (self.now)();
Ok(self
.expires_at(key)?
.is_none_or(|expiration| expiration < now))
}
/// The number of stored entries, expired ones included until they are evicted.
pub fn len(&self) -> Result<usize, Error> {
Ok(self
.state
.lock()
.map_err(|_| Error::Unavailable)?
.values
.len())
}
pub fn is_empty(&self) -> Result<bool, Error> {
Ok(self.len()? == 0)
}
/// Entries in the expiration heap, stale ones included; bounded by eviction.
pub fn expiration_heap_len(&self) -> Result<usize, Error> {
Ok(self
.state
.lock()
.map_err(|_| Error::Unavailable)?
.expiration_heap
.len())
}
pub fn max_size_in_memory(&self) -> usize {
self.max_size_in_memory
}
@ -172,7 +226,9 @@ impl<V: Clone> InMemoryCache<V> {
Ok(())
}
fn evict(state: &mut CacheState<V>, capacity: usize, now: Duration, key: &str) {
/// Writing an existing `key` never evicts another entry, unlike Python, which pops the
/// earliest-expiring entry whenever the cache is full.
fn evict(state: &mut CacheState<V>, capacity: usize, now: Duration, key: Option<&str>) {
while let Some(Reverse((expiration, key))) = state.expiration_heap.peek().cloned() {
if state.expirations.get(&key).copied() != Some(expiration) {
state.expiration_heap.pop();
@ -183,7 +239,7 @@ impl<V: Clone> InMemoryCache<V> {
break;
}
}
if state.values.contains_key(key) {
if key.is_some_and(|key| state.values.contains_key(key)) {
return;
}
while state.values.len() >= capacity {
@ -209,6 +265,40 @@ impl<V: Clone> InMemoryCache<V> {
state.values.remove(key);
state.expirations.remove(key);
}
/// `get_cache` under the held lock: an expired entry is removed and reads as missing.
fn live(state: &mut CacheState<V>, key: &str, now: Duration) -> Option<V> {
if state
.expirations
.get(key)
.is_some_and(|expiration| *expiration < now)
{
Self::remove(state, key);
}
state.values.get(key).cloned()
}
/// Python `set_cache` under the held lock: evict first (even when `key` already exists),
/// then skip oversized values, then write, keeping a live key's expiry.
fn store(
&self,
state: &mut CacheState<V>,
key: String,
value: V,
ttl: Option<Duration>,
now: Duration,
) -> Result<CacheWrite, Error> {
Self::evict(state, self.max_size_in_memory, now, None);
if !self.check_value_size(&value)? {
return Ok(CacheWrite::TooLarge);
}
let expiration = state.expirations.get(&key).copied();
if expiration.is_none_or(|expiration| expiration < now) {
Self::set_expiration(state, &key, now + ttl.unwrap_or(self.default_ttl));
}
state.values.insert(key, value);
Ok(CacheWrite::Stored)
}
}
impl<V> ClaimCache for InMemoryCache<V>
@ -227,7 +317,7 @@ where
}
let now = (self.now)();
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
Self::evict(&mut state, self.max_size_in_memory, now, key);
Self::evict(&mut state, self.max_size_in_memory, now, Some(key));
let existing = state
.values
.get(key)
@ -262,38 +352,12 @@ impl CounterCache for InMemoryCache<f64> {
}
let now = (self.now)();
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
Self::evict(&mut state, self.max_size_in_memory, now, key);
let value = state.values.get(key).copied().unwrap_or_default() + amount;
if !state.expirations.contains_key(key) {
Self::set_expiration(
&mut state,
key,
now + self.get_ttl(&context).unwrap_or(self.default_ttl),
);
}
state.values.insert(key.into(), value);
let value = Self::live(&mut state, key, now).unwrap_or_default() + amount;
self.store(&mut state, key.into(), value, self.get_ttl(&context), now)?;
Ok(value)
}
}
impl InMemoryCache<f64> {
pub async fn async_increment_pipeline(
&self,
operations: Vec<IncrementOperation>,
) -> Result<Vec<f64>, Error> {
operations
.into_iter()
.map(|operation| {
self.increment_cache(
&operation.key,
operation.amount,
ExactCacheContext { ttl: operation.ttl },
)
})
.collect()
}
}
impl<V: Clone + Send + Sync + 'static> BaseCache for InMemoryCache<V> {
type Value = V;
type Context = ExactCacheContext;
@ -315,18 +379,12 @@ impl<V: Clone + Send + Sync + 'static> BaseCache for InMemoryCache<V> {
fn get_cache(&self, key: &str, _: &ExactCacheContext) -> Result<Option<Self::Value>, Error> {
self.get_cache(key)
}
}
impl<V: Clone + Send + Sync + 'static> DisconnectCache for InMemoryCache<V> {
async fn disconnect(&self) -> Result<(), Error> {
Ok(())
}
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
Ok(CacheConnectionResult {
status: CacheConnectionStatus::Success,
message: "In-memory cache connection test successful".into(),
error: None,
})
}
}
impl<V: Clone + Send + Sync + 'static> BatchCache for InMemoryCache<V> {}
@ -367,34 +425,9 @@ where
}
let now = (self.now)();
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
Self::evict(&mut state, self.max_size_in_memory, now, key);
let mut stored = state.values.get(key).cloned().unwrap_or_default();
let mut stored = Self::live(&mut state, key, now).unwrap_or_default();
stored.extend(values.iter().cloned());
if let (Some(limit), Some(measure)) = (self.max_entry_bytes, &self.measure_value)
&& measure(&stored)? > limit
{
return Ok(values);
}
if !state.expirations.contains_key(key) {
Self::set_expiration(&mut state, key, now + ttl.unwrap_or(self.default_ttl));
}
state.values.insert(key.into(), stored);
self.store(&mut state, key.into(), stored, ttl, now)?;
Ok(values)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn repeated_increments_keep_one_heap_entry_per_expiration() {
let cache = InMemoryCache::<f64>::new(Some(4), None);
for _ in 0..100 {
cache
.increment_cache("counter", 1.0, ExactCacheContext::default())
.unwrap();
}
assert_eq!(cache.state.lock().unwrap().expiration_heap.len(), 1);
}
}

687
litellm-rust/crates/cache-memory/tests/cache.rs
View file

@ -8,111 +8,296 @@ use std::{
};
use litellm_cache::{
BaseCache, CacheBackend, CacheConnectionStatus, ClaimCache, CounterCache, DeleteCache, Error,
ExactCacheContext, IncrementOperation, SetCache, get_cache, set_cache,
BaseCache, BatchCache, BatchEntry, CacheBackend, ClaimCache, CounterCache, DeleteCache,
DisconnectCache, Error, ExactCacheContext, FlushCache, IncrementOperation, SetCache, TtlCache,
get_cache, set_cache,
};
use litellm_cache_memory::{CacheWrite, InMemoryCache};
use rstest::{fixture, rstest};
type Clock = Arc<AtomicU64>;
#[fixture]
fn clock() -> Arc<AtomicU64> {
fn clock() -> Clock {
Arc::new(AtomicU64::new(100))
}
fn cache(clock: Arc<AtomicU64>, capacity: usize) -> InMemoryCache<String> {
fn cache_with<V: Clone>(clock: &Clock, capacity: usize) -> InMemoryCache<V> {
let clock = clock.clone();
InMemoryCache::with_clock(Some(capacity), Some(Duration::from_secs(60)), move || {
Duration::from_secs(clock.load(Ordering::SeqCst))
Duration::from_millis(clock.load(Ordering::SeqCst) * 1000)
})
}
fn cache(clock: &Clock, capacity: usize) -> InMemoryCache<String> {
cache_with(clock, capacity)
}
fn at(clock: &Clock, seconds: u64) {
clock.store(seconds, Ordering::SeqCst);
}
fn secs(seconds: u64) -> Option<Duration> {
Some(Duration::from_secs(seconds))
}
fn ttl(seconds: u64) -> ExactCacheContext {
ExactCacheContext { ttl: secs(seconds) }
}
fn measured(capacity: usize) -> InMemoryCache<String> {
InMemoryCache::with_clock_and_size_measurement(
Some(capacity),
secs(60),
Some(4),
Some(Arc::new(|value: &String| {
if value.is_empty() {
return Err(Error::InvalidEntry);
}
Ok(value.len())
})),
|| Duration::from_secs(100),
)
}
#[rstest]
fn default_explicit_and_override_ttls_follow_python_rules(clock: Arc<AtomicU64>) {
let cache = cache(clock.clone(), 4);
fn default_explicit_and_override_ttls_follow_python_rules(clock: Clock) {
let cache = cache(&clock, 4);
cache.set_cache("key", "first".into(), None).unwrap();
assert_eq!(
cache.expires_at("key").unwrap(),
Some(Duration::from_secs(160))
);
cache
.set_cache("key", "second".into(), Some(Duration::from_secs(10)))
.unwrap();
assert_eq!(
cache.expires_at("key").unwrap(),
Some(Duration::from_secs(160))
);
clock.store(160, Ordering::SeqCst);
assert_eq!(cache.expires_at("key").unwrap(), secs(160));
cache.set_cache("key", "second".into(), secs(10)).unwrap();
assert_eq!(cache.expires_at("key").unwrap(), secs(160));
at(&clock, 160);
assert_eq!(cache.get_cache("key").unwrap(), Some("second".into()));
clock.store(161, Ordering::SeqCst);
at(&clock, 161);
assert_eq!(cache.get_cache("key").unwrap(), None);
cache
.set_cache("key", "third".into(), Some(Duration::from_secs(10)))
.unwrap();
assert_eq!(
cache.expires_at("key").unwrap(),
Some(Duration::from_secs(171))
);
assert_eq!(cache.expires_at("key").unwrap(), None);
cache.set_cache("key", "third".into(), secs(10)).unwrap();
assert_eq!(cache.expires_at("key").unwrap(), secs(171));
}
#[rstest]
fn write_at_expiry_boundary_refreshes_ttl(clock: Arc<AtomicU64>) {
let cache = cache(clock.clone(), 4);
cache
.set_cache("key", "first".into(), Some(Duration::from_secs(10)))
.unwrap();
clock.store(110, Ordering::SeqCst);
cache
.set_cache("key", "second".into(), Some(Duration::from_secs(10)))
.unwrap();
assert_eq!(
cache.expires_at("key").unwrap(),
Some(Duration::from_secs(120))
);
clock.store(115, Ordering::SeqCst);
#[case::unset(None, secs(600))]
#[case::zero_falls_back_like_python_or(Some(Duration::ZERO), secs(600))]
#[case::explicit(secs(5), secs(5))]
fn default_ttl_falls_back_to_ten_minutes(
#[case] default_ttl: Option<Duration>,
#[case] expected: Option<Duration>,
) {
let cache = InMemoryCache::<String>::with_clock(None, default_ttl, || Duration::ZERO);
assert_eq!(cache.get_ttl(&ExactCacheContext::default()), expected);
cache.set_cache("key", "value".into(), None).unwrap();
assert_eq!(cache.expires_at("key").unwrap(), expected);
assert_eq!(cache.max_size_in_memory(), 200);
}
#[rstest]
fn write_at_expiry_boundary_refreshes_ttl(clock: Clock) {
let cache = cache(&clock, 4);
cache.set_cache("key", "first".into(), secs(10)).unwrap();
at(&clock, 110);
cache.set_cache("key", "second".into(), secs(10)).unwrap();
assert_eq!(cache.expires_at("key").unwrap(), secs(120));
at(&clock, 115);
assert_eq!(cache.get_cache("key").unwrap(), Some("second".into()));
}
#[rstest]
fn capacity_evicts_earliest_and_ignores_stale_heap_entries(clock: Arc<AtomicU64>) {
let cache = cache(clock, 2);
cache
.set_cache("early", "a".into(), Some(Duration::from_secs(10)))
.unwrap();
cache
.set_cache("late", "b".into(), Some(Duration::from_secs(20)))
.unwrap();
fn expired_key_without_a_read_allows_a_ttl_override(clock: Clock) {
let cache = cache(&clock, 4);
cache.set_cache("key", "first".into(), secs(1)).unwrap();
assert_eq!(cache.allow_ttl_override("key"), Ok(false));
at(&clock, 102);
assert_eq!(cache.allow_ttl_override("key"), Ok(true));
cache.set_cache("key", "second".into(), secs(1)).unwrap();
assert_eq!(cache.expires_at("key").unwrap(), secs(103));
assert_eq!(cache.allow_ttl_override("missing"), Ok(true));
}
#[rstest]
fn capacity_evicts_earliest_and_ignores_stale_heap_entries(clock: Clock) {
let cache = cache(&clock, 2);
cache.set_cache("early", "a".into(), secs(10)).unwrap();
cache.set_cache("late", "b".into(), secs(20)).unwrap();
cache.delete_cache("early").unwrap();
cache
.set_cache("new", "c".into(), Some(Duration::from_secs(30)))
.unwrap();
cache.set_cache("new", "c".into(), secs(30)).unwrap();
assert_eq!(cache.get_cache("late").unwrap(), Some("b".into()));
cache
.set_cache("last", "d".into(), Some(Duration::from_secs(40)))
.unwrap();
cache.set_cache("last", "d".into(), secs(40)).unwrap();
assert_eq!(cache.get_cache("late").unwrap(), None);
}
#[test]
fn disabled_size_limited_and_validated_writes_are_observable() {
let cache = |capacity| {
InMemoryCache::with_clock_and_size_measurement(
Some(capacity),
Some(Duration::from_secs(60)),
Some(4),
Some(Arc::new(|value: &String| {
if value.is_empty() {
return Err(Error::InvalidEntry);
}
Ok(value.len())
})),
|| Duration::from_secs(100),
)
};
let disabled = cache(0);
#[rstest]
fn max_size_is_respected_when_every_item_has_a_long_ttl(clock: Clock) {
let cache = cache(&clock, 3);
for index in 0..3 {
at(&clock, 100 + index);
cache
.set_cache(
format!("key_{index}"),
format!("value_{index}"),
secs(86_400),
)
.unwrap();
}
assert_eq!(cache.len(), Ok(3));
cache
.set_cache("key_3", "value_3".into(), secs(86_400))
.unwrap();
assert_eq!(cache.len(), Ok(3));
assert_eq!(cache.get_cache("key_0").unwrap(), None);
assert_eq!(cache.expires_at("key_0").unwrap(), None);
for key in ["key_1", "key_2", "key_3"] {
assert!(cache.get_cache(key).unwrap().is_some(), "{key}");
}
}
#[rstest]
fn expired_items_are_evicted_before_live_ones(clock: Clock) {
let cache = cache(&clock, 3);
cache.set_cache("expired_1", "1".into(), secs(1)).unwrap();
cache.set_cache("expired_2", "2".into(), secs(1)).unwrap();
cache
.set_cache("long_lived", "3".into(), secs(86_400))
.unwrap();
assert_eq!(cache.len(), Ok(3));
at(&clock, 102);
cache
.set_cache("new_item", "4".into(), secs(86_400))
.unwrap();
assert_eq!(cache.len(), Ok(2));
assert_eq!(cache.get_cache("long_lived").unwrap(), Some("3".into()));
assert_eq!(cache.get_cache("new_item").unwrap(), Some("4".into()));
for key in ["expired_1", "expired_2"] {
assert_eq!(cache.expires_at(key).unwrap(), None, "{key}");
}
}
#[rstest]
fn injected_clock_controls_expiry_and_eviction(clock: Clock) {
let cache = cache(&clock, 2);
at(&clock, 0);
cache
.set_cache("first", "original".into(), secs(10))
.unwrap();
at(&clock, 9);
cache.set_cache("second", "survivor".into(), None).unwrap();
assert_eq!(cache.get_cache("first").unwrap(), Some("original".into()));
at(&clock, 11);
assert_eq!(cache.get_cache("first").unwrap(), None);
cache
.set_cache("third", "replacement".into(), None)
.unwrap();
assert_eq!(cache.get_cache("second").unwrap(), Some("survivor".into()));
at(&clock, 70);
cache.set_cache("fourth", "new".into(), None).unwrap();
assert_eq!(cache.get_cache("second").unwrap(), None);
assert_eq!(
disabled.set_cache("a", "x".into(), None).unwrap(),
cache.get_cache("third").unwrap(),
Some("replacement".into())
);
assert_eq!(cache.get_cache("fourth").unwrap(), Some("new".into()));
}
#[rstest]
fn rewriting_one_key_keeps_one_heap_entry(clock: Clock) {
let cache = cache(&clock, 10);
for index in 0..1_000 {
cache
.set_cache("hot_key", format!("value_{index}"), secs(60))
.unwrap();
}
assert_eq!(cache.expiration_heap_len(), Ok(1));
}
#[rstest]
fn repeated_increments_keep_one_heap_entry_per_expiration() {
let cache = InMemoryCache::<f64>::new(Some(4), None);
for _ in 0..100 {
cache
.increment_cache("counter", 1.0, ExactCacheContext::default())
.unwrap();
}
assert_eq!(cache.expiration_heap_len(), Ok(1));
}
#[rstest]
fn reinserting_expired_keys_below_capacity_prunes_the_heap(clock: Clock) {
let cache = cache(&clock, 200);
for cycle in 0..3 {
for index in 0..5 {
cache
.set_cache(format!("key_{index}"), format!("value_{cycle}"), secs(1))
.unwrap();
}
at(&clock, 100 + 2 * (cycle + 1));
}
for index in 0..5 {
cache
.set_cache(format!("key_{index}"), "final".into(), secs(1))
.unwrap();
}
assert_eq!(cache.len(), Ok(5));
assert_eq!(cache.expiration_heap_len(), Ok(5));
}
#[rstest]
fn evict_cache_drops_expired_entries_then_makes_room(clock: Clock) {
let cache = cache(&clock, 2);
assert_eq!(cache.is_empty(), Ok(true));
cache.set_cache("short", "a".into(), secs(1)).unwrap();
cache.set_cache("long", "b".into(), secs(50)).unwrap();
at(&clock, 102);
cache.evict_cache().unwrap();
assert_eq!(cache.len(), Ok(1));
assert_eq!(cache.expires_at("short").unwrap(), None);
cache.set_cache("longer", "c".into(), secs(90)).unwrap();
cache.evict_cache().unwrap();
assert_eq!(cache.len(), Ok(1));
assert_eq!(cache.get_cache("long").unwrap(), None);
assert_eq!(cache.get_cache("longer").unwrap(), Some("c".into()));
}
#[rstest]
fn evict_element_if_expired_reports_removal(clock: Clock) {
let cache = cache(&clock, 4);
cache.set_cache("key", "value".into(), secs(10)).unwrap();
assert_eq!(cache.evict_element_if_expired("key"), Ok(false));
assert_eq!(cache.evict_element_if_expired("missing"), Ok(false));
at(&clock, 110);
assert_eq!(cache.evict_element_if_expired("key"), Ok(false));
at(&clock, 111);
assert_eq!(cache.evict_element_if_expired("key"), Ok(true));
assert_eq!(cache.len(), Ok(0));
assert_eq!(cache.expires_at("key").unwrap(), None);
}
#[rstest]
#[case::fits("ok", Ok(true))]
#[case::at_limit("four", Ok(true))]
#[case::too_large("oversized", Ok(false))]
#[case::measure_error("", Err(Error::InvalidEntry))]
fn check_value_size_applies_the_entry_limit(
#[case] value: &str,
#[case] expected: Result<bool, Error>,
) {
assert_eq!(measured(2).check_value_size(&value.to_string()), expected);
}
#[rstest]
fn values_are_unbounded_without_a_measure() {
let cache = InMemoryCache::<String>::default();
assert_eq!(cache.max_entry_bytes(), None);
assert_eq!(cache.check_value_size(&"x".repeat(1 << 20)), Ok(true));
}
#[rstest]
fn disabled_size_limited_and_validated_writes_are_observable() {
assert_eq!(
measured(0).set_cache("a", "x".into(), None).unwrap(),
CacheWrite::Disabled
);
let cache = cache(2);
let cache = measured(2);
assert_eq!(cache.max_entry_bytes(), Some(4));
assert_eq!(
cache.set_cache("large", "oversized".into(), None).unwrap(),
CacheWrite::TooLarge
@ -132,30 +317,21 @@ fn disabled_size_limited_and_validated_writes_are_observable() {
assert_eq!(cache.get_cache("small").unwrap(), None);
}
#[rstest]
#[tokio::test]
async fn connection_test_matches_python_result_contract() {
async fn disconnect_is_a_no_op_that_keeps_entries() {
let cache = InMemoryCache::<String>::default();
let result = BaseCache::test_connection(&cache).await.unwrap();
assert_eq!(result.status, CacheConnectionStatus::Success);
assert_eq!(result.message, "In-memory cache connection test successful");
assert_eq!(result.error, None);
assert_eq!(
serde_json::to_value(result).unwrap(),
serde_json::json!({
"status": "success",
"message": "In-memory cache connection test successful"
})
);
cache.set_cache("key", "value".into(), None).unwrap();
cache.disconnect().await.unwrap();
assert_eq!(cache.get_cache("key").unwrap(), Some("value".into()));
}
#[rstest]
#[tokio::test]
async fn generic_consumers_share_typed_values_and_honor_expiration() {
let clock = clock();
let cache: CacheBackend<InMemoryCache<String>> = Arc::new(cache(clock.clone(), 4));
async fn generic_consumers_share_typed_values_and_honor_expiration(clock: Clock) {
let cache: CacheBackend<InMemoryCache<String>> = Arc::new(self::cache(&clock, 4));
let reader = Arc::clone(&cache);
let context = ExactCacheContext {
ttl: Some(Duration::from_secs(5)),
};
let context = ttl(5);
set_cache(cache.as_ref(), "sync", "first".into(), &context).unwrap();
assert_eq!(
get_cache(reader.as_ref(), "sync", &context).unwrap(),
@ -181,7 +357,7 @@ async fn generic_consumers_share_typed_values_and_honor_expiration() {
reader.async_get_cache("async", &context).await.unwrap(),
None
);
clock.store(106, Ordering::SeqCst);
at(&clock, 106);
assert_eq!(get_cache(reader.as_ref(), "sync", &context).unwrap(), None);
assert_eq!(
reader.async_get_cache("batch", &context).await.unwrap(),
@ -189,34 +365,95 @@ async fn generic_consumers_share_typed_values_and_honor_expiration() {
);
}
#[test]
fn claims_are_atomic_and_refresh_eligible_winners() {
let clock = clock();
let cache = InMemoryCache::with_clock(Some(4), Some(Duration::from_secs(60)), {
let clock = clock.clone();
move || Duration::from_secs(clock.load(Ordering::SeqCst))
});
let context = ExactCacheContext {
ttl: Some(Duration::from_secs(10)),
};
#[rstest]
#[case::context_ttl(ttl(5), secs(105))]
#[case::default_ttl(ExactCacheContext::default(), secs(160))]
#[tokio::test]
async fn pipeline_writes_use_the_context_ttl_or_the_default(
clock: Clock,
#[case] context: ExactCacheContext,
#[case] expected: Option<Duration>,
) {
let cache = cache(&clock, 4);
cache
.async_set_cache_pipeline(
vec![("a".into(), "1".into()), ("b".into(), "2".into())],
context,
)
.await
.unwrap();
assert_eq!(cache.expires_at("a").unwrap(), expected);
assert_eq!(cache.expires_at("b").unwrap(), expected);
}
#[rstest]
#[tokio::test]
async fn batch_reads_return_one_entry_per_key_and_drop_expired_ones(clock: Clock) {
let cache = cache(&clock, 4);
cache.set_cache("short", "a".into(), secs(1)).unwrap();
cache.set_cache("long", "b".into(), secs(50)).unwrap();
let keys = vec!["short".to_string(), "missing".into(), "long".into()];
assert_eq!(
cache
.batch_get_cache(&keys, &ExactCacheContext::default())
.unwrap(),
[
BatchEntry::Hit("a".to_string()),
BatchEntry::Miss,
BatchEntry::Hit("b".into()),
]
);
at(&clock, 102);
assert_eq!(
cache
.async_batch_get_cache(keys, ExactCacheContext::default())
.await
.unwrap(),
[
BatchEntry::Miss,
BatchEntry::Miss,
BatchEntry::Hit("b".into())
]
);
}
#[rstest]
#[tokio::test]
async fn flush_clears_values_and_expirations(clock: Clock) {
let cache = cache(&clock, 4);
cache.set_cache("a", "1".into(), None).unwrap();
cache.set_cache("b", "2".into(), None).unwrap();
cache.flush_cache().unwrap();
assert_eq!(cache.len(), Ok(0));
assert_eq!(cache.expiration_heap_len(), Ok(0));
cache.set_cache("c", "3".into(), None).unwrap();
FlushCache::async_flush_cache(&cache).await.unwrap();
assert_eq!(cache.is_empty(), Ok(true));
assert_eq!(
cache.async_get_oldest_n_keys(5).await.unwrap(),
Vec::<String>::new()
);
}
#[rstest]
fn claims_are_atomic_and_refresh_eligible_winners(clock: Clock) {
let cache = cache(&clock, 4);
let context = ttl(10);
assert_eq!(
cache
.claim_cache("affinity", "first".to_string(), &[], context.clone())
.unwrap(),
"first"
);
clock.store(103, Ordering::SeqCst);
at(&clock, 103);
assert_eq!(
cache
.claim_cache("affinity", "second".to_string(), &[], context.clone())
.unwrap(),
"first"
);
assert_eq!(
cache.expires_at("affinity").unwrap(),
Some(Duration::from_secs(110))
);
clock.store(105, Ordering::SeqCst);
assert_eq!(cache.expires_at("affinity").unwrap(), secs(110));
at(&clock, 105);
assert_eq!(
cache
.claim_cache(
@ -228,13 +465,10 @@ fn claims_are_atomic_and_refresh_eligible_winners() {
.unwrap(),
"first"
);
assert_eq!(
cache.expires_at("affinity").unwrap(),
Some(Duration::from_secs(115))
);
assert_eq!(cache.expires_at("affinity").unwrap(), secs(115));
}
#[test]
#[rstest]
fn counters_increment_under_one_lock() {
let cache = InMemoryCache::<f64>::default();
assert_eq!(
@ -250,44 +484,107 @@ fn counters_increment_under_one_lock() {
}
#[rstest]
fn rewriting_an_existing_key_at_capacity_keeps_other_entries(clock: Arc<AtomicU64>) {
let cache = cache(clock, 2);
cache
.set_cache("hot", "1".into(), Some(Duration::from_secs(10)))
.unwrap();
cache
.set_cache("cold", "2".into(), Some(Duration::from_secs(20)))
.unwrap();
fn concurrent_increments_are_atomic() {
let cache = Arc::new(InMemoryCache::<f64>::default());
cache.set_cache("counter", 1000.0, None).unwrap();
let threads = (0..8)
.map(|_| {
let cache = cache.clone();
std::thread::spawn(move || {
cache
.increment_cache("counter", 1.0, ExactCacheContext::default())
.unwrap()
})
})
.collect::<Vec<_>>();
for thread in threads {
thread.join().unwrap();
}
assert_eq!(cache.get_cache("counter").unwrap(), Some(1008.0));
}
#[rstest]
#[case::window_semantics(false)]
#[case::refresh_ttl_is_ignored(true)]
#[tokio::test]
async fn async_increment_delegates_to_the_locked_sync_path(
clock: Clock,
#[case] refresh_ttl: bool,
) {
let cache = cache_with::<f64>(&clock, 4);
assert_eq!(
cache
.async_increment("counter", 2.0, ttl(10), refresh_ttl)
.await,
Ok(2.0)
);
at(&clock, 105);
assert_eq!(
cache
.async_increment("counter", 3.0, ttl(10), refresh_ttl)
.await,
Ok(5.0)
);
assert_eq!(cache.get_cache("counter").unwrap(), Some(5.0));
assert_eq!(cache.expires_at("counter").unwrap(), secs(110));
}
#[rstest]
fn expired_counters_restart_from_zero_with_a_new_ttl(clock: Clock) {
let cache = cache_with::<f64>(&clock, 4);
cache.increment_cache("counter", 2.0, ttl(10)).unwrap();
at(&clock, 111);
assert_eq!(cache.increment_cache("counter", 1.0, ttl(10)), Ok(1.0));
assert_eq!(cache.expires_at("counter").unwrap(), secs(121));
}
/// Python `InMemoryCache.set_cache` runs `evict_cache()` before every insert, and step 2 evicts
/// the earliest expiry while `len(cache_dict) >= max_size_in_memory`, even when the key being
/// written already exists.
#[rstest]
fn overwriting_an_existing_key_at_capacity_evicts_the_earliest_expiry_like_python(clock: Clock) {
let cache = cache(&clock, 2);
cache.set_cache("hot", "1".into(), secs(10)).unwrap();
cache.set_cache("cold", "2".into(), secs(20)).unwrap();
cache.set_cache("cold", "3".into(), None).unwrap();
assert_eq!(cache.get_cache("hot").unwrap(), Some("1".into()));
assert_eq!(cache.get_cache("hot").unwrap(), None);
assert_eq!(cache.get_cache("cold").unwrap(), Some("3".into()));
}
/// `claim_cache` has no Python counterpart; it never evicts another entry for a key it holds.
#[rstest]
fn claiming_an_existing_key_at_capacity_keeps_other_entries(clock: Clock) {
let cache = cache(&clock, 2);
cache.set_cache("hot", "1".into(), secs(10)).unwrap();
cache.set_cache("cold", "2".into(), secs(20)).unwrap();
cache
.claim_cache("cold", "4".into(), &[], ExactCacheContext::default())
.unwrap();
assert_eq!(cache.get_cache("hot").unwrap(), Some("1".into()));
cache.set_cache("new", "5".into(), None).unwrap();
assert_eq!(cache.get_cache("hot").unwrap(), None);
assert_eq!(cache.get_cache("cold").unwrap(), Some("3".into()));
assert_eq!(cache.get_cache("new").unwrap(), Some("5".into()));
assert_eq!(cache.get_cache("cold").unwrap(), Some("2".into()));
}
#[test]
fn incrementing_an_existing_counter_at_capacity_keeps_every_counter() {
let cache = InMemoryCache::<f64>::new(Some(2), None);
/// Python `increment_cache` is `get_cache` then `set_cache`, so at capacity the write evicts
/// the earliest expiry first: equal expiries tie-break on the key, and the value read before
/// eviction is the one written back.
#[rstest]
fn incrementing_at_capacity_evicts_the_earliest_expiry_like_python(clock: Clock) {
let cache = cache_with::<f64>(&clock, 2);
for key in ["a", "b", "a", "b"] {
cache
.increment_cache(key, 1.0, ExactCacheContext::default())
.unwrap();
}
assert_eq!(cache.get_cache("a").unwrap(), Some(2.0));
assert_eq!(cache.get_cache("a").unwrap(), None);
assert_eq!(cache.get_cache("b").unwrap(), Some(2.0));
}
#[test]
fn disabled_cache_does_not_retain_claims_or_counters() {
#[rstest]
#[tokio::test]
async fn disabled_cache_does_not_retain_claims_counters_or_sets() {
let claims = InMemoryCache::<String>::new(Some(0), None);
assert_eq!(
claims
@ -305,62 +602,128 @@ fn disabled_cache_does_not_retain_claims_or_counters() {
2.0
);
assert_eq!(counters.get_cache("key").unwrap(), None);
let sets = InMemoryCache::<HashSet<String>>::new(Some(0), None);
assert_eq!(
sets.async_set_cache_sadd("key", vec!["a".into()], None)
.await
.unwrap(),
["a"]
);
assert_eq!(sets.get_cache("key").unwrap(), None);
}
#[rstest]
#[tokio::test]
async fn ttl_and_oldest_key_operations_use_the_stored_expirations() {
let clock = Arc::new(AtomicU64::new(100));
let cache = cache(clock, 3);
cache
.set_cache("later", "2".into(), Some(Duration::from_secs(20)))
.unwrap();
cache
.set_cache("first", "1".into(), Some(Duration::from_secs(10)))
.unwrap();
async fn ttl_and_oldest_key_operations_use_the_stored_expirations(clock: Clock) {
let cache = cache(&clock, 3);
cache.set_cache("later", "2".into(), secs(20)).unwrap();
cache.set_cache("first", "1".into(), secs(10)).unwrap();
cache.set_cache("latest", "3".into(), secs(30)).unwrap();
assert_eq!(cache.async_get_ttl("first").await.unwrap(), secs(110));
assert_eq!(
cache.async_get_ttl("first").await.unwrap(),
Some(Duration::from_secs(110))
TtlCache::async_get_ttl(&cache, "later").await.unwrap(),
secs(120)
);
assert_eq!(cache.async_get_oldest_n_keys(1).await.unwrap(), ["first"]);
assert_eq!(
cache.async_get_oldest_n_keys(10).await.unwrap(),
["first", "later", "latest"]
);
assert_eq!(
cache.async_get_oldest_n_keys(0).await.unwrap(),
Vec::<String>::new()
);
assert_eq!(cache.async_get_ttl("missing").await.unwrap(), None);
}
#[rstest]
#[tokio::test]
async fn increment_pipeline_preserves_operation_order() {
let cache = InMemoryCache::<f64>::new(Some(3), None);
async fn increment_pipeline_preserves_operation_order(clock: Clock) {
let cache = cache_with::<f64>(&clock, 3);
let operation = |key: &str, amount, ttl| IncrementOperation {
key: key.into(),
amount,
ttl: secs(ttl),
};
assert_eq!(
cache
.async_increment_pipeline(vec![
IncrementOperation {
key: "a".into(),
amount: 1.0,
ttl: Some(Duration::from_secs(10)),
},
IncrementOperation {
key: "a".into(),
amount: 2.0,
ttl: Some(Duration::from_secs(20)),
},
operation("a", 1.0, 10),
operation("b", 5.0, 30),
operation("a", 2.0, 20),
])
.await
.unwrap(),
[1.0, 3.0]
[1.0, 5.0, 3.0]
);
assert_eq!(cache.get_cache("a").unwrap(), Some(3.0));
assert_eq!(cache.expires_at("a").unwrap(), secs(110));
assert_eq!(cache.expires_at("b").unwrap(), secs(130));
assert_eq!(
cache.async_increment_pipeline(Vec::new()).await.unwrap(),
Vec::<f64>::new()
);
}
#[rstest]
#[tokio::test]
async fn set_capability_preserves_python_result_and_deduplicates_storage() {
let cache = InMemoryCache::<HashSet<String>>::new(None, None);
let inserted = vec!["a".into(), "a".into(), "b".into()];
async fn set_capability_preserves_python_result_and_deduplicates_storage(clock: Clock) {
let cache = cache_with::<HashSet<String>>(&clock, 4);
let inserted = vec!["a".to_string(), "a".into(), "b".into()];
assert_eq!(
cache
.async_set_cache_sadd("members", inserted.clone(), None)
.async_set_cache_sadd("members", inserted.clone(), secs(10))
.await
.unwrap(),
inserted
);
assert_eq!(
cache
.async_set_cache_sadd("members", vec!["c".into()], secs(99))
.await
.unwrap(),
["c"]
);
assert_eq!(
cache.get_cache("members").unwrap(),
Some(HashSet::from(["a".into(), "b".into(), "c".into()]))
);
assert_eq!(cache.expires_at("members").unwrap(), secs(110));
at(&clock, 111);
cache
.async_set_cache_sadd("members", vec!["d".into()], None)
.await
.unwrap();
assert_eq!(
cache.get_cache("members").unwrap(),
Some(HashSet::from(["d".into()]))
);
assert_eq!(cache.expires_at("members").unwrap(), secs(171));
}
#[rstest]
#[tokio::test]
async fn oversized_set_additions_are_not_stored() {
let cache = InMemoryCache::<HashSet<String>>::with_clock_and_size_measurement(
Some(4),
None,
Some(2),
Some(Arc::new(|value: &HashSet<String>| Ok(value.len()))),
|| Duration::ZERO,
);
cache
.async_set_cache_sadd("members", vec!["a".into(), "b".into()], None)
.await
.unwrap();
assert_eq!(
cache
.async_set_cache_sadd("members", vec!["c".into()], None)
.await
.unwrap(),
["c"]
);
assert_eq!(
cache.get_cache("members").unwrap(),
Some(HashSet::from(["a".into(), "b".into()]))

98
litellm-rust/crates/cache-memory/tests/contract.rs Normal file
View file

@ -0,0 +1,98 @@
use std::time::Duration;
use litellm_cache::ExactCacheContext;
use litellm_cache_memory::InMemoryCache;
use litellm_cache_testing as contract;
use rstest::{fixture, rstest};
#[fixture]
fn strings() -> InMemoryCache<String> {
InMemoryCache::new(Some(16), None)
}
#[fixture]
fn counters() -> InMemoryCache<f64> {
InMemoryCache::new(Some(16), None)
}
#[fixture]
fn context() -> ExactCacheContext {
ExactCacheContext {
ttl: Some(Duration::from_secs(60)),
}
}
#[rstest]
#[tokio::test]
async fn hit_and_miss(strings: InMemoryCache<String>, context: ExactCacheContext) {
contract::hit_and_miss(&strings, context, "memory:", "value".into()).await;
}
#[rstest]
#[tokio::test]
async fn sync_async_equivalence(strings: InMemoryCache<String>, context: ExactCacheContext) {
contract::sync_async_equivalence(
&strings,
context,
"memory:",
"first".into(),
"second".into(),
)
.await;
}
#[rstest]
#[tokio::test]
async fn overwrite_replaces(strings: InMemoryCache<String>, context: ExactCacheContext) {
contract::overwrite_replaces(
&strings,
context,
"memory:",
"first".into(),
"second".into(),
)
.await;
}
#[rstest]
#[tokio::test]
async fn pipeline_writes_every_entry(strings: InMemoryCache<String>, context: ExactCacheContext) {
contract::pipeline_writes_every_entry(
&strings,
context,
"memory:",
vec!["a".into(), "b".into(), "c".into()],
)
.await;
}
#[rstest]
#[tokio::test]
async fn batch_preserves_order(strings: InMemoryCache<String>, context: ExactCacheContext) {
contract::batch_preserves_order(
&strings,
context,
"memory:",
"first".into(),
"second".into(),
)
.await;
}
#[rstest]
#[tokio::test]
async fn delete_removes_key(strings: InMemoryCache<String>, context: ExactCacheContext) {
contract::delete_removes_key(&strings, context, "memory:", "value".into()).await;
}
#[rstest]
#[tokio::test]
async fn flush_clears(strings: InMemoryCache<String>, context: ExactCacheContext) {
contract::flush_clears(&strings, context, "memory:", "value".into()).await;
}
#[rstest]
#[tokio::test]
async fn counter_accumulates(counters: InMemoryCache<f64>, context: ExactCacheContext) {
contract::counter_accumulates(&counters, context, "memory:").await;
}

3
litellm-rust/crates/cache-qdrant-semantic/Cargo.toml
View file

@ -17,7 +17,8 @@ tokio.workspace = true
uuid.workspace = true
[dev-dependencies]
litellm-cache-response.workspace = true
futures-executor = "0.3"
litellm-cache-testing.workspace = true
rstest.workspace = true
tonic = "0.14"
tonic-prost = "0.14"

105
litellm-rust/crates/cache-qdrant-semantic/src/semantic.rs → litellm-rust/crates/cache-qdrant-semantic/src/cache.rs
View file

@ -1,7 +1,8 @@
use std::future::Future;
use futures_util::future::try_join_all;
use litellm_cache::{BaseCache, CacheCodec, CacheConnectionResult, Error, SemanticCacheContext};
use litellm_cache::{
BaseCache, CacheCodec, Error, SemanticCacheContext,
semantic::{Embedder, SemanticCache, SemanticLookup, prompt_from_messages},
};
use qdrant_client::{
Payload, Qdrant,
qdrant::{
@ -14,26 +15,7 @@ use qdrant_client::{
use serde_json::{Map, Value, json};
use uuid::Uuid;
use crate::prompt_from_messages;
pub trait Embedder: Send + Sync + 'static {
fn model(&self) -> &str;
fn embed(&self, input: &str) -> impl Future<Output = Result<Vec<f32>, Error>> + Send;
}
#[derive(Clone, Debug, PartialEq)]
pub enum Quantization {
Binary,
Scalar,
Product,
}
pub struct QdrantSemanticConfig {
pub collection_name: String,
pub similarity_threshold: f64,
pub vector_size: u64,
pub quantization: Quantization,
}
use crate::{QdrantSemanticConfig, Quantization};
pub struct QdrantSemanticCache<E: Embedder, C: CacheCodec> {
client: Qdrant,
@ -100,14 +82,9 @@ impl<E: Embedder, C: CacheCodec> QdrantSemanticCache<E, C> {
&self.embedder
}
/// Python reads `kwargs["messages"]` unguarded, so a request without messages fails.
fn prompt(context: &SemanticCacheContext) -> Result<String, Error> {
let Some(messages) = context.messages.as_ref().and_then(Value::as_array) else {
return Err(Error::MissingPrompt);
};
if messages.is_empty() {
return Err(Error::MissingPrompt);
}
Ok(prompt_from_messages(messages))
prompt_from_messages(context).ok_or(Error::MissingPrompt)
}
async fn set(
@ -117,7 +94,10 @@ impl<E: Embedder, C: CacheCodec> QdrantSemanticCache<E, C> {
context: &SemanticCacheContext,
) -> Result<(), Error> {
let prompt = Self::prompt(context)?;
let vector = self.embedder.embed(&prompt).await?;
let vector = self
.embedder
.async_embed(&prompt, context.metadata.as_ref())
.await?;
let response =
String::from_utf8(self.codec.encode(&value)?).map_err(|_| Error::InvalidEntry)?;
let payload = Payload::try_from(json!({
@ -147,9 +127,12 @@ impl<E: Embedder, C: CacheCodec> QdrantSemanticCache<E, C> {
&self,
key: &str,
context: &SemanticCacheContext,
) -> Result<Option<C::Value>, Error> {
) -> Result<SemanticLookup<C::Value>, Error> {
let prompt = Self::prompt(context)?;
let vector = self.embedder.embed(&prompt).await?;
let vector = self
.embedder
.async_embed(&prompt, context.metadata.as_ref())
.await?;
let result = self
.client
.search_points(
@ -171,20 +154,27 @@ impl<E: Embedder, C: CacheCodec> QdrantSemanticCache<E, C> {
.await
.map_err(|_| Error::Unavailable)?;
let Some(point) = result.result.into_iter().next() else {
return Ok(None);
return Ok(SemanticLookup::miss(Some(0.0)));
};
let payload: Map<String, Value> = Payload::from(point.payload).into();
if payload.get("litellm_cache_key").and_then(Value::as_str) != Some(key) {
return Ok(None);
if !payload
.get("litellm_cache_key")
.is_some_and(|cached| python_str(cached).as_deref() == Some(key))
{
return Ok(SemanticLookup::miss(Some(0.0)));
}
if f64::from(point.score) < self.config.similarity_threshold {
return Ok(None);
let similarity = f64::from(point.score);
if similarity < self.config.similarity_threshold {
return Ok(SemanticLookup::miss(Some(similarity)));
}
let response = payload
.get("response")
.and_then(Value::as_str)
.ok_or(Error::InvalidEntry)?;
self.codec.decode(response.as_bytes()).map(Some)
Ok(SemanticLookup {
value: Some(self.codec.decode(response.as_bytes())?),
similarity: Some(similarity),
})
}
}
@ -219,7 +209,8 @@ impl<E: Embedder, C: CacheCodec> BaseCache for QdrantSemanticCache<E, C> {
}
fn get_cache(&self, key: &str, context: &Self::Context) -> Result<Option<Self::Value>, Error> {
self.runtime.block_on(self.get(key, context))
self.get_cache_with_similarity(key, context)
.map(|lookup| lookup.value)
}
async fn async_set_cache(
@ -236,7 +227,7 @@ impl<E: Embedder, C: CacheCodec> BaseCache for QdrantSemanticCache<E, C> {
key: &str,
context: &Self::Context,
) -> Result<Option<Self::Value>, Error> {
self.get(key, context).await
self.get(key, context).await.map(|lookup| lookup.value)
}
async fn async_set_cache_pipeline(
@ -251,12 +242,36 @@ impl<E: Embedder, C: CacheCodec> BaseCache for QdrantSemanticCache<E, C> {
.await
.map(|_| ())
}
}
async fn disconnect(&self) -> Result<(), Error> {
Ok(())
/// Python stamps the top point's score, even below the threshold, and `0.0` when there is no
/// point or it belongs to another key. A request without messages fails before any search.
impl<E: Embedder, C: CacheCodec> SemanticCache for QdrantSemanticCache<E, C> {
fn get_cache_with_similarity(
&self,
key: &str,
context: &Self::Context,
) -> Result<SemanticLookup<Self::Value>, Error> {
self.runtime.block_on(self.get(key, context))
}
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
Err(Error::UnsupportedOperation)
async fn async_get_cache_with_similarity(
&self,
key: &str,
context: &Self::Context,
) -> Result<SemanticLookup<Self::Value>, Error> {
self.get(key, context).await
}
}
/// `str(value)` for the scalar payload values `_payload_matches_cache_key` compares; `None` for
/// null (a pre-isolation point without a key) and for containers, which never equal a key.
fn python_str(value: &Value) -> Option<String> {
match value {
Value::String(text) => Some(text.clone()),
Value::Number(number) => Some(number.to_string()),
Value::Bool(true) => Some("True".into()),
Value::Bool(false) => Some("False".into()),
Value::Null | Value::Array(_) | Value::Object(_) => None,
}
}

13
litellm-rust/crates/cache-qdrant-semantic/src/config.rs Normal file
View file

@ -0,0 +1,13 @@
#[derive(Clone, Debug, PartialEq)]
pub enum Quantization {
Binary,
Scalar,
Product,
}
pub struct QdrantSemanticConfig {
pub collection_name: String,
pub similarity_threshold: f64,
pub vector_size: u64,
pub quantization: Quantization,
}

14
litellm-rust/crates/cache-qdrant-semantic/src/embedder.rs
View file

@ -1,11 +1,9 @@
use std::time::Duration;
use litellm_cache::Error;
use litellm_cache::{Error, semantic::Embedder};
use reqwest::Client;
use serde_json::Value;
use crate::Embedder;
pub struct OpenAiEmbedder {
client: Client,
api_base: String,
@ -31,14 +29,16 @@ impl OpenAiEmbedder {
timeout: config.timeout,
}
}
}
impl Embedder for OpenAiEmbedder {
fn model(&self) -> &str {
pub fn model(&self) -> &str {
&self.model
}
}
async fn embed(&self, input: &str) -> Result<Vec<f32>, Error> {
/// An OpenAI-compatible `/embeddings` call. It has no router to route on, so `metadata` is
/// unused, and it only embeds asynchronously: sync cache calls block on the cache's runtime.
impl Embedder for OpenAiEmbedder {
async fn async_embed(&self, input: &str, _metadata: Option<&Value>) -> Result<Vec<f32>, Error> {
let request = self
.client
.post(format!("{}/embeddings", self.api_base))

8
litellm-rust/crates/cache-qdrant-semantic/src/lib.rs
View file

@ -1,7 +1,7 @@
mod cache;
mod config;
mod embedder;
mod prompt;
mod semantic;
pub use cache::QdrantSemanticCache;
pub use config::{QdrantSemanticConfig, Quantization};
pub use embedder::{OpenAiEmbedder, OpenAiEmbedderConfig};
pub use prompt::prompt_from_messages;
pub use semantic::{Embedder, QdrantSemanticCache, QdrantSemanticConfig, Quantization};

59
litellm-rust/crates/cache-qdrant-semantic/src/prompt.rs
View file

@ -1,59 +0,0 @@
use serde_json::Value;
fn search_results_text(search_results: Option<&Value>) -> String {
let Some(Value::Array(results)) = search_results else {
return String::new();
};
results
.iter()
.filter_map(Value::as_object)
.flat_map(|result| {
let source = result
.get("source")
.and_then(Value::as_str)
.map(str::to_owned);
let title = result
.get("title")
.and_then(Value::as_str)
.map(str::to_owned);
let content = result
.get("content")
.and_then(Value::as_array)
.into_iter()
.flatten()
.filter_map(Value::as_object)
.filter_map(|block| block.get("text").and_then(Value::as_str).map(str::to_owned));
let citations = result
.get("citations")
.filter(|value| !value.is_null())
.map(|value| serde_json::to_string(value).unwrap_or_default());
source
.into_iter()
.chain(title)
.chain(content)
.chain(citations)
})
.collect()
}
pub fn prompt_from_messages(messages: &[Value]) -> String {
messages
.iter()
.filter_map(Value::as_object)
.map(|message| {
let content = match message.get("content") {
Some(Value::String(content)) => content.clone(),
Some(Value::Array(parts)) => parts
.iter()
.filter_map(Value::as_object)
.filter_map(|part| part.get("text").and_then(Value::as_str))
.collect(),
_ => String::new(),
};
format!(
"{content}{}",
search_results_text(message.get("search_results"))
)
})
.collect()
}

91
litellm-rust/crates/cache-qdrant-semantic/tests/contract.rs Normal file
View file

@ -0,0 +1,91 @@
//! `overwrite_replaces` does not apply: like Python, every write upserts a new `uuid4` point,
//! so a second write with the same prompt adds a tie instead of replacing the first.
mod support;
use std::future::Future;
use litellm_cache::{JsonCodec, SemanticCacheContext, semantic::PreparedEmbedding};
use litellm_cache_qdrant_semantic::{QdrantSemanticCache, QdrantSemanticConfig, Quantization};
use litellm_cache_testing as contract;
use qdrant_client::Qdrant;
use rstest::{fixture, rstest};
use serde_json::{Value, json};
use support::{FakeQdrant, FakeState};
type Cache = QdrantSemanticCache<PreparedEmbedding, JsonCodec<Value>>;
const PREFIX: &str = "contract:";
#[fixture]
fn context() -> SemanticCacheContext {
SemanticCacheContext {
messages: Some(json!([{"role": "user", "content": "contract prompt"}])),
..Default::default()
}
}
/// Runs a contract against a fresh fake Qdrant. The sync cache methods block on the runtime, so
/// the contract is polled on a blocking thread outside the runtime's own executor.
async fn run<F, Fut>(check: F)
where
F: FnOnce(Cache) -> Fut + Send + 'static,
Fut: Future<Output = ()>,
{
let server = FakeQdrant::start(FakeState::default()).await;
let runtime = tokio::runtime::Handle::current();
let cache = QdrantSemanticCache::connect(
Qdrant::from_url(&server.url()).build().unwrap(),
PreparedEmbedding(vec![0.6, 0.8]),
JsonCodec::new(),
QdrantSemanticConfig {
collection_name: "contract".to_owned(),
similarity_threshold: 0.9,
vector_size: 2,
quantization: Quantization::Binary,
},
runtime.clone(),
)
.await
.unwrap();
tokio::task::spawn_blocking(move || {
let _guard = runtime.enter();
futures_executor::block_on(check(cache));
})
.await
.unwrap();
server.stop();
}
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn hit_and_miss(context: SemanticCacheContext) {
run(|cache| async move {
contract::hit_and_miss(&cache, context, PREFIX, json!({"answer": 42})).await;
})
.await;
}
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn sync_async_equivalence(context: SemanticCacheContext) {
run(|cache| async move {
contract::sync_async_equivalence(&cache, context, PREFIX, json!("first"), json!([2])).await;
})
.await;
}
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn pipeline_writes_every_entry(context: SemanticCacheContext) {
run(|cache| async move {
contract::pipeline_writes_every_entry(
&cache,
context,
PREFIX,
vec![json!("a"), json!(2), json!({"c": true})],
)
.await;
})
.await;
}

91
litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs
View file

@ -3,9 +3,10 @@ use std::{
time::Duration,
};
use litellm_cache::Error;
use litellm_cache_qdrant_semantic::{Embedder, OpenAiEmbedder, OpenAiEmbedderConfig};
use serde_json::Value;
use litellm_cache::{Error, semantic::Embedder};
use litellm_cache_qdrant_semantic::{OpenAiEmbedder, OpenAiEmbedderConfig};
use rstest::rstest;
use serde_json::{Value, json};
use tokio::{
io::{AsyncReadExt, AsyncWriteExt},
net::TcpListener,
@ -98,6 +99,7 @@ fn config(base: String, timeout: Option<Duration>) -> OpenAiEmbedderConfig {
}
}
#[rstest]
#[tokio::test]
async fn posts_embeddings_request_and_parses_vector() {
let server = TestHttpServer::response("200 OK", r#"{"data":[{"embedding":[0.1,0.2]}]}"#).await;
@ -108,7 +110,14 @@ async fn posts_embeddings_request_and_parses_vector() {
Some(Duration::from_secs(1)),
),
);
assert_eq!(embedder.embed("hello").await.unwrap(), vec![0.1, 0.2]);
assert_eq!(embedder.model(), "test-model");
assert_eq!(
embedder
.async_embed("hello", Some(&json!({"ignored": true})))
.await
.unwrap(),
vec![0.1, 0.2]
);
let request = server.request.lock().unwrap().clone().unwrap();
let request_text = String::from_utf8(request).unwrap();
assert!(request_text.starts_with("POST /embeddings HTTP/1.1\r\n"));
@ -120,37 +129,50 @@ async fn posts_embeddings_request_and_parses_vector() {
assert_eq!(body["encoding_format"], "float");
}
#[rstest]
#[case::error_status("500 Internal Server Error", "{}", 0, None, Err(Error::Unavailable))]
#[case::timed_out(
"200 OK",
r#"{"data":[{"embedding":[0.1,0.2]}]}"#,
500,
Some(Duration::from_millis(200)),
Err(Error::Unavailable)
)]
#[case::within_timeout(
"200 OK",
r#"{"data":[{"embedding":[0.1,0.2]}]}"#,
100,
Some(Duration::from_secs(1)),
Ok(vec![0.1, 0.2])
)]
#[case::missing_embedding("200 OK", r#"{"data":[]}"#, 0, None, Err(Error::Unavailable))]
#[tokio::test]
async fn status_and_timeout_errors_are_unavailable() {
let server = TestHttpServer::response("500 Internal Server Error", "{}").await;
let embedder = OpenAiEmbedder::new(reqwest::Client::new(), config(server.base_url(), None));
assert_eq!(embedder.embed("hello").await, Err(Error::Unavailable));
let server = TestHttpServer::response_after(
"200 OK",
r#"{"data":[{"embedding":[0.1,0.2]}]}"#,
Duration::from_millis(500),
)
.await;
let embedder = OpenAiEmbedder::new(
reqwest::Client::new(),
config(server.base_url(), Some(Duration::from_millis(200))),
);
assert_eq!(embedder.embed("hello").await, Err(Error::Unavailable));
let server = TestHttpServer::response_after(
"200 OK",
r#"{"data":[{"embedding":[0.1,0.2]}]}"#,
Duration::from_millis(100),
)
.await;
let embedder = OpenAiEmbedder::new(
reqwest::Client::new(),
config(server.base_url(), Some(Duration::from_secs(1))),
);
assert_eq!(embedder.embed("hello").await.unwrap(), vec![0.1, 0.2]);
async fn status_timeout_and_body_errors_are_unavailable(
#[case] status: &str,
#[case] body: &str,
#[case] delay_ms: u64,
#[case] timeout: Option<Duration>,
#[case] expected: Result<Vec<f32>, Error>,
) {
let server =
TestHttpServer::response_after(status, body, Duration::from_millis(delay_ms)).await;
let embedder = OpenAiEmbedder::new(reqwest::Client::new(), config(server.base_url(), timeout));
assert_eq!(embedder.async_embed("hello", None).await, expected);
}
#[rstest]
fn sync_embedding_is_unsupported() {
let embedder = OpenAiEmbedder::new(
reqwest::Client::new(),
config("http://127.0.0.1:9".to_owned(), None),
);
assert_eq!(
embedder.embed("hello", None),
Err(Error::UnsupportedOperation)
);
}
#[rstest]
#[tokio::test]
async fn uses_the_injected_client() {
let server = TestHttpServer::response("200 OK", r#"{"data":[{"embedding":[0.1,0.2]}]}"#).await;
@ -159,7 +181,10 @@ async fn uses_the_injected_client() {
.build()
.unwrap();
let embedder = OpenAiEmbedder::new(client, config(server.base_url(), None));
assert_eq!(embedder.embed("hello").await.unwrap(), vec![0.1, 0.2]);
assert_eq!(
embedder.async_embed("hello", None).await.unwrap(),
vec![0.1, 0.2]
);
let request = server.request.lock().unwrap().clone().unwrap();
let request_text = String::from_utf8(request).unwrap();
assert!(request_text.contains("\r\nuser-agent: litellm-embedder-test\r\n"));

38
litellm-rust/crates/cache-qdrant-semantic/tests/prompt.rs
View file

@ -1,38 +0,0 @@
use litellm_cache_qdrant_semantic::prompt_from_messages;
use serde_json::json;
#[test]
fn prompt_matches_python_message_content_rules() {
let messages = vec![
json!({"role": "user", "content": "hello"}),
json!({
"role": "user",
"content": [
{"type": "text", "text": "world"},
{"type": "image_url", "image_url": {"url": "ignored"}},
{"type": "text", "text": "!"},
],
}),
];
assert_eq!(prompt_from_messages(&messages), "helloworld!");
}
#[test]
fn prompt_includes_search_result_text_and_compact_citations() {
let messages = vec![json!({
"role": "tool",
"content": null,
"search_results": [{
"source": "source",
"title": "title",
"content": [{"text": "body"}],
"citations": {"page": 1, "section": "intro"},
}],
})];
assert_eq!(
prompt_from_messages(&messages),
r#"sourcetitlebody{"page":1,"section":"intro"}"#
);
}

595
litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs
View file

@ -1,27 +1,32 @@
#[path = "support/mod.rs"]
mod support;
use std::{collections::HashMap, sync::Arc, time::Duration};
use std::{
collections::HashMap,
sync::{Arc, Mutex},
time::Duration,
};
use litellm_cache::{BaseCache, CacheCodec, CacheContext, Error, SemanticCacheContext};
use litellm_cache_qdrant_semantic::{
Embedder, QdrantSemanticCache, QdrantSemanticConfig, Quantization,
use litellm_cache::{
BaseCache, CacheContext, Error, JsonCodec, SemanticCacheContext,
semantic::{Embedder, SemanticCache, SemanticLookup},
};
use litellm_cache_response::{
CacheEntry, CacheKeyInput, ResponseCache, ResponseCacheCodec, ResponseCacheRequest,
};
use qdrant_client::Payload;
use litellm_cache_qdrant_semantic::{QdrantSemanticCache, QdrantSemanticConfig, Quantization};
use qdrant_client::{
Qdrant,
Payload, Qdrant,
qdrant::{self, CompressionRatio, Distance, PointId, QuantizationType, Value, VectorParams},
};
use rstest::{fixture, rstest};
use serde_json::{Value as JsonValue, json};
use support::{FakeQdrant, FakeState, StoredPoint};
type Calls = Arc<Mutex<Vec<(String, Option<JsonValue>)>>>;
type Cache = QdrantSemanticCache<FixedEmbedder, JsonCodec<JsonValue>>;
/// Embeds known prompts, fails on anything else, and records every call.
#[derive(Clone)]
struct FixedEmbedder {
vectors: Arc<HashMap<String, Vec<f32>>>,
calls: Calls,
}
impl FixedEmbedder {
@ -33,16 +38,21 @@ impl FixedEmbedder {
.map(|(prompt, vector)| (prompt.to_owned(), vector))
.collect(),
),
calls: Calls::default(),
}
}
}
impl Embedder for FixedEmbedder {
fn model(&self) -> &str {
"fixed"
}
async fn embed(&self, input: &str) -> Result<Vec<f32>, Error> {
async fn async_embed(
&self,
input: &str,
metadata: Option<&JsonValue>,
) -> Result<Vec<f32>, Error> {
self.calls
.lock()
.unwrap()
.push((input.to_owned(), metadata.cloned()));
self.vectors.get(input).cloned().ok_or(Error::Unavailable)
}
}
@ -63,22 +73,20 @@ fn context(prompt: &str) -> SemanticCacheContext {
}
}
fn value(response: JsonValue) -> CacheEntry {
CacheEntry {
timestamp: Some(1.0),
response,
}
#[fixture]
fn entry() -> JsonValue {
json!({"timestamp": 1.0, "response": {"answer": 42}})
}
async fn connect(
server: &FakeQdrant,
vectors: impl IntoIterator<Item = (&'static str, Vec<f32>)>,
) -> QdrantSemanticCache<FixedEmbedder, ResponseCacheCodec> {
) -> Cache {
let client = Qdrant::from_url(&server.url()).build().unwrap();
QdrantSemanticCache::connect(
client,
FixedEmbedder::new(vectors),
ResponseCacheCodec,
JsonCodec::new(),
config(Quantization::Binary),
tokio::runtime::Handle::current(),
)
@ -86,72 +94,70 @@ async fn connect(
.unwrap()
}
#[rstest]
#[case::binary(Quantization::Binary)]
#[case::scalar(Quantization::Scalar)]
#[case::product(Quantization::Product)]
#[tokio::test(flavor = "multi_thread")]
#[expect(
deprecated,
reason = "the test verifies Qdrant's legacy always_ram quantization contract"
)]
async fn connect_sets_collection_quantization_and_index() {
for (quantization, expected) in [
(Quantization::Binary, 0),
(Quantization::Scalar, 1),
(Quantization::Product, 2),
] {
let server = FakeQdrant::start(FakeState::default()).await;
let client = Qdrant::from_url(&server.url()).build().unwrap();
QdrantSemanticCache::connect(
client,
FixedEmbedder::new([]),
ResponseCacheCodec,
config(quantization),
tokio::runtime::Handle::current(),
)
.await
async fn connect_sets_collection_quantization_and_index(#[case] quantization: Quantization) {
let server = FakeQdrant::start(FakeState::default()).await;
let client = Qdrant::from_url(&server.url()).build().unwrap();
QdrantSemanticCache::connect(
client,
FixedEmbedder::new([]),
JsonCodec::<JsonValue>::new(),
config(quantization.clone()),
tokio::runtime::Handle::current(),
)
.await
.unwrap();
let state = server.state.lock().unwrap();
let request = &state.created_collections[0];
let Some(qdrant::vectors_config::Config::Params(VectorParams { size, distance, .. })) = request
.vectors_config
.as_ref()
.and_then(|config| config.config.clone())
else {
panic!("missing vector params");
};
assert_eq!(size, 2);
assert_eq!(distance, Distance::Cosine as i32);
let quantization_config = request
.quantization_config
.as_ref()
.unwrap()
.quantization
.unwrap();
let state = server.state.lock().unwrap();
let request = &state.created_collections[0];
let Some(qdrant::vectors_config::Config::Params(VectorParams { size, distance, .. })) =
request
.vectors_config
.as_ref()
.and_then(|config| config.config.clone())
else {
panic!("missing vector params");
};
assert_eq!(size, 2);
assert_eq!(distance, Distance::Cosine as i32);
let quantization_config = request
.quantization_config
.as_ref()
.unwrap()
.quantization
.unwrap();
match (expected, quantization_config) {
(0, qdrant::quantization_config::Quantization::Binary(binary)) => {
assert_eq!(binary.always_ram, Some(false));
}
(1, qdrant::quantization_config::Quantization::Scalar(scalar)) => {
assert_eq!(scalar.r#type, QuantizationType::Int8 as i32);
assert_eq!(scalar.quantile, Some(0.99));
assert_eq!(scalar.always_ram, Some(false));
}
(2, qdrant::quantization_config::Quantization::Product(product)) => {
assert_eq!(product.compression, CompressionRatio::X16 as i32);
assert_eq!(product.always_ram, Some(false));
}
_ => panic!("unexpected quantization"),
#[expect(
deprecated,
reason = "the test verifies Qdrant's legacy always_ram quantization contract"
)]
match (quantization, quantization_config) {
(Quantization::Binary, qdrant::quantization_config::Quantization::Binary(binary)) => {
assert_eq!(binary.always_ram, Some(false));
}
assert!(state.index_creations >= 1);
assert_eq!(state.field_indexes[0].collection_name, "semantic");
assert_eq!(state.field_indexes[0].field_name, "litellm_cache_key");
assert_eq!(
state.field_indexes[0].field_type,
Some(qdrant::FieldType::Keyword as i32)
);
server.stop();
(Quantization::Scalar, qdrant::quantization_config::Quantization::Scalar(scalar)) => {
assert_eq!(scalar.r#type, QuantizationType::Int8 as i32);
assert_eq!(scalar.quantile, Some(0.99));
assert_eq!(scalar.always_ram, Some(false));
}
(Quantization::Product, qdrant::quantization_config::Quantization::Product(product)) => {
assert_eq!(product.compression, CompressionRatio::X16 as i32);
assert_eq!(product.always_ram, Some(false));
}
_ => panic!("unexpected quantization"),
}
assert!(state.index_creations >= 1);
assert_eq!(state.field_indexes[0].collection_name, "semantic");
assert_eq!(state.field_indexes[0].field_name, "litellm_cache_key");
assert_eq!(
state.field_indexes[0].field_type,
Some(qdrant::FieldType::Keyword as i32)
);
server.stop();
}
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn existing_collection_skips_create_and_index_failure_is_non_fatal() {
let server = FakeQdrant::start(FakeState {
@ -160,19 +166,25 @@ async fn existing_collection_skips_create_and_index_failure_is_non_fatal() {
..Default::default()
})
.await;
let _cache = connect(&server, [("hello", vec![1.0, 0.0])]).await;
let cache = connect(&server, [("hello", vec![1.0, 0.0])]).await;
assert_eq!(cache.collection_name(), "semantic");
assert_eq!(cache.similarity_threshold(), 0.9);
assert_eq!(cache.vector_size(), 2);
let state = server.state.lock().unwrap();
assert!(state.created_collections.is_empty());
assert!(state.index_creations >= 1);
server.stop();
}
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn async_and_sync_set_get_store_exact_payload() {
async fn async_and_sync_set_get_store_exact_payload(entry: JsonValue) {
let server = FakeQdrant::start(FakeState::default()).await;
let cache = Arc::new(connect(&server, [("hello", vec![1.0, 0.0])]).await);
let ctx = context("hello");
let entry = value(json!({"answer": 42}));
let ctx = SemanticCacheContext {
metadata: Some(json!({"tenant": "team"})),
..context("hello")
};
cache
.async_set_cache("key", entry.clone(), ctx.clone())
.await
@ -188,10 +200,8 @@ async fn async_and_sync_set_get_store_exact_payload() {
payload_keys.sort();
assert_eq!(payload_keys, ["litellm_cache_key", "response", "text"]);
assert_eq!(payload["litellm_cache_key"], Value::from("key"));
assert_eq!(
payload["response"],
Value::from(String::from_utf8(ResponseCacheCodec.encode(&entry).unwrap()).unwrap())
);
assert_eq!(payload["text"], Value::from("hello"));
assert_eq!(payload["response"], Value::from(entry.to_string()));
}
let sync_entry = entry.clone();
let sync_cache = cache.clone();
@ -207,208 +217,276 @@ async fn async_and_sync_set_get_store_exact_payload() {
})
.await
.unwrap();
assert_eq!(
*cache.embedder().calls.lock().unwrap(),
vec![("hello".to_owned(), ctx.metadata.clone()); 4]
);
server.stop();
}
#[rstest]
#[case::content_parts_skip_images(
json!([
{"role": "user", "content": "hello"},
{
"role": "user",
"content": [
{"type": "text", "text": "world"},
{"type": "image_url", "image_url": {"url": "ignored"}},
{"type": "text", "text": "!"},
],
},
]),
"helloworld!"
)]
#[case::search_results_and_compact_citations(
json!([{
"role": "tool",
"content": null,
"search_results": [{
"source": "source",
"title": "title",
"content": [{"text": "body"}],
"citations": {"page": 1, "section": "intro"},
}],
}]),
r#"sourcetitlebody{"page":1,"section":"intro"}"#
)]
#[tokio::test(flavor = "multi_thread")]
async fn misses_and_payload_validation_are_safe() {
async fn prompt_matches_python_message_rules(
#[case] messages: JsonValue,
#[case] prompt: &'static str,
entry: JsonValue,
) {
let server = FakeQdrant::start(FakeState::default()).await;
let cache = connect(&server, [(prompt, vec![1.0, 0.0])]).await;
let context = SemanticCacheContext {
messages: Some(messages),
..Default::default()
};
cache.async_set_cache("key", entry, context).await.unwrap();
assert_eq!(cache.embedder().calls.lock().unwrap()[0].0, prompt);
assert_eq!(
server.state.lock().unwrap().points[0].payload["text"],
Value::from(prompt)
);
server.stop();
}
#[rstest]
#[case::no_messages(SemanticCacheContext::default())]
#[case::empty_messages(SemanticCacheContext { messages: Some(json!([])), ..Default::default() })]
#[case::responses_input_is_not_read(SemanticCacheContext { input: Some(json!("hello")), ..Default::default() })]
#[tokio::test(flavor = "multi_thread")]
async fn requests_without_messages_are_missing_a_prompt(
#[case] context: SemanticCacheContext,
entry: JsonValue,
) {
let server = FakeQdrant::start(FakeState::default()).await;
let cache = connect(&server, [("hello", vec![1.0, 0.0])]).await;
assert_eq!(
cache.async_set_cache("key", entry, context.clone()).await,
Err(Error::MissingPrompt)
);
assert_eq!(
cache.async_get_cache("key", &context).await,
Err(Error::MissingPrompt)
);
assert!(cache.embedder().calls.lock().unwrap().is_empty());
server.stop();
}
#[rstest]
#[case::other_key("other", "hello", None)]
#[case::below_similarity_threshold("key", "near", None)]
#[tokio::test(flavor = "multi_thread")]
async fn misses_and_payload_validation_are_safe(
#[case] key: &str,
#[case] prompt: &str,
#[case] numeric_key_point: Option<u64>,
entry: JsonValue,
) {
let server = FakeQdrant::start(FakeState::default()).await;
let cache = connect(
&server,
[("hello", vec![1.0, 0.0]), ("near", vec![0.7, 0.71414286])],
)
.await;
let entry = value(json!({"answer": 1}));
cache
.async_set_cache("key", entry, context("hello"))
.await
.unwrap();
if let Some(id) = numeric_key_point {
server.insert_point(StoredPoint {
id: Some(PointId::from(id)),
vector: vec![1.0, 0.0],
payload: Payload::try_from(json!({
"litellm_cache_key": id,
"response": "{}",
}))
.unwrap()
.into(),
});
}
assert_eq!(
cache
.async_get_cache("other", &context("hello"))
.await
.unwrap(),
None
);
assert_eq!(
cache
.async_get_cache("key", &context("near"))
.await
.unwrap(),
None
);
server.insert_point(StoredPoint {
id: Some(PointId::from(99_u64)),
vector: vec![1.0, 0.0],
payload: Payload::try_from(json!({
"litellm_cache_key": 99,
"response": "{}",
}))
.unwrap()
.into(),
});
assert_eq!(
cache
.async_get_cache("99", &context("hello"))
.await
.unwrap(),
cache.async_get_cache(key, &context(prompt)).await.unwrap(),
None
);
server.stop();
}
#[rstest]
#[case::hit("key", context("hello"), Ok((true, Some(1.0))))]
#[case::below_similarity_threshold("key", context("near"), Ok((false, Some(0.7))))]
#[case::no_results("other", context("hello"), Ok((false, Some(0.0))))]
#[case::no_prompt("key", SemanticCacheContext::default(), Err(Error::MissingPrompt))]
#[tokio::test(flavor = "multi_thread")]
async fn decoding_errors_missing_prompt_pipeline_and_ttl_behave_as_required() {
async fn lookup_reports_python_semantic_similarity(
#[case] key: &'static str,
#[case] context: SemanticCacheContext,
#[case] expected: Result<(bool, Option<f64>), Error>,
#[values(false, true)] use_async: bool,
entry: JsonValue,
) {
let server = FakeQdrant::start(FakeState::default()).await;
let cache = connect(&server, [("one", vec![1.0, 0.0]), ("two", vec![0.0, 1.0])]).await;
let empty = SemanticCacheContext::default();
assert_eq!(
cache
.async_set_cache("key", value(json!({})), empty.clone())
.await,
Err(Error::MissingPrompt)
let cache = Arc::new(
connect(
&server,
[("hello", vec![1.0, 0.0]), ("near", vec![0.7, 0.71414286])],
)
.await,
);
cache
.async_set_cache("key", entry.clone(), self::context("hello"))
.await
.unwrap();
server.insert_point(StoredPoint {
id: Some(PointId::from(99_u64)),
vector: vec![1.0, 0.0],
payload: Payload::try_from(json!({"litellm_cache_key": 99, "response": "{}"}))
.unwrap()
.into(),
});
let lookup = if use_async {
cache.async_get_cache_with_similarity(key, &context).await
} else {
let cache = Arc::clone(&cache);
tokio::task::spawn_blocking(move || cache.get_cache_with_similarity(key, &context))
.await
.unwrap()
};
match (lookup, expected) {
(Ok(SemanticLookup { value, similarity }), Ok((hit, expected))) => {
assert_eq!(value, hit.then_some(entry));
assert_eq!(similarity.is_some(), expected.is_some());
if let (Some(similarity), Some(expected)) = (similarity, expected) {
assert!((similarity - expected).abs() < 1e-6, "{similarity}");
}
}
(lookup, expected) => assert_eq!(lookup.map(|_| ()), expected.map(|_| ())),
}
server.stop();
}
#[rstest]
#[case::codec_decodes_the_payload(Some(json!("{\"a\":1}")), Ok(Some(json!({"a": 1}))))]
#[case::undecodable_response(Some(json!("not json")), Err(Error::InvalidEntry))]
#[case::non_string_response(Some(json!(1)), Err(Error::InvalidEntry))]
#[case::missing_response(None, Err(Error::InvalidEntry))]
#[tokio::test(flavor = "multi_thread")]
async fn stored_responses_go_through_the_codec(
#[case] response: Option<JsonValue>,
#[case] expected: Result<Option<JsonValue>, Error>,
) {
let server = FakeQdrant::start(FakeState::default()).await;
let cache = connect(&server, [("hello", vec![1.0, 0.0])]).await;
let mut payload = serde_json::Map::new();
payload.insert("litellm_cache_key".to_owned(), json!("key"));
if let Some(response) = response {
payload.insert("response".to_owned(), response);
}
server.insert_point(StoredPoint {
id: Some(PointId::from(1_u64)),
vector: vec![1.0, 0.0],
payload: Payload::try_from(JsonValue::Object(payload))
.unwrap()
.into(),
});
assert_eq!(
cache.async_get_cache("key", &empty).await,
Err(Error::MissingPrompt)
cache.async_get_cache("key", &context("hello")).await,
expected
);
server.stop();
}
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn embedding_failures_propagate() {
let server = FakeQdrant::start(FakeState::default()).await;
let cache = connect(&server, []).await;
assert_eq!(
cache.async_get_cache("key", &context("unknown")).await,
Err(Error::Unavailable)
);
server.stop();
}
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn ttl_is_ignored_and_entries_do_not_expire(entry: JsonValue) {
let server = FakeQdrant::start(FakeState::default()).await;
let cache = connect(&server, [("one", vec![1.0, 0.0])]).await;
let ctx = context("one").with_ttl(Some(Duration::from_secs(1)));
assert_eq!(cache.get_ttl(&ctx), None);
cache
.async_set_cache(
"ttl",
value(json!({"ttl": true})),
context("one").with_ttl(Some(Duration::from_secs(1))),
)
.async_set_cache("ttl", entry, ctx.clone())
.await
.unwrap();
tokio::time::sleep(Duration::from_millis(1_100)).await;
assert!(
cache
.async_get_cache(
"ttl",
&context("one").with_ttl(Some(Duration::from_secs(1))),
)
.await
.unwrap()
.is_some()
);
assert!(cache.async_get_cache("ttl", &ctx).await.unwrap().is_some());
server.stop();
}
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn pipeline_upserts_each_entry_and_waits_for_indexing() {
let server = FakeQdrant::start(FakeState::default()).await;
let cache = connect(&server, [("one", vec![1.0, 0.0])]).await;
cache
.async_set_cache_pipeline(
vec![
("one".to_owned(), value(json!({"n": 1}))),
("two".to_owned(), value(json!({"n": 2}))),
("one".to_owned(), json!({"n": 1})),
("two".to_owned(), json!({"n": 2})),
],
context("one"),
)
.await
.unwrap();
assert!(
cache
.async_get_cache("one", &context("one"))
.await
.unwrap()
.is_some()
);
assert!(
cache
.async_get_cache("two", &context("one"))
.await
.unwrap()
.is_some()
);
for (key, value) in [("one", json!({"n": 1})), ("two", json!({"n": 2}))] {
assert_eq!(
cache.async_get_cache(key, &context("one")).await.unwrap(),
Some(value)
);
}
assert_eq!(
server.state.lock().unwrap().upsert_waits,
vec![Some(true), Some(true), Some(true)]
);
assert_eq!(cache.get_ttl(&context("one")), None);
assert_eq!(
cache.test_connection().await,
Err(Error::UnsupportedOperation)
);
server.stop();
}
#[tokio::test(flavor = "multi_thread")]
async fn response_payloads_decode_and_invalid_entries_fail() {
let server = FakeQdrant::start(FakeState::default()).await;
let cache = connect(&server, [("hello", vec![1.0, 0.0])]).await;
for (key, response) in [
("python", json!("{'timestamp': 1.0, 'response': {'a': 1}}")),
("garbage", json!("not json")),
("missing", json!("unused")),
] {
let mut payload = serde_json::Map::new();
payload.insert("litellm_cache_key".to_owned(), json!(key));
if key != "missing" {
payload.insert("response".to_owned(), response);
}
server.insert_point(StoredPoint {
id: Some(PointId::from(key.len() as u64)),
vector: vec![1.0, 0.0],
payload: Payload::try_from(JsonValue::Object(payload))
.unwrap()
.into(),
});
}
assert_eq!(
cache
.async_get_cache("python", &context("hello"))
.await
.unwrap(),
Some(value(json!({"a": 1})))
);
assert_eq!(
cache.async_get_cache("garbage", &context("hello")).await,
Err(Error::InvalidEntry)
);
assert_eq!(
cache.async_get_cache("missing", &context("hello")).await,
Err(Error::InvalidEntry)
);
server.stop();
}
#[tokio::test(flavor = "multi_thread")]
async fn response_cache_facade_turns_invalid_entry_into_miss() {
let server = FakeQdrant::start(FakeState::default()).await;
let cache = Arc::new(connect(&server, [("hello", vec![1.0, 0.0])]).await);
let request = ResponseCacheRequest::<SemanticCacheContext>::new(CacheKeyInput {
preset: Some("key".to_owned()),
..Default::default()
})
.with_context(context("hello"));
let response = json!({"answer": 42});
let facade = ResponseCache::new(cache.clone());
facade
.async_store(&request, response.clone(), Duration::from_secs(1))
.await
.unwrap();
assert_eq!(
facade
.async_lookup(&request, Duration::from_secs(1))
.await
.unwrap(),
Some(response)
);
{
let mut state = server.state.lock().unwrap();
state.points[0]
.payload
.insert("response".to_owned(), Value::from("not json"));
}
assert_eq!(
facade
.async_lookup(&request, Duration::from_secs(1))
.await
.unwrap(),
None
vec![Some(true), Some(true)]
);
server.stop();
}
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn stopped_qdrant_server_maps_to_unavailable() {
let server = FakeQdrant::start(FakeState::default()).await;
@ -420,3 +498,28 @@ async fn stopped_qdrant_server_maps_to_unavailable() {
Err(Error::Unavailable)
);
}
/// `_payload_matches_cache_key` compares `str(cached_key) == str(key)`, so a point whose stored
/// key is the number 99 answers a lookup for `"99"`.
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn numeric_stored_cache_keys_match_like_python_str() {
let server = FakeQdrant::start(FakeState::default()).await;
let cache = connect(&server, [("hello", vec![1.0, 0.0])]).await;
server.insert_point(StoredPoint {
id: Some(PointId::from(99_u64)),
vector: vec![1.0, 0.0],
payload: Payload::try_from(json!({"litellm_cache_key": 99, "response": "{}"}))
.unwrap()
.into(),
});
let lookup = cache
.async_get_cache_with_similarity("99", &context("hello"))
.await
.unwrap();
assert_eq!(lookup.value, Some(json!({})));
assert!((lookup.similarity.unwrap() - 1.0).abs() < 1e-6);
server.stop();
}

5
litellm-rust/crates/cache-qdrant-semantic/tests/support/mod.rs
View file

@ -1,15 +1,16 @@
#![allow(dead_code)]
use std::{
collections::{HashMap, HashSet},
net::SocketAddr,
sync::{Arc, Mutex},
};
use qdrant_client::qdrant::collections_server::CollectionsServer;
use qdrant_client::qdrant::{
self, CollectionExists, CollectionExistsRequest, CollectionExistsResponse,
CollectionOperationResponse, CreateCollection, CreateFieldIndexCollection, Filter, PointId,
PointsOperationResponse, ScoredPoint, SearchPoints, SearchResponse, Value, Vector, Vectors,
collections_server::Collections,
collections_server::{Collections, CollectionsServer},
points_server::{Points, PointsServer},
};
use tokio::sync::oneshot;

6
litellm-rust/crates/cache-redis-semantic/Cargo.toml
View file

@ -8,14 +8,12 @@ repository.workspace = true
[dependencies]
litellm-cache.workspace = true
litellm-cache-redis.workspace = true
litellm-cache-response.workspace = true
redis = { version = "1.7.0", features = ["tls-rustls"] }
r2d2 = "0.8.10"
serde_json.workspace = true
sha2.workspace = true
tokio.workspace = true
[dev-dependencies]
litellm-cache-testing.workspace = true
redis-test = "1.0.4"
rstest.workspace = true
serde_json.workspace = true
tokio.workspace = true

568
litellm-rust/crates/cache-redis-semantic/src/cache.rs
View file

@ -1,105 +1,36 @@
use std::{
future::Future,
sync::{Arc, OnceLock},
sync::Arc,
time::{Duration, SystemTime, UNIX_EPOCH},
};
use litellm_cache::{
BaseCache, CacheCodec, CacheConnectionResult, CacheConnectionStatus, Error,
SemanticCacheContext,
BaseCache, CacheCodec, Error, SemanticCacheContext,
semantic::{Embedder, SemanticCache, SemanticLookup, prompt_from_context},
};
use litellm_cache_redis::{
RedisTopology,
connection::{ConnectionRef, Connections},
};
use litellm_cache_response::{CacheEntry, ResponseCacheCodec};
use serde_json::Value;
use sha2::{Digest, Sha256};
use crate::prompt::prompt_from_context;
const CACHE_KEY_FIELD: &str = "litellm_cache_key";
const VECTOR_FIELD: &str = "prompt_vector";
pub trait Embedder: Send + Sync + 'static {
fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result<Vec<f32>, Error>;
fn async_embed(
&self,
prompt: &str,
metadata: Option<&Value>,
) -> impl Future<Output = Result<Vec<f32>, Error>> + Send;
}
#[derive(Clone, Debug)]
pub struct RedisSemanticConfig {
pub index_name: String,
pub similarity_threshold: f32,
}
use crate::{
RedisSemanticConfig,
index::{CACHE_KEY_FIELD, Index, VECTOR_FIELD},
reply::{bytes_field, first_document, number_field, string_field},
};
struct Inner {
index_name: String,
index: Index,
distance_threshold: f64,
resolved_index: OnceLock<String>,
codec: ResponseCacheCodec,
clock: fn() -> f64,
}
impl Inner {
fn new(config: RedisSemanticConfig) -> Self {
fn new(config: RedisSemanticConfig, clock: fn() -> f64) -> Self {
Self {
index_name: config.index_name,
index: Index::new(config.index_name),
distance_threshold: 1.0 - f64::from(config.similarity_threshold),
resolved_index: OnceLock::new(),
codec: ResponseCacheCodec,
clock: timestamp,
}
}
fn ensure_index(
&self,
connection: &mut ConnectionRef<'_>,
dims: usize,
) -> Result<String, Error> {
if let Some(name) = self.resolved_index.get() {
return Ok(name.clone());
}
let name = match index_compatible(connection, &self.index_name, dims)? {
Some(true) => self.index_name.clone(),
Some(false) => self.isolated_index(connection, dims)?,
None => match create_index(connection, &self.index_name, dims) {
Ok(()) => self.index_name.clone(),
Err(_) => match index_compatible(connection, &self.index_name, dims)? {
Some(true) => self.index_name.clone(),
Some(false) => self.isolated_index(connection, dims)?,
None => return Err(Error::Unavailable),
},
},
};
let _ = self.resolved_index.set(name.clone());
Ok(name)
}
fn isolated_index(
&self,
connection: &mut ConnectionRef<'_>,
dims: usize,
) -> Result<String, Error> {
let name = format!("{}_isolated", self.index_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)
}
clock,
}
}
@ -107,15 +38,14 @@ impl Inner {
&self,
connection: &mut ConnectionRef<'_>,
tag: &str,
value: &CacheEntry,
response: Vec<u8>,
prompt: &str,
vector: &[f32],
ttl: Option<Duration>,
) -> Result<(), Error> {
let index = self.ensure_index(connection, vector.len())?;
let index = self.index.ensure(connection, vector.len())?;
let entry_id = entry_id(prompt, tag);
let hash_key = format!("{index}:{entry_id}");
let response = self.codec.encode(value)?;
redis::cmd("HSET")
.arg(&hash_key)
.arg("entry_id")
@ -149,8 +79,8 @@ impl Inner {
connection: &mut ConnectionRef<'_>,
tag: &str,
vector: &[f32],
) -> Result<Option<CacheEntry>, Error> {
let index = self.ensure_index(connection, vector.len())?;
) -> 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)
@ -183,57 +113,80 @@ impl Inner {
.query::<redis::Value>(connection)
.map_err(|_| Error::Unavailable)?;
let Some(fields) = first_document(&result) else {
return Ok(None);
return Ok(SemanticLookup::miss(Some(0.0)));
};
if string_field(fields, CACHE_KEY_FIELD).as_deref() != Some(tag) {
return Ok(None);
return Ok(SemanticLookup::miss(Some(0.0)));
}
if number_field(fields, "vector_distance")
.is_none_or(|distance| distance > self.distance_threshold)
{
return Ok(None);
}
let Some(response) = bytes_field(fields, "response") else {
return Ok(None);
// 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)));
};
self.codec.decode(&response).map(Some)
}
}
pub struct RedisSemanticCache<E: Embedder, C = redis::Connection> {
connections: Arc<Connections<C>>,
embedder: E,
inner: Arc<Inner>,
}
impl<E: Embedder> RedisSemanticCache<E> {
pub fn new(url: &str, embedder: E, config: RedisSemanticConfig) -> Result<Self, Error> {
Ok(Self {
connections: Arc::new(Connections::open(url, &RedisTopology::Standalone)?),
embedder,
inner: Arc::new(Inner::new(config)),
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),
})
}
}
impl<E: Embedder, C: redis::ConnectionLike + Send + 'static> RedisSemanticCache<E, C> {
pub fn with_connection(connection: C, embedder: E, config: RedisSemanticConfig) -> Self {
/// `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,
inner: Arc::new(Inner::new(config)),
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 {
index_name: self.inner.index_name.clone(),
distance_threshold: self.inner.distance_threshold,
resolved_index: OnceLock::new(),
codec: self.inner.codec,
clock,
}),
inner: Arc::new(Inner::new(config, clock)),
..self
}
}
@ -243,7 +196,7 @@ impl<E: Embedder, C: redis::ConnectionLike + Send + 'static> RedisSemanticCache<
}
pub fn index_name(&self) -> &str {
&self.inner.index_name
self.inner.index.name()
}
pub fn similarity_threshold(&self) -> f32 {
@ -253,12 +206,25 @@ impl<E: Embedder, C: redis::ConnectionLike + Send + 'static> RedisSemanticCache<
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: Embedder, C: redis::ConnectionLike + Send + 'static> BaseCache
for RedisSemanticCache<E, C>
impl<E, S, C> BaseCache for RedisSemanticCache<E, S, C>
where
E: Embedder,
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
type Value = CacheEntry;
type Value = S::Value;
type Context = SemanticCacheContext;
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
@ -274,22 +240,18 @@ impl<E: Embedder, C: redis::ConnectionLike + Send + 'static> BaseCache
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).to_string();
let tag = Self::tag(key, context);
self.connections.execute(|connection| {
self.inner
.store(connection, &tag, &value, &prompt, &vector, context.ttl)
.store(connection, tag, response, &prompt, &vector, context.ttl)
})
}
fn get_cache(&self, key: &str, context: &Self::Context) -> Result<Option<Self::Value>, Error> {
let Some(prompt) = prompt_from_context(context) else {
return Ok(None);
};
let vector = self.embedder.embed(&prompt, context.metadata.as_ref())?;
let tag = Self::tag(key, context).to_string();
self.connections
.execute(|connection| self.inner.lookup(connection, &tag, &vector))
self.get_cache_with_similarity(key, context)
.map(|lookup| lookup.value)
}
async fn async_set_cache(
@ -301,14 +263,15 @@ impl<E: Embedder, C: redis::ConnectionLike + Send + 'static> BaseCache
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_string();
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, &value, &prompt, &vector, context.ttl)
inner.store(connection, &tag, response, &prompt, &vector, context.ttl)
})
.await
}
@ -318,49 +281,54 @@ impl<E: Embedder, C: redis::ConnectionLike + Send + 'static> BaseCache
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(None);
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_string();
let tag = Self::tag(key, context).to_owned();
let inner = Arc::clone(&self.inner);
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
let lookup = Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
inner.lookup(connection, &tag, &vector)
})
.await
}
async fn disconnect(&self) -> Result<(), Error> {
Ok(())
}
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
match Connections::run_blocking(Arc::clone(&self.connections), |connection| {
Ok(match redis::cmd("PING").query::<String>(connection) {
Ok(_) => CacheConnectionResult {
status: CacheConnectionStatus::Success,
message: "Redis cache connection test successful".into(),
error: None,
},
Err(error) => CacheConnectionResult {
status: CacheConnectionStatus::Failed,
message: format!("Redis connection failed: {error}"),
error: Some(error.to_string()),
},
})
})
.await
{
Ok(result) => Ok(result),
Err(error) => Ok(CacheConnectionResult {
status: CacheConnectionStatus::Failed,
message: format!("Redis connection failed: {error}"),
error: Some(error.to_string()),
}),
}
.await?;
self.decode(lookup)
}
}
@ -387,228 +355,46 @@ fn vector_buffer(vector: &[f32]) -> Vec<u8> {
}
fn escape_tag(value: &str) -> String {
value
.chars()
.flat_map(|ch| {
if matches!(
ch,
',' | '.'
| '<'
| '>'
| '{'
| '}'
| '['
| ']'
| '\\'
| '"'
| '\''
| ':'
| ';'
| '!'
| '@'
| '#'
| '$'
| '%'
| '^'
| '&'
| '*'
| '('
| ')'
| '-'
| '+'
| '='
| '~'
| '|'
| '/'
| ' '
| '?'
) {
vec!['\\', ch]
} else {
vec![ch]
}
})
.collect()
}
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(_) => return Err(Error::Unavailable),
};
Ok(Some(schema_compatible(&info, dims)))
}
fn unknown_index(error: &redis::RedisError) -> bool {
let message = error.to_string().to_lowercase();
message.contains("unknown") && message.contains("index")
}
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()
.map(|attribute| {
let redis::Value::Array(attribute) = attribute else {
return (None, None, None, None, None);
};
let mut name = None;
let mut field_type = None;
let mut dim = None;
let mut data_type = None;
let mut distance_metric = None;
for pair in attribute.as_chunks::<2>().0 {
match string_value(&pair[0]).as_deref() {
Some("identifier") => name = string_value(&pair[1]),
Some("type") => field_type = string_value(&pair[1]),
Some("dim") => dim = number_value(&pair[1]),
Some("data_type") => data_type = string_value(&pair[1]),
Some("distance_metric") => distance_metric = string_value(&pair[1]),
_ => {}
}
}
(name, field_type, dim, data_type, distance_metric)
})
.collect::<Vec<_>>();
let has_field = |name: &str, field_type: &str| {
fields
.iter()
.any(|(n, t, ..)| n.as_deref() == Some(name) && t.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(|(n, t, d, data, metric)| {
n.as_deref() == Some(VECTOR_FIELD)
&& t.as_deref() == Some("VECTOR")
&& *d == Some(dims as f64)
&& data
.as_deref()
.is_some_and(|data| data.eq_ignore_ascii_case("float32"))
&& metric
.as_deref()
.is_some_and(|metric| metric.eq_ignore_ascii_case("cosine"))
})
}
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,
}
}
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()),
}
}
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])
}
fn string_field(fields: &[redis::Value], name: &str) -> Option<String> {
field_value(fields, name).and_then(string_value)
}
fn number_field(fields: &[redis::Value], name: &str) -> Option<f64> {
field_value(fields, name).and_then(number_value)
}
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,
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 {

8
litellm-rust/crates/cache-redis-semantic/src/config.rs Normal file
View file

@ -0,0 +1,8 @@
/// `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,
}

205
litellm-rust/crates/cache-redis-semantic/src/index.rs Normal file
View file

@ -0,0 +1,205 @@
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(_) => match index_compatible(connection, &self.name, dims)? {
Some(true) => self.name.clone(),
Some(false) => self.isolated(connection, dims)?,
None => return Err(Error::Unavailable),
},
},
};
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(_) => return Err(Error::Unavailable),
};
Ok(Some(schema_compatible(&info, dims)))
}
fn unknown_index(error: &redis::RedisError) -> bool {
let message = error.to_string().to_lowercase();
message.contains("unknown") && message.contains("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"))
})
}

Some files were not shown because too many files have changed in this diff Show more