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

chore: merge main into litellm_background_response_read_billed_once
Some checks failed
LiteLLM Rust / rust-lint (push) Has been cancelled
Details
LiteLLM Rust / rust-test (push) Has been cancelled
Details
LiteLLM Rust / rust-wheel (push) Has been cancelled
Details
Terraform Modules / fmt, validate, test (aws) (push) Has been cancelled
Details
Terraform Modules / fmt, validate, test (gcp) (push) Has been cancelled
Details
Terraform Provider / gofmt, vet, build, test (push) Has been cancelled
Details
Terraform Provider / Provider endpoints vs proxy OpenAPI schema (push) Has been cancelled
Details

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
jesus 2026-09-21 22:41:21 +00:00
commit 38cc3e6715
585 changed files with 45119 additions and 8301 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

52
.circleci/config.yml
View file

@ -6,6 +6,9 @@ parameters:
migration_candidate_image:
type: string
default: ""
migration_baseline_image:
type: string
default: "ghcr.io/berriai/litellm-database:v1.102.0"
migration_source_sha:
type: string
default: ""
@ -2946,7 +2949,10 @@ jobs:
parameters:
suite:
type: enum
enum: [startup, recovery, legacy]
enum: [startup, recovery, legacy, upgrade, shaped]
baseline:
type: boolean
default: false
machine:
image: ubuntu-2204:2024.04.1
resource_class: large
@ -2954,6 +2960,7 @@ jobs:
environment:
LITELLM_MIGRATION_TESTS: "1"
LITELLM_MIGRATION_TEST_IMAGE: litellm-docker-database:ci
LITELLM_MIGRATION_BASELINE_IMAGE: << pipeline.parameters.migration_baseline_image >>
MIGRATION_TEST_ADMIN_URL: postgresql://postgres:postgres@127.0.0.1:5432/postgres
MIGRATION_TEST_CONTAINER_ADMIN_URL: postgresql://postgres:postgres@host.docker.internal:5432/postgres
MIGRATION_TEST_OUTPUT: /tmp/migration-results
@ -2981,6 +2988,16 @@ jobs:
- wait_for_service:
url: tcp://localhost:5432
timeout: "60"
- when:
condition: << parameters.baseline >>
steps:
- run:
name: Pull the baseline release the upgrade starts from
environment:
BASELINE_IMAGE: << pipeline.parameters.migration_baseline_image >>
command: |
[[ "$BASELINE_IMAGE" =~ ^ghcr.io/berriai/[a-z0-9._/-]+(@sha256:[0-9a-f]{64}|:v[0-9][0-9a-z.-]*)$ ]] || exit 1
docker pull "$BASELINE_IMAGE"
- run:
name: Run migration startup regressions
environment:
@ -3033,28 +3050,29 @@ jobs:
- run:
name: Run Docker container with bad DATABASE_URL
command: |
set +e
docker run --name my-app \
-p 4000:4000 \
-e LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY=true \
-e DEFAULT_NUM_WORKERS_LITELLM_PROXY=1 \
-e DATABASE_URL="postgresql://wrong:wrong@wrong:5432/wrong" \
myapp:latest \
--port 4000 > docker_output.log 2>&1 || true
--port 4000 > docker_output.log 2>&1
echo "$?" > docker_exit_code
set -e
- run:
name: Display Docker logs
command: cat docker_output.log
- run:
name: Check for expected error
name: Proxy must refuse to serve on an unreachable database
command: |
if grep -q "Error: P1001: Can't reach database server at" docker_output.log && \
(grep -q "Database setup failed after multiple retries" docker_output.log || \
grep -q "ERROR: Application startup failed. Exiting." docker_output.log); then
echo "Expected error found. Test passed."
else
echo "Expected error not found. Test failed."
cat docker_output.log
exit 1
fi
fail() { echo "FAILED: $1"; cat docker_output.log; exit 1; }
exit_code="$(cat docker_exit_code)"
[ "$exit_code" -ne 0 ] || fail "proxy exited 0 with an unreachable database"
grep -q "P1001" docker_output.log || fail "log does not name the unreachable database server"
! grep -q "Application startup complete" docker_output.log || fail "proxy reached serving state"
! docker exec my-app true 2>/dev/null || fail "container is still running"
echo "Proxy refused to serve (exit $exit_code) and never reached startup. Test passed."
provider_replay_harness:
docker:
@ -3188,6 +3206,16 @@ workflows:
name: migration-legacy-and-pooling
suite: legacy
requires: [build_docker_database_image]
- migration_startup_tests:
name: migration-upgrade
suite: upgrade
baseline: true
requires: [build_docker_database_image]
- migration_startup_tests:
name: migration-upgrade-shaped
suite: shaped
baseline: true
requires: [build_docker_database_image]
migration_startup_scheduled:
triggers:
- schedule:

4
.circleci/scripts/run_integration.sh
View file

@ -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")

3
.circleci/scripts/run_migration_tests.py
View file

@ -13,6 +13,8 @@ SUITES: Final = {
"startup": (("test_startup.py",), 12),
"recovery": (("test_recovery.py",), 15),
"legacy": (("test_legacy.py", "test_pooling.py"), 11),
"upgrade": (("test_upgrade.py", "test_rolling_upgrade.py"), 5),
"shaped": (("test_shaped_database.py",), 1),
}
@ -93,6 +95,7 @@ def main() -> int:
{
**metadata,
"suite": suite,
"baseline_image": os.environ.get("LITELLM_MIGRATION_BASELINE_IMAGE", ""),
"expected_cases": expected,
"passed": passed,
"pytest_exit_code": result.returncode,

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

@ -205,7 +205,7 @@ def main(
native_module: Final = load_native_module(native_path)
native_module_loads: Final = native_module is not None
panic_test_hook_absent: Final = native_module is not None and not hasattr(native_module, "_panic_for_test")
native_size_limit: Final = 25_000_000
native_size_limit: Final = 40_000_000
native_size_within_limit: Final = native_member.file_size <= native_size_limit
validations: Final = (
(f"Python tag is {EXPECTED_PYTHON_TAG}", python_tag == EXPECTED_PYTHON_TAG),
@ -222,7 +222,7 @@ def main(
("Python extension entry point is present", extension_entry_point_present),
("Native module loads", native_module_loads),
("Production module omits the panic test hook", panic_test_hook_absent),
("Native extension does not exceed 25 MB", native_size_within_limit),
(f"Native extension does not exceed {native_size_limit / 1_000_000:.0f} MB", native_size_within_limit),
("Wheel contents are valid", not unexpected_members),
)
@ -267,7 +267,8 @@ def main(
),
(
not native_size_within_limit,
f"native extension exceeds 20 MB: {native_member.file_size / 1_000_000:.2f} MB",
f"native extension exceeds {native_size_limit / 1_000_000:.0f} MB: "
f"{native_member.file_size / 1_000_000:.2f} MB",
),
(bool(unexpected_members), f"wheel contains unexpected build artifacts: {', '.join(unexpected_members)}"),
)

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

@ -130,6 +130,10 @@ jobs:
echo "File content around line 43:"
head -50 litellm/litellm_core_utils/custom_logger_registry.py | tail -10
- name: Check MCP operation boundary
if: steps.changes.outputs.decision != 'skip'
run: uv run --no-sync python scripts/check_mcp_operation_boundary.py
- name: Run Ruff linting
if: steps.changes.outputs.decision != 'skip'
run: |

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

@ -130,7 +130,7 @@ jobs:
- 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
for features in '' aws google azure cyberark aws,google aws,azure google,azure aws,google,azure aws,google,cyberark aws,google,azure,cyberark; do
cargo test -p litellm-secrets --locked --no-default-features --features "$features"
done

1
Makefile
View file

@ -164,6 +164,7 @@ lint-format-check-changed: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
# Linting targets
lint-ruff: $(LINT_DEP_INSTALL)
$(UV_RUN) python scripts/check_mcp_operation_boundary.py
cd litellm && $(UV_RUN) ruff check . && cd ..
$(UV_RUN) ruff check --config ruff-tests.toml tests

3
README.md
View file

@ -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) | ✅ | ✅ | ✅ | | | | | | | |

1
backend/routes/allowlist.py
View file

@ -51,6 +51,7 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = (
"/cache_settings",
"/coordination_redis/",
"/cost_tracking",
"/cost_optimization/",
"/cost/",
"/credentials",
"/credential",

121
enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
View file

@ -2,10 +2,11 @@
Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if the cost has been tracked.
"""
from collections.abc import Sequence
from dataclasses import replace as dataclasses_replace
from datetime import datetime, timedelta, timezone
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, List, Literal, Optional, Tuple, cast
from typing import TYPE_CHECKING, Final, List, Literal, Optional, Protocol, Tuple, cast
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
@ -18,8 +19,8 @@ if TYPE_CHECKING:
from prisma import models as prisma_models
from litellm.integrations.prometheus import PrometheusLogger
from litellm.proxy._types import LiteLLM_ManagedObjectTable
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.repositories.prisma_protocols import TableActions
from litellm.router import Router
from litellm.types.router import Deployment
from litellm.types.utils import LiteLLMBatch
@ -41,6 +42,42 @@ TERMINAL_MANAGED_OBJECT_STATUSES: Final[Tuple[str, ...]] = (
)
class _ManagedObjectRow(Protocol):
@property
def id(self) -> str: ...
@property
def unified_object_id(self) -> str: ...
@property
def created_by(self) -> str | None: ...
@property
def file_object(self) -> object: ...
def _managed_object_table(prisma_client: "PrismaClient") -> "TableActions[_ManagedObjectRow]":
table: Final[TableActions[_ManagedObjectRow]] = prisma_client.db.litellm_managedobjecttable
return table
def _user_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.LiteLLM_UserTable]":
table: Final[TableActions[prisma_models.LiteLLM_UserTable]] = prisma_client.db.litellm_usertable
return table
def _token_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.LiteLLM_VerificationToken]":
table: Final[TableActions[prisma_models.LiteLLM_VerificationToken]] = (
prisma_client.db.litellm_verificationtoken
)
return table
def _team_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.LiteLLM_TeamTable]":
table: Final[TableActions[prisma_models.LiteLLM_TeamTable]] = prisma_client.db.litellm_teamtable
return table
class CheckBatchCost:
def __init__(
self,
@ -73,7 +110,7 @@ class CheckBatchCost:
inline for a batch the first poll cycle then accounts again.
"""
try:
await self.prisma_client.db.litellm_managedobjecttable.find_first(
await _managed_object_table(self.prisma_client).find_first(
where={"file_purpose": "batch", "batch_processed": False}
)
except Exception as probe_err:
@ -97,10 +134,8 @@ class CheckBatchCost:
if not user_id:
return {}
try:
user_row: prisma_models.LiteLLM_UserTable | None = (
await self.prisma_client.db.litellm_usertable.find_unique(
where={"user_id": user_id}
)
user_row: prisma_models.LiteLLM_UserTable | None = await _user_table(self.prisma_client).find_unique(
where={"user_id": user_id}
)
if user_row is None:
return {}
@ -117,11 +152,9 @@ class CheckBatchCost:
if not api_key:
return None
try:
key_row: prisma_models.LiteLLM_VerificationToken | None = (
await self.prisma_client.db.litellm_verificationtoken.find_unique(
where={"token": api_key}
)
)
key_row: prisma_models.LiteLLM_VerificationToken | None = await _token_table(
self.prisma_client
).find_unique(where={"token": api_key})
return getattr(key_row, "key_alias", None) if key_row is not None else None
except Exception as e:
verbose_proxy_logger.error(f"CheckBatchCost: could not look up key alias for batch {batch_id}: {e}")
@ -132,17 +165,15 @@ class CheckBatchCost:
if not team_id:
return None
try:
team_row: prisma_models.LiteLLM_TeamTable | None = (
await self.prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": team_id}
)
team_row: prisma_models.LiteLLM_TeamTable | None = await _team_table(self.prisma_client).find_unique(
where={"team_id": team_id}
)
return getattr(team_row, "team_alias", None) if team_row is not None else None
except Exception as e:
verbose_proxy_logger.error(f"CheckBatchCost: could not look up team alias for team {team_id}: {e}")
return None
async def _get_org_id(self, job: "LiteLLM_ManagedObjectTable", batch_id: str) -> str | None:
async def _get_org_id(self, job: "_ManagedObjectRow", batch_id: str) -> str | None:
org_id = getattr(job, "org_id", None)
if org_id:
return org_id
@ -150,11 +181,9 @@ class CheckBatchCost:
team_id = getattr(job, "team_id", None)
if api_key:
try:
key_row: prisma_models.LiteLLM_VerificationToken | None = (
await self.prisma_client.db.litellm_verificationtoken.find_unique(
where={"token": api_key}
)
)
key_row: prisma_models.LiteLLM_VerificationToken | None = await _token_table(
self.prisma_client
).find_unique(where={"token": api_key})
key_org_id = getattr(key_row, "organization_id", None) if key_row is not None else None
if key_org_id:
return key_org_id
@ -166,10 +195,8 @@ class CheckBatchCost:
if not team_id:
return None
try:
team_row: prisma_models.LiteLLM_TeamTable | None = (
await self.prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": team_id}
)
team_row: prisma_models.LiteLLM_TeamTable | None = await _team_table(self.prisma_client).find_unique(
where={"team_id": team_id}
)
return getattr(team_row, "organization_id", None) if team_row is not None else None
except Exception as e:
@ -177,7 +204,7 @@ class CheckBatchCost:
return None
async def _build_creator_attribution_metadata(
self, job: "LiteLLM_ManagedObjectTable", batch_id: str
self, job: "_ManagedObjectRow", batch_id: str
) -> dict[str, object]:
"""
Rebuild the spend-tracking metadata for the key, team, and tags that created the
@ -225,7 +252,7 @@ class CheckBatchCost:
should not be polled.
"""
cutoff: Final = datetime.now(timezone.utc) - timedelta(days=MANAGED_OBJECT_STALENESS_CUTOFF_DAYS)
result: Final = await self.prisma_client.db.litellm_managedobjecttable.update_many(
result: Final = await _managed_object_table(self.prisma_client).update_many(
where={
"file_purpose": "batch",
"status": {"not_in": list(TERMINAL_MANAGED_OBJECT_STATUSES)},
@ -244,7 +271,7 @@ class CheckBatchCost:
# A row already in a terminal status is never rewritten by the sweep above, so
# without this it keeps a poll-page slot forever and starves newer batches.
retired: Final = await self.prisma_client.db.litellm_managedobjecttable.update_many(
retired: Final = await _managed_object_table(self.prisma_client).update_many(
where={
"file_purpose": "batch",
"batch_processed": False,
@ -259,9 +286,9 @@ class CheckBatchCost:
f"{MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days that were never costed"
)
async def _fallback_find_jobs(self) -> list:
async def _fallback_find_jobs(self) -> "Sequence[_ManagedObjectRow]":
"""Query batch jobs without the batch_processed filter (for older schemas)."""
return await self.prisma_client.db.litellm_managedobjecttable.find_many(
return await _managed_object_table(self.prisma_client).find_many(
where={
"file_purpose": "batch",
"status": {
@ -279,7 +306,7 @@ class CheckBatchCost:
order={"created_at": "asc"},
)
async def _retire_job(self, job: "LiteLLM_ManagedObjectTable", reason: str) -> None:
async def _retire_job(self, job: "_ManagedObjectRow", reason: str) -> None:
"""
Take a row that can never be costed out of the poll page. Leaving it selectable
would burn one of the MAX_OBJECTS_PER_POLL_CYCLE slots on every future cycle, and
@ -292,7 +319,7 @@ class CheckBatchCost:
else {"status": "stale_expired"}
)
try:
await self.prisma_client.db.litellm_managedobjecttable.update(
await _managed_object_table(self.prisma_client).update(
where={"id": job.id},
data=data,
)
@ -306,7 +333,7 @@ class CheckBatchCost:
"so it will no longer be polled"
)
async def _claim_job_for_costing(self, job: "LiteLLM_ManagedObjectTable") -> bool:
async def _claim_job_for_costing(self, job: "_ManagedObjectRow") -> bool:
"""
Atomically flip batch_processed from false to true, returning whether this pod won
the row. Every pod and uvicorn worker schedules its own poller against the shared
@ -321,7 +348,7 @@ class CheckBatchCost:
if not self._has_batch_processed_column:
return True
try:
claimed: Final = await self.prisma_client.db.litellm_managedobjecttable.update_many(
claimed: Final = await _managed_object_table(self.prisma_client).update_many(
where={"id": job.id, "batch_processed": False},
data={"batch_processed": True},
)
@ -332,7 +359,7 @@ class CheckBatchCost:
return False
return claimed > 0
async def _release_job_claim(self, job: "LiteLLM_ManagedObjectTable") -> None:
async def _release_job_claim(self, job: "_ManagedObjectRow") -> None:
"""Give a claimed row back once billing it failed, so a later poll cycle retries it.
Safe to match on batch_processed=True: while this poller is active the retrieve
@ -342,7 +369,7 @@ class CheckBatchCost:
if not self._has_batch_processed_column:
return
try:
await self.prisma_client.db.litellm_managedobjecttable.update_many(
await _managed_object_table(self.prisma_client).update_many(
where={"id": job.id, "batch_processed": True},
data={"batch_processed": False},
)
@ -353,7 +380,7 @@ class CheckBatchCost:
)
@staticmethod
def _has_unified_id_without_model(job: "LiteLLM_ManagedObjectTable") -> bool:
def _has_unified_id_without_model(job: "_ManagedObjectRow") -> bool:
"""A unified id that decodes but carries no model_id can never be routed."""
from litellm.proxy.openai_files_endpoints.common_utils import (
convert_b64_uid_to_unified_uid,
@ -402,7 +429,7 @@ class CheckBatchCost:
return isinstance(error, (NotFoundError, openai.NotFoundError)) and output_file_id in str(error)
async def _finalize_unbilled_terminal_job(
self, job: "prisma_models.LiteLLM_ManagedObjectTable", response: "LiteLLMBatch"
self, job: "_ManagedObjectRow", response: "LiteLLMBatch"
) -> None:
"""Persist a terminal batch that has nothing billable, converting any raw
provider file ids to managed ids, and take it out of the poll page."""
@ -426,7 +453,7 @@ class CheckBatchCost:
"file_object": response.model_dump_json(),
**({"batch_processed": True} if self._has_batch_processed_column else {}),
}
await self.prisma_client.db.litellm_managedobjecttable.update(
await _managed_object_table(self.prisma_client).update(
where={"id": job.id},
data=update_data,
)
@ -447,7 +474,7 @@ class CheckBatchCost:
def _resolve_job_routing(
self,
job: "LiteLLM_ManagedObjectTable",
job: "_ManagedObjectRow",
prom_logger: Optional["PrometheusLogger"],
) -> Optional[Tuple[str, str]]:
"""
@ -524,7 +551,7 @@ class CheckBatchCost:
def _resolve_unmanaged_provider_routing(
self,
job: "LiteLLM_ManagedObjectTable",
job: "_ManagedObjectRow",
prom_logger: Optional["PrometheusLogger"],
llm_provider: str,
bare_model_name: str,
@ -620,7 +647,7 @@ class CheckBatchCost:
@classmethod
def _get_managed_file_model_name(
cls,
job: "LiteLLM_ManagedObjectTable",
job: "_ManagedObjectRow",
deployment_info: "Deployment",
) -> Optional[str]:
"""
@ -640,7 +667,7 @@ class CheckBatchCost:
)
@staticmethod
def _get_input_file_id(job: "LiteLLM_ManagedObjectTable") -> Optional[str]:
def _get_input_file_id(job: "_ManagedObjectRow") -> Optional[str]:
import json
from litellm.types.utils import LiteLLMBatch
@ -660,7 +687,7 @@ class CheckBatchCost:
async def _track_completed_batch_cost(
self,
job: "LiteLLM_ManagedObjectTable",
job: "_ManagedObjectRow",
response: "LiteLLMBatch",
model_id: str,
batch_id: str,
@ -936,7 +963,7 @@ class CheckBatchCost:
# endpoint may transition a batch to "complete" before
# CheckBatchCost runs. The batch_processed=False filter
# already prevents reprocessing finished batches.
jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many(
jobs = await _managed_object_table(self.prisma_client).find_many(
where={
"file_purpose": "batch",
"batch_processed": False,
@ -1038,7 +1065,7 @@ class CheckBatchCost:
}
if self._has_batch_processed_column:
update_data["batch_processed"] = True
await self.prisma_client.db.litellm_managedobjecttable.update(
await _managed_object_table(self.prisma_client).update(
where={"id": job.id},
data=update_data,
)

22
enterprise/litellm_enterprise/proxy/hooks/managed_files.py
View file

@ -481,10 +481,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
"""
if self.prisma_client is None:
return
managed_object = (
await self.prisma_client.db.litellm_managedobjecttable.find_first(
where={"OR": [{"unified_object_id": object_id}, {"model_object_id": object_id}]}
)
managed_object = await _managed_object_table(self.prisma_client).find_first(
where={"OR": [{"unified_object_id": object_id}, {"model_object_id": object_id}]}
)
if managed_object is None:
return
@ -509,10 +507,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
"""
if self.prisma_client is None:
return
managed_file = (
await self.prisma_client.db.litellm_managedfiletable.find_first(
where={"OR": [{"unified_file_id": file_id}, {"flat_model_file_ids": {"has": file_id}}]}
)
managed_file = await _managed_file_table(self.prisma_client).find_first(
where={"OR": [{"unified_file_id": file_id}, {"flat_model_file_ids": {"has": file_id}}]}
)
if managed_file is None:
return
@ -535,8 +531,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
provider_file_ids = tuple(
file_id
for file_id in (
getattr(response, "output_file_id", None),
getattr(response, "error_file_id", None),
response.output_file_id,
response.error_file_id,
)
if file_id and not _is_base64_encoded_unified_file_id(file_id)
)
@ -544,10 +540,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
return
if self.prisma_client is None:
return
batch_row = (
await self.prisma_client.db.litellm_managedobjecttable.find_first(
where={"unified_object_id": response.id}
)
batch_row = await _managed_object_table(self.prisma_client).find_first(
where={"unified_object_id": response.id}
)
if batch_row is None or (
batch_row.created_by is None and batch_row.team_id is None

1
litellm-proxy-extras/litellm_proxy_extras/migrations/20260920041500_add_policy_attachment_is_default/migration.sql Normal file
View file

@ -0,0 +1 @@
ALTER TABLE "LiteLLM_PolicyAttachmentTable" ADD COLUMN IF NOT EXISTS "is_default" BOOLEAN NOT NULL DEFAULT false;

1
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -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

261
litellm-rust/Cargo.lock generated
View file

@ -115,6 +115,28 @@ dependencies = [
"pin-project-lite",
]
[[package]]
name = "async-stream"
version = "0.3.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0b5a71a6f37880a80d1d7f19efd781e4b5de42c88f0722cc13bcb6cc2cfe8476"
dependencies = [
"async-stream-impl",
"futures-core",
"pin-project-lite",
]
[[package]]
name = "async-stream-impl"
version = "0.3.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c7c24de15d275a1ecfd47a380fb4d5ec9bfe0933f309ed5e705b775596a3574d"
dependencies = [
"proc-macro2",
"quote",
"syn 2.0.119",
]
[[package]]
name = "async-trait"
version = "0.1.91"
@ -599,6 +621,37 @@ dependencies = [
"url",
]
[[package]]
name = "azure_storage_blob"
version = "1.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "17b10207ecf7d666df6940b50051f433b3cd5d2b9b1dd190613208d7a84e7eed"
dependencies = [
"async-stream",
"async-trait",
"azure_core",
"azure_storage_common",
"bytes",
"futures",
"percent-encoding",
"pin-project",
"serde",
"serde_json",
"time",
"tokio",
]
[[package]]
name = "azure_storage_common"
version = "1.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0af2e6aeb8d76b17fc998f453c320913f73787b944e3cc29509d19411fa0321d"
dependencies = [
"azure_core",
"serde",
"time",
]
[[package]]
name = "base64"
version = "0.13.1"
@ -927,6 +980,12 @@ dependencies = [
"libc",
]
[[package]]
name = "crc16"
version = "0.4.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "338089f42c427b86394a5ee60ff321da23a5c89c9d89514c829687b26359fcff"
[[package]]
name = "crc32fast"
version = "1.5.1"
@ -2460,8 +2519,25 @@ dependencies = [
"rstest",
"serde",
"serde_json",
"sha2 0.10.9",
"thiserror 2.0.19",
"tokio",
]
[[package]]
name = "litellm-cache-azure-blob"
version = "0.1.0"
dependencies = [
"async-trait",
"azure_core",
"azure_storage_blob",
"futures-util",
"litellm-auth-azure",
"litellm-auth-types",
"litellm-cache",
"litellm-cache-response",
"serde_json",
"tokio",
"url",
]
[[package]]
@ -2479,12 +2555,29 @@ name = "litellm-cache-redis"
version = "0.1.0"
dependencies = [
"litellm-cache",
"r2d2",
"redis",
"redis-test",
"serde_json",
"tokio",
]
[[package]]
name = "litellm-cache-response"
version = "0.1.0"
dependencies = [
"litellm-cache",
"litellm-cache-memory",
"litellm-cache-redis",
"py_literal",
"redis",
"redis-test",
"serde",
"serde_json",
"sha2 0.10.9",
"tokio",
]
[[package]]
name = "litellm-callbacks-legacy-python"
version = "0.1.0"
@ -2648,6 +2741,11 @@ dependencies = [
"futures-util",
"litellm-auth",
"litellm-auth-gcp",
"litellm-cache",
"litellm-cache-azure-blob",
"litellm-cache-memory",
"litellm-cache-redis",
"litellm-cache-response",
"litellm-callbacks-legacy-python",
"litellm-core",
"litellm-core-utils",
@ -2659,7 +2757,9 @@ dependencies = [
"pyo3",
"pyo3-async-runtimes",
"rstest",
"serde",
"serde_json",
"serde_with",
"tokio",
"tokio-tungstenite",
]
@ -2675,6 +2775,8 @@ dependencies = [
"jsonwebtoken",
"litellm-core-utils",
"litellm-secrets-aws",
"litellm-secrets-azure",
"litellm-secrets-cyberark",
"litellm-secrets-google",
"litellm-secrets-types",
"moka",
@ -2709,6 +2811,46 @@ dependencies = [
"wiremock",
]
[[package]]
name = "litellm-secrets-azure"
version = "0.1.0"
dependencies = [
"litellm-auth-azure",
"litellm-auth-types",
"litellm-core-utils",
"litellm-secrets-types",
"percent-encoding",
"reqwest 0.12.28",
"rstest",
"serde",
"serde_json",
"sha2 0.10.9",
"thiserror 2.0.19",
"tokio",
"veil",
"wiremock",
]
[[package]]
name = "litellm-secrets-cyberark"
version = "0.1.0"
dependencies = [
"base64 0.22.1",
"litellm-core-utils",
"litellm-secrets-types",
"moka",
"percent-encoding",
"reqwest 0.12.28",
"rstest",
"serde",
"serde_json",
"thiserror 2.0.19",
"tokio",
"tracing",
"veil",
"wiremock",
]
[[package]]
name = "litellm-secrets-google"
version = "0.1.0"
@ -2947,6 +3089,16 @@ dependencies = [
"minimal-lexical",
]
[[package]]
name = "num-bigint"
version = "0.4.8"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c89e69e7e0f03bea5ef08013795c25018e101932225a656383bd384495ecc367"
dependencies = [
"num-integer",
"num-traits",
]
[[package]]
name = "num-bigint"
version = "0.5.1"
@ -2957,6 +3109,15 @@ dependencies = [
"num-traits",
]
[[package]]
name = "num-complex"
version = "0.4.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495"
dependencies = [
"num-traits",
]
[[package]]
name = "num-conv"
version = "0.2.2"
@ -3130,6 +3291,48 @@ version = "2.3.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220"
[[package]]
name = "pest"
version = "2.9.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6d45aeb61b4bf818e12d4205f2466f8c4748f85f4fce0146d1c03d69d753f0ad"
dependencies = [
"memchr",
"ucd-trie",
]
[[package]]
name = "pest_derive"
version = "2.9.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "89cc5a242e25ed4e7704d0be240f2cfbe20a8c27e7e252d94835be93d92dc39f"
dependencies = [
"pest",
"pest_generator",
]
[[package]]
name = "pest_generator"
version = "2.9.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7abf21475cc3820fe4b2ca2dc2142902f67a02189f3b5b3a229f4febc01a43e5"
dependencies = [
"pest",
"pest_meta",
"proc-macro2",
"quote",
"syn 2.0.119",
]
[[package]]
name = "pest_meta"
version = "2.9.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "adba4db388f687393c18c51348d44a41d870ca9df71a2c98172ea3035dc6936e"
dependencies = [
"pest",
]
[[package]]
name = "pin-project"
version = "1.1.13"
@ -3304,6 +3507,19 @@ dependencies = [
"prost",
]
[[package]]
name = "py_literal"
version = "0.4.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "102df7a3d46db9d3891f178dcc826dc270a6746277a9ae6436f8d29fd490a8e1"
dependencies = [
"num-bigint 0.4.8",
"num-complex",
"num-traits",
"pest",
"pest_derive",
]
[[package]]
name = "pyo3"
version = "0.29.2"
@ -3391,6 +3607,16 @@ version = "1.2.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a1d01941d82fa2ab50be1e79e6714289dd7cde78eba4c074bc5a4374f650dfe0"
[[package]]
name = "quick-xml"
version = "0.41.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e660451e55124f798a69a5af3f49ccfbefbd41910eefd25caf2393e1f3473ec1"
dependencies = [
"memchr",
"serde",
]
[[package]]
name = "quinn"
version = "0.11.11"
@ -3469,6 +3695,17 @@ version = "6.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf"
[[package]]
name = "r2d2"
version = "0.8.10"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "51de85fb3fb6524929c8a2eb85e6b6d363de4e8c48f9e2c2eac4944abc181c93"
dependencies = [
"log",
"parking_lot",
"scheduled-thread-pool",
]
[[package]]
name = "rand"
version = "0.8.7"
@ -3602,9 +3839,13 @@ checksum = "2acbc41a996f7652b2ddd9dfd98cc4ff602cfd742ae35382f07f608405ab50ed"
dependencies = [
"arcstr",
"combine",
"crc16",
"itoa",
"num-bigint",
"num-bigint 0.5.1",
"percent-encoding",
"rand 0.10.2",
"rustls 0.23.42",
"rustls-native-certs",
"ryu",
"sha1_smol",
"socket2 0.6.5",
@ -4001,6 +4242,15 @@ dependencies = [
"windows-sys 0.61.2",
]
[[package]]
name = "scheduled-thread-pool"
version = "0.2.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3cbc66816425a074528352f5789333ecff06ca41b36b0b0efdfbb29edc391a19"
dependencies = [
"parking_lot",
]
[[package]]
name = "schemars"
version = "0.9.0"
@ -4912,6 +5162,7 @@ dependencies = [
"base64 0.22.1",
"bytes",
"futures",
"quick-xml",
"serde",
"serde_json",
"url",
@ -4954,6 +5205,12 @@ dependencies = [
"syn 2.0.119",
]
[[package]]
name = "ucd-trie"
version = "0.1.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2896d95c02a80c6d6a5d6e953d479f5ddf2dfdb6a244441010e373ac0fb88971"
[[package]]
name = "unarray"
version = "0.1.4"

5
litellm-rust/Cargo.toml
View file

@ -22,12 +22,17 @@ 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-secrets-azure = { path = "crates/secrets-azure" }
litellm-secrets-cyberark = { path = "crates/secrets-cyberark" }
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-azure-blob = { path = "crates/cache-azure-blob" }
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" }

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

@ -4,4 +4,4 @@ mod resolve;
mod types;
pub use resolve::AzureAuthService;
pub use types::AzureAuthInputs;
pub use types::{AzureAuthInputs, ConfigValue};

15
litellm-rust/crates/auth-azure/src/types.rs
View file

@ -51,6 +51,21 @@ pub struct AzureAuthInputs {
}
impl AzureAuthInputs {
pub fn default_credential_for_scope(scope: &str) -> Self {
Self {
azure_scope: ConfigValue::Value(Sourced::new(
scope.to_string(),
InputSource::Deployment,
)),
azure_credential: ConfigValue::Value(Sourced::new(
"DefaultAzureCredential".to_string(),
InputSource::Deployment,
)),
enable_azure_ad_token_refresh: Sourced::new(true, InputSource::Deployment),
..Self::default()
}
}
pub fn or_configured_token_refresh(self, enabled: bool) -> Self {
if *self.enable_azure_ad_token_refresh.value() || !enabled {
return self;

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

@ -0,0 +1,22 @@
[package]
name = "litellm-cache-azure-blob"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
litellm-auth-azure.workspace = true
litellm-auth-types.workspace = true
litellm-cache.workspace = true
async-trait = "0.1"
azure_core = "1.1.0"
azure_storage_blob = "1.1.0"
futures-util.workspace = true
tokio.workspace = true
url.workspace = true
[dev-dependencies]
litellm-cache-response.workspace = true
serde_json.workspace = true

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

@ -0,0 +1,254 @@
use std::{sync::Arc, time::Duration};
use azure_core::{
credentials::TokenCredential,
error::ErrorKind,
http::{ClientOptions, RequestContent},
};
use azure_storage_blob::{
BlobContainerClient, BlobContainerClientOptions,
models::{BlobClientUploadOptions, StorageErrorCode},
};
use futures_util::{TryStreamExt, future::try_join_all};
use litellm_cache::{
BaseCache, BatchCache, CacheCodec, CacheConnectionResult, CacheConnectionStatus, Error,
ExactCacheContext, FlushCache,
};
use tokio::runtime::Handle;
use url::Url;
use crate::credential::AzureBlobCredential;
pub struct AzureBlobCache<C> {
container: BlobContainerClient,
codec: C,
runtime: Handle,
account_url: String,
container_name: String,
}
impl<C: CacheCodec> AzureBlobCache<C> {
pub async fn connect(
account_url: &str,
container: &str,
codec: C,
runtime: Handle,
) -> Result<Self, Error> {
Self::connect_with_options(
account_url,
container,
Some(Arc::new(AzureBlobCredential::default())),
ClientOptions::default(),
codec,
runtime,
)
.await
}
pub async fn connect_with_options(
account_url: &str,
container: &str,
credential: Option<Arc<dyn TokenCredential>>,
client_options: ClientOptions,
codec: C,
runtime: Handle,
) -> Result<Self, Error> {
let parsed = Url::parse(account_url).map_err(|_| Error::Unavailable)?;
let account_url = parsed.as_str().trim_end_matches('/').to_string();
let container_url = {
let mut url = parsed;
url.path_segments_mut()
.map_err(|()| Error::Unavailable)?
.pop_if_empty()
.push(container);
url
};
let client = BlobContainerClient::new(
container_url,
credential,
Some(BlobContainerClientOptions {
client_options,
..BlobContainerClientOptions::default()
}),
)
.map_err(|_| Error::Unavailable)?;
let cache = Self {
container: client,
codec,
runtime,
account_url,
container_name: container.to_string(),
};
cache.create_container().await?;
Ok(cache)
}
pub fn account_url(&self) -> &str {
&self.account_url
}
pub fn container_name(&self) -> &str {
&self.container_name
}
async fn create_container(&self) -> Result<(), Error> {
match self.container.create(None).await {
Ok(_) => Ok(()),
Err(error) if is_storage_error(&error, StorageErrorCode::ContainerAlreadyExists) => {
Ok(())
}
Err(_) => Err(Error::Unavailable),
}
}
async fn upload(&self, key: &str, value: &C::Value, overwrite: bool) -> Result<(), Error> {
let payload = self.codec.encode(value)?;
let options = (!overwrite).then(|| BlobClientUploadOptions::default().if_not_exists());
match self
.container
.blob_client(key)
.upload(RequestContent::from(payload), options)
.await
{
Ok(_) => Ok(()),
Err(error) if !overwrite && is_already_present(&error) => Ok(()),
Err(_) => Err(Error::Unavailable),
}
}
async fn download(&self, key: &str) -> Result<Option<C::Value>, Error> {
let response = match self.container.blob_client(key).download(None).await {
Ok(response) => response,
Err(error) if is_storage_error(&error, StorageErrorCode::BlobNotFound) => {
return Ok(None);
}
Err(_) => return Err(Error::Unavailable),
};
let bytes = response
.body
.collect()
.await
.map_err(|_| Error::Unavailable)?;
self.codec.decode(&bytes).map(Some)
}
async fn delete_all_blobs(&self) -> Result<(), Error> {
let mut pages = self
.container
.list_blobs(None)
.map_err(|_| Error::Unavailable)?
.into_pages();
while let Some(page) = pages.try_next().await.map_err(|_| Error::Unavailable)? {
let page = page.into_model().map_err(|_| Error::Unavailable)?;
for name in page.blob_items.into_iter().filter_map(|item| item.name) {
self.container
.blob_client(&name)
.delete(None)
.await
.map_err(|_| Error::Unavailable)?;
}
}
Ok(())
}
fn block_on<T>(&self, future: impl Future<Output = T>) -> T {
self.runtime.block_on(future)
}
}
fn is_already_present(error: &azure_core::Error) -> bool {
is_storage_error(error, StorageErrorCode::BlobAlreadyExists)
|| is_storage_error(error, StorageErrorCode::ConditionNotMet)
}
fn is_storage_error(error: &azure_core::Error, code: StorageErrorCode) -> bool {
matches!(
error.kind(),
ErrorKind::HttpResponse {
error_code: Some(error_code),
..
} if error_code == code.as_ref()
)
}
impl<C: CacheCodec> BaseCache for AzureBlobCache<C> {
type Value = C::Value;
type Context = ExactCacheContext;
fn get_ttl(&self, _: &ExactCacheContext) -> Option<Duration> {
None
}
fn set_cache(&self, key: &str, value: C::Value, _: &ExactCacheContext) -> Result<(), Error> {
self.block_on(self.upload(key, &value, false))
}
fn get_cache(&self, key: &str, _: &ExactCacheContext) -> Result<Option<C::Value>, Error> {
self.block_on(self.download(key))
}
async fn async_set_cache(
&self,
key: &str,
value: C::Value,
_: ExactCacheContext,
) -> Result<(), Error> {
self.upload(key, &value, true).await
}
async fn async_get_cache(
&self,
key: &str,
_: &ExactCacheContext,
) -> Result<Option<C::Value>, Error> {
self.download(key).await
}
async fn async_set_cache_pipeline(
&self,
entries: Vec<(String, C::Value)>,
_: ExactCacheContext,
) -> Result<(), Error> {
try_join_all(
entries
.iter()
.map(|(key, value)| self.upload(key, value, true)),
)
.await
.map(drop)
}
async fn disconnect(&self) -> Result<(), Error> {
Ok(())
}
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
Ok(match self.container.get_properties(None).await {
Ok(_) => CacheConnectionResult {
status: CacheConnectionStatus::Success,
message: "Azure Blob cache connection test successful".into(),
error: None,
},
Err(error) => CacheConnectionResult {
status: CacheConnectionStatus::Failed,
message: format!("Azure Blob connection failed: {error}"),
error: Some(error.to_string()),
},
})
}
}
impl<C: CacheCodec> BatchCache for AzureBlobCache<C> {}
impl<C: CacheCodec> FlushCache for AzureBlobCache<C> {
fn flush_cache(&self) -> Result<(), Error> {
self.block_on(self.delete_all_blobs())
}
async fn async_flush_cache(&self) -> Result<(), Error> {
self.delete_all_blobs().await
}
}
#[cfg(test)]
mod tests;

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

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

84
litellm-rust/crates/cache-azure-blob/src/credential.rs Normal file
View file

@ -0,0 +1,84 @@
use std::{
fmt,
sync::Arc,
time::{Duration, SystemTime},
};
use azure_core::{
credentials::{AccessToken, TokenCredential, TokenRequestOptions},
error::ErrorKind,
time::OffsetDateTime,
};
use litellm_auth_azure::{AzureAuthInputs, AzureAuthService};
use litellm_auth_types::ResolvedCredential;
const STATIC_TOKEN_LIFETIME: Duration = Duration::from_secs(300);
const LLM_TOKEN_ENV: &str = "AZURE_AD_TOKEN";
type EnvLookup = Arc<dyn Fn(&str) -> Option<String> + Send + Sync>;
pub struct AzureBlobCredential {
service: AzureAuthService,
env_lookup: EnvLookup,
}
impl fmt::Debug for AzureBlobCredential {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str("AzureBlobCredential")
}
}
impl Default for AzureBlobCredential {
fn default() -> Self {
Self::new(
AzureAuthService::default(),
Arc::new(|name| std::env::var(name).ok()),
)
}
}
impl AzureBlobCredential {
pub fn new(service: AzureAuthService, env_lookup: EnvLookup) -> Self {
Self {
service,
env_lookup,
}
}
}
#[async_trait::async_trait]
impl TokenCredential for AzureBlobCredential {
async fn get_token(
&self,
scopes: &[&str],
_options: Option<TokenRequestOptions<'_>>,
) -> azure_core::Result<AccessToken> {
let env_lookup = &self.env_lookup;
let lookup = move |name: &str| (name != LLM_TOKEN_ENV).then(|| env_lookup(name)).flatten();
let credential = self
.service
.get_azure_ad_token(
&AzureAuthInputs::default_credential_for_scope(&scopes.join(" ")),
&lookup,
)
.await
.map_err(|error| {
azure_core::Error::with_message(ErrorKind::Credential, error.to_string())
})?
.ok_or_else(|| {
azure_core::Error::with_message(
ErrorKind::Credential,
"no Azure credential is available for blob storage",
)
})?;
let (token, expires_on) = match credential.into_value() {
ResolvedCredential::AccessToken { token, expires_on } => (token, expires_on),
ResolvedCredential::Static(token) => (token, None),
};
let expires_on = expires_on.unwrap_or_else(|| SystemTime::now() + STATIC_TOKEN_LIFETIME);
Ok(AccessToken::new(
token.expose().to_string(),
OffsetDateTime::from(expires_on),
))
}
}

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

@ -0,0 +1,5 @@
mod cache;
mod credential;
pub use cache::AzureBlobCache;
pub use credential::AzureBlobCredential;

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

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

@ -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

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

@ -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);
}
}

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

@ -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()]))
);
}

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

@ -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 = ["cluster", "tls-rustls"] }
r2d2 = "0.8.10"
tokio.workspace = true
[dev-dependencies]
redis-test = "1.0.4"
serde_json.workspace = true

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

@ -1,58 +1,201 @@
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;
use crate::topology::RedisTopology;
mod connection;
mod operations;
pub(crate) use connection::ConnectionRef;
use connection::{ClusterConnectionManager, ConnectionManager};
pub use operations::{
RedisArg, RedisLpopOperation, RedisLpopResult, RedisRpushOperation, RedisScript,
};
const DEFAULT_TTL: Duration = Duration::from_secs(600);
const KEY_PREFIX: &str = "litellm-cache:";
const REDIS_TIMEOUT: Duration = Duration::from_secs(5);
const REDIS_POOL_SIZE: u32 = 16;
pub struct RedisCache<C = redis::Connection> {
connection: Arc<Mutex<C>>,
default_ttl: Duration,
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>),
Cluster(r2d2::Pool<ClusterConnectionManager>),
Fixed(Mutex<C>),
}
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))
}
}
impl<C> RedisCache<C>
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::Node(&mut pooled.connection));
pooled.failed = matches!(result, Err(Error::Unavailable));
result
}
Self::Cluster(pool) => {
let mut pooled = pool.get().map_err(|_| Error::Unavailable)?;
let result = operation(&mut ConnectionRef::Cluster(&mut pooled.connection));
pooled.failed = matches!(result, Err(Error::Unavailable));
result
}
Self::Fixed(connection) => {
let mut connection = connection.lock().map_err(|_| Error::Unavailable)?;
operation(&mut ConnectionRef::Node(&mut *connection))
}
}
}
}
pub struct RedisCache<S, C = redis::Connection> {
connections: Arc<Connections<C>>,
default_ttl: Duration,
codec: S,
namespace: Option<String>,
topology: RedisTopology,
}
impl<S: CacheCodec> RedisCache<S> {
pub fn new(url: &str, default_ttl: Option<Duration>, codec: S) -> Result<Self, Error> {
Self::connect(url, &RedisTopology::Standalone, default_ttl, codec)
}
pub fn connect(
url: &str,
topology: &RedisTopology,
default_ttl: Option<Duration>,
codec: S,
) -> Result<Self, Error> {
let connections = match topology {
RedisTopology::Standalone => Connections::Pool(pool(ConnectionManager::open(url)?)?),
RedisTopology::Cluster { startup_nodes } => {
Connections::Cluster(pool(ClusterConnectionManager::open(url, startup_nodes)?)?)
}
};
Ok(Self {
connections: Arc::new(connections),
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
codec,
namespace: None,
topology: topology.clone(),
})
}
}
fn pool<M: r2d2::ManageConnection>(manager: M) -> Result<r2d2::Pool<M>, Error> {
r2d2::Pool::builder()
.max_size(REDIS_POOL_SIZE)
.min_idle(Some(0))
.connection_timeout(REDIS_TIMEOUT)
.test_on_check_out(false)
.build(manager)
.map_err(|_| Error::Unavailable)
}
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,
topology: RedisTopology::Standalone,
}
}
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
pub fn topology(&self) -> &RedisTopology {
&self.topology
}
fn encode(value: &CacheEntry) -> Result<Vec<u8>, Error> {
serde_json::to_vec(value).map_err(|_| Error::InvalidEntry)
fn namespaced_key(&self, key: &str) -> String {
namespaced_key(self.namespace.as_deref(), key)
}
fn decode(value: Vec<u8>) -> Result<CacheEntry, Error> {
serde_json::from_slice(&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 flush_matching(connection: &mut ConnectionRef<'_>, pattern: &str) -> Result<(), Error> {
connection.scan(pattern, 1000, |connection, keys| {
if !keys.is_empty() {
connection
.del::<_, usize>(keys)
.map_err(|_| Error::Unavailable)?;
}
Ok(true)
})
}
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 +204,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)?;
}
Ok(())
.collect::<Result<Vec<_>, _>>()?;
let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl));
if entries.is_empty() {
return Ok(());
}
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
let commands = entries
.into_iter()
.map(|(key, payload)| {
let mut command = redis::cmd("SETEX");
command.arg(key).arg(ttl).arg(payload);
command
})
.collect();
connection.pipeline(commands).map(drop)
})
.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 connection.ping() {
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 +623,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 +638,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 +661,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 +680,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,

392
litellm-rust/crates/cache-redis/src/cache/connection.rs vendored Normal file
View file

@ -0,0 +1,392 @@
use std::collections::HashMap;
use litellm_cache::Error;
use redis::{
ConnectionAddr, ConnectionInfo, ConnectionLike, IntoConnectionInfo,
cluster::{ClusterClient, ClusterClientBuilder, ClusterConnection, NodeAddress},
cluster_routing::{
MultipleNodeRoutingInfo, ResponsePolicy, RoutingInfo, SingleNodeRoutingInfo, Slot,
},
};
use super::REDIS_TIMEOUT;
use crate::topology::RedisNode;
pub(super) struct PooledConnection<C> {
pub(super) connection: C,
pub(super) failed: bool,
}
/// Pools connections without a checkout PING, which would double every operation's round trips.
/// A timed-out command leaves its reply on the socket while redis still reports the connection
/// open, so any connection whose operation failed is discarded instead of being reused.
pub(super) struct ConnectionManager(redis::Client);
impl ConnectionManager {
pub(super) fn open(url: &str) -> Result<Self, Error> {
redis::Client::open(url)
.map(Self)
.map_err(|_| Error::Unavailable)
}
}
impl r2d2::ManageConnection for ConnectionManager {
type Connection = PooledConnection<redis::Connection>;
type Error = redis::RedisError;
fn connect(&self) -> Result<Self::Connection, redis::RedisError> {
let connection = self.0.get_connection()?;
connection.set_read_timeout(Some(REDIS_TIMEOUT))?;
connection.set_write_timeout(Some(REDIS_TIMEOUT))?;
Ok(PooledConnection {
connection,
failed: false,
})
}
fn is_valid(&self, connection: &mut Self::Connection) -> Result<(), redis::RedisError> {
redis::cmd("PING").query::<String>(&mut connection.connection)?;
Ok(())
}
fn has_broken(&self, connection: &mut Self::Connection) -> bool {
connection.failed || !redis::ConnectionLike::is_open(&connection.connection)
}
}
pub(super) struct ClusterConnectionManager(ClusterClient);
impl ClusterConnectionManager {
pub(super) fn open(url: &str, startup_nodes: &[RedisNode]) -> Result<Self, Error> {
if startup_nodes.is_empty() {
return Err(Error::Unavailable);
}
let info = url.into_connection_info().map_err(|_| Error::Unavailable)?;
let nodes = startup_nodes
.iter()
.map(|node| node_info(&info, node))
.collect::<Result<Vec<_>, _>>()?;
ClusterClientBuilder::new(nodes)
.connection_timeout(REDIS_TIMEOUT)
.response_timeout(REDIS_TIMEOUT)
.build()
.map(Self)
.map_err(|_| Error::Unavailable)
}
}
fn node_info(info: &ConnectionInfo, node: &RedisNode) -> Result<ConnectionInfo, Error> {
let addr = match info.addr() {
ConnectionAddr::Tcp(..) => ConnectionAddr::Tcp(node.host.clone(), node.port),
ConnectionAddr::TcpTls {
insecure,
tls_params,
..
} => ConnectionAddr::TcpTls {
host: node.host.clone(),
port: node.port,
insecure: *insecure,
tls_params: tls_params.clone(),
},
_ => return Err(Error::Unavailable),
};
Ok(info.clone().set_addr(addr))
}
impl r2d2::ManageConnection for ClusterConnectionManager {
type Connection = PooledConnection<ClusterConnection>;
type Error = redis::RedisError;
fn connect(&self) -> Result<Self::Connection, redis::RedisError> {
let connection = self.0.get_connection()?;
connection.set_read_timeout(Some(REDIS_TIMEOUT))?;
connection.set_write_timeout(Some(REDIS_TIMEOUT))?;
Ok(PooledConnection {
connection,
failed: false,
})
}
fn is_valid(&self, connection: &mut Self::Connection) -> Result<(), redis::RedisError> {
redis::cmd("PING").query::<String>(&mut connection.connection)?;
Ok(())
}
fn has_broken(&self, connection: &mut Self::Connection) -> bool {
connection.failed || !redis::ConnectionLike::is_open(&connection.connection)
}
}
pub(crate) enum ConnectionRef<'a> {
Node(&'a mut dyn redis::ConnectionLike),
Cluster(&'a mut ClusterConnection),
}
impl redis::ConnectionLike for ConnectionRef<'_> {
fn req_packed_command(&mut self, cmd: &[u8]) -> redis::RedisResult<redis::Value> {
match self {
Self::Node(connection) => connection.req_packed_command(cmd),
Self::Cluster(connection) => connection.req_packed_command(cmd),
}
}
fn req_packed_commands(
&mut self,
cmd: &[u8],
offset: usize,
count: usize,
) -> redis::RedisResult<Vec<redis::Value>> {
match self {
Self::Node(connection) => connection.req_packed_commands(cmd, offset, count),
Self::Cluster(connection) => connection.req_packed_commands(cmd, offset, count),
}
}
fn get_db(&self) -> i64 {
match self {
Self::Node(connection) => connection.get_db(),
Self::Cluster(connection) => redis::ConnectionLike::get_db(*connection),
}
}
fn supports_pipelining(&self) -> bool {
match self {
Self::Node(connection) => connection.supports_pipelining(),
Self::Cluster(connection) => redis::ConnectionLike::supports_pipelining(*connection),
}
}
fn check_connection(&mut self) -> bool {
match self {
Self::Node(connection) => connection.check_connection(),
Self::Cluster(connection) => connection.check_connection(),
}
}
fn is_open(&self) -> bool {
match self {
Self::Node(connection) => connection.is_open(),
Self::Cluster(connection) => redis::ConnectionLike::is_open(*connection),
}
}
}
impl ConnectionRef<'_> {
pub(crate) fn pipeline(
&mut self,
commands: Vec<redis::Cmd>,
) -> Result<Vec<redis::Value>, Error> {
match self {
Self::Node(connection) => {
let mut pipeline = redis::pipe();
for command in &commands {
pipeline.add_command(command.clone());
}
pipeline
.query::<Vec<redis::Value>>(*connection)
.map_err(|_| Error::Unavailable)
}
Self::Cluster(connection) => {
let mut replies: Vec<Option<redis::Value>> = vec![None; commands.len()];
for indices in slot_groups(&commands).into_values() {
let mut pipeline = redis::pipe();
for index in &indices {
pipeline.add_command(commands[*index].clone());
}
let values = connection
.req_packed_commands(&pipeline.get_packed_pipeline(), 0, indices.len())
.map_err(|_| Error::Unavailable)?;
if values.len() != indices.len() {
return Err(Error::Unavailable);
}
for (index, value) in indices.into_iter().zip(values) {
replies[index] = Some(value);
}
}
replies
.into_iter()
.collect::<Option<Vec<_>>>()
.ok_or(Error::Unavailable)
}
}
}
pub(crate) fn scan(
&mut self,
pattern: &str,
count: usize,
mut visit: impl FnMut(&mut Self, Vec<String>) -> Result<bool, Error>,
) -> Result<(), Error> {
let pages = match self {
Self::Node(connection) => {
let page = scan_command(0, pattern, count)
.query::<ScanPage>(*connection)
.map_err(|_| Error::Unavailable)?;
vec![(None, page)]
}
Self::Cluster(connection) => connection
.route_command(
&scan_command(0, pattern, count),
RoutingInfo::MultiNode((
MultipleNodeRoutingInfo::AllMasters,
Some(ResponsePolicy::Special),
)),
)
.map_err(|_| Error::Unavailable)
.and_then(primary_pages)?
.into_iter()
.map(|(node, page)| (Some(node), page))
.collect(),
};
for (node, (mut cursor, mut keys)) in pages {
loop {
if !visit(self, keys)? {
return Ok(());
}
if cursor == 0 {
break;
}
(cursor, keys) = self.scan_page(node.as_ref(), cursor, pattern, count)?;
}
}
Ok(())
}
pub(crate) fn ping(&mut self) -> Result<bool, redis::RedisError> {
let command = redis::cmd("PING");
match self {
Self::Node(connection) => command
.query::<String>(*connection)
.map(|response| response == "PONG"),
Self::Cluster(connection) => connection
.route_command(
&command,
RoutingInfo::MultiNode((
MultipleNodeRoutingInfo::AllNodes,
Some(ResponsePolicy::AllSucceeded),
)),
)
.map(|_| true),
}
}
pub(crate) fn node_text(&mut self, command: &redis::Cmd) -> Result<String, Error> {
match self {
Self::Node(connection) => command.query(*connection).map_err(|_| Error::Unavailable),
Self::Cluster(connection) => {
let value = connection
.route_command(
command,
RoutingInfo::MultiNode((
MultipleNodeRoutingInfo::AllNodes,
Some(ResponsePolicy::Special),
)),
)
.map_err(|_| Error::Unavailable)?;
let redis::Value::Map(entries) = value else {
return Err(Error::Unavailable);
};
let mut replies = entries
.into_iter()
.map(|(node, reply)| {
Ok((
redis::from_redis_value::<String>(node)
.map_err(|_| Error::Unavailable)?,
redis::from_redis_value::<String>(reply)
.map_err(|_| Error::Unavailable)?,
))
})
.collect::<Result<Vec<(String, String)>, Error>>()?;
replies.sort();
Ok(replies
.into_iter()
.map(|(_, reply)| reply)
.collect::<Vec<_>>()
.join("\n"))
}
}
}
pub(crate) fn flushall(&mut self) -> Result<(), Error> {
let command = redis::cmd("FLUSHALL");
match self {
Self::Node(connection) => command.query(*connection).map_err(|_| Error::Unavailable),
Self::Cluster(connection) => connection
.route_command(
&command,
RoutingInfo::MultiNode((
MultipleNodeRoutingInfo::AllMasters,
Some(ResponsePolicy::AllSucceeded),
)),
)
.map(|_| ())
.map_err(|_| Error::Unavailable),
}
}
fn scan_page(
&mut self,
node: Option<&NodeAddress>,
cursor: u64,
pattern: &str,
count: usize,
) -> Result<ScanPage, Error> {
let command = scan_command(cursor, pattern, count);
match (self, node) {
(Self::Node(connection), None) => {
command.query(*connection).map_err(|_| Error::Unavailable)
}
(Self::Cluster(connection), Some(node)) => connection
.route_command(
&command,
RoutingInfo::SingleNode(SingleNodeRoutingInfo::ByAddress {
host: node.host().to_string(),
port: node.port(),
}),
)
.map_err(|_| Error::Unavailable)
.and_then(|value| redis::from_redis_value(value).map_err(|_| Error::Unavailable)),
_ => Err(Error::Unavailable),
}
}
}
type ScanPage = (u64, Vec<String>);
fn primary_pages(value: redis::Value) -> Result<Vec<(NodeAddress, ScanPage)>, Error> {
let redis::Value::Map(entries) = value else {
return Err(Error::Unavailable);
};
entries
.into_iter()
.map(|(node, page)| {
let node = redis::from_redis_value::<String>(node).map_err(|_| Error::Unavailable)?;
let node = NodeAddress::try_from(node.as_str()).map_err(|_| Error::Unavailable)?;
let page = redis::from_redis_value::<ScanPage>(page).map_err(|_| Error::Unavailable)?;
Ok((node, page))
})
.collect()
}
fn scan_command(cursor: u64, pattern: &str, count: usize) -> redis::Cmd {
let mut command = redis::cmd("SCAN");
command
.cursor_arg(cursor)
.arg("MATCH")
.arg(pattern)
.arg("COUNT")
.arg(count);
command
}
fn slot_groups(commands: &[redis::Cmd]) -> HashMap<Slot, Vec<usize>> {
let mut groups: HashMap<Slot, Vec<usize>> = HashMap::new();
for (index, command) in commands.iter().enumerate() {
let key = match command.args_iter().nth(1) {
Some(redis::Arg::Simple(key)) => key,
_ => b"",
};
groups.entry(Slot::for_key(key)).or_default().push(index);
}
groups
}

632
litellm-rust/crates/cache-redis/src/cache/operations.rs vendored Normal file
View file

@ -0,0 +1,632 @@
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| connection.ping().map_err(|_| Error::Unavailable))
}
pub async fn ping(&self) -> Result<bool, Error> {
Self::run_blocking(Arc::clone(&self.connections), |connection| {
connection.ping().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 matches = Vec::new();
connection.scan(&pattern, count, |_, keys| {
matches.extend(keys);
Ok(matches.len() < count)
})?;
matches.truncate(count);
Ok(matches)
})
.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 sadd = redis::cmd("SADD");
sadd.arg(&key).arg(values);
let mut expire = redis::cmd("EXPIRE");
expire.arg(&key).arg(ttl);
let replies = connection.pipeline(vec![sadd, expire])?;
replies
.into_iter()
.next()
.map(redis::from_redis_value::<usize>)
.transpose()
.map_err(|_| Error::Unavailable)?
.ok_or(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 commands = operations
.into_iter()
.map(|(key, values)| {
let mut command = redis::cmd("RPUSH");
command.arg(key).arg(values);
command
})
.collect();
connection
.pipeline(commands)?
.into_iter()
.map(|value| redis::from_redis_value(value).map_err(|_| Error::Unavailable))
.collect()
})
.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 commands = operations
.into_iter()
.map(|(key, count)| {
let mut command = redis::cmd("LPOP");
command.arg(key);
if let Some(count) = count {
command.arg(count);
}
command
})
.collect();
connection.pipeline(commands)
})
.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| connection.node_text(redis::cmd("CLIENT").arg("LIST")))
}
pub fn info(&self) -> Result<String, Error> {
self.connections
.execute(|connection| connection.node_text(&redis::cmd("INFO")))
}
pub fn flushall(&self) -> Result<(), Error> {
self.connections.execute(|connection| connection.flushall())
}
}
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 commands = Vec::with_capacity(operations.len() * 2);
let mut increments = Vec::with_capacity(operations.len());
for (key, amount, ttl) in operations {
let mut increment = redis::cmd("INCRBYFLOAT");
increment.arg(&key).arg(amount);
increments.push(commands.len());
commands.push(increment);
if let Some(ttl) = ttl {
let mut expire = redis::cmd("EXPIRE");
expire.arg(key).arg(ttl);
commands.push(expire);
}
}
let mut replies = connection.pipeline(commands)?;
increments
.into_iter()
.map(|index| {
redis::from_redis_value(std::mem::take(&mut replies[index]))
.map_err(|_| Error::Unavailable)
})
.collect()
})
.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,
}
}
}

6
litellm-rust/crates/cache-redis/src/lib.rs
View file

@ -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
View 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>,
},
}

701
litellm-rust/crates/cache-redis/tests/cache.rs
View file

@ -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
);
}

492
litellm-rust/crates/cache-redis/tests/cluster.rs Normal file
View file

@ -0,0 +1,492 @@
//! Contract tests against a real Redis Cluster. Set `LITELLM_TEST_REDIS_CLUSTER_NODES` to a
//! comma separated `host:port` list (for example `127.0.0.1:7000,127.0.0.1:7001`) to run them.
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use litellm_cache::{
BaseCache, BatchCache, BatchEntry, CacheConnectionStatus, CacheScript, ClaimCache,
CounterCache, DeleteCache, Error, ExactCacheContext, FlushCache, IncrementOperation, JsonCodec,
ScriptCache,
};
use litellm_cache_redis::{
RedisArg, RedisCache, RedisLpopOperation, RedisLpopResult, RedisNode, RedisRpushOperation,
RedisTopology,
};
use redis::cluster_routing::Slot;
type Cache = RedisCache<JsonCodec<serde_json::Value>>;
fn topology() -> Option<RedisTopology> {
let nodes = std::env::var("LITELLM_TEST_REDIS_CLUSTER_NODES").ok()?;
let startup_nodes = nodes
.split(',')
.map(|node| {
let (host, port) = node.trim().rsplit_once(':').expect("host:port");
RedisNode {
host: host.to_string(),
port: port.parse().expect("port"),
}
})
.collect();
Some(RedisTopology::Cluster { startup_nodes })
}
fn namespace(label: &str) -> String {
let nanos = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_nanos();
format!("cluster-test:{label}:{nanos}")
}
fn cluster_url() -> String {
std::env::var("LITELLM_TEST_REDIS_CLUSTER_URL")
.unwrap_or_else(|_| "redis://127.0.0.1:7000".into())
}
fn cluster_cache(label: &str) -> Option<Cache> {
let topology = topology()?;
Some(
Cache::connect(
&cluster_url(),
&topology,
Some(Duration::from_secs(120)),
JsonCodec::new(),
)
.expect("cluster connection")
.with_namespace(Some(namespace(label))),
)
}
fn counter_cache(label: &str) -> Option<RedisCache<JsonCodec<f64>>> {
let topology = topology()?;
Some(
RedisCache::connect(
&cluster_url(),
&topology,
Some(Duration::from_secs(60)),
JsonCodec::new(),
)
.expect("cluster connection")
.with_namespace(Some(namespace(label))),
)
}
fn multi_slot_keys(count: usize) -> Vec<String> {
let keys: Vec<String> = (0..count).map(|index| format!("key-{index}")).collect();
let slots: std::collections::HashSet<Slot> = keys.iter().map(Slot::for_key).collect();
assert!(slots.len() > 1, "keys must span multiple slots");
keys
}
macro_rules! cluster_or_skip {
($label:expr) => {
match cluster_cache($label) {
Some(cache) => cache,
None => return,
}
};
}
#[test]
fn constructor_rejects_clusters_without_startup_nodes() {
let error = Cache::connect(
"redis://127.0.0.1:7000",
&RedisTopology::Cluster {
startup_nodes: Vec::new(),
},
None,
JsonCodec::new(),
)
.err();
assert!(matches!(error, Some(Error::Unavailable)));
}
#[test]
fn constructor_rejects_unix_socket_urls_for_clusters() {
let error = Cache::connect(
"redis+unix:///tmp/redis.sock",
&RedisTopology::Cluster {
startup_nodes: vec![RedisNode {
host: "127.0.0.1".into(),
port: 7000,
}],
},
None,
JsonCodec::new(),
)
.err();
assert!(matches!(error, Some(Error::Unavailable)));
}
#[test]
fn single_key_operations_round_trip_with_ttl_rounding() {
let cache = cluster_or_skip!("single");
let context = ExactCacheContext {
ttl: Some(Duration::from_millis(1500)),
};
let keys = multi_slot_keys(12);
for (index, key) in keys.iter().enumerate() {
cache
.set_cache(key, serde_json::json!({ "index": index }), &context)
.unwrap();
}
for (index, key) in keys.iter().enumerate() {
assert_eq!(
cache.get_cache(key, &context).unwrap(),
Some(serde_json::json!({ "index": index }))
);
}
let runtime = tokio::runtime::Runtime::new().unwrap();
let ttl = runtime.block_on(cache.async_get_ttl(&keys[0])).unwrap();
assert_eq!(ttl, Some(2));
cache.delete_cache(&keys[0]).unwrap();
assert_eq!(cache.get_cache(&keys[0], &context).unwrap(), None);
assert!(cache.sync_ping().unwrap());
}
#[tokio::test]
async fn batch_reads_span_slots_and_preserve_order_with_malformed_entries() {
let cache = cluster_or_skip!("batch");
let context = ExactCacheContext::default();
let keys = multi_slot_keys(40);
for (index, key) in keys.iter().enumerate() {
if index % 5 == 0 {
continue;
}
cache
.async_set_cache(key, serde_json::json!(index), context.clone())
.await
.unwrap();
}
let mut raw = redis::cluster::ClusterClient::new(vec![cluster_url()])
.unwrap()
.get_connection()
.unwrap();
let malformed = format!("{}:{}", cache.namespace().unwrap(), keys[1]);
redis::cmd("SET")
.arg(&malformed)
.arg("not json")
.exec(&mut raw)
.unwrap();
let entries = cache
.async_batch_get_cache(keys.clone(), context.clone())
.await
.unwrap();
assert_eq!(entries.len(), keys.len());
for (index, entry) in entries.iter().enumerate() {
let expected = if index == 1 {
BatchEntry::Invalid
} else if index % 5 == 0 {
BatchEntry::Miss
} else {
BatchEntry::Hit(serde_json::json!(index))
};
assert_eq!(*entry, expected, "entry {index}");
}
let sync_entries = cache.batch_get_cache(&keys, &context).unwrap();
assert_eq!(sync_entries, entries);
cache.delete_cache_keys(keys.clone()).await.unwrap();
let entries = cache.async_batch_get_cache(keys, context).await.unwrap();
assert!(entries.iter().all(|entry| *entry == BatchEntry::Miss));
}
#[tokio::test]
async fn pipelines_group_by_slot_and_return_results_in_submission_order() {
let cache = cluster_or_skip!("pipeline");
let keys = multi_slot_keys(30);
let entries = keys
.iter()
.enumerate()
.map(|(index, key)| (key.clone(), serde_json::json!(index)))
.collect();
cache
.async_set_cache_pipeline(entries, ExactCacheContext::default())
.await
.unwrap();
let hits = cache
.async_batch_get_cache(keys.clone(), ExactCacheContext::default())
.await
.unwrap();
assert!(
hits.iter()
.enumerate()
.all(|(index, entry)| *entry == BatchEntry::Hit(serde_json::json!(index)))
);
let queues: Vec<String> = keys.iter().map(|key| format!("queue:{key}")).collect();
let pushed = cache
.async_rpush_pipeline(
queues
.iter()
.enumerate()
.map(|(index, key)| RedisRpushOperation {
key: key.clone(),
values: (0..=index)
.map(|value| RedisArg::Integer(value as i64))
.collect(),
})
.collect(),
)
.await
.unwrap();
assert_eq!(pushed, (1..=keys.len()).collect::<Vec<_>>());
let popped = cache
.async_lpop_pipeline(
queues
.iter()
.enumerate()
.map(|(index, key)| RedisLpopOperation {
key: key.clone(),
count: (index % 2 == 0).then_some(2),
})
.collect(),
)
.await
.unwrap();
for (index, result) in popped.into_iter().enumerate() {
match result {
RedisLpopResult::Value(value) => {
assert_eq!(index % 2, 1, "queue {index}");
assert_eq!(value, b"0");
}
RedisLpopResult::Values(values) => {
assert_eq!(index % 2, 0, "queue {index}");
let expected: Vec<Vec<u8>> = (0..=index)
.take(2)
.map(|value| value.to_string().into_bytes())
.collect();
assert_eq!(values, expected);
}
other => panic!("queue {index}: {other:?}"),
}
}
let counters: Vec<String> = keys.iter().map(|key| format!("counter:{key}")).collect();
let Some(counter) = counter_cache("counter") else {
return;
};
let totals = counter
.async_increment_pipeline(
counters
.iter()
.enumerate()
.map(|(index, key)| IncrementOperation {
key: key.clone(),
amount: index as f64 + 0.5,
ttl: (index % 3 == 0).then_some(Duration::from_secs(30)),
})
.collect(),
)
.await
.unwrap();
let expected: Vec<f64> = (0..keys.len()).map(|index| index as f64 + 0.5).collect();
assert_eq!(totals, expected);
assert_eq!(counter.async_get_ttl(&counters[0]).await.unwrap(), Some(30));
assert_eq!(counter.async_get_ttl(&counters[1]).await.unwrap(), None);
counter.async_flush_cache().await.unwrap();
cache.async_flush_cache().await.unwrap();
}
#[tokio::test]
async fn scan_and_scoped_flush_cover_every_primary() {
let cache = cluster_or_skip!("flush");
let other = cluster_or_skip!("other");
let context = ExactCacheContext::default();
let keys = multi_slot_keys(60);
for key in &keys {
cache
.async_set_cache(key, serde_json::json!(true), context.clone())
.await
.unwrap();
other
.async_set_cache(key, serde_json::json!(true), context.clone())
.await
.unwrap();
}
let mut scanned = cache.async_scan_iter("key-", 1000).await.unwrap();
scanned.sort();
let mut expected: Vec<String> = keys
.iter()
.map(|key| format!("{}:{key}", cache.namespace().unwrap()))
.collect();
expected.sort();
assert_eq!(scanned, expected);
assert_eq!(cache.async_scan_iter("key-", 7).await.unwrap().len(), 7);
cache.flush_cache().unwrap();
let flushed = cache
.async_batch_get_cache(keys.clone(), context.clone())
.await
.unwrap();
assert!(flushed.iter().all(|entry| *entry == BatchEntry::Miss));
let kept = other.async_batch_get_cache(keys, context).await.unwrap();
assert!(
kept.iter()
.all(|entry| *entry == BatchEntry::Hit(serde_json::json!(true)))
);
other.async_flush_cache().await.unwrap();
}
fn ping_calls_per_node(startup: &redis::Client) -> Vec<(String, u64)> {
let mut connection = startup.get_connection().unwrap();
let nodes: String = redis::cmd("CLUSTER")
.arg("NODES")
.query(&mut connection)
.unwrap();
let mut counts: Vec<(String, u64)> = nodes
.lines()
.map(|line| {
let address = line.split_whitespace().nth(1).unwrap();
let address = address.split('@').next().unwrap();
let mut node = redis::Client::open(format!("redis://{address}"))
.unwrap()
.get_connection()
.unwrap();
let stats: String = redis::cmd("INFO")
.arg("commandstats")
.query(&mut node)
.unwrap();
let calls = stats
.lines()
.find_map(|stat| stat.strip_prefix("cmdstat_ping:calls="))
.and_then(|rest| rest.split(',').next())
.map_or(0, |calls| calls.parse().unwrap());
(address.to_string(), calls)
})
.collect();
counts.sort();
counts
}
#[tokio::test]
async fn ping_reaches_every_node() {
let cache = cluster_or_skip!("ping");
let startup = redis::Client::open(cluster_url()).unwrap();
let before = ping_calls_per_node(&startup);
assert!(before.len() >= 2, "{before:?}");
assert!(cache.ping().await.unwrap());
let after = ping_calls_per_node(&startup);
for ((node, calls_before), (_, calls_after)) in before.iter().zip(&after) {
assert!(calls_after > calls_before, "{node} was not pinged");
}
assert!(cache.sync_ping().unwrap());
let result = cache.test_connection().await.unwrap();
assert_eq!(result.status, CacheConnectionStatus::Success);
}
#[tokio::test]
async fn counters_claims_scripts_and_sets_work_on_the_cluster() {
let Some(counter) = counter_cache("counter") else {
return;
};
let context = ExactCacheContext::default();
assert_eq!(
counter
.increment_cache("spend", 1.5, context.clone())
.unwrap(),
1.5
);
assert_eq!(
counter
.async_increment("spend", 2.0, context.clone())
.await
.unwrap(),
3.5
);
assert_eq!(
counter
.increment_with_floor("budget", -3, Duration::from_secs(30))
.unwrap(),
0
);
assert_eq!(
counter
.async_increment_with_floor("budget", 7, Duration::from_secs(30))
.await
.unwrap(),
7
);
assert_eq!(counter.async_set_max("peak", 4.0, None).await.unwrap(), 4.0);
assert_eq!(counter.async_set_max("peak", 2.0, None).await.unwrap(), 4.0);
counter.flush_cache().unwrap();
let cache = cluster_or_skip!("claim");
let owner = serde_json::json!("owner-a");
let rival = serde_json::json!("owner-b");
assert_eq!(
cache
.claim_cache("lock", owner.clone(), &[], context.clone())
.unwrap(),
owner
);
assert_eq!(
cache
.async_claim_cache("lock", rival.clone(), vec![owner.clone()], context.clone())
.await
.unwrap(),
owner
);
assert_eq!(
cache
.claim_cache("lock", rival.clone(), &[], context.clone())
.unwrap(),
owner
);
assert_eq!(
cache
.async_claim_cache("lock", rival.clone(), vec![rival.clone()], context.clone())
.await
.unwrap(),
rival
);
let script = cache
.async_register_script("return redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2])".into());
let reply = script
.invoke(
vec!["scripted".into()],
vec![RedisArg::Bytes(b"payload".to_vec()), RedisArg::Integer(5)],
)
.await
.unwrap();
assert_eq!(reply, redis::Value::Okay);
assert_eq!(cache.async_get_ttl("scripted").await.unwrap(), Some(5));
let evaluated: redis::Value = cache
.async_eval(
"return redis.call('GET', KEYS[1])".into(),
vec!["scripted".into()],
Vec::new(),
)
.await
.unwrap();
assert_eq!(evaluated, redis::Value::BulkString(b"payload".to_vec()));
assert_eq!(
cache
.async_set_cache_sadd(
"members",
vec![
RedisArg::Bytes(b"a".to_vec()),
RedisArg::Bytes(b"b".to_vec())
],
Some(Duration::from_secs(9)),
)
.await
.unwrap(),
2
);
assert_eq!(cache.async_get_ttl("members").await.unwrap(), Some(9));
let result = cache.test_connection().await.unwrap();
assert_eq!(result.status, CacheConnectionStatus::Success);
assert!(cache.ping().await.unwrap());
let info = cache.info().unwrap();
assert!(info.matches("redis_version").count() > 1, "{info}");
assert!(cache.client_list().unwrap().contains("id="));
cache.async_flush_cache().await.unwrap();
assert_eq!(cache.async_get_ttl("members").await.unwrap(), None);
assert_eq!(cache.get_cache("lock", &context).unwrap(), None);
}

20
litellm-rust/crates/cache-response/Cargo.toml Normal file
View 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
View 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
View 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
View 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
View 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
View 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
View 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
View 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
View 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
View 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
View file

@ -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
View file

@ -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
View 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
View file

@ -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
View 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
View 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
View 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
View file

@ -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
View file

@ -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
View file

@ -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
View 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
View 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())));
}

112
litellm-rust/crates/core-utils/src/serde_compat.rs
View file

@ -1,33 +1,91 @@
use serde::{Deserialize, Deserializer, de::Error};
use serde_json::Value;
use serde::{
Deserializer,
de::{Error, Visitor},
};
use serde_with::DeserializeAs;
pub struct LaxI64;
pub struct FiniteF64;
pub fn parse_str_bool(value: &str) -> Option<bool> {
let token = value.trim_matches(|character: char| {
character.is_whitespace() || matches!(character, '\u{1c}'..='\u{1f}')
});
if token.eq_ignore_ascii_case("true") {
return Some(true);
}
token.eq_ignore_ascii_case("false").then_some(false)
}
impl<'de> DeserializeAs<'de, i64> for LaxI64 {
fn deserialize_as<D: Deserializer<'de>>(deserializer: D) -> Result<i64, D::Error> {
match Value::deserialize(deserializer)? {
Value::Number(number) if number.is_f64() => number.as_f64().and_then(integral_float),
Value::Number(number) => number.as_i64(),
Value::String(value) => integer_string(value.trim()),
Value::Bool(value) => Some(i64::from(value)),
_ => None,
}
.ok_or_else(|| D::Error::custom("expected an integer in the i64 range"))
deserializer.deserialize_any(Self)
}
}
impl<'de> Visitor<'de> for LaxI64 {
type Value = i64;
fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter.write_str("an integer in the i64 range")
}
fn visit_i64<E: Error>(self, value: i64) -> Result<i64, E> {
Ok(value)
}
fn visit_u64<E: Error>(self, value: u64) -> Result<i64, E> {
i64::try_from(value).map_err(E::custom)
}
fn visit_f64<E: Error>(self, value: f64) -> Result<i64, E> {
integral_float(value).ok_or_else(|| E::custom("expected an integer in the i64 range"))
}
fn visit_str<E: Error>(self, value: &str) -> Result<i64, E> {
integer_string(value.trim())
.ok_or_else(|| E::custom("expected an integer in the i64 range"))
}
fn visit_bool<E: Error>(self, value: bool) -> Result<i64, E> {
Ok(i64::from(value))
}
}
impl<'de> DeserializeAs<'de, f64> for FiniteF64 {
fn deserialize_as<D: Deserializer<'de>>(deserializer: D) -> Result<f64, D::Error> {
match Value::deserialize(deserializer)? {
Value::Number(number) => number.as_f64(),
Value::String(value) => value.trim().parse::<f64>().ok(),
Value::Bool(value) => Some(f64::from(value)),
_ => None,
}
.filter(|value| value.is_finite())
.ok_or_else(|| D::Error::custom("expected a finite number"))
deserializer.deserialize_any(Self)
}
}
impl<'de> Visitor<'de> for FiniteF64 {
type Value = f64;
fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter.write_str("a finite number")
}
fn visit_i64<E: Error>(self, value: i64) -> Result<f64, E> {
Ok(value as f64)
}
fn visit_u64<E: Error>(self, value: u64) -> Result<f64, E> {
Ok(value as f64)
}
fn visit_f64<E: Error>(self, value: f64) -> Result<f64, E> {
value
.is_finite()
.then_some(value)
.ok_or_else(|| E::custom("expected a finite number"))
}
fn visit_str<E: Error>(self, value: &str) -> Result<f64, E> {
self.visit_f64(value.trim().parse::<f64>().map_err(E::custom)?)
}
fn visit_bool<E: Error>(self, value: bool) -> Result<f64, E> {
Ok(f64::from(value))
}
}
@ -66,7 +124,7 @@ fn integral_float(value: f64) -> Option<i64> {
#[cfg(test)]
mod tests {
use serde::Serialize;
use serde::{Deserialize, Serialize};
use serde_json::json;
use serde_with::serde_as;
@ -81,6 +139,22 @@ mod tests {
float: Option<f64>,
}
#[test]
fn boolean_tokens_follow_python_string_trimming_without_redis_tokens() {
for (input, expected) in [
(" True ", Some(true)),
("\u{1c}TRUE\u{1f}", Some(true)),
("\u{a0}False\u{2003}", Some(false)),
("true\u{200b}", None),
("yes", None),
("1", None),
("", None),
("unknown", None),
] {
assert_eq!(parse_str_bool(input), expected, "{input:?}");
}
}
#[test]
fn adapters_compose_and_serialize_as_numbers() {
let numbers: Numbers = serde_json::from_value(json!({

4
litellm-rust/crates/core-utils/src/settings.rs
View file

@ -1,5 +1,7 @@
use std::str::FromStr;
use crate::serde_compat::parse_str_bool;
pub trait Lookup {
fn get(&self, name: &str) -> Option<String>;
@ -9,7 +11,7 @@ pub trait Lookup {
fn enabled(&self, name: &str) -> Option<bool> {
self.get(name)
.is_some_and(|value| value.trim().eq_ignore_ascii_case("true"))
.is_some_and(|value| parse_str_bool(&value) == Some(true))
.then_some(true)
}

15
litellm-rust/crates/host-python/src/marshal.rs
View file

@ -45,7 +45,7 @@ where
fn into_pyobject(self, py: Python<'py>) -> PyResult<Self::Output> {
catch_unwind(AssertUnwindSafe(|| pythonize::pythonize(py, &self.0)))
.map_err(panic_to_pyerr)?
.map_err(|error| PyValueError::new_err(error.to_string()))
.map_err(PyErr::from)
}
}
@ -87,6 +87,19 @@ mod tests {
});
}
#[test]
fn pythonized_preserves_python_serialization_error_types() {
crate::initialize_python();
Python::attach(|py| {
let value = std::collections::BTreeMap::from([(vec![1], "value")]);
let direct = to_py(py, &value).unwrap_err();
let wrapped = Pythonized(value).into_pyobject(py).unwrap_err();
assert!(direct.is_instance_of::<pyo3::exceptions::PyTypeError>(py));
assert!(wrapped.is_instance_of::<pyo3::exceptions::PyTypeError>(py));
assert_eq!(wrapped.to_string(), direct.to_string());
});
}
#[test]
fn pythonized_maps_serializer_panics_to_a_base_exception() {
crate::initialize_python();

13
litellm-rust/crates/http/src/config.rs
View file

@ -129,6 +129,7 @@ mod tests {
use rstest::rstest;
use super::*;
use crate::TlsSource;
fn settings(ssl_verify: Option<SslVerify>, ssl_cert_file: Option<&str>) -> HttpSettings {
HttpSettings {
@ -298,7 +299,11 @@ mod tests {
};
assert!(matches!(
reqwest::ClientBuilder::try_from(&config),
Err(Error::Read { path: reported, .. }) if reported == path
Err(Error::Read {
path: reported,
tls_source: TlsSource::CaBundle,
..
}) if reported == path
));
}
@ -315,7 +320,11 @@ mod tests {
std::fs::remove_file(&path).unwrap();
assert!(matches!(
result,
Err(Error::InvalidPem { path: reported, .. }) if reported == path
Err(Error::InvalidPem {
path: reported,
tls_source: TlsSource::CaBundle,
..
}) if reported == path
));
}
}

18
litellm-rust/crates/http/src/error.rs
View file

@ -1,11 +1,25 @@
use std::path::PathBuf;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum TlsSource {
CaBundle,
ClientIdentity,
}
#[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)]
pub enum Error {
#[error("could not read {}: {message}", path.display())]
Read { path: PathBuf, message: String },
Read {
path: PathBuf,
message: String,
tls_source: TlsSource,
},
#[error("{} is not a PEM file: {message}", path.display())]
InvalidPem { path: PathBuf, message: String },
InvalidPem {
path: PathBuf,
message: String,
tls_source: TlsSource,
},
#[error("could not build the HTTP client: {0}")]
Client(String),
#[error("request body could not be serialized: {0}")]

2
litellm-rust/crates/http/src/lib.rs
View file

@ -10,7 +10,7 @@ mod tls;
pub mod transport;
pub use config::{HttpClientConfig, Resolution, Verify};
pub use error::Error;
pub use error::{Error, TlsSource};
pub use pool::{ClientVariant, HttpClientPool};
pub use proxy::EnvironmentProxies;
pub use settings::{HttpSettings, HttpSettingsLayer, SslVerify, TcpKeepalive};

48
litellm-rust/crates/http/src/media.rs
View file

@ -54,16 +54,39 @@ impl Default for UrlPolicy {
impl UrlPolicy {
fn allows(&self, host: &str, port: u16) -> bool {
let host = normalize_host(host);
let with_port = format!("{host}:{port}");
self.allowed_hosts
.iter()
.map(|entry| normalize_host(entry))
.any(|entry| entry == host || entry == with_port)
.filter_map(|entry| parse_allowed_host(entry))
.any(|(entry_host, entry_port)| {
entry_host == host && entry_port.is_none_or(|entry_port| entry_port == port)
})
}
}
fn normalize_host(host: &str) -> String {
host.to_ascii_lowercase().trim_end_matches('.').to_owned()
pub fn normalize_host(host: &str) -> String {
let host = host.trim().trim_end_matches('.');
let host = host
.strip_prefix('[')
.and_then(|host| host.strip_suffix(']'))
.unwrap_or(host);
host.to_ascii_lowercase()
}
fn parse_allowed_host(entry: &str) -> Option<(String, Option<u16>)> {
let entry = entry.trim();
if let Some(entry) = entry.strip_prefix('[') {
let (host, suffix) = entry.split_once(']')?;
let port = match suffix {
"" => None,
suffix => Some(suffix.strip_prefix(':')?.parse().ok()?),
};
return Some((normalize_host(host), port));
}
let (host, port) = match entry.rsplit_once(':') {
Some((host, port)) if !host.contains(':') => (host, Some(port.parse().ok()?)),
_ => (entry, None),
};
Some((normalize_host(host), port))
}
type ProxyMatch = Arc<dyn Fn(&Url) -> bool + Send + Sync>;
@ -670,6 +693,21 @@ mod tests {
assert!(matches!(result, Err(Error::BlockedUrl)));
}
#[test]
fn allowlist_matches_bracketed_ipv6_hosts_and_ports() {
let policy = UrlPolicy {
validate: true,
allowed_hosts: vec!["[2001:db8::1]".into(), "[2001:db8::1]:8443".into()],
};
assert!(policy.allows("2001:db8::1", 443));
assert!(policy.allows("2001:db8::1", 8443));
let port_specific = UrlPolicy {
validate: true,
allowed_hosts: vec!["[2001:db8::1]:8443".into()],
};
assert!(!port_specific.allows("2001:db8::1", 9443));
}
#[tokio::test]
async fn validation_off_fetches_private_hosts_and_follows_redirects() {
let (url, server, _) = serve_named(

19
litellm-rust/crates/http/src/settings.rs
View file

@ -3,7 +3,10 @@ use std::{
time::Duration,
};
use litellm_core_utils::settings::{Layer, Lookup, merge};
use litellm_core_utils::{
serde_compat::parse_str_bool,
settings::{Layer, Lookup, merge},
};
use crate::proxy::EnvironmentProxies;
@ -16,9 +19,9 @@ pub enum SslVerify {
impl SslVerify {
pub fn parse(value: &str) -> Self {
match value.trim().to_ascii_lowercase().as_str() {
"true" => Self::Enabled,
"false" => Self::Disabled,
match parse_str_bool(value) {
Some(true) => Self::Enabled,
Some(false) => Self::Disabled,
_ => Self::CaBundle(PathBuf::from(value)),
}
}
@ -152,9 +155,7 @@ impl HttpSettings {
Self {
ssl_verify: merged.ssl_verify,
ssl_cert_file: merged.ssl_cert_file,
ssl_certificate: merged
.ssl_certificate
.filter(|path| !path.as_os_str().is_empty()),
ssl_certificate: merged.ssl_certificate,
ssl_security_level: merged.ssl_security_level.filter(|level| !level.is_empty()),
ssl_ecdh_curve: merged.ssl_ecdh_curve.filter(|curve| !curve.is_empty()),
force_ipv4: merged.force_ipv4.unwrap_or(defaults.force_ipv4),
@ -287,7 +288,7 @@ mod tests {
}
#[test]
fn empty_environment_values_clear_the_setting_like_python_truthiness() {
fn empty_certificate_is_retained_for_validation_while_empty_tuning_is_absent() {
let configured = HttpSettingsLayer {
ssl_certificate: Some("/configured/client.pem".into()),
ssl_security_level: Some("configured".into()),
@ -300,7 +301,7 @@ mod tests {
("SSL_ECDH_CURVE", ""),
]));
let settings = HttpSettings::from_layers([environment, configured]);
assert_eq!(settings.ssl_certificate, None);
assert_eq!(settings.ssl_certificate, Some(PathBuf::new()));
assert_eq!(settings.ssl_security_level, None);
assert_eq!(settings.ssl_ecdh_curve, None);
}

49
litellm-rust/crates/http/src/tls.rs
View file

@ -9,7 +9,7 @@ use rustls::{
use crate::{
config::{HttpClientConfig, Verify},
error::Error,
error::{Error, TlsSource},
};
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
@ -197,15 +197,17 @@ impl TryFrom<&HttpClientConfig> for ClientConfig {
Verify::BuiltInRoots => builder.with_root_certificates(RootCertStore {
roots: webpki_roots::TLS_SERVER_ROOTS.to_vec(),
}),
Verify::CaBundle(path) => builder.with_root_certificates(bundle_roots(path)?),
Verify::CaBundle(path) => {
builder.with_root_certificates(bundle_roots(path, TlsSource::CaBundle)?)
}
};
let mut tls = match &config.client_certificate {
None => verified.with_no_client_auth(),
Some(path) => {
let (chain, key) = identity(path)?;
let (chain, key) = identity(path, TlsSource::ClientIdentity)?;
verified
.with_client_auth_cert(chain, key)
.map_err(|error| invalid_pem(path, error))?
.map_err(|error| invalid_pem(path, TlsSource::ClientIdentity, error))?
}
};
tls.alpn_protocols = if config.http2 {
@ -217,47 +219,52 @@ impl TryFrom<&HttpClientConfig> for ClientConfig {
}
}
fn bundle_roots(path: &Path) -> Result<RootCertStore, Error> {
let certificates = certificates(path)?;
fn bundle_roots(path: &Path, source: TlsSource) -> Result<RootCertStore, Error> {
let certificates = certificates(path, source)?;
if certificates.is_empty() {
return Err(invalid_pem(path, "no certificates found"));
return Err(invalid_pem(path, source, "no certificates found"));
}
let mut store = RootCertStore::empty();
for certificate in certificates {
store
.add(certificate)
.map_err(|error| invalid_pem(path, error))?;
.map_err(|error| invalid_pem(path, source, error))?;
}
Ok(store)
}
fn identity(path: &Path) -> Result<(Vec<CertificateDer<'static>>, PrivateKeyDer<'static>), Error> {
let chain = certificates(path)?;
fn identity(
path: &Path,
source: TlsSource,
) -> Result<(Vec<CertificateDer<'static>>, PrivateKeyDer<'static>), Error> {
let chain = certificates(path, source)?;
if chain.is_empty() {
return Err(invalid_pem(path, "no certificates found"));
return Err(invalid_pem(path, source, "no certificates found"));
}
let key =
PrivateKeyDer::from_pem_slice(&read(path)?).map_err(|error| invalid_pem(path, error))?;
let key = PrivateKeyDer::from_pem_slice(&read(path, source)?)
.map_err(|error| invalid_pem(path, source, error))?;
Ok((chain, key))
}
fn certificates(path: &Path) -> Result<Vec<CertificateDer<'static>>, Error> {
CertificateDer::pem_slice_iter(&read(path)?)
fn certificates(path: &Path, source: TlsSource) -> Result<Vec<CertificateDer<'static>>, Error> {
CertificateDer::pem_slice_iter(&read(path, source)?)
.collect::<Result<_, _>>()
.map_err(|error| invalid_pem(path, error))
.map_err(|error| invalid_pem(path, source, error))
}
fn read(path: &Path) -> Result<Vec<u8>, Error> {
fn read(path: &Path, source: TlsSource) -> Result<Vec<u8>, Error> {
std::fs::read(path).map_err(|error| Error::Read {
path: path.to_path_buf(),
message: error.to_string(),
tls_source: source,
})
}
fn invalid_pem(path: &Path, message: impl fmt::Display) -> Error {
fn invalid_pem(path: &Path, source: TlsSource, message: impl fmt::Display) -> Error {
Error::InvalidPem {
path: path.to_path_buf(),
message: message.to_string(),
tls_source: source,
}
}
@ -405,7 +412,11 @@ mod tests {
std::fs::remove_file(&path).unwrap();
assert!(matches!(
result,
Err(Error::InvalidPem { path: reported, .. }) if reported == path
Err(Error::InvalidPem {
path: reported,
tls_source: TlsSource::ClientIdentity,
..
}) if reported == path
));
}
}

8
litellm-rust/crates/python-bridge/Cargo.toml
View file

@ -20,6 +20,12 @@ tiktoken = ["litellm-token-counter/tiktoken"]
[dependencies]
bytes.workspace = true
litellm-cache.workspace = true
litellm-cache-azure-blob.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
@ -36,6 +42,8 @@ serde_json.workspace = true
tokio = { workspace = true, features = ["sync"] }
[dev-dependencies]
serde.workspace = true
serde_with.workspace = true
criterion.workspace = true
futures-util.workspace = true
rstest.workspace = true

176
litellm-rust/crates/python-bridge/python_settings.json
View file

@ -1,26 +1,154 @@
{
"http_settings": [
"ssl_verify",
"ssl_certificate",
"ssl_security_level",
"ssl_ecdh_curve",
"force_ipv4",
"http2",
"aiohttp_trust_env",
"disable_aiohttp_trust_env",
"disable_aiohttp_transport",
"user_agent"
],
"url_policy": [
"user_url_validation",
"user_url_allowed_hosts"
],
"provider_defaults": [
"vertex_project",
"vertex_location",
"enable_azure_ad_token_refresh"
],
"secret_manager": [
"readable"
]
"http_settings": {
"version": 1,
"fields": {
"ssl_verify": {
"adapter": "SslVerifyInput",
"required": true,
"precedence": "module_global",
"sensitive": false,
"shapes": [
"none",
"bool",
"str"
],
"unsupported_live": "configuration_error"
},
"ssl_certificate": {
"adapter": "OptionalStrictString",
"required": true,
"precedence": "module_global",
"sensitive": false,
"shapes": [],
"unsupported_live": null
},
"ssl_security_level": {
"adapter": "TuningString",
"required": true,
"precedence": "module_global",
"sensitive": false,
"shapes": [],
"unsupported_live": null
},
"ssl_ecdh_curve": {
"adapter": "TuningString",
"required": true,
"precedence": "module_global",
"sensitive": false,
"shapes": [],
"unsupported_live": null
},
"force_ipv4": {
"adapter": "Truthy",
"required": true,
"precedence": "module_global",
"sensitive": false,
"shapes": [],
"unsupported_live": null
},
"http2": {
"adapter": "ExactTrue",
"required": true,
"precedence": "module_global",
"sensitive": false,
"shapes": [],
"unsupported_live": null
},
"aiohttp_trust_env": {
"adapter": "Truthy",
"required": true,
"precedence": "module_global",
"sensitive": false,
"shapes": [],
"unsupported_live": null
},
"disable_aiohttp_trust_env": {
"adapter": "Truthy",
"required": true,
"precedence": "module_global",
"sensitive": false,
"shapes": [],
"unsupported_live": null
},
"disable_aiohttp_transport": {
"adapter": "ExactTrue",
"required": true,
"precedence": "module_global",
"sensitive": false,
"shapes": [],
"unsupported_live": null
},
"user_agent": {
"adapter": "StrictString",
"required": true,
"precedence": "accessor",
"sensitive": false,
"shapes": [],
"unsupported_live": null
}
}
},
"url_policy": {
"version": 1,
"fields": {
"user_url_validation": {
"adapter": "Truthy",
"required": true,
"precedence": "module_global",
"sensitive": false,
"shapes": [],
"unsupported_live": null
},
"user_url_allowed_hosts": {
"adapter": "HostCollection",
"required": true,
"precedence": "module_global",
"sensitive": false,
"shapes": [],
"unsupported_live": null
}
}
},
"provider_defaults": {
"version": 1,
"fields": {
"vertex_project": {
"adapter": "FalsyOptionalString",
"required": true,
"precedence": "module_global",
"sensitive": true,
"shapes": [],
"unsupported_live": null
},
"vertex_location": {
"adapter": "FalsyOptionalString",
"required": true,
"precedence": "module_global",
"sensitive": true,
"shapes": [],
"unsupported_live": null
},
"enable_azure_ad_token_refresh": {
"adapter": "ExactTrue",
"required": true,
"precedence": "module_global",
"sensitive": false,
"shapes": [],
"unsupported_live": null
}
}
},
"secret_manager": {
"version": 1,
"fields": {
"readable": {
"adapter": "StrictBool",
"required": true,
"precedence": "accessor",
"sensitive": false,
"shapes": [],
"unsupported_live": null
}
}
}
}

291
litellm-rust/crates/python-bridge/src/cache/binding.rs vendored Normal file
View 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
View 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)
}

861
litellm-rust/crates/python-bridge/src/cache/config.rs vendored Normal file
View file

@ -0,0 +1,861 @@
use std::time::Duration;
use litellm_cache::CacheType;
use litellm_cache_redis::{RedisNode, RedisTopology};
use pyo3::{
exceptions::{PyTypeError, PyValueError},
prelude::*,
types::{PyAny, PyDict, PyList, 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) topology: RedisTopology,
pub(super) connection: RedisConnectionConfig,
}
pub(super) struct AzureBlobCacheConfig {
pub(super) account_url: String,
pub(super) container: String,
}
struct RedisClientProjection<'py> {
topology: RedisTopology,
host: String,
port: u16,
pool_size: usize,
resolved: Bound<'py, PyDict>,
tls: Option<RedisTlsConfig>,
}
const REDIS_PY_DEFAULT_MAX_CONNECTIONS: usize = 1 << 31;
pub(super) enum CacheBackendConfig {
Memory(MemoryCacheConfig),
Redis(Box<RedisCacheConfig>),
AzureBlob(AzureBlobCacheConfig),
}
#[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::AzureBlob) => project_azure_blob(&backend).map(|backend| {
CacheConfigProjection::Native(Box::new(Self {
policy,
backend: CacheBackendConfig::AzureBlob(backend),
}))
}),
Some(
CacheType::RedisSemantic
| CacheType::ValkeySemantic
| CacheType::S3
| CacheType::Disk
| CacheType::QdrantSemantic
| CacheType::Gcs,
)
| None => Ok(CacheConfigProjection::Unsupported(
UnsupportedCacheConfig::Backend,
)),
}
}
pub(super) fn service_mismatch(&self, service: &NativeResponseCache) -> Option<&'static str> {
let default_ttl = match &self.backend {
CacheBackendConfig::Memory(config) => Some(config.default_ttl),
CacheBackendConfig::Redis(config) => Some(config.default_ttl),
CacheBackendConfig::AzureBlob(_) => None,
};
if service.default_ttl() != 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) if service.topology() != Some(&config.topology) => {
Some("facade and native backend topologies must match")
}
CacheBackendConfig::Redis(config) => (service.namespace()
!= config.namespace.as_deref())
.then_some("facade and native backend namespaces must match"),
CacheBackendConfig::AzureBlob(config) => match service.azure_blob_identity() {
None => Some("facade and native backend types must match"),
Some((account_url, container))
if account_url != config.account_url || container != config.container =>
{
Some("facade and native backend containers must match")
}
Some(_) => None,
},
}
}
}
#[inline(never)]
fn project_azure_blob(backend: &Bound<'_, PyAny>) -> PyResult<AzureBlobCacheConfig> {
let client = backend.getattr("container_client")?;
let container = client.getattr("container_name")?.extract::<String>()?;
let url = client.getattr("url")?.extract::<String>()?;
let account_url = url
.strip_suffix(container.as_str())
.and_then(|url| url.strip_suffix('/'))
.ok_or_else(|| PyValueError::new_err("Azure Blob container URL is malformed"))?;
Ok(AzureBlobCacheConfig {
account_url: account_url.to_string(),
container,
})
}
#[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, "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 projection = if has_value(&source, "startup_nodes")? {
project_cluster_client(&source, &client)?
} else {
project_standalone_client(&client)?
};
let RedisClientProjection {
topology,
host,
port,
pool_size,
resolved,
tls,
} = match projection {
Ok(projection) => projection,
Err(reason) => return Ok(Err(reason)),
};
if has_value(&resolved, "credential_provider")? {
return Ok(Err(UnsupportedCacheConfig::RedisCredentials));
}
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>()?,
topology,
connection: RedisConnectionConfig {
host,
port,
database: optional_i64(&resolved, "db")?.unwrap_or(0),
username: optional_dict_string(&resolved, "username")?,
password: optional_dict_string(&resolved, "password")?,
protocol,
pool_size,
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_standalone_client<'py>(
client: &Bound<'py, PyAny>,
) -> PyResult<Result<RedisClientProjection<'py>, UnsupportedCacheConfig>> {
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>()?;
if has_value(&resolved, "redis_connect_func")? {
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));
};
Ok(Ok(RedisClientProjection {
topology: RedisTopology::Standalone,
host: required_string(&resolved, "host")?,
port: port(required_i64(&resolved, "port")?)?,
pool_size: pool.getattr("max_connections")?.extract::<usize>()?,
resolved,
tls,
}))
}
#[inline(never)]
fn project_cluster_client<'py>(
source: &Bound<'py, PyDict>,
client: &Bound<'py, PyAny>,
) -> PyResult<Result<RedisClientProjection<'py>, UnsupportedCacheConfig>> {
let Some(startup_nodes) = startup_nodes(source)? else {
return Ok(Err(UnsupportedCacheConfig::RedisTopology));
};
if !instance_class_is(client, "redis.cluster", "RedisCluster")? {
return Ok(Err(UnsupportedCacheConfig::RedisConnection));
}
let nodes = client.getattr("nodes_manager")?;
if !class_is(
&nodes.getattr("connection_pool_class")?,
"redis.connection",
"ConnectionPool",
)? {
return Ok(Err(UnsupportedCacheConfig::RedisConnection));
}
let resolved = nodes.getattr("connection_kwargs")?.cast_into::<PyDict>()?;
if let Some(connect) = resolved.get_item("redis_connect_func")?
&& !connect.is_none()
{
let own_hook = connect
.getattr("__self__")
.is_ok_and(|owner| owner.is(client))
&& connect
.getattr("__func__")
.and_then(|function| Ok(function.is(&client.get_type().getattr("on_connect")?)))
.unwrap_or(false);
if !own_hook {
return Ok(Err(UnsupportedCacheConfig::RedisCredentials));
}
}
let tls = if optional_bool(&resolved, "ssl")?.unwrap_or(false) {
Some(project_tls(&resolved)?)
} else {
None
};
let first = &startup_nodes[0];
Ok(Ok(RedisClientProjection {
host: first.host.clone(),
port: first.port,
pool_size: optional_i64(&resolved, "max_connections")?
.map(|value| {
usize::try_from(value).map_err(|_| PyValueError::new_err("invalid Redis pool size"))
})
.transpose()?
.unwrap_or(REDIS_PY_DEFAULT_MAX_CONNECTIONS),
topology: RedisTopology::Cluster { startup_nodes },
resolved,
tls,
}))
}
#[inline(never)]
fn startup_nodes(source: &Bound<'_, PyDict>) -> PyResult<Option<Vec<RedisNode>>> {
let Some(nodes) = source.get_item("startup_nodes")? else {
return Ok(None);
};
let Ok(nodes) = nodes.cast_into::<PyList>() else {
return Ok(None);
};
if nodes.is_empty() {
return Ok(None);
}
let mut parsed = Vec::with_capacity(nodes.len());
for node in nodes.iter() {
let Ok(node) = node.cast_into::<PyDict>() else {
return Ok(None);
};
if node.len() != 2 || !has_value(&node, "host")? || !has_value(&node, "port")? {
return Ok(None);
}
let (Ok(host), Ok(port)) = (
required_string(&node, "host"),
required_i64(&node, "port").and_then(port),
) else {
return Ok(None);
};
parsed.push(RedisNode { host, port });
}
Ok(Some(parsed))
}
#[inline(never)]
fn port(value: i64) -> PyResult<u16> {
u16::try_from(value).map_err(|_| PyValueError::new_err("invalid Redis port"))
}
#[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 litellm_cache_redis::{RedisNode, RedisTopology};
use super::{
CacheBackendConfig, CacheConfigProjection, CertificateRequirement, NativeCacheConfig,
RedisProtocol,
};
use crate::cache::native::NativeResponseCache;
fn cluster_facade<'py>(py: Python<'py>, startup_nodes: &str, hook: &str) -> Bound<'py, PyAny> {
facade(
py,
&format!(
"RedisCluster = type('RedisCluster', (), {{'__module__': 'redis.cluster', 'on_connect': lambda self, connection: None}})\n\
client = RedisCluster()\n\
client.nodes_manager = SimpleNamespace(connection_pool_class=ConnectionPool, connection_kwargs={{'password': 'secret', 'redis_connect_func': {hook}, 'protocol': 3, 'ssl': True, 'ssl_cert_reqs': 'none'}})\n\
backend = SimpleNamespace(default_ttl=120, namespace='team', redis_flush_size=100, redis_kwargs={{'startup_nodes': {startup_nodes}, 'password': 'secret'}}, redis_client=client)\n\
facade = SimpleNamespace(type='redis', mode='default-on', ttl=None, namespace='team', supported_call_types=None, redis_flush_size=100, semantic_cache_scope='key', cache=backend)"
),
)
}
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");
});
}
#[test]
fn projects_cluster_startup_nodes_as_redis_topology() {
Python::initialize();
Python::attach(|py| {
let facade = cluster_facade(
py,
"[{'host': 'node-a', 'port': 7000}, {'host': 'node-b', 'port': 7001}]",
"client.on_connect",
);
let CacheConfigProjection::Native(config) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("cluster startup nodes should project natively");
};
let CacheBackendConfig::Redis(redis) = &config.backend else {
panic!("expected Redis configuration");
};
let expected = RedisTopology::Cluster {
startup_nodes: vec![
RedisNode {
host: "node-a".into(),
port: 7000,
},
RedisNode {
host: "node-b".into(),
port: 7001,
},
],
};
assert_eq!(redis.topology, expected);
assert_eq!(redis.connection.host, "node-a");
assert_eq!(redis.connection.port, 7000);
assert_eq!(redis.connection.password.as_deref(), Some("secret"));
assert_eq!(redis.connection.protocol, RedisProtocol::Resp3);
assert_eq!(
redis
.connection
.tls
.as_ref()
.unwrap()
.certificate_requirement,
CertificateRequirement::None
);
});
}
#[test]
fn malformed_startup_nodes_and_foreign_connect_hooks_stay_on_python() {
Python::initialize();
Python::attach(|py| {
for (startup_nodes, hook, message) in [
(
"[{'host': 'node-a', 'port': 7000, 'server_type': 'primary'}]",
"client.on_connect",
"native Redis topology is not implemented",
),
(
"[{'host': 'node-a', 'port': 'seven'}]",
"client.on_connect",
"native Redis topology is not implemented",
),
(
"[]",
"client.on_connect",
"native Redis topology is not implemented",
),
(
"[{'host': 'node-a', 'port': 7000}]",
"lambda connection: None",
"native Redis credentials require Python",
),
] {
let facade = cluster_facade(py, startup_nodes, hook);
let CacheConfigProjection::Unsupported(reason) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("{startup_nodes} with {hook} must stay on Python");
};
assert_eq!(reason.message(), message, "{startup_nodes} with {hook}");
}
});
}
}

393
litellm-rust/crates/python-bridge/src/cache/facade.rs vendored Normal file
View file

@ -0,0 +1,393 @@
use litellm_cache_redis::RedisTopology;
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: Option<usize>,
attributes: RedisPoolAttributes,
}
struct AzureBlobClientGuard {
sync_client: Py<PyAny>,
async_client: Py<PyAny>,
url: String,
container_name: String,
}
enum ConnectionGuard {
None,
RedisPool(RedisPoolGuard),
AzureBlob(AzureBlobClientGuard),
}
struct RedisPoolAttributes {
pool: &'static str,
connection_class: &'static str,
max_connections: Option<&'static str>,
}
const STANDALONE_POOL: RedisPoolAttributes = RedisPoolAttributes {
pool: "connection_pool",
connection_class: "connection_class",
max_connections: Some("max_connections"),
};
const CLUSTER_POOL: RedisPoolAttributes = RedisPoolAttributes {
pool: "nodes_manager",
connection_class: "connection_pool_class",
max_connections: None,
};
pub(super) struct FacadeGuard {
outer: ObjectGuard,
backend: ObjectGuard,
connection: ConnectionGuard,
}
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>, attributes: RedisPoolAttributes) -> PyResult<Self> {
let pool = backend.getattr("redis_client")?.getattr(attributes.pool)?;
Ok(Self {
reference: pool.clone().unbind(),
connection_class: pool.getattr(attributes.connection_class)?.unbind(),
connection_kwargs: pool
.getattr("connection_kwargs")?
.call_method0("copy")?
.unbind(),
max_connections: Self::max_connections(&pool, &attributes)?,
attributes,
})
}
fn max_connections(
pool: &Bound<'_, PyAny>,
attributes: &RedisPoolAttributes,
) -> PyResult<Option<usize>> {
attributes
.max_connections
.map(|name| pool.getattr(name)?.extract::<usize>())
.transpose()
}
fn matches(&self, py: Python<'_>, backend: &Bound<'_, PyAny>) -> PyResult<bool> {
let pool = backend
.getattr("redis_client")?
.getattr(self.attributes.pool)?;
Ok(self.reference.bind(py).is(&pool)
&& self
.connection_class
.bind(py)
.is(&pool.getattr(self.attributes.connection_class)?)
&& self.max_connections == Self::max_connections(&pool, &self.attributes)?
&& 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 AzureBlobClientGuard {
fn capture(backend: &Bound<'_, PyAny>) -> PyResult<Self> {
let sync_client = backend.getattr("container_client")?;
Ok(Self {
url: sync_client.getattr("url")?.extract::<String>()?,
container_name: sync_client.getattr("container_name")?.extract::<String>()?,
sync_client: sync_client.unbind(),
async_client: backend.getattr("async_container_client")?.unbind(),
})
}
fn matches(&self, py: Python<'_>, backend: &Bound<'_, PyAny>) -> PyResult<bool> {
let sync_client = backend.getattr("container_client")?;
Ok(self.sync_client.bind(py).is(&sync_client)
&& self
.async_client
.bind(py)
.is(&backend.getattr("async_container_client")?)
&& self.url == sync_client.getattr("url")?.extract::<String>()?
&& self.container_name == sync_client.getattr("container_name")?.extract::<String>()?)
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&self.sync_client)?;
visit.call(&self.async_client)
}
}
impl ConnectionGuard {
fn capture(kind: &str, cluster: bool, backend: &Bound<'_, PyAny>) -> PyResult<Self> {
Ok(match (kind, cluster) {
("redis", false) => Self::RedisPool(RedisPoolGuard::capture(backend, STANDALONE_POOL)?),
("redis", true) => Self::RedisPool(RedisPoolGuard::capture(backend, CLUSTER_POOL)?),
("azure-blob", _) => Self::AzureBlob(AzureBlobClientGuard::capture(backend)?),
_ => Self::None,
})
}
fn matches(&self, py: Python<'_>, backend: &Bound<'_, PyAny>) -> PyResult<bool> {
match self {
Self::None => Ok(true),
Self::RedisPool(guard) => guard.matches(py, backend),
Self::AzureBlob(guard) => guard.matches(py, backend),
}
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
match self {
Self::None => Ok(()),
Self::RedisPool(guard) => guard.traverse(visit),
Self::AzureBlob(guard) => guard.traverse(visit),
}
}
}
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 cluster = matches!(service.topology(), Some(RedisTopology::Cluster { .. }));
let (module, name, cache_kind) = match (kind, cluster) {
("memory", _) => ("litellm.caching.in_memory_cache", "InMemoryCache", "local"),
("redis", false) => ("litellm.caching.redis_cache", "RedisCache", "redis"),
("redis", true) => (
"litellm.caching.redis_cluster_cache",
"RedisClusterCache",
"redis",
),
("azure-blob", _) => (
"litellm.caching.azure_blob_cache",
"AzureBlobCache",
"azure-blob",
),
_ => 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",
],
)?,
connection: ConnectionGuard::capture(kind, cluster, &backend)?,
})
}
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);
}
self.connection.matches(py, &backend)
}
pub(super) fn traverse(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
self.outer.traverse(&visit)?;
self.backend.traverse(&visit)?;
self.connection.traverse(&visit)
}
}
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
View 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)
}

112
litellm-rust/crates/python-bridge/src/cache/handle.rs vendored Normal file
View file

@ -0,0 +1,112 @@
use litellm_cache_redis::{RedisNode, RedisTopology};
use litellm_host_python::{release_gil, run_sync_value};
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, startup_nodes=None))]
fn redis(
py: Python<'_>,
url: String,
ttl_seconds: f64,
namespace: Option<String>,
startup_nodes: Option<Vec<(String, u16)>>,
) -> PyResult<Self> {
let ttl = Some(duration(ttl_seconds)?);
let topology = match startup_nodes {
None => RedisTopology::Standalone,
Some(nodes) => RedisTopology::Cluster {
startup_nodes: nodes
.into_iter()
.map(|(host, port)| RedisNode { host, port })
.collect(),
},
};
let service = release_gil(py, move || {
NativeResponseCache::redis(&url, &topology, ttl, namespace)
})
.map_err(cache_error)?;
Ok(Self {
service,
guard: None,
pid: std::process::id(),
})
}
#[staticmethod]
#[pyo3(signature = (account_url, container))]
fn azure_blob(py: Python<'_>, account_url: String, container: String) -> PyResult<Self> {
let service = run_sync_value(py, async move {
NativeResponseCache::azure_blob(&account_url, &container)
.await
.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
View 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()),
}
}

243
litellm-rust/crates/python-bridge/src/cache/native.rs vendored Normal file
View file

@ -0,0 +1,243 @@
use std::{sync::Arc, time::Duration};
use litellm_cache::{CacheCodec, CacheConnectionResult, Error};
use litellm_cache_azure_blob::AzureBlobCache;
use litellm_cache_memory::InMemoryCache;
use litellm_cache_redis::{RedisCache, RedisTopology};
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>>,
},
AzureBlob(Arc<ResponseCache<AzureBlobCache<ResponseCacheCodec>>>),
}
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,
topology: &RedisTopology,
ttl: Option<Duration>,
namespace: Option<String>,
) -> Result<Self, Error> {
let backend =
RedisCache::connect(url, topology, ttl, ResponseCacheCodec)?.with_namespace(namespace);
Ok(Self::Redis {
cache: Arc::new(ResponseCache::new(Arc::new(backend))),
buffer: None,
})
}
pub async fn azure_blob(account_url: &str, container: &str) -> Result<Self, Error> {
let backend = AzureBlobCache::connect(
account_url,
container,
ResponseCacheCodec,
tokio::runtime::Handle::current(),
)
.await?;
Ok(Self::AzureBlob(Arc::new(ResponseCache::new(Arc::new(
backend,
)))))
}
pub fn azure_blob_identity(&self) -> Option<(&str, &str)> {
match self {
Self::AzureBlob(cache) => Some((
cache.backend().account_url(),
cache.backend().container_name(),
)),
Self::Memory(_) | Self::Redis { .. } => None,
}
}
}
impl NativeResponseCache {
pub fn kind(&self) -> &'static str {
match self {
Self::Memory(_) => "memory",
Self::Redis { .. } => "redis",
Self::AzureBlob(_) => "azure-blob",
}
}
pub fn default_ttl(&self) -> Option<Duration> {
match self {
Self::Memory(cache) => cache.default_ttl(),
Self::Redis { cache, .. } => cache.default_ttl(),
Self::AzureBlob(cache) => cache.default_ttl(),
}
}
pub fn namespace(&self) -> Option<&str> {
match self {
Self::Memory(_) | Self::AzureBlob(_) => None,
Self::Redis { cache, .. } => cache.backend().namespace(),
}
}
pub fn topology(&self) -> Option<&RedisTopology> {
match self {
Self::Memory(_) | Self::AzureBlob(_) => None,
Self::Redis { cache, .. } => Some(cache.backend().topology()),
}
}
pub fn capacity(&self) -> Option<usize> {
match self {
Self::Memory(cache) => Some(cache.backend().max_size_in_memory()),
Self::Redis { .. } | Self::AzureBlob(_) => None,
}
}
pub fn max_entry_bytes(&self) -> Option<usize> {
match self {
Self::Memory(cache) => cache.backend().max_entry_bytes(),
Self::Redis { .. } | Self::AzureBlob(_) => 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))),
},
other => other,
}
}
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),
Self::AzureBlob(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),
Self::AzureBlob(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),
Self::AzureBlob(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,
Self::AzureBlob(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,
Self::AzureBlob(cache) => cache.async_store(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,
Self::AzureBlob(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,
Self::AzureBlob(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
}
Self::AzureBlob(cache) => 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,
Self::AzureBlob(cache) => cache.test_connection().await,
}
}
}

48
litellm-rust/crates/python-bridge/src/cache/request.rs vendored Normal file
View 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
View 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)
}
}

231
litellm-rust/crates/python-bridge/src/coercion.rs Normal file
View file

@ -0,0 +1,231 @@
use std::collections::BTreeSet;
use litellm_core_utils::serde_compat::parse_str_bool;
use litellm_http::SslVerify;
use pyo3::{
exceptions::{PyAttributeError, PyRuntimeError, PyValueError},
prelude::*,
types::{PyBool, PyString},
};
#[derive(Debug)]
pub(crate) enum ProjectionError {
Python(PyErr),
InvalidConfiguration(String),
UnsupportedLiveObject(String),
InternalSchemaFailure(String),
}
impl From<PyErr> for ProjectionError {
fn from(error: PyErr) -> Self {
Self::Python(error)
}
}
impl From<ProjectionError> for PyErr {
fn from(error: ProjectionError) -> Self {
match error {
ProjectionError::Python(error) => error,
ProjectionError::InvalidConfiguration(message)
| ProjectionError::UnsupportedLiveObject(message) => PyValueError::new_err(message),
ProjectionError::InternalSchemaFailure(message) => PyRuntimeError::new_err(message),
}
}
}
pub(crate) struct Truthy(pub bool);
pub(crate) struct ExactTrue(pub bool);
pub(crate) struct StrBool(pub Option<bool>);
pub(crate) struct OptionalStrictString(pub Option<String>);
pub(crate) struct FalsyOptionalString(pub Option<String>);
pub(crate) struct TuningString(pub Option<String>);
pub(crate) struct StringCollection(pub Vec<String>);
pub(crate) struct SslVerifyInput(pub Option<SslVerify>);
pub(crate) struct Field<'py> {
path: &'static str,
value: Bound<'py, PyAny>,
}
impl<'py> Field<'py> {
pub(crate) fn new(path: &'static str, value: Bound<'py, PyAny>) -> Self {
Self { path, value }
}
pub(crate) fn read(
snapshot: &Bound<'py, PyAny>,
path: &'static str,
) -> Result<Self, ProjectionError> {
let name = path.rsplit('.').next().unwrap_or(path);
match snapshot.getattr(name) {
Ok(value) => Ok(Self::new(path, value)),
Err(error) if error.is_instance_of::<PyAttributeError>(snapshot.py()) => {
match Self::missing_field(snapshot, name) {
Ok(true) => Err(ProjectionError::InternalSchemaFailure(format!(
"{path}: missing snapshot field"
))),
_ => Err(error.into()),
}
}
Err(error) => Err(error.into()),
}
}
fn missing_field(snapshot: &Bound<'_, PyAny>, name: &str) -> PyResult<bool> {
let py = snapshot.py();
let object = py.import("builtins")?.getattr("object")?;
let missing = object.call0()?;
let lookup = py.import("inspect")?.getattr("getattr_static")?;
let declared = lookup.call1((snapshot, name, &missing))?;
let fallback = lookup.call1((snapshot.get_type(), "__getattr__", &missing))?;
let getter = lookup.call1((snapshot.get_type(), "__getattribute__"))?;
Ok(declared.is(&missing)
&& fallback.is(&missing)
&& getter.is(object.getattr("__getattribute__")?))
}
fn expected(&self, expected: &'static str) -> Result<String, ProjectionError> {
Ok(format!(
"{}: expected {expected}, got {}",
self.path,
self.value.get_type().name()?
))
}
fn invalid(&self, expected: &'static str) -> ProjectionError {
match self.expected(expected) {
Ok(message) => ProjectionError::InvalidConfiguration(message),
Err(error) => error,
}
}
pub(crate) fn truthy(&self) -> Result<Truthy, ProjectionError> {
Ok(Truthy(self.value.is_truthy()?))
}
pub(crate) fn exact_true(&self) -> ExactTrue {
ExactTrue(self.value.is(PyBool::new(self.value.py(), true)))
}
pub(crate) fn strict_string(&self) -> Result<String, ProjectionError> {
let value = self
.value
.cast::<PyString>()
.map_err(|_| self.invalid("a string"))?;
Ok(value.to_str()?.to_owned())
}
pub(crate) fn schema_string(&self) -> Result<String, ProjectionError> {
if !self.value.is_instance_of::<PyString>() {
return Err(ProjectionError::InternalSchemaFailure(
self.expected("a string")?,
));
}
self.strict_string()
}
pub(crate) fn schema_bool(&self) -> Result<bool, ProjectionError> {
if !self.value.is_instance_of::<PyBool>() {
return Err(ProjectionError::InternalSchemaFailure(
self.expected("a Boolean")?,
));
}
Ok(self.exact_true().0)
}
pub(crate) fn str_bool(&self) -> Result<StrBool, ProjectionError> {
if self.value.is_none() {
return Ok(StrBool(None));
}
Ok(StrBool(parse_str_bool(&self.strict_string()?)))
}
pub(crate) fn optional_strict_string(&self) -> Result<OptionalStrictString, ProjectionError> {
if self.value.is_none() {
return Ok(OptionalStrictString(None));
}
self.strict_string().map(Some).map(OptionalStrictString)
}
pub(crate) fn falsy_optional_string(&self) -> Result<FalsyOptionalString, ProjectionError> {
if !self.truthy()?.0 {
return Ok(FalsyOptionalString(None));
}
self.strict_string().map(Some).map(FalsyOptionalString)
}
pub(crate) fn tuning_string(&self) -> Result<TuningString, ProjectionError> {
if !self.truthy()?.0 || !self.value.is_instance_of::<PyString>() {
return Ok(TuningString(None));
}
self.strict_string().map(Some).map(TuningString)
}
pub(crate) fn string_collection(&self) -> Result<StringCollection, ProjectionError> {
if !self.truthy()?.0 {
return Ok(StringCollection(Vec::new()));
}
if self.value.is_instance_of::<PyString>() {
return self
.strict_string()
.map(|value| StringCollection(vec![value]));
}
let values = self
.value
.try_iter()?
.filter_map(|item| {
let member = match item {
Ok(value) => Self::new(self.path, value),
Err(error) => return Some(Err(error.into())),
};
match member.truthy() {
Ok(Truthy(false)) => None,
Ok(Truthy(true)) => Some(member.strict_string()),
Err(error) => Some(Err(error)),
}
})
.collect::<Result<Vec<_>, ProjectionError>>()?;
Ok(StringCollection(values))
}
pub(crate) fn host_collection(&self) -> Result<StringCollection, ProjectionError> {
let values = self
.string_collection()?
.0
.into_iter()
.map(|host| litellm_http::media::normalize_host(&host))
.collect::<BTreeSet<_>>();
Ok(StringCollection(values.into_iter().collect()))
}
pub(crate) fn ssl_verify(&self) -> Result<SslVerifyInput, ProjectionError> {
if self.value.is_none() {
return Ok(SslVerifyInput(None));
}
if self.value.is_instance_of::<PyBool>() {
return Ok(SslVerifyInput(Some(if self.exact_true().0 {
SslVerify::Enabled
} else {
SslVerify::Disabled
})));
}
if self.value.is_instance_of::<PyString>() {
let parsed = match self.str_bool()?.0 {
Some(true) => SslVerify::Enabled,
Some(false) => SslVerify::Disabled,
None => SslVerify::CaBundle(self.strict_string()?.into()),
};
return Ok(SslVerifyInput(Some(parsed)));
}
let context = self.value.py().import("ssl")?.getattr("SSLContext")?;
if self.value.is_instance(&context)? {
return Err(ProjectionError::UnsupportedLiveObject(self.expected(
"a Boolean, Boolean string, CA path, or None; live SSLContext is unsupported",
)?));
}
Err(self.invalid("a Boolean, Boolean string, CA path, or None"))
}
}
#[cfg(test)]
mod tests;

372
litellm-rust/crates/python-bridge/src/coercion/tests.rs Normal file
View file

@ -0,0 +1,372 @@
use std::ffi::CString;
use pyo3::{
exceptions::{PyLookupError, PyRuntimeError, PyValueError},
types::PyDict,
};
use rstest::rstest;
use super::*;
fn evaluate<'py>(py: Python<'py>, source: &str) -> Bound<'py, PyAny> {
py.eval(&CString::new(source).unwrap(), None, None).unwrap()
}
#[rstest]
#[case("None", false, false)]
#[case("False", false, false)]
#[case("True", true, true)]
#[case("0", false, false)]
#[case("1", true, false)]
#[case("''", false, false)]
#[case("'false'", true, false)]
#[case("[]", false, false)]
#[case("[0]", true, false)]
#[case("{}", false, false)]
#[case("object()", true, false)]
fn boolean_operations_have_distinct_python_semantics(
#[case] source: &str,
#[case] truth: bool,
#[case] exact: bool,
) {
Python::initialize();
Python::attach(|py| {
let value = evaluate(py, source);
let field = Field::new("test.flag", value.clone());
assert_eq!(field.truthy().unwrap().0, truth);
assert_eq!(field.exact_true().0, exact);
assert_eq!(
field.truthy().unwrap().0,
py.import("builtins")
.unwrap()
.getattr("bool")
.unwrap()
.call1((value,))
.unwrap()
.extract::<bool>()
.unwrap()
);
});
}
#[rstest]
#[case("None", Ok(None), Ok(None), Ok(None))]
#[case("''", Ok(Some("")), Ok(None), Ok(None))]
#[case(
"' value '",
Ok(Some(" value ")),
Ok(Some(" value ")),
Ok(Some(" value "))
)]
#[case("[]", Err(()), Ok(None), Ok(None))]
#[case("0", Err(()), Ok(None), Ok(None))]
#[case("1", Err(()), Err(()), Ok(None))]
#[case("object()", Err(()), Err(()), Ok(None))]
fn string_operations_do_not_conflate_absence_and_type_checks(
#[case] source: &str,
#[case] strict: Result<Option<&str>, ()>,
#[case] fallback: Result<Option<&str>, ()>,
#[case] tuning: Result<Option<&str>, ()>,
) {
Python::initialize();
Python::attach(|py| {
let field = Field::new("test.string", evaluate(py, source));
let owned =
|expected: Result<Option<&str>, ()>| expected.map(|value| value.map(str::to_owned));
assert_eq!(
field
.optional_strict_string()
.map(|value| value.0)
.map_err(|_| ()),
owned(strict)
);
assert_eq!(
field
.falsy_optional_string()
.map(|value| value.0)
.map_err(|_| ()),
owned(fallback)
);
assert_eq!(
field.tuning_string().map(|value| value.0).map_err(|_| ()),
owned(tuning)
);
});
}
#[rstest]
#[case("None", None)]
#[case("' True '", Some(true))]
#[case("' fAlSe '", Some(false))]
#[case("'yes'", None)]
#[case("'1'", None)]
#[case("'unknown'", None)]
fn string_boolean_tokens_remain_separate_from_truthiness(
#[case] source: &str,
#[case] expected: Option<bool>,
) {
Python::initialize();
Python::attach(|py| {
assert_eq!(
Field::new("test.flag", evaluate(py, source))
.str_bool()
.unwrap()
.0,
expected
);
});
}
#[rstest]
#[case("'EXAMPLE.TEST.'", vec!["example.test"])]
#[case("['B.test', '', None, 0, [], 'A.test.', 'b.test']", vec!["a.test", "b.test"])]
#[case("('B.test', 'a.test')", vec!["a.test", "b.test"])]
#[case("{'B.test', 'a.test'}", vec!["a.test", "b.test"])]
#[case("(host for host in ['B.test', 'a.test'])", vec!["a.test", "b.test"])]
#[case("None", vec![])]
#[case("False", vec![])]
fn host_collection_is_owned_normalized_and_deterministic(
#[case] source: &str,
#[case] expected: Vec<&str>,
) {
Python::initialize();
Python::attach(|py| {
assert_eq!(
Field::new("url_policy.user_url_allowed_hosts", evaluate(py, source))
.host_collection()
.unwrap()
.0,
expected
);
});
}
#[test]
fn protocol_errors_preserve_exception_identity_traceback_cause_and_context() {
Python::initialize();
Python::attach(|py| {
let locals = PyDict::new(py);
py.run(
c"
failure = LookupError('protocol failed')
cause = ValueError('cause')
context = RuntimeError('context')
def fail():
try:
raise context
except RuntimeError:
raise failure from cause
class Bool:
def __bool__(self): return fail()
class Length:
def __len__(self): return fail()
class Iter:
def __iter__(self): return fail()
class Next:
def __iter__(self): return self
def __next__(self): return fail()
class Descriptor:
@property
def flag(self): return fail()
values = (Bool(), Length(), Iter(), Next(), [Bool()])
descriptor = Descriptor()
",
Some(&locals),
Some(&locals),
)
.unwrap();
let values = locals.get_item("values").unwrap().unwrap();
for value in values.try_iter().unwrap() {
let error = Field::new("test.flag", value.unwrap())
.host_collection()
.err()
.unwrap();
let error = PyErr::from(error);
assert!(
error
.value(py)
.is(locals.get_item("failure").unwrap().unwrap())
);
assert!(error.is_instance_of::<PyLookupError>(py));
assert!(error.traceback(py).is_some());
assert!(
error
.value(py)
.getattr("__cause__")
.unwrap()
.is(locals.get_item("cause").unwrap().unwrap())
);
assert!(
error
.value(py)
.getattr("__context__")
.unwrap()
.is(locals.get_item("context").unwrap().unwrap())
);
}
let error = Field::read(
&locals.get_item("descriptor").unwrap().unwrap(),
"test.flag",
)
.err()
.unwrap();
assert!(
PyErr::from(error)
.value(py)
.is(locals.get_item("failure").unwrap().unwrap())
);
});
}
#[test]
fn identity_and_string_contents_do_not_invoke_unrelated_protocols() {
Python::initialize();
Python::attach(|py| {
let locals = PyDict::new(py);
py.run(
c"
class Hostile:
def __bool__(self): raise AssertionError('bool called')
def __eq__(self, other): raise AssertionError('eq called')
def __str__(self): raise AssertionError('str called')
class Text(str):
def __str__(self): raise AssertionError('str called')
def strip(self): raise AssertionError('strip called')
def lower(self): raise AssertionError('lower called')
hostile = Hostile()
text = Text(' False ')
",
Some(&locals),
Some(&locals),
)
.unwrap();
let hostile = Field::new("test.flag", locals.get_item("hostile").unwrap().unwrap());
assert!(!hostile.exact_true().0);
assert!(matches!(
hostile.strict_string(),
Err(ProjectionError::InvalidConfiguration(_))
));
let text = Field::new("test.flag", locals.get_item("text").unwrap().unwrap());
assert_eq!(text.strict_string().unwrap(), " False ");
assert_eq!(text.str_bool().unwrap().0, Some(false));
});
}
#[test]
fn missing_snapshot_fields_and_descriptor_attribute_errors_are_distinct() {
Python::initialize();
Python::attach(|py| {
let locals = PyDict::new(py);
py.run(
c"
failure = AttributeError('descriptor failed')
class Snapshot:
@property
def flag(self): raise failure
snapshot = Snapshot()
class Dynamic:
def __getattr__(self, name): raise failure
class Intercepted:
def __getattribute__(self, name): raise failure
dynamic = Dynamic()
intercepted = Intercepted()
",
Some(&locals),
Some(&locals),
)
.unwrap();
let snapshot = locals.get_item("snapshot").unwrap().unwrap();
let descriptor = PyErr::from(Field::read(&snapshot, "test.flag").err().unwrap());
assert!(
descriptor
.value(py)
.is(locals.get_item("failure").unwrap().unwrap())
);
for name in ["dynamic", "intercepted"] {
let value = locals.get_item(name).unwrap().unwrap();
let error = PyErr::from(Field::read(&value, "test.flag").err().unwrap());
assert!(
error
.value(py)
.is(locals.get_item("failure").unwrap().unwrap())
);
}
let missing = PyErr::from(Field::read(&snapshot, "test.missing").err().unwrap());
assert!(missing.is_instance_of::<PyRuntimeError>(py));
assert!(missing.to_string().contains("test.missing"));
});
}
#[test]
fn configuration_errors_name_fields_without_exposing_values() {
Python::initialize();
Python::attach(|py| {
for source in [
"{'secret': 'do-not-print'}",
"['host.test', {'secret': 'do-not-print'}]",
] {
let field = Field::new("test.setting", evaluate(py, source));
let error = PyErr::from(field.falsy_optional_string().err().unwrap());
assert!(error.is_instance_of::<PyValueError>(py));
assert!(error.to_string().contains("test.setting"));
assert!(!error.to_string().contains("do-not-print"));
}
let hosts = Field::new(
"url_policy.user_url_allowed_hosts",
evaluate(py, "['host.test', 1]"),
);
assert!(matches!(
hosts.host_collection(),
Err(ProjectionError::InvalidConfiguration(_))
));
assert!(matches!(
Field::new("test.flag", evaluate(py, "1")).str_bool(),
Err(ProjectionError::InvalidConfiguration(_))
));
});
}
#[test]
fn projection_releases_the_source_collection() {
Python::initialize();
Python::attach(|py| {
let source = evaluate(py, "['A.test']");
let projected = Field::new("test.hosts", source.clone())
.host_collection()
.unwrap()
.0;
source.call_method1("append", ("b.test",)).unwrap();
assert_eq!(projected, ["a.test"]);
assert_eq!(
Field::new("test.hosts", source)
.host_collection()
.unwrap()
.0,
["a.test", "b.test"]
);
});
}
#[rstest]
#[case("True", Some(true))]
#[case("False", Some(false))]
#[case("1", None)]
#[case("None", None)]
#[case("[]", None)]
fn accessor_booleans_are_strict_schema_values(
#[case] source: &str,
#[case] expected: Option<bool>,
) {
Python::initialize();
Python::attach(|py| {
let result = Field::new("secret_manager.readable", evaluate(py, source)).schema_bool();
match expected {
Some(expected) => assert_eq!(result.unwrap(), expected),
None => {
let error = PyErr::from(result.unwrap_err());
assert!(error.is_instance_of::<PyRuntimeError>(py));
assert!(error.to_string().contains("secret_manager.readable"));
}
}
});
}

229
litellm-rust/crates/python-bridge/src/http.rs
View file

@ -7,12 +7,12 @@ use std::{
use litellm_core_utils::settings::ProcessEnvironment;
use litellm_http::{
HttpClientConfig, HttpClientPool, HttpSettings, HttpSettingsLayer, Resolution, SslVerify,
Unsupported,
TlsSource, Unsupported,
media::{PublicDnsResolver, UrlPolicy},
};
use pyo3::{prelude::*, types::PyDict};
use pyo3::{exceptions::PyValueError, prelude::*, types::PyDict};
use crate::{errors::RustBridgeDeclined, python_settings::PythonSettings};
use crate::{coercion::Field, python_settings::PythonSettings};
static POOL: LazyLock<HttpClientPool> =
LazyLock::new(|| HttpClientPool::new(Arc::new(PublicDnsResolver)));
@ -41,6 +41,30 @@ pub(crate) fn call_config(
Ok(resolution.config)
}
pub(crate) fn client_error(error: litellm_http::Error) -> PyErr {
match error {
litellm_http::Error::Read {
tls_source: TlsSource::ClientIdentity,
..
}
| litellm_http::Error::InvalidPem {
tls_source: TlsSource::ClientIdentity,
..
} => PyValueError::new_err(
"http_settings.ssl_certificate: expected a readable PEM certificate and private key",
),
litellm_http::Error::Read {
tls_source: TlsSource::CaBundle,
..
}
| litellm_http::Error::InvalidPem {
tls_source: TlsSource::CaBundle,
..
} => PyValueError::new_err("http_settings.ssl_verify: expected a readable PEM CA bundle"),
_ => PyValueError::new_err("http_settings: native HTTP client configuration is invalid"),
}
}
fn unreported(
reported: &Mutex<HashSet<Unsupported>>,
unsupported: Vec<Unsupported>,
@ -53,25 +77,25 @@ fn unreported(
}
pub(crate) fn url_policy(py: Python<'_>) -> PyResult<UrlPolicy> {
let policy: PythonUrlPolicy =
PythonSettings::UrlPolicy
.read(py)?
.extract()
.map_err(|error: PyErr| {
RustBridgeDeclined::new_err(format!(
"litellm URL policy cannot be used by the Rust route: {error}"
))
})?;
project_url_policy(&PythonSettings::UrlPolicy.read(py)?)
}
fn project_url_policy(value: &Bound<'_, PyAny>) -> PyResult<UrlPolicy> {
Ok(UrlPolicy {
validate: policy.user_url_validation,
allowed_hosts: policy.user_url_allowed_hosts,
validate: Field::read(value, "url_policy.user_url_validation")?
.truthy()?
.0,
allowed_hosts: Field::read(value, "url_policy.user_url_allowed_hosts")?
.host_collection()?
.0,
})
}
fn call_ssl_verify(kwargs: &Bound<'_, PyDict>) -> PyResult<Option<SslVerify>> {
Ok(kwargs
.get_item("ssl_verify")?
.and_then(|value| ssl_verify(&value)))
match kwargs.get_item("ssl_verify")? {
Some(value) => Ok(Field::new("request.ssl_verify", value).ssl_verify()?.0),
None => Ok(None),
}
}
fn for_call(call_ssl_verify: Option<SslVerify>, asynchronous: bool) -> HttpSettingsLayer {
@ -82,64 +106,47 @@ fn for_call(call_ssl_verify: Option<SslVerify>, asynchronous: bool) -> HttpSetti
}
}
#[derive(FromPyObject)]
struct PythonUrlPolicy {
user_url_validation: bool,
user_url_allowed_hosts: Vec<String>,
}
#[derive(FromPyObject)]
struct PythonHttpSettings<'py> {
ssl_verify: Bound<'py, PyAny>,
ssl_certificate: Option<String>,
ssl_security_level: Option<String>,
ssl_ecdh_curve: Option<String>,
force_ipv4: bool,
http2: bool,
aiohttp_trust_env: bool,
disable_aiohttp_trust_env: bool,
disable_aiohttp_transport: bool,
user_agent: String,
}
fn configured(value: &Bound<'_, PyAny>) -> PyResult<HttpSettingsLayer> {
let python: PythonHttpSettings = value.extract().map_err(|error: PyErr| {
RustBridgeDeclined::new_err(format!(
"litellm HTTP settings cannot be used by the Rust route: {error}"
))
})?;
Ok(HttpSettingsLayer {
ssl_verify: ssl_verify(&python.ssl_verify),
ssl_certificate: python.ssl_certificate.map(PathBuf::from),
ssl_security_level: python.ssl_security_level,
ssl_ecdh_curve: python.ssl_ecdh_curve,
force_ipv4: Some(python.force_ipv4),
http2: Some(python.http2),
aiohttp_trust_env: Some(python.aiohttp_trust_env),
disable_aiohttp_trust_env: Some(python.disable_aiohttp_trust_env),
disable_aiohttp_transport: Some(python.disable_aiohttp_transport),
user_agent: Some(python.user_agent),
ssl_verify: Field::read(value, "http_settings.ssl_verify")?
.ssl_verify()?
.0,
ssl_certificate: Field::read(value, "http_settings.ssl_certificate")?
.optional_strict_string()?
.0
.map(PathBuf::from),
ssl_security_level: Field::read(value, "http_settings.ssl_security_level")?
.tuning_string()?
.0,
ssl_ecdh_curve: Field::read(value, "http_settings.ssl_ecdh_curve")?
.tuning_string()?
.0,
force_ipv4: Some(Field::read(value, "http_settings.force_ipv4")?.truthy()?.0),
http2: Some(Field::read(value, "http_settings.http2")?.exact_true().0),
aiohttp_trust_env: Some(
Field::read(value, "http_settings.aiohttp_trust_env")?
.truthy()?
.0,
),
disable_aiohttp_trust_env: Some(
Field::read(value, "http_settings.disable_aiohttp_trust_env")?
.truthy()?
.0,
),
disable_aiohttp_transport: Some(
Field::read(value, "http_settings.disable_aiohttp_transport")?
.exact_true()
.0,
),
user_agent: Some(Field::read(value, "http_settings.user_agent")?.schema_string()?),
..HttpSettingsLayer::default()
})
}
fn ssl_verify(value: &Bound<'_, PyAny>) -> Option<SslVerify> {
if let Ok(enabled) = value.extract::<bool>() {
return Some(if enabled {
SslVerify::Enabled
} else {
SslVerify::Disabled
});
}
value
.extract::<String>()
.ok()
.map(|path| SslVerify::parse(&path))
}
#[cfg(test)]
mod tests {
use litellm_http::Verify;
use pyo3::exceptions::PyRuntimeError;
use rstest::rstest;
use super::*;
@ -163,7 +170,7 @@ defaults = dict(
user_agent='litellm/test',
)
defaults.update(dict({overrides}))
settings = types.SimpleNamespace(**{{name: defaults[name] for name in json.loads(contract)['http_settings']}})
settings = types.SimpleNamespace(**{{name: defaults[name] for name in json.loads(contract)['http_settings']['fields']}})
"
);
let locals = PyDict::new(py);
@ -189,6 +196,33 @@ settings = types.SimpleNamespace(**{{name: defaults[name] for name in json.loads
});
}
#[test]
fn client_error_uses_tls_source_when_paths_match() {
Python::initialize();
Python::attach(|py| {
let path = PathBuf::from("/shared.pem");
let ca_error = client_error(litellm_http::Error::InvalidPem {
path: path.clone(),
message: "invalid".into(),
tls_source: TlsSource::CaBundle,
});
assert_eq!(
ca_error.to_string(),
"ValueError: http_settings.ssl_verify: expected a readable PEM CA bundle"
);
let client_error = client_error(litellm_http::Error::InvalidPem {
path,
message: "invalid".into(),
tls_source: TlsSource::ClientIdentity,
});
assert!(client_error.is_instance_of::<PyValueError>(py));
assert_eq!(
client_error.to_string(),
"ValueError: http_settings.ssl_certificate: expected a readable PEM certificate and private key"
);
});
}
#[test]
fn python_settings_flow_into_the_configured_layer() {
Python::initialize();
@ -259,12 +293,16 @@ user_agent='litellm/9.9.9',
});
}
#[test]
fn ssl_context_global_is_ignored_so_environment_and_defaults_apply() {
#[rstest]
#[case("ssl_verify=object()")]
#[case("ssl_verify=__import__('ssl').SSLContext(__import__('ssl').PROTOCOL_TLS_CLIENT)")]
#[case("ssl_certificate=1")]
fn invalid_http_configuration_is_terminal(#[case] overrides: &str) {
Python::initialize();
Python::attach(|py| {
let layer = configured(&python_settings(py, "ssl_verify=object()")).unwrap();
assert_eq!(layer.ssl_verify, None);
let error = configured(&python_settings(py, overrides)).unwrap_err();
assert!(error.is_instance_of::<PyValueError>(py));
assert!(error.to_string().contains("http_settings.ssl_"));
});
}
@ -281,11 +319,21 @@ user_agent='litellm/9.9.9',
}
#[test]
fn mistyped_python_settings_decline_instead_of_raising() {
fn mutable_globals_use_their_consumer_operations() {
Python::initialize();
Python::attach(|py| {
let error = configured(&python_settings(py, "force_ipv4='yes'")).unwrap_err();
assert!(error.is_instance_of::<RustBridgeDeclined>(py));
let layer = configured(&python_settings(py,
"force_ipv4='yes', http2=1, disable_aiohttp_transport=1, aiohttp_trust_env=[1], disable_aiohttp_trust_env=[], ssl_security_level=1, ssl_ecdh_curve=[]"
)).unwrap();
assert_eq!(layer.force_ipv4, Some(true));
assert_eq!(layer.http2, Some(false));
assert_eq!(layer.disable_aiohttp_transport, Some(false));
assert_eq!(layer.aiohttp_trust_env, Some(true));
assert_eq!(layer.disable_aiohttp_trust_env, Some(false));
assert_eq!(layer.ssl_security_level, None);
assert_eq!(layer.ssl_ecdh_curve, None);
let error = configured(&python_settings(py, "user_agent=1")).unwrap_err();
assert!(error.is_instance_of::<PyRuntimeError>(py));
});
}
@ -323,17 +371,36 @@ user_agent='litellm/9.9.9',
}
#[test]
fn live_ssl_context_argument_is_ignored_so_the_configured_value_applies() {
fn live_ssl_context_argument_raises_instead_of_using_another_layer() {
Python::initialize();
Python::attach(|py| {
let kwargs = PyDict::new(py);
kwargs
.set_item("ssl_verify", py.eval(c"object()", None, None).unwrap())
let ssl = py.import("ssl").unwrap();
let context = ssl
.getattr("SSLContext")
.unwrap()
.call1((ssl.getattr("PROTOCOL_TLS_CLIENT").unwrap(),))
.unwrap();
let call = for_call(call_ssl_verify(&kwargs).unwrap(), true);
let settings =
HttpSettings::from_layers([call, configured_ssl_verify(SslVerify::Disabled)]);
assert_eq!(settings.ssl_verify, Some(SslVerify::Disabled));
kwargs.set_item("ssl_verify", context).unwrap();
let error = call_ssl_verify(&kwargs).unwrap_err();
assert!(error.is_instance_of::<PyValueError>(py));
assert!(error.to_string().contains("request.ssl_verify"));
assert!(error.to_string().contains("SSLContext"));
});
}
#[test]
fn url_policy_uses_truthiness_and_normalized_owned_hosts() {
Python::initialize();
Python::attach(|py| {
let value = py.eval(c"__import__('types').SimpleNamespace(user_url_validation=[], user_url_allowed_hosts=['B.test', 'a.test.', 'b.test'])", None, None).unwrap();
assert_eq!(
project_url_policy(&value).unwrap(),
UrlPolicy {
validate: false,
allowed_hosts: vec!["a.test".into(), "b.test".into()],
}
);
});
}

13
litellm-rust/crates/python-bridge/src/lib.rs
View file

@ -1,3 +1,5 @@
mod cache;
mod coercion;
mod credentials;
mod diagnostics;
mod errors;
@ -9,6 +11,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 +35,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::*;

54
litellm-rust/crates/python-bridge/src/marshal.rs
View file

@ -172,6 +172,60 @@ mod tests {
request_input_sources(&kwargs, names.iter().copied())
}
#[serde_with::serde_as]
#[derive(Debug, serde::Deserialize, serde::Serialize, PartialEq)]
struct Numbers {
#[serde_as(deserialize_as = "Option<Vec<litellm_core_utils::serde_compat::LaxI64>>")]
integers: Option<Vec<i64>>,
#[serde_as(deserialize_as = "Option<litellm_core_utils::serde_compat::FiniteF64>")]
float: Option<f64>,
}
#[test]
fn numeric_adapters_agree_across_json_and_python_boundaries() {
Python::initialize();
Python::attach(|py| {
for input in [
json!({}),
json!({"integers": null, "float": null}),
json!({"integers": [i64::MIN, i64::MAX, "9007199254740993.0", " +1_000.00 ", true, 3.0], "float": " 1.25 "}),
json!({"integers": [u64::MAX]}),
json!({"integers": ["1.0000000000000001"]}),
json!({"integers": [2.5]}),
json!({"float": "NaN"}),
json!({"float": "inf"}),
json!({"float": "1e999"}),
json!({"float": true}),
json!({"float": u64::MAX}),
] {
let expected = serde_json::from_value::<Numbers>(input.clone());
let python = litellm_host_python::to_py(py, &input).unwrap();
let actual = from_py::<Numbers>(python.bind(py));
match (expected, actual) {
(Ok(expected), Ok(actual)) => {
assert_eq!(actual, expected);
let serialized = litellm_host_python::to_py(py, &actual).unwrap();
assert_eq!(
from_py::<Value>(serialized.bind(py)).unwrap(),
serde_json::to_value(expected).unwrap()
);
}
(Err(_), Err(_)) => {}
mismatch => panic!("boundary mismatch for {input}: {mismatch:?}"),
}
}
for source in [
c"{'float': float('nan')}",
c"{'float': float('inf')}",
c"{'integers': [float('inf')]}",
c"{'integers': [2 ** 100]}",
] {
let value = py.eval(source, None, None).unwrap();
assert!(from_py::<Numbers>(&value).is_err());
}
});
}
#[test]
fn argument_converters_keep_nested_values_and_accept_explicit_none() {
Python::initialize();

216
litellm-rust/crates/python-bridge/src/python_settings.rs
View file

@ -43,32 +43,204 @@ pub(crate) const CONTRACT: &str = include_str!("../python_settings.json");
#[cfg(test)]
mod tests {
use std::{collections::BTreeSet, ffi::CString};
use pyo3::{prelude::*, types::PyDict};
use super::{CONTRACT, PythonSettings};
use pyo3::prelude::*;
use serde_json::{Value, json};
struct SettingSpec {
group: &'static str,
name: &'static str,
adapter: &'static str,
precedence: &'static str,
sensitive: bool,
shapes: &'static [&'static str],
unsupported_live: Option<&'static str>,
}
const SETTINGS: &[SettingSpec] = &[
SettingSpec {
group: "http_settings",
name: "ssl_verify",
adapter: "SslVerifyInput",
precedence: "module_global",
sensitive: false,
shapes: &["none", "bool", "str"],
unsupported_live: Some("configuration_error"),
},
SettingSpec {
group: "http_settings",
name: "ssl_certificate",
adapter: "OptionalStrictString",
precedence: "module_global",
sensitive: false,
shapes: &[],
unsupported_live: None,
},
SettingSpec {
group: "http_settings",
name: "ssl_security_level",
adapter: "TuningString",
precedence: "module_global",
sensitive: false,
shapes: &[],
unsupported_live: None,
},
SettingSpec {
group: "http_settings",
name: "ssl_ecdh_curve",
adapter: "TuningString",
precedence: "module_global",
sensitive: false,
shapes: &[],
unsupported_live: None,
},
SettingSpec {
group: "http_settings",
name: "force_ipv4",
adapter: "Truthy",
precedence: "module_global",
sensitive: false,
shapes: &[],
unsupported_live: None,
},
SettingSpec {
group: "http_settings",
name: "http2",
adapter: "ExactTrue",
precedence: "module_global",
sensitive: false,
shapes: &[],
unsupported_live: None,
},
SettingSpec {
group: "http_settings",
name: "aiohttp_trust_env",
adapter: "Truthy",
precedence: "module_global",
sensitive: false,
shapes: &[],
unsupported_live: None,
},
SettingSpec {
group: "http_settings",
name: "disable_aiohttp_trust_env",
adapter: "Truthy",
precedence: "module_global",
sensitive: false,
shapes: &[],
unsupported_live: None,
},
SettingSpec {
group: "http_settings",
name: "disable_aiohttp_transport",
adapter: "ExactTrue",
precedence: "module_global",
sensitive: false,
shapes: &[],
unsupported_live: None,
},
SettingSpec {
group: "http_settings",
name: "user_agent",
adapter: "StrictString",
precedence: "accessor",
sensitive: false,
shapes: &[],
unsupported_live: None,
},
SettingSpec {
group: "url_policy",
name: "user_url_validation",
adapter: "Truthy",
precedence: "module_global",
sensitive: false,
shapes: &[],
unsupported_live: None,
},
SettingSpec {
group: "url_policy",
name: "user_url_allowed_hosts",
adapter: "HostCollection",
precedence: "module_global",
sensitive: false,
shapes: &[],
unsupported_live: None,
},
SettingSpec {
group: "provider_defaults",
name: "vertex_project",
adapter: "FalsyOptionalString",
precedence: "module_global",
sensitive: true,
shapes: &[],
unsupported_live: None,
},
SettingSpec {
group: "provider_defaults",
name: "vertex_location",
adapter: "FalsyOptionalString",
precedence: "module_global",
sensitive: true,
shapes: &[],
unsupported_live: None,
},
SettingSpec {
group: "provider_defaults",
name: "enable_azure_ad_token_refresh",
adapter: "ExactTrue",
precedence: "module_global",
sensitive: false,
shapes: &[],
unsupported_live: None,
},
SettingSpec {
group: "secret_manager",
name: "readable",
adapter: "StrictBool",
precedence: "accessor",
sensitive: false,
shapes: &[],
unsupported_live: None,
},
];
#[test]
fn every_settings_group_is_in_the_python_contract() {
Python::initialize();
Python::attach(|py| {
let locals = PyDict::new(py);
locals.set_item("contract", CONTRACT).unwrap();
let source = CString::new("import json\nkeys = list(json.loads(contract))").unwrap();
py.run(&source, Some(&locals), Some(&locals)).unwrap();
let declared: BTreeSet<String> = locals
.get_item("keys")
fn settings_manifest_matches_the_semantic_contract() {
pyo3::Python::initialize();
let manifest: Value = pyo3::Python::attach(|py| {
let value = py
.import("json")
.unwrap()
.unwrap()
.extract::<Vec<String>>()
.unwrap()
.into_iter()
.collect();
let read: BTreeSet<String> = PythonSettings::ALL
.map(|group| group.name().to_owned())
.into();
assert_eq!(read, declared);
.call_method1("loads", (CONTRACT,))
.unwrap();
litellm_host_python::from_py(&value).unwrap()
});
let expected: serde_json::Map<String, Value> = PythonSettings::ALL
.into_iter()
.map(|group| {
let fields: serde_json::Map<String, Value> = SETTINGS
.iter()
.filter(|spec| spec.group == group.name())
.map(|spec| {
(
spec.name.to_owned(),
json!({
"adapter": spec.adapter,
"required": true,
"precedence": spec.precedence,
"sensitive": spec.sensitive,
"shapes": spec.shapes,
"unsupported_live": spec.unsupported_live,
}),
)
})
.collect();
(
group.name().to_owned(),
json!({"version": 1, "fields": fields}),
)
})
.collect();
assert_eq!(manifest, Value::Object(expected));
}
}

73
litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs
View file

@ -19,7 +19,7 @@ use pyo3::{
types::{PyDict, PyTuple},
};
use crate::{errors::RustBridgeDeclined, http, python_settings::PythonSettings};
use crate::{coercion::Field, errors::RustBridgeDeclined, http, python_settings::PythonSettings};
const SURFACE: LegacySurface = LegacySurface {
call_type: "ocr",
@ -51,7 +51,7 @@ fn run_ocr(
ocr_settings(py)?,
secrets,
)
.map_err(|error| RustBridgeDeclined::new_err(error.to_string()))?;
.map_err(http::client_error)?;
run_legacy_call(
py,
if asynchronous { ASYNC_SURFACE } else { SURFACE },
@ -62,14 +62,8 @@ fn run_ocr(
)
}
#[derive(FromPyObject)]
struct PythonSecretManager {
readable: bool,
}
fn process_environment_secrets(secret_manager: &Bound<'_, PyAny>) -> PyResult<Secrets> {
let manager: PythonSecretManager = secret_manager.extract()?;
if manager.readable {
if Field::read(secret_manager, "secret_manager.readable")?.schema_bool()? {
return Err(RustBridgeDeclined::new_err(
"a readable secret manager is configured and the Rust route only reads the process environment",
));
@ -77,26 +71,24 @@ fn process_environment_secrets(secret_manager: &Bound<'_, PyAny>) -> PyResult<Se
Ok(Arc::new(ProcessEnvironment))
}
#[derive(FromPyObject)]
struct PythonProviderDefaults {
vertex_project: Option<String>,
vertex_location: Option<String>,
enable_azure_ad_token_refresh: Option<bool>,
fn ocr_settings(py: Python<'_>) -> PyResult<OcrSettings> {
project_provider_defaults(&PythonSettings::ProviderDefaults.read(py)?)
}
fn ocr_settings(py: Python<'_>) -> PyResult<OcrSettings> {
let defaults: PythonProviderDefaults = PythonSettings::ProviderDefaults
.read(py)?
.extract()
.map_err(|error: PyErr| {
RustBridgeDeclined::new_err(format!(
"litellm provider defaults cannot be used by the Rust route: {error}"
))
})?;
fn project_provider_defaults(value: &Bound<'_, PyAny>) -> PyResult<OcrSettings> {
Ok(OcrSettings {
vertex_project: defaults.vertex_project,
vertex_location: defaults.vertex_location,
enable_azure_ad_token_refresh: defaults.enable_azure_ad_token_refresh == Some(true),
vertex_project: Field::read(value, "provider_defaults.vertex_project")?
.falsy_optional_string()?
.0,
vertex_location: Field::read(value, "provider_defaults.vertex_location")?
.falsy_optional_string()?
.0,
enable_azure_ad_token_refresh: Field::read(
value,
"provider_defaults.enable_azure_ad_token_refresh",
)?
.exact_true()
.0,
..OcrSettings::from_environment(&ProcessEnvironment)
})
}
@ -140,6 +132,35 @@ mod tests {
locals.get_item("manager").unwrap().unwrap()
}
#[test]
fn provider_defaults_distinguish_falsey_values_and_exact_true() {
Python::initialize();
Python::attach(|py| {
let value = py.eval(c"__import__('types').SimpleNamespace(vertex_project=[], vertex_location=0, enable_azure_ad_token_refresh=1)", None, None).unwrap();
let projected = super::project_provider_defaults(&value).unwrap();
assert_eq!(projected.vertex_project, None);
assert_eq!(projected.vertex_location, None);
assert!(!projected.enable_azure_ad_token_refresh);
value.setattr("vertex_project", "project").unwrap();
value.setattr("vertex_location", "region").unwrap();
value
.setattr("enable_azure_ad_token_refresh", true)
.unwrap();
let next = super::project_provider_defaults(&value).unwrap();
assert_eq!(next.vertex_project.as_deref(), Some("project"));
assert_eq!(next.vertex_location.as_deref(), Some("region"));
assert!(next.enable_azure_ad_token_refresh);
value.setattr("vertex_project", 1).unwrap();
let error = super::project_provider_defaults(&value).err().unwrap();
assert!(error.is_instance_of::<pyo3::exceptions::PyValueError>(py));
assert!(
error
.to_string()
.contains("provider_defaults.vertex_project")
);
});
}
#[test]
fn a_readable_secret_manager_sends_the_call_back_to_python() {
Python::initialize();

24
litellm-rust/crates/secrets-azure/Cargo.toml Normal file
View file

@ -0,0 +1,24 @@
[package]
name = "litellm-secrets-azure"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
litellm-auth-azure.workspace = true
litellm-auth-types.workspace = true
litellm-secrets-types.workspace = true
litellm-core-utils.workspace = true
reqwest.workspace = true
serde.workspace = true
thiserror.workspace = true
veil.workspace = true
percent-encoding = "2.3"
[dev-dependencies]
tokio.workspace = true
wiremock = "0.6.5"
rstest.workspace = true
serde_json.workspace = true
sha2.workspace = true

25
litellm-rust/crates/secrets-azure/src/error.rs Normal file
View file

@ -0,0 +1,25 @@
#[derive(thiserror::Error, veil::Redact)]
pub enum Error {
#[error("{0} environment variable is missing")]
MissingEnvironment(&'static str),
#[error("AZURE_KEY_VAULT_URI is not a valid https vault URL")]
VaultUri,
#[error("Azure Key Vault credentials are not configured")]
MissingCredentials,
#[error(transparent)]
Auth(
#[from]
#[redact]
litellm_auth_types::Error,
),
#[error("Azure Key Vault request failed")]
Http(
#[source]
#[redact]
reqwest::Error,
),
#[error("Azure Key Vault returned HTTP {0}")]
Status(u16),
#[error("Azure Key Vault response is missing the secret value")]
MissingValue,
}

118
litellm-rust/crates/secrets-azure/src/key_vault.rs Normal file
View file

@ -0,0 +1,118 @@
use std::sync::Arc;
use litellm_auth_azure::{AzureAuthInputs, AzureAuthService, ConfigValue};
use litellm_auth_types::{InputSource, Sourced};
use litellm_core_utils::settings::Lookup;
use litellm_secrets_types::{Secret, SecretValue};
use percent_encoding::{AsciiSet, NON_ALPHANUMERIC};
use serde::Deserialize;
use crate::Error;
const AZURE_KEY_VAULT_URI: &str = "AZURE_KEY_VAULT_URI";
const API_VERSION: &str = "7.4";
const PATH_SEGMENT: &AsciiSet = &NON_ALPHANUMERIC
.remove(b'-')
.remove(b'.')
.remove(b'_')
.remove(b'~');
#[derive(Clone)]
pub struct AzureKeyVault {
client: reqwest::Client,
vault: reqwest::Url,
auth: Arc<AzureAuthService>,
inputs: Arc<AzureAuthInputs>,
environment: Arc<dyn Lookup + Send + Sync>,
}
#[derive(Deserialize)]
struct SecretResponse {
value: Option<String>,
}
impl AzureKeyVault {
pub fn with_client(
client: reqwest::Client,
vault: reqwest::Url,
environment: Arc<dyn Lookup + Send + Sync>,
) -> Result<Self, Error> {
if vault.host_str().is_none() {
return Err(Error::VaultUri);
}
let inputs = AzureAuthInputs {
azure_scope: ConfigValue::Value(Sourced::new(
scope_for(&vault),
InputSource::Deployment,
)),
enable_azure_ad_token_refresh: Sourced::new(true, InputSource::Deployment),
..AzureAuthInputs::default()
};
Ok(Self {
client,
vault,
auth: Arc::new(AzureAuthService::default()),
inputs: Arc::new(inputs),
environment,
})
}
pub fn new(environment: Arc<dyn Lookup + Send + Sync>) -> Result<Self, Error> {
let value = environment
.get(AZURE_KEY_VAULT_URI)
.ok_or(Error::MissingEnvironment(AZURE_KEY_VAULT_URI))?;
let vault = reqwest::Url::parse(&value).map_err(|_| Error::VaultUri)?;
if vault.scheme() != "https" || vault.host_str().is_none() {
return Err(Error::VaultUri);
}
Self::with_client(reqwest::Client::new(), vault, environment)
}
pub fn scope(&self) -> &str {
self.inputs
.azure_scope
.as_value()
.map(|value| value.value().as_str())
.unwrap_or_default()
}
pub async fn get_secret_from_azure_key_vault(
&self,
name: &str,
) -> Result<Option<Secret>, Error> {
let token = self
.auth
.get_azure_ad_token(&self.inputs, &|key| self.environment.get(key))
.await?
.ok_or(Error::MissingCredentials)?;
let encoded_name = percent_encoding::utf8_percent_encode(name, PATH_SEGMENT);
let url = self
.vault
.join(&format!("secrets/{encoded_name}?api-version={API_VERSION}"))
.map_err(|_| Error::VaultUri)?;
let response = self
.client
.get(url)
.bearer_auth(token.value().secret().expose())
.send()
.await
.map_err(Error::Http)?;
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 payload: SecretResponse = response.json().await.map_err(Error::Http)?;
let value = payload.value.ok_or(Error::MissingValue)?;
Ok(Some(Secret::String(SecretValue::new(value))))
}
}
fn scope_for(vault: &reqwest::Url) -> String {
let host = vault.host_str().unwrap_or_default();
let resource = host
.split_once('.')
.map_or(host, |(_, remainder)| remainder);
format!("https://{resource}/.default")
}

7
litellm-rust/crates/secrets-azure/src/lib.rs Normal file
View file

@ -0,0 +1,7 @@
#![forbid(unsafe_code)]
mod error;
mod key_vault;
pub use error::Error;
pub use key_vault::AzureKeyVault;

8
litellm-rust/crates/secrets-azure/tests/fixtures/key_vault_parity.json vendored Normal file
View file

@ -0,0 +1,8 @@
{
"cases": [
{"name": "plain_value", "secret_name": "OPENAI-API-KEY", "response": {"status": 200, "body": {"value": "sk-parity-1", "id": "https://example.vault.azure.net/secrets/OPENAI-API-KEY/abc"}}, "expected": {"value": "sk-parity-1"}},
{"name": "json_value_is_kept_as_string", "secret_name": "JSON-SECRET", "response": {"status": 200, "body": {"value": "{\"api_key\": \"nested\"}", "id": "https://example.vault.azure.net/secrets/JSON-SECRET/abc"}}, "expected": {"value": "{\"api_key\": \"nested\"}"}},
{"name": "missing_secret", "secret_name": "MISSING", "response": {"status": 404, "body": {"error": {"code": "SecretNotFound", "message": "not found"}}}, "expected": {"missing": true}},
{"name": "forbidden", "secret_name": "FORBIDDEN", "response": {"status": 403, "body": {"error": {"code": "Forbidden", "message": "denied"}}}, "expected": {"error": true}}
]
}

222
litellm-rust/crates/secrets-azure/tests/key_vault.rs Normal file
View file

@ -0,0 +1,222 @@
use std::sync::Arc;
use litellm_secrets_azure::{AzureKeyVault, Error};
use litellm_secrets_types::{Secret, SecretValue};
use serde::Deserialize;
use wiremock::{
Mock, MockServer, ResponseTemplate,
matchers::{header, path, query_param},
};
fn manager(server: &MockServer) -> AzureKeyVault {
AzureKeyVault::with_client(
reqwest::Client::new(),
server.uri().parse().unwrap(),
Arc::new(|name: &str| (name == "AZURE_AD_TOKEN").then(|| "fake".to_owned())),
)
.unwrap()
}
#[tokio::test]
async fn reads_secret_with_bearer_token_and_api_version() {
let server = MockServer::start().await;
Mock::given(path("/secrets/OPENAI-API-KEY"))
.and(query_param("api-version", "7.4"))
.and(header("authorization", "Bearer fake"))
.respond_with(
ResponseTemplate::new(200)
.set_body_json(serde_json::json!({"value": "s3cret", "id": "secret-id"})),
)
.expect(1)
.mount(&server)
.await;
let secret = manager(&server)
.get_secret_from_azure_key_vault("OPENAI-API-KEY")
.await
.unwrap()
.unwrap();
assert_eq!(secret, Secret::String(SecretValue::new("s3cret")));
}
#[tokio::test]
async fn percent_encodes_secret_name_path_segment() {
let server = MockServer::start().await;
Mock::given(path("/secrets/name%2Fwith%20spaces"))
.and(query_param("api-version", "7.4"))
.respond_with(
ResponseTemplate::new(200).set_body_json(serde_json::json!({"value": "value"})),
)
.expect(1)
.mount(&server)
.await;
let secret = manager(&server)
.get_secret_from_azure_key_vault("name/with spaces")
.await
.unwrap()
.unwrap();
assert_eq!(secret.as_str(), Some("value"));
}
#[rstest::rstest]
#[case::not_found(404, None)]
#[case::forbidden(403, Some(403))]
#[tokio::test]
async fn handles_statuses(#[case] status: u16, #[case] expected_status: Option<u16>) {
let server = MockServer::start().await;
Mock::given(path("/secrets/NAME"))
.respond_with(ResponseTemplate::new(status))
.expect(1)
.mount(&server)
.await;
let result = manager(&server)
.get_secret_from_azure_key_vault("NAME")
.await;
match expected_status {
None => assert_eq!(result.unwrap(), None),
Some(status) => assert!(matches!(result, Err(Error::Status(actual)) if actual == status)),
}
}
#[tokio::test]
async fn missing_value_is_an_error() {
let server = MockServer::start().await;
Mock::given(path("/secrets/NAME"))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({})))
.expect(1)
.mount(&server)
.await;
assert!(matches!(
manager(&server)
.get_secret_from_azure_key_vault("NAME")
.await,
Err(Error::MissingValue)
));
}
#[test]
fn new_validates_vault_environment() {
assert!(matches!(
AzureKeyVault::new(Arc::new(|_: &str| None)),
Err(Error::MissingEnvironment("AZURE_KEY_VAULT_URI"))
));
assert!(matches!(
AzureKeyVault::new(Arc::new(|name: &str| {
(name == "AZURE_KEY_VAULT_URI").then(|| "http://vault.example".to_owned())
})),
Err(Error::VaultUri)
));
assert!(matches!(
AzureKeyVault::new(Arc::new(|name: &str| {
(name == "AZURE_KEY_VAULT_URI").then(|| "vault.example".to_owned())
})),
Err(Error::VaultUri)
));
}
#[rstest::rstest]
#[case("https://myvault.vault.azure.net", "https://vault.azure.net/.default")]
#[case(
"https://v.vault.usgovcloudapi.net/",
"https://vault.usgovcloudapi.net/.default"
)]
#[case("http://localhost:8080", "https://localhost/.default")]
#[test]
fn derives_scope_from_vault_host(#[case] uri: &str, #[case] expected: &str) {
let manager = AzureKeyVault::with_client(
reqwest::Client::new(),
uri.parse().unwrap(),
Arc::new(|_: &str| None),
)
.unwrap();
assert_eq!(manager.scope(), expected);
}
#[tokio::test]
async fn missing_credentials_do_not_request_vault() {
let server = MockServer::start().await;
Mock::given(path("/secrets/NAME"))
.respond_with(ResponseTemplate::new(200))
.expect(0)
.mount(&server)
.await;
assert!(
manager_without_credentials(&server)
.get_secret_from_azure_key_vault("NAME")
.await
.is_err()
);
}
fn manager_without_credentials(server: &MockServer) -> AzureKeyVault {
AzureKeyVault::with_client(
reqwest::Client::new(),
server.uri().parse().unwrap(),
Arc::new(|name: &str| {
(name == "AZURE_CREDENTIAL").then(|| "ClientSecretCredential".to_owned())
}),
)
.unwrap()
}
#[derive(Deserialize)]
struct Fixture {
cases: Vec<FixtureCase>,
}
#[derive(Deserialize)]
struct FixtureCase {
secret_name: String,
response: FixtureResponse,
expected: FixtureExpected,
}
#[derive(Deserialize)]
struct FixtureResponse {
status: u16,
body: serde_json::Value,
}
#[derive(Deserialize)]
struct FixtureExpected {
value: Option<String>,
missing: Option<bool>,
error: Option<bool>,
}
#[tokio::test]
async fn parity_fixture_matches_python_backend_contract() {
let fixture: Fixture =
serde_json::from_str(include_str!("fixtures/key_vault_parity.json")).unwrap();
for case in fixture.cases {
let server = MockServer::start().await;
Mock::given(path(format!("/secrets/{}", case.secret_name)))
.respond_with(
ResponseTemplate::new(case.response.status).set_body_json(case.response.body),
)
.expect(1)
.mount(&server)
.await;
let result = manager(&server)
.get_secret_from_azure_key_vault(&case.secret_name)
.await;
if case.expected.missing == Some(true) {
assert_eq!(result.unwrap(), None);
} else if case.expected.error == Some(true) {
assert!(result.is_err());
} else {
assert_eq!(
result.unwrap().unwrap().as_str(),
case.expected.value.as_deref()
);
}
}
}

30
litellm-rust/crates/secrets-azure/tests/live.rs Normal file
View file

@ -0,0 +1,30 @@
use std::sync::Arc;
use litellm_core_utils::settings::ProcessEnvironment;
use litellm_secrets_azure::AzureKeyVault;
use litellm_secrets_types::Secret;
#[tokio::test]
#[ignore]
async fn reads_a_real_secret() {
let environment = Arc::new(ProcessEnvironment);
let manager = AzureKeyVault::new(environment).unwrap();
let name = std::env::var("AZURE_KEY_VAULT_LIVE_SECRET_NAME").unwrap();
let secret = manager
.get_secret_from_azure_key_vault(&name)
.await
.unwrap()
.unwrap();
assert!(matches!(&secret, Secret::String(_)));
let host = std::env::var("AZURE_KEY_VAULT_URI")
.unwrap()
.parse::<reqwest::Url>()
.unwrap()
.host_str()
.unwrap()
.to_owned();
let value_len = secret.as_str().unwrap().len();
println!(
"native provider=litellm-secrets-azure vault_host={host} secret={name} value_len={value_len}"
);
}

26
litellm-rust/crates/secrets-cyberark/Cargo.toml Normal file
View file

@ -0,0 +1,26 @@
[package]
name = "litellm-secrets-cyberark"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
litellm-secrets-types.workspace = true
litellm-core-utils.workspace = true
base64.workspace = true
moka.workspace = true
reqwest.workspace = true
serde_json.workspace = true
thiserror.workspace = true
veil.workspace = true
tracing = "0.1"
percent-encoding = "2.3"
tokio = { workspace = true, features = ["sync"] }
[dev-dependencies]
rstest.workspace = true
tokio.workspace = true
wiremock = "0.6.5"
serde.workspace = true
serde_json.workspace = true

27
litellm-rust/crates/secrets-cyberark/src/error.rs Normal file
View file

@ -0,0 +1,27 @@
#[derive(thiserror::Error, veil::Redact)]
pub enum Error {
#[error("CyberArk Conjur HTTP request failed")]
Http(
#[from]
#[redact]
reqwest::Error,
),
#[error("CyberArk Conjur authentication returned HTTP {0}")]
AuthStatus(u16),
#[error("CyberArk Conjur returned HTTP {0}")]
Status(u16),
#[error(
"CyberArk credentials are missing: set CYBERARK_API_KEY or both CYBERARK_CLIENT_CERT and CYBERARK_CLIENT_KEY"
)]
MissingCredentials,
#[error("CyberArk client certificate could not be loaded")]
ClientCertificate,
#[error("invalid refresh interval")]
RefreshInterval,
#[error("invalid CyberArk Conjur endpoint")]
Endpoint,
#[error("CyberArk secret manager requires an enterprise license")]
EnterpriseRequired,
#[error(transparent)]
Operation(#[from] litellm_secrets_types::Error),
}

7
litellm-rust/crates/secrets-cyberark/src/lib.rs Normal file
View file

@ -0,0 +1,7 @@
#![forbid(unsafe_code)]
mod error;
mod secret_manager;
pub use error::Error;
pub use secret_manager::{CyberArkSecretManager, DeleteOutcome};

317
litellm-rust/crates/secrets-cyberark/src/secret_manager.rs Normal file
View file

@ -0,0 +1,317 @@
use std::{fs, sync::Arc, time::Duration};
use base64::{Engine, engine::general_purpose::STANDARD};
use litellm_core_utils::settings::Lookup;
use litellm_secrets_types::{BaseSecretManager, SecretValue, validate_secret_name};
use moka::future::Cache;
use percent_encoding::{AsciiSet, NON_ALPHANUMERIC, utf8_percent_encode};
use crate::Error;
const CYBERARK_API_BASE: &str = "CYBERARK_API_BASE";
const CYBERARK_ACCOUNT: &str = "CYBERARK_ACCOUNT";
const CYBERARK_USERNAME: &str = "CYBERARK_USERNAME";
const CYBERARK_API_KEY: &str = "CYBERARK_API_KEY";
const CYBERARK_CLIENT_CERT: &str = "CYBERARK_CLIENT_CERT";
const CYBERARK_CLIENT_KEY: &str = "CYBERARK_CLIENT_KEY";
const CYBERARK_SSL_VERIFY: &str = "CYBERARK_SSL_VERIFY";
const CYBERARK_REFRESH_INTERVAL: &str = "CYBERARK_REFRESH_INTERVAL";
const DEFAULT_API_BASE: &str = "http://127.0.0.1:8080";
const DEFAULT_ACCOUNT: &str = "default";
const DEFAULT_USERNAME: &str = "admin";
const DEFAULT_REFRESH_INTERVAL: Duration = Duration::from_secs(300);
const SECRET_NAME_SAFE: &AsciiSet = &NON_ALPHANUMERIC
.remove(b'-')
.remove(b'_')
.remove(b'.')
.remove(b'~');
#[derive(Clone)]
pub struct CyberArkSecretManager {
client: reqwest::Client,
endpoint: reqwest::Url,
account: String,
username: String,
api_key: SecretValue,
token: Cache<(), SecretValue>,
secrets: Cache<String, SecretValue>,
authentication_lock: Arc<tokio::sync::Mutex<()>>,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum DeleteOutcome {
NotSupported,
}
impl CyberArkSecretManager {
pub fn with_client(
client: reqwest::Client,
endpoint: reqwest::Url,
account: String,
username: String,
api_key: SecretValue,
refresh_interval: Option<Duration>,
) -> Self {
let endpoint = normalize_endpoint(endpoint);
let ttl = refresh_interval
.filter(|interval| !interval.is_zero())
.unwrap_or(DEFAULT_REFRESH_INTERVAL);
let token = Cache::builder().time_to_live(ttl).build();
let secrets = Cache::builder().time_to_live(ttl).build();
Self {
client,
endpoint,
account,
username,
api_key,
token,
secrets,
authentication_lock: Arc::new(tokio::sync::Mutex::new(())),
}
}
pub fn new(
environment: Arc<dyn Lookup + Send + Sync>,
enterprise_enabled: bool,
) -> Result<Self, Error> {
let api_key = environment.get(CYBERARK_API_KEY).unwrap_or_default();
let cert = environment.get(CYBERARK_CLIENT_CERT).unwrap_or_default();
let key = environment.get(CYBERARK_CLIENT_KEY).unwrap_or_default();
if api_key.is_empty() && (cert.is_empty() || key.is_empty()) {
return Err(Error::MissingCredentials);
}
if !enterprise_enabled {
return Err(Error::EnterpriseRequired);
}
let verify = environment
.get(CYBERARK_SSL_VERIFY)
.map(|value| !value.trim().eq_ignore_ascii_case("false"))
.unwrap_or(true);
let mut builder = reqwest::Client::builder();
if !verify {
tracing::warn!(
"CyberArk SSL verification is disabled. This is insecure and should only be used for testing with self-signed certificates."
);
builder = builder.danger_accept_invalid_certs(true);
}
if !cert.is_empty() && !key.is_empty() {
let certificate = fs::read(cert).map_err(|_| Error::ClientCertificate)?;
let private_key = fs::read(key).map_err(|_| Error::ClientCertificate)?;
let identity = reqwest::Identity::from_pem(&[certificate, private_key].concat())
.map_err(|_| Error::ClientCertificate)?;
builder = builder.identity(identity);
}
let client = builder.build()?;
let endpoint = reqwest::Url::parse(
&environment
.get(CYBERARK_API_BASE)
.unwrap_or_else(|| DEFAULT_API_BASE.to_owned()),
)
.map_err(|_| Error::Endpoint)?;
let account = environment
.get(CYBERARK_ACCOUNT)
.unwrap_or_else(|| DEFAULT_ACCOUNT.to_owned());
let username = environment
.get(CYBERARK_USERNAME)
.unwrap_or_else(|| DEFAULT_USERNAME.to_owned());
let refresh_interval = environment
.get(CYBERARK_REFRESH_INTERVAL)
.map(|value| {
value
.parse::<u64>()
.map(Duration::from_secs)
.map_err(|_| Error::RefreshInterval)
})
.transpose()?;
Ok(Self::with_client(
client,
endpoint,
account,
username,
SecretValue::new(api_key),
refresh_interval,
))
}
fn secret_url(&self, name: &str) -> Result<reqwest::Url, Error> {
let encoded = utf8_percent_encode(name, SECRET_NAME_SAFE);
self.endpoint
.join(&format!("secrets/{}/variable/{}", self.account, encoded))
.map_err(|_| Error::Endpoint)
}
async fn authenticate(&self) -> Result<SecretValue, Error> {
if let Some(token) = self.token.get(&()).await {
return Ok(token);
}
let _guard = self.authentication_lock.lock().await;
if let Some(token) = self.token.get(&()).await {
return Ok(token);
}
let url = self
.endpoint
.join(&format!(
"authn/{}/{}/authenticate",
self.account, self.username
))
.map_err(|_| Error::Endpoint)?;
let response = self
.client
.post(url)
.body(self.api_key.expose().to_owned())
.send()
.await?;
if !response.status().is_success() {
return Err(Error::AuthStatus(response.status().as_u16()));
}
let token = SecretValue::new(STANDARD.encode(response.text().await?));
self.token.insert((), token.clone()).await;
Ok(token)
}
async fn authorization_header(&self) -> Result<String, Error> {
Ok(format!(
"Token token=\"{}\"",
self.authenticate().await?.expose()
))
}
pub async fn async_read_secret(&self, name: &str) -> Result<Option<SecretValue>, Error> {
if let Some(value) = self.secrets.get(name).await {
return Ok(Some(value));
}
let response = self
.client
.get(self.secret_url(name)?)
.header("Authorization", self.authorization_header().await?)
.send()
.await?;
if response.status() == reqwest::StatusCode::NOT_FOUND {
return Ok(None);
}
if !response.status().is_success() {
return Err(Error::Status(response.status().as_u16()));
}
let value = SecretValue::new(response.text().await?);
self.secrets.insert(name.to_owned(), value.clone()).await;
Ok(Some(value))
}
pub async fn async_write_secret(
&self,
name: &str,
value: &SecretValue,
_description: Option<&str>,
) -> Result<(), Error> {
validate_secret_name(name)?;
self.ensure_variable_exists(name).await;
let response = self
.client
.post(self.secret_url(name)?)
.header("Authorization", self.authorization_header().await?)
.body(value.expose().to_owned())
.send()
.await?;
if !response.status().is_success() {
return Err(Error::Status(response.status().as_u16()));
}
self.secrets.insert(name.to_owned(), value.clone()).await;
Ok(())
}
async fn ensure_variable_exists(&self, name: &str) {
let policy_url = self
.endpoint
.join(&format!("policies/{}/policy/root", self.account));
let Ok(policy_url) = policy_url else {
tracing::warn!("Could not build CyberArk policy endpoint");
return;
};
let Ok(authorization) = self.authorization_header().await else {
tracing::warn!("Could not authenticate while ensuring CyberArk variable exists");
return;
};
let body = format!(
"- !variable {}\n",
serde_json::to_string(name).expect("serializing a string cannot fail")
);
let response = self
.client
.post(policy_url)
.header("Authorization", authorization)
.header("Content-Type", "application/x-yaml")
.body(body)
.send()
.await;
match response {
Ok(response) if response.status().is_success() => {}
Ok(response)
if matches!(
response.status(),
reqwest::StatusCode::CONFLICT | reqwest::StatusCode::UNPROCESSABLE_ENTITY
) =>
{
tracing::debug!(
"CyberArk variable policy already exists or conflicts: {}",
response.status()
);
}
Ok(response) => {
tracing::warn!(
"Could not ensure CyberArk variable exists: {}",
response.status()
);
}
Err(error) => {
tracing::warn!("Error ensuring CyberArk variable exists: {error}");
}
}
}
pub async fn async_delete_secret(
&self,
name: &str,
_recovery_window_in_days: i64,
) -> Result<DeleteOutcome, Error> {
tracing::warn!(
"CyberArk Conjur does not support direct secret deletion. Secrets must be removed through policy updates."
);
self.secrets.invalidate(name).await;
Ok(DeleteOutcome::NotSupported)
}
}
impl BaseSecretManager for CyberArkSecretManager {
type Error = Error;
type WriteResponse = ();
type DeleteResponse = DeleteOutcome;
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<(), Error> {
self.async_write_secret(name, value, description).await
}
async fn async_delete_secret(
&self,
name: &str,
recovery_window_in_days: i64,
) -> Result<DeleteOutcome, Error> {
self.async_delete_secret(name, recovery_window_in_days)
.await
}
}
fn normalize_endpoint(mut endpoint: reqwest::Url) -> reqwest::Url {
if !endpoint.path().ends_with('/') {
endpoint.set_path(&format!("{}/", endpoint.path()));
}
endpoint
}

32
litellm-rust/crates/secrets-cyberark/tests/fixtures/parity.json vendored Normal file
View file

@ -0,0 +1,32 @@
{
"endpoint": "http://conjur.test:8080",
"account": "acct",
"username": "admin",
"api_key": "k3y",
"authenticate_path": "/authn/acct/admin/authenticate",
"token_json": "{\"protected\":\"p\",\"payload\":\"q\",\"signature\":\"s\"}",
"authorization_header": "Token token=\"eyJwcm90ZWN0ZWQiOiJwIiwicGF5bG9hZCI6InEiLCJzaWduYXR1cmUiOiJzIn0=\"",
"policy_path": "/policies/acct/policy/root",
"secrets": [
{
"name": "OPENAI_API_KEY",
"path": "/secrets/acct/variable/OPENAI_API_KEY",
"policy_body": "- !variable \"OPENAI_API_KEY\"\n"
},
{
"name": "team/app/key",
"path": "/secrets/acct/variable/team%2Fapp%2Fkey",
"policy_body": "- !variable \"team/app/key\"\n"
},
{
"name": "a b+c.d-e_f~g",
"path": "/secrets/acct/variable/a%20b%2Bc.d-e_f~g",
"policy_body": "- !variable \"a b+c.d-e_f~g\"\n"
},
{
"name": "needs \"quote\"",
"path": "/secrets/acct/variable/needs%20%22quote%22",
"policy_body": "- !variable \"needs \\\"quote\\\"\"\n"
}
]
}

516
litellm-rust/crates/secrets-cyberark/tests/secret_manager.rs Normal file
View file

@ -0,0 +1,516 @@
use std::{sync::Arc, time::Duration};
use base64::{Engine, engine::general_purpose::STANDARD};
use litellm_secrets_cyberark::{CyberArkSecretManager, DeleteOutcome, Error};
use litellm_secrets_types::SecretValue;
use serde::Deserialize;
use wiremock::{
Match, Mock, MockServer, Request, ResponseTemplate,
matchers::{body_string, header, method, path},
};
const TOKEN_JSON: &str = r#"{"protected":"p","payload":"q","signature":"s"}"#;
#[derive(Deserialize)]
struct ParityFixture {
endpoint: String,
account: String,
username: String,
api_key: String,
authenticate_path: String,
token_json: String,
authorization_header: String,
policy_path: String,
secrets: Vec<ParitySecret>,
}
#[derive(Deserialize)]
struct ParitySecret {
name: String,
path: String,
policy_body: String,
}
#[derive(Debug)]
struct RawPath(String);
impl Match for RawPath {
fn matches(&self, request: &Request) -> bool {
request.url.path() == self.0
}
}
fn fixture() -> ParityFixture {
serde_json::from_str(include_str!("fixtures/parity.json")).unwrap()
}
fn manager(server: &MockServer, ttl: Duration) -> CyberArkSecretManager {
CyberArkSecretManager::with_client(
reqwest::Client::new(),
server.uri().parse().unwrap(),
"acct".into(),
"admin".into(),
SecretValue::new("k3y"),
Some(ttl),
)
}
async fn mount_auth(server: &MockServer, expected: u64) {
Mock::given(method("POST"))
.and(path("/authn/acct/admin/authenticate"))
.and(body_string("k3y"))
.respond_with(ResponseTemplate::new(200).set_body_string(TOKEN_JSON))
.expect(expected)
.mount(server)
.await;
}
#[tokio::test]
async fn successful_reads_cache_auth_secret_and_redact_values() {
let server = MockServer::start().await;
mount_auth(&server, 1).await;
let token = STANDARD.encode(TOKEN_JSON);
Mock::given(path("/secrets/acct/variable/OPENAI_API_KEY"))
.and(header("authorization", format!("Token token=\"{token}\"")))
.respond_with(ResponseTemplate::new(200).set_body_string("sk-live"))
.expect(1)
.mount(&server)
.await;
let manager = manager(&server, Duration::from_secs(60));
for _ in 0..2 {
let value = manager
.async_read_secret("OPENAI_API_KEY")
.await
.unwrap()
.unwrap();
assert_eq!(value.expose(), "sk-live");
assert!(!format!("{value:?}").contains("sk-live"));
}
}
#[tokio::test]
async fn concurrent_reads_share_authentication_request() {
let server = MockServer::start().await;
Mock::given(path("/authn/acct/admin/authenticate"))
.and(body_string("k3y"))
.respond_with(
ResponseTemplate::new(200)
.set_body_string(TOKEN_JSON)
.set_delay(Duration::from_millis(20)),
)
.expect(1)
.mount(&server)
.await;
Mock::given(path("/secrets/acct/variable/key"))
.and(header(
"authorization",
format!("Token token=\"{}\"", STANDARD.encode(TOKEN_JSON)),
))
.respond_with(ResponseTemplate::new(200).set_body_string("value"))
.expect(2)
.mount(&server)
.await;
let manager = manager(&server, Duration::from_secs(60));
let (first, second) = tokio::join!(
manager.async_read_secret("key"),
manager.async_read_secret("key")
);
assert_eq!(first.unwrap().unwrap().expose(), "value");
assert_eq!(second.unwrap().unwrap().expose(), "value");
}
#[rstest::rstest]
#[case::not_found(404)]
#[case::unauthorized(401)]
#[case::forbidden(403)]
#[case::server_error(500)]
#[tokio::test]
async fn failed_reads_are_not_cached(#[case] status: u16) {
let server = MockServer::start().await;
mount_auth(&server, 1).await;
let failing = Mock::given(path("/secrets/acct/variable/key"))
.respond_with(ResponseTemplate::new(status))
.expect(1)
.mount_as_scoped(&server)
.await;
let manager = manager(&server, Duration::from_secs(60));
let result = manager.async_read_secret("key").await;
if status == 404 {
assert_eq!(result.unwrap(), None);
} else {
assert!(matches!(result, Err(Error::Status(actual)) if actual == status));
}
drop(failing);
Mock::given(path("/secrets/acct/variable/key"))
.respond_with(ResponseTemplate::new(200).set_body_string("recovered"))
.expect(1)
.mount(&server)
.await;
for _ in 0..2 {
assert_eq!(
manager
.async_read_secret("key")
.await
.unwrap()
.unwrap()
.expose(),
"recovered"
);
}
}
#[tokio::test]
async fn failed_authentication_is_not_cached_and_does_not_read_secret() {
let server = MockServer::start().await;
let failing = Mock::given(path("/authn/acct/admin/authenticate"))
.respond_with(ResponseTemplate::new(401))
.expect(1)
.mount_as_scoped(&server)
.await;
let unused_secret = Mock::given(path("/secrets/acct/variable/key"))
.respond_with(ResponseTemplate::new(200).set_body_string("value"))
.expect(0)
.mount_as_scoped(&server)
.await;
let manager = manager(&server, Duration::from_secs(60));
assert!(matches!(
manager.async_read_secret("key").await,
Err(Error::AuthStatus(401))
));
drop(unused_secret);
drop(failing);
mount_auth(&server, 1).await;
Mock::given(path("/secrets/acct/variable/key"))
.respond_with(ResponseTemplate::new(200).set_body_string("value"))
.expect(1)
.mount(&server)
.await;
assert_eq!(
manager
.async_read_secret("key")
.await
.unwrap()
.unwrap()
.expose(),
"value"
);
}
#[tokio::test]
async fn expired_tokens_and_secrets_are_fetched_again() {
let server = MockServer::start().await;
mount_auth(&server, 2).await;
Mock::given(path("/secrets/acct/variable/key"))
.respond_with(ResponseTemplate::new(200).set_body_string("value"))
.expect(2)
.mount(&server)
.await;
let manager = manager(&server, Duration::from_millis(1));
for _ in 0..2 {
assert!(manager.async_read_secret("key").await.unwrap().is_some());
tokio::time::sleep(Duration::from_millis(5)).await;
}
}
#[rstest::rstest]
#[tokio::test]
async fn secret_names_use_python_quote_encoding(
#[values("OPENAI_API_KEY", "team/app/key", "a b+c.d-e_f~g", "needs \"quote\"")] name: &str,
) {
let fixture = fixture();
let secret = fixture
.secrets
.iter()
.find(|secret| secret.name == name)
.unwrap();
let server = MockServer::start().await;
mount_auth(&server, 1).await;
Mock::given(RawPath(secret.path.clone()))
.respond_with(ResponseTemplate::new(200).set_body_string("value"))
.expect(1)
.mount(&server)
.await;
assert_eq!(
manager(&server, Duration::from_secs(60))
.async_read_secret(name)
.await
.unwrap()
.unwrap()
.expose(),
"value"
);
}
#[rstest::rstest]
#[case(201)]
#[case(409)]
#[case(422)]
#[case(500)]
#[tokio::test]
async fn writes_tolerate_policy_status_and_cache_value(#[case] policy_status: u16) {
let server = MockServer::start().await;
mount_auth(&server, 1).await;
Mock::given(path("/policies/acct/policy/root"))
.and(header("content-type", "application/x-yaml"))
.and(body_string("- !variable \"team/app\"\n"))
.respond_with(ResponseTemplate::new(policy_status))
.expect(1)
.mount(&server)
.await;
Mock::given(path("/secrets/acct/variable/team%2Fapp"))
.and(body_string("v"))
.respond_with(ResponseTemplate::new(200))
.expect(1)
.mount(&server)
.await;
let manager = manager(&server, Duration::from_secs(60));
manager
.async_write_secret("team/app", &SecretValue::new("v"), None)
.await
.unwrap();
assert_eq!(
manager
.async_read_secret("team/app")
.await
.unwrap()
.unwrap()
.expose(),
"v"
);
}
#[tokio::test]
async fn failed_value_write_is_not_cached() {
let server = MockServer::start().await;
mount_auth(&server, 1).await;
Mock::given(path("/policies/acct/policy/root"))
.respond_with(ResponseTemplate::new(409))
.mount(&server)
.await;
Mock::given(path("/secrets/acct/variable/key"))
.and(body_string("v"))
.respond_with(ResponseTemplate::new(403))
.expect(1)
.mount(&server)
.await;
Mock::given(path("/secrets/acct/variable/key"))
.respond_with(ResponseTemplate::new(200).set_body_string("recovered"))
.expect(1)
.mount(&server)
.await;
let manager = manager(&server, Duration::from_secs(60));
assert!(matches!(
manager
.async_write_secret("key", &SecretValue::new("v"), None)
.await,
Err(Error::Status(403))
));
assert_eq!(
manager
.async_read_secret("key")
.await
.unwrap()
.unwrap()
.expose(),
"recovered"
);
}
#[tokio::test]
async fn unsafe_names_fail_before_http_calls() {
let server = MockServer::start().await;
let manager = manager(&server, Duration::from_secs(60));
assert!(matches!(
manager
.async_write_secret("../etc", &SecretValue::new("v"), None)
.await,
Err(Error::Operation(
litellm_secrets_types::Error::UnsafeSecretName
))
));
}
#[tokio::test]
async fn delete_invalidates_cache_and_reports_not_supported() {
let server = MockServer::start().await;
mount_auth(&server, 1).await;
Mock::given(path("/secrets/acct/variable/key"))
.respond_with(ResponseTemplate::new(200).set_body_string("v"))
.expect(2)
.mount(&server)
.await;
let manager = manager(&server, Duration::from_secs(60));
assert_eq!(
manager
.async_read_secret("key")
.await
.unwrap()
.unwrap()
.expose(),
"v"
);
assert_eq!(
manager.async_delete_secret("key", 7).await.unwrap(),
DeleteOutcome::NotSupported
);
assert_eq!(
manager
.async_read_secret("key")
.await
.unwrap()
.unwrap()
.expose(),
"v"
);
}
#[test]
fn new_validates_credentials_before_license_and_configuration() {
let empty: Arc<dyn litellm_core_utils::settings::Lookup + Send + Sync> =
Arc::new(|_: &str| None);
assert!(matches!(
CyberArkSecretManager::new(empty, true),
Err(Error::MissingCredentials)
));
assert!(matches!(
CyberArkSecretManager::new(
Arc::new(|name: &str| (name == "CYBERARK_API_KEY").then(|| "k3y".into())),
false
),
Err(Error::EnterpriseRequired)
));
assert!(matches!(
CyberArkSecretManager::new(
Arc::new(|name: &str| (name == "CYBERARK_CLIENT_CERT").then(|| "cert".into())),
true
),
Err(Error::MissingCredentials)
));
assert!(matches!(
CyberArkSecretManager::new(
Arc::new(|name: &str| match name {
"CYBERARK_API_KEY" => Some("k3y".into()),
"CYBERARK_REFRESH_INTERVAL" => Some("abc".into()),
_ => None,
}),
true
),
Err(Error::RefreshInterval)
));
assert!(matches!(
CyberArkSecretManager::new(
Arc::new(|name: &str| match name {
"CYBERARK_API_KEY" => Some("k3y".into()),
"CYBERARK_API_BASE" => Some("not a url".into()),
_ => None,
}),
true
),
Err(Error::Endpoint)
));
}
#[tokio::test]
async fn new_reads_environment_defaults_end_to_end() {
let server = MockServer::start().await;
Mock::given(path("/authn/default/admin/authenticate"))
.and(body_string("k3y"))
.respond_with(ResponseTemplate::new(200).set_body_string(TOKEN_JSON))
.mount(&server)
.await;
Mock::given(path("/secrets/default/variable/key"))
.respond_with(ResponseTemplate::new(200).set_body_string("value"))
.mount(&server)
.await;
let endpoint = server.uri();
let manager = CyberArkSecretManager::new(
Arc::new(move |name: &str| match name {
"CYBERARK_API_BASE" => Some(endpoint.clone()),
"CYBERARK_API_KEY" => Some("k3y".into()),
_ => None,
}),
true,
)
.unwrap();
assert_eq!(
manager
.async_read_secret("key")
.await
.unwrap()
.unwrap()
.expose(),
"value"
);
}
#[test]
fn new_reports_missing_client_certificate_files() {
assert!(matches!(
CyberArkSecretManager::new(
Arc::new(|name: &str| match name {
"CYBERARK_CLIENT_CERT" => Some("/missing/cert".into()),
"CYBERARK_CLIENT_KEY" => Some("/missing/key".into()),
_ => None,
}),
true
),
Err(Error::ClientCertificate)
));
}
#[tokio::test]
async fn trailing_slash_endpoint_preserves_base_path() {
let server = MockServer::start().await;
Mock::given(path("/prefix/authn/acct/admin/authenticate"))
.and(body_string("k3y"))
.respond_with(ResponseTemplate::new(200).set_body_string(TOKEN_JSON))
.expect(1)
.mount(&server)
.await;
Mock::given(path("/prefix/secrets/acct/variable/key"))
.respond_with(ResponseTemplate::new(200).set_body_string("value"))
.mount(&server)
.await;
let endpoint = format!("{}/prefix/", server.uri()).parse().unwrap();
let manager = CyberArkSecretManager::with_client(
reqwest::Client::new(),
endpoint,
"acct".into(),
"admin".into(),
SecretValue::new("k3y"),
Some(Duration::from_secs(60)),
);
assert_eq!(
manager
.async_read_secret("key")
.await
.unwrap()
.unwrap()
.expose(),
"value"
);
}
#[test]
fn parity_fixture_matches_authentication_contract() {
let fixture = fixture();
assert_eq!(fixture.endpoint, "http://conjur.test:8080");
assert_eq!(fixture.account, "acct");
assert_eq!(fixture.username, "admin");
assert_eq!(fixture.api_key, "k3y");
assert_eq!(fixture.authenticate_path, "/authn/acct/admin/authenticate");
assert_eq!(fixture.token_json, TOKEN_JSON);
assert_eq!(
fixture.authorization_header,
format!("Token token=\"{}\"", STANDARD.encode(TOKEN_JSON))
);
assert_eq!(fixture.policy_path, "/policies/acct/policy/root");
assert_eq!(fixture.secrets.len(), 4);
assert_eq!(
fixture.secrets[1].policy_body,
"- !variable \"team/app/key\"\n"
);
}

4
litellm-rust/crates/secrets/Cargo.toml
View file

@ -9,11 +9,15 @@ repository.workspace = true
default = []
aws = ["dep:litellm-secrets-aws"]
google = ["dep:litellm-secrets-google"]
azure = ["dep:litellm-secrets-azure"]
cyberark = ["dep:litellm-secrets-cyberark"]
[dependencies]
litellm-secrets-types.workspace = true
litellm-secrets-aws = { workspace = true, optional = true }
litellm-secrets-google = { workspace = true, optional = true }
litellm-secrets-azure = { workspace = true, optional = true }
litellm-secrets-cyberark = { workspace = true, optional = true }
litellm-core-utils.workspace = true
base64.workspace = true
serde.workspace = true

6
litellm-rust/crates/secrets/src/error.rs
View file

@ -30,4 +30,10 @@ pub enum Error {
#[cfg(feature = "google")]
#[error(transparent)]
Google(#[from] litellm_secrets_google::Error),
#[cfg(feature = "azure")]
#[error(transparent)]
Azure(#[from] litellm_secrets_azure::Error),
#[cfg(feature = "cyberark")]
#[error(transparent)]
Cyberark(#[from] litellm_secrets_cyberark::Error),
}

19
litellm-rust/crates/secrets/src/handler.rs
View file

@ -13,6 +13,10 @@ pub enum SecretManager {
GoogleKms(crate::google::GoogleKms),
#[cfg(feature = "google")]
GoogleSecretManager(crate::google::GoogleSecretManager),
#[cfg(feature = "azure")]
AzureKeyVault(crate::azure::AzureKeyVault),
#[cfg(feature = "cyberark")]
Cyberark(crate::cyberark::CyberArkSecretManager),
}
impl SecretManager {
@ -27,6 +31,10 @@ impl SecretManager {
Self::GoogleKms(_) => KeyManagementSystem::GoogleKms,
#[cfg(feature = "google")]
Self::GoogleSecretManager(_) => KeyManagementSystem::GoogleSecretManager,
#[cfg(feature = "azure")]
Self::AzureKeyVault(_) => KeyManagementSystem::AzureKeyVault,
#[cfg(feature = "cyberark")]
Self::Cyberark(_) => KeyManagementSystem::Cyberark,
}
}
}
@ -78,6 +86,17 @@ pub async fn get_secret_from_manager(
.get_secret_from_google_secret_manager(secret_name)
.await
.map_err(Error::from),
#[cfg(feature = "azure")]
SecretManager::AzureKeyVault(client) => client
.get_secret_from_azure_key_vault(secret_name)
.await
.map_err(Error::from),
#[cfg(feature = "cyberark")]
SecretManager::Cyberark(client) => client
.async_read_secret(secret_name)
.await
.map(|value| value.map(Secret::String))
.map_err(Error::from),
}
}

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