mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-25 01:02:15 +00:00
Merge remote-tracking branch 'origin/main' into litellm_lit_5534_cost_map_metadata_echo
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> # Conflicts: # litellm/proxy/management_endpoints/model_management_endpoints.py
This commit is contained in:
commit
099a8f66eb
1140 changed files with 41631 additions and 7496 deletions
|
|
@ -1508,7 +1508,7 @@ jobs:
|
|||
- run:
|
||||
name: Run tests
|
||||
command: |
|
||||
uv run --no-sync python -m pytest -vv tests/local_testing/test_basic_python_version.py -k "not v2_resolver"
|
||||
uv run --no-sync python -m pytest -vv tests/local_testing/test_basic_python_version.py -k "not legacy_resolver"
|
||||
|
||||
installing_litellm_on_python_3_13:
|
||||
docker:
|
||||
|
|
@ -1532,7 +1532,7 @@ jobs:
|
|||
- run:
|
||||
name: Run tests
|
||||
command: |
|
||||
uv run --no-sync python -m pytest -v tests/local_testing/test_basic_python_version.py -k "not v2_resolver"
|
||||
uv run --no-sync python -m pytest -v tests/local_testing/test_basic_python_version.py -k "not legacy_resolver"
|
||||
|
||||
installing_litellm_on_python_v2_migration_resolver:
|
||||
docker:
|
||||
|
|
@ -1561,10 +1561,11 @@ jobs:
|
|||
url: tcp://localhost:5432
|
||||
timeout: "60"
|
||||
- run:
|
||||
name: Run v2 migration resolver proxy smoke test
|
||||
name: Run both migration resolvers against Postgres
|
||||
command: |
|
||||
uv run --no-sync python -m pytest -vv \
|
||||
tests/local_testing/test_basic_python_version.py::test_litellm_proxy_server_config_no_general_settings_v2_resolver
|
||||
tests/local_testing/test_basic_python_version.py::test_litellm_proxy_server_config_no_general_settings \
|
||||
tests/local_testing/test_basic_python_version.py::test_litellm_proxy_server_config_no_general_settings_legacy_resolver
|
||||
|
||||
helm_chart_testing:
|
||||
machine:
|
||||
|
|
|
|||
|
|
@ -121,6 +121,10 @@ start_proxy() {
|
|||
"LITELLM_MODEL_COST_MAP_URL=$INTEGRATION_UPSTREAM_URL/_cost_map"
|
||||
"MODEL_COST_MAP_MIN_MODEL_COUNT=1"
|
||||
"MODEL_COST_MAP_MAX_SHRINK_RATIO=0"
|
||||
"GEMINI_API_BASE=$INTEGRATION_UPSTREAM_URL"
|
||||
"ANTHROPIC_API_BASE=$INTEGRATION_UPSTREAM_URL"
|
||||
"GEMINI_API_KEY=sk-scripted-provider"
|
||||
"ANTHROPIC_API_KEY=sk-scripted-provider"
|
||||
)
|
||||
else
|
||||
cost_map_env=("LITELLM_LOCAL_MODEL_COST_MAP=True")
|
||||
|
|
|
|||
83
.github/e2e-stack/redact_output.py
vendored
Normal file
83
.github/e2e-stack/redact_output.py
vendored
Normal file
|
|
@ -0,0 +1,83 @@
|
|||
import argparse
|
||||
import os
|
||||
import sys
|
||||
from functools import reduce
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
from xml.sax.saxutils import escape
|
||||
|
||||
from pydantic import JsonValue, TypeAdapter, ValidationError
|
||||
from secrets_to_env import MIN_MASKED_LENGTH
|
||||
|
||||
REDACTED: Final = "***"
|
||||
json_adapter: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
|
||||
|
||||
|
||||
def string_leaves(node: JsonValue) -> tuple[str, ...]:
|
||||
match node:
|
||||
case str():
|
||||
return (node,)
|
||||
case list():
|
||||
return tuple(leaf for child in node for leaf in string_leaves(child))
|
||||
case dict():
|
||||
return tuple(leaf for child in node.values() for leaf in string_leaves(child))
|
||||
return ()
|
||||
|
||||
|
||||
def field_lines(value: str) -> tuple[str, ...]:
|
||||
try:
|
||||
return tuple(line for leaf in string_leaves(json_adapter.validate_json(value)) for line in leaf.splitlines())
|
||||
except ValidationError:
|
||||
return ()
|
||||
|
||||
|
||||
def masked_values(values_files: tuple[Path, ...]) -> tuple[str, ...]:
|
||||
values: Final = frozenset(
|
||||
line.split("=", 1)[1].strip().strip("'")
|
||||
for path in values_files
|
||||
for line in path.read_text().splitlines()
|
||||
if "=" in line
|
||||
)
|
||||
texts: Final = frozenset(text for value in values for text in (value, *field_lines(value)))
|
||||
renderings: Final = frozenset(
|
||||
rendering
|
||||
for text in texts
|
||||
if len(text) >= MIN_MASKED_LENGTH
|
||||
for rendering in (text, escape(text), escape(text, {'"': """}))
|
||||
)
|
||||
return tuple(sorted(renderings, key=lambda rendering: (-len(rendering), rendering)))
|
||||
|
||||
|
||||
def redact(text: str, values: tuple[str, ...]) -> str:
|
||||
return reduce(lambda redacted, value: redacted.replace(value, REDACTED), values, text)
|
||||
|
||||
|
||||
def write_redacted(source: Path, out_dir: Path, values: tuple[str, ...]) -> None:
|
||||
target: Final = out_dir / source.name
|
||||
with os.fdopen(os.open(target, os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW, 0o600), "w") as handle:
|
||||
_ = handle.write(redact(source.read_text(errors="replace"), values))
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser: Final = argparse.ArgumentParser()
|
||||
_ = parser.add_argument("--values", action="append", type=Path, required=True)
|
||||
_ = parser.add_argument("--out", type=Path, required=True)
|
||||
_ = parser.add_argument("files", nargs="*", type=Path)
|
||||
args: Final = parser.parse_args()
|
||||
values_files: Final = tuple(args.values)
|
||||
out_dir: Final[Path] = args.out
|
||||
sources: Final = tuple(args.files)
|
||||
try:
|
||||
values: Final = masked_values(values_files)
|
||||
out_dir.mkdir(mode=0o700, exist_ok=True)
|
||||
for source in sources:
|
||||
write_redacted(source, out_dir, values)
|
||||
except OSError as error:
|
||||
_ = sys.stderr.write(f"could not redact {error.filename}\n")
|
||||
return 1
|
||||
_ = sys.stdout.write(f"redacted {len(sources)} file(s) into {out_dir}\n")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
2
.github/e2e-stack/up.sh
vendored
2
.github/e2e-stack/up.sh
vendored
|
|
@ -143,7 +143,7 @@ env "${SERVER_ENV[@]}" uv run --no-sync python migrations/run.py >"${LOGS_DIR}/m
|
|||
|
||||
start_server() {
|
||||
local name="$1"; shift
|
||||
env "${SERVER_ENV[@]}" "$@" >"${LOGS_DIR}/${name}.log" 2>&1 &
|
||||
env -u AWS_ROLE_NAME "${SERVER_ENV[@]}" "$@" >"${LOGS_DIR}/${name}.log" 2>&1 &
|
||||
echo $! > "${PIDS_DIR}/${name}.pid"
|
||||
}
|
||||
|
||||
|
|
|
|||
23
.github/workflows/test-e2e-changed.yml
vendored
23
.github/workflows/test-e2e-changed.yml
vendored
|
|
@ -175,6 +175,8 @@ jobs:
|
|||
env:
|
||||
TESTS: ${{ needs.detect.outputs.tests }}
|
||||
E2E_FIXTURE_MODE: live
|
||||
E2E_PROVIDER_EDGE_HOST_REACHABLE: '1'
|
||||
COLUMNS: '400'
|
||||
run: |
|
||||
umask 077
|
||||
read -r -a test_files <<< "${TESTS}"
|
||||
|
|
@ -189,6 +191,7 @@ jobs:
|
|||
uv run --no-sync python .github/e2e-stack/assert_tests_ran.py "${report}" "${test_files[@]}"
|
||||
verified=$?
|
||||
set -e
|
||||
grep -E '^(FAILED|ERROR) ' "${log}" || true
|
||||
grep -E '^=+ .* in [0-9.]+s( \([0-9:]+\))? =+$' "${log}" | tail -n 1
|
||||
echo "::endgroup::"
|
||||
if [ "${status}" = "5" ]; then
|
||||
|
|
@ -206,6 +209,24 @@ jobs:
|
|||
echo "pass ${pass} of 3 passed"
|
||||
done
|
||||
|
||||
- name: Redact the pytest output
|
||||
if: always() && steps.boot.outcome == 'success'
|
||||
run: |
|
||||
umask 077
|
||||
shopt -s nullglob
|
||||
uv run --no-sync python .github/e2e-stack/redact_output.py \
|
||||
--values tests/e2e/.env --values "${RUNNER_TEMP}/litellm-e2e-stack/stack.env" \
|
||||
--out "${RUNNER_TEMP}/e2e-redacted" "${RUNNER_TEMP}"/e2e-pass-*.log "${RUNNER_TEMP}"/e2e-pass-*.xml
|
||||
|
||||
- name: Keep the redacted pytest output
|
||||
if: always() && steps.boot.outcome == 'success'
|
||||
uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1
|
||||
with:
|
||||
name: e2e-changed-pytest-output-${{ github.run_attempt }}
|
||||
path: ${{ runner.temp }}/e2e-redacted
|
||||
retention-days: 14
|
||||
if-no-files-found: ignore
|
||||
|
||||
- name: Stop the stack
|
||||
if: always() && steps.boot.outcome != 'skipped'
|
||||
run: bash .github/e2e-stack/down.sh
|
||||
|
|
@ -214,7 +235,7 @@ jobs:
|
|||
if: always()
|
||||
run: |
|
||||
rm -f tests/e2e/.env "${RUNNER_TEMP}/e2e-boot.log" "${RUNNER_TEMP}"/e2e-pass-*.log "${RUNNER_TEMP}"/e2e-pass-*.xml
|
||||
rm -rf "${RUNNER_TEMP}/litellm-e2e-stack"
|
||||
rm -rf "${RUNNER_TEMP}/litellm-e2e-stack" "${RUNNER_TEMP}/e2e-redacted"
|
||||
|
||||
gate:
|
||||
name: e2e-changed-tests
|
||||
|
|
|
|||
14
.github/workflows/test-rust.yml
vendored
14
.github/workflows/test-rust.yml
vendored
|
|
@ -120,6 +120,20 @@ jobs:
|
|||
|
||||
- run: cargo test --workspace --doc --locked
|
||||
|
||||
- name: Test token counter feature combinations
|
||||
run: |
|
||||
for features in '' fast huggingface tiktoken fast,huggingface fast,tiktoken huggingface,tiktoken fast,huggingface,tiktoken; do
|
||||
cargo test -p litellm-token-counter --locked --no-default-features --features "$features"
|
||||
cargo check -p litellm-python-bridge --locked --no-default-features --features "abi3${features:+,$features}"
|
||||
done
|
||||
|
||||
- name: Test secret manager feature combinations
|
||||
run: |
|
||||
cargo test -p litellm-auth-gcp --locked --no-default-features
|
||||
for features in '' aws google aws,google; do
|
||||
cargo test -p litellm-secrets --locked --no-default-features --features "$features"
|
||||
done
|
||||
|
||||
rust-wheel:
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
|
|
|||
3
.github/workflows/test-unit.yml
vendored
3
.github/workflows/test-unit.yml
vendored
|
|
@ -51,7 +51,7 @@ jobs:
|
|||
include:
|
||||
- shard: mcp-integration
|
||||
artifact-name: mcp-integration
|
||||
test-path: "tests/mcp_tests"
|
||||
test-path: "tests/mcp_tests tests/test_litellm/experimental_mcp_client"
|
||||
workers: 2
|
||||
reruns: 0
|
||||
timeout-minutes: 20
|
||||
|
|
@ -113,7 +113,6 @@ jobs:
|
|||
tests/test_litellm/compression
|
||||
tests/test_litellm/containers
|
||||
tests/test_litellm/endpoints
|
||||
tests/test_litellm/experimental_mcp_client
|
||||
tests/test_litellm/models
|
||||
tests/test_litellm/repositories
|
||||
tests/test_litellm/images
|
||||
|
|
|
|||
|
|
@ -307,6 +307,7 @@ For MCP OAuth, an upstream may advertise dynamic client registration but refuse
|
|||
| [Deepgram (`deepgram`)](https://docs.litellm.ai/docs/providers/deepgram) | ✅ | ✅ | ✅ | | | ✅ | | | | |
|
||||
| [DeepInfra (`deepinfra`)](https://docs.litellm.ai/docs/providers/deepinfra) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [Deepseek (`deepseek`)](https://docs.litellm.ai/docs/providers/deepseek) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [Eden AI (`edenai`)](https://docs.litellm.ai/docs/providers/edenai) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | | | |
|
||||
| [ElevenLabs (`elevenlabs`)](https://docs.litellm.ai/docs/providers/elevenlabs) | ✅ | ✅ | ✅ | | | ✅ | ✅ | | | |
|
||||
| [Empower (`empower`)](https://docs.litellm.ai/docs/providers/empower) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [Fal AI (`fal_ai`)](https://docs.litellm.ai/docs/providers/fal_ai) | ✅ | ✅ | ✅ | | ✅ | | | | | |
|
||||
|
|
@ -356,7 +357,7 @@ For MCP OAuth, an upstream may advertise dynamic client registration but refuse
|
|||
| [Petals (`petals`)](https://docs.litellm.ai/docs/providers/petals) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [Pinstripes (`pinstripes`)](https://docs.litellm.ai/docs/providers/pinstripes) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [Predibase (`predibase`)](https://docs.litellm.ai/docs/providers/predibase) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [Qwen AI Platform (`qwen_ai_platform`)](https://docs.litellm.ai/docs/providers/qwencloud) | ✅ | ✅ | ✅ | ✅ | ✅ | | | | | ✅ |
|
||||
| [Qianwen AI Platform (`qwen_ai_platform`)](https://docs.litellm.ai/docs/providers/qwencloud) | ✅ | ✅ | ✅ | ✅ | ✅ | | | | | ✅ |
|
||||
| [QwenCloud (`qwencloud`)](https://docs.litellm.ai/docs/providers/qwencloud) | ✅ | ✅ | ✅ | ✅ | ✅ | | | | | ✅ |
|
||||
| [Recraft (`recraft`)](https://docs.litellm.ai/docs/providers/recraft) | | | | | ✅ | | | | | |
|
||||
| [Replicate (`replicate`)](https://docs.litellm.ai/docs/providers/replicate) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
|
|
|
|||
|
|
@ -51,6 +51,7 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = (
|
|||
"/cache_settings",
|
||||
"/coordination_redis/",
|
||||
"/cost_tracking",
|
||||
"/cost_optimization/",
|
||||
"/cost/",
|
||||
"/credentials",
|
||||
"/credential",
|
||||
|
|
|
|||
|
|
@ -7,6 +7,9 @@ metadata:
|
|||
{{- include "litellm.commonLabels" . | nindent 4 }}
|
||||
app.kubernetes.io/component: backend
|
||||
spec:
|
||||
{{- if and (not .Values.backend.hpa.enabled) (not (kindIs "invalid" .Values.backend.replicaCount)) }}
|
||||
replicas: {{ .Values.backend.replicaCount }}
|
||||
{{- end }}
|
||||
{{- with .Values.backend.strategy }}
|
||||
strategy:
|
||||
{{- toYaml . | nindent 4 }}
|
||||
|
|
|
|||
|
|
@ -7,6 +7,9 @@ metadata:
|
|||
{{- include "litellm.commonLabels" . | nindent 4 }}
|
||||
app.kubernetes.io/component: gateway
|
||||
spec:
|
||||
{{- if and (not .Values.gateway.hpa.enabled) (not (kindIs "invalid" .Values.gateway.replicaCount)) }}
|
||||
replicas: {{ .Values.gateway.replicaCount }}
|
||||
{{- end }}
|
||||
{{- with .Values.gateway.strategy }}
|
||||
strategy:
|
||||
{{- toYaml . | nindent 4 }}
|
||||
|
|
|
|||
|
|
@ -7,6 +7,9 @@ metadata:
|
|||
{{- include "litellm.commonLabels" . | nindent 4 }}
|
||||
app.kubernetes.io/component: ui
|
||||
spec:
|
||||
{{- if and (not .Values.ui.hpa.enabled) (not (kindIs "invalid" .Values.ui.replicaCount)) }}
|
||||
replicas: {{ .Values.ui.replicaCount }}
|
||||
{{- end }}
|
||||
{{- with .Values.ui.strategy }}
|
||||
strategy:
|
||||
{{- toYaml . | nindent 4 }}
|
||||
|
|
|
|||
100
helm/litellm/tests/replica_count_tests.yaml
Normal file
100
helm/litellm/tests/replica_count_tests.yaml
Normal file
|
|
@ -0,0 +1,100 @@
|
|||
suite: test fixed replica count when HPA is disabled
|
||||
templates:
|
||||
- gateway/deployment.yaml
|
||||
- gateway/configmap.yaml
|
||||
- backend/deployment.yaml
|
||||
- ui/deployment.yaml
|
||||
values:
|
||||
- ./values/required.yaml
|
||||
tests:
|
||||
- it: gateway renders replicaCount into spec.replicas when its HPA is disabled
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
gateway.hpa.enabled: false
|
||||
gateway.replicaCount: 3
|
||||
asserts:
|
||||
- isKind:
|
||||
of: Deployment
|
||||
- equal:
|
||||
path: spec.replicas
|
||||
value: 3
|
||||
|
||||
- it: backend renders replicaCount into spec.replicas when its HPA is disabled
|
||||
template: backend/deployment.yaml
|
||||
set:
|
||||
backend.hpa.enabled: false
|
||||
backend.replicaCount: 2
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.replicas
|
||||
value: 2
|
||||
|
||||
- it: ui renders replicaCount into spec.replicas when its HPA is disabled
|
||||
template: ui/deployment.yaml
|
||||
set:
|
||||
ui.hpa.enabled: false
|
||||
ui.replicaCount: 2
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.replicas
|
||||
value: 2
|
||||
|
||||
- it: replicaCount 0 scales the gateway to zero instead of being treated as unset
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
gateway.hpa.enabled: false
|
||||
gateway.replicaCount: 0
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.replicas
|
||||
value: 0
|
||||
|
||||
- it: a component with HPA disabled but no replicaCount set keeps omitting spec.replicas, so upgrades do not reset a hand-scaled Deployment
|
||||
set:
|
||||
gateway.hpa.enabled: false
|
||||
backend.hpa.enabled: false
|
||||
ui.hpa.enabled: false
|
||||
asserts:
|
||||
- notExists:
|
||||
path: spec.replicas
|
||||
template: gateway/deployment.yaml
|
||||
- notExists:
|
||||
path: spec.replicas
|
||||
template: backend/deployment.yaml
|
||||
- notExists:
|
||||
path: spec.replicas
|
||||
template: ui/deployment.yaml
|
||||
|
||||
- it: every component omits spec.replicas when its HPA is enabled, so the autoscaler owns the count
|
||||
set:
|
||||
gateway.hpa.enabled: true
|
||||
gateway.replicaCount: 3
|
||||
backend.hpa.enabled: true
|
||||
backend.replicaCount: 3
|
||||
ui.hpa.enabled: true
|
||||
ui.replicaCount: 3
|
||||
asserts:
|
||||
- notExists:
|
||||
path: spec.replicas
|
||||
template: gateway/deployment.yaml
|
||||
- notExists:
|
||||
path: spec.replicas
|
||||
template: backend/deployment.yaml
|
||||
- notExists:
|
||||
path: spec.replicas
|
||||
template: ui/deployment.yaml
|
||||
|
||||
- it: a component with HPA disabled renders replicas while a sibling with HPA enabled does not
|
||||
set:
|
||||
gateway.hpa.enabled: false
|
||||
gateway.replicaCount: 4
|
||||
backend.hpa.enabled: true
|
||||
backend.replicaCount: 4
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.replicas
|
||||
value: 4
|
||||
template: gateway/deployment.yaml
|
||||
- notExists:
|
||||
path: spec.replicas
|
||||
template: backend/deployment.yaml
|
||||
|
|
@ -397,6 +397,11 @@ gateway:
|
|||
# failureThreshold: 30
|
||||
# periodSeconds: 10
|
||||
startupProbe: {}
|
||||
# Optional fixed pod count, rendered into the Deployment's spec.replicas only
|
||||
# when hpa.enabled is false. Unset by default so an existing Deployment keeps
|
||||
# its current count; with the HPA on, the autoscaler owns the count, e.g.:
|
||||
# replicaCount: 3
|
||||
replicaCount:
|
||||
hpa:
|
||||
enabled: true
|
||||
minReplicas: 1
|
||||
|
|
@ -524,6 +529,8 @@ backend:
|
|||
strategy: {}
|
||||
# Optional startupProbe; same shape as gateway.startupProbe. Empty by default.
|
||||
startupProbe: {}
|
||||
# Same semantics as gateway.replicaCount.
|
||||
replicaCount:
|
||||
hpa:
|
||||
enabled: true
|
||||
minReplicas: 1
|
||||
|
|
@ -590,6 +597,8 @@ ui:
|
|||
strategy: {}
|
||||
# Optional startupProbe; same shape as gateway.startupProbe. Empty by default.
|
||||
startupProbe: {}
|
||||
# Same semantics as gateway.replicaCount.
|
||||
replicaCount:
|
||||
hpa:
|
||||
enabled: false
|
||||
minReplicas: 1
|
||||
|
|
|
|||
|
|
@ -0,0 +1 @@
|
|||
ALTER TABLE "LiteLLM_PolicyAttachmentTable" ADD COLUMN IF NOT EXISTS "is_default" BOOLEAN NOT NULL DEFAULT false;
|
||||
|
|
@ -1419,6 +1419,7 @@ model LiteLLM_PolicyAttachmentTable {
|
|||
models String[] @default([]) // Model names or patterns
|
||||
tags String[] @default([]) // Tag patterns (e.g., ["healthcare", "prod-*"])
|
||||
priority Int? // Explicit execution order
|
||||
is_default Boolean @default(false) // Applied only when no non-default attachment matches
|
||||
created_at DateTime @default(now())
|
||||
created_by String?
|
||||
updated_at DateTime @default(now()) @updatedAt
|
||||
|
|
|
|||
1021
litellm-rust/Cargo.lock
generated
1021
litellm-rust/Cargo.lock
generated
File diff suppressed because it is too large
Load diff
|
|
@ -14,20 +14,32 @@ litellm-host = { path = "crates/host" }
|
|||
litellm-callbacks-legacy-python = { path = "crates/callbacks-legacy-python" }
|
||||
litellm-framing = { path = "crates/framer" }
|
||||
litellm-auth = { path = "crates/auth" }
|
||||
litellm-auth-types = { path = "crates/auth-types" }
|
||||
litellm-auth-aws = { path = "crates/auth-aws" }
|
||||
litellm-auth-azure = { path = "crates/auth-azure" }
|
||||
litellm-auth-gcp = { path = "crates/auth-gcp" }
|
||||
litellm-secrets = { path = "crates/secrets" }
|
||||
litellm-secrets-types = { path = "crates/secrets-types" }
|
||||
litellm-secrets-aws = { path = "crates/secrets-aws" }
|
||||
litellm-secrets-google = { path = "crates/secrets-google" }
|
||||
litellm-http = { path = "crates/http" }
|
||||
litellm-llms = { path = "crates/llms" }
|
||||
litellm-types = { path = "crates/types" }
|
||||
litellm-core-utils = { path = "crates/core-utils" }
|
||||
litellm-cache = { path = "crates/cache" }
|
||||
litellm-cache-memory = { path = "crates/cache-memory" }
|
||||
litellm-cache-redis = { path = "crates/cache-redis" }
|
||||
litellm-cache-response = { path = "crates/cache-response" }
|
||||
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" }
|
||||
litellm-token-counter-tiktoken = { path = "crates/token-counter-tiktoken" }
|
||||
litellm-host-python = { path = "crates/host-python" }
|
||||
|
||||
bytes = "1"
|
||||
http = "1"
|
||||
google-cloud-auth = { version = "1.16.0", default-features = false }
|
||||
jsonwebtoken = { version = "11.1.0", default-features = false }
|
||||
hyper-util = { version = "0.1.20", default-features = false, features = ["client-proxy"] }
|
||||
proptest = "1.7.0"
|
||||
pyo3 = "0.29.2"
|
||||
|
|
@ -45,6 +57,8 @@ serde_with = { version = "=3.16.1", default-features = false, features = ["std",
|
|||
sha2 = "0.10"
|
||||
subtle = "2"
|
||||
thiserror = "2.0"
|
||||
tokenizers = { version = "0.23.1", default-features = false, features = ["onig"] }
|
||||
tiktoken-rs = "0.12.0"
|
||||
tokio = { version = "1", features = ["rt-multi-thread", "macros", "time", "net"] }
|
||||
tokio-tungstenite = { version = "0.24", default-features = false, features = ["connect", "rustls-tls-native-roots"] }
|
||||
futures-util = { version = "0.3", default-features = false, features = ["sink", "std"] }
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ license.workspace = true
|
|||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
litellm-auth.workspace = true
|
||||
litellm-auth-types.workspace = true
|
||||
litellm-http.workspace = true
|
||||
|
||||
moka = { workspace = true, features = ["sync"] }
|
||||
|
|
|
|||
|
|
@ -3,6 +3,8 @@ pub const AWS_SECRET_ACCESS_KEY: &str = "AWS_SECRET_ACCESS_KEY";
|
|||
pub const AWS_SESSION_TOKEN: &str = "AWS_SESSION_TOKEN";
|
||||
pub const AWS_REGION_NAME: &str = "AWS_REGION_NAME";
|
||||
pub const AWS_REGION: &str = "AWS_REGION";
|
||||
pub const AWS_DEFAULT_REGION: &str = "AWS_DEFAULT_REGION";
|
||||
pub const AWS_BEDROCK_RUNTIME_ENDPOINT: &str = "AWS_BEDROCK_RUNTIME_ENDPOINT";
|
||||
pub const AWS_SESSION_NAME: &str = "AWS_SESSION_NAME";
|
||||
pub const AWS_PROFILE_NAME: &str = "AWS_PROFILE_NAME";
|
||||
pub const AWS_ROLE_NAME: &str = "AWS_ROLE_NAME";
|
||||
|
|
|
|||
|
|
@ -22,7 +22,7 @@ pub enum Error {
|
|||
AwsMissingWebIdentityCredentials,
|
||||
}
|
||||
|
||||
impl From<Error> for litellm_auth::Error {
|
||||
impl From<Error> for litellm_auth_types::Error {
|
||||
fn from(error: Error) -> Self {
|
||||
Self::ProviderAuthentication(error.to_string())
|
||||
}
|
||||
|
|
@ -34,11 +34,11 @@ mod tests {
|
|||
|
||||
#[test]
|
||||
fn converts_to_shared_auth_error_without_losing_context() {
|
||||
let error = litellm_auth::Error::from(Error::AwsProfile("profile not found".into()));
|
||||
let error = litellm_auth_types::Error::from(Error::AwsProfile("profile not found".into()));
|
||||
|
||||
assert_eq!(
|
||||
error,
|
||||
litellm_auth::Error::ProviderAuthentication(
|
||||
litellm_auth_types::Error::ProviderAuthentication(
|
||||
"AWS profile credentials failed: profile not found".into()
|
||||
)
|
||||
);
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ license.workspace = true
|
|||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
litellm-auth.workspace = true
|
||||
litellm-auth-types.workspace = true
|
||||
|
||||
moka.workspace = true
|
||||
serde_json.workspace = true
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ use std::sync::Arc;
|
|||
use azure_core::credentials::TokenCredential;
|
||||
use moka::future::Cache;
|
||||
|
||||
use litellm_auth::Error;
|
||||
use litellm_auth_types::Error;
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
|
||||
pub(crate) struct AzureCredentialProviderCacheKey {
|
||||
|
|
|
|||
|
|
@ -12,8 +12,8 @@ use azure_identity::{
|
|||
};
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
use litellm_auth::Error;
|
||||
use litellm_auth::{InputSource, ResolvedCredential, SecretValue, Sourced};
|
||||
use litellm_auth_types::Error;
|
||||
use litellm_auth_types::{InputSource, ResolvedCredential, SecretValue, Sourced};
|
||||
|
||||
use super::credential_provider_cache::{
|
||||
AzureCredentialProviderCache, AzureCredentialProviderCacheKey,
|
||||
|
|
@ -484,7 +484,7 @@ mod tests {
|
|||
use azure_core::{Bytes, Result};
|
||||
|
||||
use super::{NativeAzureRequest, NativeAzureTokenAcquirer, ValidatedAzureRequest};
|
||||
use litellm_auth::{InputSource, SecretValue, Sourced};
|
||||
use litellm_auth_types::{InputSource, SecretValue, Sourced};
|
||||
|
||||
fn deployment<T>(value: T) -> Sourced<T> {
|
||||
Sourced::new(value, InputSource::Deployment)
|
||||
|
|
@ -649,7 +649,7 @@ mod tests {
|
|||
|
||||
assert!(matches!(
|
||||
error,
|
||||
litellm_auth::Error::MixedAzureCredentialSources
|
||||
litellm_auth_types::Error::MixedAzureCredentialSources
|
||||
));
|
||||
}
|
||||
|
||||
|
|
@ -679,7 +679,10 @@ mod tests {
|
|||
authority,
|
||||
))
|
||||
.unwrap_err();
|
||||
assert!(matches!(error, litellm_auth::Error::InvalidAzureAuthority));
|
||||
assert!(matches!(
|
||||
error,
|
||||
litellm_auth_types::Error::InvalidAzureAuthority
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
use litellm_auth::Error;
|
||||
use litellm_auth::{
|
||||
use litellm_auth_types::Error;
|
||||
use litellm_auth_types::{
|
||||
CredentialFileRef, CredentialLookup, CredentialRef, InputSource, ResolvedCredential,
|
||||
SecretValue, Sourced, TokenProviderHandle,
|
||||
};
|
||||
|
|
@ -451,9 +451,9 @@ mod tests {
|
|||
};
|
||||
use crate::native::ValidatedAzureRequest;
|
||||
use crate::types::AzureAuthInputs;
|
||||
use litellm_auth::Error;
|
||||
use litellm_auth::ResolvedCredential;
|
||||
use litellm_auth::{
|
||||
use litellm_auth_types::Error;
|
||||
use litellm_auth_types::ResolvedCredential;
|
||||
use litellm_auth_types::{
|
||||
CredentialFileRef, CredentialLookup, CredentialLookupFuture, CredentialRef,
|
||||
CredentialResolver, CredentialResolverHandle, InputSource, SecretValue, Sourced,
|
||||
};
|
||||
|
|
@ -661,8 +661,8 @@ mod tests {
|
|||
#[derive(Debug)]
|
||||
struct CallerToken(&'static str);
|
||||
|
||||
impl litellm_auth::TokenProvider for CallerToken {
|
||||
fn acquire(&self) -> litellm_auth::TokenFuture<'_> {
|
||||
impl litellm_auth_types::TokenProvider for CallerToken {
|
||||
fn acquire(&self) -> litellm_auth_types::TokenFuture<'_> {
|
||||
Box::pin(async move {
|
||||
Ok(ResolvedCredential::AccessToken {
|
||||
token: SecretValue::new(self.0),
|
||||
|
|
@ -675,7 +675,7 @@ mod tests {
|
|||
fn caller_inputs(token: &'static str) -> AzureAuthInputs {
|
||||
let params = json!({"azure_ad_token": "static-token"});
|
||||
AzureAuthInputs {
|
||||
azure_ad_token_provider: Some(litellm_auth::TokenProviderHandle::new(Arc::new(
|
||||
azure_ad_token_provider: Some(litellm_auth_types::TokenProviderHandle::new(Arc::new(
|
||||
CallerToken(token),
|
||||
))),
|
||||
..AzureAuthInputs::from_optional_params(params.as_object().unwrap()).unwrap()
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
use std::collections::BTreeMap;
|
||||
|
||||
use litellm_auth::{
|
||||
use litellm_auth_types::{
|
||||
CredentialResolverHandle, Error, InputSource, SecretValue, Sourced, TokenProviderHandle,
|
||||
};
|
||||
use serde_json::{Map, Value};
|
||||
|
|
@ -126,7 +126,7 @@ fn source_for(sources: &BTreeMap<String, InputSource>, name: &str) -> InputSourc
|
|||
mod tests {
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use litellm_auth::{InputSource, Sourced};
|
||||
use litellm_auth_types::{InputSource, Sourced};
|
||||
use serde_json::json;
|
||||
|
||||
use super::{AzureAuthInputs, AzureCredentialType, ConfigValue};
|
||||
|
|
|
|||
|
|
@ -5,8 +5,11 @@ edition.workspace = true
|
|||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[features]
|
||||
google-sdk = ["dep:google-cloud-auth", "dep:http"]
|
||||
|
||||
[dependencies]
|
||||
litellm-auth.workspace = true
|
||||
litellm-auth-types.workspace = true
|
||||
|
||||
moka.workspace = true
|
||||
serde_json.workspace = true
|
||||
|
|
@ -14,3 +17,5 @@ sha2.workspace = true
|
|||
tokio.workspace = true
|
||||
|
||||
gcp_auth = "0.12.7"
|
||||
google-cloud-auth = { workspace = true, optional = true }
|
||||
http = { workspace = true, optional = true }
|
||||
|
|
|
|||
|
|
@ -1,13 +1,18 @@
|
|||
use std::{collections::BTreeMap, future::Future, path::Path, pin::Pin, sync::Arc};
|
||||
|
||||
use gcp_auth::{CustomServiceAccount, TokenProvider};
|
||||
use litellm_auth::{
|
||||
use litellm_auth_types::{
|
||||
CredentialPlacement, Error, InputSource, SecretValue, Sourced, http::apply_credential,
|
||||
};
|
||||
use moka::future::Cache;
|
||||
use serde_json::{Map, Value};
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
#[cfg(feature = "google-sdk")]
|
||||
mod sdk;
|
||||
#[cfg(feature = "google-sdk")]
|
||||
pub use sdk::GoogleCredentials;
|
||||
|
||||
const CLOUD_PLATFORM_SCOPE: &str = "https://www.googleapis.com/auth/cloud-platform";
|
||||
const GOOGLE_OAUTH_TOKEN_ENDPOINT: &str = "https://oauth2.googleapis.com/token";
|
||||
const GOOGLE_APPLICATION_CREDENTIALS_ENV: &str = "GOOGLE_APPLICATION_CREDENTIALS";
|
||||
|
|
@ -26,19 +31,31 @@ pub struct VertexConfig {
|
|||
}
|
||||
|
||||
impl VertexConfig {
|
||||
pub fn new(
|
||||
credentials: Option<Sourced<SecretValue>>,
|
||||
project_id: Option<String>,
|
||||
location: Option<String>,
|
||||
) -> Self {
|
||||
Self {
|
||||
credentials: credentials.filter(|value| !value.value().expose().trim().is_empty()),
|
||||
project_id: project_id.filter(|value| !value.trim().is_empty()),
|
||||
location: location.filter(|value| !value.trim().is_empty()),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn from_sourced_optional_params(
|
||||
params: &Map<String, Value>,
|
||||
sources: &BTreeMap<String, InputSource>,
|
||||
) -> Result<Self, Error> {
|
||||
Ok(Self {
|
||||
credentials: optional_credentials(
|
||||
Ok(Self::new(
|
||||
optional_credentials(
|
||||
params,
|
||||
sources,
|
||||
&["vertex_credentials", "vertex_ai_credentials"],
|
||||
)?,
|
||||
project_id: optional_string(params, &["vertex_project", "vertex_ai_project"])?,
|
||||
location: optional_string(params, &["vertex_location", "vertex_ai_location"])?,
|
||||
})
|
||||
optional_string(params, &["vertex_project", "vertex_ai_project"])?,
|
||||
optional_string(params, &["vertex_location", "vertex_ai_location"])?,
|
||||
))
|
||||
}
|
||||
|
||||
pub fn or_configured(self, project_id: Option<&str>, location: Option<&str>) -> Self {
|
||||
|
|
@ -469,6 +486,39 @@ mod tests {
|
|||
assert_eq!(config.location(), Some("alias-location"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn typed_config_preserves_source_and_empty_value_fallback() {
|
||||
let configured = VertexConfig::new(
|
||||
Some(Sourced::new(
|
||||
SecretValue::new("inline-json"),
|
||||
InputSource::Request,
|
||||
)),
|
||||
Some("project".into()),
|
||||
Some("location".into()),
|
||||
);
|
||||
assert!(matches!(
|
||||
credential_source(&configured, &|_| Some("environment-json".into())),
|
||||
CredentialSource::Inline(value) if value.expose() == "inline-json"
|
||||
));
|
||||
let empty = VertexConfig::new(
|
||||
Some(Sourced::new(SecretValue::new(" "), InputSource::Request)),
|
||||
Some(" ".into()),
|
||||
Some(" ".into()),
|
||||
);
|
||||
assert!(matches!(
|
||||
credential_source(&empty, &|_| None),
|
||||
CredentialSource::Adc
|
||||
));
|
||||
assert_eq!(
|
||||
get_vertex_ai_project(&empty, &|_| Some("env-project".into())).as_deref(),
|
||||
Some("env-project")
|
||||
);
|
||||
assert_eq!(
|
||||
get_vertex_ai_location(&empty, &|_| Some("env-location".into())).as_deref(),
|
||||
Some("env-location")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn project_and_location_prefer_input_then_environment() {
|
||||
let configured =
|
||||
|
|
|
|||
106
litellm-rust/crates/auth-gcp/src/sdk.rs
Normal file
106
litellm-rust/crates/auth-gcp/src/sdk.rs
Normal file
|
|
@ -0,0 +1,106 @@
|
|||
use std::sync::Arc;
|
||||
|
||||
use google_cloud_auth::credentials::{CacheableResource, CredentialsProvider, EntityTag};
|
||||
use google_cloud_auth::errors::CredentialsError;
|
||||
use http::{Extensions, HeaderMap, HeaderName, HeaderValue};
|
||||
use litellm_auth_types::Error;
|
||||
|
||||
use crate::{VertexAuth, VertexConfig};
|
||||
|
||||
type EnvironmentLookup = dyn Fn(&str) -> Option<String> + Send + Sync;
|
||||
|
||||
pub struct GoogleCredentials {
|
||||
auth: VertexAuth,
|
||||
config: VertexConfig,
|
||||
environment: Arc<EnvironmentLookup>,
|
||||
}
|
||||
|
||||
impl GoogleCredentials {
|
||||
pub fn new(config: VertexConfig, environment: Arc<EnvironmentLookup>) -> Self {
|
||||
Self {
|
||||
auth: VertexAuth::default(),
|
||||
config,
|
||||
environment,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn request_headers(&self) -> Result<HeaderMap, Error> {
|
||||
let response = self
|
||||
.auth
|
||||
.validate_environment(Vec::new(), None, &self.config, &|name| {
|
||||
(self.environment)(name)
|
||||
})
|
||||
.await?;
|
||||
response
|
||||
.headers
|
||||
.into_iter()
|
||||
.map(|(key, value)| {
|
||||
let name =
|
||||
HeaderName::from_bytes(key.as_bytes()).map_err(|_| Error::InvalidHeader)?;
|
||||
let value = HeaderValue::from_str(&value).map_err(|_| Error::InvalidHeader)?;
|
||||
Ok((name, value))
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
impl CredentialsProvider for GoogleCredentials {
|
||||
async fn headers(
|
||||
&self,
|
||||
_: Extensions,
|
||||
) -> Result<CacheableResource<HeaderMap>, CredentialsError> {
|
||||
self.request_headers()
|
||||
.await
|
||||
.map(|data| CacheableResource::New {
|
||||
entity_tag: EntityTag::new(),
|
||||
data,
|
||||
})
|
||||
.map_err(|_| CredentialsError::from_msg(false, "Google authentication failed"))
|
||||
}
|
||||
|
||||
async fn universe_domain(&self) -> Option<String> {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for GoogleCredentials {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("GoogleCredentials").finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn sdk_and_http_credentials_share_token_resolution_and_redaction() {
|
||||
let credentials = GoogleCredentials::new(
|
||||
VertexConfig::new(None, Some("project".into()), None),
|
||||
Arc::new(|name| (name == "VERTEX_AI_API_KEY").then(|| "private-token".into())),
|
||||
);
|
||||
let direct = credentials.request_headers().await.unwrap();
|
||||
let CacheableResource::New { data, .. } =
|
||||
credentials.headers(Extensions::new()).await.unwrap()
|
||||
else {
|
||||
panic!("first request did not return headers");
|
||||
};
|
||||
assert_eq!(direct, data);
|
||||
assert_eq!(data[http::header::AUTHORIZATION], "Bearer private-token");
|
||||
assert!(!format!("{credentials:?}").contains("private-token"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn invalid_token_headers_return_a_redacted_sdk_error() {
|
||||
let credentials = GoogleCredentials::new(
|
||||
VertexConfig::new(None, Some("project".into()), None),
|
||||
Arc::new(|name| (name == "VERTEX_AI_API_KEY").then(|| "private\nvalue".into())),
|
||||
);
|
||||
assert_eq!(
|
||||
credentials.request_headers().await.unwrap_err(),
|
||||
Error::InvalidHeader
|
||||
);
|
||||
let error = credentials.headers(Extensions::new()).await.unwrap_err();
|
||||
assert!(!format!("{error:?}").contains("private"));
|
||||
}
|
||||
}
|
||||
15
litellm-rust/crates/auth-types/Cargo.toml
Normal file
15
litellm-rust/crates/auth-types/Cargo.toml
Normal file
|
|
@ -0,0 +1,15 @@
|
|||
[package]
|
||||
name = "litellm-auth-types"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
serde.workspace = true
|
||||
subtle.workspace = true
|
||||
thiserror.workspace = true
|
||||
veil.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
tokio.workspace = true
|
||||
|
|
@ -5,9 +5,7 @@ use std::sync::Arc;
|
|||
|
||||
use veil::Redact;
|
||||
|
||||
use crate::Error;
|
||||
|
||||
use super::{ResolvedCredential, SecretValue, TokenProviderHandle};
|
||||
use crate::{Error, ResolvedCredential, SecretValue, TokenProviderHandle};
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub enum CredentialFileRef {
|
||||
|
|
@ -40,9 +40,6 @@ pub fn apply_credential(
|
|||
)
|
||||
}
|
||||
|
||||
/// How the upstream call is authenticated. API-key strategies become headers
|
||||
/// in `prepare`; SigV4 covers the serialized body, so it is applied where the
|
||||
/// outbound request is built.
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub enum RequestAuth {
|
||||
Header {
|
||||
57
litellm-rust/crates/auth-types/src/lib.rs
Normal file
57
litellm-rust/crates/auth-types/src/lib.rs
Normal file
|
|
@ -0,0 +1,57 @@
|
|||
#![forbid(unsafe_code)]
|
||||
|
||||
mod credential;
|
||||
mod error;
|
||||
pub mod http;
|
||||
mod policy;
|
||||
mod secret;
|
||||
mod token;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, Hash, PartialEq, Serialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum InputSource {
|
||||
Request,
|
||||
#[default]
|
||||
Deployment,
|
||||
Environment,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default, Eq, PartialEq)]
|
||||
pub struct Sourced<T> {
|
||||
value: T,
|
||||
source: InputSource,
|
||||
}
|
||||
|
||||
impl<T> Sourced<T> {
|
||||
pub fn new(value: T, source: InputSource) -> Self {
|
||||
Self { value, source }
|
||||
}
|
||||
|
||||
pub fn value(&self) -> &T {
|
||||
&self.value
|
||||
}
|
||||
|
||||
pub fn source(&self) -> InputSource {
|
||||
self.source
|
||||
}
|
||||
|
||||
pub fn into_value(self) -> T {
|
||||
self.value
|
||||
}
|
||||
|
||||
pub fn map<U>(self, map: impl FnOnce(T) -> U) -> Sourced<U> {
|
||||
Sourced::new(map(self.value), self.source)
|
||||
}
|
||||
}
|
||||
|
||||
pub use credential::{
|
||||
CredentialFileRef, CredentialLookup, CredentialLookupFuture, CredentialPlan,
|
||||
CredentialPlanResolution, CredentialRef, CredentialResolver, CredentialResolverHandle,
|
||||
};
|
||||
pub use error::Error;
|
||||
pub use http::{CredentialPlacement, RequestAuth};
|
||||
pub use policy::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy};
|
||||
pub use secret::SecretValue;
|
||||
pub use token::{ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle};
|
||||
|
|
@ -1,7 +1,5 @@
|
|||
use crate::Error;
|
||||
|
||||
use super::http::apply_credential;
|
||||
use super::{CredentialPlacement, ResolvedCredential};
|
||||
use crate::http::apply_credential;
|
||||
use crate::{CredentialPlacement, Error, ResolvedCredential};
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum CredentialPlanKind {
|
||||
|
|
@ -5,9 +5,7 @@ use std::time::SystemTime;
|
|||
|
||||
use veil::Redact;
|
||||
|
||||
use crate::Error;
|
||||
|
||||
use super::secret::SecretValue;
|
||||
use crate::{Error, SecretValue};
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub enum ResolvedCredential {
|
||||
|
|
@ -5,11 +5,14 @@ edition.workspace = true
|
|||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
serde.workspace = true
|
||||
subtle.workspace = true
|
||||
thiserror.workspace = true
|
||||
veil.workspace = true
|
||||
[features]
|
||||
default = []
|
||||
aws = ["dep:litellm-auth-aws"]
|
||||
azure = ["dep:litellm-auth-azure"]
|
||||
gcp = ["dep:litellm-auth-gcp"]
|
||||
|
||||
[dev-dependencies]
|
||||
tokio.workspace = true
|
||||
[dependencies]
|
||||
litellm-auth-types.workspace = true
|
||||
litellm-auth-aws = { workspace = true, optional = true }
|
||||
litellm-auth-azure = { workspace = true, optional = true }
|
||||
litellm-auth-gcp = { workspace = true, optional = true }
|
||||
|
|
|
|||
|
|
@ -1,55 +1,10 @@
|
|||
mod credential;
|
||||
mod error;
|
||||
pub mod http;
|
||||
mod policy;
|
||||
mod secret;
|
||||
mod token;
|
||||
#![forbid(unsafe_code)]
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
pub use litellm_auth_types::*;
|
||||
|
||||
#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, Hash, PartialEq, Serialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum InputSource {
|
||||
Request,
|
||||
#[default]
|
||||
Deployment,
|
||||
Environment,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default, Eq, PartialEq)]
|
||||
pub struct Sourced<T> {
|
||||
value: T,
|
||||
source: InputSource,
|
||||
}
|
||||
|
||||
impl<T> Sourced<T> {
|
||||
pub fn new(value: T, source: InputSource) -> Self {
|
||||
Self { value, source }
|
||||
}
|
||||
|
||||
pub fn value(&self) -> &T {
|
||||
&self.value
|
||||
}
|
||||
|
||||
pub fn source(&self) -> InputSource {
|
||||
self.source
|
||||
}
|
||||
|
||||
pub fn into_value(self) -> T {
|
||||
self.value
|
||||
}
|
||||
|
||||
pub fn map<U>(self, map: impl FnOnce(T) -> U) -> Sourced<U> {
|
||||
Sourced::new(map(self.value), self.source)
|
||||
}
|
||||
}
|
||||
|
||||
pub use credential::{
|
||||
CredentialFileRef, CredentialLookup, CredentialLookupFuture, CredentialPlan,
|
||||
CredentialPlanResolution, CredentialRef, CredentialResolver, CredentialResolverHandle,
|
||||
};
|
||||
pub use error::Error;
|
||||
pub use http::{CredentialPlacement, RequestAuth};
|
||||
pub use policy::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy};
|
||||
pub use secret::SecretValue;
|
||||
pub use token::{ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle};
|
||||
#[cfg(feature = "aws")]
|
||||
pub use litellm_auth_aws as aws;
|
||||
#[cfg(feature = "azure")]
|
||||
pub use litellm_auth_azure as azure;
|
||||
#[cfg(feature = "gcp")]
|
||||
pub use litellm_auth_gcp as gcp;
|
||||
|
|
|
|||
33
litellm-rust/crates/auth/tests/facade.rs
Normal file
33
litellm-rust/crates/auth/tests/facade.rs
Normal file
|
|
@ -0,0 +1,33 @@
|
|||
use litellm_auth::{
|
||||
CredentialPlacement, CredentialPlanKind, CredentialRule, ExistingHeaderBehavior,
|
||||
ProviderAuthPolicy, ResolvedCredential, SecretValue,
|
||||
};
|
||||
|
||||
const RULES: &[CredentialRule] = &[CredentialRule {
|
||||
kind: CredentialPlanKind::Static,
|
||||
placement: CredentialPlacement::Header("x-api-key"),
|
||||
}];
|
||||
|
||||
#[test]
|
||||
fn facade_applies_shared_auth_policy() {
|
||||
let policy = ProviderAuthPolicy {
|
||||
rules: RULES,
|
||||
accepted_existing_headers: &["x-api-key"],
|
||||
existing_header_behavior: ExistingHeaderBehavior::Preserve,
|
||||
scope: None,
|
||||
audience: None,
|
||||
};
|
||||
|
||||
let headers = policy
|
||||
.apply(
|
||||
Vec::new(),
|
||||
CredentialPlanKind::Static,
|
||||
&ResolvedCredential::Static(SecretValue::new("secret")),
|
||||
)
|
||||
.expect("facade policy applies");
|
||||
|
||||
assert_eq!(
|
||||
headers,
|
||||
vec![("x-api-key".to_string(), "secret".to_string())]
|
||||
);
|
||||
}
|
||||
|
|
@ -7,8 +7,8 @@ repository.workspace = true
|
|||
|
||||
[dependencies]
|
||||
litellm-cache.workspace = true
|
||||
serde_json.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
serde_json.workspace = true
|
||||
rstest.workspace = true
|
||||
tokio.workspace = true
|
||||
|
|
|
|||
|
|
@ -1,18 +1,20 @@
|
|||
use std::cmp::Reverse;
|
||||
use std::collections::{BinaryHeap, HashMap};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||
use std::{
|
||||
cmp::Reverse,
|
||||
collections::{BinaryHeap, HashMap, HashSet},
|
||||
hash::Hash,
|
||||
sync::{Arc, Mutex},
|
||||
time::{Duration, SystemTime, UNIX_EPOCH},
|
||||
};
|
||||
|
||||
use litellm_cache::{
|
||||
BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheEntry, CacheFuture, CacheKwargs,
|
||||
Error,
|
||||
BaseCache, BatchCache, CacheConnectionResult, CacheConnectionStatus, ClaimCache, CounterCache,
|
||||
DeleteCache, Error, ExactCacheContext, FlushCache, IncrementOperation, SetCache, TtlCache,
|
||||
};
|
||||
|
||||
const DEFAULT_MAX_SIZE_IN_MEMORY: usize = 200;
|
||||
const DEFAULT_TTL: Duration = Duration::from_secs(600);
|
||||
|
||||
type ValueMeasure<V> = Arc<dyn Fn(&V) -> Result<usize, Error> + Send + Sync>;
|
||||
type ValueValidator<V> = Arc<dyn Fn(&V) -> Result<(), Error> + Send + Sync>;
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum CacheWrite {
|
||||
|
|
@ -33,7 +35,6 @@ pub struct InMemoryCache<V: Clone> {
|
|||
default_ttl: Duration,
|
||||
max_entry_bytes: Option<usize>,
|
||||
measure_value: Option<ValueMeasure<V>>,
|
||||
validate_value: Option<ValueValidator<V>>,
|
||||
now: Arc<dyn Fn() -> Duration + Send + Sync>,
|
||||
}
|
||||
|
||||
|
|
@ -77,7 +78,6 @@ impl<V: Clone> InMemoryCache<V> {
|
|||
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
|
||||
max_entry_bytes,
|
||||
measure_value,
|
||||
validate_value: None,
|
||||
now: Arc::new(now),
|
||||
}
|
||||
}
|
||||
|
|
@ -91,9 +91,6 @@ impl<V: Clone> InMemoryCache<V> {
|
|||
if self.max_size_in_memory == 0 {
|
||||
return Ok(CacheWrite::Disabled);
|
||||
}
|
||||
if let Some(validate) = &self.validate_value {
|
||||
validate(&value)?;
|
||||
}
|
||||
if let (Some(limit), Some(measure)) = (self.max_entry_bytes, &self.measure_value)
|
||||
&& measure(&value)? > limit
|
||||
{
|
||||
|
|
@ -101,15 +98,13 @@ impl<V: Clone> InMemoryCache<V> {
|
|||
}
|
||||
let now = (self.now)();
|
||||
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
|
||||
Self::evict(&mut state, self.max_size_in_memory, now);
|
||||
let key = key.into();
|
||||
state.values.insert(key.clone(), value);
|
||||
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) {
|
||||
let expiration = now + ttl.unwrap_or(self.default_ttl);
|
||||
state.expirations.insert(key.clone(), expiration);
|
||||
state.expiration_heap.push(Reverse((expiration, key)));
|
||||
Self::set_expiration(&mut state, &key, now + ttl.unwrap_or(self.default_ttl));
|
||||
}
|
||||
state.values.insert(key, value);
|
||||
Ok(CacheWrite::Stored)
|
||||
}
|
||||
|
||||
|
|
@ -126,6 +121,14 @@ impl<V: Clone> InMemoryCache<V> {
|
|||
Ok(state.values.get(key).cloned())
|
||||
}
|
||||
|
||||
pub fn max_size_in_memory(&self) -> usize {
|
||||
self.max_size_in_memory
|
||||
}
|
||||
|
||||
pub fn max_entry_bytes(&self) -> Option<usize> {
|
||||
self.max_entry_bytes
|
||||
}
|
||||
|
||||
pub fn expires_at(&self, key: &str) -> Result<Option<Duration>, Error> {
|
||||
Ok(self
|
||||
.state
|
||||
|
|
@ -136,6 +139,25 @@ impl<V: Clone> InMemoryCache<V> {
|
|||
.copied())
|
||||
}
|
||||
|
||||
pub async fn async_get_ttl(&self, key: &str) -> Result<Option<Duration>, Error> {
|
||||
self.expires_at(key)
|
||||
}
|
||||
|
||||
pub async fn async_get_oldest_n_keys(&self, count: usize) -> Result<Vec<String>, Error> {
|
||||
let state = self.state.lock().map_err(|_| Error::Unavailable)?;
|
||||
let mut expirations = state
|
||||
.expirations
|
||||
.iter()
|
||||
.map(|(key, expiration)| (key.clone(), *expiration))
|
||||
.collect::<Vec<_>>();
|
||||
expirations.sort_unstable_by_key(|(_, expiration)| *expiration);
|
||||
Ok(expirations
|
||||
.into_iter()
|
||||
.take(count)
|
||||
.map(|(key, _)| key)
|
||||
.collect())
|
||||
}
|
||||
|
||||
pub fn delete_cache(&self, key: &str) -> Result<(), Error> {
|
||||
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
|
||||
Self::remove(&mut state, key);
|
||||
|
|
@ -150,7 +172,7 @@ impl<V: Clone> InMemoryCache<V> {
|
|||
Ok(())
|
||||
}
|
||||
|
||||
fn evict(state: &mut CacheState<V>, capacity: usize, now: Duration) {
|
||||
fn evict(state: &mut CacheState<V>, capacity: usize, now: Duration, key: &str) {
|
||||
while let Some(Reverse((expiration, key))) = state.expiration_heap.peek().cloned() {
|
||||
if state.expirations.get(&key).copied() != Some(expiration) {
|
||||
state.expiration_heap.pop();
|
||||
|
|
@ -161,6 +183,9 @@ impl<V: Clone> InMemoryCache<V> {
|
|||
break;
|
||||
}
|
||||
}
|
||||
if state.values.contains_key(key) {
|
||||
return;
|
||||
}
|
||||
while state.values.len() >= capacity {
|
||||
let Some(Reverse((expiration, key))) = state.expiration_heap.pop() else {
|
||||
break;
|
||||
|
|
@ -171,84 +196,205 @@ impl<V: Clone> InMemoryCache<V> {
|
|||
}
|
||||
}
|
||||
|
||||
fn set_expiration(state: &mut CacheState<V>, key: &str, expiration: Duration) {
|
||||
if state.expirations.get(key).copied() != Some(expiration) {
|
||||
state.expirations.insert(key.into(), expiration);
|
||||
state
|
||||
.expiration_heap
|
||||
.push(Reverse((expiration, key.into())));
|
||||
}
|
||||
}
|
||||
|
||||
fn remove(state: &mut CacheState<V>, key: &str) {
|
||||
state.values.remove(key);
|
||||
state.expirations.remove(key);
|
||||
}
|
||||
}
|
||||
|
||||
impl InMemoryCache<CacheEntry> {
|
||||
pub fn response_cache(capacity: usize, ttl: Duration, max_entry_bytes: usize) -> Self {
|
||||
Self::response_cache_with_clock(capacity, ttl, max_entry_bytes, || {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
})
|
||||
}
|
||||
|
||||
pub fn response_cache_with_clock(
|
||||
capacity: usize,
|
||||
ttl: Duration,
|
||||
max_entry_bytes: usize,
|
||||
now: impl Fn() -> Duration + Send + Sync + 'static,
|
||||
) -> Self {
|
||||
let mut cache = Self::with_clock_and_size_measurement(
|
||||
Some(capacity),
|
||||
Some(ttl),
|
||||
Some(max_entry_bytes),
|
||||
Some(Arc::new(|entry: &CacheEntry| {
|
||||
serde_json::to_vec(entry)
|
||||
.map(|bytes| bytes.len())
|
||||
.map_err(|_| Error::InvalidEntry)
|
||||
})),
|
||||
now,
|
||||
impl<V> ClaimCache for InMemoryCache<V>
|
||||
where
|
||||
V: Clone + PartialEq + Send + Sync + 'static,
|
||||
{
|
||||
fn claim_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
candidate: V,
|
||||
eligible: &[V],
|
||||
context: ExactCacheContext,
|
||||
) -> Result<V, Error> {
|
||||
if self.max_size_in_memory == 0 {
|
||||
return Ok(candidate);
|
||||
}
|
||||
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 existing = state
|
||||
.values
|
||||
.get(key)
|
||||
.filter(|existing| eligible.is_empty() || eligible.contains(existing))
|
||||
.cloned();
|
||||
if let Some(existing) = &existing
|
||||
&& eligible.is_empty()
|
||||
&& *existing != candidate
|
||||
{
|
||||
return Ok(existing.clone());
|
||||
}
|
||||
let winner = existing.unwrap_or(candidate);
|
||||
Self::set_expiration(
|
||||
&mut state,
|
||||
key,
|
||||
now + self.get_ttl(&context).unwrap_or(self.default_ttl),
|
||||
);
|
||||
cache.validate_value = Some(Arc::new(|entry: &CacheEntry| {
|
||||
entry
|
||||
.timestamp
|
||||
.is_finite()
|
||||
.then_some(())
|
||||
.ok_or(Error::InvalidEntry)
|
||||
}));
|
||||
cache
|
||||
state.values.insert(key.into(), winner.clone());
|
||||
Ok(winner)
|
||||
}
|
||||
}
|
||||
|
||||
impl BaseCache for InMemoryCache<CacheEntry> {
|
||||
type Value = CacheEntry;
|
||||
impl CounterCache for InMemoryCache<f64> {
|
||||
fn increment_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
amount: f64,
|
||||
context: ExactCacheContext,
|
||||
) -> Result<f64, Error> {
|
||||
if self.max_size_in_memory == 0 {
|
||||
return Ok(amount);
|
||||
}
|
||||
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);
|
||||
Ok(value)
|
||||
}
|
||||
}
|
||||
|
||||
fn default_ttl(&self) -> Duration {
|
||||
self.default_ttl
|
||||
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;
|
||||
|
||||
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
|
||||
context.ttl.or(Some(self.default_ttl))
|
||||
}
|
||||
|
||||
fn set_cache(&self, key: &str, value: Self::Value, kwargs: CacheKwargs) -> Result<(), Error> {
|
||||
let ttl = self.get_ttl(&kwargs);
|
||||
fn set_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
value: Self::Value,
|
||||
context: &ExactCacheContext,
|
||||
) -> Result<(), Error> {
|
||||
let ttl = self.get_ttl(context).unwrap_or(self.default_ttl);
|
||||
self.set_cache(key, value, Some(ttl)).map(|_| ())
|
||||
}
|
||||
|
||||
fn get_cache(&self, key: &str, _: &CacheKwargs) -> Result<Option<Self::Value>, Error> {
|
||||
fn get_cache(&self, key: &str, _: &ExactCacheContext) -> Result<Option<Self::Value>, Error> {
|
||||
self.get_cache(key)
|
||||
}
|
||||
|
||||
fn delete_cache(&self, key: &str) -> Result<(), Error> {
|
||||
self.delete_cache(key)
|
||||
async fn disconnect(&self) -> Result<(), Error> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn flush_cache(&self) -> Result<(), Error> {
|
||||
self.flush_cache()
|
||||
}
|
||||
|
||||
fn disconnect(&self) -> CacheFuture<'_, ()> {
|
||||
Box::pin(async { Ok(()) })
|
||||
}
|
||||
|
||||
fn test_connection(&self) -> CacheFuture<'_, CacheConnectionResult> {
|
||||
Box::pin(async {
|
||||
Ok(CacheConnectionResult {
|
||||
status: CacheConnectionStatus::Success,
|
||||
message: "In-memory cache connection test successful".into(),
|
||||
error: None,
|
||||
})
|
||||
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> {}
|
||||
|
||||
impl<V: Clone + Send + Sync + 'static> DeleteCache for InMemoryCache<V> {
|
||||
fn delete_cache(&self, key: &str) -> Result<(), Error> {
|
||||
InMemoryCache::delete_cache(self, key)
|
||||
}
|
||||
}
|
||||
|
||||
impl<V: Clone + Send + Sync + 'static> FlushCache for InMemoryCache<V> {
|
||||
fn flush_cache(&self) -> Result<(), Error> {
|
||||
InMemoryCache::flush_cache(self)
|
||||
}
|
||||
}
|
||||
|
||||
impl<V: Clone + Send + Sync + 'static> TtlCache for InMemoryCache<V> {
|
||||
async fn async_get_ttl(&self, key: &str) -> Result<Option<Duration>, Error> {
|
||||
InMemoryCache::async_get_ttl(self, key).await
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> SetCache for InMemoryCache<HashSet<T>>
|
||||
where
|
||||
T: Clone + Eq + Hash + Send + Sync + 'static,
|
||||
{
|
||||
type SetValue = T;
|
||||
type SetResult = Vec<T>;
|
||||
|
||||
async fn async_set_cache_sadd(
|
||||
&self,
|
||||
key: &str,
|
||||
values: Vec<Self::SetValue>,
|
||||
ttl: Option<Duration>,
|
||||
) -> Result<Self::SetResult, Error> {
|
||||
if self.max_size_in_memory == 0 {
|
||||
return Ok(values);
|
||||
}
|
||||
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();
|
||||
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);
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,8 +1,16 @@
|
|||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::time::Duration;
|
||||
use std::{
|
||||
collections::HashSet,
|
||||
sync::{
|
||||
Arc,
|
||||
atomic::{AtomicU64, Ordering},
|
||||
},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use litellm_cache::{BaseCache, CacheConnectionStatus, CacheEntry, Error};
|
||||
use litellm_cache::{
|
||||
BaseCache, CacheBackend, CacheConnectionStatus, ClaimCache, CounterCache, DeleteCache, Error,
|
||||
ExactCacheContext, IncrementOperation, SetCache, get_cache, set_cache,
|
||||
};
|
||||
use litellm_cache_memory::{CacheWrite, InMemoryCache};
|
||||
use rstest::{fixture, rstest};
|
||||
|
||||
|
|
@ -84,66 +92,49 @@ fn capacity_evicts_earliest_and_ignores_stale_heap_entries(clock: Arc<AtomicU64>
|
|||
}
|
||||
|
||||
#[test]
|
||||
fn disabled_size_limited_and_synchronized_response_writes_are_observable() {
|
||||
let disabled = InMemoryCache::<CacheEntry>::response_cache(0, Duration::from_secs(60), 80);
|
||||
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);
|
||||
assert_eq!(
|
||||
disabled
|
||||
.set_cache(
|
||||
"a",
|
||||
CacheEntry {
|
||||
timestamp: 1.0,
|
||||
response: serde_json::json!("x")
|
||||
},
|
||||
None
|
||||
)
|
||||
.unwrap(),
|
||||
disabled.set_cache("a", "x".into(), None).unwrap(),
|
||||
CacheWrite::Disabled
|
||||
);
|
||||
let cache = InMemoryCache::<CacheEntry>::response_cache(2, Duration::from_secs(60), 80);
|
||||
let cache = cache(2);
|
||||
assert_eq!(
|
||||
cache
|
||||
.set_cache(
|
||||
"large",
|
||||
CacheEntry {
|
||||
timestamp: 1.0,
|
||||
response: serde_json::json!("x".repeat(100))
|
||||
},
|
||||
None
|
||||
)
|
||||
.unwrap(),
|
||||
cache.set_cache("large", "oversized".into(), None).unwrap(),
|
||||
CacheWrite::TooLarge
|
||||
);
|
||||
cache
|
||||
.set_cache(
|
||||
"small",
|
||||
CacheEntry {
|
||||
timestamp: 1.0,
|
||||
response: serde_json::json!("ok"),
|
||||
},
|
||||
None,
|
||||
)
|
||||
.unwrap();
|
||||
assert!(cache.get_cache("small").unwrap().is_some());
|
||||
assert_eq!(cache.get_cache("large").unwrap(), None);
|
||||
assert_eq!(
|
||||
cache
|
||||
.set_cache(
|
||||
"invalid",
|
||||
CacheEntry {
|
||||
timestamp: f64::NAN,
|
||||
response: serde_json::json!("bad"),
|
||||
},
|
||||
None,
|
||||
)
|
||||
.unwrap_err(),
|
||||
Error::InvalidEntry
|
||||
cache.set_cache("small", "ok".into(), None).unwrap(),
|
||||
CacheWrite::Stored
|
||||
);
|
||||
assert_eq!(cache.get_cache("small").unwrap(), Some("ok".into()));
|
||||
assert_eq!(
|
||||
cache.set_cache("invalid", String::new(), None),
|
||||
Err(Error::InvalidEntry)
|
||||
);
|
||||
assert_eq!(cache.get_cache("invalid").unwrap(), None);
|
||||
cache.delete_cache("small").unwrap();
|
||||
cache.flush_cache().unwrap();
|
||||
assert_eq!(cache.get_cache("small").unwrap(), None);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn connection_test_matches_python_result_contract() {
|
||||
let cache = InMemoryCache::<CacheEntry>::default();
|
||||
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");
|
||||
|
|
@ -156,3 +147,222 @@ async fn connection_test_matches_python_result_contract() {
|
|||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[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));
|
||||
let reader = Arc::clone(&cache);
|
||||
let context = ExactCacheContext {
|
||||
ttl: Some(Duration::from_secs(5)),
|
||||
};
|
||||
set_cache(cache.as_ref(), "sync", "first".into(), &context).unwrap();
|
||||
assert_eq!(
|
||||
get_cache(reader.as_ref(), "sync", &context).unwrap(),
|
||||
Some("first".into())
|
||||
);
|
||||
cache
|
||||
.batch_cache_write("async", "second".into(), context.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
cache
|
||||
.async_set_cache_pipeline(vec![("batch".into(), "third".into())], context.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
drop(cache);
|
||||
for (key, value) in [("sync", "first"), ("async", "second"), ("batch", "third")] {
|
||||
assert_eq!(
|
||||
reader.async_get_cache(key, &context).await.unwrap(),
|
||||
Some(value.into())
|
||||
);
|
||||
}
|
||||
reader.async_delete_cache("async").await.unwrap();
|
||||
assert_eq!(
|
||||
reader.async_get_cache("async", &context).await.unwrap(),
|
||||
None
|
||||
);
|
||||
clock.store(106, Ordering::SeqCst);
|
||||
assert_eq!(get_cache(reader.as_ref(), "sync", &context).unwrap(), None);
|
||||
assert_eq!(
|
||||
reader.async_get_cache("batch", &context).await.unwrap(),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[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)),
|
||||
};
|
||||
assert_eq!(
|
||||
cache
|
||||
.claim_cache("affinity", "first".to_string(), &[], context.clone())
|
||||
.unwrap(),
|
||||
"first"
|
||||
);
|
||||
clock.store(103, Ordering::SeqCst);
|
||||
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
|
||||
.claim_cache(
|
||||
"affinity",
|
||||
"second".to_string(),
|
||||
&["first".to_string(), "second".to_string()],
|
||||
context,
|
||||
)
|
||||
.unwrap(),
|
||||
"first"
|
||||
);
|
||||
assert_eq!(
|
||||
cache.expires_at("affinity").unwrap(),
|
||||
Some(Duration::from_secs(115))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn counters_increment_under_one_lock() {
|
||||
let cache = InMemoryCache::<f64>::default();
|
||||
assert_eq!(
|
||||
CounterCache::increment_cache(&cache, "counter", 1.5, ExactCacheContext::default())
|
||||
.unwrap(),
|
||||
1.5
|
||||
);
|
||||
assert_eq!(
|
||||
CounterCache::increment_cache(&cache, "counter", 2.0, ExactCacheContext::default())
|
||||
.unwrap(),
|
||||
3.5
|
||||
);
|
||||
}
|
||||
|
||||
#[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();
|
||||
|
||||
cache.set_cache("cold", "3".into(), None).unwrap();
|
||||
assert_eq!(cache.get_cache("hot").unwrap(), Some("1".into()));
|
||||
assert_eq!(cache.get_cache("cold").unwrap(), Some("3".into()));
|
||||
|
||||
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()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn incrementing_an_existing_counter_at_capacity_keeps_every_counter() {
|
||||
let cache = InMemoryCache::<f64>::new(Some(2), None);
|
||||
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("b").unwrap(), Some(2.0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn disabled_cache_does_not_retain_claims_or_counters() {
|
||||
let claims = InMemoryCache::<String>::new(Some(0), None);
|
||||
assert_eq!(
|
||||
claims
|
||||
.claim_cache("key", "first".into(), &[], ExactCacheContext::default())
|
||||
.unwrap(),
|
||||
"first"
|
||||
);
|
||||
assert_eq!(claims.get_cache("key").unwrap(), None);
|
||||
|
||||
let counters = InMemoryCache::<f64>::new(Some(0), None);
|
||||
assert_eq!(
|
||||
counters
|
||||
.increment_cache("key", 2.0, ExactCacheContext::default())
|
||||
.unwrap(),
|
||||
2.0
|
||||
);
|
||||
assert_eq!(counters.get_cache("key").unwrap(), None);
|
||||
}
|
||||
|
||||
#[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();
|
||||
|
||||
assert_eq!(
|
||||
cache.async_get_ttl("first").await.unwrap(),
|
||||
Some(Duration::from_secs(110))
|
||||
);
|
||||
assert_eq!(cache.async_get_oldest_n_keys(1).await.unwrap(), ["first"]);
|
||||
assert_eq!(cache.async_get_ttl("missing").await.unwrap(), None);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn increment_pipeline_preserves_operation_order() {
|
||||
let cache = InMemoryCache::<f64>::new(Some(3), None);
|
||||
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)),
|
||||
},
|
||||
])
|
||||
.await
|
||||
.unwrap(),
|
||||
[1.0, 3.0]
|
||||
);
|
||||
assert_eq!(cache.get_cache("a").unwrap(), Some(3.0));
|
||||
}
|
||||
|
||||
#[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()];
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_set_cache_sadd("members", inserted.clone(), None)
|
||||
.await
|
||||
.unwrap(),
|
||||
inserted
|
||||
);
|
||||
assert_eq!(
|
||||
cache.get_cache("members").unwrap(),
|
||||
Some(HashSet::from(["a".into(), "b".into()]))
|
||||
);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -7,9 +7,10 @@ repository.workspace = true
|
|||
|
||||
[dependencies]
|
||||
litellm-cache.workspace = true
|
||||
redis = "1.7.0"
|
||||
serde_json.workspace = true
|
||||
redis = { version = "1.7.0", features = ["tls-rustls"] }
|
||||
r2d2 = "0.8.10"
|
||||
tokio.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
redis-test = "1.0.4"
|
||||
serde_json.workspace = true
|
||||
|
|
|
|||
|
|
@ -1,58 +1,243 @@
|
|||
use std::sync::{Arc, Mutex, MutexGuard};
|
||||
use std::time::Duration;
|
||||
use std::{
|
||||
sync::{Arc, Mutex},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use litellm_cache::{
|
||||
BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheEntry, CacheFuture, CacheKwargs,
|
||||
Error,
|
||||
BaseCache, BatchCache, BatchEntry, CacheCodec, CacheConnectionResult, CacheConnectionStatus,
|
||||
ClaimCache, CounterCache, DeleteCache, Error, ExactCacheContext, FlushCache,
|
||||
};
|
||||
use redis::Commands;
|
||||
|
||||
const DEFAULT_TTL: Duration = Duration::from_secs(600);
|
||||
const KEY_PREFIX: &str = "litellm-cache:";
|
||||
mod operations;
|
||||
|
||||
pub struct RedisCache<C = redis::Connection> {
|
||||
connection: Arc<Mutex<C>>,
|
||||
default_ttl: Duration,
|
||||
pub use operations::{
|
||||
RedisArg, RedisLpopOperation, RedisLpopResult, RedisRpushOperation, RedisScript,
|
||||
};
|
||||
|
||||
const DEFAULT_TTL: Duration = Duration::from_secs(600);
|
||||
const REDIS_TIMEOUT: Duration = Duration::from_secs(5);
|
||||
const REDIS_POOL_SIZE: u32 = 16;
|
||||
|
||||
struct PooledConnection {
|
||||
connection: redis::Connection,
|
||||
failed: bool,
|
||||
}
|
||||
|
||||
impl RedisCache<redis::Connection> {
|
||||
pub fn new(url: &str, default_ttl: Option<Duration>) -> Result<Self, Error> {
|
||||
let client = redis::Client::open(url).map_err(|_| Error::Unavailable)?;
|
||||
let connection = client.get_connection().map_err(|_| Error::Unavailable)?;
|
||||
Ok(Self::with_connection(connection, default_ttl))
|
||||
/// Pools connections without a checkout PING, which would double every operation's round trips.
|
||||
/// A timed-out command leaves its reply on the socket while redis still reports the connection
|
||||
/// open, so any connection whose operation failed is discarded instead of being reused.
|
||||
struct ConnectionManager(redis::Client);
|
||||
|
||||
impl r2d2::ManageConnection for ConnectionManager {
|
||||
type Connection = PooledConnection;
|
||||
type Error = redis::RedisError;
|
||||
|
||||
fn connect(&self) -> Result<PooledConnection, redis::RedisError> {
|
||||
let connection = self.0.get_connection()?;
|
||||
connection.set_read_timeout(Some(REDIS_TIMEOUT))?;
|
||||
connection.set_write_timeout(Some(REDIS_TIMEOUT))?;
|
||||
Ok(PooledConnection {
|
||||
connection,
|
||||
failed: false,
|
||||
})
|
||||
}
|
||||
|
||||
fn is_valid(&self, connection: &mut PooledConnection) -> Result<(), redis::RedisError> {
|
||||
redis::cmd("PING").query::<String>(&mut connection.connection)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn has_broken(&self, connection: &mut PooledConnection) -> bool {
|
||||
connection.failed || !redis::ConnectionLike::is_open(&connection.connection)
|
||||
}
|
||||
}
|
||||
|
||||
impl<C> RedisCache<C>
|
||||
const INCREMENT_SCRIPT: &str = concat!(
|
||||
"local value = redis.call('INCRBYFLOAT', KEYS[1], ARGV[1]); ",
|
||||
"if redis.call('TTL', KEYS[1]) == -1 then ",
|
||||
"redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return value"
|
||||
);
|
||||
|
||||
const CLAIM_SCRIPT: &str = concat!(
|
||||
"local current = redis.call('GET', KEYS[1]); ",
|
||||
"if ARGV[1] == '' then if current ~= false and current ~= '' then return 0; end; ",
|
||||
"elseif current ~= ARGV[1] then return 0; end; ",
|
||||
"if ARGV[3] ~= '' then redis.call('SET', KEYS[1], ARGV[3], 'EX', ARGV[2]); ",
|
||||
"elseif ARGV[4] == '1' then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return 1"
|
||||
);
|
||||
const CLAIM_ATTEMPTS: usize = 8;
|
||||
|
||||
enum Connections<C> {
|
||||
Pool(r2d2::Pool<ConnectionManager>),
|
||||
Fixed(Mutex<C>),
|
||||
}
|
||||
|
||||
struct ConnectionRef<'a>(&'a mut dyn redis::ConnectionLike);
|
||||
|
||||
impl redis::ConnectionLike for ConnectionRef<'_> {
|
||||
fn req_packed_command(&mut self, cmd: &[u8]) -> redis::RedisResult<redis::Value> {
|
||||
self.0.req_packed_command(cmd)
|
||||
}
|
||||
|
||||
fn req_packed_commands(
|
||||
&mut self,
|
||||
cmd: &[u8],
|
||||
offset: usize,
|
||||
count: usize,
|
||||
) -> redis::RedisResult<Vec<redis::Value>> {
|
||||
self.0.req_packed_commands(cmd, offset, count)
|
||||
}
|
||||
|
||||
fn get_db(&self) -> i64 {
|
||||
self.0.get_db()
|
||||
}
|
||||
|
||||
fn supports_pipelining(&self) -> bool {
|
||||
self.0.supports_pipelining()
|
||||
}
|
||||
|
||||
fn check_connection(&mut self) -> bool {
|
||||
self.0.check_connection()
|
||||
}
|
||||
|
||||
fn is_open(&self) -> bool {
|
||||
self.0.is_open()
|
||||
}
|
||||
}
|
||||
|
||||
impl<C> Connections<C>
|
||||
where
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
fn with_connection(connection: C, default_ttl: Option<Duration>) -> Self {
|
||||
Self {
|
||||
connection: Arc::new(Mutex::new(connection)),
|
||||
fn execute<T>(
|
||||
&self,
|
||||
operation: impl FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error>,
|
||||
) -> Result<T, Error> {
|
||||
match self {
|
||||
Self::Pool(pool) => {
|
||||
let mut pooled = pool.get().map_err(|_| Error::Unavailable)?;
|
||||
let result = operation(&mut ConnectionRef(&mut pooled.connection));
|
||||
pooled.failed = matches!(result, Err(Error::Unavailable));
|
||||
result
|
||||
}
|
||||
Self::Fixed(connection) => {
|
||||
let mut connection = connection.lock().map_err(|_| Error::Unavailable)?;
|
||||
operation(&mut ConnectionRef(&mut *connection))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct RedisCache<S, C = redis::Connection> {
|
||||
connections: Arc<Connections<C>>,
|
||||
default_ttl: Duration,
|
||||
codec: S,
|
||||
namespace: Option<String>,
|
||||
}
|
||||
|
||||
impl<S: CacheCodec> RedisCache<S> {
|
||||
pub fn new(url: &str, default_ttl: Option<Duration>, codec: S) -> Result<Self, Error> {
|
||||
let client = redis::Client::open(url).map_err(|_| Error::Unavailable)?;
|
||||
let pool = r2d2::Pool::builder()
|
||||
.max_size(REDIS_POOL_SIZE)
|
||||
.min_idle(Some(0))
|
||||
.connection_timeout(REDIS_TIMEOUT)
|
||||
.test_on_check_out(false)
|
||||
.build(ConnectionManager(client))
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
Ok(Self {
|
||||
connections: Arc::new(Connections::Pool(pool)),
|
||||
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
|
||||
codec,
|
||||
namespace: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, C> RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
pub fn with_connection(connection: C, default_ttl: Option<Duration>, codec: S) -> Self {
|
||||
Self {
|
||||
connections: Arc::new(Connections::Fixed(Mutex::new(connection))),
|
||||
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
|
||||
codec,
|
||||
namespace: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn connection(&self) -> Result<MutexGuard<'_, C>, Error> {
|
||||
self.connection.lock().map_err(|_| Error::Unavailable)
|
||||
pub fn with_namespace(self, namespace: Option<String>) -> Self {
|
||||
Self {
|
||||
namespace: namespace.filter(|value| !value.is_empty()),
|
||||
..self
|
||||
}
|
||||
}
|
||||
|
||||
fn namespaced_key(key: &str) -> String {
|
||||
format!("{KEY_PREFIX}{key}")
|
||||
pub fn namespace(&self) -> Option<&str> {
|
||||
self.namespace.as_deref()
|
||||
}
|
||||
|
||||
fn namespaced_pattern() -> &'static str {
|
||||
const PATTERN: &str = "litellm-cache:*";
|
||||
PATTERN
|
||||
fn namespaced_key(&self, key: &str) -> String {
|
||||
namespaced_key(self.namespace.as_deref(), key)
|
||||
}
|
||||
|
||||
fn encode(value: &CacheEntry) -> Result<Vec<u8>, Error> {
|
||||
serde_json::to_vec(value).map_err(|_| Error::InvalidEntry)
|
||||
fn namespaced_pattern(&self) -> Result<String, Error> {
|
||||
let namespace = self.namespace.as_ref().ok_or(Error::UnscopedFlush)?;
|
||||
let escaped: String = namespace
|
||||
.chars()
|
||||
.flat_map(|ch| {
|
||||
if matches!(ch, '*' | '?' | '[' | ']' | '\\') {
|
||||
vec!['\\', ch]
|
||||
} else {
|
||||
vec![ch]
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
Ok(format!("{escaped}:*"))
|
||||
}
|
||||
|
||||
fn decode(value: Vec<u8>) -> Result<CacheEntry, Error> {
|
||||
serde_json::from_slice(&value).map_err(|_| Error::InvalidEntry)
|
||||
fn flush_matching(connection: &mut ConnectionRef<'_>, pattern: &str) -> Result<(), Error> {
|
||||
let mut cursor = 0u64;
|
||||
loop {
|
||||
let (next_cursor, keys): (u64, Vec<String>) = redis::cmd("SCAN")
|
||||
.cursor_arg(cursor)
|
||||
.arg("MATCH")
|
||||
.arg(pattern)
|
||||
.arg("COUNT")
|
||||
.arg(1000)
|
||||
.query(connection)
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
if !keys.is_empty() {
|
||||
connection
|
||||
.del::<_, usize>(keys)
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
}
|
||||
if next_cursor == 0 {
|
||||
return Ok(());
|
||||
}
|
||||
cursor = next_cursor;
|
||||
}
|
||||
}
|
||||
|
||||
fn decode_response(&self, value: redis::Value) -> Result<Option<S::Value>, Error> {
|
||||
match value {
|
||||
redis::Value::Nil => Ok(None),
|
||||
redis::Value::BulkString(bytes) => self.codec.decode(&bytes).map(Some),
|
||||
redis::Value::SimpleString(text) => self.codec.decode(text.as_bytes()).map(Some),
|
||||
_ => Err(Error::InvalidEntry),
|
||||
}
|
||||
}
|
||||
|
||||
fn decode_batch_response(&self, value: redis::Value) -> Result<BatchEntry<S::Value>, Error> {
|
||||
match self.decode_response(value) {
|
||||
Ok(Some(value)) => Ok(BatchEntry::Hit(value)),
|
||||
Ok(None) => Ok(BatchEntry::Miss),
|
||||
Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid),
|
||||
Err(error) => Err(error),
|
||||
}
|
||||
}
|
||||
|
||||
fn ttl_seconds(ttl: Duration) -> u64 {
|
||||
|
|
@ -61,196 +246,418 @@ where
|
|||
.max(1)
|
||||
}
|
||||
|
||||
fn run_blocking<T, F>(connection: Arc<Mutex<C>>, operation: F) -> CacheFuture<'static, T>
|
||||
async fn run_blocking<T, F>(connections: Arc<Connections<C>>, operation: F) -> Result<T, Error>
|
||||
where
|
||||
T: Send + 'static,
|
||||
F: FnOnce(&mut C) -> Result<T, Error> + Send + 'static,
|
||||
F: FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error> + Send + 'static,
|
||||
{
|
||||
Box::pin(async move {
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let mut connection = connection.lock().map_err(|_| Error::Unavailable)?;
|
||||
operation(&mut connection)
|
||||
})
|
||||
tokio::task::spawn_blocking(move || connections.execute(operation))
|
||||
.await
|
||||
.map_err(|_| Error::Unavailable)?
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl<C> BaseCache for RedisCache<C>
|
||||
fn namespaced_key(namespace: Option<&str>, key: &str) -> String {
|
||||
match namespace {
|
||||
Some(namespace) if !key.starts_with(&format!("{namespace}:")) => {
|
||||
format!("{namespace}:{key}")
|
||||
}
|
||||
_ => key.into(),
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, C> BaseCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
type Value = CacheEntry;
|
||||
type Value = S::Value;
|
||||
type Context = ExactCacheContext;
|
||||
|
||||
fn default_ttl(&self) -> Duration {
|
||||
self.default_ttl
|
||||
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
|
||||
context.ttl.or(Some(self.default_ttl))
|
||||
}
|
||||
|
||||
fn set_cache(&self, key: &str, value: Self::Value, kwargs: CacheKwargs) -> Result<(), Error> {
|
||||
let payload = Self::encode(&value)?;
|
||||
let ttl = Self::ttl_seconds(self.get_ttl(&kwargs));
|
||||
self.connection()?
|
||||
.set_ex::<_, _, ()>(Self::namespaced_key(key), payload, ttl)
|
||||
.map_err(|_| Error::Unavailable)
|
||||
}
|
||||
|
||||
fn get_cache(&self, key: &str, _: &CacheKwargs) -> Result<Option<Self::Value>, Error> {
|
||||
self.connection()?
|
||||
.get::<_, Option<Vec<u8>>>(Self::namespaced_key(key))
|
||||
.map_err(|_| Error::Unavailable)?
|
||||
.map(Self::decode)
|
||||
.transpose()
|
||||
}
|
||||
|
||||
fn delete_cache(&self, key: &str) -> Result<(), Error> {
|
||||
self.connection()?
|
||||
.del::<_, ()>(Self::namespaced_key(key))
|
||||
.map_err(|_| Error::Unavailable)
|
||||
}
|
||||
|
||||
fn flush_cache(&self) -> Result<(), Error> {
|
||||
let mut connection = self.connection()?;
|
||||
let keys = connection
|
||||
.scan_match(Self::namespaced_pattern())
|
||||
.map_err(|_| Error::Unavailable)?
|
||||
.collect::<redis::RedisResult<Vec<String>>>()
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
if keys.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
connection
|
||||
.del::<_, usize>(keys)
|
||||
.map(|_| ())
|
||||
.map_err(|_| Error::Unavailable)
|
||||
}
|
||||
|
||||
fn async_set_cache<'a>(
|
||||
&'a self,
|
||||
key: &'a str,
|
||||
fn set_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
value: Self::Value,
|
||||
kwargs: CacheKwargs,
|
||||
) -> CacheFuture<'a, ()> {
|
||||
let payload = Self::encode(&value);
|
||||
let key = Self::namespaced_key(key);
|
||||
let ttl = Self::ttl_seconds(self.get_ttl(&kwargs));
|
||||
Self::run_blocking(Arc::clone(&self.connection), move |connection| {
|
||||
context: &ExactCacheContext,
|
||||
) -> Result<(), Error> {
|
||||
let payload = self.codec.encode(&value)?;
|
||||
let ttl = Self::ttl_seconds(self.get_ttl(context).unwrap_or(self.default_ttl));
|
||||
let key = self.namespaced_key(key);
|
||||
self.connections.execute(|connection| {
|
||||
connection
|
||||
.set_ex::<_, _, ()>(key, payload?, ttl)
|
||||
.set_ex::<_, _, ()>(key, payload, ttl)
|
||||
.map_err(|_| Error::Unavailable)
|
||||
})
|
||||
}
|
||||
|
||||
fn async_get_cache<'a>(
|
||||
&'a self,
|
||||
key: &'a str,
|
||||
_: &'a CacheKwargs,
|
||||
) -> CacheFuture<'a, Option<Self::Value>> {
|
||||
let key = Self::namespaced_key(key);
|
||||
Box::pin(async move {
|
||||
Self::run_blocking(Arc::clone(&self.connection), move |connection| {
|
||||
connection
|
||||
.get::<_, Option<Vec<u8>>>(key)
|
||||
.map_err(|_| Error::Unavailable)
|
||||
})
|
||||
.await?
|
||||
.map(Self::decode)
|
||||
.transpose()
|
||||
})
|
||||
fn get_cache(&self, key: &str, _: &ExactCacheContext) -> Result<Option<Self::Value>, Error> {
|
||||
let key = self.namespaced_key(key);
|
||||
let value = self.connections.execute(|connection| {
|
||||
connection
|
||||
.get::<_, redis::Value>(key)
|
||||
.map_err(|_| Error::Unavailable)
|
||||
})?;
|
||||
self.decode_response(value)
|
||||
}
|
||||
|
||||
fn async_set_cache_pipeline<'a>(
|
||||
&'a self,
|
||||
async fn async_set_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
value: Self::Value,
|
||||
context: ExactCacheContext,
|
||||
) -> Result<(), Error> {
|
||||
let payload = self.codec.encode(&value)?;
|
||||
let key = self.namespaced_key(key);
|
||||
let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl));
|
||||
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
|
||||
connection
|
||||
.set_ex::<_, _, ()>(key, payload, ttl)
|
||||
.map_err(|_| Error::Unavailable)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
async fn async_get_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
_: &ExactCacheContext,
|
||||
) -> Result<Option<Self::Value>, Error> {
|
||||
let key = self.namespaced_key(key);
|
||||
let value = Self::run_blocking(Arc::clone(&self.connections), move |connection| {
|
||||
connection
|
||||
.get::<_, redis::Value>(key)
|
||||
.map_err(|_| Error::Unavailable)
|
||||
})
|
||||
.await?;
|
||||
self.decode_response(value)
|
||||
}
|
||||
|
||||
async fn async_set_cache_pipeline(
|
||||
&self,
|
||||
cache_list: Vec<(String, Self::Value)>,
|
||||
kwargs: CacheKwargs,
|
||||
) -> CacheFuture<'a, ()> {
|
||||
context: ExactCacheContext,
|
||||
) -> Result<(), Error> {
|
||||
let entries = cache_list
|
||||
.into_iter()
|
||||
.map(|(key, value)| {
|
||||
Self::encode(&value).map(|payload| (Self::namespaced_key(&key), payload))
|
||||
self.codec
|
||||
.encode(&value)
|
||||
.map(|payload| (self.namespaced_key(&key), payload))
|
||||
})
|
||||
.collect::<Result<Vec<_>, _>>();
|
||||
let ttl = Self::ttl_seconds(self.get_ttl(&kwargs));
|
||||
Self::run_blocking(Arc::clone(&self.connection), move |connection| {
|
||||
for (key, payload) in entries? {
|
||||
connection
|
||||
.set_ex::<_, _, ()>(key, payload, ttl)
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl));
|
||||
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
|
||||
let mut pipeline = redis::pipe();
|
||||
for (key, payload) in entries {
|
||||
pipeline
|
||||
.cmd("SETEX")
|
||||
.arg(key)
|
||||
.arg(ttl)
|
||||
.arg(payload)
|
||||
.ignore();
|
||||
}
|
||||
Ok(())
|
||||
pipeline
|
||||
.query::<()>(connection)
|
||||
.map_err(|_| Error::Unavailable)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
fn async_delete_cache<'a>(&'a self, key: &'a str) -> CacheFuture<'a, ()> {
|
||||
let key = Self::namespaced_key(key);
|
||||
Self::run_blocking(Arc::clone(&self.connection), move |connection| {
|
||||
async fn disconnect(&self) -> Result<(), Error> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
|
||||
match Self::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()),
|
||||
}),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, C> BatchCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
fn batch_get_cache(
|
||||
&self,
|
||||
keys: &[String],
|
||||
_: &ExactCacheContext,
|
||||
) -> Result<Vec<BatchEntry<Self::Value>>, Error> {
|
||||
let keys = keys
|
||||
.iter()
|
||||
.map(|key| self.namespaced_key(key))
|
||||
.collect::<Vec<_>>();
|
||||
let values = self.connections.execute(|connection| {
|
||||
redis::cmd("MGET")
|
||||
.arg(keys)
|
||||
.query::<Vec<redis::Value>>(connection)
|
||||
.map_err(|_| Error::Unavailable)
|
||||
})?;
|
||||
values
|
||||
.into_iter()
|
||||
.map(|value| self.decode_batch_response(value))
|
||||
.collect()
|
||||
}
|
||||
|
||||
async fn async_batch_get_cache(
|
||||
&self,
|
||||
keys: Vec<String>,
|
||||
_: ExactCacheContext,
|
||||
) -> Result<Vec<BatchEntry<Self::Value>>, Error> {
|
||||
let keys = keys
|
||||
.iter()
|
||||
.map(|key| self.namespaced_key(key))
|
||||
.collect::<Vec<_>>();
|
||||
let values = Self::run_blocking(Arc::clone(&self.connections), move |connection| {
|
||||
redis::cmd("MGET")
|
||||
.arg(keys)
|
||||
.query::<Vec<redis::Value>>(connection)
|
||||
.map_err(|_| Error::Unavailable)
|
||||
})
|
||||
.await?;
|
||||
values
|
||||
.into_iter()
|
||||
.map(|value| self.decode_batch_response(value))
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, C> DeleteCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
fn delete_cache(&self, key: &str) -> Result<(), Error> {
|
||||
let key = self.namespaced_key(key);
|
||||
self.connections
|
||||
.execute(|connection| connection.del::<_, ()>(key).map_err(|_| Error::Unavailable))
|
||||
}
|
||||
|
||||
async fn async_delete_cache(&self, key: &str) -> Result<(), Error> {
|
||||
let key = self.namespaced_key(key);
|
||||
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
|
||||
connection.del::<_, ()>(key).map_err(|_| Error::Unavailable)
|
||||
})
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, C> FlushCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
fn flush_cache(&self) -> Result<(), Error> {
|
||||
let pattern = self.namespaced_pattern()?;
|
||||
self.connections
|
||||
.execute(|connection| Self::flush_matching(connection, &pattern))
|
||||
}
|
||||
|
||||
fn disconnect(&self) -> CacheFuture<'_, ()> {
|
||||
Box::pin(async { Ok(()) })
|
||||
}
|
||||
|
||||
fn test_connection(&self) -> CacheFuture<'_, CacheConnectionResult> {
|
||||
Box::pin(async move {
|
||||
Self::run_blocking(Arc::clone(&self.connection), |connection| {
|
||||
redis::cmd("PING")
|
||||
.query::<String>(connection)
|
||||
.map_err(|_| Error::Unavailable)
|
||||
})
|
||||
.await?;
|
||||
Ok(CacheConnectionResult {
|
||||
status: CacheConnectionStatus::Success,
|
||||
message: "Redis cache connection test successful".into(),
|
||||
error: None,
|
||||
})
|
||||
async fn async_flush_cache(&self) -> Result<(), Error> {
|
||||
let pattern = self.namespaced_pattern()?;
|
||||
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
|
||||
Self::flush_matching(connection, &pattern)
|
||||
})
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, C> CounterCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec<Value = f64>,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
fn increment_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
amount: f64,
|
||||
context: ExactCacheContext,
|
||||
) -> Result<f64, Error> {
|
||||
let key = self.namespaced_key(key);
|
||||
let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl));
|
||||
self.connections
|
||||
.execute(|connection| increment(connection, key, amount, ttl))
|
||||
}
|
||||
|
||||
async fn async_increment(
|
||||
&self,
|
||||
key: &str,
|
||||
amount: f64,
|
||||
context: ExactCacheContext,
|
||||
) -> Result<f64, Error> {
|
||||
let key = self.namespaced_key(key);
|
||||
let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl));
|
||||
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
|
||||
increment(connection, key, amount, ttl)
|
||||
})
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
fn increment(
|
||||
connection: &mut ConnectionRef<'_>,
|
||||
key: String,
|
||||
amount: f64,
|
||||
ttl: u64,
|
||||
) -> Result<f64, Error> {
|
||||
redis::cmd("EVAL")
|
||||
.arg(INCREMENT_SCRIPT)
|
||||
.arg(1)
|
||||
.arg(key)
|
||||
.arg(amount)
|
||||
.arg(ttl)
|
||||
.query(connection)
|
||||
.map_err(|_| Error::Unavailable)
|
||||
}
|
||||
|
||||
fn stored_bytes(value: redis::Value) -> Result<Option<Vec<u8>>, Error> {
|
||||
match value {
|
||||
redis::Value::Nil => Ok(None),
|
||||
redis::Value::BulkString(bytes) => Ok(Some(bytes)),
|
||||
redis::Value::SimpleString(text) => Ok(Some(text.into_bytes())),
|
||||
_ => Err(Error::InvalidEntry),
|
||||
}
|
||||
}
|
||||
|
||||
/// Eligibility is decided on decoded values, so a pin written by another encoder (Python's
|
||||
/// `json.dumps` spacing or key order) still matches. The write is a compare-and-set on the
|
||||
/// bytes that decision was made on, retried when another claimant wins the race.
|
||||
fn claim<S: CacheCodec>(
|
||||
connection: &mut ConnectionRef<'_>,
|
||||
codec: &S,
|
||||
key: &str,
|
||||
candidate: S::Value,
|
||||
eligible: &[S::Value],
|
||||
ttl: u64,
|
||||
) -> Result<S::Value, Error>
|
||||
where
|
||||
S::Value: PartialEq,
|
||||
{
|
||||
let payload = codec.encode(&candidate)?;
|
||||
if payload.is_empty() {
|
||||
return Err(Error::InvalidEntry);
|
||||
}
|
||||
for _ in 0..CLAIM_ATTEMPTS {
|
||||
let current = stored_bytes(
|
||||
connection
|
||||
.get::<_, redis::Value>(key)
|
||||
.map_err(|_| Error::Unavailable)?,
|
||||
)?
|
||||
.filter(|bytes| !bytes.is_empty());
|
||||
let existing = current
|
||||
.as_deref()
|
||||
.and_then(|bytes| codec.decode(bytes).ok())
|
||||
.filter(|existing| eligible.is_empty() || eligible.contains(existing));
|
||||
let refresh = existing
|
||||
.as_ref()
|
||||
.is_some_and(|existing| !eligible.is_empty() || *existing == candidate);
|
||||
let write: &[u8] = if existing.is_some() { b"" } else { &payload };
|
||||
let applied = redis::cmd("EVAL")
|
||||
.arg(CLAIM_SCRIPT)
|
||||
.arg(1)
|
||||
.arg(key)
|
||||
.arg(current.as_deref().unwrap_or_default())
|
||||
.arg(ttl)
|
||||
.arg(write)
|
||||
.arg(u8::from(refresh))
|
||||
.query::<bool>(connection)
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
if applied {
|
||||
return Ok(existing.unwrap_or(candidate));
|
||||
}
|
||||
}
|
||||
Err(Error::Unavailable)
|
||||
}
|
||||
|
||||
impl<S, C> ClaimCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec + Clone + 'static,
|
||||
S::Value: PartialEq,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
fn claim_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
candidate: S::Value,
|
||||
eligible: &[S::Value],
|
||||
context: ExactCacheContext,
|
||||
) -> Result<S::Value, Error> {
|
||||
let key = self.namespaced_key(key);
|
||||
let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl));
|
||||
self.connections
|
||||
.execute(|connection| claim(connection, &self.codec, &key, candidate, eligible, ttl))
|
||||
}
|
||||
|
||||
async fn async_claim_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
candidate: S::Value,
|
||||
eligible: Vec<S::Value>,
|
||||
context: ExactCacheContext,
|
||||
) -> Result<S::Value, Error> {
|
||||
let key = self.namespaced_key(key);
|
||||
let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl));
|
||||
let codec = self.codec.clone();
|
||||
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
|
||||
claim(connection, &codec, &key, candidate, &eligible, ttl)
|
||||
})
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::RedisCache;
|
||||
use litellm_cache::{BaseCache, CacheEntry, CacheKwargs};
|
||||
use redis_test::{MockCmd, MockRedisConnection};
|
||||
use serde_json::json;
|
||||
use std::time::Duration;
|
||||
|
||||
fn entry() -> CacheEntry {
|
||||
CacheEntry {
|
||||
timestamp: 123.0,
|
||||
response: json!({"choices": [{"text": "cached"}]}),
|
||||
}
|
||||
}
|
||||
use litellm_cache::{
|
||||
BaseCache, CacheCodec, DeleteCache, ExactCacheContext, FlushCache, JsonCodec,
|
||||
};
|
||||
use redis_test::{MockCmd, MockRedisConnection};
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn cache_entries_round_trip_through_json() {
|
||||
let entry = entry();
|
||||
let encoded = RedisCache::<redis::Connection>::encode(&entry).unwrap();
|
||||
assert_eq!(
|
||||
RedisCache::<redis::Connection>::decode(encoded).unwrap(),
|
||||
entry
|
||||
);
|
||||
}
|
||||
use super::RedisCache;
|
||||
|
||||
#[test]
|
||||
fn invalid_json_is_rejected() {
|
||||
assert!(RedisCache::<redis::Connection>::decode(b"not json".to_vec()).is_err());
|
||||
fn entry() -> serde_json::Value {
|
||||
json!({"deployment": "model-a", "cooldown_seconds": 30})
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ttl_seconds_rounds_up_and_keeps_expiration_positive() {
|
||||
assert_eq!(
|
||||
RedisCache::<redis::Connection>::ttl_seconds(Duration::ZERO),
|
||||
RedisCache::<JsonCodec<serde_json::Value>>::ttl_seconds(Duration::ZERO),
|
||||
1
|
||||
);
|
||||
assert_eq!(
|
||||
RedisCache::<redis::Connection>::ttl_seconds(Duration::from_millis(1500)),
|
||||
RedisCache::<JsonCodec<serde_json::Value>>::ttl_seconds(Duration::from_millis(1500)),
|
||||
2
|
||||
);
|
||||
assert_eq!(
|
||||
RedisCache::<redis::Connection>::ttl_seconds(Duration::from_secs(15)),
|
||||
RedisCache::<JsonCodec<serde_json::Value>>::ttl_seconds(Duration::from_secs(15)),
|
||||
15
|
||||
);
|
||||
}
|
||||
|
|
@ -258,7 +665,9 @@ mod tests {
|
|||
#[test]
|
||||
fn redis_commands_round_trip_entries_and_delete_only_namespaced_keys() {
|
||||
let value = entry();
|
||||
let payload = RedisCache::<redis::Connection>::encode(&value).unwrap();
|
||||
let payload = JsonCodec::<serde_json::Value>::new()
|
||||
.encode(&value)
|
||||
.unwrap();
|
||||
let connection = MockRedisConnection::new([
|
||||
MockCmd::new(
|
||||
redis::cmd("SETEX")
|
||||
|
|
@ -271,13 +680,17 @@ mod tests {
|
|||
MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)),
|
||||
])
|
||||
.assert_all_commands_consumed();
|
||||
let cache = RedisCache::with_connection(connection, None);
|
||||
let cache =
|
||||
RedisCache::with_connection(connection, None, JsonCodec::<serde_json::Value>::new())
|
||||
.with_namespace(Some("litellm-cache".into()));
|
||||
|
||||
cache
|
||||
.set_cache("key", value.clone(), CacheKwargs::default())
|
||||
.set_cache("key", value.clone(), &ExactCacheContext::default())
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
cache.get_cache("key", &CacheKwargs::default()).unwrap(),
|
||||
cache
|
||||
.get_cache("key", &ExactCacheContext::default())
|
||||
.unwrap(),
|
||||
Some(value)
|
||||
);
|
||||
cache.delete_cache("key").unwrap();
|
||||
|
|
@ -290,13 +703,17 @@ mod tests {
|
|||
redis::cmd("SCAN")
|
||||
.cursor_arg(0)
|
||||
.arg("MATCH")
|
||||
.arg("litellm-cache:*"),
|
||||
.arg("litellm-cache:*")
|
||||
.arg("COUNT")
|
||||
.arg(1000),
|
||||
Ok(redis_test::redis_value!(["0", ["litellm-cache:key"]])),
|
||||
),
|
||||
MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)),
|
||||
])
|
||||
.assert_all_commands_consumed();
|
||||
let cache = RedisCache::with_connection(connection, None);
|
||||
let cache =
|
||||
RedisCache::with_connection(connection, None, JsonCodec::<serde_json::Value>::new())
|
||||
.with_namespace(Some("litellm-cache".into()));
|
||||
|
||||
cache.flush_cache().unwrap();
|
||||
}
|
||||
|
|
@ -305,7 +722,9 @@ mod tests {
|
|||
async fn test_connection_runs_ping_off_executor() {
|
||||
let connection = MockRedisConnection::new([MockCmd::new(redis::cmd("PING"), Ok("PONG"))])
|
||||
.assert_all_commands_consumed();
|
||||
let cache = RedisCache::with_connection(connection, None);
|
||||
let cache =
|
||||
RedisCache::with_connection(connection, None, JsonCodec::<serde_json::Value>::new())
|
||||
.with_namespace(Some("litellm-cache".into()));
|
||||
|
||||
assert_eq!(
|
||||
cache.test_connection().await.unwrap().status,
|
||||
|
|
|
|||
633
litellm-rust/crates/cache-redis/src/cache/operations.rs
vendored
Normal file
633
litellm-rust/crates/cache-redis/src/cache/operations.rs
vendored
Normal file
|
|
@ -0,0 +1,633 @@
|
|||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use litellm_cache::{
|
||||
CacheCodec, CacheScript, ClientInfoCache, Error, IncrementOperation, QueueCache, ScanCache,
|
||||
ScriptCache, SetCache, TtlCache,
|
||||
};
|
||||
use redis::Commands;
|
||||
|
||||
use super::{ConnectionRef, Connections, RedisCache, namespaced_key};
|
||||
|
||||
const INCREMENT_WITH_FLOOR_SCRIPT: &str = concat!(
|
||||
"local count = redis.call('INCRBY', KEYS[1], ARGV[1]); ",
|
||||
"if count < 0 then count = redis.call('INCRBY', KEYS[1], -count); end; ",
|
||||
"if redis.call('TTL', KEYS[1]) < 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; ",
|
||||
"return count"
|
||||
);
|
||||
const SET_MAX_SCRIPT: &str = concat!(
|
||||
"local current = redis.call('GET', KEYS[1]); ",
|
||||
"if current == false or tonumber(current) < tonumber(ARGV[1]) then ",
|
||||
"redis.call('SET', KEYS[1], ARGV[1]); ",
|
||||
"if tonumber(ARGV[2]) > 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; ",
|
||||
"return ARGV[1]; end; return current"
|
||||
);
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub enum RedisArg {
|
||||
Bytes(Vec<u8>),
|
||||
Integer(i64),
|
||||
Float(f64),
|
||||
}
|
||||
|
||||
impl From<&str> for RedisArg {
|
||||
fn from(value: &str) -> Self {
|
||||
Self::Bytes(value.as_bytes().to_vec())
|
||||
}
|
||||
}
|
||||
|
||||
impl From<String> for RedisArg {
|
||||
fn from(value: String) -> Self {
|
||||
Self::Bytes(value.into_bytes())
|
||||
}
|
||||
}
|
||||
|
||||
impl From<Vec<u8>> for RedisArg {
|
||||
fn from(value: Vec<u8>) -> Self {
|
||||
Self::Bytes(value)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<i64> for RedisArg {
|
||||
fn from(value: i64) -> Self {
|
||||
Self::Integer(value)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<f64> for RedisArg {
|
||||
fn from(value: f64) -> Self {
|
||||
Self::Float(value)
|
||||
}
|
||||
}
|
||||
|
||||
impl redis::ToRedisArgs for RedisArg {
|
||||
fn write_redis_args<W>(&self, out: &mut W)
|
||||
where
|
||||
W: ?Sized + redis::RedisWrite,
|
||||
{
|
||||
match self {
|
||||
Self::Bytes(value) => value.write_redis_args(out),
|
||||
Self::Integer(value) => value.write_redis_args(out),
|
||||
Self::Float(value) => value.write_redis_args(out),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct RedisRpushOperation {
|
||||
pub key: String,
|
||||
pub values: Vec<RedisArg>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct RedisLpopOperation {
|
||||
pub key: String,
|
||||
pub count: Option<usize>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub enum RedisLpopResult {
|
||||
Missing,
|
||||
Value(Vec<u8>),
|
||||
Values(Vec<Vec<u8>>),
|
||||
}
|
||||
|
||||
pub struct RedisScript<C> {
|
||||
connections: Arc<Connections<C>>,
|
||||
namespace: Option<String>,
|
||||
source: String,
|
||||
}
|
||||
|
||||
impl<C> CacheScript for RedisScript<C>
|
||||
where
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
type Argument = RedisArg;
|
||||
type Output = redis::Value;
|
||||
|
||||
async fn invoke(
|
||||
&self,
|
||||
keys: Vec<String>,
|
||||
arguments: Vec<Self::Argument>,
|
||||
) -> Result<Self::Output, Error> {
|
||||
let keys = keys
|
||||
.into_iter()
|
||||
.map(|key| namespaced_key(self.namespace.as_deref(), &key))
|
||||
.collect::<Vec<_>>();
|
||||
let connections = Arc::clone(&self.connections);
|
||||
let source = self.source.clone();
|
||||
tokio::task::spawn_blocking(move || {
|
||||
connections.execute(|connection| {
|
||||
redis::cmd("EVAL")
|
||||
.arg(source)
|
||||
.arg(keys.len())
|
||||
.arg(keys)
|
||||
.arg(arguments)
|
||||
.query(connection)
|
||||
.map_err(|_| Error::Unavailable)
|
||||
})
|
||||
})
|
||||
.await
|
||||
.map_err(|_| Error::Unavailable)?
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, C> RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
pub async fn delete_cache_keys(&self, keys: Vec<String>) -> Result<usize, Error> {
|
||||
if keys.is_empty() {
|
||||
return Ok(0);
|
||||
}
|
||||
let keys = keys
|
||||
.into_iter()
|
||||
.map(|key| self.namespaced_key(&key))
|
||||
.collect::<Vec<_>>();
|
||||
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
|
||||
connection.del(keys).map_err(|_| Error::Unavailable)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub fn batch_get_counts(&self, keys: &[String]) -> Result<Vec<Option<i64>>, Error> {
|
||||
let keys = keys
|
||||
.iter()
|
||||
.map(|key| self.namespaced_key(key))
|
||||
.collect::<Vec<_>>();
|
||||
let values = self.connections.execute(|connection| {
|
||||
redis::cmd("MGET")
|
||||
.arg(keys)
|
||||
.query::<Vec<redis::Value>>(connection)
|
||||
.map_err(|_| Error::Unavailable)
|
||||
})?;
|
||||
values.into_iter().map(count).collect()
|
||||
}
|
||||
|
||||
pub async fn async_batch_get_counts(
|
||||
&self,
|
||||
keys: Vec<String>,
|
||||
) -> Result<Vec<Option<i64>>, Error> {
|
||||
let keys = keys
|
||||
.iter()
|
||||
.map(|key| self.namespaced_key(key))
|
||||
.collect::<Vec<_>>();
|
||||
let values = Self::run_blocking(Arc::clone(&self.connections), move |connection| {
|
||||
redis::cmd("MGET")
|
||||
.arg(keys)
|
||||
.query::<Vec<redis::Value>>(connection)
|
||||
.map_err(|_| Error::Unavailable)
|
||||
})
|
||||
.await?;
|
||||
values.into_iter().map(count).collect()
|
||||
}
|
||||
|
||||
pub fn sync_ping(&self) -> Result<bool, Error> {
|
||||
self.connections.execute(|connection| {
|
||||
redis::cmd("PING")
|
||||
.query::<String>(connection)
|
||||
.map(|response| response == "PONG")
|
||||
.map_err(|_| Error::Unavailable)
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn ping(&self) -> Result<bool, Error> {
|
||||
Self::run_blocking(Arc::clone(&self.connections), |connection| {
|
||||
redis::cmd("PING")
|
||||
.query::<String>(connection)
|
||||
.map(|response| response == "PONG")
|
||||
.map_err(|_| Error::Unavailable)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn async_get_ttl(&self, key: &str) -> Result<Option<i64>, Error> {
|
||||
let key = self.namespaced_key(key);
|
||||
let ttl = Self::run_blocking(Arc::clone(&self.connections), move |connection| {
|
||||
redis::cmd("TTL")
|
||||
.arg(key)
|
||||
.query::<i64>(connection)
|
||||
.map_err(|_| Error::Unavailable)
|
||||
})
|
||||
.await?;
|
||||
Ok((ttl >= 0).then_some(ttl))
|
||||
}
|
||||
|
||||
pub async fn async_scan_iter(&self, pattern: &str, count: usize) -> Result<Vec<String>, Error> {
|
||||
let pattern = format!("{}*", self.namespaced_key(pattern));
|
||||
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
|
||||
let mut cursor = 0u64;
|
||||
let mut matches = Vec::new();
|
||||
loop {
|
||||
let (next_cursor, keys): (u64, Vec<String>) = redis::cmd("SCAN")
|
||||
.cursor_arg(cursor)
|
||||
.arg("MATCH")
|
||||
.arg(&pattern)
|
||||
.arg("COUNT")
|
||||
.arg(count)
|
||||
.query(connection)
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
matches.extend(keys);
|
||||
if matches.len() >= count || next_cursor == 0 {
|
||||
matches.truncate(count);
|
||||
return Ok(matches);
|
||||
}
|
||||
cursor = next_cursor;
|
||||
}
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn async_set_cache_sadd(
|
||||
&self,
|
||||
key: &str,
|
||||
values: Vec<RedisArg>,
|
||||
ttl: Option<Duration>,
|
||||
) -> Result<usize, Error> {
|
||||
if values.is_empty() {
|
||||
return Err(Error::InvalidEntry);
|
||||
}
|
||||
let key = self.namespaced_key(key);
|
||||
let ttl = Self::ttl_seconds(ttl.unwrap_or(self.default_ttl));
|
||||
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
|
||||
let mut pipeline = redis::pipe();
|
||||
pipeline.cmd("SADD").arg(&key).arg(values);
|
||||
pipeline.cmd("EXPIRE").arg(&key).arg(ttl).ignore();
|
||||
pipeline
|
||||
.query::<(usize,)>(connection)
|
||||
.map(|(added,)| added)
|
||||
.map_err(|_| Error::Unavailable)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn async_rpush(&self, key: &str, values: Vec<RedisArg>) -> Result<usize, Error> {
|
||||
if values.is_empty() {
|
||||
return Err(Error::InvalidEntry);
|
||||
}
|
||||
let key = self.namespaced_key(key);
|
||||
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
|
||||
redis::cmd("RPUSH")
|
||||
.arg(key)
|
||||
.arg(values)
|
||||
.query(connection)
|
||||
.map_err(|_| Error::Unavailable)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn async_rpush_pipeline(
|
||||
&self,
|
||||
operations: Vec<RedisRpushOperation>,
|
||||
) -> Result<Vec<usize>, Error> {
|
||||
let operations = operations
|
||||
.into_iter()
|
||||
.map(|operation| {
|
||||
if operation.values.is_empty() {
|
||||
return Err(Error::InvalidEntry);
|
||||
}
|
||||
Ok((self.namespaced_key(&operation.key), operation.values))
|
||||
})
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
if operations.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
|
||||
let mut pipeline = redis::pipe();
|
||||
for (key, values) in operations {
|
||||
pipeline.cmd("RPUSH").arg(key).arg(values);
|
||||
}
|
||||
pipeline.query(connection).map_err(|_| Error::Unavailable)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn async_lpop(
|
||||
&self,
|
||||
key: &str,
|
||||
count: Option<usize>,
|
||||
) -> Result<RedisLpopResult, Error> {
|
||||
let key = self.namespaced_key(key);
|
||||
let multiple = count.is_some();
|
||||
let value = Self::run_blocking(Arc::clone(&self.connections), move |connection| {
|
||||
let mut command = redis::cmd("LPOP");
|
||||
command.arg(key);
|
||||
if let Some(count) = count {
|
||||
command.arg(count);
|
||||
}
|
||||
command
|
||||
.query::<redis::Value>(connection)
|
||||
.map_err(|_| Error::Unavailable)
|
||||
})
|
||||
.await?;
|
||||
lpop_result(value, multiple)
|
||||
}
|
||||
|
||||
pub async fn async_lpop_pipeline(
|
||||
&self,
|
||||
operations: Vec<RedisLpopOperation>,
|
||||
) -> Result<Vec<RedisLpopResult>, Error> {
|
||||
let operations = operations
|
||||
.into_iter()
|
||||
.map(|operation| (self.namespaced_key(&operation.key), operation.count))
|
||||
.collect::<Vec<_>>();
|
||||
if operations.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let multiple = operations
|
||||
.iter()
|
||||
.map(|(_, count)| count.is_some())
|
||||
.collect::<Vec<_>>();
|
||||
let values = Self::run_blocking(Arc::clone(&self.connections), move |connection| {
|
||||
let mut pipeline = redis::pipe();
|
||||
for (key, count) in operations {
|
||||
let command = pipeline.cmd("LPOP").arg(key);
|
||||
if let Some(count) = count {
|
||||
command.arg(count);
|
||||
}
|
||||
}
|
||||
pipeline
|
||||
.query::<Vec<redis::Value>>(connection)
|
||||
.map_err(|_| Error::Unavailable)
|
||||
})
|
||||
.await?;
|
||||
values
|
||||
.into_iter()
|
||||
.zip(multiple)
|
||||
.map(|(value, multiple)| lpop_result(value, multiple))
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub async fn async_eval(
|
||||
&self,
|
||||
script: String,
|
||||
keys: Vec<String>,
|
||||
arguments: Vec<RedisArg>,
|
||||
) -> Result<redis::Value, Error> {
|
||||
let keys = keys
|
||||
.into_iter()
|
||||
.map(|key| self.namespaced_key(&key))
|
||||
.collect::<Vec<_>>();
|
||||
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
|
||||
redis::cmd("EVAL")
|
||||
.arg(script)
|
||||
.arg(keys.len())
|
||||
.arg(keys)
|
||||
.arg(arguments)
|
||||
.query(connection)
|
||||
.map_err(|_| Error::Unavailable)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub fn client_list(&self) -> Result<String, Error> {
|
||||
self.connections.execute(|connection| {
|
||||
redis::cmd("CLIENT")
|
||||
.arg("LIST")
|
||||
.query(connection)
|
||||
.map_err(|_| Error::Unavailable)
|
||||
})
|
||||
}
|
||||
|
||||
pub fn info(&self) -> Result<String, Error> {
|
||||
self.connections.execute(|connection| {
|
||||
redis::cmd("INFO")
|
||||
.query(connection)
|
||||
.map_err(|_| Error::Unavailable)
|
||||
})
|
||||
}
|
||||
|
||||
pub fn flushall(&self) -> Result<(), Error> {
|
||||
self.connections.execute(|connection| {
|
||||
redis::cmd("FLUSHALL")
|
||||
.query(connection)
|
||||
.map_err(|_| Error::Unavailable)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, C> RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec<Value = f64>,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
pub fn increment_with_floor(
|
||||
&self,
|
||||
key: &str,
|
||||
amount: i64,
|
||||
ttl: Duration,
|
||||
) -> Result<i64, Error> {
|
||||
let key = self.namespaced_key(key);
|
||||
let ttl = Self::ttl_seconds(ttl);
|
||||
self.connections
|
||||
.execute(|connection| increment_with_floor(connection, key, amount, ttl))
|
||||
}
|
||||
|
||||
pub async fn async_increment_pipeline(
|
||||
&self,
|
||||
operations: Vec<IncrementOperation>,
|
||||
) -> Result<Vec<f64>, Error> {
|
||||
let operations = operations
|
||||
.into_iter()
|
||||
.map(|operation| {
|
||||
(
|
||||
self.namespaced_key(&operation.key),
|
||||
operation.amount,
|
||||
operation.ttl.map(Self::ttl_seconds),
|
||||
)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
if operations.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
|
||||
let mut pipeline = redis::pipe();
|
||||
for (key, amount, ttl) in operations {
|
||||
pipeline.cmd("INCRBYFLOAT").arg(&key).arg(amount);
|
||||
if let Some(ttl) = ttl {
|
||||
pipeline.cmd("EXPIRE").arg(key).arg(ttl).ignore();
|
||||
}
|
||||
}
|
||||
pipeline.query(connection).map_err(|_| Error::Unavailable)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn async_increment_with_floor(
|
||||
&self,
|
||||
key: &str,
|
||||
amount: i64,
|
||||
ttl: Duration,
|
||||
) -> Result<i64, Error> {
|
||||
let key = self.namespaced_key(key);
|
||||
let ttl = Self::ttl_seconds(ttl);
|
||||
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
|
||||
increment_with_floor(connection, key, amount, ttl)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn async_set_max(
|
||||
&self,
|
||||
key: &str,
|
||||
value: f64,
|
||||
ttl: Option<Duration>,
|
||||
) -> Result<f64, Error> {
|
||||
let key = self.namespaced_key(key);
|
||||
let ttl = Self::ttl_seconds(ttl.unwrap_or(self.default_ttl));
|
||||
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
|
||||
redis::cmd("EVAL")
|
||||
.arg(SET_MAX_SCRIPT)
|
||||
.arg(1)
|
||||
.arg(key)
|
||||
.arg(value)
|
||||
.arg(ttl)
|
||||
.query(connection)
|
||||
.map_err(|_| Error::Unavailable)
|
||||
})
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
fn redis_bytes(value: redis::Value) -> Result<Vec<u8>, Error> {
|
||||
match value {
|
||||
redis::Value::BulkString(bytes) => Ok(bytes),
|
||||
redis::Value::SimpleString(text) => Ok(text.into_bytes()),
|
||||
_ => Err(Error::InvalidEntry),
|
||||
}
|
||||
}
|
||||
|
||||
fn lpop_result(value: redis::Value, multiple: bool) -> Result<RedisLpopResult, Error> {
|
||||
match value {
|
||||
redis::Value::Nil => Ok(RedisLpopResult::Missing),
|
||||
redis::Value::Array(values) if multiple => values
|
||||
.into_iter()
|
||||
.map(redis_bytes)
|
||||
.collect::<Result<Vec<_>, _>>()
|
||||
.map(RedisLpopResult::Values),
|
||||
value if !multiple => redis_bytes(value).map(RedisLpopResult::Value),
|
||||
_ => Err(Error::InvalidEntry),
|
||||
}
|
||||
}
|
||||
|
||||
fn count(value: redis::Value) -> Result<Option<i64>, Error> {
|
||||
match value {
|
||||
redis::Value::Nil => Ok(None),
|
||||
redis::Value::Int(value) => Ok(Some(value)),
|
||||
redis::Value::BulkString(value) => std::str::from_utf8(&value)
|
||||
.ok()
|
||||
.and_then(|value| value.parse().ok())
|
||||
.map(Some)
|
||||
.ok_or(Error::InvalidEntry),
|
||||
redis::Value::SimpleString(value) => {
|
||||
value.parse().map(Some).map_err(|_| Error::InvalidEntry)
|
||||
}
|
||||
_ => Err(Error::InvalidEntry),
|
||||
}
|
||||
}
|
||||
|
||||
fn increment_with_floor(
|
||||
connection: &mut ConnectionRef<'_>,
|
||||
key: String,
|
||||
amount: i64,
|
||||
ttl: u64,
|
||||
) -> Result<i64, Error> {
|
||||
redis::cmd("EVAL")
|
||||
.arg(INCREMENT_WITH_FLOOR_SCRIPT)
|
||||
.arg(1)
|
||||
.arg(key)
|
||||
.arg(amount)
|
||||
.arg(ttl)
|
||||
.query(connection)
|
||||
.map_err(|_| Error::Unavailable)
|
||||
}
|
||||
|
||||
impl<S, C> TtlCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
async fn async_get_ttl(&self, key: &str) -> Result<Option<Duration>, Error> {
|
||||
RedisCache::async_get_ttl(self, key)
|
||||
.await
|
||||
.map(|ttl| ttl.map(|seconds| Duration::from_secs(seconds as u64)))
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, C> ScanCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
async fn async_scan_iter(&self, pattern: &str, count: usize) -> Result<Vec<String>, Error> {
|
||||
RedisCache::async_scan_iter(self, pattern, count).await
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, C> ClientInfoCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
type ClientList = String;
|
||||
type Info = String;
|
||||
|
||||
fn client_list(&self) -> Result<Self::ClientList, Error> {
|
||||
RedisCache::client_list(self)
|
||||
}
|
||||
|
||||
fn info(&self) -> Result<Self::Info, Error> {
|
||||
RedisCache::info(self)
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, C> SetCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
type SetValue = RedisArg;
|
||||
type SetResult = usize;
|
||||
|
||||
async fn async_set_cache_sadd(
|
||||
&self,
|
||||
key: &str,
|
||||
values: Vec<Self::SetValue>,
|
||||
ttl: Option<Duration>,
|
||||
) -> Result<Self::SetResult, Error> {
|
||||
RedisCache::async_set_cache_sadd(self, key, values, ttl).await
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, C> QueueCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
type QueueValue = RedisArg;
|
||||
type PopResult = RedisLpopResult;
|
||||
|
||||
async fn async_rpush(&self, key: &str, values: Vec<Self::QueueValue>) -> Result<usize, Error> {
|
||||
RedisCache::async_rpush(self, key, values).await
|
||||
}
|
||||
|
||||
async fn async_lpop(&self, key: &str, count: Option<usize>) -> Result<Self::PopResult, Error> {
|
||||
RedisCache::async_lpop(self, key, count).await
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, C> ScriptCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
type Script = RedisScript<C>;
|
||||
|
||||
fn async_register_script(&self, source: String) -> Self::Script {
|
||||
RedisScript {
|
||||
connections: Arc::clone(&self.connections),
|
||||
namespace: self.namespace.clone(),
|
||||
source,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -1,3 +1,7 @@
|
|||
mod cache;
|
||||
mod topology;
|
||||
|
||||
pub use cache::RedisCache;
|
||||
pub use cache::{
|
||||
RedisArg, RedisCache, RedisLpopOperation, RedisLpopResult, RedisRpushOperation, RedisScript,
|
||||
};
|
||||
pub use topology::{RedisNode, RedisTopology};
|
||||
|
|
|
|||
14
litellm-rust/crates/cache-redis/src/topology.rs
Normal file
14
litellm-rust/crates/cache-redis/src/topology.rs
Normal file
|
|
@ -0,0 +1,14 @@
|
|||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct RedisNode {
|
||||
pub host: String,
|
||||
pub port: u16,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default, PartialEq, Eq)]
|
||||
pub enum RedisTopology {
|
||||
#[default]
|
||||
Standalone,
|
||||
Cluster {
|
||||
startup_nodes: Vec<RedisNode>,
|
||||
},
|
||||
}
|
||||
|
|
@ -1,6 +1,703 @@
|
|||
use litellm_cache_redis::RedisCache;
|
||||
use std::time::Duration;
|
||||
|
||||
use litellm_cache::{
|
||||
BaseCache, BatchCache, BatchEntry, CacheCodec, CacheConnectionStatus, CacheScript, ClaimCache,
|
||||
CounterCache, DeleteCache, Error, ExactCacheContext, FlushCache, IncrementOperation, JsonCodec,
|
||||
ScriptCache, get_cache, set_cache,
|
||||
};
|
||||
use litellm_cache_redis::{
|
||||
RedisArg, RedisCache, RedisLpopOperation, RedisLpopResult, RedisRpushOperation,
|
||||
};
|
||||
use redis_test::{MockCmd, MockRedisConnection};
|
||||
|
||||
struct TaggedByteCodec(u8);
|
||||
|
||||
impl CacheCodec for TaggedByteCodec {
|
||||
type Value = u8;
|
||||
|
||||
fn encode(&self, value: &u8) -> Result<Vec<u8>, Error> {
|
||||
if *value > 127 {
|
||||
return Err(Error::InvalidEntry);
|
||||
}
|
||||
Ok(vec![self.0, *value])
|
||||
}
|
||||
|
||||
fn decode(&self, bytes: &[u8]) -> Result<u8, Error> {
|
||||
match bytes {
|
||||
[tag, value] if *tag == self.0 => Ok(*value),
|
||||
_ => Err(Error::InvalidEntry),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn constructor_rejects_invalid_urls() {
|
||||
assert!(RedisCache::new("not a redis url", None).is_err());
|
||||
assert!(RedisCache::new("not a redis url", None, JsonCodec::<String>::new()).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn generic_helpers_use_the_injected_codec_and_ttl() {
|
||||
let connection = MockRedisConnection::new([
|
||||
MockCmd::new(
|
||||
redis::cmd("SETEX")
|
||||
.arg("counter")
|
||||
.arg(2)
|
||||
.arg([42u8, 7].as_slice()),
|
||||
Ok("OK"),
|
||||
),
|
||||
MockCmd::new(redis::cmd("GET").arg("counter"), Ok(vec![42u8, 7])),
|
||||
])
|
||||
.assert_all_commands_consumed();
|
||||
let cache = RedisCache::with_connection(connection, None, TaggedByteCodec(42));
|
||||
let context = ExactCacheContext {
|
||||
ttl: Some(Duration::from_millis(1500)),
|
||||
};
|
||||
set_cache(&cache, "counter", 7, &context).unwrap();
|
||||
assert_eq!(get_cache(&cache, "counter", &context).unwrap(), Some(7));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn async_operations_preserve_codec_ttl_and_missing_values() {
|
||||
let connection = MockRedisConnection::new([
|
||||
MockCmd::new(
|
||||
redis::cmd("SETEX")
|
||||
.arg("counter")
|
||||
.arg(9)
|
||||
.arg([42u8, 7].as_slice()),
|
||||
Ok("OK"),
|
||||
),
|
||||
MockCmd::new(redis::cmd("GET").arg("counter"), Ok(vec![42u8, 7])),
|
||||
MockCmd::new(
|
||||
redis::cmd("SETEX")
|
||||
.arg("batch")
|
||||
.arg(2)
|
||||
.arg([42u8, 8].as_slice()),
|
||||
Ok("OK"),
|
||||
),
|
||||
MockCmd::new(redis::cmd("DEL").arg("counter"), Ok(1u32)),
|
||||
MockCmd::new(redis::cmd("GET").arg("counter"), Ok(redis::Value::Nil)),
|
||||
])
|
||||
.assert_all_commands_consumed();
|
||||
let cache = RedisCache::with_connection(
|
||||
connection,
|
||||
Some(Duration::from_secs(9)),
|
||||
TaggedByteCodec(42),
|
||||
);
|
||||
let context = ExactCacheContext::default();
|
||||
cache
|
||||
.batch_cache_write("counter", 7, context.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
cache.async_get_cache("counter", &context).await.unwrap(),
|
||||
Some(7)
|
||||
);
|
||||
cache
|
||||
.async_set_cache_pipeline(
|
||||
vec![("batch".into(), 8)],
|
||||
ExactCacheContext {
|
||||
ttl: Some(Duration::from_millis(1500)),
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
cache.async_delete_cache("counter").await.unwrap();
|
||||
assert_eq!(
|
||||
cache.async_get_cache("counter", &context).await.unwrap(),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn codec_errors_propagate_without_writing_partial_batches() {
|
||||
let connection = MockRedisConnection::new([
|
||||
MockCmd::new(redis::cmd("GET").arg("invalid"), Ok(vec![99u8, 7])),
|
||||
MockCmd::new(redis::cmd("GET").arg("invalid"), Ok(vec![99u8, 7])),
|
||||
])
|
||||
.assert_all_commands_consumed();
|
||||
let cache = RedisCache::with_connection(connection, None, TaggedByteCodec(42));
|
||||
let context = ExactCacheContext::default();
|
||||
assert_eq!(
|
||||
cache.set_cache("invalid", 255, &context),
|
||||
Err(Error::InvalidEntry)
|
||||
);
|
||||
assert_eq!(
|
||||
cache.async_set_cache("invalid", 255, context.clone()).await,
|
||||
Err(Error::InvalidEntry)
|
||||
);
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_set_cache_pipeline(
|
||||
vec![("valid".into(), 7), ("invalid".into(), 255)],
|
||||
context.clone(),
|
||||
)
|
||||
.await,
|
||||
Err(Error::InvalidEntry)
|
||||
);
|
||||
assert_eq!(
|
||||
cache.get_cache("invalid", &context),
|
||||
Err(Error::InvalidEntry)
|
||||
);
|
||||
assert_eq!(
|
||||
cache.async_get_cache("invalid", &context).await,
|
||||
Err(Error::InvalidEntry)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn namespaces_are_optional_and_existing_prefixes_are_not_duplicated() {
|
||||
let connection = MockRedisConnection::new([
|
||||
MockCmd::new(redis::cmd("GET").arg("team:key"), Ok(redis::Value::Nil)),
|
||||
MockCmd::new(redis::cmd("GET").arg("team:key"), Ok(redis::Value::Nil)),
|
||||
])
|
||||
.assert_all_commands_consumed();
|
||||
let cache = RedisCache::with_connection(connection, None, JsonCodec::<String>::new())
|
||||
.with_namespace(Some("team".into()));
|
||||
assert_eq!(
|
||||
cache
|
||||
.get_cache("key", &ExactCacheContext::default())
|
||||
.unwrap(),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
cache
|
||||
.get_cache("team:key", &ExactCacheContext::default())
|
||||
.unwrap(),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn flush_requires_a_namespace_and_escapes_glob_metacharacters() {
|
||||
let unscoped = RedisCache::with_connection(
|
||||
MockRedisConnection::new([]).assert_all_commands_consumed(),
|
||||
None,
|
||||
JsonCodec::<String>::new(),
|
||||
);
|
||||
assert_eq!(unscoped.flush_cache(), Err(Error::UnscopedFlush));
|
||||
let connection = MockRedisConnection::new([
|
||||
MockCmd::new(
|
||||
redis::cmd("SCAN")
|
||||
.cursor_arg(0)
|
||||
.arg("MATCH")
|
||||
.arg("team\\*:*")
|
||||
.arg("COUNT")
|
||||
.arg(1000),
|
||||
Ok(redis_test::redis_value!(["0", ["team*:key"]])),
|
||||
),
|
||||
MockCmd::new(redis::cmd("DEL").arg("team*:key"), Ok(1u32)),
|
||||
])
|
||||
.assert_all_commands_consumed();
|
||||
let scoped = RedisCache::with_connection(connection, None, JsonCodec::<String>::new())
|
||||
.with_namespace(Some("team*".into()));
|
||||
scoped.flush_cache().unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn connection_failures_use_the_python_result_contract() {
|
||||
let error = redis::RedisError::from((redis::ErrorKind::Io, "connection refused"));
|
||||
let connection =
|
||||
MockRedisConnection::new([MockCmd::new(redis::cmd("PING"), Err::<String, _>(error))])
|
||||
.assert_all_commands_consumed();
|
||||
let cache = RedisCache::with_connection(connection, None, JsonCodec::<String>::new());
|
||||
|
||||
let result = cache.test_connection().await.unwrap();
|
||||
assert_eq!(result.status, CacheConnectionStatus::Failed);
|
||||
assert!(result.message.starts_with("Redis connection failed:"));
|
||||
assert!(result.error.is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn batch_reads_keep_order_and_treat_invalid_values_as_invalid_entries() {
|
||||
let connection = MockRedisConnection::new([MockCmd::new(
|
||||
redis::cmd("MGET").arg("hit").arg("miss").arg("invalid"),
|
||||
Ok(vec![
|
||||
redis::Value::BulkString(vec![42, 7]),
|
||||
redis::Value::Nil,
|
||||
redis::Value::BulkString(vec![99, 7]),
|
||||
]),
|
||||
)])
|
||||
.assert_all_commands_consumed();
|
||||
let cache = RedisCache::with_connection(connection, None, TaggedByteCodec(42));
|
||||
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_batch_get_cache(
|
||||
vec!["hit".into(), "miss".into(), "invalid".into()],
|
||||
ExactCacheContext::default(),
|
||||
)
|
||||
.await
|
||||
.unwrap(),
|
||||
vec![BatchEntry::Hit(7), BatchEntry::Miss, BatchEntry::Invalid]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn async_flush_deletes_each_scan_page_separately() {
|
||||
let connection = MockRedisConnection::new([
|
||||
MockCmd::new(
|
||||
redis::cmd("SCAN")
|
||||
.cursor_arg(0)
|
||||
.arg("MATCH")
|
||||
.arg("team:*")
|
||||
.arg("COUNT")
|
||||
.arg(1000),
|
||||
Ok(redis_test::redis_value!(["7", ["team:a", "team:b"]])),
|
||||
),
|
||||
MockCmd::new(redis::cmd("DEL").arg("team:a").arg("team:b"), Ok(2u32)),
|
||||
MockCmd::new(
|
||||
redis::cmd("SCAN")
|
||||
.cursor_arg(7)
|
||||
.arg("MATCH")
|
||||
.arg("team:*")
|
||||
.arg("COUNT")
|
||||
.arg(1000),
|
||||
Ok(redis_test::redis_value!(["0", ["team:c"]])),
|
||||
),
|
||||
MockCmd::new(redis::cmd("DEL").arg("team:c"), Ok(1u32)),
|
||||
])
|
||||
.assert_all_commands_consumed();
|
||||
let cache = RedisCache::with_connection(connection, None, JsonCodec::<String>::new())
|
||||
.with_namespace(Some("team".into()));
|
||||
|
||||
cache.async_flush_cache().await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn direct_redis_operations_preserve_namespace_values_and_missing_ttls() {
|
||||
let mut sadd_pipeline = redis::pipe();
|
||||
sadd_pipeline
|
||||
.cmd("SADD")
|
||||
.arg("team:members")
|
||||
.arg("a")
|
||||
.arg("b")
|
||||
.cmd("EXPIRE")
|
||||
.arg("team:members")
|
||||
.arg(600u64)
|
||||
.ignore();
|
||||
let connection = MockRedisConnection::new([
|
||||
MockCmd::new(
|
||||
redis::cmd("MGET").arg("team:count").arg("team:missing"),
|
||||
Ok(redis_test::redis_value!(["7", nil])),
|
||||
),
|
||||
MockCmd::new(
|
||||
redis::cmd("MGET").arg("team:count").arg("team:missing"),
|
||||
Ok(redis_test::redis_value!(["7", nil])),
|
||||
),
|
||||
MockCmd::new(redis::cmd("PING"), Ok("PONG")),
|
||||
MockCmd::new(redis::cmd("PING"), Ok("PONG")),
|
||||
MockCmd::new(redis::cmd("TTL").arg("team:missing"), Ok(-2i64)),
|
||||
MockCmd::new(
|
||||
redis::cmd("SCAN")
|
||||
.cursor_arg(0)
|
||||
.arg("MATCH")
|
||||
.arg("team:job-*")
|
||||
.arg("COUNT")
|
||||
.arg(25),
|
||||
Ok(redis_test::redis_value!(["4", ["team:job-a"]])),
|
||||
),
|
||||
MockCmd::new(
|
||||
redis::cmd("SCAN")
|
||||
.cursor_arg(4)
|
||||
.arg("MATCH")
|
||||
.arg("team:job-*")
|
||||
.arg("COUNT")
|
||||
.arg(25),
|
||||
Ok(redis_test::redis_value!(["0", ["team:job-b"]])),
|
||||
),
|
||||
MockCmd::new(
|
||||
redis::cmd("DEL").arg("team:job-a").arg("team:job-b"),
|
||||
Ok(2u32),
|
||||
),
|
||||
MockCmd::with_values(
|
||||
sadd_pipeline,
|
||||
Ok(vec![redis::Value::Int(2), redis::Value::Int(1)]),
|
||||
),
|
||||
MockCmd::new(
|
||||
redis::cmd("RPUSH").arg("team:queue").arg("a").arg("b"),
|
||||
Ok(2u32),
|
||||
),
|
||||
MockCmd::new(
|
||||
redis::cmd("LPOP").arg("team:queue").arg(2usize),
|
||||
Ok(redis_test::redis_value!(["a", "b"])),
|
||||
),
|
||||
MockCmd::new(
|
||||
redis::cmd("EVAL")
|
||||
.arg("return KEYS[1]")
|
||||
.arg(1usize)
|
||||
.arg("team:key"),
|
||||
Ok("team:key"),
|
||||
),
|
||||
MockCmd::new(
|
||||
redis::cmd("EVAL")
|
||||
.arg("return KEYS[1]")
|
||||
.arg(1usize)
|
||||
.arg("team:key"),
|
||||
Ok("team:key"),
|
||||
),
|
||||
MockCmd::new(redis::cmd("CLIENT").arg("LIST"), Ok("id=1")),
|
||||
MockCmd::new(redis::cmd("INFO"), Ok("redis_version:7")),
|
||||
MockCmd::new(redis::cmd("FLUSHALL"), Ok("OK")),
|
||||
])
|
||||
.assert_all_commands_consumed();
|
||||
let cache = RedisCache::with_connection(connection, None, JsonCodec::<String>::new())
|
||||
.with_namespace(Some("team".into()));
|
||||
|
||||
assert_eq!(
|
||||
cache
|
||||
.batch_get_counts(&["count".into(), "missing".into()])
|
||||
.unwrap(),
|
||||
[Some(7), None]
|
||||
);
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_batch_get_counts(vec!["count".into(), "missing".into()])
|
||||
.await
|
||||
.unwrap(),
|
||||
[Some(7), None]
|
||||
);
|
||||
assert!(cache.sync_ping().unwrap());
|
||||
assert!(cache.ping().await.unwrap());
|
||||
assert_eq!(cache.async_get_ttl("missing").await.unwrap(), None);
|
||||
assert_eq!(
|
||||
cache.async_scan_iter("job-", 25).await.unwrap(),
|
||||
["team:job-a", "team:job-b"]
|
||||
);
|
||||
assert_eq!(
|
||||
cache
|
||||
.delete_cache_keys(vec!["job-a".into(), "job-b".into()])
|
||||
.await
|
||||
.unwrap(),
|
||||
2
|
||||
);
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_set_cache_sadd("members", vec!["a".into(), "b".into()], None)
|
||||
.await
|
||||
.unwrap(),
|
||||
2
|
||||
);
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_rpush("queue", vec!["a".into(), "b".into()])
|
||||
.await
|
||||
.unwrap(),
|
||||
2
|
||||
);
|
||||
assert_eq!(
|
||||
cache.async_lpop("queue", Some(2)).await.unwrap(),
|
||||
RedisLpopResult::Values(vec![b"a".to_vec(), b"b".to_vec()])
|
||||
);
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_eval("return KEYS[1]".into(), vec!["key".into()], Vec::new())
|
||||
.await
|
||||
.unwrap(),
|
||||
redis::Value::BulkString(b"team:key".to_vec())
|
||||
);
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_register_script("return KEYS[1]".into())
|
||||
.invoke(vec!["key".into()], Vec::new())
|
||||
.await
|
||||
.unwrap(),
|
||||
redis::Value::BulkString(b"team:key".to_vec())
|
||||
);
|
||||
assert_eq!(cache.client_list().unwrap(), "id=1");
|
||||
assert_eq!(cache.info().unwrap(), "redis_version:7");
|
||||
cache.flushall().unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn direct_redis_pipelines_preserve_operation_order() {
|
||||
let mut rpush_pipeline = redis::pipe();
|
||||
rpush_pipeline
|
||||
.cmd("RPUSH")
|
||||
.arg("team:a")
|
||||
.arg("one")
|
||||
.cmd("RPUSH")
|
||||
.arg("team:b")
|
||||
.arg("two");
|
||||
let mut lpop_pipeline = redis::pipe();
|
||||
lpop_pipeline
|
||||
.cmd("LPOP")
|
||||
.arg("team:a")
|
||||
.arg(2usize)
|
||||
.cmd("LPOP")
|
||||
.arg("team:b");
|
||||
let connection = MockRedisConnection::new([
|
||||
MockCmd::with_values(
|
||||
rpush_pipeline,
|
||||
Ok(vec![redis::Value::Int(1), redis::Value::Int(2)]),
|
||||
),
|
||||
MockCmd::with_values(
|
||||
lpop_pipeline,
|
||||
Ok(vec![redis_test::redis_value!(["one"]), redis::Value::Nil]),
|
||||
),
|
||||
])
|
||||
.assert_all_commands_consumed();
|
||||
let queue = RedisCache::with_connection(connection, None, JsonCodec::<String>::new())
|
||||
.with_namespace(Some("team".into()));
|
||||
|
||||
assert_eq!(
|
||||
queue
|
||||
.async_rpush_pipeline(vec![
|
||||
RedisRpushOperation {
|
||||
key: "a".into(),
|
||||
values: vec![RedisArg::from("one")],
|
||||
},
|
||||
RedisRpushOperation {
|
||||
key: "b".into(),
|
||||
values: vec![RedisArg::from("two")],
|
||||
},
|
||||
])
|
||||
.await
|
||||
.unwrap(),
|
||||
[1, 2]
|
||||
);
|
||||
assert_eq!(
|
||||
queue
|
||||
.async_lpop_pipeline(vec![
|
||||
RedisLpopOperation {
|
||||
key: "a".into(),
|
||||
count: Some(2),
|
||||
},
|
||||
RedisLpopOperation {
|
||||
key: "b".into(),
|
||||
count: None,
|
||||
},
|
||||
])
|
||||
.await
|
||||
.unwrap(),
|
||||
[
|
||||
RedisLpopResult::Values(vec![b"one".to_vec()]),
|
||||
RedisLpopResult::Missing,
|
||||
]
|
||||
);
|
||||
|
||||
let mut increment_pipeline = redis::pipe();
|
||||
increment_pipeline
|
||||
.cmd("INCRBYFLOAT")
|
||||
.arg("team:counter")
|
||||
.arg(1.5f64)
|
||||
.cmd("EXPIRE")
|
||||
.arg("team:counter")
|
||||
.arg(10u64)
|
||||
.ignore()
|
||||
.cmd("INCRBYFLOAT")
|
||||
.arg("team:counter")
|
||||
.arg(2.0f64);
|
||||
let connection = MockRedisConnection::new([MockCmd::with_values(
|
||||
increment_pipeline,
|
||||
Ok(vec![
|
||||
redis::Value::BulkString(b"1.5".to_vec()),
|
||||
redis::Value::Int(1),
|
||||
redis::Value::BulkString(b"3.5".to_vec()),
|
||||
]),
|
||||
)])
|
||||
.assert_all_commands_consumed();
|
||||
let counters = RedisCache::with_connection(connection, None, JsonCodec::<f64>::new())
|
||||
.with_namespace(Some("team".into()));
|
||||
assert_eq!(
|
||||
counters
|
||||
.async_increment_pipeline(vec![
|
||||
IncrementOperation {
|
||||
key: "counter".into(),
|
||||
amount: 1.5,
|
||||
ttl: Some(Duration::from_secs(10)),
|
||||
},
|
||||
IncrementOperation {
|
||||
key: "counter".into(),
|
||||
amount: 2.0,
|
||||
ttl: None,
|
||||
},
|
||||
])
|
||||
.await
|
||||
.unwrap(),
|
||||
[1.5, 3.5]
|
||||
);
|
||||
}
|
||||
|
||||
const INCREMENT_WITH_FLOOR_SCRIPT: &str = concat!(
|
||||
"local count = redis.call('INCRBY', KEYS[1], ARGV[1]); ",
|
||||
"if count < 0 then count = redis.call('INCRBY', KEYS[1], -count); end; ",
|
||||
"if redis.call('TTL', KEYS[1]) < 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; ",
|
||||
"return count"
|
||||
);
|
||||
const SET_MAX_SCRIPT: &str = concat!(
|
||||
"local current = redis.call('GET', KEYS[1]); ",
|
||||
"if current == false or tonumber(current) < tonumber(ARGV[1]) then ",
|
||||
"redis.call('SET', KEYS[1], ARGV[1]); ",
|
||||
"if tonumber(ARGV[2]) > 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; ",
|
||||
"return ARGV[1]; end; return current"
|
||||
);
|
||||
|
||||
#[tokio::test]
|
||||
async fn counter_repairs_are_atomic_and_use_default_ttl() {
|
||||
let floor = || {
|
||||
redis::cmd("EVAL")
|
||||
.arg(INCREMENT_WITH_FLOOR_SCRIPT)
|
||||
.arg(1)
|
||||
.arg("team:counter")
|
||||
.arg(-2i64)
|
||||
.arg(30u64)
|
||||
.clone()
|
||||
};
|
||||
let connection = MockRedisConnection::new([
|
||||
MockCmd::new(floor(), Ok(0i64)),
|
||||
MockCmd::new(floor(), Ok(0i64)),
|
||||
MockCmd::new(
|
||||
redis::cmd("EVAL")
|
||||
.arg(SET_MAX_SCRIPT)
|
||||
.arg(1)
|
||||
.arg("team:counter")
|
||||
.arg(4.5f64)
|
||||
.arg(600u64),
|
||||
Ok("4.5"),
|
||||
),
|
||||
])
|
||||
.assert_all_commands_consumed();
|
||||
let cache = RedisCache::with_connection(connection, None, JsonCodec::<f64>::new())
|
||||
.with_namespace(Some("team".into()));
|
||||
|
||||
assert_eq!(
|
||||
cache
|
||||
.increment_with_floor("counter", -2, Duration::from_secs(30))
|
||||
.unwrap(),
|
||||
0
|
||||
);
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_increment_with_floor("counter", -2, Duration::from_secs(30))
|
||||
.await
|
||||
.unwrap(),
|
||||
0
|
||||
);
|
||||
assert_eq!(
|
||||
cache.async_set_max("counter", 4.5, None).await.unwrap(),
|
||||
4.5
|
||||
);
|
||||
}
|
||||
|
||||
const CLAIM_SCRIPT: &str = concat!(
|
||||
"local current = redis.call('GET', KEYS[1]); ",
|
||||
"if ARGV[1] == '' then if current ~= false and current ~= '' then return 0; end; ",
|
||||
"elseif current ~= ARGV[1] then return 0; end; ",
|
||||
"if ARGV[3] ~= '' then redis.call('SET', KEYS[1], ARGV[3], 'EX', ARGV[2]); ",
|
||||
"elseif ARGV[4] == '1' then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return 1"
|
||||
);
|
||||
|
||||
fn claim_eval(expected: &str, write: &str, refresh: bool) -> redis::Cmd {
|
||||
let mut cmd = redis::cmd("EVAL");
|
||||
cmd.arg(CLAIM_SCRIPT)
|
||||
.arg(1)
|
||||
.arg("pin")
|
||||
.arg(expected)
|
||||
.arg(600)
|
||||
.arg(write)
|
||||
.arg(u8::from(refresh));
|
||||
cmd
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn claims_match_eligible_values_written_by_another_encoder() {
|
||||
let python_payload = r#"{"model_id": "a", "deployment": "east"}"#;
|
||||
let stored = serde_json::json!({"deployment": "east", "model_id": "a"});
|
||||
let candidate = serde_json::json!({"model_id": "b"});
|
||||
let connection = MockRedisConnection::new([
|
||||
MockCmd::new(redis::cmd("GET").arg("pin"), Ok(python_payload)),
|
||||
MockCmd::new(claim_eval(python_payload, "", true), Ok(1)),
|
||||
])
|
||||
.assert_all_commands_consumed();
|
||||
let cache =
|
||||
RedisCache::with_connection(connection, None, JsonCodec::<serde_json::Value>::new());
|
||||
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_claim_cache(
|
||||
"pin",
|
||||
candidate,
|
||||
vec![stored.clone()],
|
||||
ExactCacheContext::default()
|
||||
)
|
||||
.await
|
||||
.unwrap(),
|
||||
stored
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn claims_retry_when_the_key_changes_and_replace_ineligible_winners() {
|
||||
let candidate = serde_json::json!({"model_id": "b"});
|
||||
let payload = r#"{"model_id":"b"}"#;
|
||||
let connection = MockRedisConnection::new([
|
||||
MockCmd::new(redis::cmd("GET").arg("pin"), Ok(redis::Value::Nil)),
|
||||
MockCmd::new(claim_eval("", payload, false), Ok(0)),
|
||||
MockCmd::new(redis::cmd("GET").arg("pin"), Ok(r#"{"model_id":"gone"}"#)),
|
||||
MockCmd::new(claim_eval(r#"{"model_id":"gone"}"#, payload, false), Ok(1)),
|
||||
])
|
||||
.assert_all_commands_consumed();
|
||||
let cache =
|
||||
RedisCache::with_connection(connection, None, JsonCodec::<serde_json::Value>::new());
|
||||
|
||||
assert_eq!(
|
||||
cache
|
||||
.claim_cache(
|
||||
"pin",
|
||||
candidate.clone(),
|
||||
&[serde_json::json!({"model_id": "a"})],
|
||||
ExactCacheContext::default()
|
||||
)
|
||||
.unwrap(),
|
||||
candidate
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn claims_without_eligible_values_keep_the_winner_without_refreshing_its_ttl() {
|
||||
let stored = r#"{"model_id": "a"}"#;
|
||||
let connection = MockRedisConnection::new([
|
||||
MockCmd::new(redis::cmd("GET").arg("pin"), Ok(stored)),
|
||||
MockCmd::new(claim_eval(stored, "", false), Ok(1)),
|
||||
])
|
||||
.assert_all_commands_consumed();
|
||||
let cache =
|
||||
RedisCache::with_connection(connection, None, JsonCodec::<serde_json::Value>::new());
|
||||
|
||||
assert_eq!(
|
||||
cache
|
||||
.claim_cache(
|
||||
"pin",
|
||||
serde_json::json!({"model_id": "b"}),
|
||||
&[],
|
||||
ExactCacheContext::default()
|
||||
)
|
||||
.unwrap(),
|
||||
serde_json::json!({"model_id": "a"})
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn async_increment_runs_the_atomic_script() {
|
||||
let mut eval = redis::cmd("EVAL");
|
||||
eval.arg(concat!(
|
||||
"local value = redis.call('INCRBYFLOAT', KEYS[1], ARGV[1]); ",
|
||||
"if redis.call('TTL', KEYS[1]) == -1 then ",
|
||||
"redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return value"
|
||||
))
|
||||
.arg(1)
|
||||
.arg("counter")
|
||||
.arg(2.5f64)
|
||||
.arg(600);
|
||||
let connection =
|
||||
MockRedisConnection::new([MockCmd::new(eval, Ok("4.5"))]).assert_all_commands_consumed();
|
||||
let cache = RedisCache::with_connection(connection, None, JsonCodec::<f64>::new());
|
||||
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_increment("counter", 2.5, ExactCacheContext::default())
|
||||
.await
|
||||
.unwrap(),
|
||||
4.5
|
||||
);
|
||||
}
|
||||
|
|
|
|||
20
litellm-rust/crates/cache-response/Cargo.toml
Normal file
20
litellm-rust/crates/cache-response/Cargo.toml
Normal file
|
|
@ -0,0 +1,20 @@
|
|||
[package]
|
||||
name = "litellm-cache-response"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
litellm-cache.workspace = true
|
||||
py_literal = "0.4.0"
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
sha2.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
litellm-cache-memory.workspace = true
|
||||
litellm-cache-redis.workspace = true
|
||||
redis = "1.7.0"
|
||||
redis-test = "1.0.4"
|
||||
tokio.workspace = true
|
||||
61
litellm-rust/crates/cache-response/README.md
Normal file
61
litellm-rust/crates/cache-response/README.md
Normal file
|
|
@ -0,0 +1,61 @@
|
|||
# Response cache foundation
|
||||
|
||||
`ResponseCache<B>` adds request keys, independent read/write controls, response envelopes, and freshness checks to any `B: BaseCache<Value = CacheEntry>`
|
||||
|
||||
## Ownership
|
||||
|
||||
`litellm-cache` defines typed storage and codec traits. Memory and Redis implement those traits without depending on response policy. Other consumers can store their own value types using the same backend implementations
|
||||
|
||||
`litellm-cache-response` owns response keys, controls, entries, the Python-compatible response codec, and `WriteBuffer`, the backend-neutral deferred-write policy. It has no runtime dependency on a specific cache backend or Python
|
||||
|
||||
The Python bridge constructs backends and selects them through its private `NativeResponseCache` enum, which only dispatches. Generic Rust callers inject their backend directly. A native gateway can construct the same generic response service in its own host
|
||||
|
||||
## Native Rust use
|
||||
|
||||
```rust
|
||||
use std::{sync::Arc, time::Duration};
|
||||
use litellm_cache_memory::InMemoryCache;
|
||||
use litellm_cache_response::{CacheKeyInput, ResponseCache, ResponseCacheRequest};
|
||||
use serde_json::json;
|
||||
|
||||
let cache = ResponseCache::new(Arc::new(InMemoryCache::default()));
|
||||
let request = ResponseCacheRequest::new(CacheKeyInput {
|
||||
preset: Some("example:key".into()),
|
||||
..Default::default()
|
||||
});
|
||||
let now = Duration::from_secs(100);
|
||||
cache.store(&request, json!({"answer": 7}), now)?;
|
||||
assert_eq!(cache.async_lookup(&request, now).await?, Some(json!({"answer": 7})));
|
||||
```
|
||||
|
||||
For Redis, inject `RedisCache::new(url, ttl, ResponseCacheCodec)` instead. Namespaces are optional and existing namespace prefixes are preserved. Sync operations check out independent connections from a bounded pool, while async callers, including counters and claims, move that blocking work off the executor. The pool skips the checkout PING and instead discards any connection whose command failed
|
||||
|
||||
Callers supply Unix time for response freshness. Backend TTL uses its own clock. A read can reject an entry through `max_age` even while the backend still retains it
|
||||
|
||||
## Python integration boundary
|
||||
|
||||
The extension keeps a private test harness for memory and Redis single and batch response lookup and storage. Batch lookup returns ordered values plus missing indices for embedding partial-hit wiring. No bridge-only cache type is part of the public API
|
||||
|
||||
Object responses are written as they are, and every other response shape is written as a serialized string, which is the pair of shapes Python reads. A string on the wire is therefore always a serialized response, so string-valued responses round trip. Typed backends such as memory never pass through the codec
|
||||
|
||||
The resolver reads the namespace's `cache` attribute each time it resolves. A captured binding retains the selected service for its operation, including background writes. `None` disables caching. Custom Python cache objects keep their original methods, arguments, returned awaitables, exceptions, and caller-task execution
|
||||
|
||||
Python callbacks use the built-in `Cache` API, so a `Cache` subclass works unchanged. A batch lookup takes one original kwargs mapping per request and returns the list of `get_cache` or gathered `async_get_cache` results, while native bindings return `{values, missing_indices}`. A batch store hands the caller's original result to `async_add_cache_pipeline`. `ping` calls `ping`, and a flush goes to the facade's backend
|
||||
|
||||
The private facade test harness checks object identity, method overrides, effective TTL, Redis namespace, memory capacity, and later configuration changes before selecting native execution. Its snapshot includes Redis connection settings, so a later `redis_kwargs` change, including an SSL option, selects Python callback execution. Buffered async writes honor `redis_flush_size`. Public activation must construct the shared native service from the initial Python Redis settings, including `litellm.default_redis_ttl` and SSL options. A buffered entry keeps the time it was produced, and a failed flush drops its batch instead of growing the buffer during an outage. The harness does not migrate entries or replace Python methods. Until activation configures one shared service, the Python facade and native test service can hold separate data. Existing public cache constructors remain on Python
|
||||
|
||||
Native cache handles must be recreated after fork. The bridge releases the GIL around native operations, and Redis runs blocking connection operations off the async executor. Native errors propagate to the host, which owns the existing fail-open and logging policy
|
||||
|
||||
The Redis backend also provides the primitives needed to preserve its direct Python surface later: TLS URLs, ping, bulk delete, counter batches, TTL, scan, set membership, raw queue push and pop, queue and counter pipelines, counter floor and maximum operations, script evaluation, client information, namespaced flush, and full flush. These are backend operations only and are not exported to Python by this PR. Memory provides TTL, oldest-key, and counter-pipeline operations
|
||||
|
||||
## Adding another backend
|
||||
|
||||
Implement `BaseCache` for the backend with its associated value type, and accept a `CacheCodec` when wire serialization is needed. `ResponseCache<B>` then works without another response implementation. Add a concrete bridge enum variant and constructor only when exposing that backend to Python
|
||||
|
||||
Verify typed values, TTL precedence, missing entries, serialization failures, namespaces, batch ordering, and sync/async behavior. Run response fixtures with `ResponseCacheCodec`, including both Python envelope encodings, before enabling a public facade
|
||||
|
||||
## Follow-up scope
|
||||
|
||||
Public SDK, Router, and proxy activation still need constructor parity, stream replay, embedding partial-batch integration, response reconstruction, callback scheduling, and failure-policy integration. This foundation does not switch those request paths
|
||||
|
||||
Redis cluster, disk, cloud stores, and semantic caching remain follow-ups. The generic dual cache takes read, write, and remote-failure policies, runs its async operations through the async L2 methods, and provides L2-first counters and atomic affinity claims. Errors propagate by default, and `RemoteFailurePolicy::UseLocal` opts key-value operations and claims into the local tier when L2 is unavailable. Claims compare decoded values, so a pin written by Python still matches. Public Router integration remains follow-up work. Reservations and pubsub still need explicit capabilities owned by their consuming features. Adding a cache backend does not establish those guarantees
|
||||
45
litellm-rust/crates/cache-response/src/buffer.rs
Normal file
45
litellm-rust/crates/cache-response/src/buffer.rs
Normal file
|
|
@ -0,0 +1,45 @@
|
|||
use std::{sync::Mutex, time::Duration};
|
||||
|
||||
use litellm_cache::{BaseCache, Error, ExactCacheContext};
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{CacheEntry, ResponseCache, ResponseCacheRequest};
|
||||
|
||||
pub struct WriteBuffer {
|
||||
flush_size: usize,
|
||||
entries: Mutex<Vec<(ResponseCacheRequest, Value, Duration)>>,
|
||||
}
|
||||
|
||||
impl WriteBuffer {
|
||||
pub fn new(flush_size: usize) -> Self {
|
||||
Self {
|
||||
flush_size: flush_size.max(1),
|
||||
entries: Mutex::new(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn async_store<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>>(
|
||||
&self,
|
||||
cache: &ResponseCache<B>,
|
||||
request: &ResponseCacheRequest,
|
||||
response: Value,
|
||||
now: Duration,
|
||||
) -> Result<(), Error> {
|
||||
let pending = {
|
||||
let mut entries = self.entries.lock().map_err(|_| Error::Unavailable)?;
|
||||
entries.push((request.clone(), response, now));
|
||||
(entries.len() >= self.flush_size).then(|| std::mem::take(&mut *entries))
|
||||
};
|
||||
// A failed flush drops its batch, as Python does. Requeueing would grow the
|
||||
// buffer and re-send an ever larger pipeline on every write during an outage.
|
||||
match pending {
|
||||
Some(pending) => cache.async_store_entries(pending).await,
|
||||
None => Ok(()),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn clear(&self) -> Result<(), Error> {
|
||||
self.entries.lock().map_err(|_| Error::Unavailable)?.clear();
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
147
litellm-rust/crates/cache-response/src/caching.rs
Normal file
147
litellm-rust/crates/cache-response/src/caching.rs
Normal file
|
|
@ -0,0 +1,147 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize, PartialEq, Eq)]
|
||||
pub enum CacheMode {
|
||||
#[default]
|
||||
#[serde(rename = "default_on")]
|
||||
DefaultOn,
|
||||
#[serde(rename = "default_off")]
|
||||
DefaultOff,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Deserialize, Serialize)]
|
||||
pub struct CacheKeyField {
|
||||
pub name: String,
|
||||
pub value: Option<String>,
|
||||
pub api_parameter: bool,
|
||||
pub internal_parameter: bool,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default, Deserialize, Serialize)]
|
||||
#[serde(default)]
|
||||
pub struct CacheKeyInput {
|
||||
pub fields: Vec<CacheKeyField>,
|
||||
pub preset: Option<String>,
|
||||
pub namespace: Option<String>,
|
||||
pub include_provider_parameters: bool,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
pub struct CacheKeyContext {
|
||||
pub model_group: Option<String>,
|
||||
pub caching_groups: Vec<(Vec<String>, String)>,
|
||||
pub file_checksum: Option<String>,
|
||||
pub file_object_name: Option<String>,
|
||||
pub metadata_file_name: Option<String>,
|
||||
pub parameters_file_name: Option<String>,
|
||||
}
|
||||
|
||||
impl CacheKeyContext {
|
||||
pub fn apply(self, input: &mut CacheKeyInput) {
|
||||
let group = self.model_group.as_ref().and_then(|model| {
|
||||
self.caching_groups
|
||||
.iter()
|
||||
.find(|(models, _)| models.contains(model))
|
||||
});
|
||||
for field in &mut input.fields {
|
||||
match field.name.as_str() {
|
||||
"model" => {
|
||||
field.value = group
|
||||
.map(|(_, formatted)| formatted.clone())
|
||||
.or_else(|| self.model_group.clone())
|
||||
.or_else(|| field.value.take())
|
||||
}
|
||||
"file" => {
|
||||
field.value = self
|
||||
.file_checksum
|
||||
.clone()
|
||||
.or_else(|| self.file_object_name.clone())
|
||||
.or_else(|| self.metadata_file_name.clone())
|
||||
.or_else(|| self.parameters_file_name.clone())
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn get_cache_key(input: &CacheKeyInput) -> String {
|
||||
cache_key(input)
|
||||
}
|
||||
|
||||
pub fn cache_key(input: &CacheKeyInput) -> String {
|
||||
if let Some(preset) = &input.preset {
|
||||
return preset.clone();
|
||||
}
|
||||
let mut digest = Sha256::new();
|
||||
for field in &input.fields {
|
||||
if (field.api_parameter || (input.include_provider_parameters && !field.internal_parameter))
|
||||
&& let Some(value) = &field.value
|
||||
{
|
||||
digest.update(field.name.as_bytes());
|
||||
digest.update(b": ");
|
||||
digest.update(value.as_bytes());
|
||||
}
|
||||
}
|
||||
let hash = format!("{:x}", digest.finalize());
|
||||
input
|
||||
.namespace
|
||||
.as_deref()
|
||||
.filter(|namespace| !namespace.is_empty())
|
||||
.map_or(hash.clone(), |namespace| format!("{namespace}:{hash}"))
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize)]
|
||||
pub struct CacheControls {
|
||||
pub supported_call_type: bool,
|
||||
pub configured: bool,
|
||||
pub native_backend: bool,
|
||||
pub default_on: bool,
|
||||
pub caching: Option<bool>,
|
||||
pub no_cache: bool,
|
||||
pub no_store: bool,
|
||||
#[serde(default)]
|
||||
pub use_cache: bool,
|
||||
}
|
||||
|
||||
impl CacheControls {
|
||||
pub fn reads(self) -> bool {
|
||||
self.supported_call_type
|
||||
&& self.configured
|
||||
&& self.caching.unwrap_or(true)
|
||||
&& !self.no_cache
|
||||
&& (self.default_on || self.use_cache)
|
||||
}
|
||||
|
||||
pub fn writes(self) -> bool {
|
||||
self.supported_call_type
|
||||
&& self.configured
|
||||
&& self.caching.unwrap_or(true)
|
||||
&& !self.no_store
|
||||
&& (self.default_on || self.use_cache)
|
||||
}
|
||||
}
|
||||
|
||||
pub fn should_use_cache(controls: CacheControls) -> bool {
|
||||
controls.reads() || controls.writes()
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct CacheEntry {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub timestamp: Option<f64>,
|
||||
pub response: Value,
|
||||
}
|
||||
|
||||
impl CacheEntry {
|
||||
pub fn fresh(&self, now: Duration, max_age: Option<Duration>) -> bool {
|
||||
self.timestamp.is_none_or(|timestamp| {
|
||||
timestamp.is_finite()
|
||||
&& max_age.is_none_or(|age| now.as_secs_f64() - timestamp <= age.as_secs_f64())
|
||||
})
|
||||
}
|
||||
}
|
||||
129
litellm-rust/crates/cache-response/src/codec.rs
Normal file
129
litellm-rust/crates/cache-response/src/codec.rs
Normal file
|
|
@ -0,0 +1,129 @@
|
|||
use litellm_cache::{CacheCodec, Error};
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::CacheEntry;
|
||||
|
||||
#[derive(Clone, Copy, Debug, Default)]
|
||||
pub struct ResponseCacheCodec;
|
||||
|
||||
impl CacheCodec for ResponseCacheCodec {
|
||||
type Value = CacheEntry;
|
||||
|
||||
fn encode(&self, value: &CacheEntry) -> Result<Vec<u8>, Error> {
|
||||
if value
|
||||
.timestamp
|
||||
.is_some_and(|timestamp| !timestamp.is_finite())
|
||||
{
|
||||
return Err(Error::InvalidEntry);
|
||||
}
|
||||
// Python reads a `response` that is either a dict or a serialized string, so every
|
||||
// other shape is written serialized. A string on the wire is therefore always a
|
||||
// serialized response, which keeps string-valued responses unambiguous.
|
||||
if value.timestamp.is_none() || value.response.is_object() {
|
||||
return serde_json::to_vec(value).map_err(|_| Error::InvalidEntry);
|
||||
}
|
||||
let response = serde_json::to_string(&value.response).map_err(|_| Error::InvalidEntry)?;
|
||||
serde_json::to_vec(&CacheEntry {
|
||||
timestamp: value.timestamp,
|
||||
response: Value::String(response),
|
||||
})
|
||||
.map_err(|_| Error::InvalidEntry)
|
||||
}
|
||||
|
||||
fn decode(&self, bytes: &[u8]) -> Result<CacheEntry, Error> {
|
||||
let text = std::str::from_utf8(bytes).map_err(|_| Error::InvalidEntry)?;
|
||||
let value = decode_value(text)?;
|
||||
let Some(timestamp) = value.get("timestamp") else {
|
||||
return Ok(CacheEntry {
|
||||
timestamp: None,
|
||||
response: value,
|
||||
});
|
||||
};
|
||||
let Some(timestamp) = timestamp.as_f64().filter(|timestamp| timestamp.is_finite()) else {
|
||||
return Err(Error::InvalidEntry);
|
||||
};
|
||||
let response = match value.get("response").ok_or(Error::InvalidEntry)? {
|
||||
Value::String(text) => decode_value(text)?,
|
||||
response => response.clone(),
|
||||
};
|
||||
Ok(CacheEntry {
|
||||
timestamp: Some(timestamp),
|
||||
response,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn decode_value(text: &str) -> Result<Value, Error> {
|
||||
if let Ok(value) = serde_json::from_str(text) {
|
||||
return Ok(value);
|
||||
}
|
||||
check_literal_depth(text)?;
|
||||
let literal: py_literal::Value = text.parse().map_err(|_| Error::InvalidEntry)?;
|
||||
literal_value(literal, 0)
|
||||
}
|
||||
|
||||
fn literal_value(value: py_literal::Value, depth: usize) -> Result<Value, Error> {
|
||||
use py_literal::Value as Literal;
|
||||
if depth > 128 {
|
||||
return Err(Error::InvalidEntry);
|
||||
}
|
||||
match value {
|
||||
Literal::String(text) => Ok(Value::String(text)),
|
||||
Literal::Boolean(value) => Ok(Value::Bool(value)),
|
||||
Literal::None => Ok(Value::Null),
|
||||
Literal::Integer(value) => {
|
||||
serde_json::from_str(&value.to_string()).map_err(|_| Error::InvalidEntry)
|
||||
}
|
||||
Literal::Float(value) => serde_json::Number::from_f64(value)
|
||||
.map(Value::Number)
|
||||
.ok_or(Error::InvalidEntry),
|
||||
Literal::List(values) | Literal::Tuple(values) => values
|
||||
.into_iter()
|
||||
.map(|value| literal_value(value, depth + 1))
|
||||
.collect::<Result<Vec<_>, _>>()
|
||||
.map(Value::Array),
|
||||
Literal::Dict(entries) => entries
|
||||
.into_iter()
|
||||
.map(|(key, value)| {
|
||||
let Literal::String(key) = key else {
|
||||
return Err(Error::InvalidEntry);
|
||||
};
|
||||
Ok((key, literal_value(value, depth + 1)?))
|
||||
})
|
||||
.collect::<Result<serde_json::Map<_, _>, _>>()
|
||||
.map(Value::Object),
|
||||
_ => Err(Error::InvalidEntry),
|
||||
}
|
||||
}
|
||||
|
||||
fn check_literal_depth(text: &str) -> Result<(), Error> {
|
||||
let mut quote = None;
|
||||
let mut escaped = false;
|
||||
let mut depth = 0usize;
|
||||
for ch in text.chars() {
|
||||
if escaped {
|
||||
escaped = false;
|
||||
continue;
|
||||
}
|
||||
if let Some(delimiter) = quote {
|
||||
if ch == '\\' {
|
||||
escaped = true;
|
||||
} else if ch == delimiter {
|
||||
quote = None;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
match ch {
|
||||
'\'' | '"' => quote = Some(ch),
|
||||
'[' | '{' | '(' => {
|
||||
depth += 1;
|
||||
if depth > 128 {
|
||||
return Err(Error::InvalidEntry);
|
||||
}
|
||||
}
|
||||
']' | '}' | ')' => depth = depth.saturating_sub(1),
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
22
litellm-rust/crates/cache-response/src/embedding.rs
Normal file
22
litellm-rust/crates/cache-response/src/embedding.rs
Normal file
|
|
@ -0,0 +1,22 @@
|
|||
use serde::Serialize;
|
||||
use serde_json::Value;
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize)]
|
||||
pub struct PartialHits {
|
||||
pub values: Vec<Option<Value>>,
|
||||
pub missing_indices: Vec<usize>,
|
||||
}
|
||||
|
||||
impl PartialHits {
|
||||
pub fn new(values: Vec<Option<Value>>) -> Self {
|
||||
let missing_indices = values
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter_map(|(index, value)| value.is_none().then_some(index))
|
||||
.collect();
|
||||
Self {
|
||||
values,
|
||||
missing_indices,
|
||||
}
|
||||
}
|
||||
}
|
||||
14
litellm-rust/crates/cache-response/src/lib.rs
Normal file
14
litellm-rust/crates/cache-response/src/lib.rs
Normal file
|
|
@ -0,0 +1,14 @@
|
|||
mod buffer;
|
||||
mod caching;
|
||||
mod codec;
|
||||
mod embedding;
|
||||
mod response;
|
||||
|
||||
pub use buffer::WriteBuffer;
|
||||
pub use caching::{
|
||||
CacheControls, CacheEntry, CacheKeyContext, CacheKeyField, CacheKeyInput, CacheMode, cache_key,
|
||||
get_cache_key, should_use_cache,
|
||||
};
|
||||
pub use codec::ResponseCacheCodec;
|
||||
pub use embedding::PartialHits;
|
||||
pub use response::{ResponseCache, ResponseCacheRequest};
|
||||
280
litellm-rust/crates/cache-response/src/response.rs
Normal file
280
litellm-rust/crates/cache-response/src/response.rs
Normal file
|
|
@ -0,0 +1,280 @@
|
|||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use litellm_cache::{
|
||||
BaseCache, BatchCache, BatchEntry, CacheConnectionResult, Error, ExactCacheContext, FlushCache,
|
||||
};
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{CacheControls, CacheEntry, CacheKeyInput, PartialHits, cache_key};
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct ResponseCacheRequest {
|
||||
pub key: CacheKeyInput,
|
||||
pub controls: CacheControls,
|
||||
pub context: ExactCacheContext,
|
||||
pub max_age: Option<Duration>,
|
||||
}
|
||||
|
||||
impl ResponseCacheRequest {
|
||||
pub fn new(key: CacheKeyInput) -> Self {
|
||||
Self {
|
||||
key,
|
||||
controls: CacheControls {
|
||||
configured: true,
|
||||
supported_call_type: true,
|
||||
native_backend: true,
|
||||
default_on: true,
|
||||
..Default::default()
|
||||
},
|
||||
context: ExactCacheContext::default(),
|
||||
max_age: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct ResponseCache<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> {
|
||||
backend: Arc<B>,
|
||||
}
|
||||
|
||||
impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCache<B> {
|
||||
pub fn new(backend: Arc<B>) -> Self {
|
||||
Self { backend }
|
||||
}
|
||||
|
||||
pub fn backend(&self) -> &B {
|
||||
&self.backend
|
||||
}
|
||||
|
||||
pub fn default_ttl(&self) -> Option<Duration> {
|
||||
self.backend.get_ttl(&ExactCacheContext::default())
|
||||
}
|
||||
|
||||
pub async fn async_flush(&self) -> Result<(), Error>
|
||||
where
|
||||
B: FlushCache,
|
||||
{
|
||||
self.backend.async_flush_cache().await
|
||||
}
|
||||
|
||||
pub async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
|
||||
self.backend.test_connection().await
|
||||
}
|
||||
|
||||
pub fn lookup(
|
||||
&self,
|
||||
request: &ResponseCacheRequest,
|
||||
now: Duration,
|
||||
) -> Result<Option<Value>, Error> {
|
||||
if !request.controls.reads() {
|
||||
return Ok(None);
|
||||
}
|
||||
let entry = match self
|
||||
.backend
|
||||
.get_cache(&cache_key(&request.key), &request.context)
|
||||
{
|
||||
Ok(entry) => entry,
|
||||
Err(Error::InvalidEntry) => None,
|
||||
Err(error) => return Err(error),
|
||||
};
|
||||
Ok(Self::fresh_or_miss(entry, now, request.max_age))
|
||||
}
|
||||
|
||||
pub async fn async_lookup(
|
||||
&self,
|
||||
request: &ResponseCacheRequest,
|
||||
now: Duration,
|
||||
) -> Result<Option<Value>, Error> {
|
||||
if !request.controls.reads() {
|
||||
return Ok(None);
|
||||
}
|
||||
let entry = match self
|
||||
.backend
|
||||
.async_get_cache(&cache_key(&request.key), &request.context)
|
||||
.await
|
||||
{
|
||||
Ok(entry) => entry,
|
||||
Err(Error::InvalidEntry) => None,
|
||||
Err(error) => return Err(error),
|
||||
};
|
||||
Ok(Self::fresh_or_miss(entry, now, request.max_age))
|
||||
}
|
||||
|
||||
pub fn lookup_batch(
|
||||
&self,
|
||||
requests: &[ResponseCacheRequest],
|
||||
now: Duration,
|
||||
) -> Result<PartialHits, Error>
|
||||
where
|
||||
B: BatchCache,
|
||||
{
|
||||
let readable = requests
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter(|(_, request)| request.controls.reads())
|
||||
.collect::<Vec<_>>();
|
||||
let keys = readable
|
||||
.iter()
|
||||
.map(|(_, request)| cache_key(&request.key))
|
||||
.collect::<Vec<_>>();
|
||||
let entries = if let Some((_, request)) = readable.first() {
|
||||
self.backend.batch_get_cache(&keys, &request.context)?
|
||||
} else {
|
||||
Vec::new()
|
||||
};
|
||||
Self::partial_hits(requests, readable, entries, now)
|
||||
}
|
||||
|
||||
pub async fn async_lookup_batch(
|
||||
&self,
|
||||
requests: &[ResponseCacheRequest],
|
||||
now: Duration,
|
||||
) -> Result<PartialHits, Error>
|
||||
where
|
||||
B: BatchCache,
|
||||
{
|
||||
let readable = requests
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter(|(_, request)| request.controls.reads())
|
||||
.collect::<Vec<_>>();
|
||||
let keys = readable
|
||||
.iter()
|
||||
.map(|(_, request)| cache_key(&request.key))
|
||||
.collect::<Vec<_>>();
|
||||
let entries = if let Some((_, request)) = readable.first() {
|
||||
self.backend
|
||||
.async_batch_get_cache(keys, request.context.clone())
|
||||
.await?
|
||||
} else {
|
||||
Vec::new()
|
||||
};
|
||||
Self::partial_hits(requests, readable, entries, now)
|
||||
}
|
||||
|
||||
pub fn store(
|
||||
&self,
|
||||
request: &ResponseCacheRequest,
|
||||
response: Value,
|
||||
now: Duration,
|
||||
) -> Result<(), Error> {
|
||||
if !request.controls.writes() {
|
||||
return Ok(());
|
||||
}
|
||||
self.backend.set_cache(
|
||||
&cache_key(&request.key),
|
||||
CacheEntry {
|
||||
timestamp: Some(now.as_secs_f64()),
|
||||
response,
|
||||
},
|
||||
&request.context,
|
||||
)
|
||||
}
|
||||
|
||||
pub async fn async_store(
|
||||
&self,
|
||||
request: &ResponseCacheRequest,
|
||||
response: Value,
|
||||
now: Duration,
|
||||
) -> Result<(), Error> {
|
||||
if !request.controls.writes() {
|
||||
return Ok(());
|
||||
}
|
||||
self.backend
|
||||
.async_set_cache(
|
||||
&cache_key(&request.key),
|
||||
CacheEntry {
|
||||
timestamp: Some(now.as_secs_f64()),
|
||||
response,
|
||||
},
|
||||
request.context.clone(),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn async_store_batch(
|
||||
&self,
|
||||
entries: Vec<(ResponseCacheRequest, Value)>,
|
||||
now: Duration,
|
||||
) -> Result<(), Error> {
|
||||
self.async_store_entries(
|
||||
entries
|
||||
.into_iter()
|
||||
.map(|(request, response)| (request, response, now))
|
||||
.collect(),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Stores entries that each carry the time they were produced, so a deferred write keeps
|
||||
/// the freshness of its original response.
|
||||
pub async fn async_store_entries(
|
||||
&self,
|
||||
entries: Vec<(ResponseCacheRequest, Value, Duration)>,
|
||||
) -> Result<(), Error> {
|
||||
let writable = entries
|
||||
.into_iter()
|
||||
.filter(|(request, _, _)| request.controls.writes())
|
||||
.map(|(request, response, now)| {
|
||||
(
|
||||
cache_key(&request.key),
|
||||
CacheEntry {
|
||||
timestamp: Some(now.as_secs_f64()),
|
||||
response,
|
||||
},
|
||||
request.context,
|
||||
)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let Some((_, _, first_kwargs)) = writable.first() else {
|
||||
return Ok(());
|
||||
};
|
||||
if writable
|
||||
.iter()
|
||||
.all(|(_, _, context)| context == first_kwargs)
|
||||
{
|
||||
let context = first_kwargs.clone();
|
||||
let cache_list = writable
|
||||
.into_iter()
|
||||
.map(|(key, entry, _)| (key, entry))
|
||||
.collect();
|
||||
return self
|
||||
.backend
|
||||
.async_set_cache_pipeline(cache_list, context)
|
||||
.await;
|
||||
}
|
||||
for (key, entry, context) in writable {
|
||||
self.backend.async_set_cache(&key, entry, context).await?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn partial_hits(
|
||||
requests: &[ResponseCacheRequest],
|
||||
readable: Vec<(usize, &ResponseCacheRequest)>,
|
||||
entries: Vec<BatchEntry<CacheEntry>>,
|
||||
now: Duration,
|
||||
) -> Result<PartialHits, Error> {
|
||||
if readable.len() != entries.len() {
|
||||
return Err(Error::Unavailable);
|
||||
}
|
||||
let mut values = vec![None; requests.len()];
|
||||
for ((index, request), entry) in readable.into_iter().zip(entries) {
|
||||
let response = match entry {
|
||||
BatchEntry::Hit(entry) => Self::fresh_or_miss(Some(entry), now, request.max_age),
|
||||
BatchEntry::Miss | BatchEntry::Invalid => None,
|
||||
};
|
||||
values[index] = response;
|
||||
}
|
||||
Ok(PartialHits::new(values))
|
||||
}
|
||||
|
||||
fn fresh_or_miss(
|
||||
entry: Option<CacheEntry>,
|
||||
now: Duration,
|
||||
max_age: Option<Duration>,
|
||||
) -> Option<Value> {
|
||||
entry
|
||||
.filter(|entry| entry.fresh(now, max_age))
|
||||
.map(|entry| entry.response)
|
||||
}
|
||||
}
|
||||
90
litellm-rust/crates/cache-response/tests/caching.rs
Normal file
90
litellm-rust/crates/cache-response/tests/caching.rs
Normal file
|
|
@ -0,0 +1,90 @@
|
|||
use litellm_cache_response::{
|
||||
CacheControls, CacheKeyContext, CacheKeyField, CacheKeyInput, cache_key, get_cache_key,
|
||||
};
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
#[test]
|
||||
fn keys_match_python_order_groups_files_presets_and_namespaces() {
|
||||
let mut input = CacheKeyInput {
|
||||
fields: vec![
|
||||
CacheKeyField {
|
||||
name: "model".into(),
|
||||
value: Some("deployment".into()),
|
||||
api_parameter: true,
|
||||
internal_parameter: false,
|
||||
},
|
||||
CacheKeyField {
|
||||
name: "file".into(),
|
||||
value: None,
|
||||
api_parameter: true,
|
||||
internal_parameter: false,
|
||||
},
|
||||
],
|
||||
namespace: Some("team".into()),
|
||||
..Default::default()
|
||||
};
|
||||
CacheKeyContext {
|
||||
model_group: Some("group".into()),
|
||||
caching_groups: vec![(vec!["group".into()], "['group']".into())],
|
||||
file_checksum: Some("checksum".into()),
|
||||
..Default::default()
|
||||
}
|
||||
.apply(&mut input);
|
||||
assert_eq!(
|
||||
cache_key(&input),
|
||||
format!(
|
||||
"team:{:x}",
|
||||
Sha256::digest(b"model: ['group']file: checksum")
|
||||
)
|
||||
);
|
||||
input.preset = Some("preset".into());
|
||||
assert_eq!(get_cache_key(&input), "preset");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_controls_honor_default_modes_and_directives() {
|
||||
let enabled = CacheControls {
|
||||
supported_call_type: true,
|
||||
configured: true,
|
||||
default_on: true,
|
||||
..Default::default()
|
||||
};
|
||||
assert!(enabled.reads());
|
||||
assert!(enabled.writes());
|
||||
assert!(
|
||||
!CacheControls {
|
||||
default_on: false,
|
||||
..enabled
|
||||
}
|
||||
.reads()
|
||||
);
|
||||
assert!(
|
||||
CacheControls {
|
||||
default_on: false,
|
||||
use_cache: true,
|
||||
..enabled
|
||||
}
|
||||
.reads()
|
||||
);
|
||||
assert!(
|
||||
!CacheControls {
|
||||
no_cache: true,
|
||||
..enabled
|
||||
}
|
||||
.reads()
|
||||
);
|
||||
assert!(
|
||||
!CacheControls {
|
||||
no_store: true,
|
||||
..enabled
|
||||
}
|
||||
.writes()
|
||||
);
|
||||
assert!(
|
||||
!CacheControls {
|
||||
caching: Some(false),
|
||||
..enabled
|
||||
}
|
||||
.writes()
|
||||
);
|
||||
}
|
||||
484
litellm-rust/crates/cache-response/tests/response.rs
Normal file
484
litellm-rust/crates/cache-response/tests/response.rs
Normal file
|
|
@ -0,0 +1,484 @@
|
|||
use std::{
|
||||
sync::{
|
||||
Arc,
|
||||
atomic::{AtomicU64, Ordering},
|
||||
},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use litellm_cache::{BaseCache, CacheCodec, Error};
|
||||
use litellm_cache_memory::InMemoryCache;
|
||||
use litellm_cache_redis::RedisCache;
|
||||
use litellm_cache_response::{
|
||||
CacheEntry, CacheKeyField, CacheKeyInput, ResponseCache, ResponseCacheCodec,
|
||||
ResponseCacheRequest, WriteBuffer,
|
||||
};
|
||||
use redis_test::{MockCmd, MockRedisConnection};
|
||||
use serde_json::json;
|
||||
|
||||
fn memory() -> Arc<ResponseCache<InMemoryCache<CacheEntry>>> {
|
||||
Arc::new(ResponseCache::new(Arc::new(InMemoryCache::new(
|
||||
Some(8),
|
||||
Some(Duration::from_secs(600)),
|
||||
))))
|
||||
}
|
||||
|
||||
fn request() -> ResponseCacheRequest {
|
||||
ResponseCacheRequest::new(CacheKeyInput {
|
||||
preset: Some("tenant:key".into()),
|
||||
..Default::default()
|
||||
})
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sync_and_async_consumers_share_keys_ttls_and_freshness() {
|
||||
let clock = Arc::new(AtomicU64::new(100));
|
||||
let backend = Arc::new(InMemoryCache::with_clock(
|
||||
Some(8),
|
||||
Some(Duration::from_secs(600)),
|
||||
{
|
||||
let clock = clock.clone();
|
||||
move || Duration::from_secs(clock.load(Ordering::SeqCst))
|
||||
},
|
||||
));
|
||||
let cache = ResponseCache::new(backend.clone());
|
||||
let mut request = request();
|
||||
request.context.ttl = Some(Duration::from_secs(10));
|
||||
request.max_age = Some(Duration::from_secs(5));
|
||||
cache
|
||||
.store(
|
||||
&request,
|
||||
json!({"choices": [1], "usage": {"total_tokens": 7}}),
|
||||
Duration::from_secs(100),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
backend.expires_at("tenant:key").unwrap(),
|
||||
Some(Duration::from_secs(110))
|
||||
);
|
||||
assert!(
|
||||
cache
|
||||
.async_lookup(&request, Duration::from_secs(105))
|
||||
.await
|
||||
.unwrap()
|
||||
.is_some()
|
||||
);
|
||||
assert_eq!(
|
||||
cache.lookup(&request, Duration::from_secs(106)).unwrap(),
|
||||
None
|
||||
);
|
||||
request.max_age = None;
|
||||
assert_eq!(
|
||||
cache
|
||||
.lookup(&request, Duration::from_secs(106))
|
||||
.unwrap()
|
||||
.unwrap()["usage"]["total_tokens"],
|
||||
7
|
||||
);
|
||||
clock.store(111, Ordering::SeqCst);
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_lookup(&request, Duration::from_secs(111))
|
||||
.await
|
||||
.unwrap(),
|
||||
None
|
||||
);
|
||||
cache
|
||||
.async_store(&request, json!({"choices": [2]}), Duration::from_secs(111))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
cache.lookup(&request, Duration::from_secs(111)).unwrap(),
|
||||
Some(json!({"choices": [2]}))
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn directives_skip_io_and_keep_reads_and_writes_independent() {
|
||||
let cache = memory();
|
||||
let mut request = request();
|
||||
let now = Duration::from_secs(100);
|
||||
request.controls.no_store = true;
|
||||
cache
|
||||
.async_store(&request, json!({"v": 1}), now)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(cache.lookup(&request, now).unwrap(), None);
|
||||
request.controls.no_store = false;
|
||||
request.controls.no_cache = true;
|
||||
cache.store(&request, json!({"v": 2}), now).unwrap();
|
||||
assert_eq!(cache.async_lookup(&request, now).await.unwrap(), None);
|
||||
request.controls.no_cache = false;
|
||||
assert_eq!(cache.lookup(&request, now).unwrap(), Some(json!({"v": 2})));
|
||||
request.controls.default_on = false;
|
||||
cache.store(&request, json!({"v": 3}), now).unwrap();
|
||||
assert_eq!(cache.lookup(&request, now).unwrap(), None);
|
||||
request.controls.use_cache = true;
|
||||
assert_eq!(cache.lookup(&request, now).unwrap(), Some(json!({"v": 2})));
|
||||
request.controls.supported_call_type = false;
|
||||
assert_eq!(cache.lookup(&request, now).unwrap(), None);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn redis_consumer_reads_python_sync_and_async_envelopes_and_writes_compatible_json() {
|
||||
let connection = MockRedisConnection::new([
|
||||
MockCmd::new(
|
||||
redis::cmd("GET").arg("tenant:key"),
|
||||
Ok(br#"{'timestamp': 100.0, 'response': '{"ok": true, "text": "cached"}'}"#.to_vec()),
|
||||
),
|
||||
MockCmd::new(
|
||||
redis::cmd("GET").arg("tenant:key"),
|
||||
Ok(br#"{"timestamp":100.0,"response":{"ok":true,"text":"cached"}}"#.to_vec()),
|
||||
),
|
||||
MockCmd::new(
|
||||
redis::cmd("SETEX")
|
||||
.arg("tenant:key")
|
||||
.arg(600)
|
||||
.arg(br#"{"timestamp":100.0,"response":{"ok":true,"text":"cached"}}"#.as_slice()),
|
||||
Ok("OK"),
|
||||
),
|
||||
])
|
||||
.assert_all_commands_consumed();
|
||||
let backend = RedisCache::with_connection(connection, None, ResponseCacheCodec)
|
||||
.with_namespace(Some("tenant".into()));
|
||||
let cache = ResponseCache::new(Arc::new(backend));
|
||||
let request = request();
|
||||
let expected = json!({"ok": true, "text": "cached"});
|
||||
assert_eq!(
|
||||
cache.lookup(&request, Duration::from_secs(101)).unwrap(),
|
||||
Some(expected.clone())
|
||||
);
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_lookup(&request, Duration::from_secs(101))
|
||||
.await
|
||||
.unwrap(),
|
||||
Some(expected.clone())
|
||||
);
|
||||
cache
|
||||
.async_store(&request, expected, Duration::from_secs(100))
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn captured_service_keeps_the_selected_backend_for_background_writes() {
|
||||
let original = memory();
|
||||
let captured = original.clone();
|
||||
let replacement = memory();
|
||||
let request = request();
|
||||
let writer = tokio::spawn({
|
||||
let request = request.clone();
|
||||
async move {
|
||||
captured
|
||||
.async_store(
|
||||
&request,
|
||||
json!({"selected": "original"}),
|
||||
Duration::from_secs(100),
|
||||
)
|
||||
.await
|
||||
}
|
||||
});
|
||||
writer.await.unwrap().unwrap();
|
||||
assert_eq!(
|
||||
original.lookup(&request, Duration::from_secs(100)).unwrap(),
|
||||
Some(json!({"selected":"original"}))
|
||||
);
|
||||
assert_eq!(
|
||||
replacement
|
||||
.lookup(&request, Duration::from_secs(100))
|
||||
.unwrap(),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn generated_keys_preserve_namespace_and_explicit_keys() {
|
||||
let cache = memory();
|
||||
let key = CacheKeyInput {
|
||||
fields: vec![CacheKeyField {
|
||||
name: "model".into(),
|
||||
value: Some("a".into()),
|
||||
api_parameter: true,
|
||||
internal_parameter: false,
|
||||
}],
|
||||
namespace: Some("tenant".into()),
|
||||
..Default::default()
|
||||
};
|
||||
let generated = ResponseCacheRequest::new(key.clone());
|
||||
let explicit = ResponseCacheRequest::new(CacheKeyInput {
|
||||
preset: Some(litellm_cache_response::cache_key(&key)),
|
||||
..Default::default()
|
||||
});
|
||||
cache
|
||||
.store(&generated, json!({"value": 7}), Duration::from_secs(100))
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
cache.lookup(&explicit, Duration::from_secs(100)).unwrap(),
|
||||
Some(json!({"value":7}))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn response_codec_accepts_python_literals_without_executing_code() {
|
||||
let bytes = br#"{'timestamp': 100.0, 'response': {'text': 'hello \\ world', 'flag': True, 'empty': None, 'list': [1, 2.5]}}"#;
|
||||
let entry = ResponseCacheCodec.decode(bytes).unwrap();
|
||||
assert_eq!(
|
||||
entry.response,
|
||||
json!({"text": "hello \\ world", "flag": true, "empty": null, "list": [1, 2.5]})
|
||||
);
|
||||
for bytes in [
|
||||
b"__import__('os').system('false')".as_slice(),
|
||||
b"{'timestamp': 'invalid', 'response': {}}",
|
||||
b"{'timestamp': 1e9999, 'response': {}}",
|
||||
] {
|
||||
assert_eq!(
|
||||
ResponseCacheCodec.decode(bytes).unwrap_err(),
|
||||
Error::InvalidEntry
|
||||
);
|
||||
}
|
||||
let deep = format!("{}None{}", "[".repeat(1000), "]".repeat(1000));
|
||||
assert_eq!(
|
||||
ResponseCacheCodec.decode(deep.as_bytes()).unwrap_err(),
|
||||
Error::InvalidEntry
|
||||
);
|
||||
assert_eq!(
|
||||
ResponseCacheCodec
|
||||
.encode(&CacheEntry {
|
||||
timestamp: Some(f64::NAN),
|
||||
response: json!({})
|
||||
})
|
||||
.unwrap_err(),
|
||||
Error::InvalidEntry
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn invalid_entries_are_misses_and_disabled_reads_do_not_touch_redis() {
|
||||
let connection = MockRedisConnection::new([MockCmd::new(
|
||||
redis::cmd("GET").arg("tenant:key"),
|
||||
Ok(b"invalid".to_vec()),
|
||||
)])
|
||||
.assert_all_commands_consumed();
|
||||
let backend = RedisCache::with_connection(connection, None, ResponseCacheCodec);
|
||||
let cache = ResponseCache::new(Arc::new(backend));
|
||||
let mut request = request();
|
||||
request.controls.no_cache = true;
|
||||
assert_eq!(cache.lookup(&request, Duration::ZERO).unwrap(), None);
|
||||
request.controls.no_cache = false;
|
||||
assert_eq!(
|
||||
cache.async_lookup(&request, Duration::ZERO).await.unwrap(),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn string_responses_round_trip_through_typed_and_wire_backends() {
|
||||
let cache = ResponseCache::new(Arc::new(InMemoryCache::default()));
|
||||
let now = Duration::from_secs(100);
|
||||
for response in [json!("hello world"), json!("123"), json!("null")] {
|
||||
cache.store(&request(), response.clone(), now).unwrap();
|
||||
assert_eq!(
|
||||
cache.lookup(&request(), now).unwrap(),
|
||||
Some(response.clone())
|
||||
);
|
||||
|
||||
let wire = ResponseCacheCodec
|
||||
.encode(&CacheEntry {
|
||||
timestamp: Some(100.0),
|
||||
response: response.clone(),
|
||||
})
|
||||
.unwrap();
|
||||
assert_eq!(ResponseCacheCodec.decode(&wire).unwrap().response, response);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn non_object_responses_are_written_as_python_readable_serialized_strings() {
|
||||
let wire = ResponseCacheCodec
|
||||
.encode(&CacheEntry {
|
||||
timestamp: Some(100.0),
|
||||
response: json!([1, 2]),
|
||||
})
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
serde_json::from_slice::<serde_json::Value>(&wire).unwrap(),
|
||||
json!({"timestamp": 100.0, "response": "[1,2]"})
|
||||
);
|
||||
assert_eq!(
|
||||
ResponseCacheCodec.decode(&wire).unwrap().response,
|
||||
json!([1, 2])
|
||||
);
|
||||
assert_eq!(
|
||||
ResponseCacheCodec.decode(br#"{"timestamp": 100.0, "response": "not serialized"}"#),
|
||||
Err(Error::InvalidEntry)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn response_entries_preserve_the_existing_json_representation() {
|
||||
let codec = ResponseCacheCodec;
|
||||
let entry = CacheEntry {
|
||||
timestamp: Some(123.0),
|
||||
response: json!({"choices": [{"text": "cached"}]}),
|
||||
};
|
||||
let bytes = codec.encode(&entry).unwrap();
|
||||
assert_eq!(bytes, serde_json::to_vec(&entry).unwrap());
|
||||
assert_eq!(codec.decode(&bytes).unwrap(), entry);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn response_codec_preserves_values_without_timestamps() {
|
||||
let codec = ResponseCacheCodec;
|
||||
let raw = json!({"choices": [{"text": "legacy"}]});
|
||||
let entry = codec.decode(&serde_json::to_vec(&raw).unwrap()).unwrap();
|
||||
assert_eq!(entry.timestamp, None);
|
||||
assert_eq!(entry.response, raw);
|
||||
|
||||
let backend = Arc::new(InMemoryCache::default());
|
||||
BaseCache::set_cache(backend.as_ref(), "tenant:key", entry, &Default::default()).unwrap();
|
||||
let cache = ResponseCache::new(backend);
|
||||
assert_eq!(
|
||||
cache.lookup(&request(), Duration::from_secs(100)).unwrap(),
|
||||
Some(json!({"choices": [{"text": "legacy"}]}))
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn batch_lookup_reports_partial_hits_and_batch_store_populates_misses() {
|
||||
let cache = memory();
|
||||
let requests = ["hit", "miss", "disabled"].map(|key| {
|
||||
ResponseCacheRequest::new(CacheKeyInput {
|
||||
preset: Some(key.into()),
|
||||
..Default::default()
|
||||
})
|
||||
});
|
||||
cache
|
||||
.store(&requests[0], json!({"value": 1}), Duration::from_secs(100))
|
||||
.unwrap();
|
||||
let mut requests = requests.to_vec();
|
||||
requests[2].controls.caching = Some(false);
|
||||
|
||||
let partial = cache
|
||||
.async_lookup_batch(&requests, Duration::from_secs(100))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(partial.values, vec![Some(json!({"value": 1})), None, None]);
|
||||
assert_eq!(partial.missing_indices, vec![1, 2]);
|
||||
|
||||
cache
|
||||
.async_store_batch(
|
||||
vec![
|
||||
(requests[1].clone(), json!({"value": 2})),
|
||||
(requests[2].clone(), json!({"value": 3})),
|
||||
],
|
||||
Duration::from_secs(100),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
cache
|
||||
.lookup(&requests[1], Duration::from_secs(100))
|
||||
.unwrap(),
|
||||
Some(json!({"value": 2}))
|
||||
);
|
||||
requests[2].controls.caching = None;
|
||||
assert_eq!(
|
||||
cache
|
||||
.lookup(&requests[2], Duration::from_secs(100))
|
||||
.unwrap(),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn deferred_entries_keep_the_time_they_were_produced() {
|
||||
let cache = ResponseCache::new(Arc::new(InMemoryCache::default()));
|
||||
let mut request = request();
|
||||
request.max_age = Some(Duration::from_secs(10));
|
||||
cache
|
||||
.async_store_entries(vec![(
|
||||
request.clone(),
|
||||
json!({"answer": 7}),
|
||||
Duration::from_secs(100),
|
||||
)])
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
cache.lookup(&request, Duration::from_secs(110)).unwrap(),
|
||||
Some(json!({"answer": 7}))
|
||||
);
|
||||
assert_eq!(
|
||||
cache.lookup(&request, Duration::from_secs(111)).unwrap(),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn write_buffer_flushes_at_its_size_and_keeps_each_produced_time() {
|
||||
let cache = ResponseCache::new(Arc::new(InMemoryCache::default()));
|
||||
let buffer = WriteBuffer::new(2);
|
||||
let mut first = request();
|
||||
first.max_age = Some(Duration::from_secs(10));
|
||||
let mut second = request();
|
||||
second.key.preset = Some("tenant:other".into());
|
||||
|
||||
buffer
|
||||
.async_store(
|
||||
&cache,
|
||||
&first,
|
||||
json!({"answer": 7}),
|
||||
Duration::from_secs(100),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
cache.lookup(&first, Duration::from_secs(100)).unwrap(),
|
||||
None
|
||||
);
|
||||
|
||||
buffer
|
||||
.async_store(
|
||||
&cache,
|
||||
&second,
|
||||
json!({"answer": 8}),
|
||||
Duration::from_secs(200),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
cache.lookup(&first, Duration::from_secs(110)).unwrap(),
|
||||
Some(json!({"answer": 7}))
|
||||
);
|
||||
assert_eq!(
|
||||
cache.lookup(&first, Duration::from_secs(111)).unwrap(),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
cache.lookup(&second, Duration::from_secs(200)).unwrap(),
|
||||
Some(json!({"answer": 8}))
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn write_buffer_clear_drops_pending_entries() {
|
||||
let cache = ResponseCache::new(Arc::new(InMemoryCache::default()));
|
||||
let buffer = WriteBuffer::new(2);
|
||||
let mut other = request();
|
||||
other.key.preset = Some("tenant:other".into());
|
||||
let now = Duration::from_secs(100);
|
||||
|
||||
buffer
|
||||
.async_store(&cache, &request(), json!({"answer": 7}), now)
|
||||
.await
|
||||
.unwrap();
|
||||
buffer.clear().unwrap();
|
||||
buffer
|
||||
.async_store(&cache, &other, json!({"answer": 8}), now)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(cache.lookup(&request(), now).unwrap(), None);
|
||||
assert_eq!(cache.lookup(&other, now).unwrap(), None);
|
||||
}
|
||||
2
litellm-rust/crates/cache/Cargo.toml
vendored
2
litellm-rust/crates/cache/Cargo.toml
vendored
|
|
@ -8,8 +8,8 @@ repository.workspace = true
|
|||
[dependencies]
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
sha2.workspace = true
|
||||
thiserror.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
rstest.workspace = true
|
||||
tokio.workspace = true
|
||||
|
|
|
|||
121
litellm-rust/crates/cache/src/base_cache.rs
vendored
121
litellm-rust/crates/cache/src/base_cache.rs
vendored
|
|
@ -1,18 +1,35 @@
|
|||
use std::future::Future;
|
||||
use std::pin::Pin;
|
||||
use std::time::Duration;
|
||||
use std::{future::Future, time::Duration};
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::Error;
|
||||
|
||||
pub type CacheFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, Error>> + Send + 'a>>;
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub enum BatchEntry<V> {
|
||||
Hit(V),
|
||||
Miss,
|
||||
Invalid,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default, PartialEq)]
|
||||
pub struct CacheKwargs {
|
||||
pub trait CacheContext: Clone + Send + Sync + 'static {
|
||||
fn ttl(&self) -> Option<Duration>;
|
||||
|
||||
fn with_ttl(&self, ttl: Option<Duration>) -> Self;
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default, PartialEq, Eq)]
|
||||
pub struct ExactCacheContext {
|
||||
pub ttl: Option<Duration>,
|
||||
pub extras: Map<String, Value>,
|
||||
}
|
||||
|
||||
impl CacheContext for ExactCacheContext {
|
||||
fn ttl(&self) -> Option<Duration> {
|
||||
self.ttl
|
||||
}
|
||||
|
||||
fn with_ttl(&self, ttl: Option<Duration>) -> Self {
|
||||
Self { ttl }
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Deserialize, Serialize, PartialEq, Eq)]
|
||||
|
|
@ -32,67 +49,59 @@ pub struct CacheConnectionResult {
|
|||
|
||||
pub trait BaseCache: Send + Sync {
|
||||
type Value: Clone + Send + Sync + 'static;
|
||||
type Context: CacheContext;
|
||||
|
||||
fn default_ttl(&self) -> Duration {
|
||||
Duration::from_secs(60)
|
||||
}
|
||||
fn get_ttl(&self, context: &Self::Context) -> Option<Duration>;
|
||||
|
||||
fn get_ttl(&self, kwargs: &CacheKwargs) -> Duration {
|
||||
kwargs.ttl.unwrap_or_else(|| self.default_ttl())
|
||||
}
|
||||
|
||||
fn set_cache(&self, key: &str, value: Self::Value, kwargs: CacheKwargs) -> Result<(), Error>;
|
||||
|
||||
fn get_cache(&self, key: &str, kwargs: &CacheKwargs) -> Result<Option<Self::Value>, Error>;
|
||||
|
||||
fn async_set_cache<'a>(
|
||||
&'a self,
|
||||
key: &'a str,
|
||||
fn set_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
value: Self::Value,
|
||||
kwargs: CacheKwargs,
|
||||
) -> CacheFuture<'a, ()> {
|
||||
Box::pin(async move { self.set_cache(key, value, kwargs) })
|
||||
context: &Self::Context,
|
||||
) -> Result<(), Error>;
|
||||
|
||||
fn get_cache(&self, key: &str, context: &Self::Context) -> Result<Option<Self::Value>, Error>;
|
||||
|
||||
fn async_set_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
value: Self::Value,
|
||||
context: Self::Context,
|
||||
) -> impl Future<Output = Result<(), Error>> + Send {
|
||||
async move { self.set_cache(key, value, &context) }
|
||||
}
|
||||
|
||||
fn async_get_cache<'a>(
|
||||
&'a self,
|
||||
key: &'a str,
|
||||
kwargs: &'a CacheKwargs,
|
||||
) -> CacheFuture<'a, Option<Self::Value>> {
|
||||
Box::pin(async move { self.get_cache(key, kwargs) })
|
||||
fn async_get_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
context: &Self::Context,
|
||||
) -> impl Future<Output = Result<Option<Self::Value>, Error>> + Send {
|
||||
async move { self.get_cache(key, context) }
|
||||
}
|
||||
|
||||
fn async_set_cache_pipeline<'a>(
|
||||
&'a self,
|
||||
cache_list: Vec<(String, Self::Value)>,
|
||||
kwargs: CacheKwargs,
|
||||
) -> CacheFuture<'a, ()> {
|
||||
Box::pin(async move {
|
||||
for (key, value) in cache_list {
|
||||
self.set_cache(&key, value, kwargs.clone())?;
|
||||
fn async_set_cache_pipeline(
|
||||
&self,
|
||||
entries: Vec<(String, Self::Value)>,
|
||||
context: Self::Context,
|
||||
) -> impl Future<Output = Result<(), Error>> + Send {
|
||||
async move {
|
||||
for (key, value) in entries {
|
||||
self.async_set_cache(&key, value, context.clone()).await?;
|
||||
}
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn batch_cache_write<'a>(
|
||||
&'a self,
|
||||
key: &'a str,
|
||||
fn batch_cache_write(
|
||||
&self,
|
||||
key: &str,
|
||||
value: Self::Value,
|
||||
kwargs: CacheKwargs,
|
||||
) -> CacheFuture<'a, ()> {
|
||||
self.async_set_cache(key, value, kwargs)
|
||||
context: Self::Context,
|
||||
) -> impl Future<Output = Result<(), Error>> + Send {
|
||||
self.async_set_cache(key, value, context)
|
||||
}
|
||||
|
||||
fn delete_cache(&self, key: &str) -> Result<(), Error>;
|
||||
fn disconnect(&self) -> impl Future<Output = Result<(), Error>> + Send;
|
||||
|
||||
fn async_delete_cache<'a>(&'a self, key: &'a str) -> CacheFuture<'a, ()> {
|
||||
Box::pin(async move { self.delete_cache(key) })
|
||||
}
|
||||
|
||||
fn flush_cache(&self) -> Result<(), Error>;
|
||||
|
||||
fn disconnect(&self) -> CacheFuture<'_, ()>;
|
||||
|
||||
fn test_connection(&self) -> CacheFuture<'_, CacheConnectionResult>;
|
||||
fn test_connection(&self) -> impl Future<Output = Result<CacheConnectionResult, Error>> + Send;
|
||||
}
|
||||
|
|
|
|||
85
litellm-rust/crates/cache/src/cache_type.rs
vendored
Normal file
85
litellm-rust/crates/cache/src/cache_type.rs
vendored
Normal file
|
|
@ -0,0 +1,85 @@
|
|||
use serde::{Deserialize, Serialize};
|
||||
|
||||
#[derive(Clone, Copy, Debug, Deserialize, Serialize, PartialEq, Eq, Hash)]
|
||||
pub enum CacheType {
|
||||
#[serde(rename = "local")]
|
||||
Local,
|
||||
#[serde(rename = "redis")]
|
||||
Redis,
|
||||
#[serde(rename = "redis-semantic")]
|
||||
RedisSemantic,
|
||||
#[serde(rename = "valkey-semantic")]
|
||||
ValkeySemantic,
|
||||
#[serde(rename = "s3")]
|
||||
S3,
|
||||
#[serde(rename = "disk")]
|
||||
Disk,
|
||||
#[serde(rename = "qdrant-semantic")]
|
||||
QdrantSemantic,
|
||||
#[serde(rename = "azure-blob")]
|
||||
AzureBlob,
|
||||
#[serde(rename = "gcs")]
|
||||
Gcs,
|
||||
}
|
||||
|
||||
impl CacheType {
|
||||
pub const ALL: [Self; 9] = [
|
||||
Self::Local,
|
||||
Self::Redis,
|
||||
Self::RedisSemantic,
|
||||
Self::ValkeySemantic,
|
||||
Self::S3,
|
||||
Self::Disk,
|
||||
Self::QdrantSemantic,
|
||||
Self::AzureBlob,
|
||||
Self::Gcs,
|
||||
];
|
||||
|
||||
pub const fn as_python_name(self) -> &'static str {
|
||||
match self {
|
||||
Self::Local => "local",
|
||||
Self::Redis => "redis",
|
||||
Self::RedisSemantic => "redis-semantic",
|
||||
Self::ValkeySemantic => "valkey-semantic",
|
||||
Self::S3 => "s3",
|
||||
Self::Disk => "disk",
|
||||
Self::QdrantSemantic => "qdrant-semantic",
|
||||
Self::AzureBlob => "azure-blob",
|
||||
Self::Gcs => "gcs",
|
||||
}
|
||||
}
|
||||
|
||||
pub fn from_python_name(value: &str) -> Option<Self> {
|
||||
Self::ALL
|
||||
.into_iter()
|
||||
.find(|cache_type| cache_type.as_python_name() == value)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::CacheType;
|
||||
|
||||
#[test]
|
||||
fn every_python_cache_type_has_one_round_trip_identity() {
|
||||
let names = CacheType::ALL.map(CacheType::as_python_name);
|
||||
assert_eq!(
|
||||
names,
|
||||
[
|
||||
"local",
|
||||
"redis",
|
||||
"redis-semantic",
|
||||
"valkey-semantic",
|
||||
"s3",
|
||||
"disk",
|
||||
"qdrant-semantic",
|
||||
"azure-blob",
|
||||
"gcs",
|
||||
]
|
||||
);
|
||||
assert_eq!(
|
||||
names.map(CacheType::from_python_name),
|
||||
CacheType::ALL.map(Some)
|
||||
);
|
||||
}
|
||||
}
|
||||
167
litellm-rust/crates/cache/src/caching.rs
vendored
167
litellm-rust/crates/cache/src/caching.rs
vendored
|
|
@ -1,166 +1,23 @@
|
|||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
use crate::{BaseCache, CacheKwargs, Error};
|
||||
|
||||
pub use crate::BaseCache as Cache;
|
||||
use crate::{BaseCache, Error};
|
||||
|
||||
#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize, PartialEq, Eq)]
|
||||
pub enum CacheMode {
|
||||
#[default]
|
||||
#[serde(rename = "default_on")]
|
||||
DefaultOn,
|
||||
#[serde(rename = "default_off")]
|
||||
DefaultOff,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Deserialize, Serialize)]
|
||||
pub struct CacheKeyField {
|
||||
pub name: String,
|
||||
pub value: Option<String>,
|
||||
pub api_parameter: bool,
|
||||
pub internal_parameter: bool,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default, Deserialize, Serialize)]
|
||||
pub struct CacheKeyInput {
|
||||
pub fields: Vec<CacheKeyField>,
|
||||
pub preset: Option<String>,
|
||||
pub namespace: Option<String>,
|
||||
pub include_provider_parameters: bool,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
pub struct CacheKeyContext {
|
||||
pub model_group: Option<String>,
|
||||
pub caching_groups: Vec<(Vec<String>, String)>,
|
||||
pub file_checksum: Option<String>,
|
||||
pub file_object_name: Option<String>,
|
||||
pub metadata_file_name: Option<String>,
|
||||
pub parameters_file_name: Option<String>,
|
||||
}
|
||||
|
||||
impl CacheKeyContext {
|
||||
pub fn apply(self, input: &mut CacheKeyInput) {
|
||||
let group = self.model_group.as_ref().and_then(|model| {
|
||||
self.caching_groups
|
||||
.iter()
|
||||
.find(|(models, _)| models.contains(model))
|
||||
});
|
||||
for field in &mut input.fields {
|
||||
match field.name.as_str() {
|
||||
"model" => {
|
||||
field.value = group
|
||||
.map(|(_, formatted)| formatted.clone())
|
||||
.or_else(|| self.model_group.clone())
|
||||
.or_else(|| field.value.take())
|
||||
}
|
||||
"file" => {
|
||||
field.value = self
|
||||
.file_checksum
|
||||
.clone()
|
||||
.or_else(|| self.file_object_name.clone())
|
||||
.or_else(|| self.metadata_file_name.clone())
|
||||
.or_else(|| self.parameters_file_name.clone())
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn get_cache_key(input: &CacheKeyInput) -> String {
|
||||
cache_key(input)
|
||||
}
|
||||
|
||||
pub fn cache_key(input: &CacheKeyInput) -> String {
|
||||
if let Some(preset) = &input.preset {
|
||||
return preset.clone();
|
||||
}
|
||||
let mut digest = Sha256::new();
|
||||
for field in &input.fields {
|
||||
if (field.api_parameter || (input.include_provider_parameters && !field.internal_parameter))
|
||||
&& let Some(value) = &field.value
|
||||
{
|
||||
digest.update(field.name.as_bytes());
|
||||
digest.update(b": ");
|
||||
digest.update(value.as_bytes());
|
||||
}
|
||||
}
|
||||
let hash = format!("{:x}", digest.finalize());
|
||||
input
|
||||
.namespace
|
||||
.as_deref()
|
||||
.filter(|namespace| !namespace.is_empty())
|
||||
.map_or(hash.clone(), |namespace| format!("{namespace}:{hash}"))
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize)]
|
||||
pub struct CacheControls {
|
||||
pub supported_call_type: bool,
|
||||
pub configured: bool,
|
||||
pub native_backend: bool,
|
||||
pub default_on: bool,
|
||||
pub caching: Option<bool>,
|
||||
pub no_cache: bool,
|
||||
pub no_store: bool,
|
||||
#[serde(default)]
|
||||
pub use_cache: bool,
|
||||
}
|
||||
|
||||
impl CacheControls {
|
||||
pub fn reads(self) -> bool {
|
||||
self.supported_call_type
|
||||
&& self.configured
|
||||
&& self.caching.unwrap_or(true)
|
||||
&& !self.no_cache
|
||||
&& (self.default_on || self.use_cache)
|
||||
}
|
||||
|
||||
pub fn writes(self) -> bool {
|
||||
self.supported_call_type
|
||||
&& self.configured
|
||||
&& !self.no_store
|
||||
&& (self.default_on || self.use_cache)
|
||||
}
|
||||
}
|
||||
|
||||
pub fn should_use_cache(controls: CacheControls) -> bool {
|
||||
controls.reads() || controls.writes()
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct CacheEntry {
|
||||
pub timestamp: f64,
|
||||
pub response: Value,
|
||||
}
|
||||
|
||||
impl CacheEntry {
|
||||
pub fn fresh(&self, now: Duration, max_age: Option<Duration>) -> bool {
|
||||
self.timestamp.is_finite()
|
||||
&& max_age.is_none_or(|age| now.as_secs_f64() - self.timestamp <= age.as_secs_f64())
|
||||
}
|
||||
}
|
||||
|
||||
pub fn get_cache(
|
||||
cache: &dyn BaseCache<Value = CacheEntry>,
|
||||
pub fn get_cache<B: BaseCache>(
|
||||
cache: &B,
|
||||
key: &str,
|
||||
kwargs: &CacheKwargs,
|
||||
) -> Result<Option<CacheEntry>, Error> {
|
||||
cache.get_cache(key, kwargs)
|
||||
context: &B::Context,
|
||||
) -> Result<Option<B::Value>, Error> {
|
||||
cache.get_cache(key, context)
|
||||
}
|
||||
|
||||
pub fn set_cache(
|
||||
cache: &dyn BaseCache<Value = CacheEntry>,
|
||||
pub fn set_cache<B: BaseCache>(
|
||||
cache: &B,
|
||||
key: &str,
|
||||
entry: CacheEntry,
|
||||
kwargs: CacheKwargs,
|
||||
value: B::Value,
|
||||
context: &B::Context,
|
||||
) -> Result<(), Error> {
|
||||
cache.set_cache(key, entry, kwargs)
|
||||
cache.set_cache(key, value, context)
|
||||
}
|
||||
|
||||
pub type CacheBackend = Arc<dyn BaseCache<Value = CacheEntry>>;
|
||||
pub type CacheBackend<B> = Arc<B>;
|
||||
|
|
|
|||
169
litellm-rust/crates/cache/src/capabilities.rs
vendored
Normal file
169
litellm-rust/crates/cache/src/capabilities.rs
vendored
Normal file
|
|
@ -0,0 +1,169 @@
|
|||
use std::{future::Future, time::Duration};
|
||||
|
||||
use crate::{BaseCache, BatchEntry, Error};
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct IncrementOperation {
|
||||
pub key: String,
|
||||
pub amount: f64,
|
||||
pub ttl: Option<Duration>,
|
||||
}
|
||||
|
||||
pub trait BatchCache: BaseCache {
|
||||
fn batch_get_cache(
|
||||
&self,
|
||||
keys: &[String],
|
||||
context: &Self::Context,
|
||||
) -> Result<Vec<BatchEntry<Self::Value>>, Error> {
|
||||
keys.iter()
|
||||
.map(|key| match self.get_cache(key, context) {
|
||||
Ok(Some(value)) => Ok(BatchEntry::Hit(value)),
|
||||
Ok(None) => Ok(BatchEntry::Miss),
|
||||
Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid),
|
||||
Err(error) => Err(error),
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn async_batch_get_cache(
|
||||
&self,
|
||||
keys: Vec<String>,
|
||||
context: Self::Context,
|
||||
) -> impl Future<Output = Result<Vec<BatchEntry<Self::Value>>, Error>> + Send {
|
||||
async move {
|
||||
let mut entries = Vec::with_capacity(keys.len());
|
||||
for key in keys {
|
||||
entries.push(match self.async_get_cache(&key, &context).await {
|
||||
Ok(Some(value)) => BatchEntry::Hit(value),
|
||||
Ok(None) => BatchEntry::Miss,
|
||||
Err(Error::InvalidEntry) => BatchEntry::Invalid,
|
||||
Err(error) => return Err(error),
|
||||
});
|
||||
}
|
||||
Ok(entries)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub trait DeleteCache: BaseCache {
|
||||
fn delete_cache(&self, key: &str) -> Result<(), Error>;
|
||||
|
||||
fn async_delete_cache(&self, key: &str) -> impl Future<Output = Result<(), Error>> + Send {
|
||||
async move { self.delete_cache(key) }
|
||||
}
|
||||
}
|
||||
|
||||
pub trait FlushCache: BaseCache {
|
||||
fn flush_cache(&self) -> Result<(), Error>;
|
||||
|
||||
fn async_flush_cache(&self) -> impl Future<Output = Result<(), Error>> + Send {
|
||||
async move { self.flush_cache() }
|
||||
}
|
||||
}
|
||||
|
||||
pub trait CounterCache: BaseCache<Value = f64> {
|
||||
fn increment_cache(&self, key: &str, amount: f64, context: Self::Context)
|
||||
-> Result<f64, Error>;
|
||||
|
||||
fn async_increment(
|
||||
&self,
|
||||
key: &str,
|
||||
amount: f64,
|
||||
context: Self::Context,
|
||||
) -> impl Future<Output = Result<f64, Error>> + Send {
|
||||
async move { self.increment_cache(key, amount, context) }
|
||||
}
|
||||
}
|
||||
|
||||
pub trait ClaimCache: BaseCache
|
||||
where
|
||||
Self::Value: PartialEq,
|
||||
{
|
||||
fn claim_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
candidate: Self::Value,
|
||||
eligible: &[Self::Value],
|
||||
context: Self::Context,
|
||||
) -> Result<Self::Value, Error>;
|
||||
|
||||
fn async_claim_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
candidate: Self::Value,
|
||||
eligible: Vec<Self::Value>,
|
||||
context: Self::Context,
|
||||
) -> impl Future<Output = Result<Self::Value, Error>> + Send {
|
||||
async move { self.claim_cache(key, candidate, &eligible, context) }
|
||||
}
|
||||
}
|
||||
|
||||
pub trait TtlCache: BaseCache {
|
||||
fn async_get_ttl(
|
||||
&self,
|
||||
key: &str,
|
||||
) -> impl Future<Output = Result<Option<Duration>, Error>> + Send;
|
||||
}
|
||||
|
||||
pub trait SetCache: BaseCache {
|
||||
type SetValue: Clone + Send + Sync + 'static;
|
||||
type SetResult: Send + Sync + 'static;
|
||||
|
||||
fn async_set_cache_sadd(
|
||||
&self,
|
||||
key: &str,
|
||||
values: Vec<Self::SetValue>,
|
||||
ttl: Option<Duration>,
|
||||
) -> impl Future<Output = Result<Self::SetResult, Error>> + Send;
|
||||
}
|
||||
|
||||
pub trait QueueCache: BaseCache {
|
||||
type QueueValue: Clone + Send + Sync + 'static;
|
||||
type PopResult: Send + Sync + 'static;
|
||||
|
||||
fn async_rpush(
|
||||
&self,
|
||||
key: &str,
|
||||
values: Vec<Self::QueueValue>,
|
||||
) -> impl Future<Output = Result<usize, Error>> + Send;
|
||||
|
||||
fn async_lpop(
|
||||
&self,
|
||||
key: &str,
|
||||
count: Option<usize>,
|
||||
) -> impl Future<Output = Result<Self::PopResult, Error>> + Send;
|
||||
}
|
||||
|
||||
pub trait ScanCache: BaseCache {
|
||||
fn async_scan_iter(
|
||||
&self,
|
||||
pattern: &str,
|
||||
count: usize,
|
||||
) -> impl Future<Output = Result<Vec<String>, Error>> + Send;
|
||||
}
|
||||
|
||||
pub trait ClientInfoCache: BaseCache {
|
||||
type ClientList: Send + Sync + 'static;
|
||||
type Info: Send + Sync + 'static;
|
||||
|
||||
fn client_list(&self) -> Result<Self::ClientList, Error>;
|
||||
|
||||
fn info(&self) -> Result<Self::Info, Error>;
|
||||
}
|
||||
|
||||
pub trait CacheScript: Send + Sync + 'static {
|
||||
type Argument: Clone + Send + Sync + 'static;
|
||||
type Output: Send + Sync + 'static;
|
||||
|
||||
fn invoke(
|
||||
&self,
|
||||
keys: Vec<String>,
|
||||
arguments: Vec<Self::Argument>,
|
||||
) -> impl Future<Output = Result<Self::Output, Error>> + Send;
|
||||
}
|
||||
|
||||
pub trait ScriptCache: BaseCache {
|
||||
type Script: CacheScript;
|
||||
|
||||
fn async_register_script(&self, source: String) -> Self::Script;
|
||||
}
|
||||
50
litellm-rust/crates/cache/src/codec.rs
vendored
Normal file
50
litellm-rust/crates/cache/src/codec.rs
vendored
Normal file
|
|
@ -0,0 +1,50 @@
|
|||
use std::marker::PhantomData;
|
||||
|
||||
use serde::{Serialize, de::DeserializeOwned};
|
||||
|
||||
use crate::Error;
|
||||
|
||||
pub trait CacheCodec: Send + Sync {
|
||||
type Value: Clone + Send + Sync + 'static;
|
||||
|
||||
fn encode(&self, value: &Self::Value) -> Result<Vec<u8>, Error>;
|
||||
|
||||
fn decode(&self, bytes: &[u8]) -> Result<Self::Value, Error>;
|
||||
}
|
||||
|
||||
pub struct JsonCodec<V>(PhantomData<fn() -> V>);
|
||||
|
||||
impl<V> Clone for JsonCodec<V> {
|
||||
fn clone(&self) -> Self {
|
||||
*self
|
||||
}
|
||||
}
|
||||
|
||||
impl<V> Copy for JsonCodec<V> {}
|
||||
|
||||
impl<V> Default for JsonCodec<V> {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl<V> JsonCodec<V> {
|
||||
pub const fn new() -> Self {
|
||||
Self(PhantomData)
|
||||
}
|
||||
}
|
||||
|
||||
impl<V> CacheCodec for JsonCodec<V>
|
||||
where
|
||||
V: Clone + Send + Sync + Serialize + DeserializeOwned + 'static,
|
||||
{
|
||||
type Value = V;
|
||||
|
||||
fn encode(&self, value: &Self::Value) -> Result<Vec<u8>, Error> {
|
||||
serde_json::to_vec(value).map_err(|_| Error::InvalidEntry)
|
||||
}
|
||||
|
||||
fn decode(&self, bytes: &[u8]) -> Result<Self::Value, Error> {
|
||||
serde_json::from_slice(bytes).map_err(|_| Error::InvalidEntry)
|
||||
}
|
||||
}
|
||||
390
litellm-rust/crates/cache/src/dual.rs
vendored
Normal file
390
litellm-rust/crates/cache/src/dual.rs
vendored
Normal file
|
|
@ -0,0 +1,390 @@
|
|||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use crate::{
|
||||
BaseCache, BatchCache, BatchEntry, CacheConnectionResult, CacheContext, ClaimCache,
|
||||
CounterCache, DeleteCache, Error, FlushCache,
|
||||
};
|
||||
|
||||
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
|
||||
pub enum ReadPolicy {
|
||||
#[default]
|
||||
LocalThenRemote,
|
||||
LocalOnly,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
|
||||
pub enum WritePolicy {
|
||||
#[default]
|
||||
Both,
|
||||
LocalOnly,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
|
||||
pub enum RemoteFailurePolicy {
|
||||
#[default]
|
||||
Propagate,
|
||||
UseLocal,
|
||||
}
|
||||
|
||||
pub struct DualCache<L1, L2> {
|
||||
l1: Arc<L1>,
|
||||
l2: Arc<L2>,
|
||||
read_policy: ReadPolicy,
|
||||
write_policy: WritePolicy,
|
||||
remote_failure_policy: RemoteFailurePolicy,
|
||||
promotion_ttl: Option<Duration>,
|
||||
}
|
||||
|
||||
impl<L1, L2> DualCache<L1, L2> {
|
||||
pub fn new(l1: Arc<L1>, l2: Arc<L2>) -> Self {
|
||||
Self {
|
||||
l1,
|
||||
l2,
|
||||
read_policy: ReadPolicy::default(),
|
||||
write_policy: WritePolicy::default(),
|
||||
remote_failure_policy: RemoteFailurePolicy::default(),
|
||||
promotion_ttl: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_read_policy(self, read_policy: ReadPolicy) -> Self {
|
||||
Self {
|
||||
read_policy,
|
||||
..self
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_write_policy(self, write_policy: WritePolicy) -> Self {
|
||||
Self {
|
||||
write_policy,
|
||||
..self
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_remote_failure_policy(self, remote_failure_policy: RemoteFailurePolicy) -> Self {
|
||||
Self {
|
||||
remote_failure_policy,
|
||||
..self
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_promotion_ttl(self, promotion_ttl: Duration) -> Self {
|
||||
Self {
|
||||
promotion_ttl: Some(promotion_ttl),
|
||||
..self
|
||||
}
|
||||
}
|
||||
|
||||
fn reads_remote(&self) -> bool {
|
||||
self.read_policy == ReadPolicy::LocalThenRemote
|
||||
}
|
||||
|
||||
fn writes_remote(&self) -> bool {
|
||||
self.write_policy == WritePolicy::Both
|
||||
}
|
||||
|
||||
fn remote<T>(&self, result: Result<T, Error>) -> Result<Option<T>, Error> {
|
||||
match result {
|
||||
Ok(value) => Ok(Some(value)),
|
||||
Err(Error::Unavailable)
|
||||
if self.remote_failure_policy == RemoteFailurePolicy::UseLocal =>
|
||||
{
|
||||
Ok(None)
|
||||
}
|
||||
Err(error) => Err(error),
|
||||
}
|
||||
}
|
||||
|
||||
fn promotion_context<C: CacheContext>(&self, context: &C) -> C {
|
||||
context.with_ttl(self.promotion_ttl.or(context.ttl()))
|
||||
}
|
||||
}
|
||||
|
||||
impl<V, C, L1, L2> DualCache<L1, L2>
|
||||
where
|
||||
V: Clone + Send + Sync + 'static,
|
||||
C: CacheContext,
|
||||
L1: BaseCache<Value = V, Context = C>,
|
||||
L2: BaseCache<Value = V, Context = C>,
|
||||
{
|
||||
fn missing(entries: &[BatchEntry<V>]) -> Vec<usize> {
|
||||
entries
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter_map(|(index, entry)| (!matches!(entry, BatchEntry::Hit(_))).then_some(index))
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn merge_batch(
|
||||
&self,
|
||||
keys: &[String],
|
||||
context: &C,
|
||||
mut entries: Vec<BatchEntry<V>>,
|
||||
missing: Vec<usize>,
|
||||
remote: Vec<BatchEntry<V>>,
|
||||
) -> Result<Vec<BatchEntry<V>>, Error> {
|
||||
if missing.len() != remote.len() {
|
||||
return Err(Error::Unavailable);
|
||||
}
|
||||
for (index, entry) in missing.into_iter().zip(remote) {
|
||||
if let BatchEntry::Hit(value) = &entry {
|
||||
let promotion_context = self.promotion_context(context);
|
||||
self.l1
|
||||
.set_cache(&keys[index], value.clone(), &promotion_context)?;
|
||||
}
|
||||
entries[index] = entry;
|
||||
}
|
||||
Ok(entries)
|
||||
}
|
||||
}
|
||||
|
||||
impl<V, C, L1, L2> BaseCache for DualCache<L1, L2>
|
||||
where
|
||||
V: Clone + Send + Sync + 'static,
|
||||
C: CacheContext,
|
||||
L1: BaseCache<Value = V, Context = C>,
|
||||
L2: BaseCache<Value = V, Context = C>,
|
||||
{
|
||||
type Value = V;
|
||||
type Context = C;
|
||||
|
||||
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
|
||||
self.l2.get_ttl(context)
|
||||
}
|
||||
|
||||
fn set_cache(&self, key: &str, value: V, context: &C) -> Result<(), Error> {
|
||||
if self.writes_remote() {
|
||||
self.remote(self.l2.set_cache(key, value.clone(), context))?;
|
||||
}
|
||||
self.l1.set_cache(key, value, context)
|
||||
}
|
||||
|
||||
fn get_cache(&self, key: &str, context: &C) -> Result<Option<V>, Error> {
|
||||
if let Some(value) = self.l1.get_cache(key, context)? {
|
||||
return Ok(Some(value));
|
||||
}
|
||||
if !self.reads_remote() {
|
||||
return Ok(None);
|
||||
}
|
||||
let value = self.remote(self.l2.get_cache(key, context))?.flatten();
|
||||
if let Some(value) = &value {
|
||||
let promotion_context = self.promotion_context(context);
|
||||
self.l1.set_cache(key, value.clone(), &promotion_context)?;
|
||||
}
|
||||
Ok(value)
|
||||
}
|
||||
|
||||
async fn async_set_cache(&self, key: &str, value: V, context: C) -> Result<(), Error> {
|
||||
if self.writes_remote() {
|
||||
self.remote(
|
||||
self.l2
|
||||
.async_set_cache(key, value.clone(), context.clone())
|
||||
.await,
|
||||
)?;
|
||||
}
|
||||
self.l1.async_set_cache(key, value, context).await
|
||||
}
|
||||
|
||||
async fn async_get_cache(&self, key: &str, context: &C) -> Result<Option<V>, Error> {
|
||||
if let Some(value) = self.l1.async_get_cache(key, context).await? {
|
||||
return Ok(Some(value));
|
||||
}
|
||||
if !self.reads_remote() {
|
||||
return Ok(None);
|
||||
}
|
||||
let value = self
|
||||
.remote(self.l2.async_get_cache(key, context).await)?
|
||||
.flatten();
|
||||
if let Some(value) = &value {
|
||||
self.l1
|
||||
.async_set_cache(key, value.clone(), self.promotion_context(context))
|
||||
.await?;
|
||||
}
|
||||
Ok(value)
|
||||
}
|
||||
|
||||
async fn async_set_cache_pipeline(
|
||||
&self,
|
||||
entries: Vec<(String, V)>,
|
||||
context: C,
|
||||
) -> Result<(), Error> {
|
||||
if self.writes_remote() {
|
||||
self.remote(
|
||||
self.l2
|
||||
.async_set_cache_pipeline(entries.clone(), context.clone())
|
||||
.await,
|
||||
)?;
|
||||
}
|
||||
self.l1.async_set_cache_pipeline(entries, context).await
|
||||
}
|
||||
|
||||
async fn disconnect(&self) -> Result<(), Error> {
|
||||
self.l2.disconnect().await?;
|
||||
self.l1.disconnect().await
|
||||
}
|
||||
|
||||
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
|
||||
self.l2.test_connection().await
|
||||
}
|
||||
}
|
||||
|
||||
impl<V, C, L1, L2> BatchCache for DualCache<L1, L2>
|
||||
where
|
||||
V: Clone + Send + Sync + 'static,
|
||||
C: CacheContext,
|
||||
L1: BatchCache<Value = V, Context = C>,
|
||||
L2: BatchCache<Value = V, Context = C>,
|
||||
{
|
||||
fn batch_get_cache(&self, keys: &[String], context: &C) -> Result<Vec<BatchEntry<V>>, Error> {
|
||||
let entries = self.l1.batch_get_cache(keys, context)?;
|
||||
let missing = Self::missing(&entries);
|
||||
if missing.is_empty() || !self.reads_remote() {
|
||||
return Ok(entries);
|
||||
}
|
||||
let remote_keys = missing
|
||||
.iter()
|
||||
.map(|index| keys[*index].clone())
|
||||
.collect::<Vec<_>>();
|
||||
match self.remote(self.l2.batch_get_cache(&remote_keys, context))? {
|
||||
Some(remote) => self.merge_batch(keys, context, entries, missing, remote),
|
||||
None => Ok(entries),
|
||||
}
|
||||
}
|
||||
|
||||
async fn async_batch_get_cache(
|
||||
&self,
|
||||
keys: Vec<String>,
|
||||
context: C,
|
||||
) -> Result<Vec<BatchEntry<V>>, Error> {
|
||||
let entries = self
|
||||
.l1
|
||||
.async_batch_get_cache(keys.clone(), context.clone())
|
||||
.await?;
|
||||
let missing = Self::missing(&entries);
|
||||
if missing.is_empty() || !self.reads_remote() {
|
||||
return Ok(entries);
|
||||
}
|
||||
let remote_keys = missing.iter().map(|index| keys[*index].clone()).collect();
|
||||
match self.remote(
|
||||
self.l2
|
||||
.async_batch_get_cache(remote_keys, context.clone())
|
||||
.await,
|
||||
)? {
|
||||
Some(remote) => self.merge_batch(&keys, &context, entries, missing, remote),
|
||||
None => Ok(entries),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<V, C, L1, L2> DeleteCache for DualCache<L1, L2>
|
||||
where
|
||||
V: Clone + Send + Sync + 'static,
|
||||
C: CacheContext,
|
||||
L1: DeleteCache<Value = V, Context = C>,
|
||||
L2: DeleteCache<Value = V, Context = C>,
|
||||
{
|
||||
fn delete_cache(&self, key: &str) -> Result<(), Error> {
|
||||
if self.writes_remote() {
|
||||
self.remote(self.l2.delete_cache(key))?;
|
||||
}
|
||||
self.l1.delete_cache(key)
|
||||
}
|
||||
|
||||
async fn async_delete_cache(&self, key: &str) -> Result<(), Error> {
|
||||
if self.writes_remote() {
|
||||
self.remote(self.l2.async_delete_cache(key).await)?;
|
||||
}
|
||||
self.l1.async_delete_cache(key).await
|
||||
}
|
||||
}
|
||||
|
||||
impl<V, C, L1, L2> FlushCache for DualCache<L1, L2>
|
||||
where
|
||||
V: Clone + Send + Sync + 'static,
|
||||
C: CacheContext,
|
||||
L1: FlushCache<Value = V, Context = C>,
|
||||
L2: FlushCache<Value = V, Context = C>,
|
||||
{
|
||||
fn flush_cache(&self) -> Result<(), Error> {
|
||||
if self.writes_remote() {
|
||||
self.remote(self.l2.flush_cache())?;
|
||||
}
|
||||
self.l1.flush_cache()
|
||||
}
|
||||
|
||||
async fn async_flush_cache(&self) -> Result<(), Error> {
|
||||
if self.writes_remote() {
|
||||
self.remote(self.l2.async_flush_cache().await)?;
|
||||
}
|
||||
self.l1.async_flush_cache().await
|
||||
}
|
||||
}
|
||||
|
||||
impl<C, L1, L2> CounterCache for DualCache<L1, L2>
|
||||
where
|
||||
C: CacheContext,
|
||||
L1: BaseCache<Value = f64, Context = C>,
|
||||
L2: CounterCache<Context = C>,
|
||||
{
|
||||
fn increment_cache(&self, key: &str, amount: f64, context: C) -> Result<f64, Error> {
|
||||
let value = self.l2.increment_cache(key, amount, context.clone())?;
|
||||
self.l1.set_cache(key, value, &context)?;
|
||||
Ok(value)
|
||||
}
|
||||
|
||||
async fn async_increment(&self, key: &str, amount: f64, context: C) -> Result<f64, Error> {
|
||||
let value = self
|
||||
.l2
|
||||
.async_increment(key, amount, context.clone())
|
||||
.await?;
|
||||
self.l1.async_set_cache(key, value, context).await?;
|
||||
Ok(value)
|
||||
}
|
||||
}
|
||||
|
||||
impl<V, C, L1, L2> ClaimCache for DualCache<L1, L2>
|
||||
where
|
||||
V: Clone + PartialEq + Send + Sync + 'static,
|
||||
C: CacheContext,
|
||||
L1: ClaimCache<Value = V, Context = C>,
|
||||
L2: ClaimCache<Value = V, Context = C>,
|
||||
{
|
||||
fn claim_cache(&self, key: &str, candidate: V, eligible: &[V], context: C) -> Result<V, Error> {
|
||||
match self.remote(
|
||||
self.l2
|
||||
.claim_cache(key, candidate.clone(), eligible, context.clone()),
|
||||
)? {
|
||||
Some(winner) => {
|
||||
self.l1.set_cache(key, winner.clone(), &context)?;
|
||||
Ok(winner)
|
||||
}
|
||||
None => self.l1.claim_cache(key, candidate, eligible, context),
|
||||
}
|
||||
}
|
||||
|
||||
async fn async_claim_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
candidate: V,
|
||||
eligible: Vec<V>,
|
||||
context: C,
|
||||
) -> Result<V, Error> {
|
||||
match self.remote(
|
||||
self.l2
|
||||
.async_claim_cache(key, candidate.clone(), eligible.clone(), context.clone())
|
||||
.await,
|
||||
)? {
|
||||
Some(winner) => {
|
||||
self.l1
|
||||
.async_set_cache(key, winner.clone(), context)
|
||||
.await?;
|
||||
Ok(winner)
|
||||
}
|
||||
None => {
|
||||
self.l1
|
||||
.async_claim_cache(key, candidate, eligible, context)
|
||||
.await
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
2
litellm-rust/crates/cache/src/error.rs
vendored
2
litellm-rust/crates/cache/src/error.rs
vendored
|
|
@ -4,4 +4,6 @@ pub enum Error {
|
|||
Unavailable,
|
||||
#[error("invalid cache entry")]
|
||||
InvalidEntry,
|
||||
#[error("flushing Redis requires an explicit namespace")]
|
||||
UnscopedFlush,
|
||||
}
|
||||
|
|
|
|||
17
litellm-rust/crates/cache/src/lib.rs
vendored
17
litellm-rust/crates/cache/src/lib.rs
vendored
|
|
@ -1,12 +1,21 @@
|
|||
mod base_cache;
|
||||
mod cache_type;
|
||||
mod caching;
|
||||
mod capabilities;
|
||||
mod codec;
|
||||
mod dual;
|
||||
mod error;
|
||||
|
||||
pub use base_cache::{
|
||||
BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheFuture, CacheKwargs,
|
||||
BaseCache, BatchEntry, CacheConnectionResult, CacheConnectionStatus, CacheContext,
|
||||
ExactCacheContext,
|
||||
};
|
||||
pub use caching::{
|
||||
Cache, CacheBackend, CacheControls, CacheEntry, CacheKeyContext, CacheKeyField, CacheKeyInput,
|
||||
CacheMode, cache_key, get_cache, get_cache_key, set_cache, should_use_cache,
|
||||
pub use cache_type::CacheType;
|
||||
pub use caching::{Cache, CacheBackend, get_cache, set_cache};
|
||||
pub use capabilities::{
|
||||
BatchCache, CacheScript, ClaimCache, ClientInfoCache, CounterCache, DeleteCache, FlushCache,
|
||||
IncrementOperation, QueueCache, ScanCache, ScriptCache, SetCache, TtlCache,
|
||||
};
|
||||
pub use codec::{CacheCodec, JsonCodec};
|
||||
pub use dual::{DualCache, ReadPolicy, RemoteFailurePolicy, WritePolicy};
|
||||
pub use error::Error;
|
||||
|
|
|
|||
208
litellm-rust/crates/cache/tests/caching.rs
vendored
208
litellm-rust/crates/cache/tests/caching.rs
vendored
|
|
@ -1,42 +1,97 @@
|
|||
use std::{sync::Mutex, time::Duration};
|
||||
|
||||
use litellm_cache::{
|
||||
BaseCache, CacheConnectionResult, CacheControls, CacheEntry, CacheFuture, CacheKeyContext,
|
||||
CacheKeyField, CacheKeyInput, CacheKwargs, Error, cache_key, get_cache_key,
|
||||
BaseCache, CacheConnectionResult, CacheContext, Error, ExactCacheContext, get_cache,
|
||||
};
|
||||
use sha2::{Digest, Sha256};
|
||||
use std::time::Duration;
|
||||
|
||||
struct TestCache {
|
||||
default_ttl: Duration,
|
||||
writes: Mutex<Vec<(String, String, ExactCacheContext)>>,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct SemanticContext {
|
||||
ttl: Option<Duration>,
|
||||
query: String,
|
||||
}
|
||||
|
||||
impl CacheContext for SemanticContext {
|
||||
fn ttl(&self) -> Option<Duration> {
|
||||
self.ttl
|
||||
}
|
||||
|
||||
fn with_ttl(&self, ttl: Option<Duration>) -> Self {
|
||||
Self {
|
||||
ttl,
|
||||
query: self.query.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct SemanticCache;
|
||||
|
||||
impl BaseCache for SemanticCache {
|
||||
type Value = String;
|
||||
type Context = SemanticContext;
|
||||
|
||||
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
|
||||
context.ttl
|
||||
}
|
||||
|
||||
fn set_cache(&self, _: &str, _: Self::Value, _: &Self::Context) -> Result<(), Error> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn get_cache(&self, _: &str, context: &Self::Context) -> Result<Option<Self::Value>, Error> {
|
||||
Ok((context.query == "matching prompt").then(|| "semantic hit".into()))
|
||||
}
|
||||
|
||||
async fn disconnect(&self) -> Result<(), Error> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
|
||||
unreachable!()
|
||||
}
|
||||
}
|
||||
|
||||
impl BaseCache for TestCache {
|
||||
type Value = CacheEntry;
|
||||
type Value = String;
|
||||
type Context = ExactCacheContext;
|
||||
|
||||
fn default_ttl(&self) -> Duration {
|
||||
self.default_ttl
|
||||
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
|
||||
context.ttl.or(Some(self.default_ttl))
|
||||
}
|
||||
|
||||
fn set_cache(&self, _: &str, _: Self::Value, _: CacheKwargs) -> Result<(), Error> {
|
||||
fn set_cache(&self, _: &str, _: Self::Value, _: &ExactCacheContext) -> Result<(), Error> {
|
||||
Err(Error::Unavailable)
|
||||
}
|
||||
|
||||
async fn async_set_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
value: Self::Value,
|
||||
context: ExactCacheContext,
|
||||
) -> Result<(), Error> {
|
||||
if key == "unavailable" {
|
||||
return Err(Error::Unavailable);
|
||||
}
|
||||
self.writes
|
||||
.lock()
|
||||
.unwrap()
|
||||
.push((key.into(), value, context));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn get_cache(&self, _: &str, _: &CacheKwargs) -> Result<Option<Self::Value>, Error> {
|
||||
fn get_cache(&self, _: &str, _: &ExactCacheContext) -> Result<Option<Self::Value>, Error> {
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
fn delete_cache(&self, _: &str) -> Result<(), Error> {
|
||||
async fn disconnect(&self) -> Result<(), Error> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn flush_cache(&self) -> Result<(), Error> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn disconnect(&self) -> CacheFuture<'_, ()> {
|
||||
Box::pin(async { Ok(()) })
|
||||
}
|
||||
|
||||
fn test_connection(&self) -> CacheFuture<'_, CacheConnectionResult> {
|
||||
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
|
||||
unreachable!()
|
||||
}
|
||||
}
|
||||
|
|
@ -45,95 +100,64 @@ impl BaseCache for TestCache {
|
|||
fn ttl_uses_default_and_allows_per_call_override() {
|
||||
let cache = TestCache {
|
||||
default_ttl: Duration::from_secs(60),
|
||||
writes: Mutex::default(),
|
||||
};
|
||||
assert_eq!(
|
||||
cache.get_ttl(&CacheKwargs::default()),
|
||||
Duration::from_secs(60)
|
||||
cache.get_ttl(&ExactCacheContext::default()),
|
||||
Some(Duration::from_secs(60))
|
||||
);
|
||||
assert_eq!(
|
||||
cache.get_ttl(&CacheKwargs {
|
||||
cache.get_ttl(&ExactCacheContext {
|
||||
ttl: Some(Duration::from_secs(5)),
|
||||
..Default::default()
|
||||
}),
|
||||
Duration::from_secs(5)
|
||||
Some(Duration::from_secs(5))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn keys_match_python_order_groups_files_presets_and_namespaces() {
|
||||
let mut input = CacheKeyInput {
|
||||
fields: vec![
|
||||
CacheKeyField {
|
||||
name: "model".into(),
|
||||
value: Some("deployment".into()),
|
||||
api_parameter: true,
|
||||
internal_parameter: false,
|
||||
},
|
||||
CacheKeyField {
|
||||
name: "file".into(),
|
||||
value: None,
|
||||
api_parameter: true,
|
||||
internal_parameter: false,
|
||||
},
|
||||
],
|
||||
namespace: Some("team".into()),
|
||||
..Default::default()
|
||||
fn associated_context_preserves_backend_specific_lookup_inputs() {
|
||||
let context = SemanticContext {
|
||||
ttl: None,
|
||||
query: "matching prompt".into(),
|
||||
};
|
||||
CacheKeyContext {
|
||||
model_group: Some("group".into()),
|
||||
caching_groups: vec![(vec!["group".into()], "['group']".into())],
|
||||
file_checksum: Some("checksum".into()),
|
||||
..Default::default()
|
||||
}
|
||||
.apply(&mut input);
|
||||
assert_eq!(
|
||||
cache_key(&input),
|
||||
format!(
|
||||
"team:{:x}",
|
||||
Sha256::digest(b"model: ['group']file: checksum")
|
||||
)
|
||||
get_cache(&SemanticCache, "shared-key", &context).unwrap(),
|
||||
Some("semantic hit".into())
|
||||
);
|
||||
input.preset = Some("preset".into());
|
||||
assert_eq!(get_cache_key(&input), "preset");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_controls_honor_default_modes_and_directives() {
|
||||
let enabled = CacheControls {
|
||||
supported_call_type: true,
|
||||
configured: true,
|
||||
default_on: true,
|
||||
..Default::default()
|
||||
#[tokio::test]
|
||||
async fn default_batch_operations_use_async_writes_and_stop_on_failure() {
|
||||
let cache = TestCache {
|
||||
default_ttl: Duration::from_secs(60),
|
||||
writes: Mutex::default(),
|
||||
};
|
||||
assert!(enabled.reads());
|
||||
assert!(enabled.writes());
|
||||
assert!(
|
||||
!CacheControls {
|
||||
default_on: false,
|
||||
..enabled
|
||||
}
|
||||
.reads()
|
||||
let entry = String::from("cached");
|
||||
let context = ExactCacheContext {
|
||||
ttl: Some(Duration::from_secs(5)),
|
||||
};
|
||||
cache
|
||||
.batch_cache_write("single", entry.clone(), context.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_set_cache_pipeline(
|
||||
vec![
|
||||
("first".into(), entry.clone()),
|
||||
("unavailable".into(), entry.clone()),
|
||||
("skipped".into(), entry.clone()),
|
||||
],
|
||||
context.clone(),
|
||||
)
|
||||
.await,
|
||||
Err(Error::Unavailable)
|
||||
);
|
||||
assert!(
|
||||
CacheControls {
|
||||
default_on: false,
|
||||
use_cache: true,
|
||||
..enabled
|
||||
}
|
||||
.reads()
|
||||
);
|
||||
assert!(
|
||||
!CacheControls {
|
||||
no_cache: true,
|
||||
..enabled
|
||||
}
|
||||
.reads()
|
||||
);
|
||||
assert!(
|
||||
!CacheControls {
|
||||
no_store: true,
|
||||
..enabled
|
||||
}
|
||||
.writes()
|
||||
assert_eq!(
|
||||
*cache.writes.lock().unwrap(),
|
||||
vec![
|
||||
("single".into(), entry.clone(), context.clone()),
|
||||
("first".into(), entry, context),
|
||||
]
|
||||
);
|
||||
}
|
||||
|
|
|
|||
41
litellm-rust/crates/cache/tests/codec.rs
vendored
Normal file
41
litellm-rust/crates/cache/tests/codec.rs
vendored
Normal file
|
|
@ -0,0 +1,41 @@
|
|||
use std::collections::BTreeMap;
|
||||
|
||||
use litellm_cache::{CacheCodec, Error, JsonCodec};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::json;
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
struct RoutingState {
|
||||
deployment: String,
|
||||
cooldown_seconds: u64,
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn json_codec_round_trips_typed_domain_values() {
|
||||
let codec = JsonCodec::<RoutingState>::new();
|
||||
let value = RoutingState {
|
||||
deployment: "deployment-a".into(),
|
||||
cooldown_seconds: 30,
|
||||
};
|
||||
let bytes = codec.encode(&value).unwrap();
|
||||
assert_eq!(codec.decode(&bytes).unwrap(), value);
|
||||
assert_eq!(
|
||||
serde_json::from_slice::<serde_json::Value>(&bytes).unwrap(),
|
||||
json!({"deployment": "deployment-a", "cooldown_seconds": 30})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn json_codec_rejects_malformed_and_wrongly_typed_entries() {
|
||||
let codec = JsonCodec::<RoutingState>::new();
|
||||
for bytes in [b"not json".as_slice(), br#"{"deployment":12}"#.as_slice()] {
|
||||
assert_eq!(codec.decode(bytes).unwrap_err(), Error::InvalidEntry);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn json_codec_propagates_encoding_errors() {
|
||||
let codec = JsonCodec::<BTreeMap<(u8, u8), String>>::new();
|
||||
let value = BTreeMap::from([((1, 2), "invalid JSON object key".into())]);
|
||||
assert_eq!(codec.encode(&value).unwrap_err(), Error::InvalidEntry);
|
||||
}
|
||||
385
litellm-rust/crates/cache/tests/dual.rs
vendored
Normal file
385
litellm-rust/crates/cache/tests/dual.rs
vendored
Normal file
|
|
@ -0,0 +1,385 @@
|
|||
use std::{
|
||||
sync::{Arc, Mutex},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use litellm_cache::{
|
||||
BaseCache, BatchCache, CacheConnectionResult, ClaimCache, CounterCache, DeleteCache, DualCache,
|
||||
Error, ExactCacheContext, FlushCache, ReadPolicy, RemoteFailurePolicy, WritePolicy,
|
||||
};
|
||||
|
||||
struct TestCache<V> {
|
||||
value: Mutex<Option<V>>,
|
||||
fail: bool,
|
||||
}
|
||||
|
||||
impl<V> TestCache<V> {
|
||||
fn new(value: Option<V>, fail: bool) -> Self {
|
||||
Self {
|
||||
value: Mutex::new(value),
|
||||
fail,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<V> BaseCache for TestCache<V>
|
||||
where
|
||||
V: Clone + Send + Sync + 'static,
|
||||
{
|
||||
type Value = V;
|
||||
type Context = ExactCacheContext;
|
||||
|
||||
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
|
||||
context.ttl.or(Some(Duration::from_secs(60)))
|
||||
}
|
||||
|
||||
fn set_cache(&self, _: &str, value: V, _: &ExactCacheContext) -> Result<(), Error> {
|
||||
*self.value.lock().unwrap() = Some(value);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn get_cache(&self, _: &str, _: &ExactCacheContext) -> Result<Option<V>, Error> {
|
||||
Ok(self.value.lock().unwrap().clone())
|
||||
}
|
||||
|
||||
async fn disconnect(&self) -> Result<(), Error> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
|
||||
unreachable!()
|
||||
}
|
||||
}
|
||||
|
||||
impl<V> BatchCache for TestCache<V> where V: Clone + Send + Sync + 'static {}
|
||||
|
||||
impl<V> DeleteCache for TestCache<V>
|
||||
where
|
||||
V: Clone + Send + Sync + 'static,
|
||||
{
|
||||
fn delete_cache(&self, _: &str) -> Result<(), Error> {
|
||||
*self.value.lock().unwrap() = None;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl<V> FlushCache for TestCache<V>
|
||||
where
|
||||
V: Clone + Send + Sync + 'static,
|
||||
{
|
||||
fn flush_cache(&self) -> Result<(), Error> {
|
||||
*self.value.lock().unwrap() = None;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl CounterCache for TestCache<f64> {
|
||||
fn increment_cache(&self, _: &str, amount: f64, _: ExactCacheContext) -> Result<f64, Error> {
|
||||
if self.fail {
|
||||
return Err(Error::Unavailable);
|
||||
}
|
||||
let mut value = self.value.lock().unwrap();
|
||||
let incremented = value.unwrap_or_default() + amount;
|
||||
*value = Some(incremented);
|
||||
Ok(incremented)
|
||||
}
|
||||
}
|
||||
|
||||
impl<V> ClaimCache for TestCache<V>
|
||||
where
|
||||
V: Clone + PartialEq + Send + Sync + 'static,
|
||||
{
|
||||
fn claim_cache(
|
||||
&self,
|
||||
_: &str,
|
||||
candidate: V,
|
||||
eligible: &[V],
|
||||
_: ExactCacheContext,
|
||||
) -> Result<V, Error> {
|
||||
if self.fail {
|
||||
return Err(Error::Unavailable);
|
||||
}
|
||||
let mut value = self.value.lock().unwrap();
|
||||
let winner = match value.as_ref() {
|
||||
Some(existing) if eligible.is_empty() || eligible.contains(existing) => {
|
||||
existing.clone()
|
||||
}
|
||||
_ => candidate,
|
||||
};
|
||||
*value = Some(winner.clone());
|
||||
Ok(winner)
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn failed_l2_increment_leaves_l1_unchanged() {
|
||||
let l1 = Arc::new(TestCache::new(Some(10.0), false));
|
||||
let cache = DualCache::new(l1.clone(), Arc::new(TestCache::new(Some(20.0), true)));
|
||||
|
||||
assert_eq!(
|
||||
cache.increment_cache("counter", 2.0, ExactCacheContext::default()),
|
||||
Err(Error::Unavailable)
|
||||
);
|
||||
assert_eq!(
|
||||
l1.get_cache("counter", &ExactCacheContext::default())
|
||||
.unwrap(),
|
||||
Some(10.0)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn claim_uses_l1_fallback_without_overwriting_an_eligible_winner() {
|
||||
let l1 = Arc::new(TestCache::new(Some("first".to_string()), false));
|
||||
let cache = DualCache::new(l1, Arc::new(TestCache::new(None, true)))
|
||||
.with_remote_failure_policy(RemoteFailurePolicy::UseLocal);
|
||||
|
||||
assert_eq!(
|
||||
cache
|
||||
.claim_cache(
|
||||
"affinity",
|
||||
"second".into(),
|
||||
&["first".into(), "second".into()],
|
||||
ExactCacheContext {
|
||||
ttl: Some(Duration::from_secs(60)),
|
||||
},
|
||||
)
|
||||
.unwrap(),
|
||||
"first"
|
||||
);
|
||||
}
|
||||
|
||||
struct SyncPanics(TestCache<String>);
|
||||
|
||||
impl BaseCache for SyncPanics {
|
||||
type Value = String;
|
||||
type Context = ExactCacheContext;
|
||||
|
||||
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
|
||||
self.0.get_ttl(context)
|
||||
}
|
||||
|
||||
fn set_cache(&self, _: &str, _: String, _: &ExactCacheContext) -> Result<(), Error> {
|
||||
panic!("sync L2 write on an async path")
|
||||
}
|
||||
|
||||
fn get_cache(&self, _: &str, _: &ExactCacheContext) -> Result<Option<String>, Error> {
|
||||
panic!("sync L2 read on an async path")
|
||||
}
|
||||
|
||||
async fn async_set_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
value: String,
|
||||
context: ExactCacheContext,
|
||||
) -> Result<(), Error> {
|
||||
self.0.set_cache(key, value, &context)
|
||||
}
|
||||
|
||||
async fn async_get_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
context: &ExactCacheContext,
|
||||
) -> Result<Option<String>, Error> {
|
||||
self.0.get_cache(key, context)
|
||||
}
|
||||
|
||||
async fn async_set_cache_pipeline(
|
||||
&self,
|
||||
cache_list: Vec<(String, String)>,
|
||||
context: ExactCacheContext,
|
||||
) -> Result<(), Error> {
|
||||
for (key, value) in cache_list {
|
||||
self.0.set_cache(&key, value, &context)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn disconnect(&self) -> Result<(), Error> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
|
||||
unreachable!()
|
||||
}
|
||||
}
|
||||
|
||||
impl BatchCache for SyncPanics {
|
||||
async fn async_batch_get_cache(
|
||||
&self,
|
||||
keys: Vec<String>,
|
||||
context: ExactCacheContext,
|
||||
) -> Result<Vec<litellm_cache::BatchEntry<String>>, Error> {
|
||||
assert_eq!(keys, ["missing"]);
|
||||
Ok(vec![match self.0.get_cache("missing", &context)? {
|
||||
Some(value) => litellm_cache::BatchEntry::Hit(value),
|
||||
None => litellm_cache::BatchEntry::Miss,
|
||||
}])
|
||||
}
|
||||
}
|
||||
|
||||
impl DeleteCache for SyncPanics {
|
||||
fn delete_cache(&self, _: &str) -> Result<(), Error> {
|
||||
panic!("sync L2 delete on an async path")
|
||||
}
|
||||
|
||||
async fn async_delete_cache(&self, key: &str) -> Result<(), Error> {
|
||||
self.0.delete_cache(key)
|
||||
}
|
||||
}
|
||||
|
||||
impl FlushCache for SyncPanics {
|
||||
fn flush_cache(&self) -> Result<(), Error> {
|
||||
panic!("sync L2 flush on an async path")
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn async_operations_use_the_async_l2_methods() {
|
||||
let l1 = Arc::new(TestCache::new(None, false));
|
||||
let cache = DualCache::new(
|
||||
l1.clone(),
|
||||
Arc::new(SyncPanics(TestCache::new(
|
||||
Some("remote".to_string()),
|
||||
false,
|
||||
))),
|
||||
);
|
||||
let context = ExactCacheContext::default();
|
||||
|
||||
assert_eq!(
|
||||
cache.async_get_cache("missing", &context).await.unwrap(),
|
||||
Some("remote".into())
|
||||
);
|
||||
assert_eq!(
|
||||
l1.get_cache("missing", &context).unwrap(),
|
||||
Some("remote".into())
|
||||
);
|
||||
|
||||
l1.delete_cache("missing").unwrap();
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_batch_get_cache(vec!["missing".into()], context.clone())
|
||||
.await
|
||||
.unwrap(),
|
||||
[litellm_cache::BatchEntry::Hit("remote".to_string())]
|
||||
);
|
||||
cache
|
||||
.async_set_cache("missing", "written".into(), context.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
cache
|
||||
.async_set_cache_pipeline(vec![("missing".into(), "piped".into())], context.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
cache.async_delete_cache("missing").await.unwrap();
|
||||
assert_eq!(
|
||||
cache.async_get_cache("missing", &context).await.unwrap(),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
struct Unavailable;
|
||||
|
||||
impl BaseCache for Unavailable {
|
||||
type Value = String;
|
||||
type Context = ExactCacheContext;
|
||||
|
||||
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
|
||||
context.ttl
|
||||
}
|
||||
|
||||
fn set_cache(&self, _: &str, _: String, _: &ExactCacheContext) -> Result<(), Error> {
|
||||
Err(Error::Unavailable)
|
||||
}
|
||||
|
||||
fn get_cache(&self, _: &str, _: &ExactCacheContext) -> Result<Option<String>, Error> {
|
||||
Err(Error::Unavailable)
|
||||
}
|
||||
|
||||
async fn disconnect(&self) -> Result<(), Error> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
|
||||
unreachable!()
|
||||
}
|
||||
}
|
||||
|
||||
impl BatchCache for Unavailable {}
|
||||
|
||||
impl DeleteCache for Unavailable {
|
||||
fn delete_cache(&self, _: &str) -> Result<(), Error> {
|
||||
Err(Error::Unavailable)
|
||||
}
|
||||
}
|
||||
|
||||
impl FlushCache for Unavailable {
|
||||
fn flush_cache(&self) -> Result<(), Error> {
|
||||
Err(Error::Unavailable)
|
||||
}
|
||||
}
|
||||
|
||||
impl ClaimCache for Unavailable {
|
||||
fn claim_cache(
|
||||
&self,
|
||||
_: &str,
|
||||
_: String,
|
||||
_: &[String],
|
||||
_: ExactCacheContext,
|
||||
) -> Result<String, Error> {
|
||||
Err(Error::InvalidEntry)
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn remote_failure_policy_selects_propagation_or_the_local_tier() {
|
||||
let context = ExactCacheContext::default();
|
||||
let strict = DualCache::new(Arc::new(TestCache::new(None, false)), Arc::new(Unavailable));
|
||||
assert_eq!(
|
||||
strict.set_cache("key", "value".into(), &context),
|
||||
Err(Error::Unavailable)
|
||||
);
|
||||
assert_eq!(strict.get_cache("key", &context), Err(Error::Unavailable));
|
||||
|
||||
let l1 = Arc::new(TestCache::new(None, false));
|
||||
let degraded = DualCache::new(l1.clone(), Arc::new(Unavailable))
|
||||
.with_remote_failure_policy(RemoteFailurePolicy::UseLocal);
|
||||
assert_eq!(degraded.get_cache("key", &context), Ok(None));
|
||||
degraded.set_cache("key", "value".into(), &context).unwrap();
|
||||
assert_eq!(
|
||||
degraded.get_cache("key", &context),
|
||||
Ok(Some("value".into()))
|
||||
);
|
||||
degraded.delete_cache("key").unwrap();
|
||||
assert_eq!(l1.get_cache("key", &context), Ok(None));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn claim_fallback_does_not_hide_non_availability_errors() {
|
||||
let cache = DualCache::new(
|
||||
Arc::new(TestCache::new(Some("first".to_string()), false)),
|
||||
Arc::new(Unavailable),
|
||||
)
|
||||
.with_remote_failure_policy(RemoteFailurePolicy::UseLocal);
|
||||
assert_eq!(
|
||||
cache.claim_cache(
|
||||
"affinity",
|
||||
"second".into(),
|
||||
&[],
|
||||
ExactCacheContext::default()
|
||||
),
|
||||
Err(Error::InvalidEntry)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn local_only_policies_never_touch_l2() {
|
||||
let l2 = Arc::new(TestCache::new(Some("remote".to_string()), false));
|
||||
let cache = DualCache::new(Arc::new(TestCache::new(None, false)), l2.clone())
|
||||
.with_read_policy(ReadPolicy::LocalOnly)
|
||||
.with_write_policy(WritePolicy::LocalOnly);
|
||||
let context = ExactCacheContext::default();
|
||||
|
||||
assert_eq!(cache.get_cache("key", &context), Ok(None));
|
||||
cache.set_cache("key", "local".into(), &context).unwrap();
|
||||
assert_eq!(l2.get_cache("key", &context), Ok(Some("remote".into())));
|
||||
}
|
||||
|
|
@ -10,13 +10,21 @@ name = "_native"
|
|||
crate-type = ["cdylib"]
|
||||
|
||||
[features]
|
||||
default = ["abi3"]
|
||||
default = ["abi3", "fast"]
|
||||
abi3 = ["pyo3/abi3-py310"]
|
||||
extension-module = ["pyo3/extension-module"]
|
||||
panic-test = []
|
||||
fast = ["litellm-token-counter/fast"]
|
||||
huggingface = ["litellm-token-counter/huggingface"]
|
||||
tiktoken = ["litellm-token-counter/tiktoken"]
|
||||
|
||||
[dependencies]
|
||||
bytes.workspace = true
|
||||
litellm-cache.workspace = true
|
||||
litellm-cache-memory.workspace = true
|
||||
litellm-cache-redis.workspace = true
|
||||
litellm-cache-response.workspace = true
|
||||
serde.workspace = true
|
||||
litellm-auth.workspace = true
|
||||
litellm-callbacks-legacy-python.workspace = true
|
||||
litellm-core.workspace = true
|
||||
|
|
@ -26,7 +34,7 @@ litellm-http.workspace = true
|
|||
litellm-llms.workspace = true
|
||||
litellm-types.workspace = true
|
||||
litellm-host-python.workspace = true
|
||||
litellm-token-counter.workspace = true
|
||||
litellm-token-counter = { path = "../token-counter", default-features = false }
|
||||
pyo3.workspace = true
|
||||
pyo3-async-runtimes.workspace = true
|
||||
serde_json.workspace = true
|
||||
|
|
|
|||
291
litellm-rust/crates/python-bridge/src/cache/binding.rs
vendored
Normal file
291
litellm-rust/crates/python-bridge/src/cache/binding.rs
vendored
Normal file
|
|
@ -0,0 +1,291 @@
|
|||
use litellm_cache_response::PartialHits;
|
||||
use litellm_host_python::{ExecutionStep, from_py, release_gil, run_async, to_py};
|
||||
use pyo3::{
|
||||
PyTraverseError, PyVisit,
|
||||
exceptions::{PyRuntimeError, PyValueError},
|
||||
prelude::*,
|
||||
types::PyDict,
|
||||
};
|
||||
use serde_json::Value;
|
||||
|
||||
use super::{
|
||||
cache_error,
|
||||
callback::PythonCallback,
|
||||
future::{ready_none, ready_value},
|
||||
native::NativeResponseCache,
|
||||
request::{now, request, requests},
|
||||
};
|
||||
|
||||
pub(super) enum CacheBinding {
|
||||
Disabled,
|
||||
Native(NativeResponseCache),
|
||||
PythonCallback(PythonCallback),
|
||||
}
|
||||
|
||||
#[pyclass(frozen, name = "_CacheTestBinding")]
|
||||
pub(crate) struct ResolvedCache {
|
||||
binding: CacheBinding,
|
||||
pid: u32,
|
||||
}
|
||||
|
||||
impl ResolvedCache {
|
||||
pub(super) fn new(binding: CacheBinding) -> Self {
|
||||
Self {
|
||||
binding,
|
||||
pid: std::process::id(),
|
||||
}
|
||||
}
|
||||
|
||||
fn check_process(&self) -> PyResult<()> {
|
||||
if matches!(self.binding, CacheBinding::Native(_)) && self.pid != std::process::id() {
|
||||
return Err(PyRuntimeError::new_err(
|
||||
"native cache bindings must be resolved again after fork",
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) fn lookup_step(
|
||||
&self,
|
||||
py: Python<'_>,
|
||||
input: &Bound<'_, PyAny>,
|
||||
kwargs: Option<&Bound<'_, PyDict>>,
|
||||
) -> PyResult<ExecutionStep> {
|
||||
self.check_process()?;
|
||||
let awaitable = match &self.binding {
|
||||
CacheBinding::Disabled => ready_none(py)?,
|
||||
CacheBinding::Native(service) => {
|
||||
let request = request(input)?;
|
||||
let service = service.clone();
|
||||
run_async(
|
||||
py,
|
||||
async move { service.async_lookup(&request, now()).await },
|
||||
cache_error,
|
||||
)?
|
||||
}
|
||||
CacheBinding::PythonCallback(callback) => callback.async_lookup(py, kwargs)?,
|
||||
};
|
||||
Ok(ExecutionStep::Await(awaitable.unbind()))
|
||||
}
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl ResolvedCache {
|
||||
#[getter]
|
||||
fn kind(&self) -> &'static str {
|
||||
match self.binding {
|
||||
CacheBinding::Disabled => "disabled",
|
||||
CacheBinding::Native(_) => "native",
|
||||
CacheBinding::PythonCallback(_) => "python_callback",
|
||||
}
|
||||
}
|
||||
|
||||
#[pyo3(signature = (request, *, callback_kwargs=None))]
|
||||
fn lookup(
|
||||
&self,
|
||||
py: Python<'_>,
|
||||
request: &Bound<'_, PyAny>,
|
||||
callback_kwargs: Option<&Bound<'_, PyDict>>,
|
||||
) -> PyResult<Py<PyAny>> {
|
||||
self.check_process()?;
|
||||
match &self.binding {
|
||||
CacheBinding::Disabled => Ok(py.None()),
|
||||
CacheBinding::Native(service) => {
|
||||
let request = self::request(request)?;
|
||||
let service = service.clone();
|
||||
let response = release_gil(py, move || service.lookup(&request, now()))
|
||||
.map_err(cache_error)?;
|
||||
to_py(py, &response)
|
||||
}
|
||||
CacheBinding::PythonCallback(callback) => {
|
||||
callback.lookup(py, callback_kwargs).map(Bound::unbind)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[pyo3(signature = (request, response, *, callback_kwargs=None))]
|
||||
fn store(
|
||||
&self,
|
||||
py: Python<'_>,
|
||||
request: &Bound<'_, PyAny>,
|
||||
response: &Bound<'_, PyAny>,
|
||||
callback_kwargs: Option<&Bound<'_, PyDict>>,
|
||||
) -> PyResult<()> {
|
||||
self.check_process()?;
|
||||
match &self.binding {
|
||||
CacheBinding::Disabled => Ok(()),
|
||||
CacheBinding::Native(service) => {
|
||||
let request = self::request(request)?;
|
||||
let response: Value = from_py(response)?;
|
||||
let service = service.clone();
|
||||
release_gil(py, move || service.store(&request, response, now()))
|
||||
.map_err(cache_error)
|
||||
}
|
||||
CacheBinding::PythonCallback(callback) => callback.store(py, response, callback_kwargs),
|
||||
}
|
||||
}
|
||||
|
||||
#[pyo3(signature = (requests, *, callback_kwargs=None))]
|
||||
fn lookup_batch(
|
||||
&self,
|
||||
py: Python<'_>,
|
||||
requests: &Bound<'_, PyAny>,
|
||||
callback_kwargs: Option<&Bound<'_, PyAny>>,
|
||||
) -> PyResult<Py<PyAny>> {
|
||||
self.check_process()?;
|
||||
match &self.binding {
|
||||
CacheBinding::Disabled => {
|
||||
let requests = self::requests(requests)?;
|
||||
to_py(py, &PartialHits::new(vec![None; requests.len()]))
|
||||
}
|
||||
CacheBinding::Native(service) => {
|
||||
let requests = self::requests(requests)?;
|
||||
let service = service.clone();
|
||||
let response = release_gil(py, move || service.lookup_batch(&requests, now()))
|
||||
.map_err(cache_error)?;
|
||||
to_py(py, &response)
|
||||
}
|
||||
CacheBinding::PythonCallback(callback) => callback
|
||||
.lookup_batch(py, requests, callback_kwargs)
|
||||
.map(Bound::unbind),
|
||||
}
|
||||
}
|
||||
|
||||
#[pyo3(signature = (request, *, callback_kwargs=None))]
|
||||
fn async_lookup<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
request: &Bound<'py, PyAny>,
|
||||
callback_kwargs: Option<&Bound<'py, PyDict>>,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let ExecutionStep::Await(awaitable) = self.lookup_step(py, request, callback_kwargs)?
|
||||
else {
|
||||
unreachable!()
|
||||
};
|
||||
Ok(awaitable.into_bound(py))
|
||||
}
|
||||
|
||||
#[pyo3(signature = (request, response, *, callback_kwargs=None))]
|
||||
fn async_store<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
request: &Bound<'py, PyAny>,
|
||||
response: &Bound<'py, PyAny>,
|
||||
callback_kwargs: Option<&Bound<'py, PyDict>>,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
self.check_process()?;
|
||||
match &self.binding {
|
||||
CacheBinding::Disabled => ready_none(py),
|
||||
CacheBinding::Native(service) => {
|
||||
let request = self::request(request)?;
|
||||
let response: Value = from_py(response)?;
|
||||
let service = service.clone();
|
||||
run_async(
|
||||
py,
|
||||
async move { service.async_store(&request, response, now()).await },
|
||||
cache_error,
|
||||
)
|
||||
}
|
||||
CacheBinding::PythonCallback(callback) => {
|
||||
callback.async_store(py, response, callback_kwargs)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[pyo3(signature = (requests, *, callback_kwargs=None))]
|
||||
fn async_lookup_batch<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
requests: &Bound<'py, PyAny>,
|
||||
callback_kwargs: Option<&Bound<'py, PyAny>>,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
self.check_process()?;
|
||||
match &self.binding {
|
||||
CacheBinding::Disabled => {
|
||||
let requests = self::requests(requests)?;
|
||||
ready_value(py, &PartialHits::new(vec![None; requests.len()]))
|
||||
}
|
||||
CacheBinding::Native(service) => {
|
||||
let requests = self::requests(requests)?;
|
||||
let service = service.clone();
|
||||
run_async(
|
||||
py,
|
||||
async move { service.async_lookup_batch(&requests, now()).await },
|
||||
cache_error,
|
||||
)
|
||||
}
|
||||
CacheBinding::PythonCallback(callback) => {
|
||||
callback.async_lookup_batch(py, requests, callback_kwargs)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[pyo3(signature = (requests, responses, *, callback_result=None, callback_kwargs=None))]
|
||||
fn async_store_batch<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
requests: &Bound<'py, PyAny>,
|
||||
responses: &Bound<'py, PyAny>,
|
||||
callback_result: Option<&Bound<'py, PyAny>>,
|
||||
callback_kwargs: Option<&Bound<'py, PyDict>>,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
self.check_process()?;
|
||||
match &self.binding {
|
||||
CacheBinding::Disabled => ready_none(py),
|
||||
CacheBinding::Native(service) => {
|
||||
let requests = self::requests(requests)?;
|
||||
let responses: Vec<Value> = from_py(responses)?;
|
||||
if requests.len() != responses.len() {
|
||||
return Err(PyValueError::new_err(
|
||||
"batch cache requests and responses must have equal lengths",
|
||||
));
|
||||
}
|
||||
let entries = requests.into_iter().zip(responses).collect();
|
||||
let service = service.clone();
|
||||
run_async(
|
||||
py,
|
||||
async move { service.async_store_batch(entries, now()).await },
|
||||
cache_error,
|
||||
)
|
||||
}
|
||||
CacheBinding::PythonCallback(callback) => {
|
||||
callback.async_store_batch(py, callback_result, callback_kwargs)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn async_flush<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
|
||||
self.check_process()?;
|
||||
match &self.binding {
|
||||
CacheBinding::Disabled => ready_none(py),
|
||||
CacheBinding::Native(service) => {
|
||||
let service = service.clone();
|
||||
run_async(py, async move { service.async_flush().await }, cache_error)
|
||||
}
|
||||
CacheBinding::PythonCallback(callback) => callback.async_flush(py),
|
||||
}
|
||||
}
|
||||
|
||||
fn ping<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
|
||||
self.check_process()?;
|
||||
match &self.binding {
|
||||
CacheBinding::Disabled => ready_none(py),
|
||||
CacheBinding::Native(service) => {
|
||||
let service = service.clone();
|
||||
run_async(
|
||||
py,
|
||||
async move { service.test_connection().await },
|
||||
cache_error,
|
||||
)
|
||||
}
|
||||
CacheBinding::PythonCallback(callback) => callback.ping(py),
|
||||
}
|
||||
}
|
||||
|
||||
fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
|
||||
if let CacheBinding::PythonCallback(callback) = &self.binding {
|
||||
callback.traverse(&visit)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
162
litellm-rust/crates/python-bridge/src/cache/callback.rs
vendored
Normal file
162
litellm-rust/crates/python-bridge/src/cache/callback.rs
vendored
Normal file
|
|
@ -0,0 +1,162 @@
|
|||
use pyo3::{
|
||||
PyTraverseError, PyVisit,
|
||||
exceptions::{PyTypeError, PyValueError},
|
||||
prelude::*,
|
||||
types::{PyDict, PyList, PyTuple},
|
||||
};
|
||||
|
||||
use super::future::ready_none;
|
||||
|
||||
pub(super) struct PythonCallback(Py<PyAny>);
|
||||
|
||||
impl PythonCallback {
|
||||
pub(super) fn new(object: Py<PyAny>) -> Self {
|
||||
Self(object)
|
||||
}
|
||||
|
||||
pub(super) fn lookup<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
kwargs: Option<&Bound<'py, PyDict>>,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
self.0
|
||||
.bind(py)
|
||||
.call_method("get_cache", (), Some(callback_kwargs(kwargs)?))
|
||||
}
|
||||
|
||||
pub(super) fn async_lookup<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
kwargs: Option<&Bound<'py, PyDict>>,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
self.0
|
||||
.bind(py)
|
||||
.call_method("async_get_cache", (), Some(callback_kwargs(kwargs)?))
|
||||
}
|
||||
|
||||
pub(super) fn store(
|
||||
&self,
|
||||
py: Python<'_>,
|
||||
response: &Bound<'_, PyAny>,
|
||||
kwargs: Option<&Bound<'_, PyDict>>,
|
||||
) -> PyResult<()> {
|
||||
self.0
|
||||
.bind(py)
|
||||
.call_method("add_cache", (response,), Some(callback_kwargs(kwargs)?))
|
||||
.map(|_| ())
|
||||
}
|
||||
|
||||
pub(super) fn async_store<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
response: &Bound<'py, PyAny>,
|
||||
kwargs: Option<&Bound<'py, PyDict>>,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
self.0.bind(py).call_method(
|
||||
"async_add_cache",
|
||||
(response,),
|
||||
Some(callback_kwargs(kwargs)?),
|
||||
)
|
||||
}
|
||||
|
||||
pub(super) fn lookup_batch<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
requests: &Bound<'py, PyAny>,
|
||||
kwargs: Option<&Bound<'py, PyAny>>,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let results = PyList::empty(py);
|
||||
for kwargs in batch_callback_kwargs(requests, kwargs)? {
|
||||
results.append(
|
||||
self.0
|
||||
.bind(py)
|
||||
.call_method("get_cache", (), Some(&kwargs))?,
|
||||
)?;
|
||||
}
|
||||
Ok(results.into_any())
|
||||
}
|
||||
|
||||
pub(super) fn async_lookup_batch<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
requests: &Bound<'py, PyAny>,
|
||||
kwargs: Option<&Bound<'py, PyAny>>,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let awaitables = batch_callback_kwargs(requests, kwargs)?
|
||||
.iter()
|
||||
.map(|kwargs| {
|
||||
self.0
|
||||
.bind(py)
|
||||
.call_method("async_get_cache", (), Some(kwargs))
|
||||
})
|
||||
.collect::<PyResult<Vec<_>>>()?;
|
||||
py.import("asyncio")?
|
||||
.call_method1("gather", PyTuple::new(py, awaitables)?)
|
||||
}
|
||||
|
||||
pub(super) fn async_store_batch<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
result: Option<&Bound<'py, PyAny>>,
|
||||
kwargs: Option<&Bound<'py, PyDict>>,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let result = result.ok_or_else(|| {
|
||||
PyTypeError::new_err("Python cache callbacks require their original callback_result")
|
||||
})?;
|
||||
self.0.bind(py).call_method(
|
||||
"async_add_cache_pipeline",
|
||||
(result,),
|
||||
Some(callback_kwargs(kwargs)?),
|
||||
)
|
||||
}
|
||||
|
||||
pub(super) fn async_flush<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
|
||||
let object = self.0.bind(py);
|
||||
let backend = match object.getattr_opt("cache")? {
|
||||
Some(backend) if !backend.is_none() => backend,
|
||||
_ => object.clone(),
|
||||
};
|
||||
if backend.hasattr("async_flush_cache")? {
|
||||
return backend.call_method0("async_flush_cache");
|
||||
}
|
||||
backend.call_method0("flush_cache")?;
|
||||
ready_none(py)
|
||||
}
|
||||
|
||||
pub(super) fn ping<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
|
||||
self.0.bind(py).call_method0("ping")
|
||||
}
|
||||
|
||||
pub(super) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
|
||||
visit.call(&self.0)
|
||||
}
|
||||
}
|
||||
|
||||
fn callback_kwargs<'a, 'py>(
|
||||
kwargs: Option<&'a Bound<'py, PyDict>>,
|
||||
) -> PyResult<&'a Bound<'py, PyDict>> {
|
||||
kwargs.ok_or_else(|| {
|
||||
PyTypeError::new_err("Python cache callbacks require their original callback_kwargs")
|
||||
})
|
||||
}
|
||||
|
||||
fn batch_callback_kwargs<'py>(
|
||||
requests: &Bound<'py, PyAny>,
|
||||
kwargs: Option<&Bound<'py, PyAny>>,
|
||||
) -> PyResult<Vec<Bound<'py, PyDict>>> {
|
||||
let kwargs = kwargs
|
||||
.ok_or_else(|| {
|
||||
PyTypeError::new_err(
|
||||
"Python cache callbacks require one original callback_kwargs mapping per request",
|
||||
)
|
||||
})?
|
||||
.try_iter()?
|
||||
.map(|item| Ok(item?.cast_into::<PyDict>()?))
|
||||
.collect::<PyResult<Vec<_>>>()?;
|
||||
if kwargs.len() != requests.len()? {
|
||||
return Err(PyValueError::new_err(
|
||||
"batch cache requests and callback_kwargs must have equal lengths",
|
||||
));
|
||||
}
|
||||
Ok(kwargs)
|
||||
}
|
||||
594
litellm-rust/crates/python-bridge/src/cache/config.rs
vendored
Normal file
594
litellm-rust/crates/python-bridge/src/cache/config.rs
vendored
Normal file
|
|
@ -0,0 +1,594 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use litellm_cache::CacheType;
|
||||
use pyo3::{
|
||||
exceptions::{PyTypeError, PyValueError},
|
||||
prelude::*,
|
||||
types::{PyAny, PyDict, PyString},
|
||||
};
|
||||
|
||||
use super::{native::NativeResponseCache, request::duration};
|
||||
|
||||
#[allow(dead_code, reason = "consumed by the cache activation follow-up")]
|
||||
pub(super) struct CachePolicy {
|
||||
pub(super) mode: String,
|
||||
pub(super) ttl: Option<Duration>,
|
||||
pub(super) namespace: Option<String>,
|
||||
pub(super) supported_call_types: Option<Vec<String>>,
|
||||
pub(super) redis_flush_size: Option<usize>,
|
||||
pub(super) semantic_cache_scope: String,
|
||||
}
|
||||
|
||||
pub(super) struct MemoryCacheConfig {
|
||||
pub(super) default_ttl: Duration,
|
||||
pub(super) capacity: usize,
|
||||
pub(super) max_entry_bytes: usize,
|
||||
}
|
||||
|
||||
#[derive(Debug, PartialEq)]
|
||||
pub(super) enum RedisProtocol {
|
||||
Resp2,
|
||||
Resp3,
|
||||
}
|
||||
|
||||
#[derive(Debug, PartialEq)]
|
||||
pub(super) enum CertificateRequirement {
|
||||
None,
|
||||
Optional,
|
||||
Required,
|
||||
}
|
||||
|
||||
#[allow(dead_code, reason = "consumed by the cache activation follow-up")]
|
||||
pub(super) struct RedisTlsConfig {
|
||||
pub(super) certificate_requirement: CertificateRequirement,
|
||||
pub(super) check_hostname: bool,
|
||||
pub(super) ca_certificate: Option<String>,
|
||||
pub(super) ca_data: Option<String>,
|
||||
pub(super) client_certificate: Option<String>,
|
||||
pub(super) client_key: Option<String>,
|
||||
}
|
||||
|
||||
#[allow(dead_code, reason = "consumed by the cache activation follow-up")]
|
||||
pub(super) struct RedisConnectionConfig {
|
||||
pub(super) host: String,
|
||||
pub(super) port: u16,
|
||||
pub(super) database: i64,
|
||||
pub(super) username: Option<String>,
|
||||
pub(super) password: Option<String>,
|
||||
pub(super) protocol: RedisProtocol,
|
||||
pub(super) pool_size: usize,
|
||||
pub(super) read_timeout: Option<Duration>,
|
||||
pub(super) connect_timeout: Option<Duration>,
|
||||
pub(super) socket_keepalive: Option<bool>,
|
||||
pub(super) health_check_interval: Duration,
|
||||
pub(super) client_name: Option<String>,
|
||||
pub(super) tls: Option<RedisTlsConfig>,
|
||||
}
|
||||
|
||||
#[allow(dead_code, reason = "consumed by the cache activation follow-up")]
|
||||
pub(super) struct RedisCacheConfig {
|
||||
pub(super) default_ttl: Duration,
|
||||
pub(super) namespace: Option<String>,
|
||||
pub(super) flush_size: usize,
|
||||
pub(super) connection: RedisConnectionConfig,
|
||||
}
|
||||
|
||||
pub(super) enum CacheBackendConfig {
|
||||
Memory(MemoryCacheConfig),
|
||||
Redis(Box<RedisCacheConfig>),
|
||||
}
|
||||
|
||||
#[allow(dead_code, reason = "consumed by the cache activation follow-up")]
|
||||
pub(super) struct NativeCacheConfig {
|
||||
pub(super) policy: CachePolicy,
|
||||
pub(super) backend: CacheBackendConfig,
|
||||
}
|
||||
|
||||
pub(super) enum UnsupportedCacheConfig {
|
||||
Backend,
|
||||
RedisTopology,
|
||||
RedisCredentials,
|
||||
RedisConnection,
|
||||
RedisOption,
|
||||
}
|
||||
|
||||
impl UnsupportedCacheConfig {
|
||||
pub(super) fn message(&self) -> &'static str {
|
||||
match self {
|
||||
Self::Backend => "native cache backend is not implemented",
|
||||
Self::RedisTopology => "native Redis topology is not implemented",
|
||||
Self::RedisCredentials => "native Redis credentials require Python",
|
||||
Self::RedisConnection => "native Redis connection type is not implemented",
|
||||
Self::RedisOption => "native Redis configuration requires Python",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) enum CacheConfigProjection {
|
||||
Native(Box<NativeCacheConfig>),
|
||||
Unsupported(UnsupportedCacheConfig),
|
||||
}
|
||||
|
||||
impl NativeCacheConfig {
|
||||
#[inline(never)]
|
||||
pub(super) fn project(facade: &Bound<'_, PyAny>) -> PyResult<CacheConfigProjection> {
|
||||
let backend_name = facade.getattr("type")?.extract::<String>()?;
|
||||
let policy = CachePolicy {
|
||||
mode: facade.getattr("mode")?.extract::<String>()?,
|
||||
ttl: optional_duration(facade.getattr("ttl")?)?,
|
||||
namespace: optional_string(facade.getattr("namespace")?)?,
|
||||
supported_call_types: facade
|
||||
.getattr("supported_call_types")?
|
||||
.extract::<Option<Vec<String>>>()?,
|
||||
redis_flush_size: facade
|
||||
.getattr("redis_flush_size")?
|
||||
.extract::<Option<usize>>()?,
|
||||
semantic_cache_scope: facade
|
||||
.getattr("semantic_cache_scope")?
|
||||
.extract::<String>()?,
|
||||
};
|
||||
let backend = facade.getattr("cache")?;
|
||||
match CacheType::from_python_name(&backend_name) {
|
||||
Some(CacheType::Local) => project_memory(&backend).map(|backend| {
|
||||
CacheConfigProjection::Native(Box::new(Self {
|
||||
policy,
|
||||
backend: CacheBackendConfig::Memory(backend),
|
||||
}))
|
||||
}),
|
||||
Some(CacheType::Redis) => match project_redis(&backend)? {
|
||||
Ok(backend) => Ok(CacheConfigProjection::Native(Box::new(Self {
|
||||
policy,
|
||||
backend: CacheBackendConfig::Redis(Box::new(backend)),
|
||||
}))),
|
||||
Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)),
|
||||
},
|
||||
Some(
|
||||
CacheType::RedisSemantic
|
||||
| CacheType::ValkeySemantic
|
||||
| CacheType::S3
|
||||
| CacheType::Disk
|
||||
| CacheType::QdrantSemantic
|
||||
| CacheType::AzureBlob
|
||||
| CacheType::Gcs,
|
||||
)
|
||||
| None => Ok(CacheConfigProjection::Unsupported(
|
||||
UnsupportedCacheConfig::Backend,
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn service_mismatch(&self, service: &NativeResponseCache) -> Option<&'static str> {
|
||||
if service.default_ttl()
|
||||
!= Some(match &self.backend {
|
||||
CacheBackendConfig::Memory(config) => config.default_ttl,
|
||||
CacheBackendConfig::Redis(config) => config.default_ttl,
|
||||
})
|
||||
{
|
||||
return Some("facade and native backend default TTLs must match");
|
||||
}
|
||||
match &self.backend {
|
||||
CacheBackendConfig::Memory(config) if service.kind() != "memory" => {
|
||||
Some("facade and native backend types must match")
|
||||
}
|
||||
CacheBackendConfig::Memory(config) if service.capacity() != Some(config.capacity) => {
|
||||
Some("facade and native backend capacities must match")
|
||||
}
|
||||
CacheBackendConfig::Memory(config)
|
||||
if service.max_entry_bytes() != Some(config.max_entry_bytes) =>
|
||||
{
|
||||
Some("facade and native backend item limits must match")
|
||||
}
|
||||
CacheBackendConfig::Memory(_) => None,
|
||||
CacheBackendConfig::Redis(_) if service.kind() != "redis" => {
|
||||
Some("facade and native backend types must match")
|
||||
}
|
||||
CacheBackendConfig::Redis(config) => (service.namespace()
|
||||
!= config.namespace.as_deref())
|
||||
.then_some("facade and native backend namespaces must match"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[inline(never)]
|
||||
fn project_memory(backend: &Bound<'_, PyAny>) -> PyResult<MemoryCacheConfig> {
|
||||
let max_size_kib = backend.getattr("max_size_per_item")?.extract::<usize>()?;
|
||||
Ok(MemoryCacheConfig {
|
||||
default_ttl: duration(backend.getattr("default_ttl")?.extract::<f64>()?)?,
|
||||
capacity: backend.getattr("max_size_in_memory")?.extract::<usize>()?,
|
||||
max_entry_bytes: max_size_kib
|
||||
.checked_mul(1024)
|
||||
.ok_or_else(|| PyValueError::new_err("memory cache item limit is too large"))?,
|
||||
})
|
||||
}
|
||||
|
||||
#[inline(never)]
|
||||
fn project_redis(
|
||||
backend: &Bound<'_, PyAny>,
|
||||
) -> PyResult<Result<RedisCacheConfig, UnsupportedCacheConfig>> {
|
||||
let source = backend.getattr("redis_kwargs")?.cast_into::<PyDict>()?;
|
||||
if has_value(&source, "startup_nodes")? {
|
||||
return Ok(Err(UnsupportedCacheConfig::RedisTopology));
|
||||
}
|
||||
if has_value(&source, "sentinel_nodes")? {
|
||||
return Ok(Err(UnsupportedCacheConfig::RedisTopology));
|
||||
}
|
||||
for key in ["credential_provider", "redis_connect_func"] {
|
||||
if has_value(&source, key)? {
|
||||
return Ok(Err(UnsupportedCacheConfig::RedisCredentials));
|
||||
}
|
||||
}
|
||||
if has_value(&source, "connection_pool")? {
|
||||
return Ok(Err(UnsupportedCacheConfig::RedisConnection));
|
||||
}
|
||||
for key in [
|
||||
"retry",
|
||||
"retry_on_error",
|
||||
"socket_keepalive_options",
|
||||
"unix_socket_path",
|
||||
"cache",
|
||||
"cache_config",
|
||||
"event_dispatcher",
|
||||
"ssl_ca_path",
|
||||
"ssl_password",
|
||||
"ssl_min_version",
|
||||
"ssl_ciphers",
|
||||
"ssl_validate_ocsp",
|
||||
"ssl_validate_ocsp_stapled",
|
||||
"ssl_ocsp_context",
|
||||
"ssl_ocsp_expected_cert",
|
||||
] {
|
||||
if has_value(&source, key)? {
|
||||
return Ok(Err(UnsupportedCacheConfig::RedisOption));
|
||||
}
|
||||
}
|
||||
for key in ["retry_on_timeout", "single_connection_client"] {
|
||||
if optional_coerced_bool(&source, key)?.unwrap_or(false) {
|
||||
return Ok(Err(UnsupportedCacheConfig::RedisOption));
|
||||
}
|
||||
}
|
||||
|
||||
let client = backend.getattr("redis_client")?;
|
||||
let pool = client.getattr("connection_pool")?;
|
||||
if !instance_class_is(&pool, "redis.connection", "ConnectionPool")? {
|
||||
return Ok(Err(UnsupportedCacheConfig::RedisConnection));
|
||||
}
|
||||
let resolved = pool.getattr("connection_kwargs")?.cast_into::<PyDict>()?;
|
||||
for key in ["credential_provider", "redis_connect_func"] {
|
||||
if has_value(&resolved, key)? {
|
||||
return Ok(Err(UnsupportedCacheConfig::RedisCredentials));
|
||||
}
|
||||
}
|
||||
let connection_class = resolved
|
||||
.get_item("connection_class")?
|
||||
.unwrap_or(pool.getattr("connection_class")?);
|
||||
let tls = if class_is(&connection_class, "redis.connection", "Connection")? {
|
||||
None
|
||||
} else if class_is(&connection_class, "redis.connection", "SSLConnection")? {
|
||||
Some(project_tls(&resolved)?)
|
||||
} else {
|
||||
return Ok(Err(UnsupportedCacheConfig::RedisConnection));
|
||||
};
|
||||
|
||||
let protocol = match optional_i64(&resolved, "protocol")?.unwrap_or(2) {
|
||||
2 => RedisProtocol::Resp2,
|
||||
3 => RedisProtocol::Resp3,
|
||||
_ => return Err(PyValueError::new_err("unsupported Redis protocol version")),
|
||||
};
|
||||
let health_check_interval =
|
||||
duration(optional_f64(&resolved, "health_check_interval")?.unwrap_or(0.0))?;
|
||||
Ok(Ok(RedisCacheConfig {
|
||||
default_ttl: duration(backend.getattr("default_ttl")?.extract::<f64>()?)?,
|
||||
namespace: optional_attribute_string(backend, "namespace")?,
|
||||
flush_size: backend.getattr("redis_flush_size")?.extract::<usize>()?,
|
||||
connection: RedisConnectionConfig {
|
||||
host: required_string(&resolved, "host")?,
|
||||
port: u16::try_from(required_i64(&resolved, "port")?)
|
||||
.map_err(|_| PyValueError::new_err("invalid Redis port"))?,
|
||||
database: optional_i64(&resolved, "db")?.unwrap_or(0),
|
||||
username: optional_dict_string(&resolved, "username")?,
|
||||
password: optional_dict_string(&resolved, "password")?,
|
||||
protocol,
|
||||
pool_size: pool.getattr("max_connections")?.extract::<usize>()?,
|
||||
read_timeout: optional_dict_duration(&resolved, "socket_timeout")?,
|
||||
connect_timeout: optional_dict_duration(&resolved, "socket_connect_timeout")?,
|
||||
socket_keepalive: optional_bool(&resolved, "socket_keepalive")?,
|
||||
health_check_interval,
|
||||
client_name: optional_dict_string(&resolved, "client_name")?,
|
||||
tls,
|
||||
},
|
||||
}))
|
||||
}
|
||||
|
||||
#[inline(never)]
|
||||
fn project_tls(values: &Bound<'_, PyDict>) -> PyResult<RedisTlsConfig> {
|
||||
Ok(RedisTlsConfig {
|
||||
certificate_requirement: certificate_requirement(values)?,
|
||||
check_hostname: optional_bool(values, "ssl_check_hostname")?.unwrap_or(false),
|
||||
ca_certificate: optional_dict_string(values, "ssl_ca_certs")?,
|
||||
ca_data: optional_dict_string(values, "ssl_ca_data")?,
|
||||
client_certificate: optional_dict_string(values, "ssl_certfile")?,
|
||||
client_key: optional_dict_string(values, "ssl_keyfile")?,
|
||||
})
|
||||
}
|
||||
|
||||
#[inline(never)]
|
||||
fn certificate_requirement(values: &Bound<'_, PyDict>) -> PyResult<CertificateRequirement> {
|
||||
let Some(value) = values.get_item("ssl_cert_reqs")? else {
|
||||
return Ok(CertificateRequirement::Required);
|
||||
};
|
||||
if value.is_none() {
|
||||
return Ok(CertificateRequirement::Required);
|
||||
}
|
||||
if let Ok(number) = value.extract::<i32>() {
|
||||
return match number {
|
||||
0 => Ok(CertificateRequirement::None),
|
||||
1 => Ok(CertificateRequirement::Optional),
|
||||
2 => Ok(CertificateRequirement::Required),
|
||||
_ => Err(PyValueError::new_err(
|
||||
"invalid Redis TLS certificate requirement",
|
||||
)),
|
||||
};
|
||||
}
|
||||
let text = value.str()?;
|
||||
let text = text.to_str()?;
|
||||
if text.eq_ignore_ascii_case("none") || text.eq_ignore_ascii_case("cert_none") {
|
||||
return Ok(CertificateRequirement::None);
|
||||
}
|
||||
if text.eq_ignore_ascii_case("optional") || text.eq_ignore_ascii_case("cert_optional") {
|
||||
return Ok(CertificateRequirement::Optional);
|
||||
}
|
||||
if text.eq_ignore_ascii_case("required") || text.eq_ignore_ascii_case("cert_required") {
|
||||
return Ok(CertificateRequirement::Required);
|
||||
}
|
||||
Err(PyValueError::new_err(
|
||||
"invalid Redis TLS certificate requirement",
|
||||
))
|
||||
}
|
||||
|
||||
#[inline(never)]
|
||||
fn instance_class_is(value: &Bound<'_, PyAny>, module: &str, name: &str) -> PyResult<bool> {
|
||||
class_is(value.get_type().as_any(), module, name)
|
||||
}
|
||||
|
||||
#[inline(never)]
|
||||
fn class_is(value: &Bound<'_, PyAny>, module: &str, name: &str) -> PyResult<bool> {
|
||||
Ok(value
|
||||
.getattr("__module__")?
|
||||
.cast_into::<PyString>()?
|
||||
.to_str()?
|
||||
== module
|
||||
&& value
|
||||
.getattr("__qualname__")?
|
||||
.cast_into::<PyString>()?
|
||||
.to_str()?
|
||||
== name)
|
||||
}
|
||||
|
||||
#[inline(never)]
|
||||
fn optional_duration(value: Bound<'_, PyAny>) -> PyResult<Option<Duration>> {
|
||||
value.extract::<Option<f64>>()?.map(duration).transpose()
|
||||
}
|
||||
|
||||
#[inline(never)]
|
||||
fn optional_attribute_string(value: &Bound<'_, PyAny>, name: &str) -> PyResult<Option<String>> {
|
||||
match value.getattr(name) {
|
||||
Ok(value) => optional_string(value),
|
||||
Err(error) if error.is_instance_of::<pyo3::exceptions::PyAttributeError>(value.py()) => {
|
||||
Ok(None)
|
||||
}
|
||||
Err(error) => Err(error),
|
||||
}
|
||||
}
|
||||
|
||||
#[inline(never)]
|
||||
fn optional_string(value: Bound<'_, PyAny>) -> PyResult<Option<String>> {
|
||||
Ok(value
|
||||
.extract::<Option<String>>()?
|
||||
.filter(|value| !value.is_empty()))
|
||||
}
|
||||
|
||||
#[inline(never)]
|
||||
fn has_value(values: &Bound<'_, PyDict>, key: &str) -> PyResult<bool> {
|
||||
Ok(values.get_item(key)?.is_some_and(|value| !value.is_none()))
|
||||
}
|
||||
|
||||
#[inline(never)]
|
||||
fn required_string(values: &Bound<'_, PyDict>, key: &str) -> PyResult<String> {
|
||||
values
|
||||
.get_item(key)?
|
||||
.ok_or_else(|| PyTypeError::new_err("Redis connection is incomplete"))?
|
||||
.extract::<String>()
|
||||
}
|
||||
|
||||
#[inline(never)]
|
||||
fn required_i64(values: &Bound<'_, PyDict>, key: &str) -> PyResult<i64> {
|
||||
values
|
||||
.get_item(key)?
|
||||
.ok_or_else(|| PyTypeError::new_err("Redis connection is incomplete"))?
|
||||
.extract::<i64>()
|
||||
}
|
||||
|
||||
#[inline(never)]
|
||||
fn optional_dict_string(values: &Bound<'_, PyDict>, key: &str) -> PyResult<Option<String>> {
|
||||
match values.get_item(key)? {
|
||||
Some(value) if !value.is_none() => optional_string(value),
|
||||
_ => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
#[inline(never)]
|
||||
fn optional_f64(values: &Bound<'_, PyDict>, key: &str) -> PyResult<Option<f64>> {
|
||||
match values.get_item(key)? {
|
||||
Some(value) => value.extract::<Option<f64>>(),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
#[inline(never)]
|
||||
fn optional_i64(values: &Bound<'_, PyDict>, key: &str) -> PyResult<Option<i64>> {
|
||||
match values.get_item(key)? {
|
||||
Some(value) => value.extract::<Option<i64>>(),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
#[inline(never)]
|
||||
fn optional_bool(values: &Bound<'_, PyDict>, key: &str) -> PyResult<Option<bool>> {
|
||||
match values.get_item(key)? {
|
||||
Some(value) => value.extract::<Option<bool>>(),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
#[inline(never)]
|
||||
fn optional_coerced_bool(values: &Bound<'_, PyDict>, key: &str) -> PyResult<Option<bool>> {
|
||||
let Some(value) = values.get_item(key)? else {
|
||||
return Ok(None);
|
||||
};
|
||||
if value.is_none() {
|
||||
return Ok(None);
|
||||
}
|
||||
if let Ok(text) = value.extract::<String>() {
|
||||
return Ok(Some(
|
||||
text == "1" || text.eq_ignore_ascii_case("true") || text.eq_ignore_ascii_case("yes"),
|
||||
));
|
||||
}
|
||||
value.extract::<bool>().map(Some)
|
||||
}
|
||||
|
||||
#[inline(never)]
|
||||
fn optional_dict_duration(values: &Bound<'_, PyDict>, key: &str) -> PyResult<Option<Duration>> {
|
||||
optional_f64(values, key)?.map(duration).transpose()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::ffi::CString;
|
||||
|
||||
use pyo3::{prelude::*, types::PyDict};
|
||||
|
||||
use super::{
|
||||
CacheBackendConfig, CacheConfigProjection, CertificateRequirement, NativeCacheConfig,
|
||||
RedisProtocol,
|
||||
};
|
||||
use crate::cache::native::NativeResponseCache;
|
||||
|
||||
fn facade<'py>(py: Python<'py>, body: &str) -> Bound<'py, PyAny> {
|
||||
let locals = PyDict::new(py);
|
||||
py.run(
|
||||
&CString::new(format!(
|
||||
"from types import SimpleNamespace\n\
|
||||
ConnectionPool = type('ConnectionPool', (), {{'__module__': 'redis.connection'}})\n\
|
||||
Connection = type('Connection', (), {{'__module__': 'redis.connection'}})\n\
|
||||
SSLConnection = type('SSLConnection', (), {{'__module__': 'redis.connection'}})\n\
|
||||
{body}"
|
||||
))
|
||||
.unwrap(),
|
||||
None,
|
||||
Some(&locals),
|
||||
)
|
||||
.unwrap();
|
||||
locals.get_item("facade").unwrap().unwrap()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn projects_effective_memory_configuration() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let facade = facade(
|
||||
py,
|
||||
"backend = SimpleNamespace(default_ttl=913, max_size_in_memory=37, max_size_per_item=8)\n\
|
||||
facade = SimpleNamespace(type='local', mode='default-on', ttl=11.5, namespace=None, supported_call_types=['completion'], redis_flush_size=None, semantic_cache_scope='key', cache=backend)",
|
||||
);
|
||||
let CacheConfigProjection::Native(config) =
|
||||
NativeCacheConfig::project(&facade).unwrap()
|
||||
else {
|
||||
panic!("memory cache should be supported");
|
||||
};
|
||||
assert_eq!(
|
||||
config.policy.ttl.unwrap(),
|
||||
std::time::Duration::from_secs_f64(11.5)
|
||||
);
|
||||
let CacheBackendConfig::Memory(memory) = config.backend else {
|
||||
panic!("expected memory configuration");
|
||||
};
|
||||
assert_eq!(memory.default_ttl, std::time::Duration::from_secs(913));
|
||||
assert_eq!(memory.capacity, 37);
|
||||
assert_eq!(memory.max_entry_bytes, 8192);
|
||||
let matching =
|
||||
NativeResponseCache::memory(37, std::time::Duration::from_secs(913), 8192);
|
||||
let mismatched =
|
||||
NativeResponseCache::memory(37, std::time::Duration::from_secs(913), 8191);
|
||||
let matching_config = NativeCacheConfig {
|
||||
policy: config.policy,
|
||||
backend: CacheBackendConfig::Memory(memory),
|
||||
};
|
||||
assert_eq!(matching_config.service_mismatch(&matching), None);
|
||||
assert_eq!(
|
||||
matching_config.service_mismatch(&mismatched),
|
||||
Some("facade and native backend item limits must match")
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn projects_resolved_redis_tls_configuration() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let facade = facade(
|
||||
py,
|
||||
"pool = ConnectionPool()\n\
|
||||
pool.connection_class = SSLConnection\n\
|
||||
pool.max_connections = 29\n\
|
||||
pool.connection_kwargs = {'host': 'cache.internal', 'port': 6380, 'db': 4, 'username': 'user', 'password': 'secret', 'protocol': 3, 'socket_timeout': 7.5, 'socket_connect_timeout': 2, 'socket_keepalive': True, 'health_check_interval': 15, 'client_name': 'litellm', 'ssl_cert_reqs': 'optional', 'ssl_check_hostname': True, 'ssl_ca_certs': '/ca.pem', 'ssl_ca_data': 'CA DATA', 'ssl_certfile': '/client.pem', 'ssl_keyfile': '/client.key'}\n\
|
||||
client = SimpleNamespace(connection_pool=pool)\n\
|
||||
backend = SimpleNamespace(default_ttl=777, namespace='team', redis_flush_size=31, redis_kwargs={}, redis_client=client)\n\
|
||||
facade = SimpleNamespace(type='redis', mode='default-off', ttl=None, namespace='team', supported_call_types=None, redis_flush_size=31, semantic_cache_scope='key', cache=backend)",
|
||||
);
|
||||
let CacheConfigProjection::Native(config) =
|
||||
NativeCacheConfig::project(&facade).unwrap()
|
||||
else {
|
||||
panic!("Redis cache should be supported");
|
||||
};
|
||||
let CacheBackendConfig::Redis(redis) = config.backend else {
|
||||
panic!("expected Redis configuration");
|
||||
};
|
||||
assert_eq!(redis.default_ttl, std::time::Duration::from_secs(777));
|
||||
assert_eq!(redis.namespace.as_deref(), Some("team"));
|
||||
assert_eq!(redis.flush_size, 31);
|
||||
assert_eq!(redis.connection.host, "cache.internal");
|
||||
assert_eq!(redis.connection.port, 6380);
|
||||
assert_eq!(redis.connection.database, 4);
|
||||
assert_eq!(redis.connection.protocol, RedisProtocol::Resp3);
|
||||
assert_eq!(redis.connection.pool_size, 29);
|
||||
let tls = redis.connection.tls.unwrap();
|
||||
assert_eq!(
|
||||
tls.certificate_requirement,
|
||||
CertificateRequirement::Optional
|
||||
);
|
||||
assert!(tls.check_hostname);
|
||||
assert_eq!(tls.ca_certificate.as_deref(), Some("/ca.pem"));
|
||||
assert_eq!(tls.ca_data.as_deref(), Some("CA DATA"));
|
||||
assert_eq!(tls.client_certificate.as_deref(), Some("/client.pem"));
|
||||
assert_eq!(tls.client_key.as_deref(), Some("/client.key"));
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dynamic_redis_auth_stays_on_python() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let facade = facade(
|
||||
py,
|
||||
"backend = SimpleNamespace(redis_kwargs={'credential_provider': object()})\n\
|
||||
facade = SimpleNamespace(type='redis', mode='default-on', ttl=None, namespace=None, supported_call_types=[], redis_flush_size=None, semantic_cache_scope='key', cache=backend)",
|
||||
);
|
||||
let CacheConfigProjection::Unsupported(reason) =
|
||||
NativeCacheConfig::project(&facade).unwrap()
|
||||
else {
|
||||
panic!("dynamic authentication must stay on Python");
|
||||
};
|
||||
assert_eq!(reason.message(), "native Redis credentials require Python");
|
||||
});
|
||||
}
|
||||
}
|
||||
293
litellm-rust/crates/python-bridge/src/cache/facade.rs
vendored
Normal file
293
litellm-rust/crates/python-bridge/src/cache/facade.rs
vendored
Normal file
|
|
@ -0,0 +1,293 @@
|
|||
use litellm_host_python::from_py;
|
||||
use pyo3::{
|
||||
PyTraverseError, PyVisit,
|
||||
exceptions::PyTypeError,
|
||||
prelude::*,
|
||||
types::{PyDict, PyTuple, PyType},
|
||||
};
|
||||
use serde_json::Value;
|
||||
|
||||
use super::{
|
||||
config::{CacheConfigProjection, NativeCacheConfig},
|
||||
handle::CacheTestHandle,
|
||||
native::NativeResponseCache,
|
||||
};
|
||||
|
||||
struct ClassGuard {
|
||||
class: Py<PyType>,
|
||||
attributes: Vec<(String, Py<PyAny>)>,
|
||||
}
|
||||
|
||||
struct ObjectGuard {
|
||||
reference: Py<PyAny>,
|
||||
classes: Vec<ClassGuard>,
|
||||
config_names: &'static [&'static str],
|
||||
config: Vec<Value>,
|
||||
}
|
||||
|
||||
struct RedisPoolGuard {
|
||||
reference: Py<PyAny>,
|
||||
connection_class: Py<PyAny>,
|
||||
connection_kwargs: Py<PyAny>,
|
||||
max_connections: usize,
|
||||
}
|
||||
|
||||
pub(super) struct FacadeGuard {
|
||||
outer: ObjectGuard,
|
||||
backend: ObjectGuard,
|
||||
redis_pool: Option<RedisPoolGuard>,
|
||||
}
|
||||
|
||||
impl ObjectGuard {
|
||||
fn capture(
|
||||
py: Python<'_>,
|
||||
object: &Bound<'_, PyAny>,
|
||||
config_names: &'static [&'static str],
|
||||
) -> PyResult<Self> {
|
||||
let classes = object
|
||||
.get_type()
|
||||
.getattr("__mro__")?
|
||||
.cast_into::<PyTuple>()?
|
||||
.iter()
|
||||
.map(|class| {
|
||||
let class = class.cast_into::<PyType>()?;
|
||||
let attributes = class
|
||||
.getattr("__dict__")?
|
||||
.call_method0("items")?
|
||||
.try_iter()?
|
||||
.map(|item| item?.extract::<(String, Py<PyAny>)>())
|
||||
.collect::<PyResult<Vec<_>>>()?;
|
||||
Ok(ClassGuard {
|
||||
class: class.unbind(),
|
||||
attributes,
|
||||
})
|
||||
})
|
||||
.collect::<PyResult<Vec<_>>>()?;
|
||||
let guard = Self {
|
||||
reference: py
|
||||
.import("weakref")?
|
||||
.getattr("ref")?
|
||||
.call1((object,))?
|
||||
.unbind(),
|
||||
classes,
|
||||
config_names,
|
||||
config: Self::config(object, config_names)?,
|
||||
};
|
||||
if !guard.matches(py, object)? {
|
||||
return Err(PyTypeError::new_err(
|
||||
"native facade registration requires unmodified built-in methods",
|
||||
));
|
||||
}
|
||||
Ok(guard)
|
||||
}
|
||||
|
||||
fn config(object: &Bound<'_, PyAny>, names: &[&str]) -> PyResult<Vec<Value>> {
|
||||
names
|
||||
.iter()
|
||||
.map(|name| match object.getattr(*name) {
|
||||
Ok(value) => from_py(&value),
|
||||
Err(error)
|
||||
if error.is_instance_of::<pyo3::exceptions::PyAttributeError>(object.py()) =>
|
||||
{
|
||||
Ok(Value::Null)
|
||||
}
|
||||
Err(error) => Err(error),
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn matches(&self, py: Python<'_>, object: &Bound<'_, PyAny>) -> PyResult<bool> {
|
||||
if !self.reference.bind(py).call0()?.is(object) {
|
||||
return Ok(false);
|
||||
}
|
||||
let mro = object
|
||||
.get_type()
|
||||
.getattr("__mro__")?
|
||||
.cast_into::<PyTuple>()?;
|
||||
if mro.len() != self.classes.len() {
|
||||
return Ok(false);
|
||||
}
|
||||
let instance = object.getattr("__dict__")?.cast_into::<PyDict>()?;
|
||||
for (class, expected) in mro.iter().zip(&self.classes) {
|
||||
if !class.is(expected.class.bind(py)) {
|
||||
return Ok(false);
|
||||
}
|
||||
let attributes = class.getattr("__dict__")?;
|
||||
if attributes.len()? != expected.attributes.len() {
|
||||
return Ok(false);
|
||||
}
|
||||
for (name, value) in &expected.attributes {
|
||||
if instance.contains(name)? || !attributes.get_item(name)?.is(value.bind(py)) {
|
||||
return Ok(false);
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(Self::config(object, self.config_names)? == self.config)
|
||||
}
|
||||
|
||||
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
|
||||
visit.call(&self.reference)?;
|
||||
for class in &self.classes {
|
||||
visit.call(&class.class)?;
|
||||
for (_, value) in &class.attributes {
|
||||
visit.call(value)?;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl RedisPoolGuard {
|
||||
fn capture(backend: &Bound<'_, PyAny>) -> PyResult<Self> {
|
||||
let pool = backend
|
||||
.getattr("redis_client")?
|
||||
.getattr("connection_pool")?;
|
||||
Ok(Self {
|
||||
reference: pool.clone().unbind(),
|
||||
connection_class: pool.getattr("connection_class")?.unbind(),
|
||||
connection_kwargs: pool
|
||||
.getattr("connection_kwargs")?
|
||||
.call_method0("copy")?
|
||||
.unbind(),
|
||||
max_connections: pool.getattr("max_connections")?.extract::<usize>()?,
|
||||
})
|
||||
}
|
||||
|
||||
fn matches(&self, py: Python<'_>, backend: &Bound<'_, PyAny>) -> PyResult<bool> {
|
||||
let pool = backend
|
||||
.getattr("redis_client")?
|
||||
.getattr("connection_pool")?;
|
||||
Ok(self.reference.bind(py).is(&pool)
|
||||
&& self
|
||||
.connection_class
|
||||
.bind(py)
|
||||
.is(&pool.getattr("connection_class")?)
|
||||
&& self.max_connections == pool.getattr("max_connections")?.extract::<usize>()?
|
||||
&& self
|
||||
.connection_kwargs
|
||||
.bind(py)
|
||||
.eq(pool.getattr("connection_kwargs")?)?)
|
||||
}
|
||||
|
||||
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
|
||||
visit.call(&self.reference)?;
|
||||
visit.call(&self.connection_class)?;
|
||||
visit.call(&self.connection_kwargs)
|
||||
}
|
||||
}
|
||||
|
||||
impl FacadeGuard {
|
||||
pub(super) fn capture(
|
||||
py: Python<'_>,
|
||||
facade: &Bound<'_, PyAny>,
|
||||
service: &NativeResponseCache,
|
||||
) -> PyResult<Self> {
|
||||
let kind = service.kind();
|
||||
let cache_type = py.import("litellm.caching.caching")?.getattr("Cache")?;
|
||||
if !facade.get_type().is(&cache_type) {
|
||||
return Err(PyTypeError::new_err(
|
||||
"only exact built-in Cache facades can be registered",
|
||||
));
|
||||
}
|
||||
let (module, name, cache_kind) = match kind {
|
||||
"memory" => ("litellm.caching.in_memory_cache", "InMemoryCache", "local"),
|
||||
"redis" => ("litellm.caching.redis_cache", "RedisCache", "redis"),
|
||||
_ => unreachable!(),
|
||||
};
|
||||
let backend = facade.getattr("cache")?;
|
||||
if facade.getattr("type")?.extract::<String>()? != cache_kind
|
||||
|| !backend.get_type().is(&py.import(module)?.getattr(name)?)
|
||||
{
|
||||
return Err(PyTypeError::new_err(
|
||||
"facade and native backend types must match",
|
||||
));
|
||||
}
|
||||
let config = match NativeCacheConfig::project(facade)? {
|
||||
CacheConfigProjection::Native(config) => *config,
|
||||
CacheConfigProjection::Unsupported(reason) => {
|
||||
return Err(PyTypeError::new_err(reason.message()));
|
||||
}
|
||||
};
|
||||
if let Some(message) = config.service_mismatch(service) {
|
||||
return Err(PyTypeError::new_err(message));
|
||||
}
|
||||
Ok(Self {
|
||||
outer: ObjectGuard::capture(
|
||||
py,
|
||||
facade,
|
||||
&[
|
||||
"type",
|
||||
"mode",
|
||||
"ttl",
|
||||
"namespace",
|
||||
"supported_call_types",
|
||||
"redis_flush_size",
|
||||
"semantic_cache_scope",
|
||||
],
|
||||
)?,
|
||||
backend: ObjectGuard::capture(
|
||||
py,
|
||||
&backend,
|
||||
&[
|
||||
"namespace",
|
||||
"default_ttl",
|
||||
"max_size_in_memory",
|
||||
"max_size_per_item",
|
||||
"redis_kwargs",
|
||||
"redis_flush_size",
|
||||
],
|
||||
)?,
|
||||
redis_pool: (kind == "redis")
|
||||
.then(|| RedisPoolGuard::capture(&backend))
|
||||
.transpose()?,
|
||||
})
|
||||
}
|
||||
|
||||
fn matches(&self, py: Python<'_>, facade: &Bound<'_, PyAny>) -> PyResult<bool> {
|
||||
if !self.outer.matches(py, facade)? {
|
||||
return Ok(false);
|
||||
}
|
||||
let backend = facade.getattr("cache")?;
|
||||
if !self.backend.matches(py, &backend)? {
|
||||
return Ok(false);
|
||||
}
|
||||
match &self.redis_pool {
|
||||
Some(guard) => guard.matches(py, &backend),
|
||||
None => Ok(true),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn traverse(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
|
||||
self.outer.traverse(&visit)?;
|
||||
self.backend.traverse(&visit)?;
|
||||
if let Some(guard) = &self.redis_pool {
|
||||
guard.traverse(&visit)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn resolve(
|
||||
py: Python<'_>,
|
||||
facade: &Bound<'_, PyAny>,
|
||||
) -> PyResult<Option<NativeResponseCache>> {
|
||||
let Ok(dict) = facade
|
||||
.getattr("__dict__")
|
||||
.and_then(|dict| dict.cast_into::<PyDict>().map_err(Into::into))
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(handle) = dict.get_item("_native_cache_handle")? else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Ok(handle) = handle.extract::<PyRef<'_, CacheTestHandle>>() else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(guard) = &handle.guard else {
|
||||
return Ok(None);
|
||||
};
|
||||
if !guard.matches(py, facade).unwrap_or(false) {
|
||||
return Ok(None);
|
||||
}
|
||||
handle.service().map(Some)
|
||||
}
|
||||
18
litellm-rust/crates/python-bridge/src/cache/future.rs
vendored
Normal file
18
litellm-rust/crates/python-bridge/src/cache/future.rs
vendored
Normal file
|
|
@ -0,0 +1,18 @@
|
|||
use litellm_host_python::to_py;
|
||||
use pyo3::prelude::*;
|
||||
|
||||
pub(super) fn ready_none(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
|
||||
ready_value(py, &())
|
||||
}
|
||||
|
||||
pub(super) fn ready_value<'py, T: serde::Serialize>(
|
||||
py: Python<'py>,
|
||||
value: &T,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let future = py
|
||||
.import("asyncio")?
|
||||
.call_method0("get_running_loop")?
|
||||
.call_method0("create_future")?;
|
||||
future.call_method1("set_result", (to_py(py, value)?,))?;
|
||||
Ok(future)
|
||||
}
|
||||
84
litellm-rust/crates/python-bridge/src/cache/handle.rs
vendored
Normal file
84
litellm-rust/crates/python-bridge/src/cache/handle.rs
vendored
Normal file
|
|
@ -0,0 +1,84 @@
|
|||
use litellm_host_python::release_gil;
|
||||
use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*};
|
||||
|
||||
use super::{cache_error, facade::FacadeGuard, native::NativeResponseCache, request::duration};
|
||||
|
||||
#[pyclass(frozen, name = "_CacheTestHandle")]
|
||||
pub(crate) struct CacheTestHandle {
|
||||
service: NativeResponseCache,
|
||||
pub(super) guard: Option<FacadeGuard>,
|
||||
pid: u32,
|
||||
}
|
||||
|
||||
impl CacheTestHandle {
|
||||
pub(super) fn service(&self) -> PyResult<NativeResponseCache> {
|
||||
if self.pid != std::process::id() {
|
||||
return Err(PyRuntimeError::new_err(
|
||||
"native cache handles must be recreated after fork",
|
||||
));
|
||||
}
|
||||
Ok(self.service.clone())
|
||||
}
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl CacheTestHandle {
|
||||
#[staticmethod]
|
||||
#[pyo3(signature = (*, capacity=200, ttl_seconds=600.0, max_entry_bytes=1048576))]
|
||||
fn memory(capacity: usize, ttl_seconds: f64, max_entry_bytes: usize) -> PyResult<Self> {
|
||||
Ok(Self {
|
||||
service: NativeResponseCache::memory(capacity, duration(ttl_seconds)?, max_entry_bytes),
|
||||
guard: None,
|
||||
pid: std::process::id(),
|
||||
})
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
#[pyo3(signature = (url, *, ttl_seconds=60.0, namespace=None))]
|
||||
fn redis(
|
||||
py: Python<'_>,
|
||||
url: String,
|
||||
ttl_seconds: f64,
|
||||
namespace: Option<String>,
|
||||
) -> PyResult<Self> {
|
||||
let ttl = Some(duration(ttl_seconds)?);
|
||||
let service = release_gil(py, move || NativeResponseCache::redis(&url, ttl, namespace))
|
||||
.map_err(cache_error)?;
|
||||
Ok(Self {
|
||||
service,
|
||||
guard: None,
|
||||
pid: std::process::id(),
|
||||
})
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn backend(&self) -> &'static str {
|
||||
self.service.kind()
|
||||
}
|
||||
|
||||
fn _bind_facade(&self, py: Python<'_>, facade: &Bound<'_, PyAny>) -> PyResult<()> {
|
||||
let service = self.service()?;
|
||||
let guard = FacadeGuard::capture(py, facade, &service)?;
|
||||
let service = service.with_redis_flush_size(
|
||||
facade
|
||||
.getattr("redis_flush_size")?
|
||||
.extract::<Option<usize>>()?,
|
||||
);
|
||||
let handle = Py::new(
|
||||
py,
|
||||
Self {
|
||||
service,
|
||||
guard: Some(guard),
|
||||
pid: self.pid,
|
||||
},
|
||||
)?;
|
||||
facade.setattr("_native_cache_handle", handle)
|
||||
}
|
||||
|
||||
fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
|
||||
if let Some(guard) = &self.guard {
|
||||
guard.traverse(visit)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
26
litellm-rust/crates/python-bridge/src/cache/mod.rs
vendored
Normal file
26
litellm-rust/crates/python-bridge/src/cache/mod.rs
vendored
Normal file
|
|
@ -0,0 +1,26 @@
|
|||
mod binding;
|
||||
mod callback;
|
||||
mod config;
|
||||
mod facade;
|
||||
mod future;
|
||||
mod handle;
|
||||
mod native;
|
||||
mod request;
|
||||
mod resolver;
|
||||
|
||||
use litellm_cache::Error;
|
||||
use pyo3::{
|
||||
exceptions::{PyRuntimeError, PyValueError},
|
||||
prelude::*,
|
||||
};
|
||||
|
||||
pub(crate) use self::{
|
||||
binding::ResolvedCache, handle::CacheTestHandle, resolver::CacheTestResolver,
|
||||
};
|
||||
|
||||
fn cache_error(error: Error) -> PyErr {
|
||||
match error {
|
||||
Error::InvalidEntry => PyValueError::new_err(error.to_string()),
|
||||
_ => PyRuntimeError::new_err(error.to_string()),
|
||||
}
|
||||
}
|
||||
198
litellm-rust/crates/python-bridge/src/cache/native.rs
vendored
Normal file
198
litellm-rust/crates/python-bridge/src/cache/native.rs
vendored
Normal file
|
|
@ -0,0 +1,198 @@
|
|||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use litellm_cache::{CacheCodec, CacheConnectionResult, Error};
|
||||
use litellm_cache_memory::InMemoryCache;
|
||||
use litellm_cache_redis::RedisCache;
|
||||
use litellm_cache_response::{
|
||||
CacheEntry, PartialHits, ResponseCache, ResponseCacheCodec, ResponseCacheRequest, WriteBuffer,
|
||||
};
|
||||
use serde_json::Value;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(super) enum NativeResponseCache {
|
||||
Memory(Arc<ResponseCache<InMemoryCache<CacheEntry>>>),
|
||||
Redis {
|
||||
cache: Arc<ResponseCache<RedisCache<ResponseCacheCodec>>>,
|
||||
buffer: Option<Arc<WriteBuffer>>,
|
||||
},
|
||||
}
|
||||
|
||||
impl NativeResponseCache {
|
||||
pub fn memory(capacity: usize, ttl: Duration, max_entry_bytes: usize) -> Self {
|
||||
Self::Memory(Arc::new(ResponseCache::new(Arc::new(
|
||||
InMemoryCache::with_clock_and_size_measurement(
|
||||
Some(capacity),
|
||||
Some(ttl),
|
||||
Some(max_entry_bytes),
|
||||
Some(Arc::new(|entry| {
|
||||
ResponseCacheCodec.encode(entry).map(|bytes| bytes.len())
|
||||
})),
|
||||
super::request::now,
|
||||
),
|
||||
))))
|
||||
}
|
||||
|
||||
pub fn redis(
|
||||
url: &str,
|
||||
ttl: Option<Duration>,
|
||||
namespace: Option<String>,
|
||||
) -> Result<Self, Error> {
|
||||
let backend = RedisCache::new(url, ttl, ResponseCacheCodec)?.with_namespace(namespace);
|
||||
Ok(Self::Redis {
|
||||
cache: Arc::new(ResponseCache::new(Arc::new(backend))),
|
||||
buffer: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl NativeResponseCache {
|
||||
pub fn kind(&self) -> &'static str {
|
||||
match self {
|
||||
Self::Memory(_) => "memory",
|
||||
Self::Redis { .. } => "redis",
|
||||
}
|
||||
}
|
||||
|
||||
pub fn default_ttl(&self) -> Option<Duration> {
|
||||
match self {
|
||||
Self::Memory(cache) => cache.default_ttl(),
|
||||
Self::Redis { cache, .. } => cache.default_ttl(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn namespace(&self) -> Option<&str> {
|
||||
match self {
|
||||
Self::Memory(_) => None,
|
||||
Self::Redis { cache, .. } => cache.backend().namespace(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn capacity(&self) -> Option<usize> {
|
||||
match self {
|
||||
Self::Memory(cache) => Some(cache.backend().max_size_in_memory()),
|
||||
Self::Redis { .. } => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn max_entry_bytes(&self) -> Option<usize> {
|
||||
match self {
|
||||
Self::Memory(cache) => cache.backend().max_entry_bytes(),
|
||||
Self::Redis { .. } => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_redis_flush_size(self, flush_size: Option<usize>) -> Self {
|
||||
match self {
|
||||
Self::Redis { cache, .. } => Self::Redis {
|
||||
cache,
|
||||
buffer: flush_size.map(|flush_size| Arc::new(WriteBuffer::new(flush_size))),
|
||||
},
|
||||
memory => memory,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn lookup(
|
||||
&self,
|
||||
request: &ResponseCacheRequest,
|
||||
now: Duration,
|
||||
) -> Result<Option<Value>, Error> {
|
||||
match self {
|
||||
Self::Memory(cache) => cache.lookup(request, now),
|
||||
Self::Redis { cache, .. } => cache.lookup(request, now),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn store(
|
||||
&self,
|
||||
request: &ResponseCacheRequest,
|
||||
response: Value,
|
||||
now: Duration,
|
||||
) -> Result<(), Error> {
|
||||
match self {
|
||||
Self::Memory(cache) => cache.store(request, response, now),
|
||||
Self::Redis { cache, .. } => cache.store(request, response, now),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn lookup_batch(
|
||||
&self,
|
||||
requests: &[ResponseCacheRequest],
|
||||
now: Duration,
|
||||
) -> Result<PartialHits, Error> {
|
||||
match self {
|
||||
Self::Memory(cache) => cache.lookup_batch(requests, now),
|
||||
Self::Redis { cache, .. } => cache.lookup_batch(requests, now),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn async_lookup(
|
||||
&self,
|
||||
request: &ResponseCacheRequest,
|
||||
now: Duration,
|
||||
) -> Result<Option<Value>, Error> {
|
||||
match self {
|
||||
Self::Memory(cache) => cache.async_lookup(request, now).await,
|
||||
Self::Redis { cache, .. } => cache.async_lookup(request, now).await,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn async_store(
|
||||
&self,
|
||||
request: &ResponseCacheRequest,
|
||||
response: Value,
|
||||
now: Duration,
|
||||
) -> Result<(), Error> {
|
||||
match self {
|
||||
Self::Memory(cache) => cache.async_store(request, response, now).await,
|
||||
Self::Redis {
|
||||
cache,
|
||||
buffer: None,
|
||||
} => cache.async_store(request, response, now).await,
|
||||
Self::Redis {
|
||||
cache,
|
||||
buffer: Some(buffer),
|
||||
} => buffer.async_store(cache, request, response, now).await,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn async_lookup_batch(
|
||||
&self,
|
||||
requests: &[ResponseCacheRequest],
|
||||
now: Duration,
|
||||
) -> Result<PartialHits, Error> {
|
||||
match self {
|
||||
Self::Memory(cache) => cache.async_lookup_batch(requests, now).await,
|
||||
Self::Redis { cache, .. } => cache.async_lookup_batch(requests, now).await,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn async_store_batch(
|
||||
&self,
|
||||
entries: Vec<(ResponseCacheRequest, Value)>,
|
||||
now: Duration,
|
||||
) -> Result<(), Error> {
|
||||
match self {
|
||||
Self::Memory(cache) => cache.async_store_batch(entries, now).await,
|
||||
Self::Redis { cache, .. } => cache.async_store_batch(entries, now).await,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn async_flush(&self) -> Result<(), Error> {
|
||||
match self {
|
||||
Self::Memory(cache) => cache.async_flush().await,
|
||||
Self::Redis { cache, buffer } => {
|
||||
if let Some(buffer) = buffer {
|
||||
buffer.clear()?;
|
||||
}
|
||||
cache.async_flush().await
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
|
||||
match self {
|
||||
Self::Memory(cache) => cache.test_connection().await,
|
||||
Self::Redis { cache, .. } => cache.test_connection().await,
|
||||
}
|
||||
}
|
||||
}
|
||||
48
litellm-rust/crates/python-bridge/src/cache/request.rs
vendored
Normal file
48
litellm-rust/crates/python-bridge/src/cache/request.rs
vendored
Normal file
|
|
@ -0,0 +1,48 @@
|
|||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||
|
||||
use litellm_cache_response::{CacheControls, CacheKeyInput, ResponseCacheRequest};
|
||||
use litellm_host_python::from_py;
|
||||
use pyo3::{exceptions::PyValueError, prelude::*};
|
||||
use serde::Deserialize;
|
||||
|
||||
#[derive(Deserialize)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
struct RequestInput {
|
||||
key: CacheKeyInput,
|
||||
controls: Option<CacheControls>,
|
||||
ttl_seconds: Option<f64>,
|
||||
max_age_seconds: Option<f64>,
|
||||
}
|
||||
|
||||
pub(super) fn request(value: &Bound<'_, PyAny>) -> PyResult<ResponseCacheRequest> {
|
||||
let input: RequestInput = from_py(value)?;
|
||||
request_input(input)
|
||||
}
|
||||
|
||||
fn request_input(input: RequestInput) -> PyResult<ResponseCacheRequest> {
|
||||
let mut request = ResponseCacheRequest::new(input.key);
|
||||
if let Some(controls) = input.controls {
|
||||
request.controls = controls;
|
||||
}
|
||||
request.context.ttl = input.ttl_seconds.map(duration).transpose()?;
|
||||
request.max_age = input.max_age_seconds.map(duration).transpose()?;
|
||||
Ok(request)
|
||||
}
|
||||
|
||||
pub(super) fn requests(value: &Bound<'_, PyAny>) -> PyResult<Vec<ResponseCacheRequest>> {
|
||||
from_py::<Vec<RequestInput>>(value)?
|
||||
.into_iter()
|
||||
.map(request_input)
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub(super) fn duration(seconds: f64) -> PyResult<Duration> {
|
||||
Duration::try_from_secs_f64(seconds)
|
||||
.map_err(|_| PyValueError::new_err("cache durations must be finite and nonnegative"))
|
||||
}
|
||||
|
||||
pub(super) fn now() -> Duration {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
}
|
||||
39
litellm-rust/crates/python-bridge/src/cache/resolver.rs
vendored
Normal file
39
litellm-rust/crates/python-bridge/src/cache/resolver.rs
vendored
Normal file
|
|
@ -0,0 +1,39 @@
|
|||
use pyo3::{PyTraverseError, PyVisit, prelude::*};
|
||||
|
||||
use super::{
|
||||
binding::{CacheBinding, ResolvedCache},
|
||||
callback::PythonCallback,
|
||||
facade,
|
||||
handle::CacheTestHandle,
|
||||
};
|
||||
|
||||
#[pyclass(frozen, name = "_CacheTestResolver")]
|
||||
pub(crate) struct CacheTestResolver {
|
||||
namespace: Py<PyAny>,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl CacheTestResolver {
|
||||
#[new]
|
||||
fn new(namespace: Py<PyAny>) -> Self {
|
||||
Self { namespace }
|
||||
}
|
||||
|
||||
pub(crate) fn resolve(&self, py: Python<'_>) -> PyResult<ResolvedCache> {
|
||||
let object = self.namespace.bind(py).getattr("cache")?;
|
||||
let binding = if object.is_none() {
|
||||
CacheBinding::Disabled
|
||||
} else if let Ok(handle) = object.extract::<PyRef<'_, CacheTestHandle>>() {
|
||||
CacheBinding::Native(handle.service()?)
|
||||
} else if let Some(service) = facade::resolve(py, &object)? {
|
||||
CacheBinding::Native(service)
|
||||
} else {
|
||||
CacheBinding::PythonCallback(PythonCallback::new(object.unbind()))
|
||||
};
|
||||
Ok(ResolvedCache::new(binding))
|
||||
}
|
||||
|
||||
fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
|
||||
visit.call(&self.namespace)
|
||||
}
|
||||
}
|
||||
|
|
@ -1,3 +1,4 @@
|
|||
mod cache;
|
||||
mod credentials;
|
||||
mod diagnostics;
|
||||
mod errors;
|
||||
|
|
@ -9,6 +10,7 @@ mod token_counter;
|
|||
|
||||
#[pymodule(gil_used = true)]
|
||||
mod _native {
|
||||
use crate::cache::{CacheTestHandle, CacheTestResolver, ResolvedCache};
|
||||
#[cfg(feature = "panic-test")]
|
||||
#[pymodule_export]
|
||||
use crate::diagnostics::_panic_for_test;
|
||||
|
|
@ -32,6 +34,16 @@ mod _native {
|
|||
use crate::token_counter::TokenCounter;
|
||||
#[pymodule_export]
|
||||
use litellm_host_python::{ForkedAfterNativeRuntimeStarted, ProcessReservedForForking};
|
||||
use pyo3::{prelude::*, types::PyModule};
|
||||
|
||||
#[pymodule_init]
|
||||
fn init(module: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
let py = module.py();
|
||||
let dict = module.dict();
|
||||
dict.set_item("_CacheTestHandle", py.get_type::<CacheTestHandle>())?;
|
||||
dict.set_item("_CacheTestResolver", py.get_type::<CacheTestResolver>())?;
|
||||
dict.set_item("_CacheTestBinding", py.get_type::<ResolvedCache>())
|
||||
}
|
||||
}
|
||||
|
||||
use pyo3::prelude::*;
|
||||
|
|
|
|||
|
|
@ -1,6 +1,11 @@
|
|||
use std::{num::NonZero, sync::Arc, thread::available_parallelism};
|
||||
use std::sync::Arc;
|
||||
|
||||
use litellm_host_python::{release_gil, run_async};
|
||||
#[cfg(any(feature = "fast", feature = "huggingface", feature = "tiktoken"))]
|
||||
use std::{num::NonZero, thread::available_parallelism};
|
||||
|
||||
#[cfg(any(feature = "fast", feature = "huggingface", feature = "tiktoken"))]
|
||||
use litellm_host_python::release_gil;
|
||||
use litellm_host_python::run_async;
|
||||
use litellm_token_counter::{
|
||||
CountableRequest, Error, InputTokenCount, TokenCounter as CoreTokenCounter,
|
||||
};
|
||||
|
|
@ -28,17 +33,66 @@ pub(crate) struct TokenCounter {
|
|||
impl TokenCounter {
|
||||
#[new]
|
||||
fn new(py: Python<'_>, tokenizer_json: &str) -> PyResult<Self> {
|
||||
Self::load(py, || CoreTokenCounter::from_json(tokenizer_json))
|
||||
#[cfg(feature = "fast")]
|
||||
{
|
||||
Self::load(py, || CoreTokenCounter::from_json_fast(tokenizer_json))
|
||||
}
|
||||
#[cfg(all(not(feature = "fast"), feature = "huggingface"))]
|
||||
{
|
||||
Self::load(py, || CoreTokenCounter::from_json(tokenizer_json))
|
||||
}
|
||||
#[cfg(not(any(feature = "fast", feature = "huggingface")))]
|
||||
{
|
||||
let _ = (py, tokenizer_json);
|
||||
Err(RustBridgeDeclined::new_err(
|
||||
"tokenizer backend requires the fast or huggingface feature",
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
fn from_cl100k_ranks(py: Python<'_>, rank_file: &str) -> PyResult<Self> {
|
||||
Self::load(py, || CoreTokenCounter::from_cl100k_ranks(rank_file))
|
||||
#[cfg(feature = "fast")]
|
||||
{
|
||||
Self::load(py, || CoreTokenCounter::from_cl100k_ranks(rank_file))
|
||||
}
|
||||
#[cfg(not(feature = "fast"))]
|
||||
{
|
||||
let _ = (py, rank_file);
|
||||
Err(RustBridgeDeclined::new_err(
|
||||
"tokenizer backend requires the fast feature",
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
fn from_o200k_ranks(py: Python<'_>, rank_file: &str) -> PyResult<Self> {
|
||||
Self::load(py, || CoreTokenCounter::from_o200k_ranks(rank_file))
|
||||
#[cfg(feature = "fast")]
|
||||
{
|
||||
Self::load(py, || CoreTokenCounter::from_o200k_ranks(rank_file))
|
||||
}
|
||||
#[cfg(not(feature = "fast"))]
|
||||
{
|
||||
let _ = (py, rank_file);
|
||||
Err(RustBridgeDeclined::new_err(
|
||||
"tokenizer backend requires the fast feature",
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
fn from_tiktoken(py: Python<'_>, encoding: &str) -> PyResult<Self> {
|
||||
#[cfg(feature = "tiktoken")]
|
||||
{
|
||||
Self::load(py, || CoreTokenCounter::from_tiktoken(encoding))
|
||||
}
|
||||
#[cfg(not(feature = "tiktoken"))]
|
||||
{
|
||||
let _ = (py, encoding);
|
||||
Err(RustBridgeDeclined::new_err(
|
||||
"tokenizer backend requires the tiktoken feature",
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
fn acount_request<'py>(&self, py: Python<'py>, body: &[u8]) -> PyResult<Bound<'py, PyAny>> {
|
||||
|
|
@ -62,6 +116,7 @@ impl TokenCounter {
|
|||
}
|
||||
|
||||
impl TokenCounter {
|
||||
#[cfg(any(feature = "fast", feature = "huggingface", feature = "tiktoken"))]
|
||||
fn load(
|
||||
py: Python<'_>,
|
||||
load: impl FnOnce() -> Result<CoreTokenCounter, Error> + Send,
|
||||
|
|
@ -74,6 +129,7 @@ impl TokenCounter {
|
|||
}
|
||||
}
|
||||
|
||||
#[cfg(any(feature = "fast", feature = "huggingface", feature = "tiktoken"))]
|
||||
fn encode_parallelism() -> usize {
|
||||
available_parallelism().map_or(1, NonZero::get)
|
||||
}
|
||||
|
|
@ -86,7 +142,10 @@ fn count_body(counter: &CoreTokenCounter, body: &[u8]) -> Result<InputTokenCount
|
|||
fn token_count_error_to_pyerr(error: Error) -> PyErr {
|
||||
let message = error.to_string();
|
||||
match error {
|
||||
Error::Load(_) | Error::Ranks(_) | Error::UnicodeClasses => PyValueError::new_err(message),
|
||||
Error::Load(_)
|
||||
| Error::Ranks(_)
|
||||
| Error::UnicodeClasses
|
||||
| Error::UnsupportedTokenizer(_) => PyValueError::new_err(message),
|
||||
Error::RequestParse(_)
|
||||
| Error::MissingInput
|
||||
| Error::FloatText
|
||||
|
|
|
|||
24
litellm-rust/crates/secrets-aws/Cargo.toml
Normal file
24
litellm-rust/crates/secrets-aws/Cargo.toml
Normal file
|
|
@ -0,0 +1,24 @@
|
|||
[package]
|
||||
name = "litellm-secrets-aws"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
litellm-auth-aws.workspace = true
|
||||
litellm-secrets-types.workspace = true
|
||||
litellm-core-utils.workspace = true
|
||||
serde_json.workspace = true
|
||||
thiserror.workspace = true
|
||||
tracing = "0.1"
|
||||
veil.workspace = true
|
||||
aws-sdk-kms = "1.120.0"
|
||||
aws-sdk-secretsmanager = "1.117.0"
|
||||
aws-credential-types = "1.3.0"
|
||||
|
||||
[dev-dependencies]
|
||||
base64.workspace = true
|
||||
rstest.workspace = true
|
||||
tokio.workspace = true
|
||||
wiremock = "0.6.5"
|
||||
79
litellm-rust/crates/secrets-aws/src/auth.rs
Normal file
79
litellm-rust/crates/secrets-aws/src/auth.rs
Normal file
|
|
@ -0,0 +1,79 @@
|
|||
use std::sync::Arc;
|
||||
|
||||
use aws_credential_types::provider::{ProvideCredentials, error::CredentialsError, future};
|
||||
use litellm_auth_aws::{
|
||||
AwsAuthConfig,
|
||||
constants::{AWS_DEFAULT_REGION, AWS_REGION, AWS_REGION_NAME},
|
||||
resolve_credentials,
|
||||
};
|
||||
use litellm_core_utils::settings::Lookup;
|
||||
use litellm_secrets_types::KeyManagementSettings;
|
||||
|
||||
use crate::Error;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct Credentials {
|
||||
config: AwsAuthConfig,
|
||||
environment: Arc<dyn Lookup + Send + Sync>,
|
||||
}
|
||||
|
||||
impl Credentials {
|
||||
pub(crate) fn new(
|
||||
settings: &KeyManagementSettings,
|
||||
environment: Arc<dyn Lookup + Send + Sync>,
|
||||
) -> Self {
|
||||
Self {
|
||||
config: AwsAuthConfig {
|
||||
region_name: region(settings, environment.as_ref()).ok(),
|
||||
role_name: settings.aws_role_name.clone(),
|
||||
session_name: settings.aws_session_name.clone(),
|
||||
external_id: settings
|
||||
.aws_external_id
|
||||
.as_ref()
|
||||
.map(|v| v.expose().to_owned()),
|
||||
profile_name: settings.aws_profile_name.clone(),
|
||||
web_identity_token: settings
|
||||
.aws_web_identity_token
|
||||
.as_ref()
|
||||
.map(|v| v.expose().to_owned()),
|
||||
sts_endpoint: settings.aws_sts_endpoint.clone(),
|
||||
..Default::default()
|
||||
},
|
||||
environment,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ProvideCredentials for Credentials {
|
||||
fn provide_credentials<'a>(&'a self) -> future::ProvideCredentials<'a>
|
||||
where
|
||||
Self: 'a,
|
||||
{
|
||||
future::ProvideCredentials::new(async {
|
||||
resolve_credentials(self.config.clone(), &|name| self.environment.get(name))
|
||||
.await
|
||||
.map_err(|_| {
|
||||
CredentialsError::provider_error("secret manager authentication failed")
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn region(
|
||||
settings: &KeyManagementSettings,
|
||||
environment: &dyn Lookup,
|
||||
) -> Result<String, Error> {
|
||||
settings
|
||||
.aws_region_name
|
||||
.clone()
|
||||
.or_else(|| environment.get(AWS_REGION_NAME))
|
||||
.or_else(|| environment.get(AWS_REGION))
|
||||
.or_else(|| environment.get(AWS_DEFAULT_REGION))
|
||||
.ok_or(Error::MissingRegion)
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for Credentials {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("Credentials").finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
31
litellm-rust/crates/secrets-aws/src/error.rs
Normal file
31
litellm-rust/crates/secrets-aws/src/error.rs
Normal file
|
|
@ -0,0 +1,31 @@
|
|||
use aws_sdk_secretsmanager::error::SdkError;
|
||||
|
||||
#[derive(thiserror::Error, veil::Redact)]
|
||||
pub enum Error {
|
||||
#[error("AWS authentication failed")]
|
||||
Auth(#[from] #[redact] litellm_auth_aws::Error),
|
||||
#[error("AWS region is not configured")]
|
||||
MissingRegion,
|
||||
#[error("KMS response has no plaintext")]
|
||||
MissingPlaintext,
|
||||
#[error("AWS request timed out")]
|
||||
Timeout,
|
||||
#[error("AWS KMS decrypt failed")]
|
||||
Decrypt(#[from] #[redact] Box<SdkError<aws_sdk_kms::operation::decrypt::DecryptError>>),
|
||||
#[error("AWS Secrets Manager read failed")]
|
||||
Read(#[from] #[redact] Box<SdkError<aws_sdk_secretsmanager::operation::get_secret_value::GetSecretValueError>>),
|
||||
#[error("AWS Secrets Manager create failed")]
|
||||
Create(#[from] #[redact] Box<SdkError<aws_sdk_secretsmanager::operation::create_secret::CreateSecretError>>),
|
||||
#[error("AWS Secrets Manager update failed")]
|
||||
Put(#[from] #[redact] Box<SdkError<aws_sdk_secretsmanager::operation::put_secret_value::PutSecretValueError>>),
|
||||
#[error("AWS Secrets Manager delete failed")]
|
||||
Delete(#[from] #[redact] Box<SdkError<aws_sdk_secretsmanager::operation::delete_secret::DeleteSecretError>>),
|
||||
#[error("AWS Secrets Manager replication failed")]
|
||||
Replicate(#[from] #[redact] Box<SdkError<aws_sdk_secretsmanager::operation::replicate_secret_to_regions::ReplicateSecretToRegionsError>>),
|
||||
#[error("AWS Secrets Manager response has no string payload")]
|
||||
MissingString,
|
||||
#[error("primary secret is not a JSON object")]
|
||||
PrimarySecret,
|
||||
#[error(transparent)]
|
||||
Operation(#[from] litellm_secrets_types::Error),
|
||||
}
|
||||
63
litellm-rust/crates/secrets-aws/src/kms.rs
Normal file
63
litellm-rust/crates/secrets-aws/src/kms.rs
Normal file
|
|
@ -0,0 +1,63 @@
|
|||
use litellm_auth_aws::constants::AWS_REGION_NAME;
|
||||
use std::sync::Arc;
|
||||
|
||||
use aws_sdk_kms::{
|
||||
Client,
|
||||
config::{BehaviorVersion, Region},
|
||||
primitives::Blob,
|
||||
};
|
||||
use litellm_core_utils::settings::Lookup;
|
||||
use litellm_secrets_types::KeyManagementSettings;
|
||||
|
||||
use crate::{Error, auth};
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct AwsKms {
|
||||
client: Client,
|
||||
}
|
||||
|
||||
impl AwsKms {
|
||||
pub fn new(client: Client) -> Self {
|
||||
Self { client }
|
||||
}
|
||||
|
||||
pub async fn decrypt(&self, ciphertext: Vec<u8>) -> Result<Vec<u8>, Error> {
|
||||
let response = self
|
||||
.client
|
||||
.decrypt()
|
||||
.ciphertext_blob(Blob::new(ciphertext))
|
||||
.send()
|
||||
.await
|
||||
.map_err(|error| Error::Decrypt(Box::new(error)))?;
|
||||
Ok(response
|
||||
.plaintext
|
||||
.ok_or(Error::MissingPlaintext)?
|
||||
.into_inner())
|
||||
}
|
||||
}
|
||||
|
||||
pub fn validate_environment(environment: &dyn Lookup) -> Result<(), Error> {
|
||||
environment
|
||||
.get(AWS_REGION_NAME)
|
||||
.map(|_| ())
|
||||
.ok_or(Error::MissingRegion)
|
||||
}
|
||||
|
||||
pub fn load_aws_kms(
|
||||
use_aws_kms: Option<bool>,
|
||||
settings: &KeyManagementSettings,
|
||||
environment: Arc<dyn Lookup + Send + Sync>,
|
||||
) -> Result<Option<AwsKms>, Error> {
|
||||
if use_aws_kms != Some(true) {
|
||||
return Ok(None);
|
||||
}
|
||||
if settings.aws_region_name.is_none() {
|
||||
validate_environment(environment.as_ref())?;
|
||||
}
|
||||
let config = aws_sdk_kms::Config::builder()
|
||||
.behavior_version(BehaviorVersion::latest())
|
||||
.region(Region::new(auth::region(settings, environment.as_ref())?))
|
||||
.credentials_provider(auth::Credentials::new(settings, environment))
|
||||
.build();
|
||||
Ok(Some(AwsKms::new(Client::from_conf(config))))
|
||||
}
|
||||
10
litellm-rust/crates/secrets-aws/src/lib.rs
Normal file
10
litellm-rust/crates/secrets-aws/src/lib.rs
Normal file
|
|
@ -0,0 +1,10 @@
|
|||
#![forbid(unsafe_code)]
|
||||
|
||||
mod auth;
|
||||
mod error;
|
||||
pub mod kms;
|
||||
pub mod secret_manager;
|
||||
|
||||
pub use error::Error;
|
||||
pub use kms::{AwsKms, load_aws_kms};
|
||||
pub use secret_manager::{AwsSecretWriteSettings, AwsSecretsManagerV2, RotationResponse};
|
||||
287
litellm-rust/crates/secrets-aws/src/secret_manager.rs
Normal file
287
litellm-rust/crates/secrets-aws/src/secret_manager.rs
Normal file
|
|
@ -0,0 +1,287 @@
|
|||
use litellm_auth_aws::constants::AWS_BEDROCK_RUNTIME_ENDPOINT;
|
||||
use std::{collections::BTreeMap, sync::Arc};
|
||||
|
||||
use aws_sdk_secretsmanager::{
|
||||
Client,
|
||||
config::{BehaviorVersion, Region},
|
||||
operation::{
|
||||
create_secret::CreateSecretOutput, delete_secret::DeleteSecretOutput,
|
||||
put_secret_value::PutSecretValueOutput,
|
||||
replicate_secret_to_regions::ReplicateSecretToRegionsOutput,
|
||||
},
|
||||
types::{ReplicaRegionType, Tag},
|
||||
};
|
||||
use litellm_auth_aws::constants::{
|
||||
AWS_ACCESS_KEY_ID, AWS_REGION, AWS_REGION_NAME, AWS_SECRET_ACCESS_KEY,
|
||||
};
|
||||
use litellm_core_utils::settings::Lookup;
|
||||
use litellm_secrets_types::{
|
||||
BaseSecretManager, KeyManagementSettings, Secret, SecretValue, async_rotate_secret,
|
||||
};
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{Error, auth};
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct AwsSecretsManagerV2 {
|
||||
client: Client,
|
||||
write_settings: AwsSecretWriteSettings,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default)]
|
||||
pub struct AwsSecretWriteSettings {
|
||||
pub kms_key_id: Option<String>,
|
||||
pub tags: Option<BTreeMap<String, String>>,
|
||||
pub replica_regions: Option<Vec<String>>,
|
||||
}
|
||||
|
||||
impl From<&KeyManagementSettings> for AwsSecretWriteSettings {
|
||||
fn from(settings: &KeyManagementSettings) -> Self {
|
||||
Self {
|
||||
kms_key_id: settings.kms_key_id.clone(),
|
||||
tags: settings.tags.clone(),
|
||||
replica_regions: settings.replica_regions.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum RotationResponse {
|
||||
Created(CreateSecretOutput),
|
||||
Updated(PutSecretValueOutput),
|
||||
}
|
||||
|
||||
impl AwsSecretsManagerV2 {
|
||||
pub fn new(client: Client, write_settings: AwsSecretWriteSettings) -> Self {
|
||||
Self {
|
||||
client,
|
||||
write_settings,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn load_aws_secret_manager(
|
||||
use_aws_secret_manager: Option<bool>,
|
||||
settings: KeyManagementSettings,
|
||||
environment: Arc<dyn Lookup + Send + Sync>,
|
||||
) -> Result<Option<Self>, Error> {
|
||||
if use_aws_secret_manager != Some(true) {
|
||||
return Ok(None);
|
||||
}
|
||||
let builder = aws_sdk_secretsmanager::Config::builder()
|
||||
.behavior_version(BehaviorVersion::latest())
|
||||
.region(Region::new(auth::region(&settings, environment.as_ref())?))
|
||||
.credentials_provider(auth::Credentials::new(&settings, environment.clone()));
|
||||
let config = match environment.get(AWS_BEDROCK_RUNTIME_ENDPOINT) {
|
||||
Some(url) => builder
|
||||
.endpoint_url(url.replace("bedrock-runtime", "secretsmanager"))
|
||||
.build(),
|
||||
None => builder.build(),
|
||||
};
|
||||
Ok(Some(Self::new(
|
||||
Client::from_conf(config),
|
||||
(&settings).into(),
|
||||
)))
|
||||
}
|
||||
|
||||
pub async fn read_secret_for_resolver(
|
||||
&self,
|
||||
name: &str,
|
||||
primary_name: Option<&str>,
|
||||
environment: &(dyn Lookup + Sync),
|
||||
) -> Result<Option<Secret>, Error> {
|
||||
if bootstrap_key(name) {
|
||||
return Ok(environment
|
||||
.get(name)
|
||||
.map(SecretValue::new)
|
||||
.map(Secret::String));
|
||||
}
|
||||
match primary_name.filter(|name| !name.is_empty()) {
|
||||
None => self
|
||||
.async_read_secret(name)
|
||||
.await
|
||||
.map(|value| value.map(Secret::String)),
|
||||
Some(primary) => {
|
||||
let value = if bootstrap_key(primary) {
|
||||
environment.get(primary).map(SecretValue::new)
|
||||
} else {
|
||||
self.async_read_secret(primary).await?
|
||||
};
|
||||
let Some(value) = value else {
|
||||
return Ok(None);
|
||||
};
|
||||
let object: Value =
|
||||
serde_json::from_str(value.expose()).map_err(|_| Error::PrimarySecret)?;
|
||||
let object = object.as_object().ok_or(Error::PrimarySecret)?;
|
||||
Ok(object.get(name).cloned().map(Secret::from_json))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn async_read_secret(&self, name: &str) -> Result<Option<SecretValue>, Error> {
|
||||
match self.client.get_secret_value().secret_id(name).send().await {
|
||||
Ok(response) => response
|
||||
.secret_string
|
||||
.map(SecretValue::new)
|
||||
.map(Some)
|
||||
.ok_or(Error::MissingString),
|
||||
Err(error)
|
||||
if matches!(
|
||||
&error,
|
||||
aws_sdk_secretsmanager::error::SdkError::TimeoutError(_)
|
||||
) || matches!(&error, aws_sdk_secretsmanager::error::SdkError::DispatchFailure(failure) if failure.is_timeout()) =>
|
||||
{
|
||||
Err(Error::Timeout)
|
||||
}
|
||||
Err(error)
|
||||
if error
|
||||
.as_service_error()
|
||||
.is_some_and(|error| error.is_resource_not_found_exception()) =>
|
||||
{
|
||||
Ok(None)
|
||||
}
|
||||
Err(error) => Err(Error::Read(Box::new(error))),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn async_write_secret(
|
||||
&self,
|
||||
name: &str,
|
||||
value: &SecretValue,
|
||||
description: Option<&str>,
|
||||
) -> Result<CreateSecretOutput, Error> {
|
||||
let response = self
|
||||
.client
|
||||
.create_secret()
|
||||
.name(name)
|
||||
.secret_string(value.expose())
|
||||
.set_description(description.filter(|v| !v.is_empty()).map(str::to_owned))
|
||||
.set_kms_key_id(
|
||||
self.write_settings
|
||||
.kms_key_id
|
||||
.clone()
|
||||
.filter(|v| !v.is_empty()),
|
||||
)
|
||||
.set_tags(self.write_settings.tags.as_ref().map(|tags| {
|
||||
tags.iter()
|
||||
.map(|(key, value)| Tag::builder().key(key).value(value).build())
|
||||
.collect()
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.map_err(|error| Error::Create(Box::new(error)))?;
|
||||
if let Some(regions) = &self.write_settings.replica_regions
|
||||
&& !regions.is_empty()
|
||||
&& self.async_replicate_secret(name, regions).await.is_err()
|
||||
{
|
||||
tracing::warn!("secret created but replication failed");
|
||||
}
|
||||
Ok(response)
|
||||
}
|
||||
|
||||
pub async fn async_replicate_secret(
|
||||
&self,
|
||||
name: &str,
|
||||
regions: &[String],
|
||||
) -> Result<Option<ReplicateSecretToRegionsOutput>, Error> {
|
||||
if regions.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
self.client
|
||||
.replicate_secret_to_regions()
|
||||
.secret_id(name)
|
||||
.set_add_replica_regions(Some(
|
||||
regions
|
||||
.iter()
|
||||
.map(|region| ReplicaRegionType::builder().region(region).build())
|
||||
.collect(),
|
||||
))
|
||||
.send()
|
||||
.await
|
||||
.map(Some)
|
||||
.map_err(|error| Error::Replicate(Box::new(error)))
|
||||
}
|
||||
|
||||
pub async fn async_put_secret_value(
|
||||
&self,
|
||||
name: &str,
|
||||
value: &SecretValue,
|
||||
) -> Result<PutSecretValueOutput, Error> {
|
||||
self.client
|
||||
.put_secret_value()
|
||||
.secret_id(name)
|
||||
.secret_string(value.expose())
|
||||
.send()
|
||||
.await
|
||||
.map_err(|error| Error::Put(Box::new(error)))
|
||||
}
|
||||
|
||||
pub async fn async_delete_secret(
|
||||
&self,
|
||||
name: &str,
|
||||
recovery_window_in_days: i64,
|
||||
) -> Result<DeleteSecretOutput, Error> {
|
||||
self.client
|
||||
.delete_secret()
|
||||
.secret_id(name)
|
||||
.recovery_window_in_days(recovery_window_in_days)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|error| Error::Delete(Box::new(error)))
|
||||
}
|
||||
|
||||
pub async fn async_rotate_secret(
|
||||
&self,
|
||||
current_name: &str,
|
||||
new_name: &str,
|
||||
value: &SecretValue,
|
||||
) -> Result<RotationResponse, Error> {
|
||||
if current_name == new_name {
|
||||
return self
|
||||
.async_put_secret_value(current_name, value)
|
||||
.await
|
||||
.map(RotationResponse::Updated);
|
||||
}
|
||||
async_rotate_secret(self, current_name, new_name, value)
|
||||
.await
|
||||
.map(RotationResponse::Created)
|
||||
}
|
||||
}
|
||||
|
||||
impl BaseSecretManager for AwsSecretsManagerV2 {
|
||||
type Error = Error;
|
||||
type WriteResponse = CreateSecretOutput;
|
||||
type DeleteResponse = DeleteSecretOutput;
|
||||
|
||||
async fn async_read_secret(&self, name: &str) -> Result<Option<SecretValue>, Error> {
|
||||
self.async_read_secret(name).await
|
||||
}
|
||||
|
||||
async fn async_write_secret(
|
||||
&self,
|
||||
name: &str,
|
||||
value: &SecretValue,
|
||||
description: Option<&str>,
|
||||
) -> Result<CreateSecretOutput, Error> {
|
||||
self.async_write_secret(name, value, description).await
|
||||
}
|
||||
|
||||
async fn async_delete_secret(
|
||||
&self,
|
||||
name: &str,
|
||||
recovery_window_in_days: i64,
|
||||
) -> Result<DeleteSecretOutput, Error> {
|
||||
self.async_delete_secret(name, recovery_window_in_days)
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
fn bootstrap_key(name: &str) -> bool {
|
||||
matches!(
|
||||
name,
|
||||
AWS_ACCESS_KEY_ID
|
||||
| AWS_SECRET_ACCESS_KEY
|
||||
| AWS_REGION_NAME
|
||||
| AWS_REGION
|
||||
| AWS_BEDROCK_RUNTIME_ENDPOINT
|
||||
)
|
||||
}
|
||||
59
litellm-rust/crates/secrets-aws/tests/kms.rs
Normal file
59
litellm-rust/crates/secrets-aws/tests/kms.rs
Normal file
|
|
@ -0,0 +1,59 @@
|
|||
use aws_sdk_kms::{
|
||||
Client,
|
||||
config::{BehaviorVersion, Credentials, Region, retry::RetryConfig},
|
||||
};
|
||||
use base64::{Engine, engine::general_purpose::STANDARD};
|
||||
use litellm_secrets_aws::AwsKms;
|
||||
use wiremock::{
|
||||
Mock, MockServer, ResponseTemplate,
|
||||
matchers::{body_json, header},
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
async fn kms_decrypt_calls_the_sdk_without_applying_lookup_policy() {
|
||||
let server = MockServer::start().await;
|
||||
let plaintext = " private-value\n";
|
||||
Mock::given(header("x-amz-target", "TrentService.Decrypt"))
|
||||
.and(body_json(
|
||||
serde_json::json!({"CiphertextBlob": STANDARD.encode("encrypted")}),
|
||||
))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200)
|
||||
.set_body_json(serde_json::json!({"Plaintext": STANDARD.encode(plaintext)})),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let client = Client::from_conf(
|
||||
aws_sdk_kms::Config::builder()
|
||||
.behavior_version(BehaviorVersion::latest())
|
||||
.region(Region::new("us-east-1"))
|
||||
.credentials_provider(Credentials::new("test", "test", None, None, "test"))
|
||||
.endpoint_url(server.uri())
|
||||
.retry_config(RetryConfig::disabled())
|
||||
.build(),
|
||||
);
|
||||
let manager = AwsKms::new(client);
|
||||
assert_eq!(
|
||||
manager.decrypt(b"encrypted".to_vec()).await.unwrap(),
|
||||
plaintext.as_bytes()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn disabled_kms_loader_does_not_require_environment_configuration() {
|
||||
use litellm_secrets_aws::load_aws_kms;
|
||||
use litellm_secrets_types::KeyManagementSettings;
|
||||
use std::sync::Arc;
|
||||
for enabled in [None, Some(false)] {
|
||||
assert!(
|
||||
load_aws_kms(
|
||||
enabled,
|
||||
&KeyManagementSettings::default(),
|
||||
Arc::new(|_: &str| None)
|
||||
)
|
||||
.unwrap()
|
||||
.is_none()
|
||||
);
|
||||
}
|
||||
}
|
||||
312
litellm-rust/crates/secrets-aws/tests/secret_manager.rs
Normal file
312
litellm-rust/crates/secrets-aws/tests/secret_manager.rs
Normal file
|
|
@ -0,0 +1,312 @@
|
|||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
|
||||
use aws_sdk_secretsmanager::{
|
||||
Client,
|
||||
config::{BehaviorVersion, Credentials, Region, retry::RetryConfig},
|
||||
};
|
||||
use litellm_secrets_aws::{AwsSecretsManagerV2, Error, RotationResponse};
|
||||
use litellm_secrets_types::{KeyManagementSettings, SecretValue};
|
||||
use serde_json::json;
|
||||
use wiremock::{
|
||||
Mock, MockServer, ResponseTemplate,
|
||||
matchers::{body_partial_json, header},
|
||||
};
|
||||
|
||||
fn manager(server: &MockServer, settings: KeyManagementSettings) -> AwsSecretsManagerV2 {
|
||||
let client = Client::from_conf(
|
||||
aws_sdk_secretsmanager::Config::builder()
|
||||
.behavior_version(BehaviorVersion::latest())
|
||||
.region(Region::new("us-east-1"))
|
||||
.credentials_provider(Credentials::new("test", "test", None, None, "test"))
|
||||
.endpoint_url(server.uri())
|
||||
.retry_config(RetryConfig::disabled())
|
||||
.build(),
|
||||
);
|
||||
AwsSecretsManagerV2::new(client, (&settings).into())
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::string_value("KEY", Some("value"))]
|
||||
#[case::missing_value("missing", None)]
|
||||
#[case::non_string_value("BOOL", None)]
|
||||
#[tokio::test]
|
||||
async fn primary_lookup_preserves_read_semantics(
|
||||
#[case] name: &str,
|
||||
#[case] expected: Option<&str>,
|
||||
) {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(header("x-amz-target", "secretsmanager.GetSecretValue"))
|
||||
.and(body_partial_json(json!({"SecretId":"primary"})))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200).set_body_json(
|
||||
json!({"SecretString":json!({"KEY":"value", "BOOL":true}).to_string()}),
|
||||
),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let manager = manager(&server, KeyManagementSettings::default());
|
||||
assert_eq!(
|
||||
manager
|
||||
.read_secret_for_resolver(name, Some("primary"), &|_: &str| None)
|
||||
.await
|
||||
.unwrap()
|
||||
.and_then(|v| v.as_str().map(str::to_owned))
|
||||
.as_deref(),
|
||||
expected
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::access_key("AWS_ACCESS_KEY_ID")]
|
||||
#[case::secret_access_key("AWS_SECRET_ACCESS_KEY")]
|
||||
#[case::region_name("AWS_REGION_NAME")]
|
||||
#[case::region("AWS_REGION")]
|
||||
#[case::bedrock_endpoint("AWS_BEDROCK_RUNTIME_ENDPOINT")]
|
||||
#[tokio::test]
|
||||
async fn bootstrap_keys_bypass_primary_lookup(#[case] name: &str) {
|
||||
let server = MockServer::start().await;
|
||||
let manager = manager(&server, KeyManagementSettings::default());
|
||||
assert_eq!(
|
||||
manager
|
||||
.read_secret_for_resolver(name, Some("primary"), &|_: &str| Some("bootstrap".into()))
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.as_str()
|
||||
.unwrap(),
|
||||
"bootstrap"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn failed_read_returns_none_but_invalid_primary_json_is_an_error() {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(body_partial_json(json!({"SecretId":"missing"})))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(400).set_body_json(json!({"__type":"ResourceNotFoundException"})),
|
||||
)
|
||||
.mount(&server)
|
||||
.await;
|
||||
Mock::given(body_partial_json(json!({"SecretId":"invalid"})))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"SecretString":"not-json"})))
|
||||
.mount(&server)
|
||||
.await;
|
||||
let manager = manager(&server, KeyManagementSettings::default());
|
||||
assert!(
|
||||
manager
|
||||
.async_read_secret("missing")
|
||||
.await
|
||||
.unwrap()
|
||||
.is_none()
|
||||
);
|
||||
assert!(matches!(
|
||||
manager
|
||||
.read_secret_for_resolver("KEY", Some("invalid"), &|_: &str| None)
|
||||
.await,
|
||||
Err(Error::PrimarySecret)
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn same_name_rotation_uses_put_and_returns_its_response() {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(header("x-amz-target", "secretsmanager.PutSecretValue"))
|
||||
.and(body_partial_json(
|
||||
json!({"SecretId":"key", "SecretString":"replacement"}),
|
||||
))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200).set_body_json(json!({"Name":"key", "VersionId":"version"})),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let response = manager(&server, KeyManagementSettings::default())
|
||||
.async_rotate_secret("key", "key", &SecretValue::new("replacement"))
|
||||
.await
|
||||
.unwrap();
|
||||
match response {
|
||||
RotationResponse::Updated(output) => assert_eq!(output.version_id(), Some("version")),
|
||||
_ => panic!("rotation created a second secret"),
|
||||
}
|
||||
assert_eq!(server.received_requests().await.unwrap().len(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn renamed_rotation_reads_creates_verifies_then_deletes() {
|
||||
let server = MockServer::start().await;
|
||||
let step = AtomicUsize::new(0);
|
||||
Mock::given(wiremock::matchers::method("POST"))
|
||||
.respond_with(move |request: &wiremock::Request| {
|
||||
let body: serde_json::Value = request.body_json().unwrap();
|
||||
let action = request
|
||||
.headers
|
||||
.get("x-amz-target")
|
||||
.unwrap()
|
||||
.to_str()
|
||||
.unwrap();
|
||||
match step.fetch_add(1, Ordering::SeqCst) {
|
||||
0 => {
|
||||
assert_eq!(action, "secretsmanager.GetSecretValue");
|
||||
assert_eq!(body["SecretId"], "old");
|
||||
ResponseTemplate::new(200).set_body_json(json!({"SecretString":"old-value"}))
|
||||
}
|
||||
1 => {
|
||||
assert_eq!(action, "secretsmanager.CreateSecret");
|
||||
assert_eq!(body["Name"], "new");
|
||||
assert_eq!(body["Description"], "Rotated from old");
|
||||
assert_eq!(body["SecretString"], "replacement");
|
||||
ResponseTemplate::new(200).set_body_json(json!({"Name":"new"}))
|
||||
}
|
||||
2 => {
|
||||
assert_eq!(action, "secretsmanager.GetSecretValue");
|
||||
assert_eq!(body["SecretId"], "new");
|
||||
ResponseTemplate::new(200).set_body_json(json!({"SecretString":"replacement"}))
|
||||
}
|
||||
3 => {
|
||||
assert_eq!(action, "secretsmanager.DeleteSecret");
|
||||
assert_eq!(body["SecretId"], "old");
|
||||
assert_eq!(body["RecoveryWindowInDays"], 7);
|
||||
ResponseTemplate::new(200).set_body_json(json!({"Name":"old"}))
|
||||
}
|
||||
_ => panic!("unexpected request"),
|
||||
}
|
||||
})
|
||||
.expect(4)
|
||||
.mount(&server)
|
||||
.await;
|
||||
assert!(matches!(
|
||||
manager(&server, KeyManagementSettings::default())
|
||||
.async_rotate_secret("old", "new", &SecretValue::new("replacement"))
|
||||
.await
|
||||
.unwrap(),
|
||||
RotationResponse::Created(_)
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn creation_passes_tags_and_kms_and_survives_replication_failure() {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(header("x-amz-target", "secretsmanager.CreateSecret"))
|
||||
.and(body_partial_json(json!({"Name":"key", "SecretString":"value", "KmsKeyId":"kms-key", "Tags":[{"Key":"stage", "Value":"test"}]})))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"Name":"key"}))).expect(1).mount(&server).await;
|
||||
Mock::given(header(
|
||||
"x-amz-target",
|
||||
"secretsmanager.ReplicateSecretToRegions",
|
||||
))
|
||||
.and(body_partial_json(
|
||||
json!({"SecretId":"key", "AddReplicaRegions":[{"Region":"replica-region"}]}),
|
||||
))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(400).set_body_json(json!({"__type":"InvalidRequestException"})),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let settings = KeyManagementSettings {
|
||||
kms_key_id: Some("kms-key".into()),
|
||||
tags: Some(std::collections::BTreeMap::from([(
|
||||
"stage".into(),
|
||||
"test".into(),
|
||||
)])),
|
||||
replica_regions: Some(vec!["replica-region".into()]),
|
||||
..Default::default()
|
||||
};
|
||||
let manager = manager(&server, settings);
|
||||
assert_eq!(
|
||||
manager
|
||||
.async_write_secret("key", &SecretValue::new("value"), None)
|
||||
.await
|
||||
.unwrap()
|
||||
.name(),
|
||||
Some("key")
|
||||
);
|
||||
assert!(
|
||||
manager
|
||||
.async_replicate_secret("key", &[])
|
||||
.await
|
||||
.unwrap()
|
||||
.is_none()
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn credential_failures_are_not_swallowed_as_missing_secrets() {
|
||||
use aws_credential_types::provider::{ProvideCredentials, error::CredentialsError, future};
|
||||
#[derive(Debug)]
|
||||
struct FailedCredentials;
|
||||
impl ProvideCredentials for FailedCredentials {
|
||||
fn provide_credentials<'a>(&'a self) -> future::ProvideCredentials<'a>
|
||||
where
|
||||
Self: 'a,
|
||||
{
|
||||
future::ProvideCredentials::ready(Err(CredentialsError::provider_error(
|
||||
"private-auth-detail",
|
||||
)))
|
||||
}
|
||||
}
|
||||
let server = MockServer::start().await;
|
||||
let config = aws_sdk_secretsmanager::Config::builder()
|
||||
.behavior_version(BehaviorVersion::latest())
|
||||
.region(Region::new("us-east-1"))
|
||||
.credentials_provider(FailedCredentials)
|
||||
.endpoint_url(server.uri())
|
||||
.retry_config(RetryConfig::disabled())
|
||||
.build();
|
||||
let manager = AwsSecretsManagerV2::new(Client::from_conf(config), Default::default());
|
||||
let error = manager.async_read_secret("key").await.unwrap_err();
|
||||
assert!(!format!("{error:?}").contains("private-auth-detail"));
|
||||
assert!(matches!(error, Error::Read(_)));
|
||||
assert!(server.received_requests().await.unwrap().is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn read_timeout_is_an_error_and_cannot_be_mistaken_for_missing() {
|
||||
use std::time::Duration;
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(wiremock::matchers::method("POST"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200)
|
||||
.set_delay(Duration::from_secs(1))
|
||||
.set_body_json(json!({"SecretString":"late"})),
|
||||
)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let config = aws_sdk_secretsmanager::Config::builder()
|
||||
.behavior_version(BehaviorVersion::latest())
|
||||
.region(Region::new("us-east-1"))
|
||||
.credentials_provider(Credentials::new("test", "test", None, None, "test"))
|
||||
.endpoint_url(server.uri())
|
||||
.retry_config(RetryConfig::disabled())
|
||||
.timeout_config(
|
||||
aws_sdk_secretsmanager::config::timeout::TimeoutConfig::builder()
|
||||
.operation_timeout(Duration::from_millis(30))
|
||||
.build(),
|
||||
)
|
||||
.build();
|
||||
let manager = AwsSecretsManagerV2::new(Client::from_conf(config), Default::default());
|
||||
assert!(matches!(
|
||||
manager.async_read_secret("key").await,
|
||||
Err(Error::Timeout)
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::denied(400, "AccessDeniedException")]
|
||||
#[case::throttled(400, "ThrottlingException")]
|
||||
#[case::unavailable(503, "ServiceUnavailableException")]
|
||||
#[tokio::test]
|
||||
async fn service_failures_remain_errors(#[case] status: u16, #[case] code: &str) {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(header("x-amz-target", "secretsmanager.GetSecretValue"))
|
||||
.respond_with(ResponseTemplate::new(status).set_body_json(json!({"__type":code})))
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
assert!(matches!(
|
||||
manager(&server, KeyManagementSettings::default())
|
||||
.async_read_secret("key")
|
||||
.await,
|
||||
Err(Error::Read(_))
|
||||
));
|
||||
}
|
||||
28
litellm-rust/crates/secrets-google/Cargo.toml
Normal file
28
litellm-rust/crates/secrets-google/Cargo.toml
Normal file
|
|
@ -0,0 +1,28 @@
|
|||
[package]
|
||||
name = "litellm-secrets-google"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
litellm-auth-gcp = { workspace = true, features = ["google-sdk"] }
|
||||
litellm-secrets-types.workspace = true
|
||||
litellm-auth-types.workspace = true
|
||||
litellm-core-utils.workspace = true
|
||||
base64.workspace = true
|
||||
serde_json.workspace = true
|
||||
thiserror.workspace = true
|
||||
moka.workspace = true
|
||||
veil.workspace = true
|
||||
google-cloud-kms-v1 = "1.14.0"
|
||||
google-cloud-gax = { version = "1.14.0", default-features = false }
|
||||
percent-encoding = "2.3"
|
||||
serde.workspace = true
|
||||
reqwest.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
google-cloud-auth.workspace = true
|
||||
rstest.workspace = true
|
||||
tokio.workspace = true
|
||||
wiremock = "0.6.5"
|
||||
21
litellm-rust/crates/secrets-google/src/auth.rs
Normal file
21
litellm-rust/crates/secrets-google/src/auth.rs
Normal file
|
|
@ -0,0 +1,21 @@
|
|||
use std::sync::Arc;
|
||||
|
||||
use litellm_auth_gcp::{GoogleCredentials, VertexConfig};
|
||||
use litellm_auth_types::{InputSource, Sourced};
|
||||
use litellm_core_utils::settings::Lookup;
|
||||
use litellm_secrets_types::SecretValue;
|
||||
|
||||
pub(crate) fn credentials(
|
||||
project: Option<String>,
|
||||
credentials: Option<SecretValue>,
|
||||
environment: Arc<dyn Lookup + Send + Sync>,
|
||||
) -> GoogleCredentials {
|
||||
GoogleCredentials::new(
|
||||
VertexConfig::new(
|
||||
credentials.map(|value| Sourced::new(value, InputSource::Environment)),
|
||||
project,
|
||||
None,
|
||||
),
|
||||
Arc::new(move |name| environment.get(name)),
|
||||
)
|
||||
}
|
||||
43
litellm-rust/crates/secrets-google/src/error.rs
Normal file
43
litellm-rust/crates/secrets-google/src/error.rs
Normal file
|
|
@ -0,0 +1,43 @@
|
|||
#[derive(thiserror::Error, veil::Redact)]
|
||||
pub enum Error {
|
||||
#[error("Google KMS client configuration failed")]
|
||||
Client(
|
||||
#[from]
|
||||
#[redact]
|
||||
google_cloud_gax::client_builder::Error,
|
||||
),
|
||||
#[error("Google authentication failed")]
|
||||
Auth(
|
||||
#[from]
|
||||
#[redact]
|
||||
litellm_auth_types::Error,
|
||||
),
|
||||
#[error("Google KMS request failed")]
|
||||
Kms(
|
||||
#[from]
|
||||
#[redact]
|
||||
google_cloud_gax::error::Error,
|
||||
),
|
||||
#[error("Google Secret Manager HTTP request failed")]
|
||||
Http(
|
||||
#[from]
|
||||
#[redact]
|
||||
reqwest::Error,
|
||||
),
|
||||
#[error("Google Secret Manager returned HTTP {0}")]
|
||||
Status(u16),
|
||||
#[error("Google Secret Manager returned no payload")]
|
||||
MissingPayload,
|
||||
#[error("required environment variable is missing: {0}")]
|
||||
MissingEnvironment(&'static str),
|
||||
#[error("invalid refresh interval")]
|
||||
RefreshInterval,
|
||||
#[error("payload is not valid base64")]
|
||||
Base64(#[from] base64::DecodeError),
|
||||
#[error("decrypted value is not UTF-8")]
|
||||
Utf8,
|
||||
#[error("invalid Google Secret Manager endpoint")]
|
||||
Endpoint,
|
||||
#[error("Google Secret Manager requires an enterprise license")]
|
||||
EnterpriseRequired,
|
||||
}
|
||||
67
litellm-rust/crates/secrets-google/src/kms.rs
Normal file
67
litellm-rust/crates/secrets-google/src/kms.rs
Normal file
|
|
@ -0,0 +1,67 @@
|
|||
use std::sync::Arc;
|
||||
|
||||
use google_cloud_kms_v1::client::KeyManagementService;
|
||||
use litellm_core_utils::settings::Lookup;
|
||||
use litellm_secrets_types::SecretValue;
|
||||
|
||||
use crate::{Error, auth};
|
||||
|
||||
const GOOGLE_APPLICATION_CREDENTIALS: &str = "GOOGLE_APPLICATION_CREDENTIALS";
|
||||
const GOOGLE_KMS_RESOURCE_NAME: &str = "GOOGLE_KMS_RESOURCE_NAME";
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct GoogleKms {
|
||||
client: KeyManagementService,
|
||||
resource_name: String,
|
||||
}
|
||||
|
||||
impl GoogleKms {
|
||||
pub fn new(client: KeyManagementService, resource_name: String) -> Self {
|
||||
Self {
|
||||
client,
|
||||
resource_name,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn decrypt(&self, ciphertext: Vec<u8>) -> Result<Vec<u8>, Error> {
|
||||
let response = self
|
||||
.client
|
||||
.decrypt()
|
||||
.set_name(&self.resource_name)
|
||||
.set_ciphertext(ciphertext)
|
||||
.send()
|
||||
.await?;
|
||||
Ok(response.plaintext.to_vec())
|
||||
}
|
||||
}
|
||||
|
||||
pub fn validate_environment(environment: &dyn Lookup) -> Result<(), Error> {
|
||||
for key in [GOOGLE_APPLICATION_CREDENTIALS, GOOGLE_KMS_RESOURCE_NAME] {
|
||||
if environment.get(key).is_none() {
|
||||
return Err(Error::MissingEnvironment(key));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn load_google_kms(
|
||||
use_google_kms: Option<bool>,
|
||||
environment: Arc<dyn Lookup + Send + Sync>,
|
||||
) -> Result<Option<GoogleKms>, Error> {
|
||||
if use_google_kms != Some(true) {
|
||||
return Ok(None);
|
||||
}
|
||||
validate_environment(environment.as_ref())?;
|
||||
let credentials = environment
|
||||
.get(GOOGLE_APPLICATION_CREDENTIALS)
|
||||
.ok_or(Error::MissingEnvironment(GOOGLE_APPLICATION_CREDENTIALS))?;
|
||||
let resource_name = environment
|
||||
.get(GOOGLE_KMS_RESOURCE_NAME)
|
||||
.ok_or(Error::MissingEnvironment(GOOGLE_KMS_RESOURCE_NAME))?;
|
||||
let credentials = auth::credentials(None, Some(SecretValue::new(credentials)), environment);
|
||||
let client = KeyManagementService::builder()
|
||||
.with_credentials(credentials)
|
||||
.build()
|
||||
.await?;
|
||||
Ok(Some(GoogleKms::new(client, resource_name)))
|
||||
}
|
||||
10
litellm-rust/crates/secrets-google/src/lib.rs
Normal file
10
litellm-rust/crates/secrets-google/src/lib.rs
Normal file
|
|
@ -0,0 +1,10 @@
|
|||
#![forbid(unsafe_code)]
|
||||
|
||||
mod auth;
|
||||
mod error;
|
||||
pub mod kms;
|
||||
pub mod secret_manager;
|
||||
|
||||
pub use error::Error;
|
||||
pub use kms::{GoogleKms, load_google_kms};
|
||||
pub use secret_manager::GoogleSecretManager;
|
||||
157
litellm-rust/crates/secrets-google/src/secret_manager.rs
Normal file
157
litellm-rust/crates/secrets-google/src/secret_manager.rs
Normal file
|
|
@ -0,0 +1,157 @@
|
|||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use base64::{Engine, engine::general_purpose::STANDARD};
|
||||
use litellm_core_utils::settings::Lookup;
|
||||
use litellm_secrets_types::{Secret, SecretValue};
|
||||
use moka::future::Cache;
|
||||
use serde::Deserialize;
|
||||
|
||||
use litellm_auth_gcp::GoogleCredentials;
|
||||
|
||||
use crate::{Error, auth};
|
||||
|
||||
const GOOGLE_SECRET_MANAGER_PROJECT_ID: &str = "GOOGLE_SECRET_MANAGER_PROJECT_ID";
|
||||
const GOOGLE_SECRET_MANAGER_REFRESH_INTERVAL: &str = "GOOGLE_SECRET_MANAGER_REFRESH_INTERVAL";
|
||||
const SECRET_MANAGER_REFRESH_INTERVAL: &str = "SECRET_MANAGER_REFRESH_INTERVAL";
|
||||
const GOOGLE_SECRET_MANAGER_ALWAYS_READ_SECRET_MANAGER: &str =
|
||||
"GOOGLE_SECRET_MANAGER_ALWAYS_READ_SECRET_MANAGER";
|
||||
const GCS_PATH_SERVICE_ACCOUNT: &str = "GCS_PATH_SERVICE_ACCOUNT";
|
||||
const DEFAULT_REFRESH_INTERVAL: Duration = Duration::from_secs(86400);
|
||||
const DEFAULT_CACHE_TTL: Duration = Duration::from_secs(600);
|
||||
const CACHE_CAPACITY: u64 = 200;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct GoogleSecretManager {
|
||||
client: reqwest::Client,
|
||||
credentials: Arc<GoogleCredentials>,
|
||||
endpoint: reqwest::Url,
|
||||
project: String,
|
||||
cache: Cache<String, SecretValue>,
|
||||
always_read: bool,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct Response {
|
||||
payload: Option<Payload>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct Payload {
|
||||
data: Option<String>,
|
||||
}
|
||||
|
||||
impl GoogleSecretManager {
|
||||
pub fn with_client(
|
||||
client: reqwest::Client,
|
||||
endpoint: reqwest::Url,
|
||||
project: String,
|
||||
environment: Arc<dyn Lookup + Send + Sync>,
|
||||
refresh_interval: Option<Duration>,
|
||||
always_read: bool,
|
||||
) -> Result<Self, Error> {
|
||||
let credentials = auth::credentials(
|
||||
Some(project.clone()),
|
||||
environment
|
||||
.get(GCS_PATH_SERVICE_ACCOUNT)
|
||||
.map(SecretValue::new),
|
||||
environment,
|
||||
);
|
||||
let ttl = refresh_interval
|
||||
.filter(|ttl| !ttl.is_zero())
|
||||
.unwrap_or(DEFAULT_CACHE_TTL);
|
||||
let cache = Cache::builder()
|
||||
.max_capacity(CACHE_CAPACITY)
|
||||
.time_to_live(ttl)
|
||||
.build();
|
||||
Ok(Self {
|
||||
client,
|
||||
credentials: Arc::new(credentials),
|
||||
endpoint,
|
||||
project,
|
||||
cache,
|
||||
always_read,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn new(
|
||||
environment: Arc<dyn Lookup + Send + Sync>,
|
||||
enterprise_enabled: bool,
|
||||
) -> Result<Self, Error> {
|
||||
if !enterprise_enabled {
|
||||
return Err(Error::EnterpriseRequired);
|
||||
}
|
||||
let project = environment
|
||||
.get(GOOGLE_SECRET_MANAGER_PROJECT_ID)
|
||||
.ok_or(Error::MissingEnvironment(GOOGLE_SECRET_MANAGER_PROJECT_ID))?;
|
||||
let ttl = environment
|
||||
.get(GOOGLE_SECRET_MANAGER_REFRESH_INTERVAL)
|
||||
.filter(|v| !v.is_empty())
|
||||
.map(|v| v.parse::<i64>().map_err(|_| Error::RefreshInterval))
|
||||
.transpose()?
|
||||
.unwrap_or(
|
||||
environment
|
||||
.get(SECRET_MANAGER_REFRESH_INTERVAL)
|
||||
.map(|v| v.parse::<i64>().map_err(|_| Error::RefreshInterval))
|
||||
.transpose()?
|
||||
.unwrap_or(DEFAULT_REFRESH_INTERVAL.as_secs() as i64),
|
||||
);
|
||||
let always_read = environment
|
||||
.get(GOOGLE_SECRET_MANAGER_ALWAYS_READ_SECRET_MANAGER)
|
||||
.is_some_and(|v| v.eq_ignore_ascii_case("true"));
|
||||
Self::with_client(
|
||||
reqwest::Client::new(),
|
||||
reqwest::Url::parse("https://secretmanager.googleapis.com").expect("static URL"),
|
||||
project,
|
||||
environment,
|
||||
Some(if ttl < 0 {
|
||||
Duration::from_nanos(1)
|
||||
} else {
|
||||
Duration::from_secs(ttl as u64)
|
||||
}),
|
||||
always_read,
|
||||
)
|
||||
}
|
||||
|
||||
pub async fn get_secret_from_google_secret_manager(
|
||||
&self,
|
||||
name: &str,
|
||||
) -> Result<Option<Secret>, Error> {
|
||||
if !self.always_read
|
||||
&& let Some(cached) = self.cache.get(name).await
|
||||
{
|
||||
return Ok(Some(Secret::String(cached)));
|
||||
}
|
||||
let url = self
|
||||
.endpoint
|
||||
.join(&format!(
|
||||
"/v1/projects/{}/secrets/{}/versions/latest:access",
|
||||
percent_encoding::utf8_percent_encode(
|
||||
&self.project,
|
||||
percent_encoding::NON_ALPHANUMERIC
|
||||
),
|
||||
percent_encoding::utf8_percent_encode(name, percent_encoding::NON_ALPHANUMERIC)
|
||||
))
|
||||
.map_err(|_| Error::Endpoint)?;
|
||||
let response = self
|
||||
.client
|
||||
.get(url)
|
||||
.headers(self.credentials.request_headers().await?)
|
||||
.send()
|
||||
.await?;
|
||||
if response.status() == reqwest::StatusCode::NOT_FOUND {
|
||||
return Ok(None);
|
||||
}
|
||||
if response.status() != reqwest::StatusCode::OK {
|
||||
return Err(Error::Status(response.status().as_u16()));
|
||||
}
|
||||
let response: Response = response.json().await?;
|
||||
let Some(data) = response.payload.and_then(|payload| payload.data) else {
|
||||
return Err(Error::MissingPayload);
|
||||
};
|
||||
let bytes = STANDARD.decode(data)?;
|
||||
let plaintext = String::from_utf8(bytes).map_err(|_| Error::Utf8)?;
|
||||
let value = SecretValue::new(plaintext);
|
||||
self.cache.insert(name.to_owned(), value.clone()).await;
|
||||
Ok(Some(Secret::String(value)))
|
||||
}
|
||||
}
|
||||
49
litellm-rust/crates/secrets-google/tests/kms.rs
Normal file
49
litellm-rust/crates/secrets-google/tests/kms.rs
Normal file
|
|
@ -0,0 +1,49 @@
|
|||
use base64::{Engine, engine::general_purpose::STANDARD};
|
||||
use google_cloud_kms_v1::client::KeyManagementService;
|
||||
use litellm_secrets_google::GoogleKms;
|
||||
use wiremock::{
|
||||
Mock, MockServer, ResponseTemplate,
|
||||
matchers::{body_json, path},
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
async fn google_kms_decrypts_using_the_configured_resource() {
|
||||
let server = MockServer::start().await;
|
||||
let resource = "projects/project/locations/global/keyRings/ring/cryptoKeys/key";
|
||||
Mock::given(path(format!("/v1/{resource}:decrypt")))
|
||||
.and(body_json(
|
||||
serde_json::json!({"ciphertext":STANDARD.encode("encrypted")}),
|
||||
))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200)
|
||||
.set_body_json(serde_json::json!({"plaintext":STANDARD.encode(" value\n")})),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let client = KeyManagementService::builder()
|
||||
.with_endpoint(server.uri())
|
||||
.with_credentials(google_cloud_auth::credentials::anonymous::Builder::new().build())
|
||||
.with_retry_policy(google_cloud_gax::retry_policy::NeverRetry)
|
||||
.build()
|
||||
.await
|
||||
.unwrap();
|
||||
let manager = GoogleKms::new(client, resource.into());
|
||||
assert_eq!(
|
||||
manager.decrypt(b"encrypted".to_vec()).await.unwrap(),
|
||||
b" value\n"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn disabled_google_kms_loader_does_not_require_environment_configuration() {
|
||||
use std::sync::Arc;
|
||||
for enabled in [None, Some(false)] {
|
||||
assert!(
|
||||
litellm_secrets_google::load_google_kms(enabled, Arc::new(|_: &str| None))
|
||||
.await
|
||||
.unwrap()
|
||||
.is_none()
|
||||
);
|
||||
}
|
||||
}
|
||||
188
litellm-rust/crates/secrets-google/tests/secret_manager.rs
Normal file
188
litellm-rust/crates/secrets-google/tests/secret_manager.rs
Normal file
|
|
@ -0,0 +1,188 @@
|
|||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use base64::{Engine, engine::general_purpose::STANDARD};
|
||||
use litellm_secrets_google::{Error, GoogleSecretManager};
|
||||
|
||||
use wiremock::{
|
||||
Mock, MockServer, ResponseTemplate,
|
||||
matchers::{header, path},
|
||||
};
|
||||
|
||||
fn manager(server: &MockServer, always_read: bool, ttl: Duration) -> GoogleSecretManager {
|
||||
GoogleSecretManager::with_client(
|
||||
reqwest::Client::new(),
|
||||
server.uri().parse().unwrap(),
|
||||
"project".into(),
|
||||
Arc::new(|name: &str| (name == "VERTEX_AI_API_KEY").then(|| "token".into())),
|
||||
Some(ttl),
|
||||
always_read,
|
||||
)
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::nonempty("private-value")]
|
||||
#[case::empty("")]
|
||||
#[tokio::test]
|
||||
async fn successful_reads_use_auth_latest_version_and_cache_including_empty_values(
|
||||
#[case] value: &str,
|
||||
) {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(path(
|
||||
"/v1/projects/project/secrets/key/versions/latest:access",
|
||||
))
|
||||
.and(header("authorization", "Bearer token"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200)
|
||||
.set_body_json(serde_json::json!({"payload":{"data":STANDARD.encode(value)}})),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let manager = manager(&server, false, Duration::from_secs(60));
|
||||
for _ in 0..2 {
|
||||
assert_eq!(
|
||||
manager
|
||||
.get_secret_from_google_secret_manager("key")
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.as_str()
|
||||
.unwrap(),
|
||||
value
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::not_found(404, serde_json::json!({}))]
|
||||
#[case::unauthorized(401, serde_json::json!({}))]
|
||||
#[case::forbidden(403, serde_json::json!({}))]
|
||||
#[case::throttled(429, serde_json::json!({}))]
|
||||
#[case::unavailable(503, serde_json::json!({}))]
|
||||
#[case::missing_payload(200, serde_json::json!({"payload":{}}))]
|
||||
#[case::invalid_base64(200, serde_json::json!({"payload":{"data":"%%%"}}))]
|
||||
#[tokio::test]
|
||||
async fn failed_or_missing_reads_are_not_cached(
|
||||
#[case] status: u16,
|
||||
#[case] body: serde_json::Value,
|
||||
) {
|
||||
let server = MockServer::start().await;
|
||||
let manager = manager(&server, false, Duration::from_secs(60));
|
||||
let failing = Mock::given(path(
|
||||
"/v1/projects/project/secrets/key/versions/latest:access",
|
||||
))
|
||||
.respond_with(ResponseTemplate::new(status).set_body_json(body))
|
||||
.expect(1)
|
||||
.mount_as_scoped(&server)
|
||||
.await;
|
||||
let result = manager.get_secret_from_google_secret_manager("key").await;
|
||||
match status {
|
||||
404 => assert_eq!(result.unwrap(), None),
|
||||
200 => assert!(matches!(
|
||||
result,
|
||||
Err(Error::MissingPayload | Error::Base64(_))
|
||||
)),
|
||||
status => assert!(matches!(result, Err(Error::Status(actual)) if actual == status)),
|
||||
}
|
||||
drop(failing);
|
||||
Mock::given(path(
|
||||
"/v1/projects/project/secrets/key/versions/latest:access",
|
||||
))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200)
|
||||
.set_body_json(serde_json::json!({"payload":{"data":STANDARD.encode("recovered")}})),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
for _ in 0..2 {
|
||||
assert_eq!(
|
||||
manager
|
||||
.get_secret_from_google_secret_manager("key")
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.as_str(),
|
||||
Some("recovered")
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::always_read(true, Duration::from_secs(60))]
|
||||
#[case::expired_cache(false, Duration::from_millis(1))]
|
||||
#[tokio::test]
|
||||
async fn always_read_and_expired_cache_fetch_again(
|
||||
#[case] always_read: bool,
|
||||
#[case] ttl: Duration,
|
||||
) {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(path(
|
||||
"/v1/projects/project/secrets/key/versions/latest:access",
|
||||
))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200)
|
||||
.set_body_json(serde_json::json!({"payload":{"data":STANDARD.encode("value")}})),
|
||||
)
|
||||
.expect(2)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let manager = manager(&server, always_read, ttl);
|
||||
for _ in 0..2 {
|
||||
tokio::time::sleep(Duration::from_millis(5)).await;
|
||||
assert!(
|
||||
manager
|
||||
.get_secret_from_google_secret_manager("key")
|
||||
.await
|
||||
.unwrap()
|
||||
.is_some()
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn google_manager_requires_host_license_and_project_configuration() {
|
||||
assert!(matches!(
|
||||
GoogleSecretManager::new(Arc::new(|_: &str| None), false),
|
||||
Err(Error::EnterpriseRequired)
|
||||
));
|
||||
assert!(matches!(
|
||||
GoogleSecretManager::new(Arc::new(|_: &str| None), true),
|
||||
Err(Error::MissingEnvironment(
|
||||
"GOOGLE_SECRET_MANAGER_PROJECT_ID"
|
||||
))
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case("true")]
|
||||
#[case("null")]
|
||||
#[case("\"text\"")]
|
||||
#[case("{\"key\":1}")]
|
||||
#[tokio::test]
|
||||
async fn cache_preserves_raw_values(#[case] raw: &str) {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(path(
|
||||
"/v1/projects/project/secrets/key/versions/latest:access",
|
||||
))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200)
|
||||
.set_body_json(serde_json::json!({"payload":{"data":STANDARD.encode(raw)}})),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let manager = manager(&server, false, Duration::from_secs(60));
|
||||
for _ in 0..2 {
|
||||
assert_eq!(
|
||||
manager
|
||||
.get_secret_from_google_secret_manager("key")
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.as_str(),
|
||||
Some(raw)
|
||||
);
|
||||
}
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue