mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
chore: merge main to restore lint checker
This commit is contained in:
commit
21fb652586
598 changed files with 53320 additions and 8418 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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
7
.github/scripts/verify_linux_native_wheel.py
vendored
|
|
@ -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
4
.github/workflows/test-linting.yml
vendored
|
|
@ -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
2
.github/workflows/test-rust.yml
vendored
|
|
@ -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 hashicorp azure cyberark aws,google aws,azure google,azure aws,google,azure aws,google,cyberark aws,google,azure,cyberark aws,google,hashicorp,azure,cyberark; do
|
||||
cargo test -p litellm-secrets --locked --no-default-features --features "$features"
|
||||
done
|
||||
|
||||
|
|
|
|||
1
Makefile
1
Makefile
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
|
|
|
|||
|
|
@ -2,6 +2,17 @@
|
|||
|
||||
This guide provides instructions for building and running the LiteLLM application using Docker and Docker Compose.
|
||||
|
||||
> **Just want to run LiteLLM?** This guide builds from source. To run the published
|
||||
> image instead, use `docker-compose.quickstart.yml` in this directory — the
|
||||
> two-service stack (gateway + Postgres) that the
|
||||
> [Docker quickstart](https://docs.litellm.ai/docs/proxy/docker_quick_start) documents:
|
||||
>
|
||||
> ```bash
|
||||
> curl -sSLO https://github.com/BerriAI/litellm/raw/main/docker/docker-compose.quickstart.yml
|
||||
> printf 'LITELLM_MASTER_KEY=sk-%s\nLITELLM_SALT_KEY=sk-%s\n' "$(openssl rand -hex 32)" "$(openssl rand -hex 32)" > .env
|
||||
> docker compose -f docker-compose.quickstart.yml up -d
|
||||
> ```
|
||||
|
||||
## Prerequisites
|
||||
|
||||
- Docker
|
||||
|
|
|
|||
41
docker/docker-compose.quickstart.yml
Normal file
41
docker/docker-compose.quickstart.yml
Normal file
|
|
@ -0,0 +1,41 @@
|
|||
# LiteLLM quickstart stack: the gateway plus a Postgres database that stores
|
||||
# models, virtual keys, and spend logs. Used by
|
||||
# https://docs.litellm.ai/docs/proxy/docker_quick_start
|
||||
#
|
||||
# curl -sSLO https://github.com/BerriAI/litellm/raw/main/docker/docker-compose.quickstart.yml
|
||||
# printf 'LITELLM_MASTER_KEY=sk-%s\nLITELLM_SALT_KEY=sk-%s\n' "$(openssl rand -hex 32)" "$(openssl rand -hex 32)" > .env
|
||||
# docker compose -f docker-compose.quickstart.yml up -d
|
||||
#
|
||||
# Compose reads .env from this directory. Keep it: regenerating LITELLM_SALT_KEY
|
||||
# makes credentials already stored in the database unreadable. For anything
|
||||
# beyond local evaluation, pin the image to a specific release tag.
|
||||
services:
|
||||
litellm:
|
||||
image: docker.litellm.ai/berriai/litellm:main-stable
|
||||
ports:
|
||||
- "4000:4000"
|
||||
environment:
|
||||
LITELLM_MASTER_KEY: ${LITELLM_MASTER_KEY:?set it in .env - see the header of this file}
|
||||
LITELLM_SALT_KEY: ${LITELLM_SALT_KEY:?set it in .env - see the header of this file}
|
||||
DATABASE_URL: postgresql://litellm:litellm@db:5432/litellm
|
||||
STORE_MODEL_IN_DB: "True"
|
||||
depends_on:
|
||||
db:
|
||||
condition: service_healthy
|
||||
|
||||
db:
|
||||
image: postgres:16
|
||||
environment:
|
||||
POSTGRES_USER: litellm
|
||||
POSTGRES_PASSWORD: litellm
|
||||
POSTGRES_DB: litellm
|
||||
healthcheck:
|
||||
test: ["CMD-SHELL", "pg_isready -U litellm"]
|
||||
interval: 5s
|
||||
timeout: 5s
|
||||
retries: 10
|
||||
volumes:
|
||||
- postgres_data:/var/lib/postgresql/data
|
||||
|
||||
volumes:
|
||||
postgres_data:
|
||||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ same route are non-inference and free.
|
|||
"""
|
||||
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import TYPE_CHECKING, Dict, Optional, cast
|
||||
from typing import TYPE_CHECKING, Dict, Final, Optional, Protocol, cast
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
|
@ -22,11 +22,31 @@ from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN
|
|||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
||||
from litellm.repositories.prisma_protocols import TableActions
|
||||
from litellm.router import Router
|
||||
|
||||
TERMINAL_RESPONSE_STATUSES = frozenset({"completed", "failed", "cancelled", "incomplete"})
|
||||
|
||||
|
||||
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
|
||||
|
||||
|
||||
class CheckResponsesCost:
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -128,7 +148,7 @@ class CheckResponsesCost:
|
|||
f"CheckResponsesCost: stale cleanup failed (poll will continue): {cleanup_err}"
|
||||
)
|
||||
|
||||
jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many(
|
||||
jobs = await _managed_object_table(self.prisma_client).find_many(
|
||||
where={
|
||||
"status": {"in": ["queued", "in_progress"]},
|
||||
"file_purpose": "response",
|
||||
|
|
@ -138,7 +158,7 @@ class CheckResponsesCost:
|
|||
)
|
||||
|
||||
verbose_proxy_logger.debug(f"Found {len(jobs)} response jobs to check")
|
||||
completed_jobs = []
|
||||
completed_jobs: Final[list[_ManagedObjectRow]] = []
|
||||
|
||||
for job in jobs:
|
||||
unified_object_id = job.unified_object_id
|
||||
|
|
@ -189,7 +209,7 @@ class CheckResponsesCost:
|
|||
|
||||
# Mark completed jobs in the database
|
||||
if len(completed_jobs) > 0:
|
||||
await self.prisma_client.db.litellm_managedobjecttable.update_many(
|
||||
await _managed_object_table(self.prisma_client).update_many(
|
||||
where={"id": {"in": [job.id for job in completed_jobs]}},
|
||||
data={"status": "completed"},
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -0,0 +1,42 @@
|
|||
CREATE TABLE IF NOT EXISTS "LiteLLM_AutoRouterUserSession" (
|
||||
"user_id" TEXT NOT NULL,
|
||||
"api_key" TEXT NOT NULL,
|
||||
"session_id" TEXT NOT NULL,
|
||||
"router_name" TEXT NOT NULL,
|
||||
"router_type" TEXT NOT NULL,
|
||||
"first_turn_at" TIMESTAMP(3) NOT NULL,
|
||||
"last_turn_at" TIMESTAMP(3) NOT NULL,
|
||||
"last_model" TEXT NOT NULL,
|
||||
"models" JSONB NOT NULL DEFAULT '{}',
|
||||
"turns" INTEGER NOT NULL DEFAULT 0,
|
||||
"unordered_turns" INTEGER NOT NULL DEFAULT 0,
|
||||
"covered_turns" INTEGER NOT NULL DEFAULT 0,
|
||||
"cache_hits" INTEGER NOT NULL DEFAULT 0,
|
||||
"same_model_turns" INTEGER NOT NULL DEFAULT 0,
|
||||
"same_model_hits" INTEGER NOT NULL DEFAULT 0,
|
||||
"first_visit_turns" INTEGER NOT NULL DEFAULT 0,
|
||||
"first_visit_hits" INTEGER NOT NULL DEFAULT 0,
|
||||
"return_turns" INTEGER NOT NULL DEFAULT 0,
|
||||
"return_hits" INTEGER NOT NULL DEFAULT 0,
|
||||
"return_expired_misses" INTEGER NOT NULL DEFAULT 0,
|
||||
"return_within_ttl_misses" INTEGER NOT NULL DEFAULT 0,
|
||||
"ttl_5m_turns" INTEGER NOT NULL DEFAULT 0,
|
||||
"ttl_1h_turns" INTEGER NOT NULL DEFAULT 0,
|
||||
"total_tokens" BIGINT NOT NULL DEFAULT 0,
|
||||
"spend" DOUBLE PRECISION NOT NULL DEFAULT 0,
|
||||
"saved_spend" DOUBLE PRECISION NOT NULL DEFAULT 0,
|
||||
"savings_estimated_turns" INTEGER NOT NULL DEFAULT 0,
|
||||
"savings_estimated_actual_spend" DOUBLE PRECISION NOT NULL DEFAULT 0,
|
||||
"savings_estimated_saved_spend" DOUBLE PRECISION NOT NULL DEFAULT 0,
|
||||
"savings_estimated_baseline_models" JSONB NOT NULL DEFAULT '{}',
|
||||
"classifier_cost" DOUBLE PRECISION NOT NULL DEFAULT 0,
|
||||
"classifier_cost_recorded_turns" INTEGER NOT NULL DEFAULT 0,
|
||||
"tier_turns" JSONB NOT NULL DEFAULT '{}',
|
||||
"baseline_models" JSONB NOT NULL DEFAULT '{}',
|
||||
|
||||
CONSTRAINT "LiteLLM_AutoRouterUserSession_pkey" PRIMARY KEY ("user_id", "api_key", "session_id", "router_name")
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS "idx_autorouter_user_session_last_turn" ON "LiteLLM_AutoRouterUserSession"("last_turn_at");
|
||||
|
||||
CREATE INDEX IF NOT EXISTS "idx_autorouter_user_session_user_last_turn" ON "LiteLLM_AutoRouterUserSession"("user_id", "last_turn_at");
|
||||
|
|
@ -0,0 +1 @@
|
|||
ALTER TABLE "LiteLLM_PolicyAttachmentTable" ADD COLUMN IF NOT EXISTS "is_default" BOOLEAN NOT NULL DEFAULT false;
|
||||
|
|
@ -0,0 +1,3 @@
|
|||
ALTER TABLE "LiteLLM_UserTable" ADD COLUMN IF NOT EXISTS "password_reset_required" BOOLEAN;
|
||||
|
||||
ALTER TABLE "LiteLLM_UserTable" ADD COLUMN IF NOT EXISTS "last_breach_check_at" TIMESTAMP(3);
|
||||
|
|
@ -246,6 +246,8 @@ model LiteLLM_UserTable {
|
|||
organization_id String?
|
||||
object_permission_id String?
|
||||
password String?
|
||||
password_reset_required Boolean?
|
||||
last_breach_check_at DateTime?
|
||||
teams String[] @default([])
|
||||
user_role String?
|
||||
max_budget Float?
|
||||
|
|
@ -1419,6 +1421,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
|
||||
|
|
@ -1620,6 +1623,47 @@ model LiteLLM_AutoRouterSession {
|
|||
@@index([last_turn_at], map: "idx_autorouter_session_last_turn")
|
||||
}
|
||||
|
||||
model LiteLLM_AutoRouterUserSession {
|
||||
user_id String
|
||||
api_key String
|
||||
session_id String
|
||||
router_name String
|
||||
router_type String
|
||||
first_turn_at DateTime
|
||||
last_turn_at DateTime
|
||||
last_model String
|
||||
models Json @default("{}")
|
||||
turns Int @default(0)
|
||||
unordered_turns Int @default(0)
|
||||
covered_turns Int @default(0)
|
||||
cache_hits Int @default(0)
|
||||
same_model_turns Int @default(0)
|
||||
same_model_hits Int @default(0)
|
||||
first_visit_turns Int @default(0)
|
||||
first_visit_hits Int @default(0)
|
||||
return_turns Int @default(0)
|
||||
return_hits Int @default(0)
|
||||
return_expired_misses Int @default(0)
|
||||
return_within_ttl_misses Int @default(0)
|
||||
ttl_5m_turns Int @default(0)
|
||||
ttl_1h_turns Int @default(0)
|
||||
total_tokens BigInt @default(0)
|
||||
spend Float @default(0)
|
||||
saved_spend Float @default(0)
|
||||
savings_estimated_turns Int @default(0)
|
||||
savings_estimated_actual_spend Float @default(0)
|
||||
savings_estimated_saved_spend Float @default(0)
|
||||
savings_estimated_baseline_models Json @default("{}")
|
||||
classifier_cost Float @default(0)
|
||||
classifier_cost_recorded_turns Int @default(0)
|
||||
tier_turns Json @default("{}")
|
||||
baseline_models Json @default("{}")
|
||||
|
||||
@@id([user_id, api_key, session_id, router_name])
|
||||
@@index([last_turn_at], map: "idx_autorouter_user_session_last_turn")
|
||||
@@index([user_id, last_turn_at], map: "idx_autorouter_user_session_user_last_turn")
|
||||
}
|
||||
|
||||
// Shadow eval: evaluation of an auto-router against one or more keys' live traffic, in
|
||||
// either direction. forward duplicates the requests the keys did not route through the
|
||||
// router through it, answering whether they should adopt it; reverse duplicates the
|
||||
|
|
|
|||
708
litellm-rust/Cargo.lock
generated
708
litellm-rust/Cargo.lock
generated
File diff suppressed because it is too large
Load diff
|
|
@ -22,12 +22,21 @@ 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-hashicorp = { path = "crates/secrets-hashicorp" }
|
||||
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-s3 = { path = "crates/cache-s3" }
|
||||
litellm-cache-gcs = { path = "crates/cache-gcs" }
|
||||
litellm-cache-disk = { path = "crates/cache-disk" }
|
||||
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" }
|
||||
|
|
@ -48,6 +57,9 @@ reqwest = { version = "0.12", default-features = false, features = ["json", "mul
|
|||
rstest = "0.26.1"
|
||||
rstest_reuse = "0.7.0"
|
||||
rustls = { version = "0.23", default-features = false, features = ["ring", "std", "tls12"] }
|
||||
rustify = "=0.7.0"
|
||||
rustify_derive = "=0.5.5"
|
||||
vaultrs = { version = "=0.8.0", default-features = false, features = ["rustls"] }
|
||||
rustls-native-certs = "0.8"
|
||||
serde = { version = "1.0", features = ["derive"] }
|
||||
serde_json = { version = "1.0", features = ["float_roundtrip"] }
|
||||
|
|
@ -64,6 +76,7 @@ base64 = "0.22"
|
|||
moka = { version = "0.12.16", features = ["future"] }
|
||||
strum = { version = "0.28.0", features = ["derive"] }
|
||||
url = "2.5.8"
|
||||
percent-encoding = "2.3"
|
||||
webpki-roots = "1"
|
||||
time = { version = "0.3.53", features = ["parsing"] }
|
||||
criterion = "0.8.2"
|
||||
|
|
|
|||
|
|
@ -4,4 +4,4 @@ mod resolve;
|
|||
mod types;
|
||||
|
||||
pub use resolve::AzureAuthService;
|
||||
pub use types::AzureAuthInputs;
|
||||
pub use types::{AzureAuthInputs, ConfigValue};
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -128,6 +128,14 @@ impl VertexAuth {
|
|||
}
|
||||
}
|
||||
|
||||
pub async fn access_token(
|
||||
&self,
|
||||
config: &VertexConfig,
|
||||
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
|
||||
) -> Result<String, Error> {
|
||||
self.load_provider(config, env_lookup).await?.token().await
|
||||
}
|
||||
|
||||
pub async fn validate_environment(
|
||||
&self,
|
||||
headers: Vec<(String, String)>,
|
||||
|
|
|
|||
22
litellm-rust/crates/cache-azure-blob/Cargo.toml
Normal file
22
litellm-rust/crates/cache-azure-blob/Cargo.toml
Normal 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
254
litellm-rust/crates/cache-azure-blob/src/cache.rs
Normal 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
746
litellm-rust/crates/cache-azure-blob/src/cache/tests.rs
vendored
Normal 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
84
litellm-rust/crates/cache-azure-blob/src/credential.rs
Normal 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
5
litellm-rust/crates/cache-azure-blob/src/lib.rs
Normal 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
0
litellm-rust/crates/cache-azure-blob/src/tests.rs
Normal file
19
litellm-rust/crates/cache-disk/Cargo.toml
Normal file
19
litellm-rust/crates/cache-disk/Cargo.toml
Normal file
|
|
@ -0,0 +1,19 @@
|
|||
[package]
|
||||
name = "litellm-cache-disk"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
litellm-cache.workspace = true
|
||||
py_literal = "0.4.0"
|
||||
rand.workspace = true
|
||||
rusqlite = { version = "0.40", features = ["bundled"] }
|
||||
serde-pickle = "1.2"
|
||||
serde_json.workspace = true
|
||||
tokio.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
rstest.workspace = true
|
||||
tempfile = "3.27.0"
|
||||
10
litellm-rust/crates/cache-disk/src/adapter.rs
Normal file
10
litellm-rust/crates/cache-disk/src/adapter.rs
Normal file
|
|
@ -0,0 +1,10 @@
|
|||
use litellm_cache::Error;
|
||||
|
||||
use crate::StoredValue;
|
||||
|
||||
pub trait ValueAdapter: Send + Sync + 'static {
|
||||
fn read(&self, value: StoredValue) -> Result<Option<Vec<u8>>, Error>;
|
||||
fn write(&self, payload: Vec<u8>) -> StoredValue;
|
||||
fn counter_seed(&self, value: Option<StoredValue>) -> Result<f64, Error>;
|
||||
fn counter_value(&self, value: f64) -> StoredValue;
|
||||
}
|
||||
301
litellm-rust/crates/cache-disk/src/cache.rs
Normal file
301
litellm-rust/crates/cache-disk/src/cache.rs
Normal file
|
|
@ -0,0 +1,301 @@
|
|||
use std::{
|
||||
path::Path,
|
||||
sync::Arc,
|
||||
time::{Duration, SystemTime, UNIX_EPOCH},
|
||||
};
|
||||
|
||||
use litellm_cache::{
|
||||
BaseCache, BatchCache, BatchEntry, CacheCodec, CacheConnectionResult, CacheConnectionStatus,
|
||||
CounterCache, DeleteCache, Error, ExactCacheContext, FlushCache,
|
||||
};
|
||||
|
||||
use crate::{DiskStore, DiskcacheSqliteStore, PythonDiskCacheAdapter, StoredValue, ValueAdapter};
|
||||
|
||||
pub struct DiskCache<S, D = DiskcacheSqliteStore, A = PythonDiskCacheAdapter> {
|
||||
store: Arc<D>,
|
||||
adapter: Arc<A>,
|
||||
codec: S,
|
||||
}
|
||||
|
||||
impl<S: CacheCodec> DiskCache<S> {
|
||||
pub fn open(directory: impl AsRef<Path>, codec: S) -> Result<Self, Error> {
|
||||
Ok(Self {
|
||||
store: Arc::new(DiskcacheSqliteStore::open(directory)?),
|
||||
adapter: Arc::new(PythonDiskCacheAdapter),
|
||||
codec,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl<S: CacheCodec, D: DiskStore> DiskCache<S, D, PythonDiskCacheAdapter> {
|
||||
pub fn with_store(store: D, codec: S) -> Self {
|
||||
Self {
|
||||
store: Arc::new(store),
|
||||
adapter: Arc::new(PythonDiskCacheAdapter),
|
||||
codec,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<S: CacheCodec, D: DiskStore, A: ValueAdapter> DiskCache<S, D, A> {
|
||||
pub fn with_adapter(store: D, adapter: A, codec: S) -> Self {
|
||||
Self {
|
||||
store: Arc::new(store),
|
||||
adapter: Arc::new(adapter),
|
||||
codec,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn directory(&self) -> &Path {
|
||||
self.store.directory()
|
||||
}
|
||||
|
||||
fn decode_stored(&self, value: StoredValue) -> Result<Option<S::Value>, Error> {
|
||||
let Some(bytes) = self.adapter.read(value)? else {
|
||||
return Ok(None);
|
||||
};
|
||||
self.codec.decode(&bytes).map(Some)
|
||||
}
|
||||
|
||||
async fn run_blocking<T, F>(store: Arc<D>, operation: F) -> Result<T, Error>
|
||||
where
|
||||
T: Send + 'static,
|
||||
F: FnOnce(&D) -> Result<T, Error> + Send + 'static,
|
||||
{
|
||||
tokio::task::spawn_blocking(move || operation(&store))
|
||||
.await
|
||||
.map_err(|_| Error::Unavailable)?
|
||||
}
|
||||
}
|
||||
|
||||
impl<S: CacheCodec, D: DiskStore, A: ValueAdapter> BaseCache for DiskCache<S, D, A> {
|
||||
type Value = S::Value;
|
||||
type Context = ExactCacheContext;
|
||||
|
||||
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
|
||||
context.ttl
|
||||
}
|
||||
|
||||
fn set_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
value: Self::Value,
|
||||
context: &Self::Context,
|
||||
) -> Result<(), Error> {
|
||||
let value = self.adapter.write(self.codec.encode(&value)?);
|
||||
let expire_time = context.ttl.map(|ttl| unix_now() + ttl.as_secs_f64());
|
||||
self.store.set(key, value, expire_time, unix_now())
|
||||
}
|
||||
|
||||
fn get_cache(&self, key: &str, _: &Self::Context) -> Result<Option<Self::Value>, Error> {
|
||||
self.store
|
||||
.get(key, unix_now())?
|
||||
.map(|value| self.decode_stored(value))
|
||||
.transpose()
|
||||
.map(|value| value.flatten())
|
||||
}
|
||||
|
||||
async fn async_set_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
value: Self::Value,
|
||||
context: ExactCacheContext,
|
||||
) -> Result<(), Error> {
|
||||
let value = self.adapter.write(self.codec.encode(&value)?);
|
||||
let ttl = context.ttl;
|
||||
let key = key.to_string();
|
||||
Self::run_blocking(Arc::clone(&self.store), move |store| {
|
||||
let expire_time = ttl.map(|ttl| unix_now() + ttl.as_secs_f64());
|
||||
store.set(&key, value, expire_time, unix_now())
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
async fn async_get_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
_: &ExactCacheContext,
|
||||
) -> Result<Option<Self::Value>, Error> {
|
||||
let key = key.to_string();
|
||||
let value = Self::run_blocking(Arc::clone(&self.store), move |store| {
|
||||
store.get(&key, unix_now())
|
||||
})
|
||||
.await?;
|
||||
value
|
||||
.map(|value| self.decode_stored(value))
|
||||
.transpose()
|
||||
.map(|value| value.flatten())
|
||||
}
|
||||
|
||||
async fn async_set_cache_pipeline(
|
||||
&self,
|
||||
entries: Vec<(String, Self::Value)>,
|
||||
context: ExactCacheContext,
|
||||
) -> Result<(), Error> {
|
||||
let entries = entries
|
||||
.into_iter()
|
||||
.map(|(key, value)| {
|
||||
self.codec
|
||||
.encode(&value)
|
||||
.map(|value| (key, self.adapter.write(value)))
|
||||
})
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
let expire_after = context.ttl;
|
||||
Self::run_blocking(Arc::clone(&self.store), move |store| {
|
||||
for (key, value) in entries {
|
||||
let expire_time = expire_after.map(|ttl| unix_now() + ttl.as_secs_f64());
|
||||
store.set(&key, value, expire_time, unix_now())?;
|
||||
}
|
||||
Ok(())
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
async fn disconnect(&self) -> Result<(), Error> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
|
||||
let result = Self::run_blocking(Arc::clone(&self.store), |store| {
|
||||
store.probe().map(|_| CacheConnectionResult {
|
||||
status: CacheConnectionStatus::Success,
|
||||
message: "Disk cache connection test successful".into(),
|
||||
error: None,
|
||||
})
|
||||
})
|
||||
.await;
|
||||
Ok(match result {
|
||||
Ok(result) => result,
|
||||
Err(error) => CacheConnectionResult {
|
||||
status: CacheConnectionStatus::Failed,
|
||||
message: format!("Disk cache connection failed: {error}"),
|
||||
error: Some(error.to_string()),
|
||||
},
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl<S: CacheCodec, D: DiskStore, A: ValueAdapter> BatchCache for DiskCache<S, D, A> {
|
||||
fn batch_get_cache(
|
||||
&self,
|
||||
keys: &[String],
|
||||
context: &ExactCacheContext,
|
||||
) -> 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()
|
||||
}
|
||||
|
||||
async fn async_batch_get_cache(
|
||||
&self,
|
||||
keys: Vec<String>,
|
||||
_: ExactCacheContext,
|
||||
) -> Result<Vec<BatchEntry<Self::Value>>, Error> {
|
||||
let values = Self::run_blocking(Arc::clone(&self.store), move |store| {
|
||||
keys.into_iter()
|
||||
.map(|key| store.get(&key, unix_now()).map(|value| (key, value)))
|
||||
.collect::<Result<Vec<_>, _>>()
|
||||
})
|
||||
.await?;
|
||||
values
|
||||
.into_iter()
|
||||
.map(|(_, value)| match value {
|
||||
None => Ok(BatchEntry::Miss),
|
||||
Some(value) => match self.decode_stored(value) {
|
||||
Ok(Some(value)) => Ok(BatchEntry::Hit(value)),
|
||||
Ok(None) => Ok(BatchEntry::Miss),
|
||||
Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid),
|
||||
Err(error) => Err(error),
|
||||
},
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
impl<S: CacheCodec, D: DiskStore, A: ValueAdapter> DeleteCache for DiskCache<S, D, A> {
|
||||
fn delete_cache(&self, key: &str) -> Result<(), Error> {
|
||||
self.store.pop(key, unix_now()).map(|_| ())
|
||||
}
|
||||
|
||||
async fn async_delete_cache(&self, key: &str) -> Result<(), Error> {
|
||||
let key = key.to_string();
|
||||
Self::run_blocking(Arc::clone(&self.store), move |store| {
|
||||
store.pop(&key, unix_now()).map(|_| ())
|
||||
})
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
impl<S: CacheCodec, D: DiskStore, A: ValueAdapter> FlushCache for DiskCache<S, D, A> {
|
||||
fn flush_cache(&self) -> Result<(), Error> {
|
||||
self.store.clear()
|
||||
}
|
||||
|
||||
async fn async_flush_cache(&self) -> Result<(), Error> {
|
||||
Self::run_blocking(Arc::clone(&self.store), |store| store.clear()).await
|
||||
}
|
||||
}
|
||||
|
||||
impl<S: CacheCodec<Value = f64>, D: DiskStore, A: ValueAdapter> CounterCache
|
||||
for DiskCache<S, D, A>
|
||||
{
|
||||
fn increment_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
amount: f64,
|
||||
context: ExactCacheContext,
|
||||
) -> Result<f64, Error> {
|
||||
increment(
|
||||
self.adapter.as_ref(),
|
||||
self.store.as_ref(),
|
||||
key,
|
||||
amount,
|
||||
context.ttl,
|
||||
)
|
||||
}
|
||||
|
||||
async fn async_increment(
|
||||
&self,
|
||||
key: &str,
|
||||
amount: f64,
|
||||
context: ExactCacheContext,
|
||||
) -> Result<f64, Error> {
|
||||
let key = key.to_string();
|
||||
let adapter = Arc::clone(&self.adapter);
|
||||
Self::run_blocking(Arc::clone(&self.store), move |store| {
|
||||
increment(adapter.as_ref(), store, &key, amount, context.ttl)
|
||||
})
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
fn increment<A: ValueAdapter, D: DiskStore>(
|
||||
adapter: &A,
|
||||
store: &D,
|
||||
key: &str,
|
||||
amount: f64,
|
||||
ttl: Option<Duration>,
|
||||
) -> Result<f64, Error> {
|
||||
let mut result = None;
|
||||
let mut apply = |current: Option<StoredValue>| {
|
||||
let initial = adapter.counter_seed(current)?;
|
||||
let value = initial + amount;
|
||||
let stored = adapter.counter_value(value);
|
||||
result = Some(value);
|
||||
Ok((stored, ttl.map(|ttl| unix_now() + ttl.as_secs_f64())))
|
||||
};
|
||||
store.update(key, unix_now(), &mut apply)?;
|
||||
result.ok_or(Error::InvalidEntry)
|
||||
}
|
||||
|
||||
fn unix_now() -> f64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs_f64()
|
||||
}
|
||||
11
litellm-rust/crates/cache-disk/src/lib.rs
Normal file
11
litellm-rust/crates/cache-disk/src/lib.rs
Normal file
|
|
@ -0,0 +1,11 @@
|
|||
mod adapter;
|
||||
mod cache;
|
||||
mod python;
|
||||
mod sqlite;
|
||||
mod store;
|
||||
|
||||
pub use adapter::ValueAdapter;
|
||||
pub use cache::DiskCache;
|
||||
pub use python::PythonDiskCacheAdapter;
|
||||
pub use sqlite::DiskcacheSqliteStore;
|
||||
pub use store::{DiskStore, StoredValue};
|
||||
77
litellm-rust/crates/cache-disk/src/python/mod.rs
Normal file
77
litellm-rust/crates/cache-disk/src/python/mod.rs
Normal file
|
|
@ -0,0 +1,77 @@
|
|||
mod value;
|
||||
|
||||
use litellm_cache::Error;
|
||||
use py_literal::Value;
|
||||
|
||||
use crate::{StoredValue, ValueAdapter};
|
||||
|
||||
#[derive(Clone, Copy, Debug, Default)]
|
||||
pub struct PythonDiskCacheAdapter;
|
||||
|
||||
impl PythonDiskCacheAdapter {
|
||||
fn python_get_cache(value: StoredValue) -> Result<Option<Value>, Error> {
|
||||
let value = match value {
|
||||
StoredValue::Bytes(value) => Value::Bytes(value),
|
||||
StoredValue::Text(value) => Value::String(value),
|
||||
StoredValue::Integer(value) => Value::Integer(value.into()),
|
||||
StoredValue::Float(value) => Value::Float(value),
|
||||
StoredValue::Pickle(value) => value::from_pickle(&value)?,
|
||||
};
|
||||
if !value::is_truthy(&value) {
|
||||
return Ok(None);
|
||||
}
|
||||
match value {
|
||||
Value::String(text) => Ok(Some(
|
||||
value::from_json_text(&text).unwrap_or(Value::String(text)),
|
||||
)),
|
||||
Value::Bytes(bytes) => match std::str::from_utf8(&bytes) {
|
||||
Ok(text) => Ok(Some(
|
||||
value::from_json_text(text).unwrap_or(Value::Bytes(bytes)),
|
||||
)),
|
||||
Err(_) => Ok(Some(Value::Bytes(bytes))),
|
||||
},
|
||||
value => Ok(Some(value)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ValueAdapter for PythonDiskCacheAdapter {
|
||||
fn read(&self, value: StoredValue) -> Result<Option<Vec<u8>>, Error> {
|
||||
match value {
|
||||
StoredValue::Text(value) => Ok((!value.is_empty()).then(|| value.into_bytes())),
|
||||
StoredValue::Bytes(value) => Ok((!value.is_empty()).then_some(value)),
|
||||
value => {
|
||||
let Some(value) = Self::python_get_cache(value)? else {
|
||||
return Ok(None);
|
||||
};
|
||||
value::to_json(&value).map(Some)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn write(&self, payload: Vec<u8>) -> StoredValue {
|
||||
StoredValue::Bytes(payload)
|
||||
}
|
||||
|
||||
fn counter_seed(&self, value: Option<StoredValue>) -> Result<f64, Error> {
|
||||
let Some(value) = value else {
|
||||
return Ok(0.0);
|
||||
};
|
||||
let Some(value) = Self::python_get_cache(value)? else {
|
||||
return Ok(0.0);
|
||||
};
|
||||
Ok(if value::is_int(&value) {
|
||||
value::to_f64(&value).unwrap_or(0.0)
|
||||
} else {
|
||||
0.0
|
||||
})
|
||||
}
|
||||
|
||||
fn counter_value(&self, value: f64) -> StoredValue {
|
||||
if value.fract() == 0.0 && value >= i64::MIN as f64 && value <= i64::MAX as f64 {
|
||||
StoredValue::Integer(value as i64)
|
||||
} else {
|
||||
StoredValue::Float(value)
|
||||
}
|
||||
}
|
||||
}
|
||||
173
litellm-rust/crates/cache-disk/src/python/value.rs
Normal file
173
litellm-rust/crates/cache-disk/src/python/value.rs
Normal file
|
|
@ -0,0 +1,173 @@
|
|||
use litellm_cache::Error;
|
||||
use py_literal::Value;
|
||||
use serde_json::{Map, Number};
|
||||
|
||||
pub(crate) fn from_pickle(bytes: &[u8]) -> Result<Value, Error> {
|
||||
let value = serde_pickle::value_from_slice(bytes, Default::default())
|
||||
.map_err(|_| Error::InvalidEntry)?;
|
||||
from_pickle_value(value)
|
||||
}
|
||||
|
||||
fn from_pickle_value(value: serde_pickle::Value) -> Result<Value, Error> {
|
||||
match value {
|
||||
serde_pickle::Value::None => Ok(Value::None),
|
||||
serde_pickle::Value::Bool(value) => Ok(Value::Boolean(value)),
|
||||
serde_pickle::Value::I64(value) => integer(value.to_string()),
|
||||
serde_pickle::Value::Int(value) => integer(value.to_string()),
|
||||
serde_pickle::Value::F64(value) => Ok(Value::Float(value)),
|
||||
serde_pickle::Value::String(value) => Ok(Value::String(value)),
|
||||
serde_pickle::Value::Bytes(value) => Ok(Value::Bytes(value)),
|
||||
serde_pickle::Value::List(values) => values
|
||||
.into_iter()
|
||||
.map(from_pickle_value)
|
||||
.collect::<Result<Vec<_>, _>>()
|
||||
.map(Value::List),
|
||||
serde_pickle::Value::Tuple(values) => values
|
||||
.into_iter()
|
||||
.map(from_pickle_value)
|
||||
.collect::<Result<Vec<_>, _>>()
|
||||
.map(Value::Tuple),
|
||||
serde_pickle::Value::Set(values) => values
|
||||
.into_iter()
|
||||
.map(from_pickle_hashable)
|
||||
.collect::<Result<Vec<_>, _>>()
|
||||
.map(Value::Set),
|
||||
serde_pickle::Value::FrozenSet(values) => values
|
||||
.into_iter()
|
||||
.map(from_pickle_hashable)
|
||||
.collect::<Result<Vec<_>, _>>()
|
||||
.map(Value::Set),
|
||||
serde_pickle::Value::Dict(values) => values
|
||||
.into_iter()
|
||||
.map(|(key, value)| Ok((from_pickle_hashable(key)?, from_pickle_value(value)?)))
|
||||
.collect::<Result<Vec<_>, Error>>()
|
||||
.map(Value::Dict),
|
||||
}
|
||||
}
|
||||
|
||||
fn from_pickle_hashable(value: serde_pickle::HashableValue) -> Result<Value, Error> {
|
||||
Ok(match value {
|
||||
serde_pickle::HashableValue::None => Value::None,
|
||||
serde_pickle::HashableValue::Bool(value) => Value::Boolean(value),
|
||||
serde_pickle::HashableValue::I64(value) => integer(value.to_string())?,
|
||||
serde_pickle::HashableValue::Int(value) => integer(value.to_string())?,
|
||||
serde_pickle::HashableValue::F64(value) => Value::Float(value),
|
||||
serde_pickle::HashableValue::Bytes(value) => Value::Bytes(value),
|
||||
serde_pickle::HashableValue::String(value) => Value::String(value),
|
||||
serde_pickle::HashableValue::Tuple(values) => Value::Tuple(
|
||||
values
|
||||
.into_iter()
|
||||
.map(from_pickle_hashable)
|
||||
.collect::<Result<Vec<_>, _>>()?,
|
||||
),
|
||||
serde_pickle::HashableValue::FrozenSet(values) => Value::Set(
|
||||
values
|
||||
.into_iter()
|
||||
.map(from_pickle_hashable)
|
||||
.collect::<Result<Vec<_>, _>>()?,
|
||||
),
|
||||
})
|
||||
}
|
||||
|
||||
fn integer(value: String) -> Result<Value, Error> {
|
||||
value.parse().map_err(|_| Error::InvalidEntry)
|
||||
}
|
||||
|
||||
pub(crate) fn from_json(value: serde_json::Value) -> Value {
|
||||
match value {
|
||||
serde_json::Value::Null => Value::None,
|
||||
serde_json::Value::Bool(value) => Value::Boolean(value),
|
||||
serde_json::Value::Number(value) => {
|
||||
if value.is_i64() || value.is_u64() {
|
||||
integer(value.to_string())
|
||||
.unwrap_or(Value::Float(value.as_f64().unwrap_or(f64::NAN)))
|
||||
} else {
|
||||
Value::Float(value.as_f64().unwrap_or(f64::NAN))
|
||||
}
|
||||
}
|
||||
serde_json::Value::String(value) => Value::String(value),
|
||||
serde_json::Value::Array(values) => {
|
||||
Value::List(values.into_iter().map(from_json).collect())
|
||||
}
|
||||
serde_json::Value::Object(values) => Value::Dict(
|
||||
values
|
||||
.into_iter()
|
||||
.map(|(key, value)| (Value::String(key), from_json(value)))
|
||||
.collect(),
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn from_json_text(value: &str) -> Result<Value, Error> {
|
||||
serde_json::from_str(value)
|
||||
.map(from_json)
|
||||
.map_err(|_| Error::InvalidEntry)
|
||||
}
|
||||
|
||||
pub(crate) fn is_truthy(value: &Value) -> bool {
|
||||
match value {
|
||||
Value::None => false,
|
||||
Value::Boolean(value) => *value,
|
||||
Value::Integer(value) => value.to_string() != "0",
|
||||
Value::Float(value) => *value != 0.0,
|
||||
Value::Complex(value) => value.re != 0.0 || value.im != 0.0,
|
||||
Value::String(value) => !value.is_empty(),
|
||||
Value::Bytes(value) => !value.is_empty(),
|
||||
Value::Tuple(value) | Value::List(value) | Value::Set(value) => !value.is_empty(),
|
||||
Value::Dict(value) => !value.is_empty(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn is_int(value: &Value) -> bool {
|
||||
matches!(value, Value::Integer(_) | Value::Boolean(_))
|
||||
}
|
||||
|
||||
pub(crate) fn to_f64(value: &Value) -> Option<f64> {
|
||||
match value {
|
||||
Value::Integer(value) => value.to_string().parse().ok(),
|
||||
Value::Boolean(value) => Some(if *value { 1.0 } else { 0.0 }),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn to_json(value: &Value) -> Result<Vec<u8>, Error> {
|
||||
serde_json::to_vec(&to_json_value(value)?).map_err(|_| Error::InvalidEntry)
|
||||
}
|
||||
|
||||
fn to_json_value(value: &Value) -> Result<serde_json::Value, Error> {
|
||||
Ok(match value {
|
||||
Value::None => serde_json::Value::Null,
|
||||
Value::Boolean(value) => serde_json::Value::Bool(*value),
|
||||
Value::Integer(value) => serde_json::Value::Number(
|
||||
value
|
||||
.to_string()
|
||||
.parse::<Number>()
|
||||
.map_err(|_| Error::InvalidEntry)?,
|
||||
),
|
||||
Value::Float(value) => {
|
||||
serde_json::Value::Number(Number::from_f64(*value).ok_or(Error::InvalidEntry)?)
|
||||
}
|
||||
Value::Complex(_) | Value::Bytes(_) => return Err(Error::InvalidEntry),
|
||||
Value::String(value) => serde_json::Value::String(value.clone()),
|
||||
Value::Tuple(values) | Value::List(values) | Value::Set(values) => {
|
||||
serde_json::Value::Array(
|
||||
values
|
||||
.iter()
|
||||
.map(to_json_value)
|
||||
.collect::<Result<Vec<_>, _>>()?,
|
||||
)
|
||||
}
|
||||
Value::Dict(values) => {
|
||||
let values = values
|
||||
.iter()
|
||||
.map(|(key, value)| {
|
||||
let Value::String(key) = key else {
|
||||
return Err(Error::InvalidEntry);
|
||||
};
|
||||
Ok((key.clone(), to_json_value(value)?))
|
||||
})
|
||||
.collect::<Result<Map<String, serde_json::Value>, _>>()?;
|
||||
serde_json::Value::Object(values)
|
||||
}
|
||||
})
|
||||
}
|
||||
817
litellm-rust/crates/cache-disk/src/sqlite.rs
Normal file
817
litellm-rust/crates/cache-disk/src/sqlite.rs
Normal file
|
|
@ -0,0 +1,817 @@
|
|||
use std::{
|
||||
collections::HashMap,
|
||||
fs::{self, OpenOptions},
|
||||
io::Write,
|
||||
path::{Path, PathBuf},
|
||||
sync::Mutex,
|
||||
};
|
||||
|
||||
use litellm_cache::Error;
|
||||
use rand::RngCore;
|
||||
use rusqlite::{Connection, OptionalExtension, params, types::Value};
|
||||
|
||||
use crate::{DiskStore, StoredValue};
|
||||
|
||||
const MODE_RAW: i64 = 1;
|
||||
const MODE_BINARY: i64 = 2;
|
||||
const MODE_TEXT: i64 = 3;
|
||||
const MODE_PICKLE: i64 = 4;
|
||||
|
||||
const DEFAULT_DISK_MIN_FILE_SIZE: i64 = 2_i64.pow(15);
|
||||
const DEFAULT_SIZE_LIMIT: i64 = 2_i64.pow(30);
|
||||
const DEFAULT_CULL_LIMIT: i64 = 10;
|
||||
|
||||
pub struct DiskcacheSqliteStore {
|
||||
directory: PathBuf,
|
||||
connection: Mutex<Connection>,
|
||||
min_file_size: usize,
|
||||
eviction_policy: String,
|
||||
size_limit: i64,
|
||||
cull_limit: i64,
|
||||
statistics: bool,
|
||||
}
|
||||
|
||||
struct StoredColumns {
|
||||
size: i64,
|
||||
mode: i64,
|
||||
filename: Option<String>,
|
||||
value: Option<Value>,
|
||||
}
|
||||
|
||||
struct Row {
|
||||
rowid: i64,
|
||||
mode: i64,
|
||||
filename: Option<String>,
|
||||
value: Value,
|
||||
}
|
||||
|
||||
impl DiskcacheSqliteStore {
|
||||
pub fn open(directory: impl AsRef<Path>) -> Result<Self, Error> {
|
||||
let directory = directory.as_ref().to_path_buf();
|
||||
fs::create_dir_all(&directory).map_err(|_| Error::Unavailable)?;
|
||||
let directory = std::path::absolute(&directory).map_err(|_| Error::Unavailable)?;
|
||||
let database = directory.join("cache.db");
|
||||
let connection = Connection::open(database).map_err(|_| Error::Unavailable)?;
|
||||
connection
|
||||
.busy_timeout(std::time::Duration::from_secs(60))
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
|
||||
let mut settings = read_settings(&connection)?;
|
||||
for (key, value) in default_settings() {
|
||||
settings.entry(key).or_insert(value);
|
||||
}
|
||||
for (key, value) in settings
|
||||
.iter()
|
||||
.filter(|(key, _)| key.starts_with("sqlite_"))
|
||||
{
|
||||
apply_pragma(&connection, key, value)?;
|
||||
}
|
||||
|
||||
connection
|
||||
.execute_batch(
|
||||
"CREATE TABLE IF NOT EXISTS Settings (
|
||||
key TEXT NOT NULL UNIQUE,
|
||||
value
|
||||
)",
|
||||
)
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
for (key, value) in &settings {
|
||||
if !matches!(key.as_str(), "count" | "size" | "hits" | "misses") {
|
||||
connection
|
||||
.execute(
|
||||
"INSERT OR REPLACE INTO Settings VALUES (?, ?)",
|
||||
params![key, value],
|
||||
)
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
}
|
||||
}
|
||||
for (key, value) in [
|
||||
("count", Value::Integer(0)),
|
||||
("size", Value::Integer(0)),
|
||||
("hits", Value::Integer(0)),
|
||||
("misses", Value::Integer(0)),
|
||||
] {
|
||||
connection
|
||||
.execute(
|
||||
"INSERT OR IGNORE INTO Settings VALUES (?, ?)",
|
||||
params![key, value],
|
||||
)
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
}
|
||||
connection
|
||||
.execute_batch(
|
||||
"CREATE TABLE IF NOT EXISTS Cache (
|
||||
rowid INTEGER PRIMARY KEY,
|
||||
key BLOB,
|
||||
raw INTEGER,
|
||||
store_time REAL,
|
||||
expire_time REAL,
|
||||
access_time REAL,
|
||||
access_count INTEGER DEFAULT 0,
|
||||
tag BLOB,
|
||||
size INTEGER DEFAULT 0,
|
||||
mode INTEGER DEFAULT 0,
|
||||
filename TEXT,
|
||||
value BLOB
|
||||
);
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS Cache_key_raw ON Cache(key, raw);
|
||||
CREATE INDEX IF NOT EXISTS Cache_expire_time ON Cache(expire_time);",
|
||||
)
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
|
||||
let eviction_policy = setting_string(&settings, "eviction_policy")
|
||||
.unwrap_or_else(|| "least-recently-stored".to_string());
|
||||
match eviction_policy.as_str() {
|
||||
"none" => {}
|
||||
"least-recently-stored" => {
|
||||
connection
|
||||
.execute_batch(
|
||||
"CREATE INDEX IF NOT EXISTS Cache_store_time ON Cache(store_time)",
|
||||
)
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
}
|
||||
"least-recently-used" => {
|
||||
connection
|
||||
.execute_batch(
|
||||
"CREATE INDEX IF NOT EXISTS Cache_access_time ON Cache(access_time)",
|
||||
)
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
}
|
||||
"least-frequently-used" => {
|
||||
connection
|
||||
.execute_batch(
|
||||
"CREATE INDEX IF NOT EXISTS Cache_access_count ON Cache(access_count)",
|
||||
)
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
}
|
||||
_ => return Err(Error::Unavailable),
|
||||
}
|
||||
connection
|
||||
.execute_batch(
|
||||
"CREATE TRIGGER IF NOT EXISTS Settings_count_insert
|
||||
AFTER INSERT ON Cache FOR EACH ROW BEGIN
|
||||
UPDATE Settings SET value = value + 1
|
||||
WHERE key = \"count\"; END;
|
||||
CREATE TRIGGER IF NOT EXISTS Settings_count_delete
|
||||
AFTER DELETE ON Cache FOR EACH ROW BEGIN
|
||||
UPDATE Settings SET value = value - 1
|
||||
WHERE key = \"count\"; END;
|
||||
CREATE TRIGGER IF NOT EXISTS Settings_size_insert
|
||||
AFTER INSERT ON Cache FOR EACH ROW BEGIN
|
||||
UPDATE Settings SET value = value + NEW.size
|
||||
WHERE key = \"size\"; END;
|
||||
CREATE TRIGGER IF NOT EXISTS Settings_size_update
|
||||
AFTER UPDATE ON Cache FOR EACH ROW BEGIN
|
||||
UPDATE Settings
|
||||
SET value = value + NEW.size - OLD.size
|
||||
WHERE key = \"size\"; END;
|
||||
CREATE TRIGGER IF NOT EXISTS Settings_size_delete
|
||||
AFTER DELETE ON Cache FOR EACH ROW BEGIN
|
||||
UPDATE Settings SET value = value - OLD.size
|
||||
WHERE key = \"size\"; END;",
|
||||
)
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
|
||||
let min_file_size = setting_i64(&settings, "disk_min_file_size")
|
||||
.unwrap_or(DEFAULT_DISK_MIN_FILE_SIZE)
|
||||
.try_into()
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
let size_limit = setting_i64(&settings, "size_limit").unwrap_or(DEFAULT_SIZE_LIMIT);
|
||||
let cull_limit = setting_i64(&settings, "cull_limit").unwrap_or(DEFAULT_CULL_LIMIT);
|
||||
let statistics = setting_i64(&settings, "statistics").unwrap_or_default() != 0;
|
||||
|
||||
Ok(Self {
|
||||
directory,
|
||||
connection: Mutex::new(connection),
|
||||
min_file_size,
|
||||
eviction_policy,
|
||||
size_limit,
|
||||
cull_limit,
|
||||
statistics,
|
||||
})
|
||||
}
|
||||
|
||||
fn set_locked(
|
||||
&self,
|
||||
connection: &Connection,
|
||||
key: &str,
|
||||
columns: StoredColumns,
|
||||
expire_time: Option<f64>,
|
||||
now: f64,
|
||||
) -> Result<Vec<String>, Error> {
|
||||
let mut cleanup = Vec::new();
|
||||
if let Some(old_filename) = connection
|
||||
.query_row(
|
||||
"SELECT filename FROM Cache WHERE key = ? AND raw = 1",
|
||||
params![key],
|
||||
|row| row.get::<_, Option<String>>(0),
|
||||
)
|
||||
.optional()
|
||||
.map_err(|_| Error::Unavailable)?
|
||||
.flatten()
|
||||
{
|
||||
cleanup.push(old_filename);
|
||||
}
|
||||
let (size, mode, filename, value) =
|
||||
(columns.size, columns.mode, columns.filename, columns.value);
|
||||
let rowid = connection
|
||||
.query_row(
|
||||
"SELECT rowid FROM Cache WHERE key = ? AND raw = 1",
|
||||
params![key],
|
||||
|row| row.get::<_, i64>(0),
|
||||
)
|
||||
.optional()
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
if let Some(rowid) = rowid {
|
||||
connection
|
||||
.execute(
|
||||
"UPDATE Cache SET store_time = ?, expire_time = ?, access_time = ?,
|
||||
access_count = 0, tag = NULL, size = ?, mode = ?, filename = ?, value = ?
|
||||
WHERE rowid = ?",
|
||||
params![now, expire_time, now, size, mode, filename, value, rowid],
|
||||
)
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
} else {
|
||||
connection
|
||||
.execute(
|
||||
"INSERT INTO Cache(
|
||||
key, raw, store_time, expire_time, access_time, access_count,
|
||||
tag, size, mode, filename, value
|
||||
) VALUES (?, 1, ?, ?, ?, 0, NULL, ?, ?, ?, ?)",
|
||||
params![key, now, expire_time, now, size, mode, filename, value],
|
||||
)
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
}
|
||||
cleanup.extend(self.cull(connection, now)?);
|
||||
Ok(cleanup)
|
||||
}
|
||||
|
||||
fn cull(&self, connection: &Connection, now: f64) -> Result<Vec<String>, Error> {
|
||||
if self.cull_limit <= 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let mut cleanup = Vec::new();
|
||||
let expired = connection
|
||||
.prepare(
|
||||
"SELECT rowid, filename FROM Cache
|
||||
WHERE expire_time IS NOT NULL AND expire_time < ?
|
||||
ORDER BY expire_time LIMIT ?",
|
||||
)
|
||||
.map_err(|_| Error::Unavailable)?
|
||||
.query_map(params![now, self.cull_limit], |row| {
|
||||
Ok((row.get::<_, i64>(0)?, row.get::<_, Option<String>>(1)?))
|
||||
})
|
||||
.map_err(|_| Error::Unavailable)?
|
||||
.collect::<Result<Vec<_>, _>>()
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
for (_, filename) in &expired {
|
||||
if let Some(filename) = filename {
|
||||
cleanup.push(filename.clone());
|
||||
}
|
||||
}
|
||||
for (rowid, _) in &expired {
|
||||
connection
|
||||
.execute("DELETE FROM Cache WHERE rowid = ?", params![rowid])
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
}
|
||||
let remaining = self.cull_limit - i64::try_from(expired.len()).unwrap_or(self.cull_limit);
|
||||
if remaining <= 0 || self.volume(connection)? < self.size_limit {
|
||||
return Ok(cleanup);
|
||||
}
|
||||
let order = match self.eviction_policy.as_str() {
|
||||
"none" => return Ok(cleanup),
|
||||
"least-recently-stored" => "store_time",
|
||||
"least-recently-used" => "access_time",
|
||||
"least-frequently-used" => "access_count",
|
||||
_ => return Err(Error::Unavailable),
|
||||
};
|
||||
let rows = connection
|
||||
.prepare(&format!(
|
||||
"SELECT rowid, filename FROM Cache ORDER BY {order} LIMIT ?"
|
||||
))
|
||||
.map_err(|_| Error::Unavailable)?
|
||||
.query_map(params![remaining], |row| {
|
||||
Ok((row.get::<_, i64>(0)?, row.get::<_, Option<String>>(1)?))
|
||||
})
|
||||
.map_err(|_| Error::Unavailable)?
|
||||
.collect::<Result<Vec<_>, _>>()
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
for (_, filename) in &rows {
|
||||
if let Some(filename) = filename {
|
||||
cleanup.push(filename.clone());
|
||||
}
|
||||
}
|
||||
for (rowid, _) in rows {
|
||||
connection
|
||||
.execute("DELETE FROM Cache WHERE rowid = ?", params![rowid])
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
}
|
||||
Ok(cleanup)
|
||||
}
|
||||
|
||||
fn volume(&self, connection: &Connection) -> Result<i64, Error> {
|
||||
let page_count: i64 = connection
|
||||
.query_row("PRAGMA page_count", [], |row| row.get(0))
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
let page_size: i64 = connection
|
||||
.query_row("PRAGMA page_size", [], |row| row.get(0))
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
let size: i64 = connection
|
||||
.query_row("SELECT value FROM Settings WHERE key = 'size'", [], |row| {
|
||||
row.get(0)
|
||||
})
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
Ok(page_count.saturating_mul(page_size).saturating_add(size))
|
||||
}
|
||||
}
|
||||
|
||||
impl DiskStore for DiskcacheSqliteStore {
|
||||
fn directory(&self) -> &Path {
|
||||
&self.directory
|
||||
}
|
||||
|
||||
fn get(&self, key: &str, now: f64) -> Result<Option<StoredValue>, Error> {
|
||||
let connection = self.connection.lock().map_err(|_| Error::Unavailable)?;
|
||||
let select = "SELECT rowid, expire_time, mode, filename, value FROM Cache
|
||||
WHERE key = ? AND raw = 1 AND (expire_time IS NULL OR expire_time > ?)";
|
||||
let row = connection
|
||||
.query_row(select, params![key, now], row_from_query)
|
||||
.optional()
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
if !self.statistics && !has_get_update(&self.eviction_policy) {
|
||||
return row
|
||||
.map(|row| fetch_row(&self.directory, row))
|
||||
.transpose()
|
||||
.map(|value| value.flatten());
|
||||
}
|
||||
transactional(&connection, |connection| {
|
||||
let row = connection
|
||||
.query_row(select, params![key, now], row_from_query)
|
||||
.optional()
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
let Some(row) = row else {
|
||||
if self.statistics {
|
||||
connection
|
||||
.execute(
|
||||
"UPDATE Settings SET value = value + 1 WHERE key = 'misses'",
|
||||
[],
|
||||
)
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
}
|
||||
return Ok(None);
|
||||
};
|
||||
let rowid = row.rowid;
|
||||
let value = fetch_row(&self.directory, row);
|
||||
let hit = value.as_ref().is_ok_and(Option::is_some);
|
||||
if hit && self.statistics {
|
||||
connection
|
||||
.execute(
|
||||
"UPDATE Settings SET value = value + 1 WHERE key = 'hits'",
|
||||
[],
|
||||
)
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
} else if !hit && self.statistics {
|
||||
connection
|
||||
.execute(
|
||||
"UPDATE Settings SET value = value + 1 WHERE key = 'misses'",
|
||||
[],
|
||||
)
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
}
|
||||
if has_get_update(&self.eviction_policy) && hit {
|
||||
let update = match self.eviction_policy.as_str() {
|
||||
"least-recently-used" => "UPDATE Cache SET access_time = ? WHERE rowid = ?",
|
||||
"least-frequently-used" => {
|
||||
"UPDATE Cache SET access_count = access_count + 1 WHERE rowid = ?"
|
||||
}
|
||||
_ => return Err(Error::Unavailable),
|
||||
};
|
||||
if self.eviction_policy == "least-recently-used" {
|
||||
connection
|
||||
.execute(update, params![now, rowid])
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
} else {
|
||||
connection
|
||||
.execute(update, params![rowid])
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
}
|
||||
}
|
||||
value
|
||||
})
|
||||
}
|
||||
|
||||
fn set(
|
||||
&self,
|
||||
key: &str,
|
||||
value: StoredValue,
|
||||
expire_time: Option<f64>,
|
||||
now: f64,
|
||||
) -> Result<(), Error> {
|
||||
let columns = store_value(&self.directory, self.min_file_size, value)?;
|
||||
let new_filename = columns.filename.clone();
|
||||
let connection = self.connection.lock().map_err(|_| Error::Unavailable)?;
|
||||
let result = transactional(&connection, |connection| {
|
||||
self.set_locked(connection, key, columns, expire_time, now)
|
||||
});
|
||||
match result {
|
||||
Ok(cleanup) => {
|
||||
cleanup_files(&self.directory, cleanup);
|
||||
Ok(())
|
||||
}
|
||||
Err(error) => {
|
||||
if let Some(filename) = new_filename {
|
||||
remove_file(&self.directory, &filename);
|
||||
}
|
||||
Err(error)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn pop(&self, key: &str, now: f64) -> Result<Option<StoredValue>, Error> {
|
||||
let connection = self.connection.lock().map_err(|_| Error::Unavailable)?;
|
||||
let selected = transactional(&connection, |connection| {
|
||||
let row = connection
|
||||
.query_row(
|
||||
"SELECT rowid, expire_time, mode, filename, value FROM Cache
|
||||
WHERE key = ? AND raw = 1
|
||||
AND (expire_time IS NULL OR expire_time > ?)",
|
||||
params![key, now],
|
||||
row_from_query,
|
||||
)
|
||||
.optional()
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
let Some(row) = row else {
|
||||
return Ok(None);
|
||||
};
|
||||
connection
|
||||
.execute("DELETE FROM Cache WHERE rowid = ?", params![row.rowid])
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
Ok(Some(row))
|
||||
})?;
|
||||
let Some(row) = selected else {
|
||||
return Ok(None);
|
||||
};
|
||||
let filename = row.filename.clone();
|
||||
let result = fetch_row(&self.directory, row)?;
|
||||
if let Some(filename) = filename {
|
||||
remove_file(&self.directory, &filename);
|
||||
}
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
fn clear(&self) -> Result<(), Error> {
|
||||
let connection = self.connection.lock().map_err(|_| Error::Unavailable)?;
|
||||
let mut last_rowid = 0_i64;
|
||||
loop {
|
||||
let batch = transactional(&connection, |connection| {
|
||||
let rows = connection
|
||||
.prepare(
|
||||
"SELECT rowid, filename FROM Cache
|
||||
WHERE rowid > ? ORDER BY rowid LIMIT 100",
|
||||
)
|
||||
.map_err(|_| Error::Unavailable)?
|
||||
.query_map(params![last_rowid], |row| {
|
||||
Ok((row.get::<_, i64>(0)?, row.get::<_, Option<String>>(1)?))
|
||||
})
|
||||
.map_err(|_| Error::Unavailable)?
|
||||
.collect::<Result<Vec<_>, _>>()
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
if rows.is_empty() {
|
||||
return Ok(rows);
|
||||
}
|
||||
let ids = rows
|
||||
.iter()
|
||||
.map(|(rowid, _)| rowid.to_string())
|
||||
.collect::<Vec<_>>()
|
||||
.join(",");
|
||||
connection
|
||||
.execute(&format!("DELETE FROM Cache WHERE rowid IN ({ids})"), [])
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
Ok(rows)
|
||||
})?;
|
||||
if batch.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
last_rowid = batch.last().map(|(rowid, _)| *rowid).unwrap_or(last_rowid);
|
||||
cleanup_files(
|
||||
&self.directory,
|
||||
batch
|
||||
.into_iter()
|
||||
.filter_map(|(_, filename)| filename)
|
||||
.collect(),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
fn update(
|
||||
&self,
|
||||
key: &str,
|
||||
now: f64,
|
||||
apply: &mut dyn FnMut(Option<StoredValue>) -> Result<(StoredValue, Option<f64>), Error>,
|
||||
) -> Result<(), Error> {
|
||||
let connection = self.connection.lock().map_err(|_| Error::Unavailable)?;
|
||||
let mut created_filename = None;
|
||||
let result = transactional(&connection, |connection| {
|
||||
let current = connection
|
||||
.query_row(
|
||||
"SELECT rowid, expire_time, mode, filename, value FROM Cache
|
||||
WHERE key = ? AND raw = 1
|
||||
AND (expire_time IS NULL OR expire_time > ?)",
|
||||
params![key, now],
|
||||
row_from_query,
|
||||
)
|
||||
.optional()
|
||||
.map_err(|_| Error::Unavailable)?
|
||||
.map(|row| fetch_row(&self.directory, row))
|
||||
.transpose()?
|
||||
.flatten();
|
||||
let (value, expire_time) = apply(current)?;
|
||||
let columns = store_value(&self.directory, self.min_file_size, value)?;
|
||||
created_filename = columns.filename.clone();
|
||||
let cleanup = self.set_locked(connection, key, columns, expire_time, now)?;
|
||||
Ok(cleanup)
|
||||
});
|
||||
match result {
|
||||
Ok(cleanup) => {
|
||||
cleanup_files(&self.directory, cleanup);
|
||||
Ok(())
|
||||
}
|
||||
Err(error) => {
|
||||
if let Some(filename) = created_filename {
|
||||
remove_file(&self.directory, &filename);
|
||||
}
|
||||
Err(error)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn probe(&self) -> Result<(), Error> {
|
||||
let connection = self.connection.lock().map_err(|_| Error::Unavailable)?;
|
||||
connection
|
||||
.query_row(
|
||||
"SELECT value FROM Settings WHERE key = 'count'",
|
||||
[],
|
||||
|row| row.get::<_, i64>(0),
|
||||
)
|
||||
.map(|_| ())
|
||||
.map_err(|_| Error::Unavailable)
|
||||
}
|
||||
}
|
||||
|
||||
fn default_settings() -> HashMap<String, Value> {
|
||||
HashMap::from([
|
||||
("statistics".to_string(), Value::Integer(0)),
|
||||
("tag_index".to_string(), Value::Integer(0)),
|
||||
(
|
||||
"eviction_policy".to_string(),
|
||||
Value::Text("least-recently-stored".to_string()),
|
||||
),
|
||||
("size_limit".to_string(), Value::Integer(DEFAULT_SIZE_LIMIT)),
|
||||
("cull_limit".to_string(), Value::Integer(DEFAULT_CULL_LIMIT)),
|
||||
("sqlite_auto_vacuum".to_string(), Value::Integer(1)),
|
||||
("sqlite_cache_size".to_string(), Value::Integer(8192)),
|
||||
(
|
||||
"sqlite_journal_mode".to_string(),
|
||||
Value::Text("wal".to_string()),
|
||||
),
|
||||
(
|
||||
"sqlite_mmap_size".to_string(),
|
||||
Value::Integer(2_i64.pow(26)),
|
||||
),
|
||||
("sqlite_synchronous".to_string(), Value::Integer(1)),
|
||||
(
|
||||
"disk_min_file_size".to_string(),
|
||||
Value::Integer(DEFAULT_DISK_MIN_FILE_SIZE),
|
||||
),
|
||||
("disk_pickle_protocol".to_string(), Value::Integer(5)),
|
||||
])
|
||||
}
|
||||
|
||||
fn read_settings(connection: &Connection) -> Result<HashMap<String, Value>, Error> {
|
||||
let mut statement = match connection.prepare("SELECT key, value FROM Settings") {
|
||||
Ok(statement) => statement,
|
||||
Err(_) => return Ok(HashMap::new()),
|
||||
};
|
||||
statement
|
||||
.query_map([], |row| Ok((row.get(0)?, row.get(1)?)))
|
||||
.map_err(|_| Error::Unavailable)?
|
||||
.collect::<Result<HashMap<_, _>, _>>()
|
||||
.map_err(|_| Error::Unavailable)
|
||||
}
|
||||
|
||||
fn apply_pragma(connection: &Connection, key: &str, value: &Value) -> Result<(), Error> {
|
||||
let pragma = key.strip_prefix("sqlite_").ok_or(Error::Unavailable)?;
|
||||
match value {
|
||||
Value::Integer(value) => connection
|
||||
.pragma_update(None, pragma, value)
|
||||
.map_err(|_| Error::Unavailable),
|
||||
Value::Text(value) => connection
|
||||
.pragma_update(None, pragma, value)
|
||||
.map_err(|_| Error::Unavailable),
|
||||
_ => Err(Error::Unavailable),
|
||||
}
|
||||
}
|
||||
|
||||
fn setting_i64(settings: &HashMap<String, Value>, key: &str) -> Option<i64> {
|
||||
match settings.get(key) {
|
||||
Some(Value::Integer(value)) => Some(*value),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn setting_string(settings: &HashMap<String, Value>, key: &str) -> Option<String> {
|
||||
match settings.get(key) {
|
||||
Some(Value::Text(value)) => Some(value.clone()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn has_get_update(policy: &str) -> bool {
|
||||
matches!(policy, "least-recently-used" | "least-frequently-used")
|
||||
}
|
||||
|
||||
fn row_from_query(row: &rusqlite::Row<'_>) -> rusqlite::Result<Row> {
|
||||
Ok(Row {
|
||||
rowid: row.get(0)?,
|
||||
mode: row.get(2)?,
|
||||
filename: row.get(3)?,
|
||||
value: row.get(4)?,
|
||||
})
|
||||
}
|
||||
|
||||
fn fetch_row(directory: &Path, row: Row) -> Result<Option<StoredValue>, Error> {
|
||||
match row.mode {
|
||||
MODE_RAW => match row.value {
|
||||
Value::Blob(value) => Ok(Some(StoredValue::Bytes(value))),
|
||||
Value::Text(value) => Ok(Some(StoredValue::Text(value))),
|
||||
Value::Integer(value) => Ok(Some(StoredValue::Integer(value))),
|
||||
Value::Real(value) => Ok(Some(StoredValue::Float(value))),
|
||||
Value::Null => Err(Error::InvalidEntry),
|
||||
},
|
||||
MODE_BINARY | MODE_PICKLE => {
|
||||
let bytes = match row.value {
|
||||
Value::Blob(value) => value,
|
||||
Value::Null => {
|
||||
let Some(value) = read_file(directory, row.filename.as_deref())? else {
|
||||
return Ok(None);
|
||||
};
|
||||
value
|
||||
}
|
||||
_ => return Err(Error::InvalidEntry),
|
||||
};
|
||||
Ok(Some(if row.mode == MODE_BINARY {
|
||||
StoredValue::Bytes(bytes)
|
||||
} else {
|
||||
StoredValue::Pickle(bytes)
|
||||
}))
|
||||
}
|
||||
MODE_TEXT => {
|
||||
let bytes = match row.value {
|
||||
Value::Null => {
|
||||
let Some(value) = read_file(directory, row.filename.as_deref())? else {
|
||||
return Ok(None);
|
||||
};
|
||||
value
|
||||
}
|
||||
Value::Blob(value) => value,
|
||||
Value::Text(value) => value.into_bytes(),
|
||||
_ => return Err(Error::InvalidEntry),
|
||||
};
|
||||
Ok(Some(StoredValue::Text(
|
||||
String::from_utf8(bytes).map_err(|_| Error::InvalidEntry)?,
|
||||
)))
|
||||
}
|
||||
_ => Err(Error::InvalidEntry),
|
||||
}
|
||||
}
|
||||
|
||||
fn read_file(directory: &Path, filename: Option<&str>) -> Result<Option<Vec<u8>>, Error> {
|
||||
let Some(filename) = filename else {
|
||||
return Err(Error::InvalidEntry);
|
||||
};
|
||||
match fs::read(directory.join(filename)) {
|
||||
Ok(value) => Ok(Some(value)),
|
||||
Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(None),
|
||||
Err(_) => Err(Error::Unavailable),
|
||||
}
|
||||
}
|
||||
|
||||
fn store_value(
|
||||
directory: &Path,
|
||||
min_file_size: usize,
|
||||
value: StoredValue,
|
||||
) -> Result<StoredColumns, Error> {
|
||||
match value {
|
||||
StoredValue::Integer(value) => Ok(StoredColumns {
|
||||
size: 0,
|
||||
mode: MODE_RAW,
|
||||
filename: None,
|
||||
value: Some(Value::Integer(value)),
|
||||
}),
|
||||
StoredValue::Float(value) => Ok(StoredColumns {
|
||||
size: 0,
|
||||
mode: MODE_RAW,
|
||||
filename: None,
|
||||
value: Some(Value::Real(value)),
|
||||
}),
|
||||
StoredValue::Text(value) if value.chars().count() < min_file_size => Ok(StoredColumns {
|
||||
size: 0,
|
||||
mode: MODE_RAW,
|
||||
filename: None,
|
||||
value: Some(Value::Text(value)),
|
||||
}),
|
||||
StoredValue::Text(value) => {
|
||||
let bytes = value.into_bytes();
|
||||
let filename = write_file(directory, &bytes)?;
|
||||
Ok(StoredColumns {
|
||||
size: i64::try_from(bytes.len()).map_err(|_| Error::Unavailable)?,
|
||||
mode: MODE_TEXT,
|
||||
filename: Some(filename),
|
||||
value: None,
|
||||
})
|
||||
}
|
||||
StoredValue::Bytes(value) if value.len() < min_file_size => Ok(StoredColumns {
|
||||
size: 0,
|
||||
mode: MODE_RAW,
|
||||
filename: None,
|
||||
value: Some(Value::Blob(value)),
|
||||
}),
|
||||
StoredValue::Bytes(value) => {
|
||||
let filename = write_file(directory, &value)?;
|
||||
Ok(StoredColumns {
|
||||
size: i64::try_from(value.len()).map_err(|_| Error::Unavailable)?,
|
||||
mode: MODE_BINARY,
|
||||
filename: Some(filename),
|
||||
value: None,
|
||||
})
|
||||
}
|
||||
StoredValue::Pickle(value) if value.len() < min_file_size => Ok(StoredColumns {
|
||||
size: 0,
|
||||
mode: MODE_PICKLE,
|
||||
filename: None,
|
||||
value: Some(Value::Blob(value)),
|
||||
}),
|
||||
StoredValue::Pickle(value) => {
|
||||
let filename = write_file(directory, &value)?;
|
||||
Ok(StoredColumns {
|
||||
size: i64::try_from(value.len()).map_err(|_| Error::Unavailable)?,
|
||||
mode: MODE_PICKLE,
|
||||
filename: Some(filename),
|
||||
value: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn write_file(directory: &Path, bytes: &[u8]) -> Result<String, Error> {
|
||||
let mut random = [0_u8; 16];
|
||||
rand::rngs::OsRng.fill_bytes(&mut random);
|
||||
let hex = random
|
||||
.iter()
|
||||
.map(|byte| format!("{byte:02x}"))
|
||||
.collect::<String>();
|
||||
let filename = format!("{}/{}/{}.val", &hex[..2], &hex[2..4], &hex[4..]);
|
||||
let path = directory.join(&filename);
|
||||
if let Some(parent) = path.parent() {
|
||||
fs::create_dir_all(parent).map_err(|_| Error::Unavailable)?;
|
||||
}
|
||||
let mut file = OpenOptions::new()
|
||||
.write(true)
|
||||
.create_new(true)
|
||||
.open(path)
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
file.write_all(bytes).map_err(|_| Error::Unavailable)?;
|
||||
Ok(filename)
|
||||
}
|
||||
|
||||
fn cleanup_files(directory: &Path, filenames: Vec<String>) {
|
||||
for filename in filenames {
|
||||
remove_file(directory, &filename);
|
||||
}
|
||||
}
|
||||
|
||||
fn remove_file(directory: &Path, filename: &str) {
|
||||
let path = directory.join(filename);
|
||||
let _ = fs::remove_file(&path);
|
||||
}
|
||||
|
||||
fn transactional<T>(
|
||||
connection: &Connection,
|
||||
operation: impl FnOnce(&Connection) -> Result<T, Error>,
|
||||
) -> Result<T, Error> {
|
||||
connection
|
||||
.execute_batch("BEGIN IMMEDIATE")
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
match operation(connection) {
|
||||
Ok(value) => {
|
||||
connection
|
||||
.execute_batch("COMMIT")
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
Ok(value)
|
||||
}
|
||||
Err(error) => {
|
||||
let _ = connection.execute_batch("ROLLBACK");
|
||||
Err(error)
|
||||
}
|
||||
}
|
||||
}
|
||||
33
litellm-rust/crates/cache-disk/src/store.rs
Normal file
33
litellm-rust/crates/cache-disk/src/store.rs
Normal file
|
|
@ -0,0 +1,33 @@
|
|||
use std::path::Path;
|
||||
|
||||
use litellm_cache::Error;
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub enum StoredValue {
|
||||
Bytes(Vec<u8>),
|
||||
Text(String),
|
||||
Integer(i64),
|
||||
Float(f64),
|
||||
Pickle(Vec<u8>),
|
||||
}
|
||||
|
||||
pub trait DiskStore: Send + Sync + 'static {
|
||||
fn directory(&self) -> &Path;
|
||||
fn get(&self, key: &str, now: f64) -> Result<Option<StoredValue>, Error>;
|
||||
fn set(
|
||||
&self,
|
||||
key: &str,
|
||||
value: StoredValue,
|
||||
expire_time: Option<f64>,
|
||||
now: f64,
|
||||
) -> Result<(), Error>;
|
||||
fn pop(&self, key: &str, now: f64) -> Result<Option<StoredValue>, Error>;
|
||||
fn clear(&self) -> Result<(), Error>;
|
||||
fn update(
|
||||
&self,
|
||||
key: &str,
|
||||
now: f64,
|
||||
apply: &mut dyn FnMut(Option<StoredValue>) -> Result<(StoredValue, Option<f64>), Error>,
|
||||
) -> Result<(), Error>;
|
||||
fn probe(&self) -> Result<(), Error>;
|
||||
}
|
||||
431
litellm-rust/crates/cache-disk/tests/cache.rs
Normal file
431
litellm-rust/crates/cache-disk/tests/cache.rs
Normal file
|
|
@ -0,0 +1,431 @@
|
|||
use std::{
|
||||
fs,
|
||||
path::{Path, PathBuf},
|
||||
sync::Arc,
|
||||
thread,
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use litellm_cache::{
|
||||
BaseCache, BatchCache, BatchEntry, CacheCodec, CounterCache, DeleteCache, ExactCacheContext,
|
||||
FlushCache, JsonCodec,
|
||||
};
|
||||
use litellm_cache_disk::{DiskCache, DiskStore, DiskcacheSqliteStore, StoredValue, ValueAdapter};
|
||||
use rstest::{fixture, rstest};
|
||||
use rusqlite::Connection;
|
||||
use serde_json::{Value, json};
|
||||
use tempfile::TempDir;
|
||||
|
||||
struct Sandbox {
|
||||
directory: TempDir,
|
||||
}
|
||||
|
||||
#[fixture]
|
||||
fn sandbox() -> Sandbox {
|
||||
Sandbox {
|
||||
directory: tempfile::tempdir().unwrap(),
|
||||
}
|
||||
}
|
||||
|
||||
impl Sandbox {
|
||||
fn store(&self) -> DiskcacheSqliteStore {
|
||||
DiskcacheSqliteStore::open(self.directory.path()).unwrap()
|
||||
}
|
||||
|
||||
fn cache<V>(&self) -> DiskCache<JsonCodec<V>>
|
||||
where
|
||||
JsonCodec<V>: CacheCodec,
|
||||
{
|
||||
DiskCache::open(self.directory.path(), JsonCodec::new()).unwrap()
|
||||
}
|
||||
|
||||
fn db(&self) -> Connection {
|
||||
Connection::open(self.directory.path().join("cache.db")).unwrap()
|
||||
}
|
||||
|
||||
fn value_files(&self) -> Vec<PathBuf> {
|
||||
fn visit(directory: &Path, files: &mut Vec<PathBuf>) {
|
||||
for entry in fs::read_dir(directory).unwrap() {
|
||||
let path = entry.unwrap().path();
|
||||
if path.is_dir() {
|
||||
visit(&path, files);
|
||||
} else if path.extension().is_some_and(|extension| extension == "val") {
|
||||
files.push(path);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let mut files = Vec::new();
|
||||
visit(self.directory.path(), &mut files);
|
||||
files
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn relative_store_directory_is_absolutized(sandbox: Sandbox) {
|
||||
let relative = PathBuf::from(format!(
|
||||
".litellm-cache-disk-{}",
|
||||
sandbox
|
||||
.directory
|
||||
.path()
|
||||
.file_name()
|
||||
.unwrap()
|
||||
.to_string_lossy()
|
||||
));
|
||||
let store = DiskcacheSqliteStore::open(&relative).unwrap();
|
||||
assert!(store.directory().is_absolute());
|
||||
assert!(store.directory().ends_with(&relative));
|
||||
let directory = store.directory().to_path_buf();
|
||||
drop(store);
|
||||
fs::remove_dir_all(directory).unwrap();
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Default)]
|
||||
struct TextAdapter;
|
||||
|
||||
impl ValueAdapter for TextAdapter {
|
||||
fn read(&self, value: StoredValue) -> Result<Option<Vec<u8>>, litellm_cache::Error> {
|
||||
match value {
|
||||
StoredValue::Text(value) => Ok(Some(value.into_bytes())),
|
||||
_ => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
fn write(&self, payload: Vec<u8>) -> StoredValue {
|
||||
StoredValue::Text(String::from_utf8(payload).unwrap())
|
||||
}
|
||||
|
||||
fn counter_seed(&self, _: Option<StoredValue>) -> Result<f64, litellm_cache::Error> {
|
||||
Ok(0.0)
|
||||
}
|
||||
|
||||
fn counter_value(&self, value: f64) -> StoredValue {
|
||||
if value.fract() == 0.0 {
|
||||
StoredValue::Integer(value as i64)
|
||||
} else {
|
||||
StoredValue::Float(value)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn roundtrip_persists_and_reopens(sandbox: Sandbox) {
|
||||
let context = ExactCacheContext::default();
|
||||
let opened = sandbox.cache::<Value>();
|
||||
opened
|
||||
.set_cache("key", json!({"answer": 42}), &context)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
opened.get_cache("key", &context).unwrap(),
|
||||
Some(json!({"answer": 42}))
|
||||
);
|
||||
drop(opened);
|
||||
let reopened = sandbox.cache::<Value>();
|
||||
assert_eq!(
|
||||
reopened.get_cache("key", &context).unwrap(),
|
||||
Some(json!({"answer": 42}))
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn ttl_and_expired_culling_match_cache_contract(sandbox: Sandbox) {
|
||||
let store = sandbox.store();
|
||||
store
|
||||
.set(
|
||||
"expired",
|
||||
StoredValue::Bytes(b"old".to_vec()),
|
||||
Some(10.0),
|
||||
0.0,
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(store.get("expired", 10.0).unwrap(), None);
|
||||
store
|
||||
.set("new", StoredValue::Bytes(b"new".to_vec()), None, 11.0)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
sandbox
|
||||
.db()
|
||||
.query_row("SELECT COUNT(*) FROM Cache", [], |row| row.get::<_, i64>(0))
|
||||
.unwrap(),
|
||||
1
|
||||
);
|
||||
assert_eq!(
|
||||
sandbox
|
||||
.db()
|
||||
.query_row(
|
||||
"SELECT value FROM Settings WHERE key = 'count'",
|
||||
[],
|
||||
|row| row.get::<_, i64>(0)
|
||||
)
|
||||
.unwrap(),
|
||||
1
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn batch_preserves_order_and_classifies_misses_and_invalid_values(sandbox: Sandbox) {
|
||||
let store = sandbox.store();
|
||||
store
|
||||
.set(
|
||||
"hit",
|
||||
StoredValue::Bytes(br#"{"ok":true}"#.to_vec()),
|
||||
None,
|
||||
0.0,
|
||||
)
|
||||
.unwrap();
|
||||
store
|
||||
.set(
|
||||
"invalid",
|
||||
StoredValue::Pickle(vec![0x80, 0x05, 0x2e]),
|
||||
None,
|
||||
0.0,
|
||||
)
|
||||
.unwrap();
|
||||
let entries = sandbox
|
||||
.cache::<Value>()
|
||||
.batch_get_cache(
|
||||
&["hit".into(), "missing".into(), "invalid".into()],
|
||||
&ExactCacheContext::default(),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
entries,
|
||||
vec![
|
||||
BatchEntry::Hit(json!({"ok": true})),
|
||||
BatchEntry::Miss,
|
||||
BatchEntry::Invalid
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case(StoredValue::Bytes(Vec::new()))]
|
||||
#[case(StoredValue::Text(String::new()))]
|
||||
#[case(StoredValue::Integer(0))]
|
||||
#[case(StoredValue::Float(0.0))]
|
||||
#[case(StoredValue::Pickle(vec![0x80, 0x05, 0x4e, 0x2e]))]
|
||||
#[case(StoredValue::Pickle(vec![0x80, 0x05, 0x89, 0x2e]))]
|
||||
#[case(StoredValue::Pickle(vec![0x80, 0x05, 0x4b, 0x00, 0x2e]))]
|
||||
#[case(StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x0a, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x47, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x2e]))]
|
||||
#[case(StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x04, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x8c, 0x00, 0x94, 0x2e]))]
|
||||
#[case(StoredValue::Pickle(vec![0x80, 0x05, 0x5d, 0x94, 0x2e]))]
|
||||
#[case(StoredValue::Pickle(vec![0x80, 0x05, 0x7d, 0x94, 0x2e]))]
|
||||
#[case(StoredValue::Pickle(vec![0x80, 0x05, 0x29, 0x2e]))]
|
||||
fn falsy_values_are_misses(sandbox: Sandbox, #[case] value: StoredValue) {
|
||||
sandbox.store().set("key", value, None, 0.0).unwrap();
|
||||
assert_eq!(
|
||||
sandbox
|
||||
.cache::<Value>()
|
||||
.get_cache("key", &ExactCacheContext::default())
|
||||
.unwrap(),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case(Some(StoredValue::Integer(2)), 1.5, 3.5, "real")]
|
||||
#[case(Some(StoredValue::Integer(2)), 1.0, 3.0, "integer")]
|
||||
#[case(Some(StoredValue::Float(3.5)), 1.0, 1.0, "integer")]
|
||||
#[case(Some(StoredValue::Text("not a number".into())), 2.0, 2.0, "integer")]
|
||||
#[case(Some(StoredValue::Text("5".into())), 2.0, 7.0, "integer")]
|
||||
#[case(Some(StoredValue::Text("3.5".into())), 2.0, 2.0, "integer")]
|
||||
#[case(Some(StoredValue::Pickle(vec![0x80, 0x05, 0x88, 0x2e])), 1.0, 2.0, "integer")]
|
||||
#[case(Some(StoredValue::Pickle(vec![0x80, 0x05, 0x4b, 0x02, 0x2e])), 1.0, 3.0, "integer")]
|
||||
#[case(Some(StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x0a, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x7d, 0x94, 0x8c, 0x01, 0x61, 0x94, 0x4b, 0x01, 0x73, 0x2e])), 1.0, 1.0, "integer")]
|
||||
#[case(Some(StoredValue::Pickle(vec![0x80, 0x05, 0x4e, 0x2e])), 4.0, 4.0, "integer")]
|
||||
fn counters_follow_python_initialization(
|
||||
sandbox: Sandbox,
|
||||
#[case] initial: Option<StoredValue>,
|
||||
#[case] amount: f64,
|
||||
#[case] expected: f64,
|
||||
#[case] sqlite_type: &str,
|
||||
) {
|
||||
if let Some(initial) = initial {
|
||||
sandbox.store().set("counter", initial, None, 0.0).unwrap();
|
||||
}
|
||||
let cache = sandbox.cache::<f64>();
|
||||
assert_eq!(
|
||||
cache
|
||||
.increment_cache("counter", amount, ExactCacheContext::default())
|
||||
.unwrap(),
|
||||
expected
|
||||
);
|
||||
assert_eq!(
|
||||
sandbox
|
||||
.db()
|
||||
.query_row(
|
||||
"SELECT typeof(value) FROM Cache WHERE key = 'counter'",
|
||||
[],
|
||||
|row| row.get::<_, String>(0)
|
||||
)
|
||||
.unwrap(),
|
||||
sqlite_type
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn counters_are_atomic_across_concurrent_callers(sandbox: Sandbox) {
|
||||
let cache = Arc::new(sandbox.cache::<f64>());
|
||||
let workers = (0..8)
|
||||
.map(|_| {
|
||||
let cache = Arc::clone(&cache);
|
||||
thread::spawn(move || {
|
||||
for _ in 0..25 {
|
||||
cache
|
||||
.increment_cache("counter", 1.0, ExactCacheContext::default())
|
||||
.unwrap();
|
||||
}
|
||||
})
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
for worker in workers {
|
||||
worker.join().unwrap();
|
||||
}
|
||||
assert_eq!(
|
||||
cache
|
||||
.increment_cache("counter", 0.0, ExactCacheContext::default())
|
||||
.unwrap(),
|
||||
200.0
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn fractional_then_integer_increment_follows_python_behavior(sandbox: Sandbox) {
|
||||
let cache = sandbox.cache::<f64>();
|
||||
assert_eq!(
|
||||
cache
|
||||
.increment_cache("counter", 3.5, ExactCacheContext::default())
|
||||
.unwrap(),
|
||||
3.5
|
||||
);
|
||||
assert_eq!(
|
||||
cache
|
||||
.increment_cache("counter", 1.0, ExactCacheContext::default())
|
||||
.unwrap(),
|
||||
1.0
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn increment_ttl_replacement_clears_expiry_without_ttl(sandbox: Sandbox) {
|
||||
let cache = sandbox.cache::<f64>();
|
||||
cache
|
||||
.increment_cache(
|
||||
"counter",
|
||||
1.0,
|
||||
ExactCacheContext {
|
||||
ttl: Some(Duration::from_secs(60)),
|
||||
},
|
||||
)
|
||||
.unwrap();
|
||||
assert!(
|
||||
sandbox
|
||||
.db()
|
||||
.query_row(
|
||||
"SELECT expire_time IS NOT NULL FROM Cache WHERE key = 'counter'",
|
||||
[],
|
||||
|row| row.get::<_, bool>(0)
|
||||
)
|
||||
.unwrap()
|
||||
);
|
||||
cache
|
||||
.increment_cache("counter", 1.0, ExactCacheContext::default())
|
||||
.unwrap();
|
||||
assert!(
|
||||
!sandbox
|
||||
.db()
|
||||
.query_row(
|
||||
"SELECT expire_time IS NOT NULL FROM Cache WHERE key = 'counter'",
|
||||
[],
|
||||
|row| row.get::<_, bool>(0)
|
||||
)
|
||||
.unwrap()
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn custom_adapter_controls_storage_and_reads(sandbox: Sandbox) {
|
||||
let cache = DiskCache::with_adapter(sandbox.store(), TextAdapter, JsonCodec::<Value>::new());
|
||||
cache
|
||||
.set_cache("key", json!({"answer": 42}), &ExactCacheContext::default())
|
||||
.unwrap();
|
||||
assert!(matches!(
|
||||
sandbox.store().get("key", 0.0).unwrap(),
|
||||
Some(StoredValue::Text(_))
|
||||
));
|
||||
assert_eq!(
|
||||
cache
|
||||
.get_cache("key", &ExactCacheContext::default())
|
||||
.unwrap(),
|
||||
Some(json!({"answer": 42}))
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn delete_flush_and_spilled_file_replacement_clean_up_storage(sandbox: Sandbox) {
|
||||
let large = vec![b'x'; 32 * 1024];
|
||||
sandbox
|
||||
.store()
|
||||
.set("large", StoredValue::Bytes(large.clone()), None, 0.0)
|
||||
.unwrap();
|
||||
assert_eq!(sandbox.value_files().len(), 1);
|
||||
sandbox
|
||||
.store()
|
||||
.set(
|
||||
"large",
|
||||
StoredValue::Bytes(vec![b'y'; 32 * 1024]),
|
||||
None,
|
||||
0.0,
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(sandbox.value_files().len(), 1);
|
||||
sandbox.store().pop("large", 0.0).unwrap();
|
||||
assert!(sandbox.value_files().is_empty());
|
||||
sandbox
|
||||
.store()
|
||||
.set("a", StoredValue::Bytes(large.clone()), None, 0.0)
|
||||
.unwrap();
|
||||
sandbox
|
||||
.store()
|
||||
.set("b", StoredValue::Bytes(large), None, 0.0)
|
||||
.unwrap();
|
||||
sandbox.store().clear().unwrap();
|
||||
assert!(sandbox.value_files().is_empty());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn async_operations_connection_and_delete_match_sync_operations(sandbox: Sandbox) {
|
||||
let cache = sandbox.cache::<Value>();
|
||||
let context = ExactCacheContext {
|
||||
ttl: Some(Duration::from_secs(60)),
|
||||
};
|
||||
cache
|
||||
.async_set_cache("a", json!(1), context.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
cache
|
||||
.async_set_cache_pipeline(
|
||||
vec![("b".into(), json!(2)), ("c".into(), json!(3))],
|
||||
context.clone(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
cache.async_get_cache("a", &context).await.unwrap(),
|
||||
Some(json!(1))
|
||||
);
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_batch_get_cache(vec!["c".into(), "missing".into()], context.clone())
|
||||
.await
|
||||
.unwrap(),
|
||||
vec![BatchEntry::Hit(json!(3)), BatchEntry::Miss]
|
||||
);
|
||||
cache.async_delete_cache("a").await.unwrap();
|
||||
cache.async_flush_cache().await.unwrap();
|
||||
assert_eq!(
|
||||
cache.test_connection().await.unwrap().status,
|
||||
litellm_cache::CacheConnectionStatus::Success
|
||||
);
|
||||
}
|
||||
113
litellm-rust/crates/cache-disk/tests/python_compat.rs
Normal file
113
litellm-rust/crates/cache-disk/tests/python_compat.rs
Normal file
|
|
@ -0,0 +1,113 @@
|
|||
use litellm_cache::Error;
|
||||
use litellm_cache_disk::{PythonDiskCacheAdapter, StoredValue, ValueAdapter};
|
||||
use rstest::rstest;
|
||||
|
||||
enum ReadExpectation {
|
||||
Bytes(&'static [u8]),
|
||||
Miss,
|
||||
Invalid,
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::pickled_dictionary_with_string_keys(
|
||||
StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x0a, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x7d, 0x94, 0x8c, 0x01, 0x61, 0x94, 0x4b, 0x01, 0x73, 0x2e]),
|
||||
ReadExpectation::Bytes(br#"{"a":1}"#)
|
||||
)]
|
||||
#[case::pickled_list_of_integers(
|
||||
StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x09, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x5d, 0x94, 0x28, 0x4b, 0x01, 0x4b, 0x02, 0x65, 0x2e]),
|
||||
ReadExpectation::Bytes(br#"[1,2]"#)
|
||||
)]
|
||||
#[case::pickled_tuple_of_integers(
|
||||
StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x07, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x4b, 0x01, 0x4b, 0x02, 0x86, 0x94, 0x2e]),
|
||||
ReadExpectation::Bytes(br#"[1,2]"#)
|
||||
)]
|
||||
#[case::pickled_set_of_integers(
|
||||
StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x09, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x8f, 0x94, 0x28, 0x4b, 0x01, 0x4b, 0x02, 0x90, 0x2e]),
|
||||
ReadExpectation::Bytes(br#"[1,2]"#)
|
||||
)]
|
||||
#[case::pickled_response_envelope(
|
||||
StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x30, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x7d, 0x94, 0x28, 0x8c, 0x09, 0x74, 0x69, 0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d, 0x70, 0x94, 0x47, 0x3f, 0xf8, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x8c, 0x08, 0x72, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x94, 0x8c, 0x08, 0x7b, 0x22, 0x61, 0x22, 0x3a, 0x20, 0x31, 0x7d, 0x94, 0x75, 0x2e]),
|
||||
ReadExpectation::Bytes(br#"{"response":"{\"a\": 1}","timestamp":1.5}"#)
|
||||
)]
|
||||
#[case::non_json_text(
|
||||
StoredValue::Text("not json".into()),
|
||||
ReadExpectation::Bytes(b"not json")
|
||||
)]
|
||||
#[case::json_text(
|
||||
StoredValue::Text("{\"a\": 1}".into()),
|
||||
ReadExpectation::Bytes(br#"{"a": 1}"#)
|
||||
)]
|
||||
#[case::non_utf8_bytes(
|
||||
StoredValue::Bytes(vec![0xff, 0xfe]),
|
||||
ReadExpectation::Bytes(&[0xff, 0xfe])
|
||||
)]
|
||||
#[case::integer_seven(StoredValue::Integer(7), ReadExpectation::Bytes(b"7"))]
|
||||
#[case::float_one_point_five(StoredValue::Float(1.5), ReadExpectation::Bytes(b"1.5"))]
|
||||
#[case::pickled_true(
|
||||
StoredValue::Pickle(vec![0x80, 0x05, 0x88, 0x2e]),
|
||||
ReadExpectation::Bytes(b"true")
|
||||
)]
|
||||
#[case::pickled_negative_integer(
|
||||
StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x06, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x4a, 0xfd, 0xff, 0xff, 0xff, 0x2e]),
|
||||
ReadExpectation::Bytes(b"-3")
|
||||
)]
|
||||
#[case::pickled_bytes(
|
||||
StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x09, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x43, 0x05, 0x62, 0x79, 0x74, 0x65, 0x73, 0x94, 0x2e]),
|
||||
ReadExpectation::Invalid
|
||||
)]
|
||||
#[case::pickled_dictionary_with_integer_key(
|
||||
StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x0a, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x7d, 0x94, 0x4b, 0x01, 0x8c, 0x01, 0x61, 0x94, 0x73, 0x2e]),
|
||||
ReadExpectation::Invalid
|
||||
)]
|
||||
#[case::pickled_complex(
|
||||
StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x2e, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x8c, 0x08, 0x62, 0x75, 0x69, 0x6c, 0x74, 0x69, 0x6e, 0x73, 0x94, 0x8c, 0x07, 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x94, 0x93, 0x94, 0x47, 0x3f, 0xf0, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x47, 0x40, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x86, 0x94, 0x52, 0x94, 0x2e]),
|
||||
ReadExpectation::Invalid
|
||||
)]
|
||||
#[case::truncated_pickle(
|
||||
StoredValue::Pickle(vec![0x80, 0x05, 0x2e]),
|
||||
ReadExpectation::Invalid
|
||||
)]
|
||||
#[case::empty_bytes(StoredValue::Bytes(Vec::new()), ReadExpectation::Miss)]
|
||||
#[case::empty_text(StoredValue::Text(String::new()), ReadExpectation::Miss)]
|
||||
#[case::zero_integer(StoredValue::Integer(0), ReadExpectation::Miss)]
|
||||
#[case::zero_float(StoredValue::Float(0.0), ReadExpectation::Miss)]
|
||||
#[case::pickled_none(StoredValue::Pickle(vec![0x80, 0x05, 0x4e, 0x2e]), ReadExpectation::Miss)]
|
||||
#[case::pickled_false(StoredValue::Pickle(vec![0x80, 0x05, 0x89, 0x2e]), ReadExpectation::Miss)]
|
||||
#[case::pickled_zero(StoredValue::Pickle(vec![0x80, 0x05, 0x4b, 0x00, 0x2e]), ReadExpectation::Miss)]
|
||||
#[case::pickled_zero_float(StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x0a, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x47, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x2e]), ReadExpectation::Miss)]
|
||||
#[case::pickled_empty_string(StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x04, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x8c, 0x00, 0x94, 0x2e]), ReadExpectation::Miss)]
|
||||
#[case::pickled_empty_list(StoredValue::Pickle(vec![0x80, 0x05, 0x5d, 0x94, 0x2e]), ReadExpectation::Miss)]
|
||||
#[case::pickled_empty_dictionary(StoredValue::Pickle(vec![0x80, 0x05, 0x7d, 0x94, 0x2e]), ReadExpectation::Miss)]
|
||||
#[case::pickled_empty_tuple(StoredValue::Pickle(vec![0x80, 0x05, 0x29, 0x2e]), ReadExpectation::Miss)]
|
||||
fn python_read_cases(#[case] row: StoredValue, #[case] expected: ReadExpectation) {
|
||||
let result = PythonDiskCacheAdapter.read(row);
|
||||
match expected {
|
||||
ReadExpectation::Bytes(expected) => assert_eq!(result.unwrap().unwrap(), expected),
|
||||
ReadExpectation::Miss => assert_eq!(result.unwrap(), None),
|
||||
ReadExpectation::Invalid => assert!(matches!(result, Err(Error::InvalidEntry))),
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::integer_two(Some(StoredValue::Integer(2)), 2.0)]
|
||||
#[case::float_three_point_five(Some(StoredValue::Float(3.5)), 0.0)]
|
||||
#[case::text_not_a_number(Some(StoredValue::Text("not a number".into())), 0.0)]
|
||||
#[case::text_five(Some(StoredValue::Text("5".into())), 5.0)]
|
||||
#[case::text_three_point_five(Some(StoredValue::Text("3.5".into())), 0.0)]
|
||||
#[case::pickled_true(Some(StoredValue::Pickle(vec![0x80, 0x05, 0x88, 0x2e])), 1.0)]
|
||||
#[case::pickled_two(Some(StoredValue::Pickle(vec![0x80, 0x05, 0x4b, 0x02, 0x2e])), 2.0)]
|
||||
#[case::pickled_dictionary(Some(StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x0a, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x7d, 0x94, 0x8c, 0x01, 0x61, 0x94, 0x4b, 0x01, 0x73, 0x2e])), 0.0)]
|
||||
#[case::missing(None, 0.0)]
|
||||
#[case::pickled_none(Some(StoredValue::Pickle(vec![0x80, 0x05, 0x4e, 0x2e])), 0.0)]
|
||||
fn python_counter_seed_cases(#[case] row: Option<StoredValue>, #[case] expected: f64) {
|
||||
assert_eq!(PythonDiskCacheAdapter.counter_seed(row).unwrap(), expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::integer_three(3.0, StoredValue::Integer(3))]
|
||||
#[case::fractional_three_point_five(3.5, StoredValue::Float(3.5))]
|
||||
#[case::negative_zero(-0.0, StoredValue::Integer(0))]
|
||||
#[case::large_float(1e300, StoredValue::Float(1e300))]
|
||||
fn python_counter_value_cases(#[case] value: f64, #[case] expected: StoredValue) {
|
||||
assert_eq!(PythonDiskCacheAdapter.counter_value(value), expected);
|
||||
}
|
||||
20
litellm-rust/crates/cache-gcs/Cargo.toml
Normal file
20
litellm-rust/crates/cache-gcs/Cargo.toml
Normal file
|
|
@ -0,0 +1,20 @@
|
|||
[package]
|
||||
name = "litellm-cache-gcs"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
futures-util.workspace = true
|
||||
litellm-auth-gcp.workspace = true
|
||||
litellm-auth-types.workspace = true
|
||||
litellm-cache.workspace = true
|
||||
percent-encoding.workspace = true
|
||||
reqwest.workspace = true
|
||||
tokio.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
serde_json.workspace = true
|
||||
tokio.workspace = true
|
||||
wiremock = "0.6.5"
|
||||
260
litellm-rust/crates/cache-gcs/src/cache.rs
Normal file
260
litellm-rust/crates/cache-gcs/src/cache.rs
Normal file
|
|
@ -0,0 +1,260 @@
|
|||
use std::{future::Future, sync::Arc, time::Duration};
|
||||
|
||||
use futures_util::future::try_join_all;
|
||||
use litellm_cache::{
|
||||
BaseCache, BatchCache, BatchEntry, CacheCodec, CacheConnectionResult, Error, ExactCacheContext,
|
||||
FlushCache,
|
||||
};
|
||||
use percent_encoding::{AsciiSet, NON_ALPHANUMERIC, percent_encode};
|
||||
use reqwest::Client;
|
||||
|
||||
use crate::{GcpTokenSource, TokenSource};
|
||||
|
||||
pub const DEFAULT_ENDPOINT: &str = "https://storage.googleapis.com";
|
||||
|
||||
const OBJECT_NAME_ENCODE_SET: &AsciiSet = &NON_ALPHANUMERIC
|
||||
.remove(b'-')
|
||||
.remove(b'_')
|
||||
.remove(b'.')
|
||||
.remove(b'~');
|
||||
|
||||
pub fn key_prefix(gcs_path: Option<&str>) -> String {
|
||||
match gcs_path {
|
||||
Some(path) if !path.is_empty() => format!("{}/", path.trim_end_matches('/')),
|
||||
_ => String::new(),
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct GcsConfig {
|
||||
pub bucket_name: String,
|
||||
pub gcs_path: Option<String>,
|
||||
pub path_service_account: Option<String>,
|
||||
pub endpoint: String,
|
||||
}
|
||||
|
||||
impl GcsConfig {
|
||||
pub fn new(bucket_name: impl Into<String>) -> Self {
|
||||
Self {
|
||||
bucket_name: bucket_name.into(),
|
||||
gcs_path: None,
|
||||
path_service_account: None,
|
||||
endpoint: DEFAULT_ENDPOINT.to_string(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct GcsCache<S: CacheCodec> {
|
||||
config: GcsConfig,
|
||||
key_prefix: String,
|
||||
client: Client,
|
||||
token: Arc<dyn TokenSource>,
|
||||
codec: S,
|
||||
}
|
||||
|
||||
impl<S: CacheCodec> GcsCache<S> {
|
||||
pub fn new(config: GcsConfig, codec: S) -> Result<Self, Error> {
|
||||
let token = Arc::new(GcpTokenSource::new(config.path_service_account.clone()));
|
||||
Self::with_token_source(config, codec, token)
|
||||
}
|
||||
|
||||
pub fn with_token_source(
|
||||
config: GcsConfig,
|
||||
codec: S,
|
||||
token: Arc<dyn TokenSource>,
|
||||
) -> Result<Self, Error> {
|
||||
let client = Client::builder().build().map_err(|_| Error::Unavailable)?;
|
||||
let key_prefix = key_prefix(config.gcs_path.as_deref());
|
||||
Ok(Self {
|
||||
config,
|
||||
key_prefix,
|
||||
client,
|
||||
token,
|
||||
codec,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn bucket_name(&self) -> &str {
|
||||
&self.config.bucket_name
|
||||
}
|
||||
|
||||
pub fn key_prefix(&self) -> &str {
|
||||
&self.key_prefix
|
||||
}
|
||||
|
||||
pub fn path_service_account(&self) -> Option<&str> {
|
||||
self.config.path_service_account.as_deref()
|
||||
}
|
||||
|
||||
pub fn object_name(&self, key: &str) -> String {
|
||||
format!("{}{}", self.key_prefix, key)
|
||||
}
|
||||
|
||||
fn encoded_object_name(&self, key: &str) -> String {
|
||||
percent_encode(self.object_name(key).as_bytes(), OBJECT_NAME_ENCODE_SET).to_string()
|
||||
}
|
||||
|
||||
fn endpoint(&self, path: &str) -> String {
|
||||
format!("{}{}", self.config.endpoint.trim_end_matches('/'), path)
|
||||
}
|
||||
|
||||
async fn async_set(&self, key: &str, value: S::Value) -> Result<(), Error> {
|
||||
let token = self.token.bearer_token().await?;
|
||||
let payload = self.codec.encode(&value)?;
|
||||
let url = self.endpoint(&format!(
|
||||
"/upload/storage/v1/b/{}/o?uploadType=media&name={}",
|
||||
self.config.bucket_name,
|
||||
self.encoded_object_name(key)
|
||||
));
|
||||
let response = self
|
||||
.client
|
||||
.post(url)
|
||||
.bearer_auth(token)
|
||||
.header(reqwest::header::CONTENT_TYPE, "application/json")
|
||||
.body(payload)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
if !response.status().is_success() {
|
||||
return Err(Error::Unavailable);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn async_get(&self, key: &str) -> Result<Option<S::Value>, Error> {
|
||||
let token = self.token.bearer_token().await?;
|
||||
let url = self.endpoint(&format!(
|
||||
"/storage/v1/b/{}/o/{}?alt=media",
|
||||
self.config.bucket_name,
|
||||
self.encoded_object_name(key)
|
||||
));
|
||||
let response = self
|
||||
.client
|
||||
.get(url)
|
||||
.bearer_auth(token)
|
||||
.header(reqwest::header::CONTENT_TYPE, "application/json")
|
||||
.send()
|
||||
.await
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
if response.status() == reqwest::StatusCode::NOT_FOUND {
|
||||
return Ok(None);
|
||||
}
|
||||
if !response.status().is_success() {
|
||||
return Err(Error::Unavailable);
|
||||
}
|
||||
let body = response.bytes().await.map_err(|_| Error::Unavailable)?;
|
||||
self.codec
|
||||
.decode(&body)
|
||||
.map(Some)
|
||||
.map_err(|_| Error::InvalidEntry)
|
||||
}
|
||||
|
||||
fn run_sync<T, F>(future: F) -> Result<T, Error>
|
||||
where
|
||||
F: Future<Output = Result<T, Error>> + Send,
|
||||
T: Send,
|
||||
{
|
||||
let run = || {
|
||||
tokio::runtime::Builder::new_current_thread()
|
||||
.enable_all()
|
||||
.build()
|
||||
.map_err(|_| Error::Unavailable)
|
||||
.and_then(|runtime| runtime.block_on(future))
|
||||
};
|
||||
if let Ok(handle) = tokio::runtime::Handle::try_current() {
|
||||
if handle.runtime_flavor() == tokio::runtime::RuntimeFlavor::MultiThread {
|
||||
return tokio::task::block_in_place(run);
|
||||
}
|
||||
return std::thread::scope(|scope| {
|
||||
scope
|
||||
.spawn(run)
|
||||
.join()
|
||||
.map_err(|_| Error::Unavailable)
|
||||
.and_then(|result| result)
|
||||
});
|
||||
}
|
||||
run()
|
||||
}
|
||||
}
|
||||
|
||||
impl<S: CacheCodec> BaseCache for GcsCache<S> {
|
||||
type Value = S::Value;
|
||||
type Context = ExactCacheContext;
|
||||
|
||||
fn get_ttl(&self, _: &Self::Context) -> Option<Duration> {
|
||||
None
|
||||
}
|
||||
|
||||
fn set_cache(&self, key: &str, value: Self::Value, _: &Self::Context) -> Result<(), Error> {
|
||||
Self::run_sync(self.async_set(key, value))
|
||||
}
|
||||
|
||||
fn get_cache(&self, key: &str, _: &Self::Context) -> Result<Option<Self::Value>, Error> {
|
||||
Self::run_sync(self.async_get(key))
|
||||
}
|
||||
|
||||
async fn async_set_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
value: Self::Value,
|
||||
_: Self::Context,
|
||||
) -> Result<(), Error> {
|
||||
self.async_set(key, value).await
|
||||
}
|
||||
|
||||
async fn async_get_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
_: &Self::Context,
|
||||
) -> Result<Option<Self::Value>, Error> {
|
||||
self.async_get(key).await
|
||||
}
|
||||
|
||||
async fn async_set_cache_pipeline(
|
||||
&self,
|
||||
entries: Vec<(String, Self::Value)>,
|
||||
context: Self::Context,
|
||||
) -> Result<(), Error> {
|
||||
try_join_all(entries.into_iter().map(|(key, value)| {
|
||||
let context = context.clone();
|
||||
async move { self.async_set_cache(&key, value, context).await }
|
||||
}))
|
||||
.await
|
||||
.map(|_| ())
|
||||
}
|
||||
|
||||
async fn disconnect(&self) -> Result<(), Error> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
|
||||
Err(Error::UnsupportedOperation)
|
||||
}
|
||||
}
|
||||
|
||||
impl<S: CacheCodec> BatchCache for GcsCache<S> {
|
||||
async fn async_batch_get_cache(
|
||||
&self,
|
||||
keys: Vec<String>,
|
||||
context: Self::Context,
|
||||
) -> Result<Vec<BatchEntry<Self::Value>>, Error> {
|
||||
try_join_all(keys.into_iter().map(|key| {
|
||||
let context = context.clone();
|
||||
async move {
|
||||
match self.async_get_cache(&key, &context).await {
|
||||
Ok(Some(value)) => Ok(BatchEntry::Hit(value)),
|
||||
Ok(None) => Ok(BatchEntry::Miss),
|
||||
Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid),
|
||||
Err(error) => Err(error),
|
||||
}
|
||||
}
|
||||
}))
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
impl<S: CacheCodec> FlushCache for GcsCache<S> {
|
||||
fn flush_cache(&self) -> Result<(), Error> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
5
litellm-rust/crates/cache-gcs/src/lib.rs
Normal file
5
litellm-rust/crates/cache-gcs/src/lib.rs
Normal file
|
|
@ -0,0 +1,5 @@
|
|||
mod cache;
|
||||
mod token;
|
||||
|
||||
pub use cache::{DEFAULT_ENDPOINT, GcsCache, GcsConfig, key_prefix};
|
||||
pub use token::{GcpTokenSource, StaticTokenSource, TokenSource};
|
||||
44
litellm-rust/crates/cache-gcs/src/token.rs
Normal file
44
litellm-rust/crates/cache-gcs/src/token.rs
Normal file
|
|
@ -0,0 +1,44 @@
|
|||
use std::{future::Future, pin::Pin};
|
||||
|
||||
use litellm_auth_gcp::{VertexAuth, VertexConfig};
|
||||
use litellm_auth_types::{InputSource, SecretValue, Sourced};
|
||||
use litellm_cache::Error;
|
||||
|
||||
pub trait TokenSource: Send + Sync + 'static {
|
||||
fn bearer_token(&self) -> Pin<Box<dyn Future<Output = Result<String, Error>> + Send + '_>>;
|
||||
}
|
||||
|
||||
pub struct GcpTokenSource {
|
||||
auth: VertexAuth,
|
||||
config: VertexConfig,
|
||||
}
|
||||
|
||||
impl GcpTokenSource {
|
||||
pub fn new(path_service_account: Option<String>) -> Self {
|
||||
let credentials = path_service_account
|
||||
.map(|path| Sourced::new(SecretValue::new(path), InputSource::Deployment));
|
||||
Self {
|
||||
auth: VertexAuth::default(),
|
||||
config: VertexConfig::new(credentials, None, None),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl TokenSource for GcpTokenSource {
|
||||
fn bearer_token(&self) -> Pin<Box<dyn Future<Output = Result<String, Error>> + Send + '_>> {
|
||||
Box::pin(async move {
|
||||
self.auth
|
||||
.access_token(&self.config, &|name| std::env::var(name).ok())
|
||||
.await
|
||||
.map_err(|_| Error::Unavailable)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub struct StaticTokenSource(pub String);
|
||||
|
||||
impl TokenSource for StaticTokenSource {
|
||||
fn bearer_token(&self) -> Pin<Box<dyn Future<Output = Result<String, Error>> + Send + '_>> {
|
||||
Box::pin(async move { Ok(self.0.clone()) })
|
||||
}
|
||||
}
|
||||
324
litellm-rust/crates/cache-gcs/tests/cache.rs
Normal file
324
litellm-rust/crates/cache-gcs/tests/cache.rs
Normal file
|
|
@ -0,0 +1,324 @@
|
|||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use litellm_cache::{
|
||||
BaseCache, BatchCache, BatchEntry, CacheContext, Error, ExactCacheContext, FlushCache,
|
||||
JsonCodec,
|
||||
};
|
||||
use litellm_cache_gcs::{GcsCache, GcsConfig, StaticTokenSource, TokenSource, key_prefix};
|
||||
use serde_json::json;
|
||||
use wiremock::{
|
||||
Mock, MockServer, ResponseTemplate,
|
||||
matchers::{body_bytes, header, method, path, query_param},
|
||||
};
|
||||
|
||||
fn config(server: &MockServer, gcs_path: Option<&str>) -> GcsConfig {
|
||||
GcsConfig {
|
||||
bucket_name: "bucket".into(),
|
||||
gcs_path: gcs_path.map(str::to_string),
|
||||
path_service_account: None,
|
||||
endpoint: server.uri(),
|
||||
}
|
||||
}
|
||||
|
||||
fn cache(server: &MockServer, gcs_path: Option<&str>) -> GcsCache<JsonCodec<serde_json::Value>> {
|
||||
GcsCache::with_token_source(
|
||||
config(server, gcs_path),
|
||||
JsonCodec::new(),
|
||||
Arc::new(StaticTokenSource("tok".into())),
|
||||
)
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn set_writes_encoded_object_and_headers() {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/upload/storage/v1/b/bucket/o"))
|
||||
.and(query_param("uploadType", "media"))
|
||||
.and(header("authorization", "Bearer tok"))
|
||||
.and(header("content-type", "application/json"))
|
||||
.and(body_bytes(br#"{"value":"entry"}"#))
|
||||
.respond_with(ResponseTemplate::new(200))
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
cache(&server, Some("cache/"))
|
||||
.set_cache(
|
||||
"team:a b/c",
|
||||
json!({"value": "entry"}),
|
||||
&ExactCacheContext::default(),
|
||||
)
|
||||
.unwrap();
|
||||
let requests = server.received_requests().await.unwrap();
|
||||
assert_eq!(requests.len(), 1);
|
||||
assert_eq!(
|
||||
requests[0].url.query(),
|
||||
Some("uploadType=media&name=cache%2Fteam%3Aa%20b%2Fc")
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn get_maps_statuses_and_decode_failures() {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/storage/v1/b/bucket/o/hit"))
|
||||
.and(query_param("alt", "media"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"value": "entry"})))
|
||||
.mount(&server)
|
||||
.await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/storage/v1/b/bucket/o/missing"))
|
||||
.respond_with(ResponseTemplate::new(404))
|
||||
.mount(&server)
|
||||
.await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/storage/v1/b/bucket/o/server-error"))
|
||||
.respond_with(ResponseTemplate::new(500))
|
||||
.mount(&server)
|
||||
.await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/storage/v1/b/bucket/o/invalid"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_string("not json"))
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let cache = cache(&server, None);
|
||||
assert_eq!(
|
||||
cache
|
||||
.get_cache("hit", &ExactCacheContext::default())
|
||||
.unwrap(),
|
||||
Some(json!({"value": "entry"}))
|
||||
);
|
||||
assert_eq!(
|
||||
cache
|
||||
.get_cache("missing", &ExactCacheContext::default())
|
||||
.unwrap(),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
cache
|
||||
.get_cache("server-error", &ExactCacheContext::default())
|
||||
.unwrap_err(),
|
||||
Error::Unavailable
|
||||
);
|
||||
assert_eq!(
|
||||
cache
|
||||
.get_cache("invalid", &ExactCacheContext::default())
|
||||
.unwrap_err(),
|
||||
Error::InvalidEntry
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn key_prefix_normalizes_paths() {
|
||||
assert_eq!(key_prefix(None), "");
|
||||
assert_eq!(key_prefix(Some("a/b/")), "a/b/");
|
||||
assert_eq!(key_prefix(Some("a/b")), "a/b/");
|
||||
assert_eq!(key_prefix(Some("")), "");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn object_names_use_python_quote_encoding() {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/upload/storage/v1/b/bucket/o"))
|
||||
.and(query_param("uploadType", "media"))
|
||||
.respond_with(ResponseTemplate::new(200))
|
||||
.expect(2)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let cache = cache(&server, Some("p/"));
|
||||
cache
|
||||
.set_cache(
|
||||
"a~b-c_d.e/f g%h",
|
||||
json!({"value": "punctuation"}),
|
||||
&ExactCacheContext::default(),
|
||||
)
|
||||
.unwrap();
|
||||
cache
|
||||
.set_cache(
|
||||
"ключ",
|
||||
json!({"value": "utf8"}),
|
||||
&ExactCacheContext::default(),
|
||||
)
|
||||
.unwrap();
|
||||
let requests = server.received_requests().await.unwrap();
|
||||
let queries: Vec<_> = requests
|
||||
.iter()
|
||||
.filter_map(|request| request.url.query())
|
||||
.collect();
|
||||
assert!(queries.contains(&"uploadType=media&name=p%2Fa~b-c_d.e%2Ff%20g%25h"));
|
||||
assert!(queries.contains(&"uploadType=media&name=p%2F%D0%BA%D0%BB%D1%8E%D1%87"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn ignores_ttl_and_writes_pipeline_concurrently() {
|
||||
let server = MockServer::start().await;
|
||||
for key in ["one", "two", "three"] {
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/upload/storage/v1/b/bucket/o"))
|
||||
.and(query_param("uploadType", "media"))
|
||||
.and(query_param("name", key))
|
||||
.respond_with(ResponseTemplate::new(200))
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
}
|
||||
let cache = cache(&server, None);
|
||||
assert_eq!(cache.get_ttl(&ExactCacheContext::default()), None);
|
||||
assert_eq!(
|
||||
cache.get_ttl(&ExactCacheContext::default().with_ttl(Some(Duration::from_secs(5)))),
|
||||
None
|
||||
);
|
||||
cache
|
||||
.async_set_cache_pipeline(
|
||||
vec![
|
||||
("one".into(), json!({"key": "one"})),
|
||||
("two".into(), json!({"key": "two"})),
|
||||
("three".into(), json!({"key": "three"})),
|
||||
],
|
||||
ExactCacheContext::default().with_ttl(Some(Duration::from_secs(5))),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn async_batch_get_preserves_hits_misses_and_invalid_entries() {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/storage/v1/b/bucket/o/hit"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"value": "entry"})))
|
||||
.mount(&server)
|
||||
.await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/storage/v1/b/bucket/o/missing"))
|
||||
.respond_with(ResponseTemplate::new(404))
|
||||
.mount(&server)
|
||||
.await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/storage/v1/b/bucket/o/invalid"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_string("not json"))
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
assert_eq!(
|
||||
cache(&server, None)
|
||||
.async_batch_get_cache(
|
||||
vec!["hit".into(), "missing".into(), "invalid".into()],
|
||||
ExactCacheContext::default(),
|
||||
)
|
||||
.await
|
||||
.unwrap(),
|
||||
vec![
|
||||
BatchEntry::Hit(json!({"value": "entry"})),
|
||||
BatchEntry::Miss,
|
||||
BatchEntry::Invalid,
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn lifecycle_operations_are_noops_and_connection_test_is_unsupported() {
|
||||
let server = MockServer::start().await;
|
||||
let cache = cache(&server, None);
|
||||
assert_eq!(cache.flush_cache(), Ok(()));
|
||||
assert_eq!(cache.disconnect().await, Ok(()));
|
||||
assert_eq!(
|
||||
cache.test_connection().await,
|
||||
Err(Error::UnsupportedOperation)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sync_operations_work_without_an_active_runtime() {
|
||||
let runtime = tokio::runtime::Builder::new_multi_thread()
|
||||
.enable_all()
|
||||
.build()
|
||||
.unwrap();
|
||||
let server = runtime.block_on(MockServer::start());
|
||||
runtime.block_on(
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/upload/storage/v1/b/bucket/o"))
|
||||
.respond_with(ResponseTemplate::new(200))
|
||||
.mount(&server),
|
||||
);
|
||||
runtime.block_on(
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/storage/v1/b/bucket/o/key"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"value": "entry"})))
|
||||
.mount(&server),
|
||||
);
|
||||
let cache = cache(&server, None);
|
||||
cache
|
||||
.set_cache(
|
||||
"key",
|
||||
json!({"value": "entry"}),
|
||||
&ExactCacheContext::default(),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
cache
|
||||
.get_cache("key", &ExactCacheContext::default())
|
||||
.unwrap(),
|
||||
Some(json!({"value": "entry"}))
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn sync_operations_work_inside_a_multi_thread_runtime() {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/upload/storage/v1/b/bucket/o"))
|
||||
.respond_with(ResponseTemplate::new(200))
|
||||
.mount(&server)
|
||||
.await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/storage/v1/b/bucket/o/key"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"value": "entry"})))
|
||||
.mount(&server)
|
||||
.await;
|
||||
let cache = cache(&server, None);
|
||||
cache
|
||||
.set_cache(
|
||||
"key",
|
||||
json!({"value": "entry"}),
|
||||
&ExactCacheContext::default(),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
cache
|
||||
.get_cache("key", &ExactCacheContext::default())
|
||||
.unwrap(),
|
||||
Some(json!({"value": "entry"}))
|
||||
);
|
||||
}
|
||||
|
||||
struct FailingTokenSource;
|
||||
|
||||
impl TokenSource for FailingTokenSource {
|
||||
fn bearer_token(
|
||||
&self,
|
||||
) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<String, Error>> + Send + '_>>
|
||||
{
|
||||
Box::pin(async { Err(Error::Unavailable) })
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn token_source_failure_skips_http() {
|
||||
let server = MockServer::start().await;
|
||||
let cache = GcsCache::with_token_source(
|
||||
config(&server, None),
|
||||
JsonCodec::<serde_json::Value>::new(),
|
||||
Arc::new(FailingTokenSource),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
cache
|
||||
.get_cache("key", &ExactCacheContext::default())
|
||||
.unwrap_err(),
|
||||
Error::Unavailable
|
||||
);
|
||||
assert_eq!(server.received_requests().await.unwrap().len(), 0);
|
||||
}
|
||||
|
|
@ -7,8 +7,8 @@ repository.workspace = true
|
|||
|
||||
[dependencies]
|
||||
litellm-cache.workspace = true
|
||||
serde_json.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
serde_json.workspace = true
|
||||
rstest.workspace = true
|
||||
tokio.workspace = true
|
||||
|
|
|
|||
|
|
@ -1,18 +1,20 @@
|
|||
use std::cmp::Reverse;
|
||||
use std::collections::{BinaryHeap, HashMap};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||
use std::{
|
||||
cmp::Reverse,
|
||||
collections::{BinaryHeap, HashMap, HashSet},
|
||||
hash::Hash,
|
||||
sync::{Arc, Mutex},
|
||||
time::{Duration, SystemTime, UNIX_EPOCH},
|
||||
};
|
||||
|
||||
use litellm_cache::{
|
||||
BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheEntry, CacheFuture, CacheKwargs,
|
||||
Error,
|
||||
BaseCache, BatchCache, CacheConnectionResult, CacheConnectionStatus, ClaimCache, CounterCache,
|
||||
DeleteCache, Error, ExactCacheContext, FlushCache, IncrementOperation, SetCache, TtlCache,
|
||||
};
|
||||
|
||||
const DEFAULT_MAX_SIZE_IN_MEMORY: usize = 200;
|
||||
const DEFAULT_TTL: Duration = Duration::from_secs(600);
|
||||
|
||||
type ValueMeasure<V> = Arc<dyn Fn(&V) -> Result<usize, Error> + Send + Sync>;
|
||||
type ValueValidator<V> = Arc<dyn Fn(&V) -> Result<(), Error> + Send + Sync>;
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum CacheWrite {
|
||||
|
|
@ -33,7 +35,6 @@ pub struct InMemoryCache<V: Clone> {
|
|||
default_ttl: Duration,
|
||||
max_entry_bytes: Option<usize>,
|
||||
measure_value: Option<ValueMeasure<V>>,
|
||||
validate_value: Option<ValueValidator<V>>,
|
||||
now: Arc<dyn Fn() -> Duration + Send + Sync>,
|
||||
}
|
||||
|
||||
|
|
@ -77,7 +78,6 @@ impl<V: Clone> InMemoryCache<V> {
|
|||
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
|
||||
max_entry_bytes,
|
||||
measure_value,
|
||||
validate_value: None,
|
||||
now: Arc::new(now),
|
||||
}
|
||||
}
|
||||
|
|
@ -91,9 +91,6 @@ impl<V: Clone> InMemoryCache<V> {
|
|||
if self.max_size_in_memory == 0 {
|
||||
return Ok(CacheWrite::Disabled);
|
||||
}
|
||||
if let Some(validate) = &self.validate_value {
|
||||
validate(&value)?;
|
||||
}
|
||||
if let (Some(limit), Some(measure)) = (self.max_entry_bytes, &self.measure_value)
|
||||
&& measure(&value)? > limit
|
||||
{
|
||||
|
|
@ -101,15 +98,13 @@ impl<V: Clone> InMemoryCache<V> {
|
|||
}
|
||||
let now = (self.now)();
|
||||
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
|
||||
Self::evict(&mut state, self.max_size_in_memory, now);
|
||||
let key = key.into();
|
||||
state.values.insert(key.clone(), value);
|
||||
Self::evict(&mut state, self.max_size_in_memory, now, &key);
|
||||
let expiration = state.expirations.get(&key).copied();
|
||||
if expiration.is_none_or(|expiration| expiration < now) {
|
||||
let expiration = now + ttl.unwrap_or(self.default_ttl);
|
||||
state.expirations.insert(key.clone(), expiration);
|
||||
state.expiration_heap.push(Reverse((expiration, key)));
|
||||
Self::set_expiration(&mut state, &key, now + ttl.unwrap_or(self.default_ttl));
|
||||
}
|
||||
state.values.insert(key, value);
|
||||
Ok(CacheWrite::Stored)
|
||||
}
|
||||
|
||||
|
|
@ -126,6 +121,14 @@ impl<V: Clone> InMemoryCache<V> {
|
|||
Ok(state.values.get(key).cloned())
|
||||
}
|
||||
|
||||
pub fn max_size_in_memory(&self) -> usize {
|
||||
self.max_size_in_memory
|
||||
}
|
||||
|
||||
pub fn max_entry_bytes(&self) -> Option<usize> {
|
||||
self.max_entry_bytes
|
||||
}
|
||||
|
||||
pub fn expires_at(&self, key: &str) -> Result<Option<Duration>, Error> {
|
||||
Ok(self
|
||||
.state
|
||||
|
|
@ -136,6 +139,25 @@ impl<V: Clone> InMemoryCache<V> {
|
|||
.copied())
|
||||
}
|
||||
|
||||
pub async fn async_get_ttl(&self, key: &str) -> Result<Option<Duration>, Error> {
|
||||
self.expires_at(key)
|
||||
}
|
||||
|
||||
pub async fn async_get_oldest_n_keys(&self, count: usize) -> Result<Vec<String>, Error> {
|
||||
let state = self.state.lock().map_err(|_| Error::Unavailable)?;
|
||||
let mut expirations = state
|
||||
.expirations
|
||||
.iter()
|
||||
.map(|(key, expiration)| (key.clone(), *expiration))
|
||||
.collect::<Vec<_>>();
|
||||
expirations.sort_unstable_by_key(|(_, expiration)| *expiration);
|
||||
Ok(expirations
|
||||
.into_iter()
|
||||
.take(count)
|
||||
.map(|(key, _)| key)
|
||||
.collect())
|
||||
}
|
||||
|
||||
pub fn delete_cache(&self, key: &str) -> Result<(), Error> {
|
||||
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
|
||||
Self::remove(&mut state, key);
|
||||
|
|
@ -150,7 +172,7 @@ impl<V: Clone> InMemoryCache<V> {
|
|||
Ok(())
|
||||
}
|
||||
|
||||
fn evict(state: &mut CacheState<V>, capacity: usize, now: Duration) {
|
||||
fn evict(state: &mut CacheState<V>, capacity: usize, now: Duration, key: &str) {
|
||||
while let Some(Reverse((expiration, key))) = state.expiration_heap.peek().cloned() {
|
||||
if state.expirations.get(&key).copied() != Some(expiration) {
|
||||
state.expiration_heap.pop();
|
||||
|
|
@ -161,6 +183,9 @@ impl<V: Clone> InMemoryCache<V> {
|
|||
break;
|
||||
}
|
||||
}
|
||||
if state.values.contains_key(key) {
|
||||
return;
|
||||
}
|
||||
while state.values.len() >= capacity {
|
||||
let Some(Reverse((expiration, key))) = state.expiration_heap.pop() else {
|
||||
break;
|
||||
|
|
@ -171,84 +196,205 @@ impl<V: Clone> InMemoryCache<V> {
|
|||
}
|
||||
}
|
||||
|
||||
fn set_expiration(state: &mut CacheState<V>, key: &str, expiration: Duration) {
|
||||
if state.expirations.get(key).copied() != Some(expiration) {
|
||||
state.expirations.insert(key.into(), expiration);
|
||||
state
|
||||
.expiration_heap
|
||||
.push(Reverse((expiration, key.into())));
|
||||
}
|
||||
}
|
||||
|
||||
fn remove(state: &mut CacheState<V>, key: &str) {
|
||||
state.values.remove(key);
|
||||
state.expirations.remove(key);
|
||||
}
|
||||
}
|
||||
|
||||
impl InMemoryCache<CacheEntry> {
|
||||
pub fn response_cache(capacity: usize, ttl: Duration, max_entry_bytes: usize) -> Self {
|
||||
Self::response_cache_with_clock(capacity, ttl, max_entry_bytes, || {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
})
|
||||
}
|
||||
|
||||
pub fn response_cache_with_clock(
|
||||
capacity: usize,
|
||||
ttl: Duration,
|
||||
max_entry_bytes: usize,
|
||||
now: impl Fn() -> Duration + Send + Sync + 'static,
|
||||
) -> Self {
|
||||
let mut cache = Self::with_clock_and_size_measurement(
|
||||
Some(capacity),
|
||||
Some(ttl),
|
||||
Some(max_entry_bytes),
|
||||
Some(Arc::new(|entry: &CacheEntry| {
|
||||
serde_json::to_vec(entry)
|
||||
.map(|bytes| bytes.len())
|
||||
.map_err(|_| Error::InvalidEntry)
|
||||
})),
|
||||
now,
|
||||
impl<V> ClaimCache for InMemoryCache<V>
|
||||
where
|
||||
V: Clone + PartialEq + Send + Sync + 'static,
|
||||
{
|
||||
fn claim_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
candidate: V,
|
||||
eligible: &[V],
|
||||
context: ExactCacheContext,
|
||||
) -> Result<V, Error> {
|
||||
if self.max_size_in_memory == 0 {
|
||||
return Ok(candidate);
|
||||
}
|
||||
let now = (self.now)();
|
||||
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
|
||||
Self::evict(&mut state, self.max_size_in_memory, now, key);
|
||||
let existing = state
|
||||
.values
|
||||
.get(key)
|
||||
.filter(|existing| eligible.is_empty() || eligible.contains(existing))
|
||||
.cloned();
|
||||
if let Some(existing) = &existing
|
||||
&& eligible.is_empty()
|
||||
&& *existing != candidate
|
||||
{
|
||||
return Ok(existing.clone());
|
||||
}
|
||||
let winner = existing.unwrap_or(candidate);
|
||||
Self::set_expiration(
|
||||
&mut state,
|
||||
key,
|
||||
now + self.get_ttl(&context).unwrap_or(self.default_ttl),
|
||||
);
|
||||
cache.validate_value = Some(Arc::new(|entry: &CacheEntry| {
|
||||
entry
|
||||
.timestamp
|
||||
.is_finite()
|
||||
.then_some(())
|
||||
.ok_or(Error::InvalidEntry)
|
||||
}));
|
||||
cache
|
||||
state.values.insert(key.into(), winner.clone());
|
||||
Ok(winner)
|
||||
}
|
||||
}
|
||||
|
||||
impl BaseCache for InMemoryCache<CacheEntry> {
|
||||
type Value = CacheEntry;
|
||||
impl CounterCache for InMemoryCache<f64> {
|
||||
fn increment_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
amount: f64,
|
||||
context: ExactCacheContext,
|
||||
) -> Result<f64, Error> {
|
||||
if self.max_size_in_memory == 0 {
|
||||
return Ok(amount);
|
||||
}
|
||||
let now = (self.now)();
|
||||
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
|
||||
Self::evict(&mut state, self.max_size_in_memory, now, key);
|
||||
let value = state.values.get(key).copied().unwrap_or_default() + amount;
|
||||
if !state.expirations.contains_key(key) {
|
||||
Self::set_expiration(
|
||||
&mut state,
|
||||
key,
|
||||
now + self.get_ttl(&context).unwrap_or(self.default_ttl),
|
||||
);
|
||||
}
|
||||
state.values.insert(key.into(), value);
|
||||
Ok(value)
|
||||
}
|
||||
}
|
||||
|
||||
fn default_ttl(&self) -> Duration {
|
||||
self.default_ttl
|
||||
impl InMemoryCache<f64> {
|
||||
pub async fn async_increment_pipeline(
|
||||
&self,
|
||||
operations: Vec<IncrementOperation>,
|
||||
) -> Result<Vec<f64>, Error> {
|
||||
operations
|
||||
.into_iter()
|
||||
.map(|operation| {
|
||||
self.increment_cache(
|
||||
&operation.key,
|
||||
operation.amount,
|
||||
ExactCacheContext { ttl: operation.ttl },
|
||||
)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
impl<V: Clone + Send + Sync + 'static> BaseCache for InMemoryCache<V> {
|
||||
type Value = V;
|
||||
type Context = ExactCacheContext;
|
||||
|
||||
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
|
||||
context.ttl.or(Some(self.default_ttl))
|
||||
}
|
||||
|
||||
fn set_cache(&self, key: &str, value: Self::Value, kwargs: CacheKwargs) -> Result<(), Error> {
|
||||
let ttl = self.get_ttl(&kwargs);
|
||||
fn set_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
value: Self::Value,
|
||||
context: &ExactCacheContext,
|
||||
) -> Result<(), Error> {
|
||||
let ttl = self.get_ttl(context).unwrap_or(self.default_ttl);
|
||||
self.set_cache(key, value, Some(ttl)).map(|_| ())
|
||||
}
|
||||
|
||||
fn get_cache(&self, key: &str, _: &CacheKwargs) -> Result<Option<Self::Value>, Error> {
|
||||
fn get_cache(&self, key: &str, _: &ExactCacheContext) -> Result<Option<Self::Value>, Error> {
|
||||
self.get_cache(key)
|
||||
}
|
||||
|
||||
fn delete_cache(&self, key: &str) -> Result<(), Error> {
|
||||
self.delete_cache(key)
|
||||
async fn disconnect(&self) -> Result<(), Error> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn flush_cache(&self) -> Result<(), Error> {
|
||||
self.flush_cache()
|
||||
}
|
||||
|
||||
fn disconnect(&self) -> CacheFuture<'_, ()> {
|
||||
Box::pin(async { Ok(()) })
|
||||
}
|
||||
|
||||
fn test_connection(&self) -> CacheFuture<'_, CacheConnectionResult> {
|
||||
Box::pin(async {
|
||||
Ok(CacheConnectionResult {
|
||||
status: CacheConnectionStatus::Success,
|
||||
message: "In-memory cache connection test successful".into(),
|
||||
error: None,
|
||||
})
|
||||
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
|
||||
Ok(CacheConnectionResult {
|
||||
status: CacheConnectionStatus::Success,
|
||||
message: "In-memory cache connection test successful".into(),
|
||||
error: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl<V: Clone + Send + Sync + 'static> BatchCache for InMemoryCache<V> {}
|
||||
|
||||
impl<V: Clone + Send + Sync + 'static> DeleteCache for InMemoryCache<V> {
|
||||
fn delete_cache(&self, key: &str) -> Result<(), Error> {
|
||||
InMemoryCache::delete_cache(self, key)
|
||||
}
|
||||
}
|
||||
|
||||
impl<V: Clone + Send + Sync + 'static> FlushCache for InMemoryCache<V> {
|
||||
fn flush_cache(&self) -> Result<(), Error> {
|
||||
InMemoryCache::flush_cache(self)
|
||||
}
|
||||
}
|
||||
|
||||
impl<V: Clone + Send + Sync + 'static> TtlCache for InMemoryCache<V> {
|
||||
async fn async_get_ttl(&self, key: &str) -> Result<Option<Duration>, Error> {
|
||||
InMemoryCache::async_get_ttl(self, key).await
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> SetCache for InMemoryCache<HashSet<T>>
|
||||
where
|
||||
T: Clone + Eq + Hash + Send + Sync + 'static,
|
||||
{
|
||||
type SetValue = T;
|
||||
type SetResult = Vec<T>;
|
||||
|
||||
async fn async_set_cache_sadd(
|
||||
&self,
|
||||
key: &str,
|
||||
values: Vec<Self::SetValue>,
|
||||
ttl: Option<Duration>,
|
||||
) -> Result<Self::SetResult, Error> {
|
||||
if self.max_size_in_memory == 0 {
|
||||
return Ok(values);
|
||||
}
|
||||
let now = (self.now)();
|
||||
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
|
||||
Self::evict(&mut state, self.max_size_in_memory, now, key);
|
||||
let mut stored = state.values.get(key).cloned().unwrap_or_default();
|
||||
stored.extend(values.iter().cloned());
|
||||
if let (Some(limit), Some(measure)) = (self.max_entry_bytes, &self.measure_value)
|
||||
&& measure(&stored)? > limit
|
||||
{
|
||||
return Ok(values);
|
||||
}
|
||||
if !state.expirations.contains_key(key) {
|
||||
Self::set_expiration(&mut state, key, now + ttl.unwrap_or(self.default_ttl));
|
||||
}
|
||||
state.values.insert(key.into(), stored);
|
||||
Ok(values)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn repeated_increments_keep_one_heap_entry_per_expiration() {
|
||||
let cache = InMemoryCache::<f64>::new(Some(4), None);
|
||||
for _ in 0..100 {
|
||||
cache
|
||||
.increment_cache("counter", 1.0, ExactCacheContext::default())
|
||||
.unwrap();
|
||||
}
|
||||
assert_eq!(cache.state.lock().unwrap().expiration_heap.len(), 1);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,8 +1,16 @@
|
|||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::time::Duration;
|
||||
use std::{
|
||||
collections::HashSet,
|
||||
sync::{
|
||||
Arc,
|
||||
atomic::{AtomicU64, Ordering},
|
||||
},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use litellm_cache::{BaseCache, CacheConnectionStatus, CacheEntry, Error};
|
||||
use litellm_cache::{
|
||||
BaseCache, CacheBackend, CacheConnectionStatus, ClaimCache, CounterCache, DeleteCache, Error,
|
||||
ExactCacheContext, IncrementOperation, SetCache, get_cache, set_cache,
|
||||
};
|
||||
use litellm_cache_memory::{CacheWrite, InMemoryCache};
|
||||
use rstest::{fixture, rstest};
|
||||
|
||||
|
|
@ -84,66 +92,49 @@ fn capacity_evicts_earliest_and_ignores_stale_heap_entries(clock: Arc<AtomicU64>
|
|||
}
|
||||
|
||||
#[test]
|
||||
fn disabled_size_limited_and_synchronized_response_writes_are_observable() {
|
||||
let disabled = InMemoryCache::<CacheEntry>::response_cache(0, Duration::from_secs(60), 80);
|
||||
fn disabled_size_limited_and_validated_writes_are_observable() {
|
||||
let cache = |capacity| {
|
||||
InMemoryCache::with_clock_and_size_measurement(
|
||||
Some(capacity),
|
||||
Some(Duration::from_secs(60)),
|
||||
Some(4),
|
||||
Some(Arc::new(|value: &String| {
|
||||
if value.is_empty() {
|
||||
return Err(Error::InvalidEntry);
|
||||
}
|
||||
Ok(value.len())
|
||||
})),
|
||||
|| Duration::from_secs(100),
|
||||
)
|
||||
};
|
||||
let disabled = cache(0);
|
||||
assert_eq!(
|
||||
disabled
|
||||
.set_cache(
|
||||
"a",
|
||||
CacheEntry {
|
||||
timestamp: 1.0,
|
||||
response: serde_json::json!("x")
|
||||
},
|
||||
None
|
||||
)
|
||||
.unwrap(),
|
||||
disabled.set_cache("a", "x".into(), None).unwrap(),
|
||||
CacheWrite::Disabled
|
||||
);
|
||||
let cache = InMemoryCache::<CacheEntry>::response_cache(2, Duration::from_secs(60), 80);
|
||||
let cache = cache(2);
|
||||
assert_eq!(
|
||||
cache
|
||||
.set_cache(
|
||||
"large",
|
||||
CacheEntry {
|
||||
timestamp: 1.0,
|
||||
response: serde_json::json!("x".repeat(100))
|
||||
},
|
||||
None
|
||||
)
|
||||
.unwrap(),
|
||||
cache.set_cache("large", "oversized".into(), None).unwrap(),
|
||||
CacheWrite::TooLarge
|
||||
);
|
||||
cache
|
||||
.set_cache(
|
||||
"small",
|
||||
CacheEntry {
|
||||
timestamp: 1.0,
|
||||
response: serde_json::json!("ok"),
|
||||
},
|
||||
None,
|
||||
)
|
||||
.unwrap();
|
||||
assert!(cache.get_cache("small").unwrap().is_some());
|
||||
assert_eq!(cache.get_cache("large").unwrap(), None);
|
||||
assert_eq!(
|
||||
cache
|
||||
.set_cache(
|
||||
"invalid",
|
||||
CacheEntry {
|
||||
timestamp: f64::NAN,
|
||||
response: serde_json::json!("bad"),
|
||||
},
|
||||
None,
|
||||
)
|
||||
.unwrap_err(),
|
||||
Error::InvalidEntry
|
||||
cache.set_cache("small", "ok".into(), None).unwrap(),
|
||||
CacheWrite::Stored
|
||||
);
|
||||
assert_eq!(cache.get_cache("small").unwrap(), Some("ok".into()));
|
||||
assert_eq!(
|
||||
cache.set_cache("invalid", String::new(), None),
|
||||
Err(Error::InvalidEntry)
|
||||
);
|
||||
assert_eq!(cache.get_cache("invalid").unwrap(), None);
|
||||
cache.delete_cache("small").unwrap();
|
||||
cache.flush_cache().unwrap();
|
||||
assert_eq!(cache.get_cache("small").unwrap(), None);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn connection_test_matches_python_result_contract() {
|
||||
let cache = InMemoryCache::<CacheEntry>::default();
|
||||
let cache = InMemoryCache::<String>::default();
|
||||
let result = BaseCache::test_connection(&cache).await.unwrap();
|
||||
assert_eq!(result.status, CacheConnectionStatus::Success);
|
||||
assert_eq!(result.message, "In-memory cache connection test successful");
|
||||
|
|
@ -156,3 +147,222 @@ async fn connection_test_matches_python_result_contract() {
|
|||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn generic_consumers_share_typed_values_and_honor_expiration() {
|
||||
let clock = clock();
|
||||
let cache: CacheBackend<InMemoryCache<String>> = Arc::new(cache(clock.clone(), 4));
|
||||
let reader = Arc::clone(&cache);
|
||||
let context = ExactCacheContext {
|
||||
ttl: Some(Duration::from_secs(5)),
|
||||
};
|
||||
set_cache(cache.as_ref(), "sync", "first".into(), &context).unwrap();
|
||||
assert_eq!(
|
||||
get_cache(reader.as_ref(), "sync", &context).unwrap(),
|
||||
Some("first".into())
|
||||
);
|
||||
cache
|
||||
.batch_cache_write("async", "second".into(), context.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
cache
|
||||
.async_set_cache_pipeline(vec![("batch".into(), "third".into())], context.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
drop(cache);
|
||||
for (key, value) in [("sync", "first"), ("async", "second"), ("batch", "third")] {
|
||||
assert_eq!(
|
||||
reader.async_get_cache(key, &context).await.unwrap(),
|
||||
Some(value.into())
|
||||
);
|
||||
}
|
||||
reader.async_delete_cache("async").await.unwrap();
|
||||
assert_eq!(
|
||||
reader.async_get_cache("async", &context).await.unwrap(),
|
||||
None
|
||||
);
|
||||
clock.store(106, Ordering::SeqCst);
|
||||
assert_eq!(get_cache(reader.as_ref(), "sync", &context).unwrap(), None);
|
||||
assert_eq!(
|
||||
reader.async_get_cache("batch", &context).await.unwrap(),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn claims_are_atomic_and_refresh_eligible_winners() {
|
||||
let clock = clock();
|
||||
let cache = InMemoryCache::with_clock(Some(4), Some(Duration::from_secs(60)), {
|
||||
let clock = clock.clone();
|
||||
move || Duration::from_secs(clock.load(Ordering::SeqCst))
|
||||
});
|
||||
let context = ExactCacheContext {
|
||||
ttl: Some(Duration::from_secs(10)),
|
||||
};
|
||||
assert_eq!(
|
||||
cache
|
||||
.claim_cache("affinity", "first".to_string(), &[], context.clone())
|
||||
.unwrap(),
|
||||
"first"
|
||||
);
|
||||
clock.store(103, Ordering::SeqCst);
|
||||
assert_eq!(
|
||||
cache
|
||||
.claim_cache("affinity", "second".to_string(), &[], context.clone())
|
||||
.unwrap(),
|
||||
"first"
|
||||
);
|
||||
assert_eq!(
|
||||
cache.expires_at("affinity").unwrap(),
|
||||
Some(Duration::from_secs(110))
|
||||
);
|
||||
clock.store(105, Ordering::SeqCst);
|
||||
assert_eq!(
|
||||
cache
|
||||
.claim_cache(
|
||||
"affinity",
|
||||
"second".to_string(),
|
||||
&["first".to_string(), "second".to_string()],
|
||||
context,
|
||||
)
|
||||
.unwrap(),
|
||||
"first"
|
||||
);
|
||||
assert_eq!(
|
||||
cache.expires_at("affinity").unwrap(),
|
||||
Some(Duration::from_secs(115))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn counters_increment_under_one_lock() {
|
||||
let cache = InMemoryCache::<f64>::default();
|
||||
assert_eq!(
|
||||
CounterCache::increment_cache(&cache, "counter", 1.5, ExactCacheContext::default())
|
||||
.unwrap(),
|
||||
1.5
|
||||
);
|
||||
assert_eq!(
|
||||
CounterCache::increment_cache(&cache, "counter", 2.0, ExactCacheContext::default())
|
||||
.unwrap(),
|
||||
3.5
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn rewriting_an_existing_key_at_capacity_keeps_other_entries(clock: Arc<AtomicU64>) {
|
||||
let cache = cache(clock, 2);
|
||||
cache
|
||||
.set_cache("hot", "1".into(), Some(Duration::from_secs(10)))
|
||||
.unwrap();
|
||||
cache
|
||||
.set_cache("cold", "2".into(), Some(Duration::from_secs(20)))
|
||||
.unwrap();
|
||||
|
||||
cache.set_cache("cold", "3".into(), None).unwrap();
|
||||
assert_eq!(cache.get_cache("hot").unwrap(), Some("1".into()));
|
||||
assert_eq!(cache.get_cache("cold").unwrap(), Some("3".into()));
|
||||
|
||||
cache
|
||||
.claim_cache("cold", "4".into(), &[], ExactCacheContext::default())
|
||||
.unwrap();
|
||||
assert_eq!(cache.get_cache("hot").unwrap(), Some("1".into()));
|
||||
|
||||
cache.set_cache("new", "5".into(), None).unwrap();
|
||||
assert_eq!(cache.get_cache("hot").unwrap(), None);
|
||||
assert_eq!(cache.get_cache("cold").unwrap(), Some("3".into()));
|
||||
assert_eq!(cache.get_cache("new").unwrap(), Some("5".into()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn incrementing_an_existing_counter_at_capacity_keeps_every_counter() {
|
||||
let cache = InMemoryCache::<f64>::new(Some(2), None);
|
||||
for key in ["a", "b", "a", "b"] {
|
||||
cache
|
||||
.increment_cache(key, 1.0, ExactCacheContext::default())
|
||||
.unwrap();
|
||||
}
|
||||
assert_eq!(cache.get_cache("a").unwrap(), Some(2.0));
|
||||
assert_eq!(cache.get_cache("b").unwrap(), Some(2.0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn disabled_cache_does_not_retain_claims_or_counters() {
|
||||
let claims = InMemoryCache::<String>::new(Some(0), None);
|
||||
assert_eq!(
|
||||
claims
|
||||
.claim_cache("key", "first".into(), &[], ExactCacheContext::default())
|
||||
.unwrap(),
|
||||
"first"
|
||||
);
|
||||
assert_eq!(claims.get_cache("key").unwrap(), None);
|
||||
|
||||
let counters = InMemoryCache::<f64>::new(Some(0), None);
|
||||
assert_eq!(
|
||||
counters
|
||||
.increment_cache("key", 2.0, ExactCacheContext::default())
|
||||
.unwrap(),
|
||||
2.0
|
||||
);
|
||||
assert_eq!(counters.get_cache("key").unwrap(), None);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn ttl_and_oldest_key_operations_use_the_stored_expirations() {
|
||||
let clock = Arc::new(AtomicU64::new(100));
|
||||
let cache = cache(clock, 3);
|
||||
cache
|
||||
.set_cache("later", "2".into(), Some(Duration::from_secs(20)))
|
||||
.unwrap();
|
||||
cache
|
||||
.set_cache("first", "1".into(), Some(Duration::from_secs(10)))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
cache.async_get_ttl("first").await.unwrap(),
|
||||
Some(Duration::from_secs(110))
|
||||
);
|
||||
assert_eq!(cache.async_get_oldest_n_keys(1).await.unwrap(), ["first"]);
|
||||
assert_eq!(cache.async_get_ttl("missing").await.unwrap(), None);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn increment_pipeline_preserves_operation_order() {
|
||||
let cache = InMemoryCache::<f64>::new(Some(3), None);
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_increment_pipeline(vec![
|
||||
IncrementOperation {
|
||||
key: "a".into(),
|
||||
amount: 1.0,
|
||||
ttl: Some(Duration::from_secs(10)),
|
||||
},
|
||||
IncrementOperation {
|
||||
key: "a".into(),
|
||||
amount: 2.0,
|
||||
ttl: Some(Duration::from_secs(20)),
|
||||
},
|
||||
])
|
||||
.await
|
||||
.unwrap(),
|
||||
[1.0, 3.0]
|
||||
);
|
||||
assert_eq!(cache.get_cache("a").unwrap(), Some(3.0));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn set_capability_preserves_python_result_and_deduplicates_storage() {
|
||||
let cache = InMemoryCache::<HashSet<String>>::new(None, None);
|
||||
let inserted = vec!["a".into(), "a".into(), "b".into()];
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_set_cache_sadd("members", inserted.clone(), None)
|
||||
.await
|
||||
.unwrap(),
|
||||
inserted
|
||||
);
|
||||
assert_eq!(
|
||||
cache.get_cache("members").unwrap(),
|
||||
Some(HashSet::from(["a".into(), "b".into()]))
|
||||
);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -7,9 +7,10 @@ repository.workspace = true
|
|||
|
||||
[dependencies]
|
||||
litellm-cache.workspace = true
|
||||
redis = "1.7.0"
|
||||
serde_json.workspace = true
|
||||
redis = { version = "1.7.0", features = ["cluster", "tls-rustls"] }
|
||||
r2d2 = "0.8.10"
|
||||
tokio.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
redis-test = "1.0.4"
|
||||
serde_json.workspace = true
|
||||
|
|
|
|||
|
|
@ -1,58 +1,220 @@
|
|||
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 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;
|
||||
|
||||
#[allow(private_interfaces)]
|
||||
pub 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)),
|
||||
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
|
||||
pub fn execute<T>(
|
||||
&self,
|
||||
operation: impl FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error>,
|
||||
) -> Result<T, Error> {
|
||||
match self {
|
||||
Self::Pool(pool) => {
|
||||
let mut pooled = pool.get().map_err(|_| Error::Unavailable)?;
|
||||
let result = operation(&mut ConnectionRef::Node(&mut pooled.connection));
|
||||
pooled.failed = matches!(result, Err(Error::Unavailable));
|
||||
result
|
||||
}
|
||||
Self::Cluster(pool) => {
|
||||
let mut pooled = pool.get().map_err(|_| Error::Unavailable)?;
|
||||
let result = operation(&mut ConnectionRef::Cluster(&mut pooled.connection));
|
||||
pooled.failed = matches!(result, Err(Error::Unavailable));
|
||||
result
|
||||
}
|
||||
Self::Fixed(connection) => {
|
||||
let mut connection = connection.lock().map_err(|_| Error::Unavailable)?;
|
||||
operation(&mut ConnectionRef::Node(&mut *connection))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn connection(&self) -> Result<MutexGuard<'_, C>, Error> {
|
||||
self.connection.lock().map_err(|_| Error::Unavailable)
|
||||
pub async fn run_blocking<T, F>(connections: Arc<Self>, operation: F) -> Result<T, Error>
|
||||
where
|
||||
T: Send + 'static,
|
||||
F: FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error> + Send + 'static,
|
||||
{
|
||||
tokio::task::spawn_blocking(move || connections.execute(operation))
|
||||
.await
|
||||
.map_err(|_| Error::Unavailable)?
|
||||
}
|
||||
|
||||
fn namespaced_key(key: &str) -> String {
|
||||
format!("{KEY_PREFIX}{key}")
|
||||
pub fn fixed(connection: C) -> Self {
|
||||
Self::Fixed(Mutex::new(connection))
|
||||
}
|
||||
|
||||
fn namespaced_pattern() -> &'static str {
|
||||
const PATTERN: &str = "litellm-cache:*";
|
||||
PATTERN
|
||||
pub fn open(url: &str, topology: &RedisTopology) -> Result<Self, Error> {
|
||||
match topology {
|
||||
RedisTopology::Standalone => Ok(Self::Pool(pool(ConnectionManager::open(url)?)?)),
|
||||
RedisTopology::Cluster { startup_nodes } => Ok(Self::Cluster(pool(
|
||||
ClusterConnectionManager::open(url, startup_nodes)?,
|
||||
)?)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
|
||||
fn encode(value: &CacheEntry) -> Result<Vec<u8>, Error> {
|
||||
serde_json::to_vec(value).map_err(|_| Error::InvalidEntry)
|
||||
pub fn connect(
|
||||
url: &str,
|
||||
topology: &RedisTopology,
|
||||
default_ttl: Option<Duration>,
|
||||
codec: S,
|
||||
) -> Result<Self, Error> {
|
||||
let connections = Connections::open(url, topology)?;
|
||||
Ok(Self {
|
||||
connections: Arc::new(connections),
|
||||
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
|
||||
codec,
|
||||
namespace: None,
|
||||
topology: topology.clone(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
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(connection)),
|
||||
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
|
||||
codec,
|
||||
namespace: None,
|
||||
topology: RedisTopology::Standalone,
|
||||
}
|
||||
}
|
||||
|
||||
fn decode(value: Vec<u8>) -> Result<CacheEntry, Error> {
|
||||
serde_json::from_slice(&value).map_err(|_| Error::InvalidEntry)
|
||||
pub fn with_namespace(self, namespace: Option<String>) -> Self {
|
||||
Self {
|
||||
namespace: namespace.filter(|value| !value.is_empty()),
|
||||
..self
|
||||
}
|
||||
}
|
||||
|
||||
pub fn namespace(&self) -> Option<&str> {
|
||||
self.namespace.as_deref()
|
||||
}
|
||||
|
||||
pub fn topology(&self) -> &RedisTopology {
|
||||
&self.topology
|
||||
}
|
||||
|
||||
fn namespaced_key(&self, key: &str) -> String {
|
||||
namespaced_key(self.namespace.as_deref(), key)
|
||||
}
|
||||
|
||||
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 {
|
||||
|
|
@ -60,197 +222,409 @@ where
|
|||
.saturating_add(u64::from(ttl.subsec_nanos() > 0))
|
||||
.max(1)
|
||||
}
|
||||
}
|
||||
|
||||
fn run_blocking<T, F>(connection: Arc<Mutex<C>>, operation: F) -> CacheFuture<'static, T>
|
||||
where
|
||||
T: Send + 'static,
|
||||
F: FnOnce(&mut C) -> 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)
|
||||
})
|
||||
.await
|
||||
.map_err(|_| Error::Unavailable)?
|
||||
})
|
||||
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<C> BaseCache for RedisCache<C>
|
||||
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));
|
||||
Connections::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 = Connections::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(());
|
||||
}
|
||||
Connections::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 Connections::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 = Connections::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);
|
||||
Connections::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()?;
|
||||
Connections::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));
|
||||
Connections::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();
|
||||
Connections::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 +632,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 +647,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 +670,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 +689,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
392
litellm-rust/crates/cache-redis/src/cache/connection.rs
vendored
Normal 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 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
632
litellm-rust/crates/cache-redis/src/cache/operations.rs
vendored
Normal 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<_>>();
|
||||
Connections::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 = Connections::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> {
|
||||
Connections::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 = Connections::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));
|
||||
Connections::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));
|
||||
Connections::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);
|
||||
Connections::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());
|
||||
}
|
||||
Connections::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 = Connections::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 = Connections::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<_>>();
|
||||
Connections::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());
|
||||
}
|
||||
Connections::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);
|
||||
Connections::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));
|
||||
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
|
||||
redis::cmd("EVAL")
|
||||
.arg(SET_MAX_SCRIPT)
|
||||
.arg(1)
|
||||
.arg(key)
|
||||
.arg(value)
|
||||
.arg(ttl)
|
||||
.query(connection)
|
||||
.map_err(|_| Error::Unavailable)
|
||||
})
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
fn redis_bytes(value: redis::Value) -> Result<Vec<u8>, Error> {
|
||||
match value {
|
||||
redis::Value::BulkString(bytes) => Ok(bytes),
|
||||
redis::Value::SimpleString(text) => Ok(text.into_bytes()),
|
||||
_ => Err(Error::InvalidEntry),
|
||||
}
|
||||
}
|
||||
|
||||
fn lpop_result(value: redis::Value, multiple: bool) -> Result<RedisLpopResult, Error> {
|
||||
match value {
|
||||
redis::Value::Nil => Ok(RedisLpopResult::Missing),
|
||||
redis::Value::Array(values) if multiple => values
|
||||
.into_iter()
|
||||
.map(redis_bytes)
|
||||
.collect::<Result<Vec<_>, _>>()
|
||||
.map(RedisLpopResult::Values),
|
||||
value if !multiple => redis_bytes(value).map(RedisLpopResult::Value),
|
||||
_ => Err(Error::InvalidEntry),
|
||||
}
|
||||
}
|
||||
|
||||
fn count(value: redis::Value) -> Result<Option<i64>, Error> {
|
||||
match value {
|
||||
redis::Value::Nil => Ok(None),
|
||||
redis::Value::Int(value) => Ok(Some(value)),
|
||||
redis::Value::BulkString(value) => std::str::from_utf8(&value)
|
||||
.ok()
|
||||
.and_then(|value| value.parse().ok())
|
||||
.map(Some)
|
||||
.ok_or(Error::InvalidEntry),
|
||||
redis::Value::SimpleString(value) => {
|
||||
value.parse().map(Some).map_err(|_| Error::InvalidEntry)
|
||||
}
|
||||
_ => Err(Error::InvalidEntry),
|
||||
}
|
||||
}
|
||||
|
||||
fn increment_with_floor(
|
||||
connection: &mut ConnectionRef<'_>,
|
||||
key: String,
|
||||
amount: i64,
|
||||
ttl: u64,
|
||||
) -> Result<i64, Error> {
|
||||
redis::cmd("EVAL")
|
||||
.arg(INCREMENT_WITH_FLOOR_SCRIPT)
|
||||
.arg(1)
|
||||
.arg(key)
|
||||
.arg(amount)
|
||||
.arg(ttl)
|
||||
.query(connection)
|
||||
.map_err(|_| Error::Unavailable)
|
||||
}
|
||||
|
||||
impl<S, C> TtlCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
async fn async_get_ttl(&self, key: &str) -> Result<Option<Duration>, Error> {
|
||||
RedisCache::async_get_ttl(self, key)
|
||||
.await
|
||||
.map(|ttl| ttl.map(|seconds| Duration::from_secs(seconds as u64)))
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, C> ScanCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
async fn async_scan_iter(&self, pattern: &str, count: usize) -> Result<Vec<String>, Error> {
|
||||
RedisCache::async_scan_iter(self, pattern, count).await
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, C> ClientInfoCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
type ClientList = String;
|
||||
type Info = String;
|
||||
|
||||
fn client_list(&self) -> Result<Self::ClientList, Error> {
|
||||
RedisCache::client_list(self)
|
||||
}
|
||||
|
||||
fn info(&self) -> Result<Self::Info, Error> {
|
||||
RedisCache::info(self)
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, C> SetCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
type SetValue = RedisArg;
|
||||
type SetResult = usize;
|
||||
|
||||
async fn async_set_cache_sadd(
|
||||
&self,
|
||||
key: &str,
|
||||
values: Vec<Self::SetValue>,
|
||||
ttl: Option<Duration>,
|
||||
) -> Result<Self::SetResult, Error> {
|
||||
RedisCache::async_set_cache_sadd(self, key, values, ttl).await
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, C> QueueCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
type QueueValue = RedisArg;
|
||||
type PopResult = RedisLpopResult;
|
||||
|
||||
async fn async_rpush(&self, key: &str, values: Vec<Self::QueueValue>) -> Result<usize, Error> {
|
||||
RedisCache::async_rpush(self, key, values).await
|
||||
}
|
||||
|
||||
async fn async_lpop(&self, key: &str, count: Option<usize>) -> Result<Self::PopResult, Error> {
|
||||
RedisCache::async_lpop(self, key, count).await
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, C> ScriptCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
type Script = RedisScript<C>;
|
||||
|
||||
fn async_register_script(&self, source: String) -> Self::Script {
|
||||
RedisScript {
|
||||
connections: Arc::clone(&self.connections),
|
||||
namespace: self.namespace.clone(),
|
||||
source,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -1,3 +1,11 @@
|
|||
mod cache;
|
||||
mod topology;
|
||||
|
||||
pub use cache::RedisCache;
|
||||
pub mod connection {
|
||||
pub use crate::cache::{ConnectionRef, Connections};
|
||||
}
|
||||
|
||||
pub use cache::{
|
||||
RedisArg, RedisCache, RedisLpopOperation, RedisLpopResult, RedisRpushOperation, RedisScript,
|
||||
};
|
||||
pub use topology::{RedisNode, RedisTopology};
|
||||
|
|
|
|||
14
litellm-rust/crates/cache-redis/src/topology.rs
Normal file
14
litellm-rust/crates/cache-redis/src/topology.rs
Normal file
|
|
@ -0,0 +1,14 @@
|
|||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct RedisNode {
|
||||
pub host: String,
|
||||
pub port: u16,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default, PartialEq, Eq)]
|
||||
pub enum RedisTopology {
|
||||
#[default]
|
||||
Standalone,
|
||||
Cluster {
|
||||
startup_nodes: Vec<RedisNode>,
|
||||
},
|
||||
}
|
||||
|
|
@ -1,6 +1,703 @@
|
|||
use litellm_cache_redis::RedisCache;
|
||||
use std::time::Duration;
|
||||
|
||||
use litellm_cache::{
|
||||
BaseCache, BatchCache, BatchEntry, CacheCodec, CacheConnectionStatus, CacheScript, ClaimCache,
|
||||
CounterCache, DeleteCache, Error, ExactCacheContext, FlushCache, IncrementOperation, JsonCodec,
|
||||
ScriptCache, get_cache, set_cache,
|
||||
};
|
||||
use litellm_cache_redis::{
|
||||
RedisArg, RedisCache, RedisLpopOperation, RedisLpopResult, RedisRpushOperation,
|
||||
};
|
||||
use redis_test::{MockCmd, MockRedisConnection};
|
||||
|
||||
struct TaggedByteCodec(u8);
|
||||
|
||||
impl CacheCodec for TaggedByteCodec {
|
||||
type Value = u8;
|
||||
|
||||
fn encode(&self, value: &u8) -> Result<Vec<u8>, Error> {
|
||||
if *value > 127 {
|
||||
return Err(Error::InvalidEntry);
|
||||
}
|
||||
Ok(vec![self.0, *value])
|
||||
}
|
||||
|
||||
fn decode(&self, bytes: &[u8]) -> Result<u8, Error> {
|
||||
match bytes {
|
||||
[tag, value] if *tag == self.0 => Ok(*value),
|
||||
_ => Err(Error::InvalidEntry),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn constructor_rejects_invalid_urls() {
|
||||
assert!(RedisCache::new("not a redis url", None).is_err());
|
||||
assert!(RedisCache::new("not a redis url", None, JsonCodec::<String>::new()).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn generic_helpers_use_the_injected_codec_and_ttl() {
|
||||
let connection = MockRedisConnection::new([
|
||||
MockCmd::new(
|
||||
redis::cmd("SETEX")
|
||||
.arg("counter")
|
||||
.arg(2)
|
||||
.arg([42u8, 7].as_slice()),
|
||||
Ok("OK"),
|
||||
),
|
||||
MockCmd::new(redis::cmd("GET").arg("counter"), Ok(vec![42u8, 7])),
|
||||
])
|
||||
.assert_all_commands_consumed();
|
||||
let cache = RedisCache::with_connection(connection, None, TaggedByteCodec(42));
|
||||
let context = ExactCacheContext {
|
||||
ttl: Some(Duration::from_millis(1500)),
|
||||
};
|
||||
set_cache(&cache, "counter", 7, &context).unwrap();
|
||||
assert_eq!(get_cache(&cache, "counter", &context).unwrap(), Some(7));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn async_operations_preserve_codec_ttl_and_missing_values() {
|
||||
let connection = MockRedisConnection::new([
|
||||
MockCmd::new(
|
||||
redis::cmd("SETEX")
|
||||
.arg("counter")
|
||||
.arg(9)
|
||||
.arg([42u8, 7].as_slice()),
|
||||
Ok("OK"),
|
||||
),
|
||||
MockCmd::new(redis::cmd("GET").arg("counter"), Ok(vec![42u8, 7])),
|
||||
MockCmd::new(
|
||||
redis::cmd("SETEX")
|
||||
.arg("batch")
|
||||
.arg(2)
|
||||
.arg([42u8, 8].as_slice()),
|
||||
Ok("OK"),
|
||||
),
|
||||
MockCmd::new(redis::cmd("DEL").arg("counter"), Ok(1u32)),
|
||||
MockCmd::new(redis::cmd("GET").arg("counter"), Ok(redis::Value::Nil)),
|
||||
])
|
||||
.assert_all_commands_consumed();
|
||||
let cache = RedisCache::with_connection(
|
||||
connection,
|
||||
Some(Duration::from_secs(9)),
|
||||
TaggedByteCodec(42),
|
||||
);
|
||||
let context = ExactCacheContext::default();
|
||||
cache
|
||||
.batch_cache_write("counter", 7, context.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
cache.async_get_cache("counter", &context).await.unwrap(),
|
||||
Some(7)
|
||||
);
|
||||
cache
|
||||
.async_set_cache_pipeline(
|
||||
vec![("batch".into(), 8)],
|
||||
ExactCacheContext {
|
||||
ttl: Some(Duration::from_millis(1500)),
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
cache.async_delete_cache("counter").await.unwrap();
|
||||
assert_eq!(
|
||||
cache.async_get_cache("counter", &context).await.unwrap(),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn codec_errors_propagate_without_writing_partial_batches() {
|
||||
let connection = MockRedisConnection::new([
|
||||
MockCmd::new(redis::cmd("GET").arg("invalid"), Ok(vec![99u8, 7])),
|
||||
MockCmd::new(redis::cmd("GET").arg("invalid"), Ok(vec![99u8, 7])),
|
||||
])
|
||||
.assert_all_commands_consumed();
|
||||
let cache = RedisCache::with_connection(connection, None, TaggedByteCodec(42));
|
||||
let context = ExactCacheContext::default();
|
||||
assert_eq!(
|
||||
cache.set_cache("invalid", 255, &context),
|
||||
Err(Error::InvalidEntry)
|
||||
);
|
||||
assert_eq!(
|
||||
cache.async_set_cache("invalid", 255, context.clone()).await,
|
||||
Err(Error::InvalidEntry)
|
||||
);
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_set_cache_pipeline(
|
||||
vec![("valid".into(), 7), ("invalid".into(), 255)],
|
||||
context.clone(),
|
||||
)
|
||||
.await,
|
||||
Err(Error::InvalidEntry)
|
||||
);
|
||||
assert_eq!(
|
||||
cache.get_cache("invalid", &context),
|
||||
Err(Error::InvalidEntry)
|
||||
);
|
||||
assert_eq!(
|
||||
cache.async_get_cache("invalid", &context).await,
|
||||
Err(Error::InvalidEntry)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn namespaces_are_optional_and_existing_prefixes_are_not_duplicated() {
|
||||
let connection = MockRedisConnection::new([
|
||||
MockCmd::new(redis::cmd("GET").arg("team:key"), Ok(redis::Value::Nil)),
|
||||
MockCmd::new(redis::cmd("GET").arg("team:key"), Ok(redis::Value::Nil)),
|
||||
])
|
||||
.assert_all_commands_consumed();
|
||||
let cache = RedisCache::with_connection(connection, None, JsonCodec::<String>::new())
|
||||
.with_namespace(Some("team".into()));
|
||||
assert_eq!(
|
||||
cache
|
||||
.get_cache("key", &ExactCacheContext::default())
|
||||
.unwrap(),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
cache
|
||||
.get_cache("team:key", &ExactCacheContext::default())
|
||||
.unwrap(),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn flush_requires_a_namespace_and_escapes_glob_metacharacters() {
|
||||
let unscoped = RedisCache::with_connection(
|
||||
MockRedisConnection::new([]).assert_all_commands_consumed(),
|
||||
None,
|
||||
JsonCodec::<String>::new(),
|
||||
);
|
||||
assert_eq!(unscoped.flush_cache(), Err(Error::UnscopedFlush));
|
||||
let connection = MockRedisConnection::new([
|
||||
MockCmd::new(
|
||||
redis::cmd("SCAN")
|
||||
.cursor_arg(0)
|
||||
.arg("MATCH")
|
||||
.arg("team\\*:*")
|
||||
.arg("COUNT")
|
||||
.arg(1000),
|
||||
Ok(redis_test::redis_value!(["0", ["team*:key"]])),
|
||||
),
|
||||
MockCmd::new(redis::cmd("DEL").arg("team*:key"), Ok(1u32)),
|
||||
])
|
||||
.assert_all_commands_consumed();
|
||||
let scoped = RedisCache::with_connection(connection, None, JsonCodec::<String>::new())
|
||||
.with_namespace(Some("team*".into()));
|
||||
scoped.flush_cache().unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn connection_failures_use_the_python_result_contract() {
|
||||
let error = redis::RedisError::from((redis::ErrorKind::Io, "connection refused"));
|
||||
let connection =
|
||||
MockRedisConnection::new([MockCmd::new(redis::cmd("PING"), Err::<String, _>(error))])
|
||||
.assert_all_commands_consumed();
|
||||
let cache = RedisCache::with_connection(connection, None, JsonCodec::<String>::new());
|
||||
|
||||
let result = cache.test_connection().await.unwrap();
|
||||
assert_eq!(result.status, CacheConnectionStatus::Failed);
|
||||
assert!(result.message.starts_with("Redis connection failed:"));
|
||||
assert!(result.error.is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn batch_reads_keep_order_and_treat_invalid_values_as_invalid_entries() {
|
||||
let connection = MockRedisConnection::new([MockCmd::new(
|
||||
redis::cmd("MGET").arg("hit").arg("miss").arg("invalid"),
|
||||
Ok(vec![
|
||||
redis::Value::BulkString(vec![42, 7]),
|
||||
redis::Value::Nil,
|
||||
redis::Value::BulkString(vec![99, 7]),
|
||||
]),
|
||||
)])
|
||||
.assert_all_commands_consumed();
|
||||
let cache = RedisCache::with_connection(connection, None, TaggedByteCodec(42));
|
||||
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_batch_get_cache(
|
||||
vec!["hit".into(), "miss".into(), "invalid".into()],
|
||||
ExactCacheContext::default(),
|
||||
)
|
||||
.await
|
||||
.unwrap(),
|
||||
vec![BatchEntry::Hit(7), BatchEntry::Miss, BatchEntry::Invalid]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn async_flush_deletes_each_scan_page_separately() {
|
||||
let connection = MockRedisConnection::new([
|
||||
MockCmd::new(
|
||||
redis::cmd("SCAN")
|
||||
.cursor_arg(0)
|
||||
.arg("MATCH")
|
||||
.arg("team:*")
|
||||
.arg("COUNT")
|
||||
.arg(1000),
|
||||
Ok(redis_test::redis_value!(["7", ["team:a", "team:b"]])),
|
||||
),
|
||||
MockCmd::new(redis::cmd("DEL").arg("team:a").arg("team:b"), Ok(2u32)),
|
||||
MockCmd::new(
|
||||
redis::cmd("SCAN")
|
||||
.cursor_arg(7)
|
||||
.arg("MATCH")
|
||||
.arg("team:*")
|
||||
.arg("COUNT")
|
||||
.arg(1000),
|
||||
Ok(redis_test::redis_value!(["0", ["team:c"]])),
|
||||
),
|
||||
MockCmd::new(redis::cmd("DEL").arg("team:c"), Ok(1u32)),
|
||||
])
|
||||
.assert_all_commands_consumed();
|
||||
let cache = RedisCache::with_connection(connection, None, JsonCodec::<String>::new())
|
||||
.with_namespace(Some("team".into()));
|
||||
|
||||
cache.async_flush_cache().await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn direct_redis_operations_preserve_namespace_values_and_missing_ttls() {
|
||||
let mut sadd_pipeline = redis::pipe();
|
||||
sadd_pipeline
|
||||
.cmd("SADD")
|
||||
.arg("team:members")
|
||||
.arg("a")
|
||||
.arg("b")
|
||||
.cmd("EXPIRE")
|
||||
.arg("team:members")
|
||||
.arg(600u64)
|
||||
.ignore();
|
||||
let connection = MockRedisConnection::new([
|
||||
MockCmd::new(
|
||||
redis::cmd("MGET").arg("team:count").arg("team:missing"),
|
||||
Ok(redis_test::redis_value!(["7", nil])),
|
||||
),
|
||||
MockCmd::new(
|
||||
redis::cmd("MGET").arg("team:count").arg("team:missing"),
|
||||
Ok(redis_test::redis_value!(["7", nil])),
|
||||
),
|
||||
MockCmd::new(redis::cmd("PING"), Ok("PONG")),
|
||||
MockCmd::new(redis::cmd("PING"), Ok("PONG")),
|
||||
MockCmd::new(redis::cmd("TTL").arg("team:missing"), Ok(-2i64)),
|
||||
MockCmd::new(
|
||||
redis::cmd("SCAN")
|
||||
.cursor_arg(0)
|
||||
.arg("MATCH")
|
||||
.arg("team:job-*")
|
||||
.arg("COUNT")
|
||||
.arg(25),
|
||||
Ok(redis_test::redis_value!(["4", ["team:job-a"]])),
|
||||
),
|
||||
MockCmd::new(
|
||||
redis::cmd("SCAN")
|
||||
.cursor_arg(4)
|
||||
.arg("MATCH")
|
||||
.arg("team:job-*")
|
||||
.arg("COUNT")
|
||||
.arg(25),
|
||||
Ok(redis_test::redis_value!(["0", ["team:job-b"]])),
|
||||
),
|
||||
MockCmd::new(
|
||||
redis::cmd("DEL").arg("team:job-a").arg("team:job-b"),
|
||||
Ok(2u32),
|
||||
),
|
||||
MockCmd::with_values(
|
||||
sadd_pipeline,
|
||||
Ok(vec![redis::Value::Int(2), redis::Value::Int(1)]),
|
||||
),
|
||||
MockCmd::new(
|
||||
redis::cmd("RPUSH").arg("team:queue").arg("a").arg("b"),
|
||||
Ok(2u32),
|
||||
),
|
||||
MockCmd::new(
|
||||
redis::cmd("LPOP").arg("team:queue").arg(2usize),
|
||||
Ok(redis_test::redis_value!(["a", "b"])),
|
||||
),
|
||||
MockCmd::new(
|
||||
redis::cmd("EVAL")
|
||||
.arg("return KEYS[1]")
|
||||
.arg(1usize)
|
||||
.arg("team:key"),
|
||||
Ok("team:key"),
|
||||
),
|
||||
MockCmd::new(
|
||||
redis::cmd("EVAL")
|
||||
.arg("return KEYS[1]")
|
||||
.arg(1usize)
|
||||
.arg("team:key"),
|
||||
Ok("team:key"),
|
||||
),
|
||||
MockCmd::new(redis::cmd("CLIENT").arg("LIST"), Ok("id=1")),
|
||||
MockCmd::new(redis::cmd("INFO"), Ok("redis_version:7")),
|
||||
MockCmd::new(redis::cmd("FLUSHALL"), Ok("OK")),
|
||||
])
|
||||
.assert_all_commands_consumed();
|
||||
let cache = RedisCache::with_connection(connection, None, JsonCodec::<String>::new())
|
||||
.with_namespace(Some("team".into()));
|
||||
|
||||
assert_eq!(
|
||||
cache
|
||||
.batch_get_counts(&["count".into(), "missing".into()])
|
||||
.unwrap(),
|
||||
[Some(7), None]
|
||||
);
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_batch_get_counts(vec!["count".into(), "missing".into()])
|
||||
.await
|
||||
.unwrap(),
|
||||
[Some(7), None]
|
||||
);
|
||||
assert!(cache.sync_ping().unwrap());
|
||||
assert!(cache.ping().await.unwrap());
|
||||
assert_eq!(cache.async_get_ttl("missing").await.unwrap(), None);
|
||||
assert_eq!(
|
||||
cache.async_scan_iter("job-", 25).await.unwrap(),
|
||||
["team:job-a", "team:job-b"]
|
||||
);
|
||||
assert_eq!(
|
||||
cache
|
||||
.delete_cache_keys(vec!["job-a".into(), "job-b".into()])
|
||||
.await
|
||||
.unwrap(),
|
||||
2
|
||||
);
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_set_cache_sadd("members", vec!["a".into(), "b".into()], None)
|
||||
.await
|
||||
.unwrap(),
|
||||
2
|
||||
);
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_rpush("queue", vec!["a".into(), "b".into()])
|
||||
.await
|
||||
.unwrap(),
|
||||
2
|
||||
);
|
||||
assert_eq!(
|
||||
cache.async_lpop("queue", Some(2)).await.unwrap(),
|
||||
RedisLpopResult::Values(vec![b"a".to_vec(), b"b".to_vec()])
|
||||
);
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_eval("return KEYS[1]".into(), vec!["key".into()], Vec::new())
|
||||
.await
|
||||
.unwrap(),
|
||||
redis::Value::BulkString(b"team:key".to_vec())
|
||||
);
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_register_script("return KEYS[1]".into())
|
||||
.invoke(vec!["key".into()], Vec::new())
|
||||
.await
|
||||
.unwrap(),
|
||||
redis::Value::BulkString(b"team:key".to_vec())
|
||||
);
|
||||
assert_eq!(cache.client_list().unwrap(), "id=1");
|
||||
assert_eq!(cache.info().unwrap(), "redis_version:7");
|
||||
cache.flushall().unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn direct_redis_pipelines_preserve_operation_order() {
|
||||
let mut rpush_pipeline = redis::pipe();
|
||||
rpush_pipeline
|
||||
.cmd("RPUSH")
|
||||
.arg("team:a")
|
||||
.arg("one")
|
||||
.cmd("RPUSH")
|
||||
.arg("team:b")
|
||||
.arg("two");
|
||||
let mut lpop_pipeline = redis::pipe();
|
||||
lpop_pipeline
|
||||
.cmd("LPOP")
|
||||
.arg("team:a")
|
||||
.arg(2usize)
|
||||
.cmd("LPOP")
|
||||
.arg("team:b");
|
||||
let connection = MockRedisConnection::new([
|
||||
MockCmd::with_values(
|
||||
rpush_pipeline,
|
||||
Ok(vec![redis::Value::Int(1), redis::Value::Int(2)]),
|
||||
),
|
||||
MockCmd::with_values(
|
||||
lpop_pipeline,
|
||||
Ok(vec![redis_test::redis_value!(["one"]), redis::Value::Nil]),
|
||||
),
|
||||
])
|
||||
.assert_all_commands_consumed();
|
||||
let queue = RedisCache::with_connection(connection, None, JsonCodec::<String>::new())
|
||||
.with_namespace(Some("team".into()));
|
||||
|
||||
assert_eq!(
|
||||
queue
|
||||
.async_rpush_pipeline(vec![
|
||||
RedisRpushOperation {
|
||||
key: "a".into(),
|
||||
values: vec![RedisArg::from("one")],
|
||||
},
|
||||
RedisRpushOperation {
|
||||
key: "b".into(),
|
||||
values: vec![RedisArg::from("two")],
|
||||
},
|
||||
])
|
||||
.await
|
||||
.unwrap(),
|
||||
[1, 2]
|
||||
);
|
||||
assert_eq!(
|
||||
queue
|
||||
.async_lpop_pipeline(vec![
|
||||
RedisLpopOperation {
|
||||
key: "a".into(),
|
||||
count: Some(2),
|
||||
},
|
||||
RedisLpopOperation {
|
||||
key: "b".into(),
|
||||
count: None,
|
||||
},
|
||||
])
|
||||
.await
|
||||
.unwrap(),
|
||||
[
|
||||
RedisLpopResult::Values(vec![b"one".to_vec()]),
|
||||
RedisLpopResult::Missing,
|
||||
]
|
||||
);
|
||||
|
||||
let mut increment_pipeline = redis::pipe();
|
||||
increment_pipeline
|
||||
.cmd("INCRBYFLOAT")
|
||||
.arg("team:counter")
|
||||
.arg(1.5f64)
|
||||
.cmd("EXPIRE")
|
||||
.arg("team:counter")
|
||||
.arg(10u64)
|
||||
.ignore()
|
||||
.cmd("INCRBYFLOAT")
|
||||
.arg("team:counter")
|
||||
.arg(2.0f64);
|
||||
let connection = MockRedisConnection::new([MockCmd::with_values(
|
||||
increment_pipeline,
|
||||
Ok(vec![
|
||||
redis::Value::BulkString(b"1.5".to_vec()),
|
||||
redis::Value::Int(1),
|
||||
redis::Value::BulkString(b"3.5".to_vec()),
|
||||
]),
|
||||
)])
|
||||
.assert_all_commands_consumed();
|
||||
let counters = RedisCache::with_connection(connection, None, JsonCodec::<f64>::new())
|
||||
.with_namespace(Some("team".into()));
|
||||
assert_eq!(
|
||||
counters
|
||||
.async_increment_pipeline(vec![
|
||||
IncrementOperation {
|
||||
key: "counter".into(),
|
||||
amount: 1.5,
|
||||
ttl: Some(Duration::from_secs(10)),
|
||||
},
|
||||
IncrementOperation {
|
||||
key: "counter".into(),
|
||||
amount: 2.0,
|
||||
ttl: None,
|
||||
},
|
||||
])
|
||||
.await
|
||||
.unwrap(),
|
||||
[1.5, 3.5]
|
||||
);
|
||||
}
|
||||
|
||||
const INCREMENT_WITH_FLOOR_SCRIPT: &str = concat!(
|
||||
"local count = redis.call('INCRBY', KEYS[1], ARGV[1]); ",
|
||||
"if count < 0 then count = redis.call('INCRBY', KEYS[1], -count); end; ",
|
||||
"if redis.call('TTL', KEYS[1]) < 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; ",
|
||||
"return count"
|
||||
);
|
||||
const SET_MAX_SCRIPT: &str = concat!(
|
||||
"local current = redis.call('GET', KEYS[1]); ",
|
||||
"if current == false or tonumber(current) < tonumber(ARGV[1]) then ",
|
||||
"redis.call('SET', KEYS[1], ARGV[1]); ",
|
||||
"if tonumber(ARGV[2]) > 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; ",
|
||||
"return ARGV[1]; end; return current"
|
||||
);
|
||||
|
||||
#[tokio::test]
|
||||
async fn counter_repairs_are_atomic_and_use_default_ttl() {
|
||||
let floor = || {
|
||||
redis::cmd("EVAL")
|
||||
.arg(INCREMENT_WITH_FLOOR_SCRIPT)
|
||||
.arg(1)
|
||||
.arg("team:counter")
|
||||
.arg(-2i64)
|
||||
.arg(30u64)
|
||||
.clone()
|
||||
};
|
||||
let connection = MockRedisConnection::new([
|
||||
MockCmd::new(floor(), Ok(0i64)),
|
||||
MockCmd::new(floor(), Ok(0i64)),
|
||||
MockCmd::new(
|
||||
redis::cmd("EVAL")
|
||||
.arg(SET_MAX_SCRIPT)
|
||||
.arg(1)
|
||||
.arg("team:counter")
|
||||
.arg(4.5f64)
|
||||
.arg(600u64),
|
||||
Ok("4.5"),
|
||||
),
|
||||
])
|
||||
.assert_all_commands_consumed();
|
||||
let cache = RedisCache::with_connection(connection, None, JsonCodec::<f64>::new())
|
||||
.with_namespace(Some("team".into()));
|
||||
|
||||
assert_eq!(
|
||||
cache
|
||||
.increment_with_floor("counter", -2, Duration::from_secs(30))
|
||||
.unwrap(),
|
||||
0
|
||||
);
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_increment_with_floor("counter", -2, Duration::from_secs(30))
|
||||
.await
|
||||
.unwrap(),
|
||||
0
|
||||
);
|
||||
assert_eq!(
|
||||
cache.async_set_max("counter", 4.5, None).await.unwrap(),
|
||||
4.5
|
||||
);
|
||||
}
|
||||
|
||||
const CLAIM_SCRIPT: &str = concat!(
|
||||
"local current = redis.call('GET', KEYS[1]); ",
|
||||
"if ARGV[1] == '' then if current ~= false and current ~= '' then return 0; end; ",
|
||||
"elseif current ~= ARGV[1] then return 0; end; ",
|
||||
"if ARGV[3] ~= '' then redis.call('SET', KEYS[1], ARGV[3], 'EX', ARGV[2]); ",
|
||||
"elseif ARGV[4] == '1' then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return 1"
|
||||
);
|
||||
|
||||
fn claim_eval(expected: &str, write: &str, refresh: bool) -> redis::Cmd {
|
||||
let mut cmd = redis::cmd("EVAL");
|
||||
cmd.arg(CLAIM_SCRIPT)
|
||||
.arg(1)
|
||||
.arg("pin")
|
||||
.arg(expected)
|
||||
.arg(600)
|
||||
.arg(write)
|
||||
.arg(u8::from(refresh));
|
||||
cmd
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn claims_match_eligible_values_written_by_another_encoder() {
|
||||
let python_payload = r#"{"model_id": "a", "deployment": "east"}"#;
|
||||
let stored = serde_json::json!({"deployment": "east", "model_id": "a"});
|
||||
let candidate = serde_json::json!({"model_id": "b"});
|
||||
let connection = MockRedisConnection::new([
|
||||
MockCmd::new(redis::cmd("GET").arg("pin"), Ok(python_payload)),
|
||||
MockCmd::new(claim_eval(python_payload, "", true), Ok(1)),
|
||||
])
|
||||
.assert_all_commands_consumed();
|
||||
let cache =
|
||||
RedisCache::with_connection(connection, None, JsonCodec::<serde_json::Value>::new());
|
||||
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_claim_cache(
|
||||
"pin",
|
||||
candidate,
|
||||
vec![stored.clone()],
|
||||
ExactCacheContext::default()
|
||||
)
|
||||
.await
|
||||
.unwrap(),
|
||||
stored
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn claims_retry_when_the_key_changes_and_replace_ineligible_winners() {
|
||||
let candidate = serde_json::json!({"model_id": "b"});
|
||||
let payload = r#"{"model_id":"b"}"#;
|
||||
let connection = MockRedisConnection::new([
|
||||
MockCmd::new(redis::cmd("GET").arg("pin"), Ok(redis::Value::Nil)),
|
||||
MockCmd::new(claim_eval("", payload, false), Ok(0)),
|
||||
MockCmd::new(redis::cmd("GET").arg("pin"), Ok(r#"{"model_id":"gone"}"#)),
|
||||
MockCmd::new(claim_eval(r#"{"model_id":"gone"}"#, payload, false), Ok(1)),
|
||||
])
|
||||
.assert_all_commands_consumed();
|
||||
let cache =
|
||||
RedisCache::with_connection(connection, None, JsonCodec::<serde_json::Value>::new());
|
||||
|
||||
assert_eq!(
|
||||
cache
|
||||
.claim_cache(
|
||||
"pin",
|
||||
candidate.clone(),
|
||||
&[serde_json::json!({"model_id": "a"})],
|
||||
ExactCacheContext::default()
|
||||
)
|
||||
.unwrap(),
|
||||
candidate
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn claims_without_eligible_values_keep_the_winner_without_refreshing_its_ttl() {
|
||||
let stored = r#"{"model_id": "a"}"#;
|
||||
let connection = MockRedisConnection::new([
|
||||
MockCmd::new(redis::cmd("GET").arg("pin"), Ok(stored)),
|
||||
MockCmd::new(claim_eval(stored, "", false), Ok(1)),
|
||||
])
|
||||
.assert_all_commands_consumed();
|
||||
let cache =
|
||||
RedisCache::with_connection(connection, None, JsonCodec::<serde_json::Value>::new());
|
||||
|
||||
assert_eq!(
|
||||
cache
|
||||
.claim_cache(
|
||||
"pin",
|
||||
serde_json::json!({"model_id": "b"}),
|
||||
&[],
|
||||
ExactCacheContext::default()
|
||||
)
|
||||
.unwrap(),
|
||||
serde_json::json!({"model_id": "a"})
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn async_increment_runs_the_atomic_script() {
|
||||
let mut eval = redis::cmd("EVAL");
|
||||
eval.arg(concat!(
|
||||
"local value = redis.call('INCRBYFLOAT', KEYS[1], ARGV[1]); ",
|
||||
"if redis.call('TTL', KEYS[1]) == -1 then ",
|
||||
"redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return value"
|
||||
))
|
||||
.arg(1)
|
||||
.arg("counter")
|
||||
.arg(2.5f64)
|
||||
.arg(600);
|
||||
let connection =
|
||||
MockRedisConnection::new([MockCmd::new(eval, Ok("4.5"))]).assert_all_commands_consumed();
|
||||
let cache = RedisCache::with_connection(connection, None, JsonCodec::<f64>::new());
|
||||
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_increment("counter", 2.5, ExactCacheContext::default())
|
||||
.await
|
||||
.unwrap(),
|
||||
4.5
|
||||
);
|
||||
}
|
||||
|
|
|
|||
492
litellm-rust/crates/cache-redis/tests/cluster.rs
Normal file
492
litellm-rust/crates/cache-redis/tests/cluster.rs
Normal 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
20
litellm-rust/crates/cache-response/Cargo.toml
Normal file
|
|
@ -0,0 +1,20 @@
|
|||
[package]
|
||||
name = "litellm-cache-response"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
litellm-cache.workspace = true
|
||||
py_literal = "0.4.0"
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
sha2.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
litellm-cache-memory.workspace = true
|
||||
litellm-cache-redis.workspace = true
|
||||
redis = "1.7.0"
|
||||
redis-test = "1.0.4"
|
||||
tokio.workspace = true
|
||||
61
litellm-rust/crates/cache-response/README.md
Normal file
61
litellm-rust/crates/cache-response/README.md
Normal file
|
|
@ -0,0 +1,61 @@
|
|||
# Response cache foundation
|
||||
|
||||
`ResponseCache<B>` adds request keys, independent read/write controls, response envelopes, and freshness checks to any `B: BaseCache<Value = CacheEntry>`
|
||||
|
||||
## Ownership
|
||||
|
||||
`litellm-cache` defines typed storage and codec traits. Memory and Redis implement those traits without depending on response policy. Other consumers can store their own value types using the same backend implementations
|
||||
|
||||
`litellm-cache-response` owns response keys, controls, entries, the Python-compatible response codec, and `WriteBuffer`, the backend-neutral deferred-write policy. It has no runtime dependency on a specific cache backend or Python
|
||||
|
||||
The Python bridge constructs backends and selects them through its private `NativeResponseCache` enum, which only dispatches. Generic Rust callers inject their backend directly. A native gateway can construct the same generic response service in its own host
|
||||
|
||||
## Native Rust use
|
||||
|
||||
```rust
|
||||
use std::{sync::Arc, time::Duration};
|
||||
use litellm_cache_memory::InMemoryCache;
|
||||
use litellm_cache_response::{CacheKeyInput, ResponseCache, ResponseCacheRequest};
|
||||
use serde_json::json;
|
||||
|
||||
let cache = ResponseCache::new(Arc::new(InMemoryCache::default()));
|
||||
let request = ResponseCacheRequest::new(CacheKeyInput {
|
||||
preset: Some("example:key".into()),
|
||||
..Default::default()
|
||||
});
|
||||
let now = Duration::from_secs(100);
|
||||
cache.store(&request, json!({"answer": 7}), now)?;
|
||||
assert_eq!(cache.async_lookup(&request, now).await?, Some(json!({"answer": 7})));
|
||||
```
|
||||
|
||||
For Redis, inject `RedisCache::new(url, ttl, ResponseCacheCodec)` instead. Namespaces are optional and existing namespace prefixes are preserved. Sync operations check out independent connections from a bounded pool, while async callers, including counters and claims, move that blocking work off the executor. The pool skips the checkout PING and instead discards any connection whose command failed
|
||||
|
||||
Callers supply Unix time for response freshness. Backend TTL uses its own clock. A read can reject an entry through `max_age` even while the backend still retains it
|
||||
|
||||
## Python integration boundary
|
||||
|
||||
The extension keeps a private test harness for memory and Redis single and batch response lookup and storage. Batch lookup returns ordered values plus missing indices for embedding partial-hit wiring. No bridge-only cache type is part of the public API
|
||||
|
||||
Object responses are written as they are, and every other response shape is written as a serialized string, which is the pair of shapes Python reads. A string on the wire is therefore always a serialized response, so string-valued responses round trip. Typed backends such as memory never pass through the codec
|
||||
|
||||
The resolver reads the namespace's `cache` attribute each time it resolves. A captured binding retains the selected service for its operation, including background writes. `None` disables caching. Custom Python cache objects keep their original methods, arguments, returned awaitables, exceptions, and caller-task execution
|
||||
|
||||
Python callbacks use the built-in `Cache` API, so a `Cache` subclass works unchanged. A batch lookup takes one original kwargs mapping per request and returns the list of `get_cache` or gathered `async_get_cache` results, while native bindings return `{values, missing_indices}`. A batch store hands the caller's original result to `async_add_cache_pipeline`. `ping` calls `ping`, and a flush goes to the facade's backend
|
||||
|
||||
The private facade test harness checks object identity, method overrides, effective TTL, Redis namespace, memory capacity, and later configuration changes before selecting native execution. Its snapshot includes Redis connection settings, so a later `redis_kwargs` change, including an SSL option, selects Python callback execution. Buffered async writes honor `redis_flush_size`. Public activation must construct the shared native service from the initial Python Redis settings, including `litellm.default_redis_ttl` and SSL options. A buffered entry keeps the time it was produced, and a failed flush drops its batch instead of growing the buffer during an outage. The harness does not migrate entries or replace Python methods. Until activation configures one shared service, the Python facade and native test service can hold separate data. Existing public cache constructors remain on Python
|
||||
|
||||
Native cache handles must be recreated after fork. The bridge releases the GIL around native operations, and Redis runs blocking connection operations off the async executor. Native errors propagate to the host, which owns the existing fail-open and logging policy
|
||||
|
||||
The Redis backend also provides the primitives needed to preserve its direct Python surface later: TLS URLs, ping, bulk delete, counter batches, TTL, scan, set membership, raw queue push and pop, queue and counter pipelines, counter floor and maximum operations, script evaluation, client information, namespaced flush, and full flush. These are backend operations only and are not exported to Python by this PR. Memory provides TTL, oldest-key, and counter-pipeline operations
|
||||
|
||||
## Adding another backend
|
||||
|
||||
Implement `BaseCache` for the backend with its associated value type, and accept a `CacheCodec` when wire serialization is needed. `ResponseCache<B>` then works without another response implementation. Add a concrete bridge enum variant and constructor only when exposing that backend to Python
|
||||
|
||||
Verify typed values, TTL precedence, missing entries, serialization failures, namespaces, batch ordering, and sync/async behavior. Run response fixtures with `ResponseCacheCodec`, including both Python envelope encodings, before enabling a public facade
|
||||
|
||||
## Follow-up scope
|
||||
|
||||
Public SDK, Router, and proxy activation still need constructor parity, stream replay, embedding partial-batch integration, response reconstruction, callback scheduling, and failure-policy integration. This foundation does not switch those request paths
|
||||
|
||||
Redis cluster, disk, and cloud stores remain follow-ups. Semantic backends plug in through `SemanticCacheContext`, which carries the prompt inputs and metadata alongside the cache TTL. The generic dual cache takes read, write, and remote-failure policies, runs its async operations through the async L2 methods, and provides L2-first counters and atomic affinity claims. Errors propagate by default, and `RemoteFailurePolicy::UseLocal` opts key-value operations and claims into the local tier when L2 is unavailable. Claims compare decoded values, so a pin written by Python still matches. Public Router integration remains follow-up work. Reservations and pubsub still need explicit capabilities owned by their consuming features. Adding a cache backend does not establish those guarantees
|
||||
45
litellm-rust/crates/cache-response/src/buffer.rs
Normal file
45
litellm-rust/crates/cache-response/src/buffer.rs
Normal file
|
|
@ -0,0 +1,45 @@
|
|||
use std::{sync::Mutex, time::Duration};
|
||||
|
||||
use litellm_cache::{BaseCache, Error, ExactCacheContext};
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{CacheEntry, ResponseCache, ResponseCacheRequest};
|
||||
|
||||
pub struct WriteBuffer {
|
||||
flush_size: usize,
|
||||
entries: Mutex<Vec<(ResponseCacheRequest, Value, Duration)>>,
|
||||
}
|
||||
|
||||
impl WriteBuffer {
|
||||
pub fn new(flush_size: usize) -> Self {
|
||||
Self {
|
||||
flush_size: flush_size.max(1),
|
||||
entries: Mutex::new(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn async_store<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>>(
|
||||
&self,
|
||||
cache: &ResponseCache<B>,
|
||||
request: &ResponseCacheRequest,
|
||||
response: Value,
|
||||
now: Duration,
|
||||
) -> Result<(), Error> {
|
||||
let pending = {
|
||||
let mut entries = self.entries.lock().map_err(|_| Error::Unavailable)?;
|
||||
entries.push((request.clone(), response, now));
|
||||
(entries.len() >= self.flush_size).then(|| std::mem::take(&mut *entries))
|
||||
};
|
||||
// A failed flush drops its batch, as Python does. Requeueing would grow the
|
||||
// buffer and re-send an ever larger pipeline on every write during an outage.
|
||||
match pending {
|
||||
Some(pending) => cache.async_store_entries(pending).await,
|
||||
None => Ok(()),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn clear(&self) -> Result<(), Error> {
|
||||
self.entries.lock().map_err(|_| Error::Unavailable)?.clear();
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
147
litellm-rust/crates/cache-response/src/caching.rs
Normal file
147
litellm-rust/crates/cache-response/src/caching.rs
Normal file
|
|
@ -0,0 +1,147 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize, PartialEq, Eq)]
|
||||
pub enum CacheMode {
|
||||
#[default]
|
||||
#[serde(rename = "default_on")]
|
||||
DefaultOn,
|
||||
#[serde(rename = "default_off")]
|
||||
DefaultOff,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Deserialize, Serialize)]
|
||||
pub struct CacheKeyField {
|
||||
pub name: String,
|
||||
pub value: Option<String>,
|
||||
pub api_parameter: bool,
|
||||
pub internal_parameter: bool,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default, Deserialize, Serialize)]
|
||||
#[serde(default)]
|
||||
pub struct CacheKeyInput {
|
||||
pub fields: Vec<CacheKeyField>,
|
||||
pub preset: Option<String>,
|
||||
pub namespace: Option<String>,
|
||||
pub include_provider_parameters: bool,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
pub struct CacheKeyContext {
|
||||
pub model_group: Option<String>,
|
||||
pub caching_groups: Vec<(Vec<String>, String)>,
|
||||
pub file_checksum: Option<String>,
|
||||
pub file_object_name: Option<String>,
|
||||
pub metadata_file_name: Option<String>,
|
||||
pub parameters_file_name: Option<String>,
|
||||
}
|
||||
|
||||
impl CacheKeyContext {
|
||||
pub fn apply(self, input: &mut CacheKeyInput) {
|
||||
let group = self.model_group.as_ref().and_then(|model| {
|
||||
self.caching_groups
|
||||
.iter()
|
||||
.find(|(models, _)| models.contains(model))
|
||||
});
|
||||
for field in &mut input.fields {
|
||||
match field.name.as_str() {
|
||||
"model" => {
|
||||
field.value = group
|
||||
.map(|(_, formatted)| formatted.clone())
|
||||
.or_else(|| self.model_group.clone())
|
||||
.or_else(|| field.value.take())
|
||||
}
|
||||
"file" => {
|
||||
field.value = self
|
||||
.file_checksum
|
||||
.clone()
|
||||
.or_else(|| self.file_object_name.clone())
|
||||
.or_else(|| self.metadata_file_name.clone())
|
||||
.or_else(|| self.parameters_file_name.clone())
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn get_cache_key(input: &CacheKeyInput) -> String {
|
||||
cache_key(input)
|
||||
}
|
||||
|
||||
pub fn cache_key(input: &CacheKeyInput) -> String {
|
||||
if let Some(preset) = &input.preset {
|
||||
return preset.clone();
|
||||
}
|
||||
let mut digest = Sha256::new();
|
||||
for field in &input.fields {
|
||||
if (field.api_parameter || (input.include_provider_parameters && !field.internal_parameter))
|
||||
&& let Some(value) = &field.value
|
||||
{
|
||||
digest.update(field.name.as_bytes());
|
||||
digest.update(b": ");
|
||||
digest.update(value.as_bytes());
|
||||
}
|
||||
}
|
||||
let hash = format!("{:x}", digest.finalize());
|
||||
input
|
||||
.namespace
|
||||
.as_deref()
|
||||
.filter(|namespace| !namespace.is_empty())
|
||||
.map_or(hash.clone(), |namespace| format!("{namespace}:{hash}"))
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize)]
|
||||
pub struct CacheControls {
|
||||
pub supported_call_type: bool,
|
||||
pub configured: bool,
|
||||
pub native_backend: bool,
|
||||
pub default_on: bool,
|
||||
pub caching: Option<bool>,
|
||||
pub no_cache: bool,
|
||||
pub no_store: bool,
|
||||
#[serde(default)]
|
||||
pub use_cache: bool,
|
||||
}
|
||||
|
||||
impl CacheControls {
|
||||
pub fn reads(self) -> bool {
|
||||
self.supported_call_type
|
||||
&& self.configured
|
||||
&& self.caching.unwrap_or(true)
|
||||
&& !self.no_cache
|
||||
&& (self.default_on || self.use_cache)
|
||||
}
|
||||
|
||||
pub fn writes(self) -> bool {
|
||||
self.supported_call_type
|
||||
&& self.configured
|
||||
&& self.caching.unwrap_or(true)
|
||||
&& !self.no_store
|
||||
&& (self.default_on || self.use_cache)
|
||||
}
|
||||
}
|
||||
|
||||
pub fn should_use_cache(controls: CacheControls) -> bool {
|
||||
controls.reads() || controls.writes()
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct CacheEntry {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub timestamp: Option<f64>,
|
||||
pub response: Value,
|
||||
}
|
||||
|
||||
impl CacheEntry {
|
||||
pub fn fresh(&self, now: Duration, max_age: Option<Duration>) -> bool {
|
||||
self.timestamp.is_none_or(|timestamp| {
|
||||
timestamp.is_finite()
|
||||
&& max_age.is_none_or(|age| now.as_secs_f64() - timestamp <= age.as_secs_f64())
|
||||
})
|
||||
}
|
||||
}
|
||||
129
litellm-rust/crates/cache-response/src/codec.rs
Normal file
129
litellm-rust/crates/cache-response/src/codec.rs
Normal file
|
|
@ -0,0 +1,129 @@
|
|||
use litellm_cache::{CacheCodec, Error};
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::CacheEntry;
|
||||
|
||||
#[derive(Clone, Copy, Debug, Default)]
|
||||
pub struct ResponseCacheCodec;
|
||||
|
||||
impl CacheCodec for ResponseCacheCodec {
|
||||
type Value = CacheEntry;
|
||||
|
||||
fn encode(&self, value: &CacheEntry) -> Result<Vec<u8>, Error> {
|
||||
if value
|
||||
.timestamp
|
||||
.is_some_and(|timestamp| !timestamp.is_finite())
|
||||
{
|
||||
return Err(Error::InvalidEntry);
|
||||
}
|
||||
// Python reads a `response` that is either a dict or a serialized string, so every
|
||||
// other shape is written serialized. A string on the wire is therefore always a
|
||||
// serialized response, which keeps string-valued responses unambiguous.
|
||||
if value.timestamp.is_none() || value.response.is_object() {
|
||||
return serde_json::to_vec(value).map_err(|_| Error::InvalidEntry);
|
||||
}
|
||||
let response = serde_json::to_string(&value.response).map_err(|_| Error::InvalidEntry)?;
|
||||
serde_json::to_vec(&CacheEntry {
|
||||
timestamp: value.timestamp,
|
||||
response: Value::String(response),
|
||||
})
|
||||
.map_err(|_| Error::InvalidEntry)
|
||||
}
|
||||
|
||||
fn decode(&self, bytes: &[u8]) -> Result<CacheEntry, Error> {
|
||||
let text = std::str::from_utf8(bytes).map_err(|_| Error::InvalidEntry)?;
|
||||
let value = decode_value(text)?;
|
||||
let Some(timestamp) = value.get("timestamp") else {
|
||||
return Ok(CacheEntry {
|
||||
timestamp: None,
|
||||
response: value,
|
||||
});
|
||||
};
|
||||
let Some(timestamp) = timestamp.as_f64().filter(|timestamp| timestamp.is_finite()) else {
|
||||
return Err(Error::InvalidEntry);
|
||||
};
|
||||
let response = match value.get("response").ok_or(Error::InvalidEntry)? {
|
||||
Value::String(text) => decode_value(text)?,
|
||||
response => response.clone(),
|
||||
};
|
||||
Ok(CacheEntry {
|
||||
timestamp: Some(timestamp),
|
||||
response,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn decode_value(text: &str) -> Result<Value, Error> {
|
||||
if let Ok(value) = serde_json::from_str(text) {
|
||||
return Ok(value);
|
||||
}
|
||||
check_literal_depth(text)?;
|
||||
let literal: py_literal::Value = text.parse().map_err(|_| Error::InvalidEntry)?;
|
||||
literal_value(literal, 0)
|
||||
}
|
||||
|
||||
fn literal_value(value: py_literal::Value, depth: usize) -> Result<Value, Error> {
|
||||
use py_literal::Value as Literal;
|
||||
if depth > 128 {
|
||||
return Err(Error::InvalidEntry);
|
||||
}
|
||||
match value {
|
||||
Literal::String(text) => Ok(Value::String(text)),
|
||||
Literal::Boolean(value) => Ok(Value::Bool(value)),
|
||||
Literal::None => Ok(Value::Null),
|
||||
Literal::Integer(value) => {
|
||||
serde_json::from_str(&value.to_string()).map_err(|_| Error::InvalidEntry)
|
||||
}
|
||||
Literal::Float(value) => serde_json::Number::from_f64(value)
|
||||
.map(Value::Number)
|
||||
.ok_or(Error::InvalidEntry),
|
||||
Literal::List(values) | Literal::Tuple(values) => values
|
||||
.into_iter()
|
||||
.map(|value| literal_value(value, depth + 1))
|
||||
.collect::<Result<Vec<_>, _>>()
|
||||
.map(Value::Array),
|
||||
Literal::Dict(entries) => entries
|
||||
.into_iter()
|
||||
.map(|(key, value)| {
|
||||
let Literal::String(key) = key else {
|
||||
return Err(Error::InvalidEntry);
|
||||
};
|
||||
Ok((key, literal_value(value, depth + 1)?))
|
||||
})
|
||||
.collect::<Result<serde_json::Map<_, _>, _>>()
|
||||
.map(Value::Object),
|
||||
_ => Err(Error::InvalidEntry),
|
||||
}
|
||||
}
|
||||
|
||||
fn check_literal_depth(text: &str) -> Result<(), Error> {
|
||||
let mut quote = None;
|
||||
let mut escaped = false;
|
||||
let mut depth = 0usize;
|
||||
for ch in text.chars() {
|
||||
if escaped {
|
||||
escaped = false;
|
||||
continue;
|
||||
}
|
||||
if let Some(delimiter) = quote {
|
||||
if ch == '\\' {
|
||||
escaped = true;
|
||||
} else if ch == delimiter {
|
||||
quote = None;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
match ch {
|
||||
'\'' | '"' => quote = Some(ch),
|
||||
'[' | '{' | '(' => {
|
||||
depth += 1;
|
||||
if depth > 128 {
|
||||
return Err(Error::InvalidEntry);
|
||||
}
|
||||
}
|
||||
']' | '}' | ')' => depth = depth.saturating_sub(1),
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
22
litellm-rust/crates/cache-response/src/embedding.rs
Normal file
22
litellm-rust/crates/cache-response/src/embedding.rs
Normal file
|
|
@ -0,0 +1,22 @@
|
|||
use serde::Serialize;
|
||||
use serde_json::Value;
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize)]
|
||||
pub struct PartialHits {
|
||||
pub values: Vec<Option<Value>>,
|
||||
pub missing_indices: Vec<usize>,
|
||||
}
|
||||
|
||||
impl PartialHits {
|
||||
pub fn new(values: Vec<Option<Value>>) -> Self {
|
||||
let missing_indices = values
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter_map(|(index, value)| value.is_none().then_some(index))
|
||||
.collect();
|
||||
Self {
|
||||
values,
|
||||
missing_indices,
|
||||
}
|
||||
}
|
||||
}
|
||||
14
litellm-rust/crates/cache-response/src/lib.rs
Normal file
14
litellm-rust/crates/cache-response/src/lib.rs
Normal file
|
|
@ -0,0 +1,14 @@
|
|||
mod buffer;
|
||||
mod caching;
|
||||
mod codec;
|
||||
mod embedding;
|
||||
mod response;
|
||||
|
||||
pub use buffer::WriteBuffer;
|
||||
pub use caching::{
|
||||
CacheControls, CacheEntry, CacheKeyContext, CacheKeyField, CacheKeyInput, CacheMode, cache_key,
|
||||
get_cache_key, should_use_cache,
|
||||
};
|
||||
pub use codec::ResponseCacheCodec;
|
||||
pub use embedding::PartialHits;
|
||||
pub use response::{ResponseCache, ResponseCacheRequest};
|
||||
291
litellm-rust/crates/cache-response/src/response.rs
Normal file
291
litellm-rust/crates/cache-response/src/response.rs
Normal file
|
|
@ -0,0 +1,291 @@
|
|||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use litellm_cache::{
|
||||
BaseCache, BatchCache, BatchEntry, CacheConnectionResult, CacheContext, Error, FlushCache,
|
||||
};
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{CacheControls, CacheEntry, CacheKeyInput, PartialHits, cache_key};
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct ResponseCacheRequest<C: CacheContext = litellm_cache::ExactCacheContext> {
|
||||
pub key: CacheKeyInput,
|
||||
pub controls: CacheControls,
|
||||
pub context: C,
|
||||
pub max_age: Option<Duration>,
|
||||
}
|
||||
|
||||
impl<C: CacheContext + Default> ResponseCacheRequest<C> {
|
||||
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: C::default(),
|
||||
max_age: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct ResponseCache<B: BaseCache<Value = CacheEntry>>
|
||||
where
|
||||
B::Context: Default + PartialEq,
|
||||
{
|
||||
backend: Arc<B>,
|
||||
}
|
||||
|
||||
impl<B> ResponseCache<B>
|
||||
where
|
||||
B: BaseCache<Value = CacheEntry>,
|
||||
B::Context: Default + PartialEq,
|
||||
{
|
||||
pub fn new(backend: Arc<B>) -> Self {
|
||||
Self { backend }
|
||||
}
|
||||
|
||||
pub fn backend(&self) -> &B {
|
||||
&self.backend
|
||||
}
|
||||
|
||||
pub fn backend_arc(&self) -> &Arc<B> {
|
||||
&self.backend
|
||||
}
|
||||
|
||||
pub fn default_ttl(&self) -> Option<Duration> {
|
||||
self.backend.get_ttl(&B::Context::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<B::Context>,
|
||||
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<B::Context>,
|
||||
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<B::Context>],
|
||||
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<B::Context>],
|
||||
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<B::Context>,
|
||||
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<B::Context>,
|
||||
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<B::Context>, 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<B::Context>, 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<B::Context>],
|
||||
readable: Vec<(usize, &ResponseCacheRequest<B::Context>)>,
|
||||
entries: Vec<BatchEntry<CacheEntry>>,
|
||||
now: Duration,
|
||||
) -> Result<PartialHits, Error> {
|
||||
if readable.len() != entries.len() {
|
||||
return Err(Error::Unavailable);
|
||||
}
|
||||
let mut values = vec![None; requests.len()];
|
||||
for ((index, request), entry) in readable.into_iter().zip(entries) {
|
||||
let response = match entry {
|
||||
BatchEntry::Hit(entry) => Self::fresh_or_miss(Some(entry), now, request.max_age),
|
||||
BatchEntry::Miss | BatchEntry::Invalid => None,
|
||||
};
|
||||
values[index] = response;
|
||||
}
|
||||
Ok(PartialHits::new(values))
|
||||
}
|
||||
|
||||
fn fresh_or_miss(
|
||||
entry: Option<CacheEntry>,
|
||||
now: Duration,
|
||||
max_age: Option<Duration>,
|
||||
) -> Option<Value> {
|
||||
entry
|
||||
.filter(|entry| entry.fresh(now, max_age))
|
||||
.map(|entry| entry.response)
|
||||
}
|
||||
}
|
||||
90
litellm-rust/crates/cache-response/tests/caching.rs
Normal file
90
litellm-rust/crates/cache-response/tests/caching.rs
Normal file
|
|
@ -0,0 +1,90 @@
|
|||
use litellm_cache_response::{
|
||||
CacheControls, CacheKeyContext, CacheKeyField, CacheKeyInput, cache_key, get_cache_key,
|
||||
};
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
#[test]
|
||||
fn keys_match_python_order_groups_files_presets_and_namespaces() {
|
||||
let mut input = CacheKeyInput {
|
||||
fields: vec![
|
||||
CacheKeyField {
|
||||
name: "model".into(),
|
||||
value: Some("deployment".into()),
|
||||
api_parameter: true,
|
||||
internal_parameter: false,
|
||||
},
|
||||
CacheKeyField {
|
||||
name: "file".into(),
|
||||
value: None,
|
||||
api_parameter: true,
|
||||
internal_parameter: false,
|
||||
},
|
||||
],
|
||||
namespace: Some("team".into()),
|
||||
..Default::default()
|
||||
};
|
||||
CacheKeyContext {
|
||||
model_group: Some("group".into()),
|
||||
caching_groups: vec![(vec!["group".into()], "['group']".into())],
|
||||
file_checksum: Some("checksum".into()),
|
||||
..Default::default()
|
||||
}
|
||||
.apply(&mut input);
|
||||
assert_eq!(
|
||||
cache_key(&input),
|
||||
format!(
|
||||
"team:{:x}",
|
||||
Sha256::digest(b"model: ['group']file: checksum")
|
||||
)
|
||||
);
|
||||
input.preset = Some("preset".into());
|
||||
assert_eq!(get_cache_key(&input), "preset");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_controls_honor_default_modes_and_directives() {
|
||||
let enabled = CacheControls {
|
||||
supported_call_type: true,
|
||||
configured: true,
|
||||
default_on: true,
|
||||
..Default::default()
|
||||
};
|
||||
assert!(enabled.reads());
|
||||
assert!(enabled.writes());
|
||||
assert!(
|
||||
!CacheControls {
|
||||
default_on: false,
|
||||
..enabled
|
||||
}
|
||||
.reads()
|
||||
);
|
||||
assert!(
|
||||
CacheControls {
|
||||
default_on: false,
|
||||
use_cache: true,
|
||||
..enabled
|
||||
}
|
||||
.reads()
|
||||
);
|
||||
assert!(
|
||||
!CacheControls {
|
||||
no_cache: true,
|
||||
..enabled
|
||||
}
|
||||
.reads()
|
||||
);
|
||||
assert!(
|
||||
!CacheControls {
|
||||
no_store: true,
|
||||
..enabled
|
||||
}
|
||||
.writes()
|
||||
);
|
||||
assert!(
|
||||
!CacheControls {
|
||||
caching: Some(false),
|
||||
..enabled
|
||||
}
|
||||
.writes()
|
||||
);
|
||||
}
|
||||
484
litellm-rust/crates/cache-response/tests/response.rs
Normal file
484
litellm-rust/crates/cache-response/tests/response.rs
Normal file
|
|
@ -0,0 +1,484 @@
|
|||
use std::{
|
||||
sync::{
|
||||
Arc,
|
||||
atomic::{AtomicU64, Ordering},
|
||||
},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use litellm_cache::{BaseCache, CacheCodec, Error};
|
||||
use litellm_cache_memory::InMemoryCache;
|
||||
use litellm_cache_redis::RedisCache;
|
||||
use litellm_cache_response::{
|
||||
CacheEntry, CacheKeyField, CacheKeyInput, ResponseCache, ResponseCacheCodec,
|
||||
ResponseCacheRequest, WriteBuffer,
|
||||
};
|
||||
use redis_test::{MockCmd, MockRedisConnection};
|
||||
use serde_json::json;
|
||||
|
||||
fn memory() -> Arc<ResponseCache<InMemoryCache<CacheEntry>>> {
|
||||
Arc::new(ResponseCache::new(Arc::new(InMemoryCache::new(
|
||||
Some(8),
|
||||
Some(Duration::from_secs(600)),
|
||||
))))
|
||||
}
|
||||
|
||||
fn request() -> ResponseCacheRequest {
|
||||
ResponseCacheRequest::new(CacheKeyInput {
|
||||
preset: Some("tenant:key".into()),
|
||||
..Default::default()
|
||||
})
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sync_and_async_consumers_share_keys_ttls_and_freshness() {
|
||||
let clock = Arc::new(AtomicU64::new(100));
|
||||
let backend = Arc::new(InMemoryCache::with_clock(
|
||||
Some(8),
|
||||
Some(Duration::from_secs(600)),
|
||||
{
|
||||
let clock = clock.clone();
|
||||
move || Duration::from_secs(clock.load(Ordering::SeqCst))
|
||||
},
|
||||
));
|
||||
let cache = ResponseCache::new(backend.clone());
|
||||
let mut request = request();
|
||||
request.context.ttl = Some(Duration::from_secs(10));
|
||||
request.max_age = Some(Duration::from_secs(5));
|
||||
cache
|
||||
.store(
|
||||
&request,
|
||||
json!({"choices": [1], "usage": {"total_tokens": 7}}),
|
||||
Duration::from_secs(100),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
backend.expires_at("tenant:key").unwrap(),
|
||||
Some(Duration::from_secs(110))
|
||||
);
|
||||
assert!(
|
||||
cache
|
||||
.async_lookup(&request, Duration::from_secs(105))
|
||||
.await
|
||||
.unwrap()
|
||||
.is_some()
|
||||
);
|
||||
assert_eq!(
|
||||
cache.lookup(&request, Duration::from_secs(106)).unwrap(),
|
||||
None
|
||||
);
|
||||
request.max_age = None;
|
||||
assert_eq!(
|
||||
cache
|
||||
.lookup(&request, Duration::from_secs(106))
|
||||
.unwrap()
|
||||
.unwrap()["usage"]["total_tokens"],
|
||||
7
|
||||
);
|
||||
clock.store(111, Ordering::SeqCst);
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_lookup(&request, Duration::from_secs(111))
|
||||
.await
|
||||
.unwrap(),
|
||||
None
|
||||
);
|
||||
cache
|
||||
.async_store(&request, json!({"choices": [2]}), Duration::from_secs(111))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
cache.lookup(&request, Duration::from_secs(111)).unwrap(),
|
||||
Some(json!({"choices": [2]}))
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn directives_skip_io_and_keep_reads_and_writes_independent() {
|
||||
let cache = memory();
|
||||
let mut request = request();
|
||||
let now = Duration::from_secs(100);
|
||||
request.controls.no_store = true;
|
||||
cache
|
||||
.async_store(&request, json!({"v": 1}), now)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(cache.lookup(&request, now).unwrap(), None);
|
||||
request.controls.no_store = false;
|
||||
request.controls.no_cache = true;
|
||||
cache.store(&request, json!({"v": 2}), now).unwrap();
|
||||
assert_eq!(cache.async_lookup(&request, now).await.unwrap(), None);
|
||||
request.controls.no_cache = false;
|
||||
assert_eq!(cache.lookup(&request, now).unwrap(), Some(json!({"v": 2})));
|
||||
request.controls.default_on = false;
|
||||
cache.store(&request, json!({"v": 3}), now).unwrap();
|
||||
assert_eq!(cache.lookup(&request, now).unwrap(), None);
|
||||
request.controls.use_cache = true;
|
||||
assert_eq!(cache.lookup(&request, now).unwrap(), Some(json!({"v": 2})));
|
||||
request.controls.supported_call_type = false;
|
||||
assert_eq!(cache.lookup(&request, now).unwrap(), None);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn redis_consumer_reads_python_sync_and_async_envelopes_and_writes_compatible_json() {
|
||||
let connection = MockRedisConnection::new([
|
||||
MockCmd::new(
|
||||
redis::cmd("GET").arg("tenant:key"),
|
||||
Ok(br#"{'timestamp': 100.0, 'response': '{"ok": true, "text": "cached"}'}"#.to_vec()),
|
||||
),
|
||||
MockCmd::new(
|
||||
redis::cmd("GET").arg("tenant:key"),
|
||||
Ok(br#"{"timestamp":100.0,"response":{"ok":true,"text":"cached"}}"#.to_vec()),
|
||||
),
|
||||
MockCmd::new(
|
||||
redis::cmd("SETEX")
|
||||
.arg("tenant:key")
|
||||
.arg(600)
|
||||
.arg(br#"{"timestamp":100.0,"response":{"ok":true,"text":"cached"}}"#.as_slice()),
|
||||
Ok("OK"),
|
||||
),
|
||||
])
|
||||
.assert_all_commands_consumed();
|
||||
let backend = RedisCache::with_connection(connection, None, ResponseCacheCodec)
|
||||
.with_namespace(Some("tenant".into()));
|
||||
let cache = ResponseCache::new(Arc::new(backend));
|
||||
let request = request();
|
||||
let expected = json!({"ok": true, "text": "cached"});
|
||||
assert_eq!(
|
||||
cache.lookup(&request, Duration::from_secs(101)).unwrap(),
|
||||
Some(expected.clone())
|
||||
);
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_lookup(&request, Duration::from_secs(101))
|
||||
.await
|
||||
.unwrap(),
|
||||
Some(expected.clone())
|
||||
);
|
||||
cache
|
||||
.async_store(&request, expected, Duration::from_secs(100))
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn captured_service_keeps_the_selected_backend_for_background_writes() {
|
||||
let original = memory();
|
||||
let captured = original.clone();
|
||||
let replacement = memory();
|
||||
let request = request();
|
||||
let writer = tokio::spawn({
|
||||
let request = request.clone();
|
||||
async move {
|
||||
captured
|
||||
.async_store(
|
||||
&request,
|
||||
json!({"selected": "original"}),
|
||||
Duration::from_secs(100),
|
||||
)
|
||||
.await
|
||||
}
|
||||
});
|
||||
writer.await.unwrap().unwrap();
|
||||
assert_eq!(
|
||||
original.lookup(&request, Duration::from_secs(100)).unwrap(),
|
||||
Some(json!({"selected":"original"}))
|
||||
);
|
||||
assert_eq!(
|
||||
replacement
|
||||
.lookup(&request, Duration::from_secs(100))
|
||||
.unwrap(),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn generated_keys_preserve_namespace_and_explicit_keys() {
|
||||
let cache = memory();
|
||||
let key = CacheKeyInput {
|
||||
fields: vec![CacheKeyField {
|
||||
name: "model".into(),
|
||||
value: Some("a".into()),
|
||||
api_parameter: true,
|
||||
internal_parameter: false,
|
||||
}],
|
||||
namespace: Some("tenant".into()),
|
||||
..Default::default()
|
||||
};
|
||||
let generated = ResponseCacheRequest::new(key.clone());
|
||||
let explicit = ResponseCacheRequest::new(CacheKeyInput {
|
||||
preset: Some(litellm_cache_response::cache_key(&key)),
|
||||
..Default::default()
|
||||
});
|
||||
cache
|
||||
.store(&generated, json!({"value": 7}), Duration::from_secs(100))
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
cache.lookup(&explicit, Duration::from_secs(100)).unwrap(),
|
||||
Some(json!({"value":7}))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn response_codec_accepts_python_literals_without_executing_code() {
|
||||
let bytes = br#"{'timestamp': 100.0, 'response': {'text': 'hello \\ world', 'flag': True, 'empty': None, 'list': [1, 2.5]}}"#;
|
||||
let entry = ResponseCacheCodec.decode(bytes).unwrap();
|
||||
assert_eq!(
|
||||
entry.response,
|
||||
json!({"text": "hello \\ world", "flag": true, "empty": null, "list": [1, 2.5]})
|
||||
);
|
||||
for bytes in [
|
||||
b"__import__('os').system('false')".as_slice(),
|
||||
b"{'timestamp': 'invalid', 'response': {}}",
|
||||
b"{'timestamp': 1e9999, 'response': {}}",
|
||||
] {
|
||||
assert_eq!(
|
||||
ResponseCacheCodec.decode(bytes).unwrap_err(),
|
||||
Error::InvalidEntry
|
||||
);
|
||||
}
|
||||
let deep = format!("{}None{}", "[".repeat(1000), "]".repeat(1000));
|
||||
assert_eq!(
|
||||
ResponseCacheCodec.decode(deep.as_bytes()).unwrap_err(),
|
||||
Error::InvalidEntry
|
||||
);
|
||||
assert_eq!(
|
||||
ResponseCacheCodec
|
||||
.encode(&CacheEntry {
|
||||
timestamp: Some(f64::NAN),
|
||||
response: json!({})
|
||||
})
|
||||
.unwrap_err(),
|
||||
Error::InvalidEntry
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn invalid_entries_are_misses_and_disabled_reads_do_not_touch_redis() {
|
||||
let connection = MockRedisConnection::new([MockCmd::new(
|
||||
redis::cmd("GET").arg("tenant:key"),
|
||||
Ok(b"invalid".to_vec()),
|
||||
)])
|
||||
.assert_all_commands_consumed();
|
||||
let backend = RedisCache::with_connection(connection, None, ResponseCacheCodec);
|
||||
let cache = ResponseCache::new(Arc::new(backend));
|
||||
let mut request = request();
|
||||
request.controls.no_cache = true;
|
||||
assert_eq!(cache.lookup(&request, Duration::ZERO).unwrap(), None);
|
||||
request.controls.no_cache = false;
|
||||
assert_eq!(
|
||||
cache.async_lookup(&request, Duration::ZERO).await.unwrap(),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn string_responses_round_trip_through_typed_and_wire_backends() {
|
||||
let cache = ResponseCache::new(Arc::new(InMemoryCache::default()));
|
||||
let now = Duration::from_secs(100);
|
||||
for response in [json!("hello world"), json!("123"), json!("null")] {
|
||||
cache.store(&request(), response.clone(), now).unwrap();
|
||||
assert_eq!(
|
||||
cache.lookup(&request(), now).unwrap(),
|
||||
Some(response.clone())
|
||||
);
|
||||
|
||||
let wire = ResponseCacheCodec
|
||||
.encode(&CacheEntry {
|
||||
timestamp: Some(100.0),
|
||||
response: response.clone(),
|
||||
})
|
||||
.unwrap();
|
||||
assert_eq!(ResponseCacheCodec.decode(&wire).unwrap().response, response);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn non_object_responses_are_written_as_python_readable_serialized_strings() {
|
||||
let wire = ResponseCacheCodec
|
||||
.encode(&CacheEntry {
|
||||
timestamp: Some(100.0),
|
||||
response: json!([1, 2]),
|
||||
})
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
serde_json::from_slice::<serde_json::Value>(&wire).unwrap(),
|
||||
json!({"timestamp": 100.0, "response": "[1,2]"})
|
||||
);
|
||||
assert_eq!(
|
||||
ResponseCacheCodec.decode(&wire).unwrap().response,
|
||||
json!([1, 2])
|
||||
);
|
||||
assert_eq!(
|
||||
ResponseCacheCodec.decode(br#"{"timestamp": 100.0, "response": "not serialized"}"#),
|
||||
Err(Error::InvalidEntry)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn response_entries_preserve_the_existing_json_representation() {
|
||||
let codec = ResponseCacheCodec;
|
||||
let entry = CacheEntry {
|
||||
timestamp: Some(123.0),
|
||||
response: json!({"choices": [{"text": "cached"}]}),
|
||||
};
|
||||
let bytes = codec.encode(&entry).unwrap();
|
||||
assert_eq!(bytes, serde_json::to_vec(&entry).unwrap());
|
||||
assert_eq!(codec.decode(&bytes).unwrap(), entry);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn response_codec_preserves_values_without_timestamps() {
|
||||
let codec = ResponseCacheCodec;
|
||||
let raw = json!({"choices": [{"text": "legacy"}]});
|
||||
let entry = codec.decode(&serde_json::to_vec(&raw).unwrap()).unwrap();
|
||||
assert_eq!(entry.timestamp, None);
|
||||
assert_eq!(entry.response, raw);
|
||||
|
||||
let backend = Arc::new(InMemoryCache::default());
|
||||
BaseCache::set_cache(backend.as_ref(), "tenant:key", entry, &Default::default()).unwrap();
|
||||
let cache = ResponseCache::new(backend);
|
||||
assert_eq!(
|
||||
cache.lookup(&request(), Duration::from_secs(100)).unwrap(),
|
||||
Some(json!({"choices": [{"text": "legacy"}]}))
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn batch_lookup_reports_partial_hits_and_batch_store_populates_misses() {
|
||||
let cache = memory();
|
||||
let requests = ["hit", "miss", "disabled"].map(|key| {
|
||||
ResponseCacheRequest::new(CacheKeyInput {
|
||||
preset: Some(key.into()),
|
||||
..Default::default()
|
||||
})
|
||||
});
|
||||
cache
|
||||
.store(&requests[0], json!({"value": 1}), Duration::from_secs(100))
|
||||
.unwrap();
|
||||
let mut requests = requests.to_vec();
|
||||
requests[2].controls.caching = Some(false);
|
||||
|
||||
let partial = cache
|
||||
.async_lookup_batch(&requests, Duration::from_secs(100))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(partial.values, vec![Some(json!({"value": 1})), None, None]);
|
||||
assert_eq!(partial.missing_indices, vec![1, 2]);
|
||||
|
||||
cache
|
||||
.async_store_batch(
|
||||
vec![
|
||||
(requests[1].clone(), json!({"value": 2})),
|
||||
(requests[2].clone(), json!({"value": 3})),
|
||||
],
|
||||
Duration::from_secs(100),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
cache
|
||||
.lookup(&requests[1], Duration::from_secs(100))
|
||||
.unwrap(),
|
||||
Some(json!({"value": 2}))
|
||||
);
|
||||
requests[2].controls.caching = None;
|
||||
assert_eq!(
|
||||
cache
|
||||
.lookup(&requests[2], Duration::from_secs(100))
|
||||
.unwrap(),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn deferred_entries_keep_the_time_they_were_produced() {
|
||||
let cache = ResponseCache::new(Arc::new(InMemoryCache::default()));
|
||||
let mut request = request();
|
||||
request.max_age = Some(Duration::from_secs(10));
|
||||
cache
|
||||
.async_store_entries(vec![(
|
||||
request.clone(),
|
||||
json!({"answer": 7}),
|
||||
Duration::from_secs(100),
|
||||
)])
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
cache.lookup(&request, Duration::from_secs(110)).unwrap(),
|
||||
Some(json!({"answer": 7}))
|
||||
);
|
||||
assert_eq!(
|
||||
cache.lookup(&request, Duration::from_secs(111)).unwrap(),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn write_buffer_flushes_at_its_size_and_keeps_each_produced_time() {
|
||||
let cache = ResponseCache::new(Arc::new(InMemoryCache::default()));
|
||||
let buffer = WriteBuffer::new(2);
|
||||
let mut first = request();
|
||||
first.max_age = Some(Duration::from_secs(10));
|
||||
let mut second = request();
|
||||
second.key.preset = Some("tenant:other".into());
|
||||
|
||||
buffer
|
||||
.async_store(
|
||||
&cache,
|
||||
&first,
|
||||
json!({"answer": 7}),
|
||||
Duration::from_secs(100),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
cache.lookup(&first, Duration::from_secs(100)).unwrap(),
|
||||
None
|
||||
);
|
||||
|
||||
buffer
|
||||
.async_store(
|
||||
&cache,
|
||||
&second,
|
||||
json!({"answer": 8}),
|
||||
Duration::from_secs(200),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
cache.lookup(&first, Duration::from_secs(110)).unwrap(),
|
||||
Some(json!({"answer": 7}))
|
||||
);
|
||||
assert_eq!(
|
||||
cache.lookup(&first, Duration::from_secs(111)).unwrap(),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
cache.lookup(&second, Duration::from_secs(200)).unwrap(),
|
||||
Some(json!({"answer": 8}))
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn write_buffer_clear_drops_pending_entries() {
|
||||
let cache = ResponseCache::new(Arc::new(InMemoryCache::default()));
|
||||
let buffer = WriteBuffer::new(2);
|
||||
let mut other = request();
|
||||
other.key.preset = Some("tenant:other".into());
|
||||
let now = Duration::from_secs(100);
|
||||
|
||||
buffer
|
||||
.async_store(&cache, &request(), json!({"answer": 7}), now)
|
||||
.await
|
||||
.unwrap();
|
||||
buffer.clear().unwrap();
|
||||
buffer
|
||||
.async_store(&cache, &other, json!({"answer": 8}), now)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(cache.lookup(&request(), now).unwrap(), None);
|
||||
assert_eq!(cache.lookup(&other, now).unwrap(), None);
|
||||
}
|
||||
20
litellm-rust/crates/cache-s3/Cargo.toml
Normal file
20
litellm-rust/crates/cache-s3/Cargo.toml
Normal file
|
|
@ -0,0 +1,20 @@
|
|||
[package]
|
||||
name = "litellm-cache-s3"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
litellm-cache.workspace = true
|
||||
litellm-auth-aws.workspace = true
|
||||
aws-sdk-s3 = { version = "1.146.1", default-features = false, features = ["rustls", "rt-tokio"] }
|
||||
aws-credential-types = "1.3.0"
|
||||
aws-smithy-types = "1.6.0"
|
||||
aws-types = "1.6.0"
|
||||
tokio.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
wiremock = "0.6.5"
|
||||
serde_json.workspace = true
|
||||
tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }
|
||||
101
litellm-rust/crates/cache-s3/src/auth.rs
Normal file
101
litellm-rust/crates/cache-s3/src/auth.rs
Normal file
|
|
@ -0,0 +1,101 @@
|
|||
use aws_credential_types::{
|
||||
Credentials as AwsCredentials,
|
||||
provider::{ProvideCredentials, error::CredentialsError, future},
|
||||
};
|
||||
use litellm_auth_aws::{AwsAuthConfig, resolve_credentials};
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct Credentials {
|
||||
config: AwsAuthConfig,
|
||||
env: fn(&str) -> Option<String>,
|
||||
}
|
||||
|
||||
impl Credentials {
|
||||
pub(crate) fn new(config: AwsAuthConfig) -> Self {
|
||||
Self::with_env(config, |name| std::env::var(name).ok())
|
||||
}
|
||||
|
||||
pub(crate) fn with_env(config: AwsAuthConfig, env: fn(&str) -> Option<String>) -> Self {
|
||||
Self { config, env }
|
||||
}
|
||||
}
|
||||
|
||||
impl ProvideCredentials for Credentials {
|
||||
fn provide_credentials<'a>(&'a self) -> future::ProvideCredentials<'a>
|
||||
where
|
||||
Self: 'a,
|
||||
{
|
||||
future::ProvideCredentials::new(async {
|
||||
if let (Some(access_key_id), Some(secret_access_key)) = (
|
||||
self.config.access_key_id.clone(),
|
||||
self.config.secret_access_key.clone(),
|
||||
) {
|
||||
return Ok(AwsCredentials::new(
|
||||
access_key_id,
|
||||
secret_access_key,
|
||||
self.config.session_token.clone(),
|
||||
None,
|
||||
"litellm-s3-cache",
|
||||
));
|
||||
}
|
||||
resolve_credentials(self.config.clone(), &self.env)
|
||||
.await
|
||||
.map_err(|_| CredentialsError::provider_error("S3 cache authentication failed"))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for Credentials {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("Credentials").finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn explicit_keys_ignore_an_ambient_session_token() {
|
||||
let provider = Credentials::with_env(
|
||||
AwsAuthConfig {
|
||||
access_key_id: Some("key".to_string()),
|
||||
secret_access_key: Some("secret".to_string()),
|
||||
region_name: Some("us-east-1".to_string()),
|
||||
..Default::default()
|
||||
},
|
||||
|name| (name == "AWS_SESSION_TOKEN").then(|| "ambient".to_string()),
|
||||
);
|
||||
let credentials = provider.provide_credentials().await.unwrap();
|
||||
assert_eq!(credentials.access_key_id(), "key");
|
||||
assert_eq!(credentials.secret_access_key(), "secret");
|
||||
assert_eq!(credentials.session_token(), None);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn explicit_keys_keep_their_session_token() {
|
||||
let provider = Credentials::new(AwsAuthConfig {
|
||||
access_key_id: Some("key".to_string()),
|
||||
secret_access_key: Some("secret".to_string()),
|
||||
session_token: Some("t".to_string()),
|
||||
region_name: Some("us-east-1".to_string()),
|
||||
..Default::default()
|
||||
});
|
||||
let credentials = provider.provide_credentials().await.unwrap();
|
||||
assert_eq!(credentials.session_token(), Some("t"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn environment_keys_resolve_with_their_session_token() {
|
||||
let provider = Credentials::with_env(AwsAuthConfig::default(), |name| match name {
|
||||
"AWS_ACCESS_KEY_ID" => Some("env-key".to_string()),
|
||||
"AWS_SECRET_ACCESS_KEY" => Some("env-secret".to_string()),
|
||||
"AWS_SESSION_TOKEN" => Some("env-token".to_string()),
|
||||
_ => None,
|
||||
});
|
||||
let credentials = provider.provide_credentials().await.unwrap();
|
||||
assert_eq!(credentials.access_key_id(), "env-key");
|
||||
assert_eq!(credentials.secret_access_key(), "env-secret");
|
||||
assert_eq!(credentials.session_token(), Some("env-token"));
|
||||
}
|
||||
}
|
||||
220
litellm-rust/crates/cache-s3/src/cache.rs
Normal file
220
litellm-rust/crates/cache-s3/src/cache.rs
Normal file
|
|
@ -0,0 +1,220 @@
|
|||
use std::{
|
||||
future::Future,
|
||||
sync::Arc,
|
||||
time::{Duration, SystemTime},
|
||||
};
|
||||
|
||||
use aws_sdk_s3::{
|
||||
config::{BehaviorVersion, Region, RequestChecksumCalculation, ResponseChecksumValidation},
|
||||
error::SdkError,
|
||||
primitives::ByteStream,
|
||||
};
|
||||
use aws_smithy_types::{DateTime, date_time::Format};
|
||||
use litellm_auth_aws::AwsAuthConfig;
|
||||
use litellm_cache::{
|
||||
BaseCache, BatchCache, CacheCodec, CacheConnectionResult, Error, ExactCacheContext, FlushCache,
|
||||
};
|
||||
use tokio::runtime::Handle;
|
||||
|
||||
use crate::auth::Credentials;
|
||||
|
||||
pub struct S3Endpoint {
|
||||
pub url: String,
|
||||
}
|
||||
|
||||
pub struct S3CacheConfig {
|
||||
pub bucket: String,
|
||||
pub key_prefix: String,
|
||||
pub region: String,
|
||||
pub endpoint: Option<S3Endpoint>,
|
||||
pub auth: AwsAuthConfig,
|
||||
}
|
||||
|
||||
pub struct S3Cache<C: CacheCodec> {
|
||||
client: aws_sdk_s3::Client,
|
||||
codec: C,
|
||||
runtime: Handle,
|
||||
bucket: Arc<str>,
|
||||
key_prefix: Arc<str>,
|
||||
region: Arc<str>,
|
||||
endpoint: Option<Arc<str>>,
|
||||
}
|
||||
|
||||
impl<C: CacheCodec> S3Cache<C> {
|
||||
pub fn new(config: S3CacheConfig, codec: C, runtime: Handle) -> Self {
|
||||
let endpoint_url: Option<String> = config.endpoint.map(|endpoint| endpoint.url);
|
||||
let base = aws_sdk_s3::Config::builder()
|
||||
.behavior_version(BehaviorVersion::latest())
|
||||
.region(Region::new(config.region.clone()))
|
||||
.credentials_provider(Credentials::new(config.auth))
|
||||
.request_checksum_calculation(RequestChecksumCalculation::WhenRequired)
|
||||
.response_checksum_validation(ResponseChecksumValidation::WhenRequired);
|
||||
let builder = match &endpoint_url {
|
||||
Some(url) => base.endpoint_url(url).force_path_style(true),
|
||||
None => base,
|
||||
};
|
||||
Self {
|
||||
client: aws_sdk_s3::Client::from_conf(builder.build()),
|
||||
codec,
|
||||
runtime,
|
||||
bucket: config.bucket.into(),
|
||||
key_prefix: config.key_prefix.into(),
|
||||
region: config.region.into(),
|
||||
endpoint: endpoint_url.map(Into::into),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn bucket(&self) -> &str {
|
||||
&self.bucket
|
||||
}
|
||||
|
||||
pub fn key_prefix(&self) -> &str {
|
||||
&self.key_prefix
|
||||
}
|
||||
|
||||
pub fn region(&self) -> &str {
|
||||
&self.region
|
||||
}
|
||||
|
||||
pub fn endpoint(&self) -> Option<&str> {
|
||||
self.endpoint.as_deref()
|
||||
}
|
||||
|
||||
pub fn to_s3_key(&self, key: &str) -> String {
|
||||
format!("{}{}", self.key_prefix, key.replace(':', "/"))
|
||||
}
|
||||
|
||||
fn block_on<F: Future>(&self, future: F) -> F::Output {
|
||||
if Handle::try_current().is_ok() {
|
||||
tokio::task::block_in_place(|| self.runtime.block_on(future))
|
||||
} else {
|
||||
self.runtime.block_on(future)
|
||||
}
|
||||
}
|
||||
|
||||
async fn put(
|
||||
&self,
|
||||
key: &str,
|
||||
value: C::Value,
|
||||
context: &ExactCacheContext,
|
||||
) -> Result<(), Error> {
|
||||
let s3_key = self.to_s3_key(key);
|
||||
let body = self.codec.encode(&value)?;
|
||||
let request = self
|
||||
.client
|
||||
.put_object()
|
||||
.bucket(self.bucket.as_ref())
|
||||
.key(&s3_key)
|
||||
.body(ByteStream::from(body))
|
||||
.content_type("application/json")
|
||||
.content_language("en")
|
||||
.content_disposition(format!("inline; filename=\"{s3_key}.json\""));
|
||||
let request = match context.ttl {
|
||||
Some(ttl) => {
|
||||
let seconds = ttl.as_secs_f64();
|
||||
request
|
||||
.cache_control(format!("immutable, max-age={seconds}, s-maxage={seconds}"))
|
||||
.expires(DateTime::from(SystemTime::now() + ttl))
|
||||
}
|
||||
None => request.cache_control("immutable, max-age=31536000, s-maxage=31536000"),
|
||||
};
|
||||
request.send().await.map_err(|_| Error::Unavailable)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn get(&self, key: &str) -> Result<Option<C::Value>, Error> {
|
||||
let output = match self
|
||||
.client
|
||||
.get_object()
|
||||
.bucket(self.bucket.as_ref())
|
||||
.key(self.to_s3_key(key))
|
||||
.send()
|
||||
.await
|
||||
{
|
||||
Ok(output) => output,
|
||||
Err(error) => {
|
||||
if let SdkError::ServiceError(service) = &error {
|
||||
let status = error
|
||||
.raw_response()
|
||||
.map(|response| response.status().as_u16());
|
||||
let not_found = service.err().is_no_such_key()
|
||||
|| service.err().meta().code() == Some("AccessDenied")
|
||||
|| status == Some(404)
|
||||
|| status == Some(403);
|
||||
if not_found {
|
||||
return Ok(None);
|
||||
}
|
||||
}
|
||||
return Err(Error::Unavailable);
|
||||
}
|
||||
};
|
||||
if let Some(expires) = output.expires_string()
|
||||
&& let Ok(expires) = DateTime::from_str(expires, Format::HttpDate)
|
||||
&& expires < DateTime::from(SystemTime::now())
|
||||
{
|
||||
return Ok(None);
|
||||
}
|
||||
let bytes = output
|
||||
.body
|
||||
.collect()
|
||||
.await
|
||||
.map_err(|_| Error::Unavailable)?
|
||||
.into_bytes();
|
||||
self.codec.decode(&bytes).map(Some)
|
||||
}
|
||||
}
|
||||
|
||||
impl<C: CacheCodec> BaseCache for S3Cache<C> {
|
||||
type Value = C::Value;
|
||||
type Context = ExactCacheContext;
|
||||
|
||||
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
|
||||
context.ttl
|
||||
}
|
||||
|
||||
fn set_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
value: Self::Value,
|
||||
context: &Self::Context,
|
||||
) -> Result<(), Error> {
|
||||
self.block_on(self.put(key, value, context))
|
||||
}
|
||||
|
||||
fn get_cache(&self, key: &str, _context: &Self::Context) -> Result<Option<Self::Value>, Error> {
|
||||
self.block_on(self.get(key))
|
||||
}
|
||||
|
||||
async fn async_set_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
value: Self::Value,
|
||||
context: Self::Context,
|
||||
) -> Result<(), Error> {
|
||||
self.put(key, value, &context).await
|
||||
}
|
||||
|
||||
async fn async_get_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
_context: &Self::Context,
|
||||
) -> Result<Option<Self::Value>, Error> {
|
||||
self.get(key).await
|
||||
}
|
||||
|
||||
async fn disconnect(&self) -> Result<(), Error> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
|
||||
Err(Error::UnsupportedOperation)
|
||||
}
|
||||
}
|
||||
|
||||
impl<C: CacheCodec> BatchCache for S3Cache<C> {}
|
||||
|
||||
impl<C: CacheCodec> FlushCache for S3Cache<C> {
|
||||
fn flush_cache(&self) -> Result<(), Error> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
4
litellm-rust/crates/cache-s3/src/lib.rs
Normal file
4
litellm-rust/crates/cache-s3/src/lib.rs
Normal file
|
|
@ -0,0 +1,4 @@
|
|||
mod auth;
|
||||
mod cache;
|
||||
|
||||
pub use cache::{S3Cache, S3CacheConfig, S3Endpoint};
|
||||
278
litellm-rust/crates/cache-s3/tests/cache.rs
Normal file
278
litellm-rust/crates/cache-s3/tests/cache.rs
Normal file
|
|
@ -0,0 +1,278 @@
|
|||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||
|
||||
use litellm_auth_aws::AwsAuthConfig;
|
||||
use litellm_cache::{
|
||||
BaseCache, BatchCache, BatchEntry, Error, ExactCacheContext, FlushCache, JsonCodec,
|
||||
};
|
||||
use litellm_cache_s3::{S3Cache, S3CacheConfig, S3Endpoint};
|
||||
use serde_json::{Value, json};
|
||||
use tokio::runtime::Handle;
|
||||
use wiremock::{
|
||||
Mock, MockServer, ResponseTemplate,
|
||||
matchers::{method, path},
|
||||
};
|
||||
|
||||
fn config(endpoint: String) -> S3CacheConfig {
|
||||
S3CacheConfig {
|
||||
bucket: "cache-bucket".to_string(),
|
||||
key_prefix: "team/".to_string(),
|
||||
region: "us-east-1".to_string(),
|
||||
endpoint: Some(S3Endpoint { url: endpoint }),
|
||||
auth: AwsAuthConfig {
|
||||
access_key_id: Some("key".to_string()),
|
||||
secret_access_key: Some("secret".to_string()),
|
||||
region_name: Some("us-east-1".to_string()),
|
||||
..Default::default()
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn cache(endpoint: &str) -> S3Cache<JsonCodec<Value>> {
|
||||
S3Cache::new(
|
||||
config(endpoint.to_string()),
|
||||
JsonCodec::<Value>::new(),
|
||||
Handle::current(),
|
||||
)
|
||||
}
|
||||
|
||||
async fn mock_server() -> MockServer {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("PUT"))
|
||||
.respond_with(ResponseTemplate::new(200).insert_header("etag", "\"etag\""))
|
||||
.mount(&server)
|
||||
.await;
|
||||
server
|
||||
}
|
||||
|
||||
fn http_date_from(headers: &wiremock::http::HeaderMap, name: &str) -> Option<SystemTime> {
|
||||
use aws_smithy_types::{DateTime, date_time::Format};
|
||||
headers
|
||||
.get(name)
|
||||
.and_then(|value| DateTime::from_str(value.to_str().ok()?, Format::HttpDate).ok())
|
||||
.map(|date| UNIX_EPOCH + Duration::new(date.secs() as u64, date.subsec_nanos()))
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn set_writes_python_metadata_with_and_without_ttl() {
|
||||
let server = mock_server().await;
|
||||
let cache = cache(&server.uri());
|
||||
let context = ExactCacheContext {
|
||||
ttl: Some(Duration::from_secs(90)),
|
||||
};
|
||||
cache
|
||||
.set_cache("alpha:beta", json!({"answer": 1}), &context)
|
||||
.unwrap();
|
||||
cache
|
||||
.set_cache("plain", json!({"answer": 2}), &ExactCacheContext::default())
|
||||
.unwrap();
|
||||
|
||||
let requests = server.received_requests().await.unwrap();
|
||||
let ttl_request = requests
|
||||
.iter()
|
||||
.find(|request| request.url.path() == "/cache-bucket/team/alpha/beta")
|
||||
.expect("ttl write should hit the converted S3 key");
|
||||
assert_eq!(
|
||||
ttl_request.headers["cache-control"].to_str().unwrap(),
|
||||
"immutable, max-age=90, s-maxage=90"
|
||||
);
|
||||
assert_eq!(
|
||||
ttl_request.headers["content-type"].to_str().unwrap(),
|
||||
"application/json"
|
||||
);
|
||||
assert_eq!(
|
||||
ttl_request.headers["content-language"].to_str().unwrap(),
|
||||
"en"
|
||||
);
|
||||
assert_eq!(
|
||||
ttl_request.headers["content-disposition"].to_str().unwrap(),
|
||||
"inline; filename=\"team/alpha/beta.json\""
|
||||
);
|
||||
let expires = http_date_from(&ttl_request.headers, "expires").expect("ttl write sets Expires");
|
||||
let remaining = expires.duration_since(SystemTime::now()).unwrap();
|
||||
assert!(remaining > Duration::from_secs(60) && remaining <= Duration::from_secs(91));
|
||||
assert_eq!(
|
||||
serde_json::from_slice::<Value>(&ttl_request.body).unwrap(),
|
||||
json!({"answer": 1})
|
||||
);
|
||||
|
||||
let plain = requests
|
||||
.iter()
|
||||
.find(|request| request.url.path() == "/cache-bucket/team/plain")
|
||||
.expect("no-ttl write should hit the converted S3 key");
|
||||
assert_eq!(
|
||||
plain.headers["cache-control"].to_str().unwrap(),
|
||||
"immutable, max-age=31536000, s-maxage=31536000"
|
||||
);
|
||||
assert!(plain.headers.get("expires").is_none());
|
||||
assert_eq!(
|
||||
plain.headers["content-disposition"].to_str().unwrap(),
|
||||
"inline; filename=\"team/plain.json\""
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn get_hit_miss_expired_and_invalid_entries() {
|
||||
let server = mock_server().await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/cache-bucket/team/hit"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"answer": 3})))
|
||||
.mount(&server)
|
||||
.await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/cache-bucket/team/missing"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(404).set_body_string("<Error><Code>NoSuchKey</Code></Error>"),
|
||||
)
|
||||
.mount(&server)
|
||||
.await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/cache-bucket/team/denied"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(403).set_body_string("<Error><Code>AccessDenied</Code></Error>"),
|
||||
)
|
||||
.mount(&server)
|
||||
.await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/cache-bucket/team/expired"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200)
|
||||
.insert_header("expires", "Thu, 01 Jan 1970 00:00:00 GMT")
|
||||
.set_body_json(json!({"answer": 4})),
|
||||
)
|
||||
.mount(&server)
|
||||
.await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/cache-bucket/team/malformed"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_string("not a cache entry"))
|
||||
.mount(&server)
|
||||
.await;
|
||||
let cache = cache(&server.uri());
|
||||
let context = ExactCacheContext::default();
|
||||
|
||||
assert_eq!(
|
||||
cache.get_cache("hit", &context).unwrap(),
|
||||
Some(json!({"answer": 3}))
|
||||
);
|
||||
assert_eq!(cache.get_cache("missing", &context).unwrap(), None);
|
||||
assert_eq!(cache.get_cache("denied", &context).unwrap(), None);
|
||||
assert_eq!(cache.get_cache("expired", &context).unwrap(), None);
|
||||
assert_eq!(
|
||||
cache.get_cache("malformed", &context),
|
||||
Err(Error::InvalidEntry)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn batch_get_preserves_order_with_hits_misses_and_invalid() {
|
||||
let server = mock_server().await;
|
||||
for (key, status, body) in [
|
||||
("first", 200, "{\"answer\": 1}"),
|
||||
("invalid", 200, "garbage"),
|
||||
] {
|
||||
Mock::given(method("GET"))
|
||||
.and(path(format!("/cache-bucket/team/{key}")))
|
||||
.respond_with(ResponseTemplate::new(status).set_body_string(body))
|
||||
.mount(&server)
|
||||
.await;
|
||||
}
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/cache-bucket/team/miss"))
|
||||
.respond_with(ResponseTemplate::new(404))
|
||||
.mount(&server)
|
||||
.await;
|
||||
let cache = cache(&server.uri());
|
||||
let context = ExactCacheContext::default();
|
||||
let keys = vec![
|
||||
"first".to_string(),
|
||||
"miss".to_string(),
|
||||
"invalid".to_string(),
|
||||
];
|
||||
|
||||
let entries = cache.batch_get_cache(&keys, &context).unwrap();
|
||||
|
||||
assert_eq!(
|
||||
entries,
|
||||
vec![
|
||||
BatchEntry::Hit(json!({"answer": 1})),
|
||||
BatchEntry::Miss,
|
||||
BatchEntry::Invalid,
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn unsupported_and_noop_capabilities_match_python() {
|
||||
let server = mock_server().await;
|
||||
let cache = cache(&server.uri());
|
||||
|
||||
assert_eq!(
|
||||
cache.test_connection().await,
|
||||
Err(Error::UnsupportedOperation)
|
||||
);
|
||||
cache.flush_cache().unwrap();
|
||||
cache.disconnect().await.unwrap();
|
||||
assert_eq!(cache.get_ttl(&ExactCacheContext::default()), None);
|
||||
assert_eq!(
|
||||
cache.get_ttl(&ExactCacheContext {
|
||||
ttl: Some(Duration::from_secs(45)),
|
||||
}),
|
||||
Some(Duration::from_secs(45))
|
||||
);
|
||||
assert!(server.received_requests().await.unwrap().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn key_conversion_prefixes_and_splits_colons() {
|
||||
let runtime = tokio::runtime::Builder::new_multi_thread()
|
||||
.worker_threads(1)
|
||||
.enable_all()
|
||||
.build()
|
||||
.unwrap();
|
||||
let _guard = runtime.enter();
|
||||
let cache = S3Cache::new(
|
||||
S3CacheConfig {
|
||||
key_prefix: "team/".to_string(),
|
||||
..config("http://localhost".to_string())
|
||||
},
|
||||
JsonCodec::<Value>::new(),
|
||||
runtime.handle().clone(),
|
||||
);
|
||||
|
||||
assert_eq!(cache.bucket(), "cache-bucket");
|
||||
assert_eq!(cache.key_prefix(), "team/");
|
||||
assert_eq!(cache.to_s3_key("a:b:c"), "team/a/b/c");
|
||||
assert_eq!(cache.to_s3_key("plain"), "team/plain");
|
||||
|
||||
let unprefixed = S3Cache::new(
|
||||
S3CacheConfig {
|
||||
key_prefix: String::new(),
|
||||
..config("http://localhost".to_string())
|
||||
},
|
||||
JsonCodec::<Value>::new(),
|
||||
runtime.handle().clone(),
|
||||
);
|
||||
assert_eq!(unprefixed.to_s3_key("a:b"), "a/b");
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn sync_methods_block_inside_and_outside_the_runtime() {
|
||||
let server = mock_server().await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/cache-bucket/team/key"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"answer": 9})))
|
||||
.mount(&server)
|
||||
.await;
|
||||
let uri = server.uri();
|
||||
let cache = tokio::task::spawn_blocking(move || {
|
||||
let cache = cache(&uri);
|
||||
let context = ExactCacheContext::default();
|
||||
cache
|
||||
.set_cache("key", json!({"answer": 9}), &context)
|
||||
.unwrap();
|
||||
cache.get_cache("key", &context).unwrap()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(cache, Some(json!({"answer": 9})));
|
||||
}
|
||||
20
litellm-rust/crates/cache-valkey-semantic/Cargo.toml
Normal file
20
litellm-rust/crates/cache-valkey-semantic/Cargo.toml
Normal file
|
|
@ -0,0 +1,20 @@
|
|||
[package]
|
||||
name = "litellm-cache-valkey-semantic"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
litellm-cache.workspace = true
|
||||
litellm-cache-redis.workspace = true
|
||||
litellm-cache-response.workspace = true
|
||||
redis = { version = "1.7.0", features = ["tls-rustls"] }
|
||||
serde_json.workspace = true
|
||||
sha2.workspace = true
|
||||
tokio.workspace = true
|
||||
uuid = { version = "1", features = ["v4"] }
|
||||
|
||||
[dev-dependencies]
|
||||
redis-test = "1.0.4"
|
||||
rstest.workspace = true
|
||||
1153
litellm-rust/crates/cache-valkey-semantic/src/lib.rs
Normal file
1153
litellm-rust/crates/cache-valkey-semantic/src/lib.rs
Normal file
File diff suppressed because it is too large
Load diff
2
litellm-rust/crates/cache/Cargo.toml
vendored
2
litellm-rust/crates/cache/Cargo.toml
vendored
|
|
@ -8,8 +8,8 @@ repository.workspace = true
|
|||
[dependencies]
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
sha2.workspace = true
|
||||
thiserror.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
rstest.workspace = true
|
||||
tokio.workspace = true
|
||||
|
|
|
|||
169
litellm-rust/crates/cache/src/base_cache.rs
vendored
169
litellm-rust/crates/cache/src/base_cache.rs
vendored
|
|
@ -1,18 +1,57 @@
|
|||
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,
|
||||
}
|
||||
|
||||
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>,
|
||||
}
|
||||
|
||||
impl CacheContext for ExactCacheContext {
|
||||
fn ttl(&self) -> Option<Duration> {
|
||||
self.ttl
|
||||
}
|
||||
|
||||
fn with_ttl(&self, ttl: Option<Duration>) -> Self {
|
||||
Self { ttl }
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default, PartialEq)]
|
||||
pub struct CacheKwargs {
|
||||
pub struct SemanticCacheContext {
|
||||
pub input: Option<serde_json::Value>,
|
||||
pub messages: Option<serde_json::Value>,
|
||||
pub metadata: Option<serde_json::Value>,
|
||||
pub scope: Option<String>,
|
||||
pub ttl: Option<Duration>,
|
||||
pub extras: Map<String, Value>,
|
||||
}
|
||||
|
||||
impl CacheContext for SemanticCacheContext {
|
||||
fn ttl(&self) -> Option<Duration> {
|
||||
self.ttl
|
||||
}
|
||||
|
||||
fn with_ttl(&self, ttl: Option<Duration>) -> Self {
|
||||
Self {
|
||||
ttl,
|
||||
..self.clone()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Deserialize, Serialize, PartialEq, Eq)]
|
||||
|
|
@ -32,67 +71,87 @@ 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;
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::time::Duration;
|
||||
|
||||
use serde_json::json;
|
||||
|
||||
use super::{CacheContext, SemanticCacheContext};
|
||||
|
||||
#[test]
|
||||
fn semantic_context_with_ttl_only_replaces_ttl() {
|
||||
let context = SemanticCacheContext {
|
||||
input: Some(json!({"input": "hello"})),
|
||||
messages: Some(json!([{"role": "user", "content": "hello"}])),
|
||||
metadata: Some(json!({"tenant": "team"})),
|
||||
scope: Some("scope".into()),
|
||||
ttl: Some(Duration::from_secs(10)),
|
||||
};
|
||||
|
||||
let updated = context.with_ttl(Some(Duration::from_secs(20)));
|
||||
|
||||
assert_eq!(updated.ttl, Some(Duration::from_secs(20)));
|
||||
assert_eq!(updated.input, context.input);
|
||||
assert_eq!(updated.messages, context.messages);
|
||||
assert_eq!(updated.metadata, context.metadata);
|
||||
assert_eq!(updated.scope, context.scope);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
85
litellm-rust/crates/cache/src/cache_type.rs
vendored
Normal file
85
litellm-rust/crates/cache/src/cache_type.rs
vendored
Normal file
|
|
@ -0,0 +1,85 @@
|
|||
use serde::{Deserialize, Serialize};
|
||||
|
||||
#[derive(Clone, Copy, Debug, Deserialize, Serialize, PartialEq, Eq, Hash)]
|
||||
pub enum CacheType {
|
||||
#[serde(rename = "local")]
|
||||
Local,
|
||||
#[serde(rename = "redis")]
|
||||
Redis,
|
||||
#[serde(rename = "redis-semantic")]
|
||||
RedisSemantic,
|
||||
#[serde(rename = "valkey-semantic")]
|
||||
ValkeySemantic,
|
||||
#[serde(rename = "s3")]
|
||||
S3,
|
||||
#[serde(rename = "disk")]
|
||||
Disk,
|
||||
#[serde(rename = "qdrant-semantic")]
|
||||
QdrantSemantic,
|
||||
#[serde(rename = "azure-blob")]
|
||||
AzureBlob,
|
||||
#[serde(rename = "gcs")]
|
||||
Gcs,
|
||||
}
|
||||
|
||||
impl CacheType {
|
||||
pub const ALL: [Self; 9] = [
|
||||
Self::Local,
|
||||
Self::Redis,
|
||||
Self::RedisSemantic,
|
||||
Self::ValkeySemantic,
|
||||
Self::S3,
|
||||
Self::Disk,
|
||||
Self::QdrantSemantic,
|
||||
Self::AzureBlob,
|
||||
Self::Gcs,
|
||||
];
|
||||
|
||||
pub const fn as_python_name(self) -> &'static str {
|
||||
match self {
|
||||
Self::Local => "local",
|
||||
Self::Redis => "redis",
|
||||
Self::RedisSemantic => "redis-semantic",
|
||||
Self::ValkeySemantic => "valkey-semantic",
|
||||
Self::S3 => "s3",
|
||||
Self::Disk => "disk",
|
||||
Self::QdrantSemantic => "qdrant-semantic",
|
||||
Self::AzureBlob => "azure-blob",
|
||||
Self::Gcs => "gcs",
|
||||
}
|
||||
}
|
||||
|
||||
pub fn from_python_name(value: &str) -> Option<Self> {
|
||||
Self::ALL
|
||||
.into_iter()
|
||||
.find(|cache_type| cache_type.as_python_name() == value)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::CacheType;
|
||||
|
||||
#[test]
|
||||
fn every_python_cache_type_has_one_round_trip_identity() {
|
||||
let names = CacheType::ALL.map(CacheType::as_python_name);
|
||||
assert_eq!(
|
||||
names,
|
||||
[
|
||||
"local",
|
||||
"redis",
|
||||
"redis-semantic",
|
||||
"valkey-semantic",
|
||||
"s3",
|
||||
"disk",
|
||||
"qdrant-semantic",
|
||||
"azure-blob",
|
||||
"gcs",
|
||||
]
|
||||
);
|
||||
assert_eq!(
|
||||
names.map(CacheType::from_python_name),
|
||||
CacheType::ALL.map(Some)
|
||||
);
|
||||
}
|
||||
}
|
||||
167
litellm-rust/crates/cache/src/caching.rs
vendored
167
litellm-rust/crates/cache/src/caching.rs
vendored
|
|
@ -1,166 +1,23 @@
|
|||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
use crate::{BaseCache, CacheKwargs, Error};
|
||||
|
||||
pub use crate::BaseCache as Cache;
|
||||
use crate::{BaseCache, Error};
|
||||
|
||||
#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize, PartialEq, Eq)]
|
||||
pub enum CacheMode {
|
||||
#[default]
|
||||
#[serde(rename = "default_on")]
|
||||
DefaultOn,
|
||||
#[serde(rename = "default_off")]
|
||||
DefaultOff,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Deserialize, Serialize)]
|
||||
pub struct CacheKeyField {
|
||||
pub name: String,
|
||||
pub value: Option<String>,
|
||||
pub api_parameter: bool,
|
||||
pub internal_parameter: bool,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default, Deserialize, Serialize)]
|
||||
pub struct CacheKeyInput {
|
||||
pub fields: Vec<CacheKeyField>,
|
||||
pub preset: Option<String>,
|
||||
pub namespace: Option<String>,
|
||||
pub include_provider_parameters: bool,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
pub struct CacheKeyContext {
|
||||
pub model_group: Option<String>,
|
||||
pub caching_groups: Vec<(Vec<String>, String)>,
|
||||
pub file_checksum: Option<String>,
|
||||
pub file_object_name: Option<String>,
|
||||
pub metadata_file_name: Option<String>,
|
||||
pub parameters_file_name: Option<String>,
|
||||
}
|
||||
|
||||
impl CacheKeyContext {
|
||||
pub fn apply(self, input: &mut CacheKeyInput) {
|
||||
let group = self.model_group.as_ref().and_then(|model| {
|
||||
self.caching_groups
|
||||
.iter()
|
||||
.find(|(models, _)| models.contains(model))
|
||||
});
|
||||
for field in &mut input.fields {
|
||||
match field.name.as_str() {
|
||||
"model" => {
|
||||
field.value = group
|
||||
.map(|(_, formatted)| formatted.clone())
|
||||
.or_else(|| self.model_group.clone())
|
||||
.or_else(|| field.value.take())
|
||||
}
|
||||
"file" => {
|
||||
field.value = self
|
||||
.file_checksum
|
||||
.clone()
|
||||
.or_else(|| self.file_object_name.clone())
|
||||
.or_else(|| self.metadata_file_name.clone())
|
||||
.or_else(|| self.parameters_file_name.clone())
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn get_cache_key(input: &CacheKeyInput) -> String {
|
||||
cache_key(input)
|
||||
}
|
||||
|
||||
pub fn cache_key(input: &CacheKeyInput) -> String {
|
||||
if let Some(preset) = &input.preset {
|
||||
return preset.clone();
|
||||
}
|
||||
let mut digest = Sha256::new();
|
||||
for field in &input.fields {
|
||||
if (field.api_parameter || (input.include_provider_parameters && !field.internal_parameter))
|
||||
&& let Some(value) = &field.value
|
||||
{
|
||||
digest.update(field.name.as_bytes());
|
||||
digest.update(b": ");
|
||||
digest.update(value.as_bytes());
|
||||
}
|
||||
}
|
||||
let hash = format!("{:x}", digest.finalize());
|
||||
input
|
||||
.namespace
|
||||
.as_deref()
|
||||
.filter(|namespace| !namespace.is_empty())
|
||||
.map_or(hash.clone(), |namespace| format!("{namespace}:{hash}"))
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize)]
|
||||
pub struct CacheControls {
|
||||
pub supported_call_type: bool,
|
||||
pub configured: bool,
|
||||
pub native_backend: bool,
|
||||
pub default_on: bool,
|
||||
pub caching: Option<bool>,
|
||||
pub no_cache: bool,
|
||||
pub no_store: bool,
|
||||
#[serde(default)]
|
||||
pub use_cache: bool,
|
||||
}
|
||||
|
||||
impl CacheControls {
|
||||
pub fn reads(self) -> bool {
|
||||
self.supported_call_type
|
||||
&& self.configured
|
||||
&& self.caching.unwrap_or(true)
|
||||
&& !self.no_cache
|
||||
&& (self.default_on || self.use_cache)
|
||||
}
|
||||
|
||||
pub fn writes(self) -> bool {
|
||||
self.supported_call_type
|
||||
&& self.configured
|
||||
&& !self.no_store
|
||||
&& (self.default_on || self.use_cache)
|
||||
}
|
||||
}
|
||||
|
||||
pub fn should_use_cache(controls: CacheControls) -> bool {
|
||||
controls.reads() || controls.writes()
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct CacheEntry {
|
||||
pub timestamp: f64,
|
||||
pub response: Value,
|
||||
}
|
||||
|
||||
impl CacheEntry {
|
||||
pub fn fresh(&self, now: Duration, max_age: Option<Duration>) -> bool {
|
||||
self.timestamp.is_finite()
|
||||
&& max_age.is_none_or(|age| now.as_secs_f64() - self.timestamp <= age.as_secs_f64())
|
||||
}
|
||||
}
|
||||
|
||||
pub fn get_cache(
|
||||
cache: &dyn BaseCache<Value = CacheEntry>,
|
||||
pub fn get_cache<B: BaseCache>(
|
||||
cache: &B,
|
||||
key: &str,
|
||||
kwargs: &CacheKwargs,
|
||||
) -> Result<Option<CacheEntry>, Error> {
|
||||
cache.get_cache(key, kwargs)
|
||||
context: &B::Context,
|
||||
) -> Result<Option<B::Value>, Error> {
|
||||
cache.get_cache(key, context)
|
||||
}
|
||||
|
||||
pub fn set_cache(
|
||||
cache: &dyn BaseCache<Value = CacheEntry>,
|
||||
pub fn set_cache<B: BaseCache>(
|
||||
cache: &B,
|
||||
key: &str,
|
||||
entry: CacheEntry,
|
||||
kwargs: CacheKwargs,
|
||||
value: B::Value,
|
||||
context: &B::Context,
|
||||
) -> Result<(), Error> {
|
||||
cache.set_cache(key, entry, kwargs)
|
||||
cache.set_cache(key, value, context)
|
||||
}
|
||||
|
||||
pub type CacheBackend = Arc<dyn BaseCache<Value = CacheEntry>>;
|
||||
pub type CacheBackend<B> = Arc<B>;
|
||||
|
|
|
|||
169
litellm-rust/crates/cache/src/capabilities.rs
vendored
Normal file
169
litellm-rust/crates/cache/src/capabilities.rs
vendored
Normal file
|
|
@ -0,0 +1,169 @@
|
|||
use std::{future::Future, time::Duration};
|
||||
|
||||
use crate::{BaseCache, BatchEntry, Error};
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct IncrementOperation {
|
||||
pub key: String,
|
||||
pub amount: f64,
|
||||
pub ttl: Option<Duration>,
|
||||
}
|
||||
|
||||
pub trait BatchCache: BaseCache {
|
||||
fn batch_get_cache(
|
||||
&self,
|
||||
keys: &[String],
|
||||
context: &Self::Context,
|
||||
) -> Result<Vec<BatchEntry<Self::Value>>, Error> {
|
||||
keys.iter()
|
||||
.map(|key| match self.get_cache(key, context) {
|
||||
Ok(Some(value)) => Ok(BatchEntry::Hit(value)),
|
||||
Ok(None) => Ok(BatchEntry::Miss),
|
||||
Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid),
|
||||
Err(error) => Err(error),
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn async_batch_get_cache(
|
||||
&self,
|
||||
keys: Vec<String>,
|
||||
context: Self::Context,
|
||||
) -> impl Future<Output = Result<Vec<BatchEntry<Self::Value>>, Error>> + Send {
|
||||
async move {
|
||||
let mut entries = Vec::with_capacity(keys.len());
|
||||
for key in keys {
|
||||
entries.push(match self.async_get_cache(&key, &context).await {
|
||||
Ok(Some(value)) => BatchEntry::Hit(value),
|
||||
Ok(None) => BatchEntry::Miss,
|
||||
Err(Error::InvalidEntry) => BatchEntry::Invalid,
|
||||
Err(error) => return Err(error),
|
||||
});
|
||||
}
|
||||
Ok(entries)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub trait DeleteCache: BaseCache {
|
||||
fn delete_cache(&self, key: &str) -> Result<(), Error>;
|
||||
|
||||
fn async_delete_cache(&self, key: &str) -> impl Future<Output = Result<(), Error>> + Send {
|
||||
async move { self.delete_cache(key) }
|
||||
}
|
||||
}
|
||||
|
||||
pub trait FlushCache: BaseCache {
|
||||
fn flush_cache(&self) -> Result<(), Error>;
|
||||
|
||||
fn async_flush_cache(&self) -> impl Future<Output = Result<(), Error>> + Send {
|
||||
async move { self.flush_cache() }
|
||||
}
|
||||
}
|
||||
|
||||
pub trait CounterCache: BaseCache<Value = f64> {
|
||||
fn increment_cache(&self, key: &str, amount: f64, context: Self::Context)
|
||||
-> Result<f64, Error>;
|
||||
|
||||
fn async_increment(
|
||||
&self,
|
||||
key: &str,
|
||||
amount: f64,
|
||||
context: Self::Context,
|
||||
) -> impl Future<Output = Result<f64, Error>> + Send {
|
||||
async move { self.increment_cache(key, amount, context) }
|
||||
}
|
||||
}
|
||||
|
||||
pub trait ClaimCache: BaseCache
|
||||
where
|
||||
Self::Value: PartialEq,
|
||||
{
|
||||
fn claim_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
candidate: Self::Value,
|
||||
eligible: &[Self::Value],
|
||||
context: Self::Context,
|
||||
) -> Result<Self::Value, Error>;
|
||||
|
||||
fn async_claim_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
candidate: Self::Value,
|
||||
eligible: Vec<Self::Value>,
|
||||
context: Self::Context,
|
||||
) -> impl Future<Output = Result<Self::Value, Error>> + Send {
|
||||
async move { self.claim_cache(key, candidate, &eligible, context) }
|
||||
}
|
||||
}
|
||||
|
||||
pub trait TtlCache: BaseCache {
|
||||
fn async_get_ttl(
|
||||
&self,
|
||||
key: &str,
|
||||
) -> impl Future<Output = Result<Option<Duration>, Error>> + Send;
|
||||
}
|
||||
|
||||
pub trait SetCache: BaseCache {
|
||||
type SetValue: Clone + Send + Sync + 'static;
|
||||
type SetResult: Send + Sync + 'static;
|
||||
|
||||
fn async_set_cache_sadd(
|
||||
&self,
|
||||
key: &str,
|
||||
values: Vec<Self::SetValue>,
|
||||
ttl: Option<Duration>,
|
||||
) -> impl Future<Output = Result<Self::SetResult, Error>> + Send;
|
||||
}
|
||||
|
||||
pub trait QueueCache: BaseCache {
|
||||
type QueueValue: Clone + Send + Sync + 'static;
|
||||
type PopResult: Send + Sync + 'static;
|
||||
|
||||
fn async_rpush(
|
||||
&self,
|
||||
key: &str,
|
||||
values: Vec<Self::QueueValue>,
|
||||
) -> impl Future<Output = Result<usize, Error>> + Send;
|
||||
|
||||
fn async_lpop(
|
||||
&self,
|
||||
key: &str,
|
||||
count: Option<usize>,
|
||||
) -> impl Future<Output = Result<Self::PopResult, Error>> + Send;
|
||||
}
|
||||
|
||||
pub trait ScanCache: BaseCache {
|
||||
fn async_scan_iter(
|
||||
&self,
|
||||
pattern: &str,
|
||||
count: usize,
|
||||
) -> impl Future<Output = Result<Vec<String>, Error>> + Send;
|
||||
}
|
||||
|
||||
pub trait ClientInfoCache: BaseCache {
|
||||
type ClientList: Send + Sync + 'static;
|
||||
type Info: Send + Sync + 'static;
|
||||
|
||||
fn client_list(&self) -> Result<Self::ClientList, Error>;
|
||||
|
||||
fn info(&self) -> Result<Self::Info, Error>;
|
||||
}
|
||||
|
||||
pub trait CacheScript: Send + Sync + 'static {
|
||||
type Argument: Clone + Send + Sync + 'static;
|
||||
type Output: Send + Sync + 'static;
|
||||
|
||||
fn invoke(
|
||||
&self,
|
||||
keys: Vec<String>,
|
||||
arguments: Vec<Self::Argument>,
|
||||
) -> impl Future<Output = Result<Self::Output, Error>> + Send;
|
||||
}
|
||||
|
||||
pub trait ScriptCache: BaseCache {
|
||||
type Script: CacheScript;
|
||||
|
||||
fn async_register_script(&self, source: String) -> Self::Script;
|
||||
}
|
||||
50
litellm-rust/crates/cache/src/codec.rs
vendored
Normal file
50
litellm-rust/crates/cache/src/codec.rs
vendored
Normal file
|
|
@ -0,0 +1,50 @@
|
|||
use std::marker::PhantomData;
|
||||
|
||||
use serde::{Serialize, de::DeserializeOwned};
|
||||
|
||||
use crate::Error;
|
||||
|
||||
pub trait CacheCodec: Send + Sync {
|
||||
type Value: Clone + Send + Sync + 'static;
|
||||
|
||||
fn encode(&self, value: &Self::Value) -> Result<Vec<u8>, Error>;
|
||||
|
||||
fn decode(&self, bytes: &[u8]) -> Result<Self::Value, Error>;
|
||||
}
|
||||
|
||||
pub struct JsonCodec<V>(PhantomData<fn() -> V>);
|
||||
|
||||
impl<V> Clone for JsonCodec<V> {
|
||||
fn clone(&self) -> Self {
|
||||
*self
|
||||
}
|
||||
}
|
||||
|
||||
impl<V> Copy for JsonCodec<V> {}
|
||||
|
||||
impl<V> Default for JsonCodec<V> {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl<V> JsonCodec<V> {
|
||||
pub const fn new() -> Self {
|
||||
Self(PhantomData)
|
||||
}
|
||||
}
|
||||
|
||||
impl<V> CacheCodec for JsonCodec<V>
|
||||
where
|
||||
V: Clone + Send + Sync + Serialize + DeserializeOwned + 'static,
|
||||
{
|
||||
type Value = V;
|
||||
|
||||
fn encode(&self, value: &Self::Value) -> Result<Vec<u8>, Error> {
|
||||
serde_json::to_vec(value).map_err(|_| Error::InvalidEntry)
|
||||
}
|
||||
|
||||
fn decode(&self, bytes: &[u8]) -> Result<Self::Value, Error> {
|
||||
serde_json::from_slice(bytes).map_err(|_| Error::InvalidEntry)
|
||||
}
|
||||
}
|
||||
390
litellm-rust/crates/cache/src/dual.rs
vendored
Normal file
390
litellm-rust/crates/cache/src/dual.rs
vendored
Normal file
|
|
@ -0,0 +1,390 @@
|
|||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use crate::{
|
||||
BaseCache, BatchCache, BatchEntry, CacheConnectionResult, CacheContext, ClaimCache,
|
||||
CounterCache, DeleteCache, Error, FlushCache,
|
||||
};
|
||||
|
||||
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
|
||||
pub enum ReadPolicy {
|
||||
#[default]
|
||||
LocalThenRemote,
|
||||
LocalOnly,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
|
||||
pub enum WritePolicy {
|
||||
#[default]
|
||||
Both,
|
||||
LocalOnly,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
|
||||
pub enum RemoteFailurePolicy {
|
||||
#[default]
|
||||
Propagate,
|
||||
UseLocal,
|
||||
}
|
||||
|
||||
pub struct DualCache<L1, L2> {
|
||||
l1: Arc<L1>,
|
||||
l2: Arc<L2>,
|
||||
read_policy: ReadPolicy,
|
||||
write_policy: WritePolicy,
|
||||
remote_failure_policy: RemoteFailurePolicy,
|
||||
promotion_ttl: Option<Duration>,
|
||||
}
|
||||
|
||||
impl<L1, L2> DualCache<L1, L2> {
|
||||
pub fn new(l1: Arc<L1>, l2: Arc<L2>) -> Self {
|
||||
Self {
|
||||
l1,
|
||||
l2,
|
||||
read_policy: ReadPolicy::default(),
|
||||
write_policy: WritePolicy::default(),
|
||||
remote_failure_policy: RemoteFailurePolicy::default(),
|
||||
promotion_ttl: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_read_policy(self, read_policy: ReadPolicy) -> Self {
|
||||
Self {
|
||||
read_policy,
|
||||
..self
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_write_policy(self, write_policy: WritePolicy) -> Self {
|
||||
Self {
|
||||
write_policy,
|
||||
..self
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_remote_failure_policy(self, remote_failure_policy: RemoteFailurePolicy) -> Self {
|
||||
Self {
|
||||
remote_failure_policy,
|
||||
..self
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_promotion_ttl(self, promotion_ttl: Duration) -> Self {
|
||||
Self {
|
||||
promotion_ttl: Some(promotion_ttl),
|
||||
..self
|
||||
}
|
||||
}
|
||||
|
||||
fn reads_remote(&self) -> bool {
|
||||
self.read_policy == ReadPolicy::LocalThenRemote
|
||||
}
|
||||
|
||||
fn writes_remote(&self) -> bool {
|
||||
self.write_policy == WritePolicy::Both
|
||||
}
|
||||
|
||||
fn remote<T>(&self, result: Result<T, Error>) -> Result<Option<T>, Error> {
|
||||
match result {
|
||||
Ok(value) => Ok(Some(value)),
|
||||
Err(Error::Unavailable)
|
||||
if self.remote_failure_policy == RemoteFailurePolicy::UseLocal =>
|
||||
{
|
||||
Ok(None)
|
||||
}
|
||||
Err(error) => Err(error),
|
||||
}
|
||||
}
|
||||
|
||||
fn promotion_context<C: CacheContext>(&self, context: &C) -> C {
|
||||
context.with_ttl(self.promotion_ttl.or(context.ttl()))
|
||||
}
|
||||
}
|
||||
|
||||
impl<V, C, L1, L2> DualCache<L1, L2>
|
||||
where
|
||||
V: Clone + Send + Sync + 'static,
|
||||
C: CacheContext,
|
||||
L1: BaseCache<Value = V, Context = C>,
|
||||
L2: BaseCache<Value = V, Context = C>,
|
||||
{
|
||||
fn missing(entries: &[BatchEntry<V>]) -> Vec<usize> {
|
||||
entries
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter_map(|(index, entry)| (!matches!(entry, BatchEntry::Hit(_))).then_some(index))
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn merge_batch(
|
||||
&self,
|
||||
keys: &[String],
|
||||
context: &C,
|
||||
mut entries: Vec<BatchEntry<V>>,
|
||||
missing: Vec<usize>,
|
||||
remote: Vec<BatchEntry<V>>,
|
||||
) -> Result<Vec<BatchEntry<V>>, Error> {
|
||||
if missing.len() != remote.len() {
|
||||
return Err(Error::Unavailable);
|
||||
}
|
||||
for (index, entry) in missing.into_iter().zip(remote) {
|
||||
if let BatchEntry::Hit(value) = &entry {
|
||||
let promotion_context = self.promotion_context(context);
|
||||
self.l1
|
||||
.set_cache(&keys[index], value.clone(), &promotion_context)?;
|
||||
}
|
||||
entries[index] = entry;
|
||||
}
|
||||
Ok(entries)
|
||||
}
|
||||
}
|
||||
|
||||
impl<V, C, L1, L2> BaseCache for DualCache<L1, L2>
|
||||
where
|
||||
V: Clone + Send + Sync + 'static,
|
||||
C: CacheContext,
|
||||
L1: BaseCache<Value = V, Context = C>,
|
||||
L2: BaseCache<Value = V, Context = C>,
|
||||
{
|
||||
type Value = V;
|
||||
type Context = C;
|
||||
|
||||
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
|
||||
self.l2.get_ttl(context)
|
||||
}
|
||||
|
||||
fn set_cache(&self, key: &str, value: V, context: &C) -> Result<(), Error> {
|
||||
if self.writes_remote() {
|
||||
self.remote(self.l2.set_cache(key, value.clone(), context))?;
|
||||
}
|
||||
self.l1.set_cache(key, value, context)
|
||||
}
|
||||
|
||||
fn get_cache(&self, key: &str, context: &C) -> Result<Option<V>, Error> {
|
||||
if let Some(value) = self.l1.get_cache(key, context)? {
|
||||
return Ok(Some(value));
|
||||
}
|
||||
if !self.reads_remote() {
|
||||
return Ok(None);
|
||||
}
|
||||
let value = self.remote(self.l2.get_cache(key, context))?.flatten();
|
||||
if let Some(value) = &value {
|
||||
let promotion_context = self.promotion_context(context);
|
||||
self.l1.set_cache(key, value.clone(), &promotion_context)?;
|
||||
}
|
||||
Ok(value)
|
||||
}
|
||||
|
||||
async fn async_set_cache(&self, key: &str, value: V, context: C) -> Result<(), Error> {
|
||||
if self.writes_remote() {
|
||||
self.remote(
|
||||
self.l2
|
||||
.async_set_cache(key, value.clone(), context.clone())
|
||||
.await,
|
||||
)?;
|
||||
}
|
||||
self.l1.async_set_cache(key, value, context).await
|
||||
}
|
||||
|
||||
async fn async_get_cache(&self, key: &str, context: &C) -> Result<Option<V>, Error> {
|
||||
if let Some(value) = self.l1.async_get_cache(key, context).await? {
|
||||
return Ok(Some(value));
|
||||
}
|
||||
if !self.reads_remote() {
|
||||
return Ok(None);
|
||||
}
|
||||
let value = self
|
||||
.remote(self.l2.async_get_cache(key, context).await)?
|
||||
.flatten();
|
||||
if let Some(value) = &value {
|
||||
self.l1
|
||||
.async_set_cache(key, value.clone(), self.promotion_context(context))
|
||||
.await?;
|
||||
}
|
||||
Ok(value)
|
||||
}
|
||||
|
||||
async fn async_set_cache_pipeline(
|
||||
&self,
|
||||
entries: Vec<(String, V)>,
|
||||
context: C,
|
||||
) -> Result<(), Error> {
|
||||
if self.writes_remote() {
|
||||
self.remote(
|
||||
self.l2
|
||||
.async_set_cache_pipeline(entries.clone(), context.clone())
|
||||
.await,
|
||||
)?;
|
||||
}
|
||||
self.l1.async_set_cache_pipeline(entries, context).await
|
||||
}
|
||||
|
||||
async fn disconnect(&self) -> Result<(), Error> {
|
||||
self.l2.disconnect().await?;
|
||||
self.l1.disconnect().await
|
||||
}
|
||||
|
||||
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
|
||||
self.l2.test_connection().await
|
||||
}
|
||||
}
|
||||
|
||||
impl<V, C, L1, L2> BatchCache for DualCache<L1, L2>
|
||||
where
|
||||
V: Clone + Send + Sync + 'static,
|
||||
C: CacheContext,
|
||||
L1: BatchCache<Value = V, Context = C>,
|
||||
L2: BatchCache<Value = V, Context = C>,
|
||||
{
|
||||
fn batch_get_cache(&self, keys: &[String], context: &C) -> Result<Vec<BatchEntry<V>>, Error> {
|
||||
let entries = self.l1.batch_get_cache(keys, context)?;
|
||||
let missing = Self::missing(&entries);
|
||||
if missing.is_empty() || !self.reads_remote() {
|
||||
return Ok(entries);
|
||||
}
|
||||
let remote_keys = missing
|
||||
.iter()
|
||||
.map(|index| keys[*index].clone())
|
||||
.collect::<Vec<_>>();
|
||||
match self.remote(self.l2.batch_get_cache(&remote_keys, context))? {
|
||||
Some(remote) => self.merge_batch(keys, context, entries, missing, remote),
|
||||
None => Ok(entries),
|
||||
}
|
||||
}
|
||||
|
||||
async fn async_batch_get_cache(
|
||||
&self,
|
||||
keys: Vec<String>,
|
||||
context: C,
|
||||
) -> Result<Vec<BatchEntry<V>>, Error> {
|
||||
let entries = self
|
||||
.l1
|
||||
.async_batch_get_cache(keys.clone(), context.clone())
|
||||
.await?;
|
||||
let missing = Self::missing(&entries);
|
||||
if missing.is_empty() || !self.reads_remote() {
|
||||
return Ok(entries);
|
||||
}
|
||||
let remote_keys = missing.iter().map(|index| keys[*index].clone()).collect();
|
||||
match self.remote(
|
||||
self.l2
|
||||
.async_batch_get_cache(remote_keys, context.clone())
|
||||
.await,
|
||||
)? {
|
||||
Some(remote) => self.merge_batch(&keys, &context, entries, missing, remote),
|
||||
None => Ok(entries),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<V, C, L1, L2> DeleteCache for DualCache<L1, L2>
|
||||
where
|
||||
V: Clone + Send + Sync + 'static,
|
||||
C: CacheContext,
|
||||
L1: DeleteCache<Value = V, Context = C>,
|
||||
L2: DeleteCache<Value = V, Context = C>,
|
||||
{
|
||||
fn delete_cache(&self, key: &str) -> Result<(), Error> {
|
||||
if self.writes_remote() {
|
||||
self.remote(self.l2.delete_cache(key))?;
|
||||
}
|
||||
self.l1.delete_cache(key)
|
||||
}
|
||||
|
||||
async fn async_delete_cache(&self, key: &str) -> Result<(), Error> {
|
||||
if self.writes_remote() {
|
||||
self.remote(self.l2.async_delete_cache(key).await)?;
|
||||
}
|
||||
self.l1.async_delete_cache(key).await
|
||||
}
|
||||
}
|
||||
|
||||
impl<V, C, L1, L2> FlushCache for DualCache<L1, L2>
|
||||
where
|
||||
V: Clone + Send + Sync + 'static,
|
||||
C: CacheContext,
|
||||
L1: FlushCache<Value = V, Context = C>,
|
||||
L2: FlushCache<Value = V, Context = C>,
|
||||
{
|
||||
fn flush_cache(&self) -> Result<(), Error> {
|
||||
if self.writes_remote() {
|
||||
self.remote(self.l2.flush_cache())?;
|
||||
}
|
||||
self.l1.flush_cache()
|
||||
}
|
||||
|
||||
async fn async_flush_cache(&self) -> Result<(), Error> {
|
||||
if self.writes_remote() {
|
||||
self.remote(self.l2.async_flush_cache().await)?;
|
||||
}
|
||||
self.l1.async_flush_cache().await
|
||||
}
|
||||
}
|
||||
|
||||
impl<C, L1, L2> CounterCache for DualCache<L1, L2>
|
||||
where
|
||||
C: CacheContext,
|
||||
L1: BaseCache<Value = f64, Context = C>,
|
||||
L2: CounterCache<Context = C>,
|
||||
{
|
||||
fn increment_cache(&self, key: &str, amount: f64, context: C) -> Result<f64, Error> {
|
||||
let value = self.l2.increment_cache(key, amount, context.clone())?;
|
||||
self.l1.set_cache(key, value, &context)?;
|
||||
Ok(value)
|
||||
}
|
||||
|
||||
async fn async_increment(&self, key: &str, amount: f64, context: C) -> Result<f64, Error> {
|
||||
let value = self
|
||||
.l2
|
||||
.async_increment(key, amount, context.clone())
|
||||
.await?;
|
||||
self.l1.async_set_cache(key, value, context).await?;
|
||||
Ok(value)
|
||||
}
|
||||
}
|
||||
|
||||
impl<V, C, L1, L2> ClaimCache for DualCache<L1, L2>
|
||||
where
|
||||
V: Clone + PartialEq + Send + Sync + 'static,
|
||||
C: CacheContext,
|
||||
L1: ClaimCache<Value = V, Context = C>,
|
||||
L2: ClaimCache<Value = V, Context = C>,
|
||||
{
|
||||
fn claim_cache(&self, key: &str, candidate: V, eligible: &[V], context: C) -> Result<V, Error> {
|
||||
match self.remote(
|
||||
self.l2
|
||||
.claim_cache(key, candidate.clone(), eligible, context.clone()),
|
||||
)? {
|
||||
Some(winner) => {
|
||||
self.l1.set_cache(key, winner.clone(), &context)?;
|
||||
Ok(winner)
|
||||
}
|
||||
None => self.l1.claim_cache(key, candidate, eligible, context),
|
||||
}
|
||||
}
|
||||
|
||||
async fn async_claim_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
candidate: V,
|
||||
eligible: Vec<V>,
|
||||
context: C,
|
||||
) -> Result<V, Error> {
|
||||
match self.remote(
|
||||
self.l2
|
||||
.async_claim_cache(key, candidate.clone(), eligible.clone(), context.clone())
|
||||
.await,
|
||||
)? {
|
||||
Some(winner) => {
|
||||
self.l1
|
||||
.async_set_cache(key, winner.clone(), context)
|
||||
.await?;
|
||||
Ok(winner)
|
||||
}
|
||||
None => {
|
||||
self.l1
|
||||
.async_claim_cache(key, candidate, eligible, context)
|
||||
.await
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
4
litellm-rust/crates/cache/src/error.rs
vendored
4
litellm-rust/crates/cache/src/error.rs
vendored
|
|
@ -4,4 +4,8 @@ pub enum Error {
|
|||
Unavailable,
|
||||
#[error("invalid cache entry")]
|
||||
InvalidEntry,
|
||||
#[error("flushing Redis requires an explicit namespace")]
|
||||
UnscopedFlush,
|
||||
#[error("operation is not supported by this cache")]
|
||||
UnsupportedOperation,
|
||||
}
|
||||
|
|
|
|||
17
litellm-rust/crates/cache/src/lib.rs
vendored
17
litellm-rust/crates/cache/src/lib.rs
vendored
|
|
@ -1,12 +1,21 @@
|
|||
mod base_cache;
|
||||
mod cache_type;
|
||||
mod caching;
|
||||
mod capabilities;
|
||||
mod codec;
|
||||
mod dual;
|
||||
mod error;
|
||||
|
||||
pub use base_cache::{
|
||||
BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheFuture, CacheKwargs,
|
||||
BaseCache, BatchEntry, CacheConnectionResult, CacheConnectionStatus, CacheContext,
|
||||
ExactCacheContext, SemanticCacheContext,
|
||||
};
|
||||
pub use caching::{
|
||||
Cache, CacheBackend, CacheControls, CacheEntry, CacheKeyContext, CacheKeyField, CacheKeyInput,
|
||||
CacheMode, cache_key, get_cache, get_cache_key, set_cache, should_use_cache,
|
||||
pub use cache_type::CacheType;
|
||||
pub use caching::{Cache, CacheBackend, get_cache, set_cache};
|
||||
pub use capabilities::{
|
||||
BatchCache, CacheScript, ClaimCache, ClientInfoCache, CounterCache, DeleteCache, FlushCache,
|
||||
IncrementOperation, QueueCache, ScanCache, ScriptCache, SetCache, TtlCache,
|
||||
};
|
||||
pub use codec::{CacheCodec, JsonCodec};
|
||||
pub use dual::{DualCache, ReadPolicy, RemoteFailurePolicy, WritePolicy};
|
||||
pub use error::Error;
|
||||
|
|
|
|||
208
litellm-rust/crates/cache/tests/caching.rs
vendored
208
litellm-rust/crates/cache/tests/caching.rs
vendored
|
|
@ -1,42 +1,97 @@
|
|||
use std::{sync::Mutex, time::Duration};
|
||||
|
||||
use litellm_cache::{
|
||||
BaseCache, CacheConnectionResult, CacheControls, CacheEntry, CacheFuture, CacheKeyContext,
|
||||
CacheKeyField, CacheKeyInput, CacheKwargs, Error, cache_key, get_cache_key,
|
||||
BaseCache, CacheConnectionResult, CacheContext, Error, ExactCacheContext, get_cache,
|
||||
};
|
||||
use sha2::{Digest, Sha256};
|
||||
use std::time::Duration;
|
||||
|
||||
struct TestCache {
|
||||
default_ttl: Duration,
|
||||
writes: Mutex<Vec<(String, String, ExactCacheContext)>>,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct SemanticContext {
|
||||
ttl: Option<Duration>,
|
||||
query: String,
|
||||
}
|
||||
|
||||
impl CacheContext for SemanticContext {
|
||||
fn ttl(&self) -> Option<Duration> {
|
||||
self.ttl
|
||||
}
|
||||
|
||||
fn with_ttl(&self, ttl: Option<Duration>) -> Self {
|
||||
Self {
|
||||
ttl,
|
||||
query: self.query.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct SemanticCache;
|
||||
|
||||
impl BaseCache for SemanticCache {
|
||||
type Value = String;
|
||||
type Context = SemanticContext;
|
||||
|
||||
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
|
||||
context.ttl
|
||||
}
|
||||
|
||||
fn set_cache(&self, _: &str, _: Self::Value, _: &Self::Context) -> Result<(), Error> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn get_cache(&self, _: &str, context: &Self::Context) -> Result<Option<Self::Value>, Error> {
|
||||
Ok((context.query == "matching prompt").then(|| "semantic hit".into()))
|
||||
}
|
||||
|
||||
async fn disconnect(&self) -> Result<(), Error> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
|
||||
unreachable!()
|
||||
}
|
||||
}
|
||||
|
||||
impl BaseCache for TestCache {
|
||||
type Value = CacheEntry;
|
||||
type Value = String;
|
||||
type Context = ExactCacheContext;
|
||||
|
||||
fn default_ttl(&self) -> Duration {
|
||||
self.default_ttl
|
||||
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
|
||||
context.ttl.or(Some(self.default_ttl))
|
||||
}
|
||||
|
||||
fn set_cache(&self, _: &str, _: Self::Value, _: CacheKwargs) -> Result<(), Error> {
|
||||
fn set_cache(&self, _: &str, _: Self::Value, _: &ExactCacheContext) -> Result<(), Error> {
|
||||
Err(Error::Unavailable)
|
||||
}
|
||||
|
||||
async fn async_set_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
value: Self::Value,
|
||||
context: ExactCacheContext,
|
||||
) -> Result<(), Error> {
|
||||
if key == "unavailable" {
|
||||
return Err(Error::Unavailable);
|
||||
}
|
||||
self.writes
|
||||
.lock()
|
||||
.unwrap()
|
||||
.push((key.into(), value, context));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn get_cache(&self, _: &str, _: &CacheKwargs) -> Result<Option<Self::Value>, Error> {
|
||||
fn get_cache(&self, _: &str, _: &ExactCacheContext) -> Result<Option<Self::Value>, Error> {
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
fn delete_cache(&self, _: &str) -> Result<(), Error> {
|
||||
async fn disconnect(&self) -> Result<(), Error> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn flush_cache(&self) -> Result<(), Error> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn disconnect(&self) -> CacheFuture<'_, ()> {
|
||||
Box::pin(async { Ok(()) })
|
||||
}
|
||||
|
||||
fn test_connection(&self) -> CacheFuture<'_, CacheConnectionResult> {
|
||||
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
|
||||
unreachable!()
|
||||
}
|
||||
}
|
||||
|
|
@ -45,95 +100,64 @@ impl BaseCache for TestCache {
|
|||
fn ttl_uses_default_and_allows_per_call_override() {
|
||||
let cache = TestCache {
|
||||
default_ttl: Duration::from_secs(60),
|
||||
writes: Mutex::default(),
|
||||
};
|
||||
assert_eq!(
|
||||
cache.get_ttl(&CacheKwargs::default()),
|
||||
Duration::from_secs(60)
|
||||
cache.get_ttl(&ExactCacheContext::default()),
|
||||
Some(Duration::from_secs(60))
|
||||
);
|
||||
assert_eq!(
|
||||
cache.get_ttl(&CacheKwargs {
|
||||
cache.get_ttl(&ExactCacheContext {
|
||||
ttl: Some(Duration::from_secs(5)),
|
||||
..Default::default()
|
||||
}),
|
||||
Duration::from_secs(5)
|
||||
Some(Duration::from_secs(5))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn keys_match_python_order_groups_files_presets_and_namespaces() {
|
||||
let mut input = CacheKeyInput {
|
||||
fields: vec![
|
||||
CacheKeyField {
|
||||
name: "model".into(),
|
||||
value: Some("deployment".into()),
|
||||
api_parameter: true,
|
||||
internal_parameter: false,
|
||||
},
|
||||
CacheKeyField {
|
||||
name: "file".into(),
|
||||
value: None,
|
||||
api_parameter: true,
|
||||
internal_parameter: false,
|
||||
},
|
||||
],
|
||||
namespace: Some("team".into()),
|
||||
..Default::default()
|
||||
fn associated_context_preserves_backend_specific_lookup_inputs() {
|
||||
let context = SemanticContext {
|
||||
ttl: None,
|
||||
query: "matching prompt".into(),
|
||||
};
|
||||
CacheKeyContext {
|
||||
model_group: Some("group".into()),
|
||||
caching_groups: vec![(vec!["group".into()], "['group']".into())],
|
||||
file_checksum: Some("checksum".into()),
|
||||
..Default::default()
|
||||
}
|
||||
.apply(&mut input);
|
||||
assert_eq!(
|
||||
cache_key(&input),
|
||||
format!(
|
||||
"team:{:x}",
|
||||
Sha256::digest(b"model: ['group']file: checksum")
|
||||
)
|
||||
get_cache(&SemanticCache, "shared-key", &context).unwrap(),
|
||||
Some("semantic hit".into())
|
||||
);
|
||||
input.preset = Some("preset".into());
|
||||
assert_eq!(get_cache_key(&input), "preset");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_controls_honor_default_modes_and_directives() {
|
||||
let enabled = CacheControls {
|
||||
supported_call_type: true,
|
||||
configured: true,
|
||||
default_on: true,
|
||||
..Default::default()
|
||||
#[tokio::test]
|
||||
async fn default_batch_operations_use_async_writes_and_stop_on_failure() {
|
||||
let cache = TestCache {
|
||||
default_ttl: Duration::from_secs(60),
|
||||
writes: Mutex::default(),
|
||||
};
|
||||
assert!(enabled.reads());
|
||||
assert!(enabled.writes());
|
||||
assert!(
|
||||
!CacheControls {
|
||||
default_on: false,
|
||||
..enabled
|
||||
}
|
||||
.reads()
|
||||
let entry = String::from("cached");
|
||||
let context = ExactCacheContext {
|
||||
ttl: Some(Duration::from_secs(5)),
|
||||
};
|
||||
cache
|
||||
.batch_cache_write("single", entry.clone(), context.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_set_cache_pipeline(
|
||||
vec![
|
||||
("first".into(), entry.clone()),
|
||||
("unavailable".into(), entry.clone()),
|
||||
("skipped".into(), entry.clone()),
|
||||
],
|
||||
context.clone(),
|
||||
)
|
||||
.await,
|
||||
Err(Error::Unavailable)
|
||||
);
|
||||
assert!(
|
||||
CacheControls {
|
||||
default_on: false,
|
||||
use_cache: true,
|
||||
..enabled
|
||||
}
|
||||
.reads()
|
||||
);
|
||||
assert!(
|
||||
!CacheControls {
|
||||
no_cache: true,
|
||||
..enabled
|
||||
}
|
||||
.reads()
|
||||
);
|
||||
assert!(
|
||||
!CacheControls {
|
||||
no_store: true,
|
||||
..enabled
|
||||
}
|
||||
.writes()
|
||||
assert_eq!(
|
||||
*cache.writes.lock().unwrap(),
|
||||
vec![
|
||||
("single".into(), entry.clone(), context.clone()),
|
||||
("first".into(), entry, context),
|
||||
]
|
||||
);
|
||||
}
|
||||
|
|
|
|||
41
litellm-rust/crates/cache/tests/codec.rs
vendored
Normal file
41
litellm-rust/crates/cache/tests/codec.rs
vendored
Normal file
|
|
@ -0,0 +1,41 @@
|
|||
use std::collections::BTreeMap;
|
||||
|
||||
use litellm_cache::{CacheCodec, Error, JsonCodec};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::json;
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
struct RoutingState {
|
||||
deployment: String,
|
||||
cooldown_seconds: u64,
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn json_codec_round_trips_typed_domain_values() {
|
||||
let codec = JsonCodec::<RoutingState>::new();
|
||||
let value = RoutingState {
|
||||
deployment: "deployment-a".into(),
|
||||
cooldown_seconds: 30,
|
||||
};
|
||||
let bytes = codec.encode(&value).unwrap();
|
||||
assert_eq!(codec.decode(&bytes).unwrap(), value);
|
||||
assert_eq!(
|
||||
serde_json::from_slice::<serde_json::Value>(&bytes).unwrap(),
|
||||
json!({"deployment": "deployment-a", "cooldown_seconds": 30})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn json_codec_rejects_malformed_and_wrongly_typed_entries() {
|
||||
let codec = JsonCodec::<RoutingState>::new();
|
||||
for bytes in [b"not json".as_slice(), br#"{"deployment":12}"#.as_slice()] {
|
||||
assert_eq!(codec.decode(bytes).unwrap_err(), Error::InvalidEntry);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn json_codec_propagates_encoding_errors() {
|
||||
let codec = JsonCodec::<BTreeMap<(u8, u8), String>>::new();
|
||||
let value = BTreeMap::from([((1, 2), "invalid JSON object key".into())]);
|
||||
assert_eq!(codec.encode(&value).unwrap_err(), Error::InvalidEntry);
|
||||
}
|
||||
385
litellm-rust/crates/cache/tests/dual.rs
vendored
Normal file
385
litellm-rust/crates/cache/tests/dual.rs
vendored
Normal file
|
|
@ -0,0 +1,385 @@
|
|||
use std::{
|
||||
sync::{Arc, Mutex},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use litellm_cache::{
|
||||
BaseCache, BatchCache, CacheConnectionResult, ClaimCache, CounterCache, DeleteCache, DualCache,
|
||||
Error, ExactCacheContext, FlushCache, ReadPolicy, RemoteFailurePolicy, WritePolicy,
|
||||
};
|
||||
|
||||
struct TestCache<V> {
|
||||
value: Mutex<Option<V>>,
|
||||
fail: bool,
|
||||
}
|
||||
|
||||
impl<V> TestCache<V> {
|
||||
fn new(value: Option<V>, fail: bool) -> Self {
|
||||
Self {
|
||||
value: Mutex::new(value),
|
||||
fail,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<V> BaseCache for TestCache<V>
|
||||
where
|
||||
V: Clone + Send + Sync + 'static,
|
||||
{
|
||||
type Value = V;
|
||||
type Context = ExactCacheContext;
|
||||
|
||||
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
|
||||
context.ttl.or(Some(Duration::from_secs(60)))
|
||||
}
|
||||
|
||||
fn set_cache(&self, _: &str, value: V, _: &ExactCacheContext) -> Result<(), Error> {
|
||||
*self.value.lock().unwrap() = Some(value);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn get_cache(&self, _: &str, _: &ExactCacheContext) -> Result<Option<V>, Error> {
|
||||
Ok(self.value.lock().unwrap().clone())
|
||||
}
|
||||
|
||||
async fn disconnect(&self) -> Result<(), Error> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
|
||||
unreachable!()
|
||||
}
|
||||
}
|
||||
|
||||
impl<V> BatchCache for TestCache<V> where V: Clone + Send + Sync + 'static {}
|
||||
|
||||
impl<V> DeleteCache for TestCache<V>
|
||||
where
|
||||
V: Clone + Send + Sync + 'static,
|
||||
{
|
||||
fn delete_cache(&self, _: &str) -> Result<(), Error> {
|
||||
*self.value.lock().unwrap() = None;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl<V> FlushCache for TestCache<V>
|
||||
where
|
||||
V: Clone + Send + Sync + 'static,
|
||||
{
|
||||
fn flush_cache(&self) -> Result<(), Error> {
|
||||
*self.value.lock().unwrap() = None;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl CounterCache for TestCache<f64> {
|
||||
fn increment_cache(&self, _: &str, amount: f64, _: ExactCacheContext) -> Result<f64, Error> {
|
||||
if self.fail {
|
||||
return Err(Error::Unavailable);
|
||||
}
|
||||
let mut value = self.value.lock().unwrap();
|
||||
let incremented = value.unwrap_or_default() + amount;
|
||||
*value = Some(incremented);
|
||||
Ok(incremented)
|
||||
}
|
||||
}
|
||||
|
||||
impl<V> ClaimCache for TestCache<V>
|
||||
where
|
||||
V: Clone + PartialEq + Send + Sync + 'static,
|
||||
{
|
||||
fn claim_cache(
|
||||
&self,
|
||||
_: &str,
|
||||
candidate: V,
|
||||
eligible: &[V],
|
||||
_: ExactCacheContext,
|
||||
) -> Result<V, Error> {
|
||||
if self.fail {
|
||||
return Err(Error::Unavailable);
|
||||
}
|
||||
let mut value = self.value.lock().unwrap();
|
||||
let winner = match value.as_ref() {
|
||||
Some(existing) if eligible.is_empty() || eligible.contains(existing) => {
|
||||
existing.clone()
|
||||
}
|
||||
_ => candidate,
|
||||
};
|
||||
*value = Some(winner.clone());
|
||||
Ok(winner)
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn failed_l2_increment_leaves_l1_unchanged() {
|
||||
let l1 = Arc::new(TestCache::new(Some(10.0), false));
|
||||
let cache = DualCache::new(l1.clone(), Arc::new(TestCache::new(Some(20.0), true)));
|
||||
|
||||
assert_eq!(
|
||||
cache.increment_cache("counter", 2.0, ExactCacheContext::default()),
|
||||
Err(Error::Unavailable)
|
||||
);
|
||||
assert_eq!(
|
||||
l1.get_cache("counter", &ExactCacheContext::default())
|
||||
.unwrap(),
|
||||
Some(10.0)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn claim_uses_l1_fallback_without_overwriting_an_eligible_winner() {
|
||||
let l1 = Arc::new(TestCache::new(Some("first".to_string()), false));
|
||||
let cache = DualCache::new(l1, Arc::new(TestCache::new(None, true)))
|
||||
.with_remote_failure_policy(RemoteFailurePolicy::UseLocal);
|
||||
|
||||
assert_eq!(
|
||||
cache
|
||||
.claim_cache(
|
||||
"affinity",
|
||||
"second".into(),
|
||||
&["first".into(), "second".into()],
|
||||
ExactCacheContext {
|
||||
ttl: Some(Duration::from_secs(60)),
|
||||
},
|
||||
)
|
||||
.unwrap(),
|
||||
"first"
|
||||
);
|
||||
}
|
||||
|
||||
struct SyncPanics(TestCache<String>);
|
||||
|
||||
impl BaseCache for SyncPanics {
|
||||
type Value = String;
|
||||
type Context = ExactCacheContext;
|
||||
|
||||
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
|
||||
self.0.get_ttl(context)
|
||||
}
|
||||
|
||||
fn set_cache(&self, _: &str, _: String, _: &ExactCacheContext) -> Result<(), Error> {
|
||||
panic!("sync L2 write on an async path")
|
||||
}
|
||||
|
||||
fn get_cache(&self, _: &str, _: &ExactCacheContext) -> Result<Option<String>, Error> {
|
||||
panic!("sync L2 read on an async path")
|
||||
}
|
||||
|
||||
async fn async_set_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
value: String,
|
||||
context: ExactCacheContext,
|
||||
) -> Result<(), Error> {
|
||||
self.0.set_cache(key, value, &context)
|
||||
}
|
||||
|
||||
async fn async_get_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
context: &ExactCacheContext,
|
||||
) -> Result<Option<String>, Error> {
|
||||
self.0.get_cache(key, context)
|
||||
}
|
||||
|
||||
async fn async_set_cache_pipeline(
|
||||
&self,
|
||||
cache_list: Vec<(String, String)>,
|
||||
context: ExactCacheContext,
|
||||
) -> Result<(), Error> {
|
||||
for (key, value) in cache_list {
|
||||
self.0.set_cache(&key, value, &context)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn disconnect(&self) -> Result<(), Error> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
|
||||
unreachable!()
|
||||
}
|
||||
}
|
||||
|
||||
impl BatchCache for SyncPanics {
|
||||
async fn async_batch_get_cache(
|
||||
&self,
|
||||
keys: Vec<String>,
|
||||
context: ExactCacheContext,
|
||||
) -> Result<Vec<litellm_cache::BatchEntry<String>>, Error> {
|
||||
assert_eq!(keys, ["missing"]);
|
||||
Ok(vec![match self.0.get_cache("missing", &context)? {
|
||||
Some(value) => litellm_cache::BatchEntry::Hit(value),
|
||||
None => litellm_cache::BatchEntry::Miss,
|
||||
}])
|
||||
}
|
||||
}
|
||||
|
||||
impl DeleteCache for SyncPanics {
|
||||
fn delete_cache(&self, _: &str) -> Result<(), Error> {
|
||||
panic!("sync L2 delete on an async path")
|
||||
}
|
||||
|
||||
async fn async_delete_cache(&self, key: &str) -> Result<(), Error> {
|
||||
self.0.delete_cache(key)
|
||||
}
|
||||
}
|
||||
|
||||
impl FlushCache for SyncPanics {
|
||||
fn flush_cache(&self) -> Result<(), Error> {
|
||||
panic!("sync L2 flush on an async path")
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn async_operations_use_the_async_l2_methods() {
|
||||
let l1 = Arc::new(TestCache::new(None, false));
|
||||
let cache = DualCache::new(
|
||||
l1.clone(),
|
||||
Arc::new(SyncPanics(TestCache::new(
|
||||
Some("remote".to_string()),
|
||||
false,
|
||||
))),
|
||||
);
|
||||
let context = ExactCacheContext::default();
|
||||
|
||||
assert_eq!(
|
||||
cache.async_get_cache("missing", &context).await.unwrap(),
|
||||
Some("remote".into())
|
||||
);
|
||||
assert_eq!(
|
||||
l1.get_cache("missing", &context).unwrap(),
|
||||
Some("remote".into())
|
||||
);
|
||||
|
||||
l1.delete_cache("missing").unwrap();
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_batch_get_cache(vec!["missing".into()], context.clone())
|
||||
.await
|
||||
.unwrap(),
|
||||
[litellm_cache::BatchEntry::Hit("remote".to_string())]
|
||||
);
|
||||
cache
|
||||
.async_set_cache("missing", "written".into(), context.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
cache
|
||||
.async_set_cache_pipeline(vec![("missing".into(), "piped".into())], context.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
cache.async_delete_cache("missing").await.unwrap();
|
||||
assert_eq!(
|
||||
cache.async_get_cache("missing", &context).await.unwrap(),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
struct Unavailable;
|
||||
|
||||
impl BaseCache for Unavailable {
|
||||
type Value = String;
|
||||
type Context = ExactCacheContext;
|
||||
|
||||
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
|
||||
context.ttl
|
||||
}
|
||||
|
||||
fn set_cache(&self, _: &str, _: String, _: &ExactCacheContext) -> Result<(), Error> {
|
||||
Err(Error::Unavailable)
|
||||
}
|
||||
|
||||
fn get_cache(&self, _: &str, _: &ExactCacheContext) -> Result<Option<String>, Error> {
|
||||
Err(Error::Unavailable)
|
||||
}
|
||||
|
||||
async fn disconnect(&self) -> Result<(), Error> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
|
||||
unreachable!()
|
||||
}
|
||||
}
|
||||
|
||||
impl BatchCache for Unavailable {}
|
||||
|
||||
impl DeleteCache for Unavailable {
|
||||
fn delete_cache(&self, _: &str) -> Result<(), Error> {
|
||||
Err(Error::Unavailable)
|
||||
}
|
||||
}
|
||||
|
||||
impl FlushCache for Unavailable {
|
||||
fn flush_cache(&self) -> Result<(), Error> {
|
||||
Err(Error::Unavailable)
|
||||
}
|
||||
}
|
||||
|
||||
impl ClaimCache for Unavailable {
|
||||
fn claim_cache(
|
||||
&self,
|
||||
_: &str,
|
||||
_: String,
|
||||
_: &[String],
|
||||
_: ExactCacheContext,
|
||||
) -> Result<String, Error> {
|
||||
Err(Error::InvalidEntry)
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn remote_failure_policy_selects_propagation_or_the_local_tier() {
|
||||
let context = ExactCacheContext::default();
|
||||
let strict = DualCache::new(Arc::new(TestCache::new(None, false)), Arc::new(Unavailable));
|
||||
assert_eq!(
|
||||
strict.set_cache("key", "value".into(), &context),
|
||||
Err(Error::Unavailable)
|
||||
);
|
||||
assert_eq!(strict.get_cache("key", &context), Err(Error::Unavailable));
|
||||
|
||||
let l1 = Arc::new(TestCache::new(None, false));
|
||||
let degraded = DualCache::new(l1.clone(), Arc::new(Unavailable))
|
||||
.with_remote_failure_policy(RemoteFailurePolicy::UseLocal);
|
||||
assert_eq!(degraded.get_cache("key", &context), Ok(None));
|
||||
degraded.set_cache("key", "value".into(), &context).unwrap();
|
||||
assert_eq!(
|
||||
degraded.get_cache("key", &context),
|
||||
Ok(Some("value".into()))
|
||||
);
|
||||
degraded.delete_cache("key").unwrap();
|
||||
assert_eq!(l1.get_cache("key", &context), Ok(None));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn claim_fallback_does_not_hide_non_availability_errors() {
|
||||
let cache = DualCache::new(
|
||||
Arc::new(TestCache::new(Some("first".to_string()), false)),
|
||||
Arc::new(Unavailable),
|
||||
)
|
||||
.with_remote_failure_policy(RemoteFailurePolicy::UseLocal);
|
||||
assert_eq!(
|
||||
cache.claim_cache(
|
||||
"affinity",
|
||||
"second".into(),
|
||||
&[],
|
||||
ExactCacheContext::default()
|
||||
),
|
||||
Err(Error::InvalidEntry)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn local_only_policies_never_touch_l2() {
|
||||
let l2 = Arc::new(TestCache::new(Some("remote".to_string()), false));
|
||||
let cache = DualCache::new(Arc::new(TestCache::new(None, false)), l2.clone())
|
||||
.with_read_policy(ReadPolicy::LocalOnly)
|
||||
.with_write_policy(WritePolicy::LocalOnly);
|
||||
let context = ExactCacheContext::default();
|
||||
|
||||
assert_eq!(cache.get_cache("key", &context), Ok(None));
|
||||
cache.set_cache("key", "local".into(), &context).unwrap();
|
||||
assert_eq!(l2.get_cache("key", &context), Ok(Some("remote".into())));
|
||||
}
|
||||
|
|
@ -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!({
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@ aws = ["dep:aws-smithy-eventstream", "dep:aws-smithy-types"]
|
|||
sse = ["dep:sse-stream"]
|
||||
|
||||
[dependencies]
|
||||
aws-smithy-eventstream = { version = "=0.61.1", optional = true }
|
||||
aws-smithy-eventstream = { version = "=0.61.4", optional = true }
|
||||
aws-smithy-types = { version = "1.6.1", optional = true }
|
||||
bytes = "1"
|
||||
futures-util.workspace = true
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
|
|
|
|||
|
|
@ -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
|
||||
));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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}")]
|
||||
|
|
|
|||
|
|
@ -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};
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -34,7 +34,7 @@ tokio = { workspace = true, features = ["sync"] }
|
|||
url.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
aws-smithy-eventstream = "=0.61.1"
|
||||
aws-smithy-eventstream = "=0.61.4"
|
||||
aws-smithy-types = "1.6.1"
|
||||
rstest.workspace = true
|
||||
tokio.workspace = true
|
||||
|
|
|
|||
|
|
@ -20,7 +20,18 @@ 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-s3.workspace = true
|
||||
litellm-cache-gcs.workspace = true
|
||||
litellm-cache-disk.workspace = true
|
||||
litellm-cache-response.workspace = true
|
||||
litellm-cache-valkey-semantic = { path = "../cache-valkey-semantic" }
|
||||
serde.workspace = true
|
||||
litellm-auth.workspace = true
|
||||
litellm-auth-aws.workspace = true
|
||||
litellm-callbacks-legacy-python.workspace = true
|
||||
litellm-core.workspace = true
|
||||
litellm-core-utils.workspace = true
|
||||
|
|
@ -32,13 +43,17 @@ litellm-host-python.workspace = true
|
|||
litellm-token-counter = { path = "../token-counter", default-features = false }
|
||||
pyo3.workspace = true
|
||||
pyo3-async-runtimes.workspace = true
|
||||
redis = { version = "1.7.0", features = ["tls-rustls"] }
|
||||
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
|
||||
sha2.workspace = true
|
||||
tokio-tungstenite.workspace = true
|
||||
|
||||
[[bench]]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
276
litellm-rust/crates/python-bridge/src/cache/binding.rs
vendored
Normal file
276
litellm-rust/crates/python-bridge/src/cache/binding.rs
vendored
Normal file
|
|
@ -0,0 +1,276 @@
|
|||
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)?;
|
||||
service.async_lookup_py(py, request)?
|
||||
}
|
||||
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)?;
|
||||
service.async_store_py(py, request, response)
|
||||
}
|
||||
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();
|
||||
service.async_store_batch_py(py, entries)
|
||||
}
|
||||
CacheBinding::PythonCallback(callback) => {
|
||||
callback.async_store_batch(py, callback_result, callback_kwargs)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn async_flush<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
|
||||
self.check_process()?;
|
||||
match &self.binding {
|
||||
CacheBinding::Disabled => ready_none(py),
|
||||
CacheBinding::Native(service) => {
|
||||
let service = service.clone();
|
||||
run_async(py, async move { service.async_flush().await }, cache_error)
|
||||
}
|
||||
CacheBinding::PythonCallback(callback) => callback.async_flush(py),
|
||||
}
|
||||
}
|
||||
|
||||
fn ping<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
|
||||
self.check_process()?;
|
||||
match &self.binding {
|
||||
CacheBinding::Disabled => ready_none(py),
|
||||
CacheBinding::Native(service) => {
|
||||
let service = service.clone();
|
||||
run_async(
|
||||
py,
|
||||
async move { service.test_connection().await },
|
||||
cache_error,
|
||||
)
|
||||
}
|
||||
CacheBinding::PythonCallback(callback) => callback.ping(py),
|
||||
}
|
||||
}
|
||||
|
||||
fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
|
||||
if let CacheBinding::PythonCallback(callback) = &self.binding {
|
||||
callback.traverse(&visit)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
162
litellm-rust/crates/python-bridge/src/cache/callback.rs
vendored
Normal file
162
litellm-rust/crates/python-bridge/src/cache/callback.rs
vendored
Normal file
|
|
@ -0,0 +1,162 @@
|
|||
use pyo3::{
|
||||
PyTraverseError, PyVisit,
|
||||
exceptions::{PyTypeError, PyValueError},
|
||||
prelude::*,
|
||||
types::{PyDict, PyList, PyTuple},
|
||||
};
|
||||
|
||||
use super::future::ready_none;
|
||||
|
||||
pub(super) struct PythonCallback(Py<PyAny>);
|
||||
|
||||
impl PythonCallback {
|
||||
pub(super) fn new(object: Py<PyAny>) -> Self {
|
||||
Self(object)
|
||||
}
|
||||
|
||||
pub(super) fn lookup<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
kwargs: Option<&Bound<'py, PyDict>>,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
self.0
|
||||
.bind(py)
|
||||
.call_method("get_cache", (), Some(callback_kwargs(kwargs)?))
|
||||
}
|
||||
|
||||
pub(super) fn async_lookup<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
kwargs: Option<&Bound<'py, PyDict>>,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
self.0
|
||||
.bind(py)
|
||||
.call_method("async_get_cache", (), Some(callback_kwargs(kwargs)?))
|
||||
}
|
||||
|
||||
pub(super) fn store(
|
||||
&self,
|
||||
py: Python<'_>,
|
||||
response: &Bound<'_, PyAny>,
|
||||
kwargs: Option<&Bound<'_, PyDict>>,
|
||||
) -> PyResult<()> {
|
||||
self.0
|
||||
.bind(py)
|
||||
.call_method("add_cache", (response,), Some(callback_kwargs(kwargs)?))
|
||||
.map(|_| ())
|
||||
}
|
||||
|
||||
pub(super) fn async_store<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
response: &Bound<'py, PyAny>,
|
||||
kwargs: Option<&Bound<'py, PyDict>>,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
self.0.bind(py).call_method(
|
||||
"async_add_cache",
|
||||
(response,),
|
||||
Some(callback_kwargs(kwargs)?),
|
||||
)
|
||||
}
|
||||
|
||||
pub(super) fn lookup_batch<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
requests: &Bound<'py, PyAny>,
|
||||
kwargs: Option<&Bound<'py, PyAny>>,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let results = PyList::empty(py);
|
||||
for kwargs in batch_callback_kwargs(requests, kwargs)? {
|
||||
results.append(
|
||||
self.0
|
||||
.bind(py)
|
||||
.call_method("get_cache", (), Some(&kwargs))?,
|
||||
)?;
|
||||
}
|
||||
Ok(results.into_any())
|
||||
}
|
||||
|
||||
pub(super) fn async_lookup_batch<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
requests: &Bound<'py, PyAny>,
|
||||
kwargs: Option<&Bound<'py, PyAny>>,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let awaitables = batch_callback_kwargs(requests, kwargs)?
|
||||
.iter()
|
||||
.map(|kwargs| {
|
||||
self.0
|
||||
.bind(py)
|
||||
.call_method("async_get_cache", (), Some(kwargs))
|
||||
})
|
||||
.collect::<PyResult<Vec<_>>>()?;
|
||||
py.import("asyncio")?
|
||||
.call_method1("gather", PyTuple::new(py, awaitables)?)
|
||||
}
|
||||
|
||||
pub(super) fn async_store_batch<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
result: Option<&Bound<'py, PyAny>>,
|
||||
kwargs: Option<&Bound<'py, PyDict>>,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let result = result.ok_or_else(|| {
|
||||
PyTypeError::new_err("Python cache callbacks require their original callback_result")
|
||||
})?;
|
||||
self.0.bind(py).call_method(
|
||||
"async_add_cache_pipeline",
|
||||
(result,),
|
||||
Some(callback_kwargs(kwargs)?),
|
||||
)
|
||||
}
|
||||
|
||||
pub(super) fn async_flush<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
|
||||
let object = self.0.bind(py);
|
||||
let backend = match object.getattr_opt("cache")? {
|
||||
Some(backend) if !backend.is_none() => backend,
|
||||
_ => object.clone(),
|
||||
};
|
||||
if backend.hasattr("async_flush_cache")? {
|
||||
return backend.call_method0("async_flush_cache");
|
||||
}
|
||||
backend.call_method0("flush_cache")?;
|
||||
ready_none(py)
|
||||
}
|
||||
|
||||
pub(super) fn ping<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
|
||||
self.0.bind(py).call_method0("ping")
|
||||
}
|
||||
|
||||
pub(super) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
|
||||
visit.call(&self.0)
|
||||
}
|
||||
}
|
||||
|
||||
fn callback_kwargs<'a, 'py>(
|
||||
kwargs: Option<&'a Bound<'py, PyDict>>,
|
||||
) -> PyResult<&'a Bound<'py, PyDict>> {
|
||||
kwargs.ok_or_else(|| {
|
||||
PyTypeError::new_err("Python cache callbacks require their original callback_kwargs")
|
||||
})
|
||||
}
|
||||
|
||||
fn batch_callback_kwargs<'py>(
|
||||
requests: &Bound<'py, PyAny>,
|
||||
kwargs: Option<&Bound<'py, PyAny>>,
|
||||
) -> PyResult<Vec<Bound<'py, PyDict>>> {
|
||||
let kwargs = kwargs
|
||||
.ok_or_else(|| {
|
||||
PyTypeError::new_err(
|
||||
"Python cache callbacks require one original callback_kwargs mapping per request",
|
||||
)
|
||||
})?
|
||||
.try_iter()?
|
||||
.map(|item| Ok(item?.cast_into::<PyDict>()?))
|
||||
.collect::<PyResult<Vec<_>>>()?;
|
||||
if kwargs.len() != requests.len()? {
|
||||
return Err(PyValueError::new_err(
|
||||
"batch cache requests and callback_kwargs must have equal lengths",
|
||||
));
|
||||
}
|
||||
Ok(kwargs)
|
||||
}
|
||||
1331
litellm-rust/crates/python-bridge/src/cache/config.rs
vendored
Normal file
1331
litellm-rust/crates/python-bridge/src/cache/config.rs
vendored
Normal file
File diff suppressed because it is too large
Load diff
63
litellm-rust/crates/python-bridge/src/cache/embedder.rs
vendored
Normal file
63
litellm-rust/crates/python-bridge/src/cache/embedder.rs
vendored
Normal file
|
|
@ -0,0 +1,63 @@
|
|||
use std::{future::Future, sync::Arc};
|
||||
|
||||
use litellm_cache::Error;
|
||||
use litellm_cache_valkey_semantic::Embedder;
|
||||
use litellm_host_python::to_py;
|
||||
use pyo3::{PyTraverseError, PyVisit, prelude::*};
|
||||
use serde_json::Value;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(super) struct PythonEmbedder {
|
||||
sync_embed: Arc<Py<PyAny>>,
|
||||
async_embed_callable: Arc<Py<PyAny>>,
|
||||
}
|
||||
|
||||
impl PythonEmbedder {
|
||||
pub(super) fn from_backend(backend: &Bound<'_, PyAny>) -> PyResult<Self> {
|
||||
Ok(Self {
|
||||
sync_embed: Arc::new(backend.getattr("_get_embedding")?.unbind()),
|
||||
async_embed_callable: Arc::new(backend.getattr("_get_async_embedding")?.unbind()),
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) fn async_embed_awaitable<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
prompt: &str,
|
||||
metadata: &Option<Value>,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let metadata = to_py(py, metadata)?;
|
||||
self.async_embed_callable.bind(py).call1((prompt, metadata))
|
||||
}
|
||||
|
||||
pub(super) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
|
||||
visit.call(&*self.sync_embed)?;
|
||||
visit.call(&*self.async_embed_callable)
|
||||
}
|
||||
}
|
||||
|
||||
impl Embedder for PythonEmbedder {
|
||||
fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result<Vec<f32>, Error> {
|
||||
let result = Python::attach(|py| -> PyResult<Vec<f64>> {
|
||||
let metadata = to_py(py, &metadata)?;
|
||||
self.sync_embed
|
||||
.bind(py)
|
||||
.call1((prompt, metadata))?
|
||||
.extract()
|
||||
})
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
Ok(result.into_iter().map(|value| value as f32).collect())
|
||||
}
|
||||
|
||||
#[expect(
|
||||
clippy::manual_async_fn,
|
||||
reason = "the shared Embedder trait uses an impl Future return"
|
||||
)]
|
||||
fn async_embed(
|
||||
&self,
|
||||
_prompt: &str,
|
||||
_metadata: Option<&Value>,
|
||||
) -> impl Future<Output = Result<Vec<f32>, Error>> + Send {
|
||||
async { Err(Error::Unavailable) }
|
||||
}
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue