mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
Merge remote-tracking branch 'upstream/main' into fix_router_retry_success_logging_49103
# Conflicts: # litellm/router_python_core.py
This commit is contained in:
commit
d35e5681ec
1154 changed files with 89689 additions and 25796 deletions
|
|
@ -128,6 +128,8 @@ fi
|
|||
sudo iptables -L integration_only -n -v -x > "$results/egress-guard.txt"
|
||||
awk '$3 == "REJECT" && $1 > 0 { rejected=1 } END { exit !rejected }' "$results/egress-guard.txt"
|
||||
|
||||
export UV_OFFLINE=true
|
||||
|
||||
if [ "$mode" = parity ]; then
|
||||
export INTEGRATION_WORKERS=1
|
||||
else
|
||||
|
|
|
|||
13
.github/PULL_REQUEST_TEMPLATE/general.md
vendored
13
.github/PULL_REQUEST_TEMPLATE/general.md
vendored
|
|
@ -6,12 +6,13 @@
|
|||
|
||||
## Handwritten 1-3 sentence description
|
||||
|
||||
<!-- The PR author rewrites this in their own words before asking for review. Being able to sum up the PR in a few
|
||||
plain sentences proves the author understands what it does, and human-written technical English is far easier
|
||||
for a reviewer to parse than AI-written English
|
||||
An AI agent may draft a starting point here, but always wraps its draft in the make-me-human-readable tags below,
|
||||
backticks included so the tags stay visible on GitHub, so nobody mistakes it for the author's words. The author rewrites it in their own words and deletes the tags.
|
||||
Tags still present at review time mean the summary is AI-written and has not been rewritten yet
|
||||
<!-- The author writes a plain summary in their own words before requesting review.
|
||||
Wrap newly AI-written summary text in the make-me-human-readable tags below, including backticks.
|
||||
The author reviews or rewrites the draft and removes the tags.
|
||||
When updating an existing PR, preserve the author's wording and removed tags.
|
||||
Changes elsewhere in the body do not make an unchanged summary an AI draft again.
|
||||
Add tags again only around newly AI-written replacement summary text, never merely to match this template.
|
||||
This preserves author review across later updates (2026-11-02, #49779).
|
||||
Say what the PR does for users and why, then anything that keeps it working over time (a sync job, a config someone owns)
|
||||
No bullets, tables, or code names here. Everything else goes inside the Details block below -->
|
||||
|
||||
|
|
|
|||
4
.github/workflows/compat-matrix-image.yml
vendored
4
.github/workflows/compat-matrix-image.yml
vendored
|
|
@ -4,7 +4,7 @@ on:
|
|||
pull_request:
|
||||
paths:
|
||||
- tests/e2e/claude_code/cron_vm/**
|
||||
- tests/e2e/claude_code/pr_gate_version_resolver.py
|
||||
- tests/e2e/cli_harness/npm_version_resolver.py
|
||||
- .github/workflows/compat-matrix-image.yml
|
||||
workflow_dispatch:
|
||||
|
||||
|
|
@ -36,7 +36,7 @@ jobs:
|
|||
whoami
|
||||
gh --version
|
||||
uv --version
|
||||
version="$(uv run --no-project --python 3.12 python /opt/litellm/tests/e2e/claude_code/pr_gate_version_resolver.py)"
|
||||
version="$(uv run --no-project --python 3.12 python /opt/litellm/tests/e2e/cli_harness/npm_version_resolver.py @anthropic-ai/claude-code)"
|
||||
/opt/litellm/tests/e2e/claude_code/cron_vm/install_claude_code.sh "${version}" /tmp/claude-cli
|
||||
/tmp/claude-cli/claude --version
|
||||
'
|
||||
|
|
|
|||
30
.github/workflows/deploy-dev-on-main.yml
vendored
30
.github/workflows/deploy-dev-on-main.yml
vendored
|
|
@ -1,30 +0,0 @@
|
|||
name: Deploy main to gateway-dev
|
||||
|
||||
on:
|
||||
push:
|
||||
branches: [main]
|
||||
|
||||
permissions: {}
|
||||
|
||||
concurrency:
|
||||
group: deploy-dev-dispatch
|
||||
cancel-in-progress: true
|
||||
|
||||
jobs:
|
||||
dispatch:
|
||||
if: github.repository == 'BerriAI/litellm'
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 5
|
||||
steps:
|
||||
- uses: actions/create-github-app-token@fee1f7d63c2ff003460e3d139729b119787bc349 # v2.2.2
|
||||
id: token
|
||||
with:
|
||||
app-id: ${{ secrets.LITELLM_OPS_DISPATCH_APP_ID }}
|
||||
private-key: ${{ secrets.LITELLM_OPS_DISPATCH_APP_PRIVATE_KEY }}
|
||||
owner: BerriAI
|
||||
repositories: litellm-ops
|
||||
permission-actions: write
|
||||
- name: Dispatch the litellm-ops dev deploy
|
||||
env:
|
||||
GH_TOKEN: ${{ steps.token.outputs.token }}
|
||||
run: gh api --method POST repos/BerriAI/litellm-ops/actions/workflows/deploy-dev.yml/dispatches -f ref=main
|
||||
1
.github/workflows/osv-scan.yml
vendored
1
.github/workflows/osv-scan.yml
vendored
|
|
@ -5,6 +5,7 @@ on:
|
|||
branches:
|
||||
- main
|
||||
- "litellm_**"
|
||||
- "rc/**"
|
||||
schedule:
|
||||
- cron: "23 6 * * *"
|
||||
workflow_dispatch:
|
||||
|
|
|
|||
2
.github/workflows/test-linting.yml
vendored
2
.github/workflows/test-linting.yml
vendored
|
|
@ -208,7 +208,7 @@ jobs:
|
|||
exit 0
|
||||
fi
|
||||
retry() { "$@" || { sleep 15; "$@"; } || { sleep 30; "$@"; }; }
|
||||
CLAUDE_VERSION="$(retry uv run --no-sync python tests/e2e/claude_code/pr_gate_version_resolver.py)"
|
||||
CLAUDE_VERSION="$(retry uv run --no-sync python tests/e2e/cli_harness/npm_version_resolver.py @anthropic-ai/claude-code)"
|
||||
tests/e2e/claude_code/cron_vm/install_claude_code.sh "$CLAUDE_VERSION" "$RUNNER_TEMP/claude-cli"
|
||||
PATH="$RUNNER_TEMP/claude-cli:$PATH" uv run --no-sync pytest -q tests/e2e_harness
|
||||
uv run --no-sync pytest --tb=short -q --noconftest -o addopts= -o "pythonpath=tests/e2e tests/e2e_harness" -p no:rerunfailures \
|
||||
|
|
|
|||
1
.github/workflows/test-unit.yml
vendored
1
.github/workflows/test-unit.yml
vendored
|
|
@ -441,6 +441,7 @@ jobs:
|
|||
tests/unit/ocr
|
||||
tests/unit/passthrough
|
||||
tests/unit/rag
|
||||
tests/unit/cookbook
|
||||
tests/unit/realtime_api
|
||||
tests/unit/rerank_api
|
||||
tests/unit/search
|
||||
|
|
|
|||
85
.github/workflows/weekly_load_anomaly.yml
vendored
85
.github/workflows/weekly_load_anomaly.yml
vendored
|
|
@ -1,85 +0,0 @@
|
|||
name: "Weekly Load Anomaly Check"
|
||||
|
||||
on:
|
||||
schedule:
|
||||
- cron: "0 12 * * 6"
|
||||
workflow_dispatch:
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
jobs:
|
||||
weekly-load-anomaly:
|
||||
if: github.event_name != 'schedule' || github.repository == 'BerriAI/litellm'
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 45
|
||||
services:
|
||||
postgres:
|
||||
image: postgres:16.6
|
||||
env:
|
||||
POSTGRES_USER: llmproxy
|
||||
POSTGRES_PASSWORD: dbpassword9090
|
||||
POSTGRES_DB: litellm
|
||||
ports:
|
||||
- 5432:5432
|
||||
options: >-
|
||||
--health-cmd "pg_isready -U llmproxy"
|
||||
--health-interval 5s
|
||||
--health-timeout 5s
|
||||
--health-retries 10
|
||||
env:
|
||||
DATABASE_URL: postgresql://llmproxy:dbpassword9090@localhost:5432/litellm
|
||||
LITELLM_MASTER_KEY: sk-weekly-anomaly-check
|
||||
ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
|
||||
AWS_BEARER_TOKEN_BEDROCK: ${{ secrets.AWS_BEARER_TOKEN_BEDROCK }}
|
||||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
with:
|
||||
python-version: "3.12"
|
||||
|
||||
- name: Set up uv
|
||||
uses: ./.github/actions/setup-uv-with-retries
|
||||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
- name: Cache the Rust build
|
||||
uses: ./.github/actions/cache-cargo-build
|
||||
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra proxy
|
||||
|
||||
- name: Cache Prisma binaries
|
||||
uses: ./.github/actions/cache-prisma-binaries
|
||||
|
||||
- name: Generate Prisma client
|
||||
run: |
|
||||
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
|
||||
|
||||
- name: Start the proxy
|
||||
run: |
|
||||
nohup uv run --no-sync litellm --config tests/e2e/perf/weekly_anomaly_config.yml --port 4000 > proxy.log 2>&1 &
|
||||
for _ in $(seq 1 90); do
|
||||
if curl -fs http://localhost:4000/health/liveliness > /dev/null; then
|
||||
exit 0
|
||||
fi
|
||||
sleep 2
|
||||
done
|
||||
echo "proxy never became live"
|
||||
tail -n 100 proxy.log
|
||||
exit 1
|
||||
|
||||
- name: Run the weekly session anomaly test
|
||||
env:
|
||||
E2E_WEEKLY_ANOMALY: "1"
|
||||
run: |
|
||||
uv run --no-sync pytest tests/e2e/perf/test_weekly_session_anomaly_e2e.py -v --tb=short -rA
|
||||
|
||||
- name: Show proxy log on failure
|
||||
if: failure()
|
||||
run: tail -n 300 proxy.log
|
||||
|
|
@ -444,6 +444,7 @@ Set `LITELLM_PROXY_API_BASE` and `LITELLM_PROXY_API_KEY` and every model call th
|
|||
| [OpenAI-like (`openai_like`)](https://docs.litellm.ai/docs/providers/openai_compatible) | | | | ✅ | | | | | | |
|
||||
| [OpenRouter (`openrouter`)](https://docs.litellm.ai/docs/providers/openrouter) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [OVHCloud AI Endpoints (`ovhcloud`)](https://docs.litellm.ai/docs/providers/ovhcloud) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [Parallel AI (`parallel_ai`)](https://docs.litellm.ai/docs/search/parallel_ai) | | | | | | | | | | |
|
||||
| [Perplexity AI (`perplexity`)](https://docs.litellm.ai/docs/providers/perplexity) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [Petals (`petals`)](https://docs.litellm.ai/docs/providers/petals) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [Pinstripes (`pinstripes`)](https://docs.litellm.ai/docs/providers/pinstripes) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
|
|
|
|||
245
cookbook/gpt_realtime_translate.py
Normal file
245
cookbook/gpt_realtime_translate.py
Normal file
|
|
@ -0,0 +1,245 @@
|
|||
#!/usr/bin/env python3
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import wave
|
||||
from collections.abc import Iterator, Sequence
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from urllib.parse import urlencode, urlsplit, urlunsplit
|
||||
|
||||
import websockets
|
||||
from websockets.asyncio.client import ClientConnection
|
||||
|
||||
SAMPLE_RATE = 24_000
|
||||
CHANNELS = 1
|
||||
SAMPLE_WIDTH = 2
|
||||
CHUNK_DURATION_SECONDS = 0.1
|
||||
CHUNK_BYTES = int(SAMPLE_RATE * CHANNELS * SAMPLE_WIDTH * CHUNK_DURATION_SECONDS)
|
||||
RESPONSE_TIMEOUT_SECONDS = 30.0
|
||||
AUDIO_EVENT_TYPES = frozenset(
|
||||
{
|
||||
"session.output_audio.delta",
|
||||
"response.audio.delta",
|
||||
"response.output_audio.delta",
|
||||
}
|
||||
)
|
||||
TRANSCRIPT_EVENT_TYPES = frozenset(
|
||||
{
|
||||
"session.output_transcript.delta",
|
||||
"response.text.delta",
|
||||
"response.output_audio_transcript.delta",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Settings:
|
||||
input_wav: Path
|
||||
output_wav: Path
|
||||
base_url: str
|
||||
model: str
|
||||
target_language: str
|
||||
trailing_silence_seconds: float
|
||||
api_key: str
|
||||
|
||||
|
||||
def write_stdout(message: str = "", *, end: str = "\n", flush: bool = False) -> None:
|
||||
sys.stdout.write(f"{message}{end}")
|
||||
if flush:
|
||||
sys.stdout.flush()
|
||||
|
||||
|
||||
def write_stderr(message: str) -> None:
|
||||
sys.stderr.write(f"{message}\n")
|
||||
|
||||
|
||||
def parse_args(argv: Sequence[str] | None = None) -> Settings | str:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Stream a 24 kHz PCM16 WAV through gpt-realtime-translate and save the translated audio",
|
||||
)
|
||||
parser.add_argument("input_wav", type=Path)
|
||||
parser.add_argument("--output", type=Path, default=Path("translated.wav"))
|
||||
parser.add_argument("--base-url", default=os.getenv("LITELLM_BASE_URL", "http://localhost:4000"))
|
||||
parser.add_argument("--model", default=os.getenv("REALTIME_TRANSLATE_MODEL", "gpt-realtime-translate"))
|
||||
parser.add_argument("--target-language", default="fr")
|
||||
parser.add_argument("--trailing-silence", type=float, default=1.5)
|
||||
parsed = parser.parse_args(argv)
|
||||
api_key = os.getenv("LITELLM_API_KEY") or os.getenv("OPENAI_API_KEY")
|
||||
if not api_key:
|
||||
return "Set LITELLM_API_KEY or OPENAI_API_KEY before running the script"
|
||||
if parsed.trailing_silence < 0:
|
||||
return "--trailing-silence must be zero or greater"
|
||||
return Settings(
|
||||
input_wav=parsed.input_wav,
|
||||
output_wav=parsed.output,
|
||||
base_url=parsed.base_url,
|
||||
model=parsed.model,
|
||||
target_language=parsed.target_language,
|
||||
trailing_silence_seconds=parsed.trailing_silence,
|
||||
api_key=api_key,
|
||||
)
|
||||
|
||||
|
||||
def translation_url(base_url: str, model: str) -> str | None:
|
||||
parsed = urlsplit(base_url.rstrip("/"))
|
||||
scheme = {"http": "ws", "https": "wss", "ws": "ws", "wss": "wss"}.get(parsed.scheme)
|
||||
if not scheme or not parsed.netloc:
|
||||
return None
|
||||
base_path = parsed.path.rstrip("/")
|
||||
realtime_path = (
|
||||
f"{base_path}/realtime/translations" if base_path.endswith("/v1") else f"{base_path}/v1/realtime/translations"
|
||||
)
|
||||
return urlunsplit((scheme, parsed.netloc, realtime_path, urlencode({"model": model}), ""))
|
||||
|
||||
|
||||
def read_pcm16_wav(path: Path) -> bytes | str:
|
||||
try:
|
||||
with wave.open(str(path), "rb") as source:
|
||||
actual_format = (
|
||||
source.getnchannels(),
|
||||
source.getsampwidth(),
|
||||
source.getframerate(),
|
||||
source.getcomptype(),
|
||||
)
|
||||
expected_format = (CHANNELS, SAMPLE_WIDTH, SAMPLE_RATE, "NONE")
|
||||
if actual_format != expected_format:
|
||||
return (
|
||||
f"{path} must be mono, 16-bit PCM, 24 kHz WAV; received "
|
||||
f"channels={actual_format[0]}, sample_width={actual_format[1]}, "
|
||||
f"sample_rate={actual_format[2]}, compression={actual_format[3]}"
|
||||
)
|
||||
return source.readframes(source.getnframes())
|
||||
except (OSError, EOFError, wave.Error) as exc:
|
||||
return f"Unable to read {path}: {exc}"
|
||||
|
||||
|
||||
def audio_chunks(audio: bytes) -> Iterator[bytes]:
|
||||
return (audio[offset : offset + CHUNK_BYTES] for offset in range(0, len(audio), CHUNK_BYTES))
|
||||
|
||||
|
||||
def audio_message(audio: bytes) -> str:
|
||||
return json.dumps(
|
||||
{
|
||||
"type": "session.input_audio_buffer.append",
|
||||
"audio": base64.b64encode(audio).decode("ascii"),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
async def configure_session(connection: ClientConnection, target_language: str) -> str | None:
|
||||
await connection.send(
|
||||
json.dumps(
|
||||
{
|
||||
"type": "session.update",
|
||||
"session": {"audio": {"output": {"language": target_language}}},
|
||||
}
|
||||
)
|
||||
)
|
||||
while True:
|
||||
raw_event = await asyncio.wait_for(connection.recv(), timeout=20)
|
||||
event = json.loads(raw_event)
|
||||
event_type = event.get("type")
|
||||
if event_type == "session.created":
|
||||
write_stdout(f"Session: {event.get('session', {}).get('id', 'created')}")
|
||||
if event_type == "session.updated":
|
||||
return None
|
||||
if event_type == "error":
|
||||
return f"Realtime API error: {json.dumps(event.get('error', event), ensure_ascii=False)}"
|
||||
|
||||
|
||||
async def send_audio(
|
||||
connection: ClientConnection, pcm: bytes, trailing_silence_seconds: float, finished: asyncio.Event
|
||||
) -> None:
|
||||
silence = bytes(round(trailing_silence_seconds * SAMPLE_RATE * CHANNELS * SAMPLE_WIDTH))
|
||||
try:
|
||||
for chunk in audio_chunks(pcm + silence):
|
||||
await connection.send(audio_message(chunk))
|
||||
await asyncio.sleep(len(chunk) / (SAMPLE_RATE * CHANNELS * SAMPLE_WIDTH))
|
||||
await connection.send(json.dumps({"type": "session.close"}))
|
||||
finally:
|
||||
finished.set()
|
||||
|
||||
|
||||
async def receive_translation(
|
||||
connection: ClientConnection, output_path: Path, sender_finished: asyncio.Event
|
||||
) -> str | None:
|
||||
audio_received = asyncio.Event()
|
||||
try:
|
||||
with wave.open(str(output_path), "wb") as output:
|
||||
output.setnchannels(CHANNELS)
|
||||
output.setsampwidth(SAMPLE_WIDTH)
|
||||
output.setframerate(SAMPLE_RATE)
|
||||
write_stdout("Translation: ", end="", flush=True)
|
||||
while True:
|
||||
try:
|
||||
raw_event = await asyncio.wait_for(connection.recv(), timeout=RESPONSE_TIMEOUT_SECONDS)
|
||||
except TimeoutError:
|
||||
if not sender_finished.is_set():
|
||||
continue
|
||||
return (
|
||||
f"The translation stream went quiet for {RESPONSE_TIMEOUT_SECONDS:g} s without closing the session"
|
||||
)
|
||||
event = json.loads(raw_event)
|
||||
event_type = event.get("type")
|
||||
if event_type in AUDIO_EVENT_TYPES:
|
||||
output.writeframes(base64.b64decode(event.get("delta", ""), validate=True))
|
||||
audio_received.set()
|
||||
elif event_type in TRANSCRIPT_EVENT_TYPES:
|
||||
write_stdout(event.get("delta", event.get("text", "")), end="", flush=True)
|
||||
elif event_type == "error":
|
||||
return f"Realtime API error: {json.dumps(event.get('error', event), ensure_ascii=False)}"
|
||||
elif event_type == "session.closed":
|
||||
write_stdout()
|
||||
return None if audio_received.is_set() else "The translation stream ended without translated audio"
|
||||
except (OSError, wave.Error) as exc:
|
||||
return f"Unable to write {output_path}: {exc}"
|
||||
|
||||
|
||||
async def translate(settings: Settings, pcm: bytes) -> str | None:
|
||||
url = translation_url(settings.base_url, settings.model)
|
||||
if not url:
|
||||
return f"Invalid --base-url: {settings.base_url}"
|
||||
sender_finished = asyncio.Event()
|
||||
try:
|
||||
async with websockets.connect(
|
||||
url,
|
||||
additional_headers={"Authorization": f"Bearer {settings.api_key}"},
|
||||
proxy=None,
|
||||
open_timeout=20,
|
||||
close_timeout=5,
|
||||
) as connection:
|
||||
configuration_error = await configure_session(connection, settings.target_language)
|
||||
if configuration_error:
|
||||
return configuration_error
|
||||
async with asyncio.TaskGroup() as tasks:
|
||||
receiver = tasks.create_task(receive_translation(connection, settings.output_wav, sender_finished))
|
||||
tasks.create_task(send_audio(connection, pcm, settings.trailing_silence_seconds, sender_finished))
|
||||
return receiver.result()
|
||||
except Exception as exc:
|
||||
return f"Translation failed: {type(exc).__name__}: {exc}"
|
||||
|
||||
|
||||
def main(argv: Sequence[str] | None = None) -> int:
|
||||
settings = parse_args(argv)
|
||||
if isinstance(settings, str):
|
||||
write_stderr(settings)
|
||||
return 2
|
||||
pcm = read_pcm16_wav(settings.input_wav)
|
||||
if isinstance(pcm, str):
|
||||
write_stderr(pcm)
|
||||
return 2
|
||||
error = asyncio.run(translate(settings, pcm))
|
||||
if error:
|
||||
write_stderr(error)
|
||||
return 1
|
||||
write_stdout(f"Translated audio: {settings.output_wav}")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
|
|
@ -3,9 +3,11 @@
|
|||
|
||||
import asyncio
|
||||
import base64
|
||||
import copy
|
||||
import json
|
||||
import time
|
||||
from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta
|
||||
from itertools import chain
|
||||
from types import MappingProxyType
|
||||
|
|
@ -34,7 +36,7 @@ from typing_extensions import ReadOnly, Unpack
|
|||
|
||||
import litellm
|
||||
from litellm import Router, verbose_logger
|
||||
from litellm._internal_context import with_service_target
|
||||
from litellm._internal_context import internal_sub_call, post_response_phase, with_service_target
|
||||
from litellm._uuid import uuid
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.constants import (
|
||||
|
|
@ -95,6 +97,7 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attributio
|
|||
request_tags_from_metadata,
|
||||
)
|
||||
from litellm.repositories.managed_file_repository import ManagedFileRepository
|
||||
from litellm.responses.streaming_iterator import spawn_logging_task
|
||||
from litellm.responses.utils import ResponsesAPIRequestUtils
|
||||
from litellm.types.llms.openai import ( # pyright: ignore[reportAttributeAccessIssue]
|
||||
AllMessageValues,
|
||||
|
|
@ -395,6 +398,134 @@ class _RouterFileCallKwargs(TypedDict, total=False):
|
|||
custom_llm_provider: ReadOnly[str | None]
|
||||
|
||||
|
||||
_MODEL_FILE_ID_MAPPING: Final = TypeAdapter(dict[str, str])
|
||||
_FILE_UPLOAD_ONLY_FIELDS: Final = frozenset({"file", "purpose", "expires_after", "extra_body"})
|
||||
_FILE_DELETE_RESERVED_FIELDS: Final = frozenset({"model", "file_id", "_litellm_internal_model_credentials"})
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _ProviderUpload:
|
||||
model_id: str
|
||||
provider_file_id: str
|
||||
upload_call_id: str | None = None
|
||||
|
||||
|
||||
def _provider_uploads(file_object: OpenAIFileObject) -> tuple[_ProviderUpload, ...]:
|
||||
file_hidden_params: Final = cast( # cast-ok: preserve mapping operations on dynamic file metadata
|
||||
dict[str, object], getattr(file_object, HIDDEN_PARAMS_ATTR)
|
||||
)
|
||||
model_file_id_mapping: Final = _MODEL_FILE_ID_MAPPING.validate_python(
|
||||
file_hidden_params.get("model_file_id_mapping") or {}
|
||||
)
|
||||
upload_call_id: Final = file_hidden_params.get("litellm_call_id")
|
||||
return tuple(
|
||||
_ProviderUpload(
|
||||
model_id=model_id,
|
||||
provider_file_id=provider_file_id,
|
||||
upload_call_id=upload_call_id if isinstance(upload_call_id, str) else None,
|
||||
)
|
||||
for model_id, provider_file_id in model_file_id_mapping.items()
|
||||
)
|
||||
|
||||
|
||||
def _keeps_one_upload_per_target(llm_router: Router, model_id: str) -> bool:
|
||||
credentials: Final = _provider_file_retrieve_credentials(llm_router=llm_router, model_id=model_id)
|
||||
return (credentials or {}).get("custom_llm_provider") == LlmProviders.LITELLM_PROXY.value
|
||||
|
||||
|
||||
def _is_kept_upload(uploads: Sequence[_ProviderUpload], index: int, per_target_model_ids: frozenset[str]) -> bool:
|
||||
upload: Final = uploads[index]
|
||||
other_uploads: Final = uploads[index + 1 :] if upload.model_id in per_target_model_ids else uploads[:index]
|
||||
return all(other.model_id != upload.model_id for other in other_uploads)
|
||||
|
||||
|
||||
def _upload_from_trailing_id_fields(decoded_unified_file_id: str) -> _ProviderUpload:
|
||||
trailing_fields: Final = decoded_unified_file_id.rpartition(";llm_output_file_id,")[2]
|
||||
provider_file_id, _, model_id = trailing_fields.partition(";llm_output_file_model_id,")
|
||||
return _ProviderUpload(model_id=model_id, provider_file_id=provider_file_id)
|
||||
|
||||
|
||||
def _carried_upload_missing_from(unified_file_id: str, model_file_ids: Mapping[str, str]) -> _ProviderUpload | None:
|
||||
decoded: Final = _is_base64_encoded_unified_file_id(unified_file_id)
|
||||
carried: Final = _upload_from_trailing_id_fields(decoded) if decoded else None
|
||||
if carried is None or model_file_ids.get(carried.model_id) in (None, carried.provider_file_id):
|
||||
return None
|
||||
return carried
|
||||
|
||||
|
||||
async def _delete_provider_file(
|
||||
llm_router: Router,
|
||||
model_id: str,
|
||||
provider_file_id: str,
|
||||
request_data: Mapping[str, object],
|
||||
) -> None:
|
||||
credentials: Final = llm_router.registry().credentials(model_id, form="with_provider")
|
||||
delete_data: Final = {
|
||||
**{k: v for k, v in request_data.items() if k not in _FILE_DELETE_RESERVED_FIELDS},
|
||||
**(
|
||||
{"_litellm_internal_model_credentials": MappingProxyType(dict(credentials))}
|
||||
if credentials is not None
|
||||
else {}
|
||||
),
|
||||
}
|
||||
await llm_router.afile_delete(
|
||||
model=model_id,
|
||||
file_id=provider_file_id,
|
||||
**cast(_RouterFileCallKwargs, delete_data),
|
||||
)
|
||||
|
||||
|
||||
def _with_own_nested_dicts(request_data: Mapping[str, object]) -> MappingProxyType[str, object]:
|
||||
return MappingProxyType({k: copy.copy(v) if isinstance(v, dict) else v for k, v in request_data.items()})
|
||||
|
||||
|
||||
async def _delete_redundant_upload(
|
||||
llm_router: Router,
|
||||
upload: _ProviderUpload,
|
||||
delete_request: Mapping[str, object],
|
||||
parent_call_id: str | None,
|
||||
) -> None:
|
||||
try:
|
||||
with internal_sub_call(parent_call_id):
|
||||
await _delete_provider_file(
|
||||
llm_router, upload.model_id, upload.provider_file_id, _with_own_nested_dicts(delete_request)
|
||||
)
|
||||
except Exception as error:
|
||||
verbose_logger.warning(
|
||||
f"Could not delete the redundant upload provider_file_id={sanitize_for_log(upload.provider_file_id)} "
|
||||
f"at deployment model_id={sanitize_for_log(upload.model_id)}: "
|
||||
f"{type(error).__name__} {sanitize_for_log(error)}"
|
||||
)
|
||||
|
||||
|
||||
async def _delete_redundant_uploads(
|
||||
llm_router: Router,
|
||||
uploads: Sequence[_ProviderUpload],
|
||||
delete_request: Mapping[str, object],
|
||||
) -> None:
|
||||
await asyncio.gather(
|
||||
*(_delete_redundant_upload(llm_router, upload, delete_request, upload.upload_call_id) for upload in uploads)
|
||||
)
|
||||
|
||||
|
||||
def _start_deleting_redundant_uploads(
|
||||
llm_router: Router,
|
||||
uploads: Sequence[_ProviderUpload],
|
||||
create_file_request: CreateFileRequest,
|
||||
) -> None:
|
||||
if not uploads:
|
||||
return
|
||||
delete_request: Final = MappingProxyType(
|
||||
{k: v for k, v in create_file_request.items() if k not in _FILE_UPLOAD_ONLY_FIELDS}
|
||||
)
|
||||
with post_response_phase():
|
||||
spawn_logging_task(
|
||||
asyncio.get_running_loop(),
|
||||
_delete_redundant_uploads(llm_router, uploads, delete_request),
|
||||
task_name="managed file repeat upload cleanup",
|
||||
)
|
||||
|
||||
|
||||
def _managed_database(prisma_client: _ManagedResourcePrismaClient) -> _ManagedResourceDatabase:
|
||||
return prisma_client.db
|
||||
|
||||
|
|
@ -2034,16 +2165,16 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
target_model_names_list=target_model_names_list,
|
||||
)
|
||||
|
||||
## STORE MODEL MAPPINGS IN DB
|
||||
model_mappings: Dict[str, str] = {}
|
||||
|
||||
for file_object in responses:
|
||||
file_hidden_params = cast( # cast-ok: preserve mapping operations on dynamic file metadata
|
||||
dict[str, object], getattr(file_object, HIDDEN_PARAMS_ATTR)
|
||||
)
|
||||
model_file_id_mapping = file_hidden_params.get("model_file_id_mapping")
|
||||
if model_file_id_mapping and isinstance(model_file_id_mapping, dict):
|
||||
model_mappings.update(model_file_id_mapping)
|
||||
uploads: Final = tuple(chain.from_iterable(map(_provider_uploads, responses)))
|
||||
per_target_model_ids: Final = frozenset(
|
||||
model_id
|
||||
for model_id in {upload.model_id for upload in uploads}
|
||||
if _keeps_one_upload_per_target(llm_router, model_id)
|
||||
)
|
||||
kept_uploads: Final = tuple(
|
||||
upload for index, upload in enumerate(uploads) if _is_kept_upload(uploads, index, per_target_model_ids)
|
||||
)
|
||||
model_mappings: Final = {upload.model_id: upload.provider_file_id for upload in kept_uploads}
|
||||
|
||||
await self.store_unified_file_id(
|
||||
file_id=response.id,
|
||||
|
|
@ -2052,6 +2183,15 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
model_mappings=model_mappings,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
_start_deleting_redundant_uploads(
|
||||
llm_router,
|
||||
tuple(
|
||||
upload
|
||||
for upload in uploads
|
||||
if upload not in kept_uploads and upload.model_id not in per_target_model_ids
|
||||
),
|
||||
create_file_request,
|
||||
)
|
||||
|
||||
# Emit Prometheus metrics for managed file creation
|
||||
prom_logger = self._get_prometheus_logger()
|
||||
|
|
@ -2692,23 +2832,13 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
specific_model_file_id_mapping: Final = model_file_id_mapping.get(file_id)
|
||||
if not specific_model_file_id_mapping:
|
||||
return
|
||||
filtered_data: Final = {
|
||||
k: v for k, v in data.items() if k not in ("model", "file_id", "_litellm_internal_model_credentials")
|
||||
}
|
||||
for model_id, model_file_id in specific_model_file_id_mapping.items():
|
||||
credentials = llm_router.registry().credentials(model_id, form="with_provider")
|
||||
delete_data = {
|
||||
**filtered_data,
|
||||
**(
|
||||
{"_litellm_internal_model_credentials": MappingProxyType(dict(credentials))}
|
||||
if credentials is not None
|
||||
else {}
|
||||
),
|
||||
}
|
||||
await llm_router.afile_delete(
|
||||
model=model_id,
|
||||
file_id=model_file_id,
|
||||
**cast(_RouterFileCallKwargs, delete_data),
|
||||
await _delete_provider_file(llm_router, model_id, model_file_id, data)
|
||||
unrecorded_upload: Final = _carried_upload_missing_from(file_id, specific_model_file_id_mapping)
|
||||
delete_call_id: Final = data.get("litellm_call_id")
|
||||
if unrecorded_upload is not None:
|
||||
await _delete_redundant_upload(
|
||||
llm_router, unrecorded_upload, data, delete_call_id if isinstance(delete_call_id, str) else None
|
||||
)
|
||||
|
||||
async def afile_content(
|
||||
|
|
|
|||
87
litellm-rust/Cargo.lock
generated
87
litellm-rust/Cargo.lock
generated
|
|
@ -3569,37 +3569,6 @@ dependencies = [
|
|||
"uuid",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-cache-redis"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"litellm-cache",
|
||||
"litellm-cache-testing",
|
||||
"litellm-error",
|
||||
"r2d2",
|
||||
"redis",
|
||||
"redis-test",
|
||||
"rstest",
|
||||
"serde_json",
|
||||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-cache-redis-semantic"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"litellm-cache",
|
||||
"litellm-cache-redis",
|
||||
"litellm-cache-testing",
|
||||
"redis",
|
||||
"redis-test",
|
||||
"rstest",
|
||||
"serde_json",
|
||||
"sha2 0.10.9",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-cache-response"
|
||||
version = "0.1.0"
|
||||
|
|
@ -3607,11 +3576,8 @@ dependencies = [
|
|||
"litellm-cache",
|
||||
"litellm-cache-gcs",
|
||||
"litellm-cache-memory",
|
||||
"litellm-cache-redis",
|
||||
"litellm-http",
|
||||
"py_literal",
|
||||
"redis",
|
||||
"redis-test",
|
||||
"rstest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
|
|
@ -3650,24 +3616,6 @@ dependencies = [
|
|||
"litellm-cache",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-cache-valkey-semantic"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"litellm-cache",
|
||||
"litellm-cache-redis",
|
||||
"litellm-cache-testing",
|
||||
"litellm-error",
|
||||
"redis",
|
||||
"redis-test",
|
||||
"rstest",
|
||||
"serde_json",
|
||||
"sha2 0.10.9",
|
||||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
"uuid",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-callbacks-legacy-python"
|
||||
version = "0.1.0"
|
||||
|
|
@ -4339,8 +4287,6 @@ dependencies = [
|
|||
"litellm-auth",
|
||||
"litellm-auth-aws",
|
||||
"litellm-cache",
|
||||
"litellm-cache-memory",
|
||||
"litellm-cache-redis",
|
||||
"litellm-cache-response",
|
||||
"litellm-callbacks-legacy-python",
|
||||
"litellm-core-utils",
|
||||
|
|
@ -4369,7 +4315,6 @@ dependencies = [
|
|||
"prost",
|
||||
"pyo3",
|
||||
"pyo3-async-runtimes",
|
||||
"redis",
|
||||
"reqwest",
|
||||
"rstest",
|
||||
"serde",
|
||||
|
|
@ -5782,17 +5727,6 @@ version = "6.0.0"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf"
|
||||
|
||||
[[package]]
|
||||
name = "r2d2"
|
||||
version = "0.8.10"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "51de85fb3fb6524929c8a2eb85e6b6d363de4e8c48f9e2c2eac4944abc181c93"
|
||||
dependencies = [
|
||||
"log",
|
||||
"parking_lot",
|
||||
"scheduled-thread-pool",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rand"
|
||||
version = "0.8.7"
|
||||
|
|
@ -5966,18 +5900,6 @@ dependencies = [
|
|||
"xxhash-rust",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "redis-test"
|
||||
version = "1.0.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "804d36862e4323b69f96440cbb13c9894fc90176abdeaf91264e21d5d77f6aca"
|
||||
dependencies = [
|
||||
"rand 0.9.5",
|
||||
"redis",
|
||||
"socket2 0.6.5",
|
||||
"tempfile",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "redox_syscall"
|
||||
version = "0.5.18"
|
||||
|
|
@ -6454,15 +6376,6 @@ dependencies = [
|
|||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "scheduled-thread-pool"
|
||||
version = "0.2.7"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3cbc66816425a074528352f5789333ecff06ca41b36b0b0efdfbb29edc391a19"
|
||||
dependencies = [
|
||||
"parking_lot",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "schemars"
|
||||
version = "0.9.0"
|
||||
|
|
|
|||
|
|
@ -57,11 +57,9 @@ litellm-db-testing = { path = "crates/db-testing" }
|
|||
litellm-cache = { path = "crates/cache" }
|
||||
litellm-cache-azure-blob = { path = "crates/cache-azure-blob" }
|
||||
litellm-cache-memory = { path = "crates/cache-memory" }
|
||||
litellm-cache-redis = { path = "crates/cache-redis" }
|
||||
litellm-cache-s3 = { path = "crates/cache-s3" }
|
||||
litellm-cache-gcs = { path = "crates/cache-gcs" }
|
||||
litellm-cache-disk = { path = "crates/cache-disk" }
|
||||
litellm-cache-redis-semantic = { path = "crates/cache-redis-semantic" }
|
||||
litellm-cache-response = { path = "crates/cache-response" }
|
||||
litellm-cache-qdrant-semantic = { path = "crates/cache-qdrant-semantic" }
|
||||
litellm-cache-testing = { path = "crates/cache-testing" }
|
||||
|
|
@ -97,7 +95,6 @@ macro_rules_attribute = "0.2.3"
|
|||
schemars = "1"
|
||||
reqwest = { version = "0.13", default-features = false, features = ["json", "multipart", "rustls", "http2", "stream", "query"] }
|
||||
redis = "1.8.0"
|
||||
redis-test = "1.0.4"
|
||||
testcontainers-modules = "0.15.0"
|
||||
qdrant-client = { version = "1.19.0", default-features = false }
|
||||
uuid = { version = "1", features = ["v4"] }
|
||||
|
|
|
|||
|
|
@ -1,19 +0,0 @@
|
|||
[package]
|
||||
name = "litellm-cache-redis-semantic"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
litellm-cache.workspace = true
|
||||
litellm-cache-redis.workspace = true
|
||||
redis = { workspace = true, features = ["tls-rustls"] }
|
||||
sha2.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
litellm-cache-testing.workspace = true
|
||||
redis-test.workspace = true
|
||||
rstest.workspace = true
|
||||
serde_json.workspace = true
|
||||
tokio.workspace = true
|
||||
|
|
@ -1,404 +0,0 @@
|
|||
use std::{
|
||||
sync::Arc,
|
||||
time::{Duration, SystemTime, UNIX_EPOCH},
|
||||
};
|
||||
|
||||
use litellm_cache::{
|
||||
BaseCache, CacheCodec, Error, SemanticCacheContext,
|
||||
semantic::{Embedder, SemanticCache, SemanticLookup, prompt_from_context},
|
||||
};
|
||||
use litellm_cache_redis::{
|
||||
RedisTopology,
|
||||
connection::{ConnectionRef, Connections},
|
||||
};
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
use crate::{
|
||||
RedisSemanticConfig,
|
||||
index::{CACHE_KEY_FIELD, Index, VECTOR_FIELD},
|
||||
reply::{bytes_field, first_document, number_field, string_field},
|
||||
};
|
||||
|
||||
struct Inner {
|
||||
index: Index,
|
||||
distance_threshold: f64,
|
||||
clock: fn() -> f64,
|
||||
}
|
||||
|
||||
impl Inner {
|
||||
fn new(config: RedisSemanticConfig, clock: fn() -> f64) -> Self {
|
||||
Self {
|
||||
index: Index::new(config.index_name),
|
||||
distance_threshold: 1.0 - f64::from(config.similarity_threshold),
|
||||
clock,
|
||||
}
|
||||
}
|
||||
|
||||
fn store(
|
||||
&self,
|
||||
connection: &mut ConnectionRef<'_>,
|
||||
tag: &str,
|
||||
response: Vec<u8>,
|
||||
prompt: &str,
|
||||
vector: &[f32],
|
||||
ttl: Option<Duration>,
|
||||
) -> Result<(), Error> {
|
||||
let index = self.index.ensure(connection, vector.len())?;
|
||||
let entry_id = entry_id(prompt, tag);
|
||||
let hash_key = format!("{index}:{entry_id}");
|
||||
redis::cmd("HSET")
|
||||
.arg(&hash_key)
|
||||
.arg("entry_id")
|
||||
.arg(&entry_id)
|
||||
.arg("prompt")
|
||||
.arg(prompt)
|
||||
.arg("response")
|
||||
.arg(response)
|
||||
.arg(VECTOR_FIELD)
|
||||
.arg(vector_buffer(vector))
|
||||
.arg("inserted_at")
|
||||
.arg(format!("{}", (self.clock)()))
|
||||
.arg("updated_at")
|
||||
.arg(format!("{}", (self.clock)()))
|
||||
.arg(CACHE_KEY_FIELD)
|
||||
.arg(tag)
|
||||
.query::<()>(connection)
|
||||
.map_err(Error::unavailable)?;
|
||||
if let Some(ttl) = ttl {
|
||||
redis::cmd("EXPIRE")
|
||||
.arg(&hash_key)
|
||||
.arg(ttl_seconds(ttl))
|
||||
.query::<()>(connection)
|
||||
.map_err(Error::unavailable)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn lookup(
|
||||
&self,
|
||||
connection: &mut ConnectionRef<'_>,
|
||||
tag: &str,
|
||||
vector: &[f32],
|
||||
) -> Result<SemanticLookup<Vec<u8>>, Error> {
|
||||
let index = self.index.ensure(connection, vector.len())?;
|
||||
let query = format!(
|
||||
"(@{CACHE_KEY_FIELD}:{{{}}})=>[KNN 1 @{VECTOR_FIELD} $vector AS vector_distance]",
|
||||
escape_tag(tag)
|
||||
);
|
||||
let result = redis::cmd("FT.SEARCH")
|
||||
.arg(&index)
|
||||
.arg(query)
|
||||
.arg("RETURN")
|
||||
.arg(8)
|
||||
.arg("entry_id")
|
||||
.arg("prompt")
|
||||
.arg("response")
|
||||
.arg("inserted_at")
|
||||
.arg("updated_at")
|
||||
.arg("metadata")
|
||||
.arg(CACHE_KEY_FIELD)
|
||||
.arg("vector_distance")
|
||||
.arg("SORTBY")
|
||||
.arg("vector_distance")
|
||||
.arg("ASC")
|
||||
.arg("DIALECT")
|
||||
.arg(2)
|
||||
.arg("LIMIT")
|
||||
.arg(0)
|
||||
.arg(1)
|
||||
.arg("PARAMS")
|
||||
.arg(2)
|
||||
.arg("vector")
|
||||
.arg(vector_buffer(vector))
|
||||
.query::<redis::Value>(connection)
|
||||
.map_err(Error::unavailable)?;
|
||||
let Some(fields) = first_document(&result) else {
|
||||
return Ok(SemanticLookup::miss(Some(0.0)));
|
||||
};
|
||||
if string_field(fields, CACHE_KEY_FIELD).as_deref() != Some(tag) {
|
||||
return Ok(SemanticLookup::miss(Some(0.0)));
|
||||
}
|
||||
// redisvl's range query only returns entries within the distance threshold, so a
|
||||
// farther hit reads as no result.
|
||||
let Some(distance) = number_field(fields, "vector_distance")
|
||||
.filter(|distance| *distance <= self.distance_threshold)
|
||||
else {
|
||||
return Ok(SemanticLookup::miss(Some(0.0)));
|
||||
};
|
||||
let Some(response) = bytes_field(fields, "response") else {
|
||||
return Ok(SemanticLookup::miss(Some(0.0)));
|
||||
};
|
||||
Ok(SemanticLookup {
|
||||
value: Some(response),
|
||||
similarity: Some(1.0 - distance),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// `RedisSemanticCache`: a redisvl-compatible semantic index on Redis Stack. Values go through
|
||||
/// the injected codec, so the response layer decides what a cached entry is.
|
||||
pub struct RedisSemanticCache<E, S, C = redis::Connection> {
|
||||
connections: Arc<Connections<C>>,
|
||||
embedder: E,
|
||||
codec: S,
|
||||
inner: Arc<Inner>,
|
||||
}
|
||||
|
||||
impl<E: Embedder, S: CacheCodec> RedisSemanticCache<E, S> {
|
||||
pub fn new(
|
||||
url: &str,
|
||||
embedder: E,
|
||||
codec: S,
|
||||
config: RedisSemanticConfig,
|
||||
) -> Result<Self, Error> {
|
||||
Ok(Self {
|
||||
connections: Arc::new(Connections::open(url, &RedisTopology::Standalone)?),
|
||||
embedder,
|
||||
codec,
|
||||
inner: Arc::new(Inner::new(config, timestamp)),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl<E, S, C> RedisSemanticCache<E, S, C>
|
||||
where
|
||||
E: Embedder,
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
pub fn with_connection(
|
||||
connection: C,
|
||||
embedder: E,
|
||||
codec: S,
|
||||
config: RedisSemanticConfig,
|
||||
) -> Self {
|
||||
Self {
|
||||
connections: Arc::new(Connections::fixed(connection)),
|
||||
embedder,
|
||||
codec,
|
||||
inner: Arc::new(Inner::new(config, timestamp)),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_clock(self, clock: fn() -> f64) -> Self {
|
||||
let config = RedisSemanticConfig {
|
||||
index_name: self.index_name().to_owned(),
|
||||
similarity_threshold: self.similarity_threshold(),
|
||||
};
|
||||
Self {
|
||||
inner: Arc::new(Inner::new(config, clock)),
|
||||
..self
|
||||
}
|
||||
}
|
||||
|
||||
pub fn embedder(&self) -> &E {
|
||||
&self.embedder
|
||||
}
|
||||
|
||||
pub fn index_name(&self) -> &str {
|
||||
self.inner.index.name()
|
||||
}
|
||||
|
||||
pub fn similarity_threshold(&self) -> f32 {
|
||||
(1.0 - self.inner.distance_threshold) as f32
|
||||
}
|
||||
|
||||
fn tag<'a>(key: &'a str, context: &'a SemanticCacheContext) -> &'a str {
|
||||
context.scope.as_deref().unwrap_or(key)
|
||||
}
|
||||
|
||||
fn decode(&self, lookup: SemanticLookup<Vec<u8>>) -> Result<SemanticLookup<S::Value>, Error> {
|
||||
Ok(SemanticLookup {
|
||||
value: lookup
|
||||
.value
|
||||
.map(|bytes| self.codec.decode(&bytes))
|
||||
.transpose()?,
|
||||
similarity: lookup.similarity,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl<E, S, C> BaseCache for RedisSemanticCache<E, S, C>
|
||||
where
|
||||
E: Embedder,
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
type Value = S::Value;
|
||||
type Context = SemanticCacheContext;
|
||||
|
||||
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
|
||||
context.ttl
|
||||
}
|
||||
|
||||
fn set_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
value: Self::Value,
|
||||
context: &Self::Context,
|
||||
) -> Result<(), Error> {
|
||||
let Some(prompt) = prompt_from_context(context) else {
|
||||
return Ok(());
|
||||
};
|
||||
let response = self.codec.encode(&value)?;
|
||||
let vector = self.embedder.embed(&prompt, context.metadata.as_ref())?;
|
||||
let tag = Self::tag(key, context);
|
||||
self.connections.execute(|connection| {
|
||||
self.inner
|
||||
.store(connection, tag, response, &prompt, &vector, context.ttl)
|
||||
})
|
||||
}
|
||||
|
||||
fn get_cache(&self, key: &str, context: &Self::Context) -> Result<Option<Self::Value>, Error> {
|
||||
self.get_cache_with_similarity(key, context)
|
||||
.map(|lookup| lookup.value)
|
||||
}
|
||||
|
||||
async fn async_set_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
value: Self::Value,
|
||||
context: Self::Context,
|
||||
) -> Result<(), Error> {
|
||||
let Some(prompt) = prompt_from_context(&context) else {
|
||||
return Ok(());
|
||||
};
|
||||
let response = self.codec.encode(&value)?;
|
||||
let vector = self
|
||||
.embedder
|
||||
.async_embed(&prompt, context.metadata.as_ref())
|
||||
.await?;
|
||||
let tag = Self::tag(key, &context).to_owned();
|
||||
let inner = Arc::clone(&self.inner);
|
||||
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
|
||||
inner.store(connection, &tag, response, &prompt, &vector, context.ttl)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
async fn async_get_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
context: &Self::Context,
|
||||
) -> Result<Option<Self::Value>, Error> {
|
||||
self.async_get_cache_with_similarity(key, context)
|
||||
.await
|
||||
.map(|lookup| lookup.value)
|
||||
}
|
||||
}
|
||||
|
||||
/// Python stamps a similarity of `0.0` when there is no prompt or no hit in the key's scope.
|
||||
impl<E, S, C> SemanticCache for RedisSemanticCache<E, S, C>
|
||||
where
|
||||
E: Embedder,
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
fn get_cache_with_similarity(
|
||||
&self,
|
||||
key: &str,
|
||||
context: &Self::Context,
|
||||
) -> Result<SemanticLookup<Self::Value>, Error> {
|
||||
let Some(prompt) = prompt_from_context(context) else {
|
||||
return Ok(SemanticLookup::miss(Some(0.0)));
|
||||
};
|
||||
let vector = self.embedder.embed(&prompt, context.metadata.as_ref())?;
|
||||
let tag = Self::tag(key, context);
|
||||
let lookup = self
|
||||
.connections
|
||||
.execute(|connection| self.inner.lookup(connection, tag, &vector))?;
|
||||
self.decode(lookup)
|
||||
}
|
||||
|
||||
async fn async_get_cache_with_similarity(
|
||||
&self,
|
||||
key: &str,
|
||||
context: &Self::Context,
|
||||
) -> Result<SemanticLookup<Self::Value>, Error> {
|
||||
let Some(prompt) = prompt_from_context(context) else {
|
||||
return Ok(SemanticLookup::miss(Some(0.0)));
|
||||
};
|
||||
let vector = self
|
||||
.embedder
|
||||
.async_embed(&prompt, context.metadata.as_ref())
|
||||
.await?;
|
||||
let tag = Self::tag(key, context).to_owned();
|
||||
let inner = Arc::clone(&self.inner);
|
||||
let lookup = Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
|
||||
inner.lookup(connection, &tag, &vector)
|
||||
})
|
||||
.await?;
|
||||
self.decode(lookup)
|
||||
}
|
||||
}
|
||||
|
||||
fn timestamp() -> f64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|duration| duration.as_secs_f64())
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
fn entry_id(prompt: &str, tag: &str) -> String {
|
||||
let mut digest = Sha256::new();
|
||||
digest.update(prompt.as_bytes());
|
||||
digest.update(CACHE_KEY_FIELD.as_bytes());
|
||||
digest.update(tag.as_bytes());
|
||||
format!("{:x}", digest.finalize())
|
||||
}
|
||||
|
||||
fn vector_buffer(vector: &[f32]) -> Vec<u8> {
|
||||
vector
|
||||
.iter()
|
||||
.flat_map(|component| component.to_le_bytes())
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn escape_tag(value: &str) -> String {
|
||||
let mut escaped = String::with_capacity(value.len());
|
||||
for ch in value.chars() {
|
||||
if matches!(
|
||||
ch,
|
||||
',' | '.'
|
||||
| '<'
|
||||
| '>'
|
||||
| '{'
|
||||
| '}'
|
||||
| '['
|
||||
| ']'
|
||||
| '\\'
|
||||
| '"'
|
||||
| '\''
|
||||
| ':'
|
||||
| ';'
|
||||
| '!'
|
||||
| '@'
|
||||
| '#'
|
||||
| '$'
|
||||
| '%'
|
||||
| '^'
|
||||
| '&'
|
||||
| '*'
|
||||
| '('
|
||||
| ')'
|
||||
| '-'
|
||||
| '+'
|
||||
| '='
|
||||
| '~'
|
||||
| '|'
|
||||
| '/'
|
||||
| ' '
|
||||
| '?'
|
||||
) {
|
||||
escaped.push('\\');
|
||||
}
|
||||
escaped.push(ch);
|
||||
}
|
||||
escaped
|
||||
}
|
||||
|
||||
fn ttl_seconds(ttl: Duration) -> u64 {
|
||||
ttl.as_secs()
|
||||
.saturating_add(u64::from(ttl.subsec_nanos() > 0))
|
||||
.max(1)
|
||||
}
|
||||
|
|
@ -1,8 +0,0 @@
|
|||
/// `RedisSemanticCache.DEFAULT_REDIS_INDEX_NAME`.
|
||||
pub const DEFAULT_INDEX_NAME: &str = "litellm_semantic_cache_index";
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct RedisSemanticConfig {
|
||||
pub index_name: String,
|
||||
pub similarity_threshold: f32,
|
||||
}
|
||||
|
|
@ -1,209 +0,0 @@
|
|||
use std::sync::OnceLock;
|
||||
|
||||
use litellm_cache::Error;
|
||||
use litellm_cache_redis::connection::ConnectionRef;
|
||||
|
||||
use crate::reply::{number_value, string_value};
|
||||
|
||||
pub(crate) const CACHE_KEY_FIELD: &str = "litellm_cache_key";
|
||||
pub(crate) const VECTOR_FIELD: &str = "prompt_vector";
|
||||
|
||||
/// The redisvl `SemanticCache` index, resolved once per cache: the configured name when its
|
||||
/// schema fits, else `<name>_isolated`, recreated when that one is stale too.
|
||||
pub(crate) struct Index {
|
||||
name: String,
|
||||
resolved: OnceLock<String>,
|
||||
}
|
||||
|
||||
impl Index {
|
||||
pub(crate) fn new(name: String) -> Self {
|
||||
Self {
|
||||
name,
|
||||
resolved: OnceLock::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn name(&self) -> &str {
|
||||
&self.name
|
||||
}
|
||||
|
||||
pub(crate) fn ensure(
|
||||
&self,
|
||||
connection: &mut ConnectionRef<'_>,
|
||||
dims: usize,
|
||||
) -> Result<String, Error> {
|
||||
if let Some(name) = self.resolved.get() {
|
||||
return Ok(name.clone());
|
||||
}
|
||||
let name = match index_compatible(connection, &self.name, dims)? {
|
||||
Some(true) => self.name.clone(),
|
||||
Some(false) => self.isolated(connection, dims)?,
|
||||
None => match create_index(connection, &self.name, dims) {
|
||||
Ok(()) => self.name.clone(),
|
||||
Err(error) => match index_compatible(connection, &self.name, dims)? {
|
||||
Some(true) => self.name.clone(),
|
||||
Some(false) => self.isolated(connection, dims)?,
|
||||
None => return Err(error),
|
||||
},
|
||||
},
|
||||
};
|
||||
let _ = self.resolved.set(name.clone());
|
||||
Ok(name)
|
||||
}
|
||||
|
||||
fn isolated(&self, connection: &mut ConnectionRef<'_>, dims: usize) -> Result<String, Error> {
|
||||
let name = format!("{}_isolated", self.name);
|
||||
match index_compatible(connection, &name, dims)? {
|
||||
Some(true) => Ok(name),
|
||||
Some(false) => {
|
||||
redis::cmd("FT.DROPINDEX")
|
||||
.arg(&name)
|
||||
.query::<()>(connection)
|
||||
.map_err(Error::unavailable)?;
|
||||
create_index(connection, &name, dims)?;
|
||||
Ok(name)
|
||||
}
|
||||
None => {
|
||||
create_index(connection, &name, dims)?;
|
||||
Ok(name)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn create_index(connection: &mut ConnectionRef<'_>, name: &str, dims: usize) -> Result<(), Error> {
|
||||
redis::cmd("FT.CREATE")
|
||||
.arg(name)
|
||||
.arg("ON")
|
||||
.arg("HASH")
|
||||
.arg("PREFIX")
|
||||
.arg(1)
|
||||
.arg(name)
|
||||
.arg("SCORE")
|
||||
.arg(1.0)
|
||||
.arg("SCHEMA")
|
||||
.arg("prompt")
|
||||
.arg("TEXT")
|
||||
.arg("WEIGHT")
|
||||
.arg(1)
|
||||
.arg("response")
|
||||
.arg("TEXT")
|
||||
.arg("WEIGHT")
|
||||
.arg(1)
|
||||
.arg("inserted_at")
|
||||
.arg("NUMERIC")
|
||||
.arg("updated_at")
|
||||
.arg("NUMERIC")
|
||||
.arg(VECTOR_FIELD)
|
||||
.arg("VECTOR")
|
||||
.arg("FLAT")
|
||||
.arg(6)
|
||||
.arg("TYPE")
|
||||
.arg("FLOAT32")
|
||||
.arg("DIM")
|
||||
.arg(dims)
|
||||
.arg("DISTANCE_METRIC")
|
||||
.arg("COSINE")
|
||||
.arg(CACHE_KEY_FIELD)
|
||||
.arg("TAG")
|
||||
.arg("SEPARATOR")
|
||||
.arg(",")
|
||||
.query::<()>(connection)
|
||||
.map_err(Error::unavailable)
|
||||
}
|
||||
|
||||
fn index_compatible(
|
||||
connection: &mut ConnectionRef<'_>,
|
||||
name: &str,
|
||||
dims: usize,
|
||||
) -> Result<Option<bool>, Error> {
|
||||
let info = match redis::cmd("FT.INFO")
|
||||
.arg(name)
|
||||
.query::<redis::Value>(connection)
|
||||
{
|
||||
Ok(info) => info,
|
||||
Err(error) if unknown_index(&error) => return Ok(None),
|
||||
Err(error) => return Err(Error::unavailable(error)),
|
||||
};
|
||||
Ok(Some(schema_compatible(&info, dims)))
|
||||
}
|
||||
|
||||
fn unknown_index(error: &redis::RedisError) -> bool {
|
||||
error
|
||||
.code()
|
||||
.is_some_and(|code| code.eq_ignore_ascii_case("unknown"))
|
||||
&& error
|
||||
.detail()
|
||||
.is_some_and(|detail| detail.to_ascii_lowercase().starts_with("index"))
|
||||
}
|
||||
|
||||
struct Attribute {
|
||||
name: Option<String>,
|
||||
field_type: Option<String>,
|
||||
dim: Option<f64>,
|
||||
data_type: Option<String>,
|
||||
distance_metric: Option<String>,
|
||||
}
|
||||
|
||||
fn attribute(value: &redis::Value) -> Option<Attribute> {
|
||||
let redis::Value::Array(pairs) = value else {
|
||||
return None;
|
||||
};
|
||||
let mut attribute = Attribute {
|
||||
name: None,
|
||||
field_type: None,
|
||||
dim: None,
|
||||
data_type: None,
|
||||
distance_metric: None,
|
||||
};
|
||||
for pair in pairs.as_chunks::<2>().0 {
|
||||
match string_value(&pair[0]).as_deref() {
|
||||
Some("identifier") => attribute.name = string_value(&pair[1]),
|
||||
Some("type") => attribute.field_type = string_value(&pair[1]),
|
||||
Some("dim") => attribute.dim = number_value(&pair[1]),
|
||||
Some("data_type") => attribute.data_type = string_value(&pair[1]),
|
||||
Some("distance_metric") => attribute.distance_metric = string_value(&pair[1]),
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
Some(attribute)
|
||||
}
|
||||
|
||||
fn schema_compatible(info: &redis::Value, dims: usize) -> bool {
|
||||
let redis::Value::Array(entries) = info else {
|
||||
return false;
|
||||
};
|
||||
let attributes = entries
|
||||
.as_chunks::<2>()
|
||||
.0
|
||||
.iter()
|
||||
.find(|pair| string_value(&pair[0]).as_deref() == Some("attributes"))
|
||||
.map(|pair| &pair[1]);
|
||||
let Some(redis::Value::Array(attributes)) = attributes else {
|
||||
return false;
|
||||
};
|
||||
let fields = attributes.iter().filter_map(attribute).collect::<Vec<_>>();
|
||||
let has_field = |name: &str, field_type: &str| {
|
||||
fields.iter().any(|field| {
|
||||
field.name.as_deref() == Some(name) && field.field_type.as_deref() == Some(field_type)
|
||||
})
|
||||
};
|
||||
has_field("prompt", "TEXT")
|
||||
&& has_field("response", "TEXT")
|
||||
&& has_field("inserted_at", "NUMERIC")
|
||||
&& has_field("updated_at", "NUMERIC")
|
||||
&& has_field(CACHE_KEY_FIELD, "TAG")
|
||||
&& fields.iter().any(|field| {
|
||||
field.name.as_deref() == Some(VECTOR_FIELD)
|
||||
&& field.field_type.as_deref() == Some("VECTOR")
|
||||
&& field.dim == Some(dims as f64)
|
||||
&& field
|
||||
.data_type
|
||||
.as_deref()
|
||||
.is_some_and(|data| data.eq_ignore_ascii_case("float32"))
|
||||
&& field
|
||||
.distance_metric
|
||||
.as_deref()
|
||||
.is_some_and(|metric| metric.eq_ignore_ascii_case("cosine"))
|
||||
})
|
||||
}
|
||||
|
|
@ -1,7 +0,0 @@
|
|||
mod cache;
|
||||
mod config;
|
||||
mod index;
|
||||
mod reply;
|
||||
|
||||
pub use cache::RedisSemanticCache;
|
||||
pub use config::{DEFAULT_INDEX_NAME, RedisSemanticConfig};
|
||||
|
|
@ -1,57 +0,0 @@
|
|||
pub(crate) fn string_value(value: &redis::Value) -> Option<String> {
|
||||
match value {
|
||||
redis::Value::BulkString(bytes) => String::from_utf8(bytes.clone()).ok(),
|
||||
redis::Value::SimpleString(text) => Some(text.clone()),
|
||||
redis::Value::VerbatimString { text, .. } => Some(text.clone()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn number_value(value: &redis::Value) -> Option<f64> {
|
||||
match value {
|
||||
redis::Value::Int(number) => Some(*number as f64),
|
||||
redis::Value::Double(number) => Some(*number),
|
||||
_ => string_value(value).and_then(|text| text.parse().ok()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn first_document(result: &redis::Value) -> Option<&[redis::Value]> {
|
||||
let redis::Value::Array(items) = result else {
|
||||
return None;
|
||||
};
|
||||
let [count, _document_id, fields, ..] = items.as_slice() else {
|
||||
return None;
|
||||
};
|
||||
if !matches!(count, redis::Value::Int(count) if *count > 0) {
|
||||
return None;
|
||||
}
|
||||
match fields {
|
||||
redis::Value::Array(fields) => Some(fields.as_slice()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn field_value<'a>(fields: &'a [redis::Value], name: &str) -> Option<&'a redis::Value> {
|
||||
fields
|
||||
.as_chunks::<2>()
|
||||
.0
|
||||
.iter()
|
||||
.find(|pair| string_value(&pair[0]).as_deref() == Some(name))
|
||||
.map(|pair| &pair[1])
|
||||
}
|
||||
|
||||
pub(crate) fn string_field(fields: &[redis::Value], name: &str) -> Option<String> {
|
||||
field_value(fields, name).and_then(string_value)
|
||||
}
|
||||
|
||||
pub(crate) fn number_field(fields: &[redis::Value], name: &str) -> Option<f64> {
|
||||
field_value(fields, name).and_then(number_value)
|
||||
}
|
||||
|
||||
pub(crate) fn bytes_field(fields: &[redis::Value], name: &str) -> Option<Vec<u8>> {
|
||||
match field_value(fields, name)? {
|
||||
redis::Value::BulkString(bytes) => Some(bytes.clone()),
|
||||
redis::Value::SimpleString(text) => Some(text.clone().into_bytes()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
|
@ -1,778 +0,0 @@
|
|||
mod support;
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use litellm_cache::{
|
||||
BaseCache, Error, JsonCodec, SemanticCacheContext,
|
||||
semantic::{SemanticCache, SemanticLookup},
|
||||
};
|
||||
use litellm_cache_redis_semantic::{DEFAULT_INDEX_NAME, RedisSemanticCache, RedisSemanticConfig};
|
||||
use redis_test::{MockCmd, MockRedisConnection};
|
||||
use rstest::{fixture, rstest};
|
||||
use serde_json::{Value, json};
|
||||
use sha2::{Digest, Sha256};
|
||||
use support::FakeEmbedder;
|
||||
|
||||
const INDEX: &str = DEFAULT_INDEX_NAME;
|
||||
const PROMPT: &str = "hello prompt";
|
||||
const CLOCK: fn() -> f64 = || 1700000000.5;
|
||||
const VECTOR: [f32; 3] = [0.1, 0.2, 0.3];
|
||||
|
||||
type MockCache = RedisSemanticCache<FakeEmbedder, JsonCodec<Value>, MockRedisConnection>;
|
||||
|
||||
#[fixture]
|
||||
fn config() -> RedisSemanticConfig {
|
||||
RedisSemanticConfig {
|
||||
index_name: INDEX.into(),
|
||||
similarity_threshold: 0.9,
|
||||
}
|
||||
}
|
||||
|
||||
#[fixture]
|
||||
fn entry() -> Value {
|
||||
json!({"timestamp": 1.0, "response": {"answer": "yes"}})
|
||||
}
|
||||
|
||||
#[fixture]
|
||||
fn context() -> SemanticCacheContext {
|
||||
messages_context(vec![json!({"role": "user", "content": PROMPT})])
|
||||
}
|
||||
|
||||
fn messages_context(messages: Vec<Value>) -> SemanticCacheContext {
|
||||
SemanticCacheContext {
|
||||
messages: Some(Value::Array(messages)),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
fn cache(commands: Vec<MockCmd>, embedder: FakeEmbedder) -> MockCache {
|
||||
RedisSemanticCache::with_connection(
|
||||
MockRedisConnection::new(commands).assert_all_commands_consumed(),
|
||||
embedder,
|
||||
JsonCodec::new(),
|
||||
config(),
|
||||
)
|
||||
.with_clock(CLOCK)
|
||||
}
|
||||
|
||||
fn encoded(value: &Value) -> Vec<u8> {
|
||||
serde_json::to_vec(value).unwrap()
|
||||
}
|
||||
|
||||
fn vector_bytes(vector: &[f32]) -> Vec<u8> {
|
||||
vector
|
||||
.iter()
|
||||
.flat_map(|component| component.to_le_bytes())
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn entry_id(prompt: &str, tag: &str) -> String {
|
||||
let mut digest = Sha256::new();
|
||||
digest.update(prompt.as_bytes());
|
||||
digest.update(b"litellm_cache_key");
|
||||
digest.update(tag.as_bytes());
|
||||
format!("{:x}", digest.finalize())
|
||||
}
|
||||
|
||||
fn s(value: &str) -> redis::Value {
|
||||
redis::Value::BulkString(value.as_bytes().to_vec())
|
||||
}
|
||||
|
||||
fn unknown_index_error() -> redis::RedisError {
|
||||
support::server_error("Unknown index name")
|
||||
}
|
||||
|
||||
fn info_missing(index: &str) -> MockCmd {
|
||||
MockCmd::new(
|
||||
redis::cmd("FT.INFO").arg(index),
|
||||
Err::<redis::Value, _>(unknown_index_error()),
|
||||
)
|
||||
}
|
||||
|
||||
fn info(index: &str, value: redis::Value) -> MockCmd {
|
||||
MockCmd::new(redis::cmd("FT.INFO").arg(index), Ok(value))
|
||||
}
|
||||
|
||||
fn attribute(name: &str, field_type: &str, extra: Vec<redis::Value>) -> redis::Value {
|
||||
let mut parts = vec![
|
||||
s("identifier"),
|
||||
s(name),
|
||||
s("attribute"),
|
||||
s(name),
|
||||
s("type"),
|
||||
s(field_type),
|
||||
];
|
||||
parts.extend(extra);
|
||||
redis::Value::Array(parts)
|
||||
}
|
||||
|
||||
fn index_info(attributes: Vec<redis::Value>) -> redis::Value {
|
||||
redis::Value::Array(vec![
|
||||
s("index_name"),
|
||||
s(INDEX),
|
||||
s("attributes"),
|
||||
redis::Value::Array(attributes),
|
||||
])
|
||||
}
|
||||
|
||||
fn vector_attribute_with(dims: i64, data_type: &str, distance_metric: &str) -> redis::Value {
|
||||
attribute(
|
||||
"prompt_vector",
|
||||
"VECTOR",
|
||||
vec![
|
||||
s("algorithm"),
|
||||
s("FLAT"),
|
||||
s("data_type"),
|
||||
s(data_type),
|
||||
s("dim"),
|
||||
redis::Value::Int(dims),
|
||||
s("distance_metric"),
|
||||
s(distance_metric),
|
||||
],
|
||||
)
|
||||
}
|
||||
|
||||
fn info_with_vector(vector: redis::Value) -> redis::Value {
|
||||
index_info(vec![
|
||||
attribute("prompt", "TEXT", vec![]),
|
||||
attribute("response", "TEXT", vec![]),
|
||||
attribute("inserted_at", "NUMERIC", vec![]),
|
||||
attribute("updated_at", "NUMERIC", vec![]),
|
||||
vector,
|
||||
attribute("litellm_cache_key", "TAG", vec![]),
|
||||
])
|
||||
}
|
||||
|
||||
fn compatible_info(dims: i64) -> redis::Value {
|
||||
info_with_vector(vector_attribute_with(dims, "FLOAT32", "COSINE"))
|
||||
}
|
||||
|
||||
fn unscoped_info(dims: i64) -> redis::Value {
|
||||
index_info(vec![
|
||||
attribute("prompt", "TEXT", vec![]),
|
||||
attribute("response", "TEXT", vec![]),
|
||||
attribute("inserted_at", "NUMERIC", vec![]),
|
||||
attribute("updated_at", "NUMERIC", vec![]),
|
||||
vector_attribute_with(dims, "FLOAT32", "COSINE"),
|
||||
])
|
||||
}
|
||||
|
||||
fn create_index(name: &str, dims: usize) -> MockCmd {
|
||||
MockCmd::new(create_index_command(name, dims), Ok("OK"))
|
||||
}
|
||||
|
||||
fn create_index_command(name: &str, dims: usize) -> redis::Cmd {
|
||||
let mut command = redis::cmd("FT.CREATE");
|
||||
command
|
||||
.arg(name)
|
||||
.arg("ON")
|
||||
.arg("HASH")
|
||||
.arg("PREFIX")
|
||||
.arg(1)
|
||||
.arg(name)
|
||||
.arg("SCORE")
|
||||
.arg(1.0)
|
||||
.arg("SCHEMA")
|
||||
.arg("prompt")
|
||||
.arg("TEXT")
|
||||
.arg("WEIGHT")
|
||||
.arg(1)
|
||||
.arg("response")
|
||||
.arg("TEXT")
|
||||
.arg("WEIGHT")
|
||||
.arg(1)
|
||||
.arg("inserted_at")
|
||||
.arg("NUMERIC")
|
||||
.arg("updated_at")
|
||||
.arg("NUMERIC")
|
||||
.arg("prompt_vector")
|
||||
.arg("VECTOR")
|
||||
.arg("FLAT")
|
||||
.arg(6)
|
||||
.arg("TYPE")
|
||||
.arg("FLOAT32")
|
||||
.arg("DIM")
|
||||
.arg(dims)
|
||||
.arg("DISTANCE_METRIC")
|
||||
.arg("COSINE")
|
||||
.arg("litellm_cache_key")
|
||||
.arg("TAG")
|
||||
.arg("SEPARATOR")
|
||||
.arg(",");
|
||||
command
|
||||
}
|
||||
|
||||
fn hset(index: &str, prompt: &str, tag: &str, vector: &[f32], value: &Value) -> MockCmd {
|
||||
MockCmd::new(
|
||||
redis::cmd("HSET")
|
||||
.arg(format!("{index}:{}", entry_id(prompt, tag)))
|
||||
.arg("entry_id")
|
||||
.arg(entry_id(prompt, tag))
|
||||
.arg("prompt")
|
||||
.arg(prompt)
|
||||
.arg("response")
|
||||
.arg(encoded(value))
|
||||
.arg("prompt_vector")
|
||||
.arg(vector_bytes(vector))
|
||||
.arg("inserted_at")
|
||||
.arg("1700000000.5")
|
||||
.arg("updated_at")
|
||||
.arg("1700000000.5")
|
||||
.arg("litellm_cache_key")
|
||||
.arg(tag),
|
||||
Ok(7),
|
||||
)
|
||||
}
|
||||
|
||||
fn search(
|
||||
index: &str,
|
||||
tag: &str,
|
||||
vector: &[f32],
|
||||
reply: redis::RedisResult<redis::Value>,
|
||||
) -> MockCmd {
|
||||
let mut command = redis::cmd("FT.SEARCH");
|
||||
command
|
||||
.arg(index)
|
||||
.arg(format!(
|
||||
"(@litellm_cache_key:{{{tag}}})=>[KNN 1 @prompt_vector $vector AS vector_distance]"
|
||||
))
|
||||
.arg("RETURN")
|
||||
.arg(8)
|
||||
.arg("entry_id")
|
||||
.arg("prompt")
|
||||
.arg("response")
|
||||
.arg("inserted_at")
|
||||
.arg("updated_at")
|
||||
.arg("metadata")
|
||||
.arg("litellm_cache_key")
|
||||
.arg("vector_distance")
|
||||
.arg("SORTBY")
|
||||
.arg("vector_distance")
|
||||
.arg("ASC")
|
||||
.arg("DIALECT")
|
||||
.arg(2)
|
||||
.arg("LIMIT")
|
||||
.arg(0)
|
||||
.arg(1)
|
||||
.arg("PARAMS")
|
||||
.arg(2)
|
||||
.arg("vector")
|
||||
.arg(vector_bytes(vector));
|
||||
MockCmd::new(command, reply)
|
||||
}
|
||||
|
||||
fn hit(tag: &str, distance: &str, response: Vec<u8>) -> redis::Value {
|
||||
redis::Value::Array(vec![
|
||||
redis::Value::Int(1),
|
||||
s("litellm_semantic_cache_index:stored-id"),
|
||||
redis::Value::Array(vec![
|
||||
s("entry_id"),
|
||||
s("stored-id"),
|
||||
s("prompt"),
|
||||
s(PROMPT),
|
||||
s("response"),
|
||||
redis::Value::BulkString(response),
|
||||
s("inserted_at"),
|
||||
s("1700000000.5"),
|
||||
s("updated_at"),
|
||||
s("1700000000.5"),
|
||||
s("litellm_cache_key"),
|
||||
s(tag),
|
||||
s("vector_distance"),
|
||||
s(distance),
|
||||
]),
|
||||
])
|
||||
}
|
||||
|
||||
fn empty_result() -> redis::Value {
|
||||
redis::Value::Array(vec![redis::Value::Int(0)])
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::creates_index_and_expires(false, Some(Duration::from_secs(5)), Some(5))]
|
||||
#[case::existing_index_without_ttl(true, None, None)]
|
||||
#[case::fractional_ttl_rounds_up(true, Some(Duration::from_millis(1500)), Some(2))]
|
||||
fn store_writes_the_redisvl_hash(
|
||||
#[case] index_exists: bool,
|
||||
#[case] ttl: Option<Duration>,
|
||||
#[case] expire: Option<u64>,
|
||||
entry: Value,
|
||||
context: SemanticCacheContext,
|
||||
) {
|
||||
let hash_key = format!("{INDEX}:{}", entry_id(PROMPT, "key1"));
|
||||
let mut commands = if index_exists {
|
||||
vec![info(INDEX, compatible_info(3))]
|
||||
} else {
|
||||
vec![info_missing(INDEX), create_index(INDEX, 3)]
|
||||
};
|
||||
commands.push(hset(INDEX, PROMPT, "key1", &VECTOR, &entry));
|
||||
commands.extend(
|
||||
expire.map(|seconds| MockCmd::new(redis::cmd("EXPIRE").arg(&hash_key).arg(seconds), Ok(1))),
|
||||
);
|
||||
let cache = cache(commands, FakeEmbedder::new(&[]));
|
||||
|
||||
cache
|
||||
.set_cache("key1", entry, &SemanticCacheContext { ttl, ..context })
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::below_distance_threshold("key1", "0.05", None, Ok(Some(entry())))]
|
||||
#[case::above_distance_threshold("key1", "0.5", None, Ok(None))]
|
||||
#[case::other_cache_key("other", "0.05", None, Ok(None))]
|
||||
#[case::malformed_response("key1", "0.05", Some(b"not json!".as_slice()), Err(Error::InvalidEntry(None)))]
|
||||
fn lookup_applies_threshold_scope_and_codec(
|
||||
#[case] stored_tag: &str,
|
||||
#[case] distance: &str,
|
||||
#[case] response: Option<&[u8]>,
|
||||
#[case] expected: Result<Option<Value>, Error>,
|
||||
entry: Value,
|
||||
context: SemanticCacheContext,
|
||||
) {
|
||||
let response = response.map_or_else(|| encoded(&entry), <[u8]>::to_vec);
|
||||
let cache = cache(
|
||||
vec![
|
||||
info(INDEX, compatible_info(3)),
|
||||
search(
|
||||
INDEX,
|
||||
"key1",
|
||||
&VECTOR,
|
||||
Ok(hit(stored_tag, distance, response)),
|
||||
),
|
||||
],
|
||||
FakeEmbedder::new(&[]),
|
||||
);
|
||||
|
||||
let kind = |result: Result<Option<Value>, Error>| {
|
||||
result.map_err(|error| std::mem::discriminant(&error))
|
||||
};
|
||||
assert_eq!(kind(cache.get_cache("key1", &context)), kind(expected));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::hit(context(), Some(hit("key1", "0.05", encoded(&entry()))), Some(entry()), Some(1.0 - 0.05))]
|
||||
#[case::beyond_distance_threshold(context(), Some(hit("key1", "0.5", encoded(&entry()))), None, Some(0.0))]
|
||||
#[case::no_results(context(), Some(empty_result()), None, Some(0.0))]
|
||||
#[case::other_cache_key(context(), Some(hit("other", "0.05", encoded(&entry()))), None, Some(0.0))]
|
||||
#[case::no_prompt(SemanticCacheContext::default(), None, None, Some(0.0))]
|
||||
#[tokio::test]
|
||||
async fn lookup_reports_python_semantic_similarity(
|
||||
#[case] context: SemanticCacheContext,
|
||||
#[case] reply: Option<redis::Value>,
|
||||
#[case] value: Option<Value>,
|
||||
#[case] similarity: Option<f64>,
|
||||
#[values(false, true)] use_async: bool,
|
||||
) {
|
||||
let commands = reply.map_or_else(Vec::new, |reply| {
|
||||
vec![
|
||||
info(INDEX, compatible_info(3)),
|
||||
search(INDEX, "key1", &VECTOR, Ok(reply)),
|
||||
]
|
||||
});
|
||||
let cache = cache(commands, FakeEmbedder::new(&[]));
|
||||
|
||||
let lookup = if use_async {
|
||||
cache
|
||||
.async_get_cache_with_similarity("key1", &context)
|
||||
.await
|
||||
} else {
|
||||
cache.get_cache_with_similarity("key1", &context)
|
||||
};
|
||||
|
||||
assert_eq!(lookup, Ok(SemanticLookup { value, similarity }));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn missing_prompt_is_a_noop_that_never_embeds(entry: Value) {
|
||||
let embedder = FakeEmbedder::new(&[]);
|
||||
let calls = embedder.calls.clone();
|
||||
let cache = cache(Vec::new(), embedder);
|
||||
let context = SemanticCacheContext::default();
|
||||
|
||||
cache.set_cache("key1", entry.clone(), &context).unwrap();
|
||||
assert_eq!(cache.get_cache("key1", &context).unwrap(), None);
|
||||
cache
|
||||
.async_set_cache("key1", entry, context.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(cache.async_get_cache("key1", &context).await.unwrap(), None);
|
||||
assert!(calls.lock().unwrap().is_empty());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn scope_overrides_key_as_filter_tag(entry: Value, context: SemanticCacheContext) {
|
||||
let cache = cache(
|
||||
vec![
|
||||
info(INDEX, compatible_info(3)),
|
||||
hset(INDEX, PROMPT, "scope-a", &VECTOR, &entry),
|
||||
search(
|
||||
INDEX,
|
||||
"scope\\-a",
|
||||
&VECTOR,
|
||||
Ok(hit("scope-a", "0.05", encoded(&entry))),
|
||||
),
|
||||
],
|
||||
FakeEmbedder::new(&[]),
|
||||
);
|
||||
let context = SemanticCacheContext {
|
||||
scope: Some("scope-a".into()),
|
||||
..context
|
||||
};
|
||||
|
||||
cache.set_cache("key1", entry.clone(), &context).unwrap();
|
||||
assert_eq!(cache.get_cache("key1", &context).unwrap(), Some(entry));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::unscoped_schema(unscoped_info(3))]
|
||||
#[case::wrong_distance_metric(info_with_vector(vector_attribute_with(3, "FLOAT32", "L2")))]
|
||||
#[case::wrong_data_type(info_with_vector(vector_attribute_with(3, "FLOAT64", "COSINE")))]
|
||||
fn incompatible_schema_falls_back_to_isolated_index(
|
||||
#[case] base_info: redis::Value,
|
||||
entry: Value,
|
||||
context: SemanticCacheContext,
|
||||
) {
|
||||
let isolated = format!("{INDEX}_isolated");
|
||||
let cache = cache(
|
||||
vec![
|
||||
info(INDEX, base_info),
|
||||
info_missing(&isolated),
|
||||
create_index(&isolated, 3),
|
||||
hset(&isolated, PROMPT, "key1", &VECTOR, &entry),
|
||||
],
|
||||
FakeEmbedder::new(&[]),
|
||||
);
|
||||
|
||||
cache.set_cache("key1", entry, &context).unwrap();
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn create_index_race_rechecks_schema_and_stores(entry: Value, context: SemanticCacheContext) {
|
||||
let cache = cache(
|
||||
vec![
|
||||
info_missing(INDEX),
|
||||
MockCmd::new(
|
||||
create_index_command(INDEX, 3),
|
||||
Err::<&str, _>(support::server_error("Index already exists")),
|
||||
),
|
||||
info(INDEX, compatible_info(3)),
|
||||
hset(INDEX, PROMPT, "key1", &VECTOR, &entry),
|
||||
],
|
||||
FakeEmbedder::new(&[]),
|
||||
);
|
||||
|
||||
cache.set_cache("key1", entry, &context).unwrap();
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::punctuation_and_spaces("a:b, c|d", "a\\:b\\,\\ c\\|d")]
|
||||
#[case::braces_and_dots("{x}.y", "\\{x\\}\\.y")]
|
||||
#[case::plain("key1", "key1")]
|
||||
fn tag_special_characters_are_escaped_in_search_filter(
|
||||
#[case] tag: &str,
|
||||
#[case] escaped: &str,
|
||||
context: SemanticCacheContext,
|
||||
) {
|
||||
let cache = cache(
|
||||
vec![
|
||||
info(INDEX, compatible_info(3)),
|
||||
search(INDEX, escaped, &VECTOR, Ok(empty_result())),
|
||||
],
|
||||
FakeEmbedder::new(&[]),
|
||||
);
|
||||
|
||||
assert_eq!(cache.get_cache(tag, &context).unwrap(), None);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::content_parts(
|
||||
SemanticCacheContext {
|
||||
messages: Some(json!([
|
||||
{"role": "user", "content": [{"type": "text", "text": "first"}, {"type": "text", "text": "second"}]},
|
||||
{"role": "assistant", "content": "reply"},
|
||||
])),
|
||||
..Default::default()
|
||||
},
|
||||
"firstsecondreply"
|
||||
)]
|
||||
#[case::responses_string_input(
|
||||
SemanticCacheContext { input: Some(json!(" plain input ")), ..Default::default() },
|
||||
"plain input"
|
||||
)]
|
||||
#[case::responses_nested_input(
|
||||
SemanticCacheContext {
|
||||
input: Some(json!([{"content": [{"type": "input_text", "text": "nested"}]}, "tail"])),
|
||||
..Default::default()
|
||||
},
|
||||
"nested\ntail"
|
||||
)]
|
||||
#[case::responses_output_text(
|
||||
SemanticCacheContext { input: Some(json!({"output_text": " result text "})), ..Default::default() },
|
||||
"result text"
|
||||
)]
|
||||
#[case::search_results(
|
||||
messages_context(vec![json!({
|
||||
"role": "user",
|
||||
"content": "question",
|
||||
"search_results": [{"source": "src", "title": "t", "content": [{"text": "found"}], "citations": {"a": 1}}],
|
||||
})]),
|
||||
"questionsrctfound{\"a\":1}"
|
||||
)]
|
||||
#[case::empty_messages_fall_back_to_input(
|
||||
SemanticCacheContext { messages: Some(json!([])), input: Some(json!("fallback")), ..Default::default() },
|
||||
"fallback"
|
||||
)]
|
||||
fn prompt_extraction_matches_python_message_and_input_shapes(
|
||||
#[case] context: SemanticCacheContext,
|
||||
#[case] prompt: &str,
|
||||
) {
|
||||
let embedder = FakeEmbedder::new(&[]);
|
||||
let calls = embedder.calls.clone();
|
||||
let cache = cache(
|
||||
vec![
|
||||
info(INDEX, compatible_info(3)),
|
||||
search(INDEX, "key1", &VECTOR, Ok(empty_result())),
|
||||
],
|
||||
embedder,
|
||||
);
|
||||
|
||||
cache.get_cache("key1", &context).unwrap();
|
||||
|
||||
assert_eq!(*calls.lock().unwrap(), vec![(prompt.to_owned(), None)]);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case(None)]
|
||||
#[case(Some(Duration::from_secs(9)))]
|
||||
fn ttl_passes_through_context_only(#[case] ttl: Option<Duration>) {
|
||||
let cache = cache(Vec::new(), FakeEmbedder::new(&[]));
|
||||
|
||||
assert_eq!(
|
||||
cache.get_ttl(&SemanticCacheContext {
|
||||
ttl,
|
||||
..Default::default()
|
||||
}),
|
||||
ttl
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn async_paths_embed_with_metadata_then_run_blocking_redis_work(
|
||||
entry: Value,
|
||||
context: SemanticCacheContext,
|
||||
) {
|
||||
let embedder = FakeEmbedder::new(&[]);
|
||||
let calls = embedder.calls.clone();
|
||||
let cache = cache(
|
||||
vec![
|
||||
info(INDEX, compatible_info(3)),
|
||||
hset(INDEX, PROMPT, "key1", &VECTOR, &entry),
|
||||
search(
|
||||
INDEX,
|
||||
"key1",
|
||||
&VECTOR,
|
||||
Ok(hit("key1", "0.05", encoded(&entry))),
|
||||
),
|
||||
],
|
||||
embedder,
|
||||
);
|
||||
let context = SemanticCacheContext {
|
||||
metadata: Some(json!({"tenant": "team"})),
|
||||
..context
|
||||
};
|
||||
|
||||
cache
|
||||
.async_set_cache("key1", entry.clone(), context.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
cache.async_get_cache("key1", &context).await.unwrap(),
|
||||
Some(entry)
|
||||
);
|
||||
assert_eq!(
|
||||
*calls.lock().unwrap(),
|
||||
vec![(PROMPT.to_owned(), context.metadata.clone()); 2]
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn accessors_report_the_config(config: RedisSemanticConfig) {
|
||||
let cache = cache(Vec::new(), FakeEmbedder::new(&[]));
|
||||
|
||||
assert_eq!(cache.index_name(), config.index_name);
|
||||
assert!((cache.similarity_threshold() - config.similarity_threshold).abs() < 1e-6);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn shared_base_index_across_dimensions_replaces_the_isolated_index(entry: Value) {
|
||||
// Pins parity with Python's `_isolated` + overwrite=True flow.
|
||||
let prompt = "shared prompt";
|
||||
let isolated = format!("{INDEX}_isolated");
|
||||
let context = || messages_context(vec![json!({"role": "user", "content": prompt})]);
|
||||
|
||||
let vector_a = vec![0.1f32; 8];
|
||||
let worker_a = cache(
|
||||
vec![
|
||||
info_missing(INDEX),
|
||||
create_index(INDEX, 8),
|
||||
hset(INDEX, prompt, "key1", &vector_a, &entry),
|
||||
],
|
||||
FakeEmbedder::new(&[(prompt, &vector_a)]),
|
||||
);
|
||||
worker_a
|
||||
.set_cache("key1", entry.clone(), &context())
|
||||
.unwrap();
|
||||
|
||||
let vector_b = vec![0.2f32; 4];
|
||||
let worker_b = cache(
|
||||
vec![
|
||||
info(INDEX, compatible_info(8)),
|
||||
info_missing(&isolated),
|
||||
create_index(&isolated, 4),
|
||||
hset(&isolated, prompt, "key1", &vector_b, &entry),
|
||||
search(
|
||||
&isolated,
|
||||
"key1",
|
||||
&vector_b,
|
||||
Ok(hit("key1", "0.0", encoded(&entry))),
|
||||
),
|
||||
search(
|
||||
&isolated,
|
||||
"key1",
|
||||
&vector_b,
|
||||
Err(redis::RedisError::from((
|
||||
redis::ErrorKind::Extension,
|
||||
"Vector dimension mismatch",
|
||||
))),
|
||||
),
|
||||
],
|
||||
FakeEmbedder::new(&[(prompt, &vector_b)]),
|
||||
);
|
||||
worker_b
|
||||
.set_cache("key1", entry.clone(), &context())
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
worker_b.get_cache("key1", &context()).unwrap(),
|
||||
Some(entry.clone())
|
||||
);
|
||||
|
||||
let vector_c = vec![0.3f32; 16];
|
||||
let worker_c = cache(
|
||||
vec![
|
||||
info(INDEX, compatible_info(8)),
|
||||
info(&isolated, compatible_info(4)),
|
||||
MockCmd::new(redis::cmd("FT.DROPINDEX").arg(&isolated), Ok("OK")),
|
||||
create_index(&isolated, 16),
|
||||
hset(&isolated, prompt, "key1", &vector_c, &entry),
|
||||
],
|
||||
FakeEmbedder::new(&[(prompt, &vector_c)]),
|
||||
);
|
||||
worker_c
|
||||
.set_cache("key1", entry.clone(), &context())
|
||||
.unwrap();
|
||||
|
||||
assert!(matches!(
|
||||
worker_b.get_cache("key1", &context()),
|
||||
Err(Error::Unavailable(Some(_)))
|
||||
));
|
||||
}
|
||||
|
||||
#[fixture]
|
||||
fn redis_stack_url() -> Option<String> {
|
||||
std::env::var("LITELLM_REDIS_STACK_URL").ok()
|
||||
}
|
||||
|
||||
fn live_cache(
|
||||
url: &str,
|
||||
index_name: &str,
|
||||
prompt: &str,
|
||||
vector: Vec<f32>,
|
||||
) -> RedisSemanticCache<FakeEmbedder, JsonCodec<Value>> {
|
||||
RedisSemanticCache::new(
|
||||
url,
|
||||
FakeEmbedder::new(&[(prompt, vector.as_slice())]),
|
||||
JsonCodec::<Value>::new(),
|
||||
RedisSemanticConfig {
|
||||
index_name: index_name.to_owned(),
|
||||
similarity_threshold: 0.9,
|
||||
},
|
||||
)
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn live_shared_index_is_replaced_across_dimensions(redis_stack_url: Option<String>, entry: Value) {
|
||||
let Some(url) = redis_stack_url else {
|
||||
return;
|
||||
};
|
||||
// Pins parity with Python's `_isolated` + overwrite=True flow.
|
||||
let base = format!("rust_semantic_shared_{}", std::process::id());
|
||||
let isolated = format!("{base}_isolated");
|
||||
let prompt = "shared live prompt";
|
||||
let context = || messages_context(vec![json!({"role": "user", "content": prompt})]);
|
||||
|
||||
let worker_a = live_cache(&url, &base, prompt, vec![0.1f32; 8]);
|
||||
worker_a
|
||||
.set_cache("key1", entry.clone(), &context())
|
||||
.unwrap();
|
||||
|
||||
let worker_b = live_cache(&url, &base, prompt, vec![0.2f32; 4]);
|
||||
worker_b
|
||||
.set_cache("key1", entry.clone(), &context())
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
worker_b.get_cache("key1", &context()).unwrap(),
|
||||
Some(entry.clone())
|
||||
);
|
||||
|
||||
let worker_c = live_cache(&url, &base, prompt, vec![0.3f32; 16]);
|
||||
worker_c
|
||||
.set_cache("key1", entry.clone(), &context())
|
||||
.unwrap();
|
||||
|
||||
assert!(matches!(
|
||||
worker_b.get_cache("key1", &context()),
|
||||
Err(Error::Unavailable(Some(_)))
|
||||
));
|
||||
|
||||
let mut connection = redis::Client::open(url).unwrap().get_connection().unwrap();
|
||||
for index in [&base, &isolated] {
|
||||
let _: Result<(), _> = redis::cmd("FT.DROPINDEX")
|
||||
.arg(index)
|
||||
.arg("DD")
|
||||
.query(&mut connection);
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn live_store_lookup_and_ttl_against_redis_stack(redis_stack_url: Option<String>, entry: Value) {
|
||||
let Some(url) = redis_stack_url else {
|
||||
return;
|
||||
};
|
||||
let prompt = "rust semantic cache live prompt";
|
||||
let index_name = format!("rust_semantic_test_{}", std::process::id());
|
||||
let cache = live_cache(&url, &index_name, prompt, vec![0.1, 0.2, 0.3, 0.4]);
|
||||
let context = SemanticCacheContext {
|
||||
ttl: Some(Duration::from_secs(120)),
|
||||
..messages_context(vec![json!({"role": "user", "content": prompt})])
|
||||
};
|
||||
|
||||
cache
|
||||
.set_cache("live-key", entry.clone(), &context)
|
||||
.unwrap();
|
||||
assert_eq!(cache.get_cache("live-key", &context).unwrap(), Some(entry));
|
||||
assert_eq!(cache.get_cache("other-key", &context).unwrap(), None);
|
||||
|
||||
let mut connection = redis::Client::open(url).unwrap().get_connection().unwrap();
|
||||
let ttl: i64 = redis::Commands::ttl(
|
||||
&mut connection,
|
||||
format!("{index_name}:{}", entry_id(prompt, "live-key")),
|
||||
)
|
||||
.unwrap();
|
||||
assert!(
|
||||
ttl > 0,
|
||||
"expected stored hash to carry an expiry, got {ttl}"
|
||||
);
|
||||
}
|
||||
|
|
@ -1,63 +0,0 @@
|
|||
mod support;
|
||||
|
||||
use litellm_cache::{JsonCodec, SemanticCacheContext, semantic::PreparedEmbedding};
|
||||
use litellm_cache_redis_semantic::{DEFAULT_INDEX_NAME, RedisSemanticCache, RedisSemanticConfig};
|
||||
use litellm_cache_testing as contract;
|
||||
use rstest::{fixture, rstest};
|
||||
use serde_json::{Value, json};
|
||||
use support::FakeSearch;
|
||||
|
||||
type Cache = RedisSemanticCache<PreparedEmbedding, JsonCodec<Value>, FakeSearch>;
|
||||
|
||||
const PREFIX: &str = "contract:";
|
||||
|
||||
#[fixture]
|
||||
fn cache() -> Cache {
|
||||
RedisSemanticCache::with_connection(
|
||||
FakeSearch::default(),
|
||||
PreparedEmbedding(vec![0.6, 0.8]),
|
||||
JsonCodec::new(),
|
||||
RedisSemanticConfig {
|
||||
index_name: DEFAULT_INDEX_NAME.into(),
|
||||
similarity_threshold: 0.9,
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
#[fixture]
|
||||
fn context() -> SemanticCacheContext {
|
||||
SemanticCacheContext {
|
||||
messages: Some(json!([{"role": "user", "content": "contract prompt"}])),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn hit_and_miss(cache: Cache, context: SemanticCacheContext) {
|
||||
contract::hit_and_miss(&cache, context, PREFIX, json!({"answer": 42})).await;
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn sync_async_equivalence(cache: Cache, context: SemanticCacheContext) {
|
||||
contract::sync_async_equivalence(&cache, context, PREFIX, json!("first"), json!([2])).await;
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn overwrite_replaces(cache: Cache, context: SemanticCacheContext) {
|
||||
contract::overwrite_replaces(&cache, context, PREFIX, json!(1), json!({"b": 2})).await;
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn pipeline_writes_every_entry(cache: Cache, context: SemanticCacheContext) {
|
||||
contract::pipeline_writes_every_entry(
|
||||
&cache,
|
||||
context,
|
||||
PREFIX,
|
||||
vec![json!("a"), json!(2), json!({"c": true})],
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
|
@ -1,302 +0,0 @@
|
|||
#![allow(dead_code)]
|
||||
|
||||
use std::{
|
||||
collections::{BTreeMap, HashMap},
|
||||
sync::{Arc, Mutex},
|
||||
};
|
||||
|
||||
use litellm_cache::{Error, semantic::Embedder};
|
||||
use serde_json::Value;
|
||||
|
||||
pub type EmbedCalls = Arc<Mutex<Vec<(String, Option<Value>)>>>;
|
||||
|
||||
/// Embeds known prompts to fixed vectors, anything else to `[0.1, 0.2, 0.3]`, and records every
|
||||
/// prompt with its metadata.
|
||||
pub struct FakeEmbedder {
|
||||
vectors: HashMap<String, Vec<f32>>,
|
||||
pub calls: EmbedCalls,
|
||||
}
|
||||
|
||||
impl FakeEmbedder {
|
||||
pub fn new(vectors: &[(&str, &[f32])]) -> Self {
|
||||
Self {
|
||||
vectors: vectors
|
||||
.iter()
|
||||
.map(|(prompt, vector)| ((*prompt).to_owned(), vector.to_vec()))
|
||||
.collect(),
|
||||
calls: EmbedCalls::default(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Embedder for FakeEmbedder {
|
||||
fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result<Vec<f32>, Error> {
|
||||
self.calls
|
||||
.lock()
|
||||
.unwrap()
|
||||
.push((prompt.to_owned(), metadata.cloned()));
|
||||
Ok(self
|
||||
.vectors
|
||||
.get(prompt)
|
||||
.cloned()
|
||||
.unwrap_or_else(|| vec![0.1, 0.2, 0.3]))
|
||||
}
|
||||
|
||||
async fn async_embed(&self, prompt: &str, metadata: Option<&Value>) -> Result<Vec<f32>, Error> {
|
||||
self.embed(prompt, metadata)
|
||||
}
|
||||
}
|
||||
|
||||
struct FakeIndex {
|
||||
prefix: Vec<u8>,
|
||||
dims: usize,
|
||||
vector_field: String,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct SearchState {
|
||||
indexes: HashMap<String, FakeIndex>,
|
||||
hashes: BTreeMap<Vec<u8>, BTreeMap<String, Vec<u8>>>,
|
||||
}
|
||||
|
||||
/// An in-memory Redis Stack speaking the `FT.*`, `HSET` and `EXPIRE` subset the semantic cache
|
||||
/// sends, with exact cosine KNN over the hashes under an index prefix.
|
||||
#[derive(Clone, Default)]
|
||||
pub struct FakeSearch {
|
||||
state: Arc<Mutex<SearchState>>,
|
||||
}
|
||||
|
||||
impl FakeSearch {
|
||||
fn run(&self, args: Vec<Vec<u8>>) -> redis::RedisResult<redis::Value> {
|
||||
let mut state = self.state.lock().unwrap();
|
||||
let text = |index: usize| String::from_utf8_lossy(&args[index]).into_owned();
|
||||
match text(0).to_uppercase().as_str() {
|
||||
"FT.CREATE" => {
|
||||
let name = text(1);
|
||||
if state.indexes.contains_key(&name) {
|
||||
return Err(server_error("Index already exists"));
|
||||
}
|
||||
let position = |token: &str| args.iter().position(|arg| arg == token.as_bytes());
|
||||
let prefix = args[position("PREFIX").unwrap() + 2].clone();
|
||||
let dims = text(position("DIM").unwrap() + 1).parse().unwrap();
|
||||
let vector_field = text(position("VECTOR").unwrap() - 1);
|
||||
state.indexes.insert(
|
||||
name,
|
||||
FakeIndex {
|
||||
prefix,
|
||||
dims,
|
||||
vector_field,
|
||||
},
|
||||
);
|
||||
Ok(redis::Value::Okay)
|
||||
}
|
||||
"FT.INFO" => {
|
||||
let index = state
|
||||
.indexes
|
||||
.get(&text(1))
|
||||
.ok_or_else(|| server_error("Unknown index name"))?;
|
||||
Ok(index_info(index))
|
||||
}
|
||||
"FT.DROPINDEX" => {
|
||||
state.indexes.remove(&text(1));
|
||||
Ok(redis::Value::Okay)
|
||||
}
|
||||
"HSET" => {
|
||||
let hash = state.hashes.entry(args[1].clone()).or_default();
|
||||
for pair in args[2..].chunks(2) {
|
||||
hash.insert(
|
||||
String::from_utf8_lossy(&pair[0]).into_owned(),
|
||||
pair[1].clone(),
|
||||
);
|
||||
}
|
||||
Ok(redis::Value::Int(((args.len() - 2) / 2) as i64))
|
||||
}
|
||||
"EXPIRE" => Ok(redis::Value::Int(i64::from(
|
||||
state.hashes.contains_key(&args[1]),
|
||||
))),
|
||||
"FT.SEARCH" => {
|
||||
let index = state
|
||||
.indexes
|
||||
.get(&text(1))
|
||||
.ok_or_else(|| server_error("no such index"))?;
|
||||
let query = text(2);
|
||||
let tag = query_tag(&query);
|
||||
let params = args.iter().position(|arg| arg == b"PARAMS").unwrap();
|
||||
let vector = floats(&args[params + 3]);
|
||||
let best = state
|
||||
.hashes
|
||||
.iter()
|
||||
.filter(|(key, _)| key.starts_with(&index.prefix))
|
||||
.filter(|(_, fields)| {
|
||||
fields.get("litellm_cache_key").map(Vec::as_slice) == Some(tag.as_bytes())
|
||||
})
|
||||
.filter_map(|(key, fields)| {
|
||||
let stored = floats(fields.get(&index.vector_field)?);
|
||||
(stored.len() == index.dims)
|
||||
.then(|| (key, fields, 1.0 - cosine(&vector, &stored)))
|
||||
})
|
||||
.min_by(|left, right| left.2.total_cmp(&right.2));
|
||||
let Some((key, fields, distance)) = best else {
|
||||
return Ok(redis::Value::Array(vec![redis::Value::Int(0)]));
|
||||
};
|
||||
let mut reply = fields
|
||||
.iter()
|
||||
.filter(|(name, _)| **name != index.vector_field)
|
||||
.flat_map(|(name, value)| [bulk(name.as_bytes()), bulk(value)])
|
||||
.collect::<Vec<_>>();
|
||||
reply.extend([
|
||||
bulk(b"vector_distance"),
|
||||
bulk(distance.to_string().as_bytes()),
|
||||
]);
|
||||
Ok(redis::Value::Array(vec![
|
||||
redis::Value::Int(1),
|
||||
bulk(key),
|
||||
redis::Value::Array(reply),
|
||||
]))
|
||||
}
|
||||
"PING" => Ok(redis::Value::SimpleString("PONG".into())),
|
||||
_ => Err(server_error("unsupported command")),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl redis::ConnectionLike for FakeSearch {
|
||||
fn req_packed_command(&mut self, command: &[u8]) -> redis::RedisResult<redis::Value> {
|
||||
let mut commands = parse_commands(command);
|
||||
self.run(commands.remove(0))
|
||||
}
|
||||
|
||||
fn req_packed_commands(
|
||||
&mut self,
|
||||
commands: &[u8],
|
||||
offset: usize,
|
||||
count: usize,
|
||||
) -> redis::RedisResult<Vec<redis::Value>> {
|
||||
let replies = parse_commands(commands)
|
||||
.into_iter()
|
||||
.map(|args| self.run(args))
|
||||
.collect::<redis::RedisResult<Vec<_>>>()?;
|
||||
Ok(replies.into_iter().skip(offset).take(count).collect())
|
||||
}
|
||||
|
||||
fn get_db(&self) -> i64 {
|
||||
0
|
||||
}
|
||||
|
||||
fn check_connection(&mut self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn is_open(&self) -> bool {
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
pub fn server_error(reply: &str) -> redis::RedisError {
|
||||
redis::parse_redis_value(format!("-{reply}\r\n").as_bytes())
|
||||
.expect("error reply")
|
||||
.extract_error()
|
||||
.expect_err("server error")
|
||||
}
|
||||
|
||||
fn bulk(bytes: &[u8]) -> redis::Value {
|
||||
redis::Value::BulkString(bytes.to_vec())
|
||||
}
|
||||
|
||||
fn index_info(index: &FakeIndex) -> redis::Value {
|
||||
let attribute = |name: &str, field_type: &str| {
|
||||
redis::Value::Array(vec![
|
||||
bulk(b"identifier"),
|
||||
bulk(name.as_bytes()),
|
||||
bulk(b"type"),
|
||||
bulk(field_type.as_bytes()),
|
||||
])
|
||||
};
|
||||
redis::Value::Array(vec![
|
||||
bulk(b"attributes"),
|
||||
redis::Value::Array(vec![
|
||||
attribute("prompt", "TEXT"),
|
||||
attribute("response", "TEXT"),
|
||||
attribute("inserted_at", "NUMERIC"),
|
||||
attribute("updated_at", "NUMERIC"),
|
||||
attribute("litellm_cache_key", "TAG"),
|
||||
redis::Value::Array(vec![
|
||||
bulk(b"identifier"),
|
||||
bulk(index.vector_field.as_bytes()),
|
||||
bulk(b"type"),
|
||||
bulk(b"VECTOR"),
|
||||
bulk(b"dim"),
|
||||
redis::Value::Int(index.dims as i64),
|
||||
bulk(b"data_type"),
|
||||
bulk(b"FLOAT32"),
|
||||
bulk(b"distance_metric"),
|
||||
bulk(b"COSINE"),
|
||||
]),
|
||||
]),
|
||||
])
|
||||
}
|
||||
|
||||
/// The tag inside `@litellm_cache_key:{...}`, with query escapes removed.
|
||||
fn query_tag(query: &str) -> String {
|
||||
let start = query.find("@litellm_cache_key:{").unwrap() + "@litellm_cache_key:{".len();
|
||||
let mut tag = String::new();
|
||||
let mut characters = query[start..].chars();
|
||||
while let Some(character) = characters.next() {
|
||||
match character {
|
||||
'\\' => tag.extend(characters.next()),
|
||||
'}' => break,
|
||||
character => tag.push(character),
|
||||
}
|
||||
}
|
||||
tag
|
||||
}
|
||||
|
||||
fn floats(bytes: &[u8]) -> Vec<f32> {
|
||||
bytes
|
||||
.as_chunks::<4>()
|
||||
.0
|
||||
.iter()
|
||||
.map(|chunk| f32::from_le_bytes(*chunk))
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn cosine(left: &[f32], right: &[f32]) -> f64 {
|
||||
let dot = left
|
||||
.iter()
|
||||
.zip(right)
|
||||
.map(|(left, right)| f64::from(*left) * f64::from(*right))
|
||||
.sum::<f64>();
|
||||
let norm = |vector: &[f32]| {
|
||||
vector
|
||||
.iter()
|
||||
.map(|value| f64::from(*value).powi(2))
|
||||
.sum::<f64>()
|
||||
.sqrt()
|
||||
};
|
||||
dot / (norm(left) * norm(right))
|
||||
}
|
||||
|
||||
/// Splits a packed RESP request into each command's arguments.
|
||||
fn parse_commands(mut bytes: &[u8]) -> Vec<Vec<Vec<u8>>> {
|
||||
let line = |bytes: &mut &[u8]| {
|
||||
let end = bytes
|
||||
.windows(2)
|
||||
.position(|window| window == b"\r\n")
|
||||
.unwrap();
|
||||
let text = String::from_utf8(bytes[1..end].to_vec()).unwrap();
|
||||
*bytes = &bytes[end + 2..];
|
||||
text.parse::<usize>().unwrap()
|
||||
};
|
||||
let mut commands = Vec::new();
|
||||
while !bytes.is_empty() {
|
||||
let count = line(&mut bytes);
|
||||
let mut args = Vec::with_capacity(count);
|
||||
for _ in 0..count {
|
||||
let length = line(&mut bytes);
|
||||
args.push(bytes[..length].to_vec());
|
||||
bytes = &bytes[length + 2..];
|
||||
}
|
||||
commands.push(args);
|
||||
}
|
||||
commands
|
||||
}
|
||||
|
|
@ -1,20 +0,0 @@
|
|||
[package]
|
||||
name = "litellm-cache-redis"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
litellm-cache.workspace = true
|
||||
litellm-error.workspace = true
|
||||
redis = { workspace = true, features = ["cluster", "tls-rustls"] }
|
||||
r2d2 = "0.8.10"
|
||||
thiserror.workspace = true
|
||||
tokio.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
litellm-cache-testing.workspace = true
|
||||
redis-test.workspace = true
|
||||
rstest.workspace = true
|
||||
serde_json.workspace = true
|
||||
|
|
@ -1,192 +0,0 @@
|
|||
use std::{
|
||||
sync::{Arc, OnceLock},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use litellm_cache::{BatchEntry, CacheCodec, Error};
|
||||
|
||||
use crate::{
|
||||
connection::{ConnectionRef, Connections},
|
||||
topology::RedisTopology,
|
||||
};
|
||||
|
||||
const DEFAULT_TTL: Duration = Duration::from_secs(600);
|
||||
|
||||
pub struct RedisCache<S, C = redis::Connection> {
|
||||
pub(crate) connections: Arc<Connections<C>>,
|
||||
pub(crate) default_ttl: Duration,
|
||||
pub(crate) codec: S,
|
||||
pub(crate) namespace: Option<String>,
|
||||
pub(crate) topology: RedisTopology,
|
||||
/// The server's major version, read from `INFO` once, like Python's `redis_version`.
|
||||
pub(crate) major_version: Arc<OnceLock<u32>>,
|
||||
}
|
||||
|
||||
impl<S: CacheCodec> RedisCache<S> {
|
||||
pub fn new(url: &str, default_ttl: Option<Duration>, codec: S) -> Result<Self, Error> {
|
||||
Self::connect(url, &RedisTopology::Standalone, default_ttl, codec)
|
||||
}
|
||||
|
||||
pub fn connect(
|
||||
url: &str,
|
||||
topology: &RedisTopology,
|
||||
default_ttl: Option<Duration>,
|
||||
codec: S,
|
||||
) -> Result<Self, Error> {
|
||||
let connections = Connections::open(url, topology)?;
|
||||
Ok(Self {
|
||||
connections: Arc::new(connections),
|
||||
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
|
||||
codec,
|
||||
namespace: None,
|
||||
topology: topology.clone(),
|
||||
major_version: Arc::default(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, C> RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
pub fn with_connection(connection: C, default_ttl: Option<Duration>, codec: S) -> Self {
|
||||
Self {
|
||||
connections: Arc::new(Connections::fixed(connection)),
|
||||
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
|
||||
codec,
|
||||
namespace: None,
|
||||
topology: RedisTopology::Standalone,
|
||||
major_version: Arc::default(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_namespace(self, namespace: Option<String>) -> Self {
|
||||
Self {
|
||||
namespace: namespace.filter(|value| !value.is_empty()),
|
||||
..self
|
||||
}
|
||||
}
|
||||
|
||||
pub fn namespace(&self) -> Option<&str> {
|
||||
self.namespace.as_deref()
|
||||
}
|
||||
|
||||
pub fn topology(&self) -> &RedisTopology {
|
||||
&self.topology
|
||||
}
|
||||
|
||||
pub(crate) fn namespaced_key(&self, key: &str) -> String {
|
||||
namespaced_key(self.namespace.as_deref(), key)
|
||||
}
|
||||
|
||||
pub(crate) fn namespaced_keys(&self, keys: &[String]) -> Vec<String> {
|
||||
keys.iter().map(|key| self.namespaced_key(key)).collect()
|
||||
}
|
||||
|
||||
/// Whole seconds for `ttl`, falling back to the default TTL like Python's `get_ttl`.
|
||||
pub(crate) fn ttl_or_default(&self, ttl: Option<Duration>) -> u64 {
|
||||
ttl_seconds(ttl.unwrap_or(self.default_ttl))
|
||||
}
|
||||
|
||||
pub(crate) fn execute<T>(
|
||||
&self,
|
||||
operation: impl FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error>,
|
||||
) -> Result<T, Error> {
|
||||
self.connections.execute(operation)
|
||||
}
|
||||
|
||||
/// `_parse_redis_major_version`: the major version from `INFO`, or
|
||||
/// `DEFAULT_REDIS_MAJOR_VERSION` when `INFO` fails or its version does not parse. The first
|
||||
/// answer is kept, as Python reads `redis_version` once at construction.
|
||||
pub(crate) async fn major_version(&self) -> u32 {
|
||||
if let Some(version) = self.major_version.get() {
|
||||
return *version;
|
||||
}
|
||||
let info = self
|
||||
.run(|connection| connection.node_text(&redis::cmd("INFO")))
|
||||
.await;
|
||||
let version = info
|
||||
.ok()
|
||||
.and_then(|info| parse_major_version(&info))
|
||||
.unwrap_or_else(default_major_version);
|
||||
*self.major_version.get_or_init(|| version)
|
||||
}
|
||||
|
||||
pub(crate) async fn run<T, F>(&self, operation: F) -> Result<T, Error>
|
||||
where
|
||||
T: Send + 'static,
|
||||
F: FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error> + Send + 'static,
|
||||
{
|
||||
Connections::run_blocking(Arc::clone(&self.connections), operation).await
|
||||
}
|
||||
|
||||
pub(crate) fn namespaced_pattern(&self) -> Result<String, Error> {
|
||||
let namespace = self.namespace.as_ref().ok_or(Error::UnscopedFlush)?;
|
||||
let escaped: String = namespace
|
||||
.chars()
|
||||
.flat_map(|ch| {
|
||||
if matches!(ch, '*' | '?' | '[' | ']' | '\\') {
|
||||
vec!['\\', ch]
|
||||
} else {
|
||||
vec![ch]
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
Ok(format!("{escaped}:*"))
|
||||
}
|
||||
|
||||
pub(crate) fn decode_response(&self, value: redis::Value) -> Result<Option<S::Value>, Error> {
|
||||
match value {
|
||||
redis::Value::Nil => Ok(None),
|
||||
redis::Value::BulkString(bytes) => self.codec.decode(&bytes).map(Some),
|
||||
redis::Value::SimpleString(text) => self.codec.decode(text.as_bytes()).map(Some),
|
||||
_ => Err(Error::InvalidEntry(None)),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn decode_batch_response(
|
||||
&self,
|
||||
value: redis::Value,
|
||||
) -> Result<BatchEntry<S::Value>, Error> {
|
||||
match self.decode_response(value) {
|
||||
Ok(Some(value)) => Ok(BatchEntry::Hit(value)),
|
||||
Ok(None) => Ok(BatchEntry::Miss),
|
||||
Err(Error::InvalidEntry(_)) => Ok(BatchEntry::Invalid),
|
||||
Err(error) => Err(error),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn namespaced_key(namespace: Option<&str>, key: &str) -> String {
|
||||
match namespace {
|
||||
Some(namespace) if !key.starts_with(&format!("{namespace}:")) => {
|
||||
format!("{namespace}:{key}")
|
||||
}
|
||||
_ => key.into(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn ttl_seconds(ttl: Duration) -> u64 {
|
||||
ttl.as_secs()
|
||||
.saturating_add(u64::from(ttl.subsec_nanos() > 0))
|
||||
.max(1)
|
||||
}
|
||||
|
||||
fn parse_major_version(info: &str) -> Option<u32> {
|
||||
let version = info
|
||||
.lines()
|
||||
.find_map(|line| line.trim().strip_prefix("redis_version:"))?
|
||||
.trim();
|
||||
match version.split_once('.') {
|
||||
Some((major, _)) => major.parse().ok(),
|
||||
None => version.parse::<f64>().ok().map(|major| major as u32),
|
||||
}
|
||||
}
|
||||
|
||||
fn default_major_version() -> u32 {
|
||||
std::env::var("DEFAULT_REDIS_MAJOR_VERSION")
|
||||
.ok()
|
||||
.and_then(|value| value.parse().ok())
|
||||
.unwrap_or(7)
|
||||
}
|
||||
|
|
@ -1,109 +0,0 @@
|
|||
use litellm_cache::{CacheCodec, ClaimCache, Error, ExactCacheContext};
|
||||
use redis::Commands;
|
||||
|
||||
use crate::{cache::RedisCache, connection::ConnectionRef};
|
||||
|
||||
const CLAIM_SCRIPT: &str = concat!(
|
||||
"local current = redis.call('GET', KEYS[1]); ",
|
||||
"if ARGV[1] == '' then if current ~= false and current ~= '' then return 0; end; ",
|
||||
"elseif current ~= ARGV[1] then return 0; end; ",
|
||||
"if ARGV[3] ~= '' then redis.call('SET', KEYS[1], ARGV[3], 'EX', ARGV[2]); ",
|
||||
"elseif ARGV[4] == '1' then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return 1"
|
||||
);
|
||||
const CLAIM_ATTEMPTS: usize = 8;
|
||||
|
||||
fn stored_bytes(value: redis::Value) -> Result<Option<Vec<u8>>, Error> {
|
||||
match value {
|
||||
redis::Value::Nil => Ok(None),
|
||||
redis::Value::BulkString(bytes) => Ok(Some(bytes)),
|
||||
redis::Value::SimpleString(text) => Ok(Some(text.into_bytes())),
|
||||
_ => Err(Error::InvalidEntry(None)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Eligibility is decided on decoded values, so a pin written by another encoder (Python's
|
||||
/// `json.dumps` spacing or key order) still matches. The write is a compare-and-set on the
|
||||
/// bytes that decision was made on, retried when another claimant wins the race.
|
||||
fn claim<S: CacheCodec>(
|
||||
connection: &mut ConnectionRef<'_>,
|
||||
codec: &S,
|
||||
key: &str,
|
||||
candidate: S::Value,
|
||||
eligible: &[S::Value],
|
||||
ttl: u64,
|
||||
) -> Result<S::Value, Error>
|
||||
where
|
||||
S::Value: PartialEq,
|
||||
{
|
||||
let payload = codec.encode(&candidate)?;
|
||||
if payload.is_empty() {
|
||||
return Err(Error::InvalidEntry(None));
|
||||
}
|
||||
for _ in 0..CLAIM_ATTEMPTS {
|
||||
let current = stored_bytes(
|
||||
connection
|
||||
.get::<_, redis::Value>(key)
|
||||
.map_err(Error::unavailable)?,
|
||||
)?
|
||||
.filter(|bytes| !bytes.is_empty());
|
||||
let existing = current
|
||||
.as_deref()
|
||||
.and_then(|bytes| codec.decode(bytes).ok())
|
||||
.filter(|existing| eligible.is_empty() || eligible.contains(existing));
|
||||
let refresh = existing
|
||||
.as_ref()
|
||||
.is_some_and(|existing| !eligible.is_empty() || *existing == candidate);
|
||||
let write: &[u8] = if existing.is_some() { b"" } else { &payload };
|
||||
let applied = redis::cmd("EVAL")
|
||||
.arg(CLAIM_SCRIPT)
|
||||
.arg(1)
|
||||
.arg(key)
|
||||
.arg(current.as_deref().unwrap_or_default())
|
||||
.arg(ttl)
|
||||
.arg(write)
|
||||
.arg(u8::from(refresh))
|
||||
.query::<bool>(connection)
|
||||
.map_err(Error::unavailable)?;
|
||||
if applied {
|
||||
return Ok(existing.unwrap_or(candidate));
|
||||
}
|
||||
}
|
||||
Err(Error::unavailable(ClaimContended))
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
#[error("claim kept losing the race after {CLAIM_ATTEMPTS} attempts")]
|
||||
struct ClaimContended;
|
||||
|
||||
impl<S, C> ClaimCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec + Clone + 'static,
|
||||
S::Value: PartialEq,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
fn claim_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
candidate: S::Value,
|
||||
eligible: &[S::Value],
|
||||
context: ExactCacheContext,
|
||||
) -> Result<S::Value, Error> {
|
||||
let key = self.namespaced_key(key);
|
||||
let ttl = self.ttl_or_default(context.ttl);
|
||||
self.execute(|connection| claim(connection, &self.codec, &key, candidate, eligible, ttl))
|
||||
}
|
||||
|
||||
async fn async_claim_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
candidate: S::Value,
|
||||
eligible: Vec<S::Value>,
|
||||
context: ExactCacheContext,
|
||||
) -> Result<S::Value, Error> {
|
||||
let key = self.namespaced_key(key);
|
||||
let ttl = self.ttl_or_default(context.ttl);
|
||||
let codec = self.codec.clone();
|
||||
self.run(move |connection| claim(connection, &codec, &key, candidate, &eligible, ttl))
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
|
@ -1,513 +0,0 @@
|
|||
use std::{
|
||||
sync::{Arc, Mutex},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use litellm_cache::Error;
|
||||
use redis::{
|
||||
ConnectionAddr, ConnectionInfo, IntoConnectionInfo,
|
||||
cluster::{
|
||||
ClusterClient, ClusterClientBuilder, ClusterConnection, ClusterPipeline, NodeAddress,
|
||||
},
|
||||
cluster_routing::{
|
||||
MultipleNodeRoutingInfo, ResponsePolicy, RoutingInfo, SingleNodeRoutingInfo,
|
||||
},
|
||||
};
|
||||
|
||||
use crate::topology::{RedisNode, RedisTopology};
|
||||
|
||||
pub(crate) const REDIS_TIMEOUT: Duration = Duration::from_secs(5);
|
||||
const REDIS_POOL_SIZE: u32 = 16;
|
||||
|
||||
#[allow(private_interfaces)]
|
||||
pub enum Connections<C> {
|
||||
Pool(r2d2::Pool<ConnectionManager>),
|
||||
Cluster(r2d2::Pool<ClusterConnectionManager>),
|
||||
Fixed(Mutex<C>),
|
||||
}
|
||||
|
||||
impl<C> Connections<C>
|
||||
where
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
pub fn execute<T>(
|
||||
&self,
|
||||
operation: impl FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error>,
|
||||
) -> Result<T, Error> {
|
||||
match self {
|
||||
Self::Pool(pool) => {
|
||||
let mut pooled = pool.get().map_err(Error::unavailable)?;
|
||||
let result = operation(&mut ConnectionRef::Node(&mut pooled.connection));
|
||||
pooled.failed = matches!(result, Err(Error::Unavailable(_)));
|
||||
result
|
||||
}
|
||||
Self::Cluster(pool) => {
|
||||
let mut pooled = pool.get().map_err(Error::unavailable)?;
|
||||
let result = operation(&mut ConnectionRef::Cluster(&mut pooled.connection));
|
||||
pooled.failed = matches!(result, Err(Error::Unavailable(_)));
|
||||
result
|
||||
}
|
||||
Self::Fixed(connection) => {
|
||||
let mut connection = connection.lock()?;
|
||||
operation(&mut ConnectionRef::Node(&mut *connection))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn run_blocking<T, F>(connections: Arc<Self>, operation: F) -> Result<T, Error>
|
||||
where
|
||||
T: Send + 'static,
|
||||
F: FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error> + Send + 'static,
|
||||
{
|
||||
tokio::task::spawn_blocking(move || connections.execute(operation))
|
||||
.await
|
||||
.map_err(Error::unavailable)?
|
||||
}
|
||||
|
||||
pub fn fixed(connection: C) -> Self {
|
||||
Self::Fixed(Mutex::new(connection))
|
||||
}
|
||||
|
||||
pub fn open(url: &str, topology: &RedisTopology) -> Result<Self, Error> {
|
||||
match topology {
|
||||
RedisTopology::Standalone => Ok(Self::Pool(pool(ConnectionManager::open(url)?)?)),
|
||||
RedisTopology::Cluster { startup_nodes } => Ok(Self::Cluster(pool(
|
||||
ClusterConnectionManager::open(url, startup_nodes)?,
|
||||
)?)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Closes every idle pooled connection; the next operation opens a fresh one. Connections
|
||||
/// checked out right now return to the pool, and a caller-owned connection stays open.
|
||||
pub fn disconnect(&self) {
|
||||
match self {
|
||||
Self::Pool(pool) => close_idle(pool),
|
||||
Self::Cluster(pool) => close_idle(pool),
|
||||
Self::Fixed(_) => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn pool<M: r2d2::ManageConnection>(manager: M) -> Result<r2d2::Pool<M>, Error> {
|
||||
r2d2::Pool::builder()
|
||||
.max_size(REDIS_POOL_SIZE)
|
||||
.min_idle(Some(0))
|
||||
.connection_timeout(REDIS_TIMEOUT)
|
||||
.test_on_check_out(false)
|
||||
.build(manager)
|
||||
.map_err(Error::unavailable)
|
||||
}
|
||||
|
||||
fn close_idle<M, T>(pool: &r2d2::Pool<M>)
|
||||
where
|
||||
M: r2d2::ManageConnection<Connection = PooledConnection<T>>,
|
||||
{
|
||||
let mut idle = Vec::new();
|
||||
while let Some(mut connection) = pool.try_get() {
|
||||
connection.failed = true;
|
||||
idle.push(connection);
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) struct PooledConnection<C> {
|
||||
connection: C,
|
||||
failed: bool,
|
||||
}
|
||||
|
||||
/// Pools connections without a checkout PING, which would double every operation's round trips.
|
||||
/// A timed-out command leaves its reply on the socket while redis still reports the connection
|
||||
/// open, so any connection whose operation failed is discarded instead of being reused.
|
||||
pub(crate) struct ConnectionManager(redis::Client);
|
||||
|
||||
impl ConnectionManager {
|
||||
fn open(url: &str) -> Result<Self, Error> {
|
||||
redis::Client::open(url)
|
||||
.map(Self)
|
||||
.map_err(Error::unavailable)
|
||||
}
|
||||
}
|
||||
|
||||
impl r2d2::ManageConnection for ConnectionManager {
|
||||
type Connection = PooledConnection<redis::Connection>;
|
||||
type Error = redis::RedisError;
|
||||
|
||||
fn connect(&self) -> Result<Self::Connection, redis::RedisError> {
|
||||
let connection = self.0.get_connection()?;
|
||||
connection.set_read_timeout(Some(REDIS_TIMEOUT))?;
|
||||
connection.set_write_timeout(Some(REDIS_TIMEOUT))?;
|
||||
Ok(PooledConnection {
|
||||
connection,
|
||||
failed: false,
|
||||
})
|
||||
}
|
||||
|
||||
fn is_valid(&self, connection: &mut Self::Connection) -> Result<(), redis::RedisError> {
|
||||
redis::cmd("PING").query::<String>(&mut connection.connection)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn has_broken(&self, connection: &mut Self::Connection) -> bool {
|
||||
connection.failed || !redis::ConnectionLike::is_open(&connection.connection)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
enum TopologyFault {
|
||||
#[error("Redis Cluster needs at least one startup node")]
|
||||
NoStartupNodes,
|
||||
#[error("Redis Cluster nodes need a TCP address")]
|
||||
NonTcpAddress,
|
||||
#[error("Redis Cluster answered a fan-out command without a per-node map")]
|
||||
FanOutReply,
|
||||
#[error("scan target does not match the connection topology")]
|
||||
ScanTarget,
|
||||
}
|
||||
|
||||
pub(crate) struct ClusterConnectionManager(ClusterClient);
|
||||
|
||||
impl ClusterConnectionManager {
|
||||
fn open(url: &str, startup_nodes: &[RedisNode]) -> Result<Self, Error> {
|
||||
if startup_nodes.is_empty() {
|
||||
return Err(Error::unavailable(TopologyFault::NoStartupNodes));
|
||||
}
|
||||
let info = url.into_connection_info().map_err(Error::unavailable)?;
|
||||
let nodes = startup_nodes
|
||||
.iter()
|
||||
.map(|node| node_info(&info, node))
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
ClusterClientBuilder::new(nodes)
|
||||
.connection_timeout(REDIS_TIMEOUT)
|
||||
.response_timeout(REDIS_TIMEOUT)
|
||||
.build()
|
||||
.map(Self)
|
||||
.map_err(Error::unavailable)
|
||||
}
|
||||
}
|
||||
|
||||
fn node_info(info: &ConnectionInfo, node: &RedisNode) -> Result<ConnectionInfo, Error> {
|
||||
let addr = match info.addr() {
|
||||
ConnectionAddr::Tcp(..) => ConnectionAddr::Tcp(node.host.clone(), node.port),
|
||||
ConnectionAddr::TcpTls {
|
||||
insecure,
|
||||
tls_params,
|
||||
..
|
||||
} => ConnectionAddr::TcpTls {
|
||||
host: node.host.clone(),
|
||||
port: node.port,
|
||||
insecure: *insecure,
|
||||
tls_params: tls_params.clone(),
|
||||
},
|
||||
_ => return Err(Error::unavailable(TopologyFault::NonTcpAddress)),
|
||||
};
|
||||
Ok(info.clone().set_addr(addr))
|
||||
}
|
||||
|
||||
impl r2d2::ManageConnection for ClusterConnectionManager {
|
||||
type Connection = PooledConnection<ClusterConnection>;
|
||||
type Error = redis::RedisError;
|
||||
|
||||
fn connect(&self) -> Result<Self::Connection, redis::RedisError> {
|
||||
let connection = self.0.get_connection()?;
|
||||
connection.set_read_timeout(Some(REDIS_TIMEOUT))?;
|
||||
connection.set_write_timeout(Some(REDIS_TIMEOUT))?;
|
||||
Ok(PooledConnection {
|
||||
connection,
|
||||
failed: false,
|
||||
})
|
||||
}
|
||||
|
||||
fn is_valid(&self, connection: &mut Self::Connection) -> Result<(), redis::RedisError> {
|
||||
redis::cmd("PING").query::<String>(&mut connection.connection)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn has_broken(&self, connection: &mut Self::Connection) -> bool {
|
||||
connection.failed || !redis::ConnectionLike::is_open(&connection.connection)
|
||||
}
|
||||
}
|
||||
|
||||
pub enum ConnectionRef<'a> {
|
||||
Node(&'a mut dyn redis::ConnectionLike),
|
||||
Cluster(&'a mut ClusterConnection),
|
||||
}
|
||||
|
||||
impl redis::ConnectionLike for ConnectionRef<'_> {
|
||||
fn req_packed_command(&mut self, cmd: &[u8]) -> redis::RedisResult<redis::Value> {
|
||||
match self {
|
||||
Self::Node(connection) => connection.req_packed_command(cmd),
|
||||
Self::Cluster(connection) => connection.req_packed_command(cmd),
|
||||
}
|
||||
}
|
||||
|
||||
fn req_packed_commands(
|
||||
&mut self,
|
||||
cmd: &[u8],
|
||||
offset: usize,
|
||||
count: usize,
|
||||
) -> redis::RedisResult<Vec<redis::Value>> {
|
||||
match self {
|
||||
Self::Node(connection) => connection.req_packed_commands(cmd, offset, count),
|
||||
Self::Cluster(connection) => connection.req_packed_commands(cmd, offset, count),
|
||||
}
|
||||
}
|
||||
|
||||
fn get_db(&self) -> i64 {
|
||||
match self {
|
||||
Self::Node(connection) => connection.get_db(),
|
||||
Self::Cluster(connection) => redis::ConnectionLike::get_db(*connection),
|
||||
}
|
||||
}
|
||||
|
||||
fn supports_pipelining(&self) -> bool {
|
||||
match self {
|
||||
Self::Node(connection) => connection.supports_pipelining(),
|
||||
Self::Cluster(connection) => redis::ConnectionLike::supports_pipelining(*connection),
|
||||
}
|
||||
}
|
||||
|
||||
fn check_connection(&mut self) -> bool {
|
||||
match self {
|
||||
Self::Node(connection) => connection.check_connection(),
|
||||
Self::Cluster(connection) => connection.check_connection(),
|
||||
}
|
||||
}
|
||||
|
||||
fn is_open(&self) -> bool {
|
||||
match self {
|
||||
Self::Node(connection) => connection.is_open(),
|
||||
Self::Cluster(connection) => redis::ConnectionLike::is_open(*connection),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ConnectionRef<'_> {
|
||||
/// Runs `pipeline` and decodes its non-ignored replies as `T`. A cluster connection refuses
|
||||
/// `Pipeline::query`, so there a transaction goes to its keys' slot as one MULTI/EXEC and
|
||||
/// anything else is split per node by `ClusterPipeline`; either way the raw replies are
|
||||
/// handed back to `pipeline` to decode.
|
||||
pub(crate) fn query_pipeline<T: redis::FromRedisValue>(
|
||||
&mut self,
|
||||
pipeline: &redis::Pipeline,
|
||||
) -> Result<T, Error> {
|
||||
match self {
|
||||
Self::Node(connection) => pipeline.query(*connection),
|
||||
Self::Cluster(connection) if pipeline.is_transaction() => {
|
||||
redis::ConnectionLike::req_packed_commands(
|
||||
*connection,
|
||||
&pipeline.get_packed_pipeline(),
|
||||
pipeline.len() + 1,
|
||||
1,
|
||||
)
|
||||
.and_then(|replies| pipeline.query(&mut Replies(Some(replies))))
|
||||
}
|
||||
Self::Cluster(connection) => {
|
||||
let mut cluster = ClusterPipeline::with_capacity(pipeline.len());
|
||||
for command in pipeline.cmd_iter() {
|
||||
cluster.add_command(command.clone());
|
||||
}
|
||||
cluster
|
||||
.query(connection)
|
||||
.and_then(|replies| pipeline.query(&mut Replies(Some(replies))))
|
||||
}
|
||||
}
|
||||
.map_err(Error::unavailable)
|
||||
}
|
||||
|
||||
pub(crate) fn scan(
|
||||
&mut self,
|
||||
pattern: &str,
|
||||
count: usize,
|
||||
mut visit: impl FnMut(&mut Self, Vec<String>) -> Result<bool, Error>,
|
||||
) -> Result<(), Error> {
|
||||
let pages = match self {
|
||||
Self::Node(connection) => {
|
||||
let page = scan_command(0, pattern, count)
|
||||
.query::<ScanPage>(*connection)
|
||||
.map_err(Error::unavailable)?;
|
||||
vec![(None, page)]
|
||||
}
|
||||
Self::Cluster(connection) => connection
|
||||
.route_command(
|
||||
&scan_command(0, pattern, count),
|
||||
RoutingInfo::MultiNode((
|
||||
MultipleNodeRoutingInfo::AllMasters,
|
||||
Some(ResponsePolicy::Special),
|
||||
)),
|
||||
)
|
||||
.map_err(Error::unavailable)
|
||||
.and_then(primary_pages)?
|
||||
.into_iter()
|
||||
.map(|(node, page)| (Some(node), page))
|
||||
.collect(),
|
||||
};
|
||||
for (node, (mut cursor, mut keys)) in pages {
|
||||
loop {
|
||||
if !visit(self, keys)? {
|
||||
return Ok(());
|
||||
}
|
||||
if cursor == 0 {
|
||||
break;
|
||||
}
|
||||
(cursor, keys) = self.scan_page(node.as_ref(), cursor, pattern, count)?;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) fn ping(&mut self) -> Result<bool, redis::RedisError> {
|
||||
let command = redis::cmd("PING");
|
||||
match self {
|
||||
Self::Node(connection) => command
|
||||
.query::<String>(*connection)
|
||||
.map(|response| response == "PONG"),
|
||||
Self::Cluster(connection) => connection
|
||||
.route_command(
|
||||
&command,
|
||||
RoutingInfo::MultiNode((
|
||||
MultipleNodeRoutingInfo::AllNodes,
|
||||
Some(ResponsePolicy::AllSucceeded),
|
||||
)),
|
||||
)
|
||||
.map(|_| true),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn node_text(&mut self, command: &redis::Cmd) -> Result<String, Error> {
|
||||
match self {
|
||||
Self::Node(connection) => command.query(*connection).map_err(Error::unavailable),
|
||||
Self::Cluster(connection) => {
|
||||
let value = connection
|
||||
.route_command(
|
||||
command,
|
||||
RoutingInfo::MultiNode((
|
||||
MultipleNodeRoutingInfo::AllNodes,
|
||||
Some(ResponsePolicy::Special),
|
||||
)),
|
||||
)
|
||||
.map_err(Error::unavailable)?;
|
||||
let redis::Value::Map(entries) = value else {
|
||||
return Err(Error::unavailable(TopologyFault::FanOutReply));
|
||||
};
|
||||
let mut replies = entries
|
||||
.into_iter()
|
||||
.map(|(node, reply)| {
|
||||
Ok((
|
||||
redis::from_redis_value::<String>(node).map_err(Error::unavailable)?,
|
||||
redis::from_redis_value::<String>(reply).map_err(Error::unavailable)?,
|
||||
))
|
||||
})
|
||||
.collect::<Result<Vec<(String, String)>, Error>>()?;
|
||||
replies.sort();
|
||||
Ok(replies
|
||||
.into_iter()
|
||||
.map(|(_, reply)| reply)
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n"))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn flushall(&mut self) -> Result<(), Error> {
|
||||
let command = redis::cmd("FLUSHALL");
|
||||
match self {
|
||||
Self::Node(connection) => command.query(*connection).map_err(Error::unavailable),
|
||||
Self::Cluster(connection) => connection
|
||||
.route_command(
|
||||
&command,
|
||||
RoutingInfo::MultiNode((
|
||||
MultipleNodeRoutingInfo::AllMasters,
|
||||
Some(ResponsePolicy::AllSucceeded),
|
||||
)),
|
||||
)
|
||||
.map(|_| ())
|
||||
.map_err(Error::unavailable),
|
||||
}
|
||||
}
|
||||
|
||||
fn scan_page(
|
||||
&mut self,
|
||||
node: Option<&NodeAddress>,
|
||||
cursor: u64,
|
||||
pattern: &str,
|
||||
count: usize,
|
||||
) -> Result<ScanPage, Error> {
|
||||
let command = scan_command(cursor, pattern, count);
|
||||
match (self, node) {
|
||||
(Self::Node(connection), None) => {
|
||||
command.query(*connection).map_err(Error::unavailable)
|
||||
}
|
||||
(Self::Cluster(connection), Some(node)) => connection
|
||||
.route_command(
|
||||
&command,
|
||||
RoutingInfo::SingleNode(SingleNodeRoutingInfo::ByAddress {
|
||||
host: node.host().to_string(),
|
||||
port: node.port(),
|
||||
}),
|
||||
)
|
||||
.map_err(Error::unavailable)
|
||||
.and_then(|value| redis::from_redis_value(value).map_err(Error::unavailable)),
|
||||
_ => Err(Error::unavailable(TopologyFault::ScanTarget)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
type ScanPage = (u64, Vec<String>);
|
||||
|
||||
fn primary_pages(value: redis::Value) -> Result<Vec<(NodeAddress, ScanPage)>, Error> {
|
||||
let redis::Value::Map(entries) = value else {
|
||||
return Err(Error::unavailable(TopologyFault::FanOutReply));
|
||||
};
|
||||
entries
|
||||
.into_iter()
|
||||
.map(|(node, page)| {
|
||||
let node = redis::from_redis_value::<String>(node).map_err(Error::unavailable)?;
|
||||
let node = NodeAddress::try_from(node.as_str()).map_err(Error::unavailable)?;
|
||||
let page = redis::from_redis_value::<ScanPage>(page).map_err(Error::unavailable)?;
|
||||
Ok((node, page))
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn scan_command(cursor: u64, pattern: &str, count: usize) -> redis::Cmd {
|
||||
let mut command = redis::cmd("SCAN");
|
||||
command
|
||||
.cursor_arg(cursor)
|
||||
.arg("MATCH")
|
||||
.arg(pattern)
|
||||
.arg("COUNT")
|
||||
.arg(count);
|
||||
command
|
||||
}
|
||||
|
||||
/// Hands already received pipeline replies to `Pipeline::query`, so it applies its own
|
||||
/// ignore and error handling to replies a cluster pipeline gathered from several nodes.
|
||||
struct Replies(Option<Vec<redis::Value>>);
|
||||
|
||||
impl redis::ConnectionLike for Replies {
|
||||
fn req_packed_command(&mut self, _: &[u8]) -> redis::RedisResult<redis::Value> {
|
||||
Err((redis::ErrorKind::Client, "replies hold a pipeline only").into())
|
||||
}
|
||||
|
||||
fn req_packed_commands(
|
||||
&mut self,
|
||||
_: &[u8],
|
||||
_: usize,
|
||||
_: usize,
|
||||
) -> redis::RedisResult<Vec<redis::Value>> {
|
||||
self.0
|
||||
.take()
|
||||
.ok_or_else(|| (redis::ErrorKind::Client, "replies were already read").into())
|
||||
}
|
||||
|
||||
fn get_db(&self) -> i64 {
|
||||
0
|
||||
}
|
||||
|
||||
fn check_connection(&mut self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn is_open(&self) -> bool {
|
||||
true
|
||||
}
|
||||
}
|
||||
|
|
@ -1,205 +0,0 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use litellm_cache::{
|
||||
BoundedCounterCache, CacheCodec, CountReadCache, CounterCache, Error, ExactCacheContext,
|
||||
IncrementOperation,
|
||||
};
|
||||
|
||||
use crate::{
|
||||
cache::{RedisCache, ttl_seconds},
|
||||
connection::ConnectionRef,
|
||||
store::mget,
|
||||
};
|
||||
|
||||
const INCREMENT_SCRIPT: &str = concat!(
|
||||
"local value = redis.call('INCRBYFLOAT', KEYS[1], ARGV[1]); ",
|
||||
"if redis.call('TTL', KEYS[1]) == -1 then ",
|
||||
"redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return value"
|
||||
);
|
||||
const INCREMENT_WITH_FLOOR_SCRIPT: &str = concat!(
|
||||
"local count = redis.call('INCRBY', KEYS[1], ARGV[1]); ",
|
||||
"if count < 0 then count = redis.call('INCRBY', KEYS[1], -count); end; ",
|
||||
"if redis.call('TTL', KEYS[1]) < 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; ",
|
||||
"return count"
|
||||
);
|
||||
const SET_MAX_SCRIPT: &str = concat!(
|
||||
"local current = redis.call('GET', KEYS[1]); ",
|
||||
"if current == false or tonumber(current) < tonumber(ARGV[1]) then ",
|
||||
"redis.call('SET', KEYS[1], ARGV[1]); ",
|
||||
"if tonumber(ARGV[2]) > 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; ",
|
||||
"return ARGV[1]; end; return current"
|
||||
);
|
||||
|
||||
impl<S, C> CounterCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
fn increment_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
amount: f64,
|
||||
context: ExactCacheContext,
|
||||
) -> Result<f64, Error> {
|
||||
let key = self.namespaced_key(key);
|
||||
let ttl = self.ttl_or_default(context.ttl);
|
||||
self.execute(|connection| increment(connection, key, amount, ttl, false))
|
||||
}
|
||||
|
||||
/// Python `_incrbyfloat_with_ttl`: without `refresh_ttl` the TTL is set only on a key that
|
||||
/// has none, in one atomic script; with it, every increment re-arms the TTL.
|
||||
async fn async_increment(
|
||||
&self,
|
||||
key: &str,
|
||||
amount: f64,
|
||||
context: ExactCacheContext,
|
||||
refresh_ttl: bool,
|
||||
) -> Result<f64, Error> {
|
||||
let key = self.namespaced_key(key);
|
||||
let ttl = self.ttl_or_default(context.ttl);
|
||||
self.run(move |connection| increment(connection, key, amount, ttl, refresh_ttl))
|
||||
.await
|
||||
}
|
||||
|
||||
async fn async_increment_pipeline(
|
||||
&self,
|
||||
operations: Vec<IncrementOperation>,
|
||||
) -> Result<Vec<f64>, Error> {
|
||||
if operations.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let mut pipeline = redis::pipe();
|
||||
for operation in operations {
|
||||
let key = self.namespaced_key(&operation.key);
|
||||
pipeline.cmd("INCRBYFLOAT").arg(&key).arg(operation.amount);
|
||||
if let Some(ttl) = operation.ttl {
|
||||
pipeline
|
||||
.cmd("EXPIRE")
|
||||
.arg(key)
|
||||
.arg(ttl_seconds(ttl))
|
||||
.ignore();
|
||||
}
|
||||
}
|
||||
self.run(move |connection| connection.query_pipeline(&pipeline))
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
fn increment(
|
||||
connection: &mut ConnectionRef<'_>,
|
||||
key: String,
|
||||
amount: f64,
|
||||
ttl: u64,
|
||||
refresh_ttl: bool,
|
||||
) -> Result<f64, Error> {
|
||||
if !refresh_ttl {
|
||||
return redis::cmd("EVAL")
|
||||
.arg(INCREMENT_SCRIPT)
|
||||
.arg(1)
|
||||
.arg(key)
|
||||
.arg(amount)
|
||||
.arg(ttl)
|
||||
.query(connection)
|
||||
.map_err(Error::unavailable);
|
||||
}
|
||||
connection
|
||||
.query_pipeline(
|
||||
redis::pipe()
|
||||
.cmd("INCRBYFLOAT")
|
||||
.arg(&key)
|
||||
.arg(amount)
|
||||
.cmd("EXPIRE")
|
||||
.arg(&key)
|
||||
.arg(ttl)
|
||||
.ignore(),
|
||||
)
|
||||
.map(|(value,)| value)
|
||||
}
|
||||
|
||||
impl<S, C> CountReadCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
fn batch_get_counts(&self, keys: &[String]) -> Result<Vec<Option<i64>>, Error> {
|
||||
let keys = self.namespaced_keys(keys);
|
||||
self.execute(|connection| mget(connection, keys))?
|
||||
.into_iter()
|
||||
.map(count)
|
||||
.collect()
|
||||
}
|
||||
|
||||
async fn async_batch_get_counts(&self, keys: Vec<String>) -> Result<Vec<Option<i64>>, Error> {
|
||||
let keys = self.namespaced_keys(&keys);
|
||||
self.run(move |connection| mget(connection, keys))
|
||||
.await?
|
||||
.into_iter()
|
||||
.map(count)
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
fn count(value: redis::Value) -> Result<Option<i64>, Error> {
|
||||
redis::from_redis_value(value).map_err(Error::invalid_entry)
|
||||
}
|
||||
|
||||
impl<S, C> BoundedCounterCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
fn increment_with_floor(&self, key: &str, amount: i64, ttl: Duration) -> Result<i64, Error> {
|
||||
let key = self.namespaced_key(key);
|
||||
let ttl = ttl_seconds(ttl);
|
||||
self.execute(|connection| increment_with_floor(connection, key, amount, ttl))
|
||||
}
|
||||
|
||||
async fn async_increment_with_floor(
|
||||
&self,
|
||||
key: &str,
|
||||
amount: i64,
|
||||
ttl: Duration,
|
||||
) -> Result<i64, Error> {
|
||||
let key = self.namespaced_key(key);
|
||||
let ttl = ttl_seconds(ttl);
|
||||
self.run(move |connection| increment_with_floor(connection, key, amount, ttl))
|
||||
.await
|
||||
}
|
||||
|
||||
async fn async_set_max(
|
||||
&self,
|
||||
key: &str,
|
||||
value: f64,
|
||||
ttl: Option<Duration>,
|
||||
) -> Result<f64, Error> {
|
||||
let key = self.namespaced_key(key);
|
||||
let ttl = self.ttl_or_default(ttl);
|
||||
self.run(move |connection| {
|
||||
redis::cmd("EVAL")
|
||||
.arg(SET_MAX_SCRIPT)
|
||||
.arg(1)
|
||||
.arg(key)
|
||||
.arg(value)
|
||||
.arg(ttl)
|
||||
.query(connection)
|
||||
.map_err(Error::unavailable)
|
||||
})
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
fn increment_with_floor(
|
||||
connection: &mut ConnectionRef<'_>,
|
||||
key: String,
|
||||
amount: i64,
|
||||
ttl: u64,
|
||||
) -> Result<i64, Error> {
|
||||
redis::cmd("EVAL")
|
||||
.arg(INCREMENT_WITH_FLOOR_SCRIPT)
|
||||
.arg(1)
|
||||
.arg(key)
|
||||
.arg(amount)
|
||||
.arg(ttl)
|
||||
.query(connection)
|
||||
.map_err(Error::unavailable)
|
||||
}
|
||||
|
|
@ -1,63 +0,0 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use litellm_cache::{CacheCodec, Error, RefreshTtlCache, ScanCache, TtlCache};
|
||||
|
||||
use crate::cache::RedisCache;
|
||||
|
||||
impl<S, C> TtlCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
async fn async_get_ttl(&self, key: &str) -> Result<Option<Duration>, Error> {
|
||||
let key = self.namespaced_key(key);
|
||||
let ttl = self
|
||||
.run(move |connection| {
|
||||
redis::cmd("TTL")
|
||||
.arg(key)
|
||||
.query::<i64>(connection)
|
||||
.map_err(Error::unavailable)
|
||||
})
|
||||
.await?;
|
||||
Ok(u64::try_from(ttl).ok().map(Duration::from_secs))
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, C> RefreshTtlCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
async fn async_refresh_ttl(&self, key: &str, ttl: Option<Duration>) -> Result<bool, Error> {
|
||||
let key = self.namespaced_key(key);
|
||||
let ttl = self.ttl_or_default(ttl);
|
||||
self.run(move |connection| {
|
||||
redis::cmd("EXPIRE")
|
||||
.arg(key)
|
||||
.arg(ttl)
|
||||
.query(connection)
|
||||
.map_err(Error::unavailable)
|
||||
})
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, C> ScanCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
async fn async_scan_iter(&self, pattern: &str, count: usize) -> Result<Vec<String>, Error> {
|
||||
let pattern = format!("{}*", self.namespaced_key(pattern));
|
||||
self.run(move |connection| {
|
||||
let mut matches = Vec::new();
|
||||
connection.scan(&pattern, count, |_, keys| {
|
||||
matches.extend(keys);
|
||||
Ok(matches.len() < count)
|
||||
})?;
|
||||
matches.truncate(count);
|
||||
Ok(matches)
|
||||
})
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
|
@ -1,15 +0,0 @@
|
|||
mod cache;
|
||||
mod claim;
|
||||
pub mod connection;
|
||||
mod counter;
|
||||
mod keys;
|
||||
mod lifecycle;
|
||||
mod queue;
|
||||
mod script;
|
||||
mod store;
|
||||
mod topology;
|
||||
|
||||
pub use cache::RedisCache;
|
||||
pub use queue::{RedisLpopOperation, RedisLpopResult, RedisRpushOperation};
|
||||
pub use script::{RedisArg, RedisScript};
|
||||
pub use topology::{RedisNode, RedisTopology};
|
||||
|
|
@ -1,86 +0,0 @@
|
|||
use litellm_cache::{
|
||||
CacheCodec, CacheConnectionResult, CacheConnectionStatus, ClientInfoCache, ConnectionCache,
|
||||
DisconnectCache, Error, PingCache,
|
||||
};
|
||||
|
||||
use crate::{cache::RedisCache, topology::RedisTopology};
|
||||
|
||||
impl<S, C> PingCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
fn sync_ping(&self) -> Result<bool, Error> {
|
||||
self.execute(|connection| connection.ping().map_err(Error::unavailable))
|
||||
}
|
||||
|
||||
async fn ping(&self) -> Result<bool, Error> {
|
||||
self.run(|connection| connection.ping().map_err(Error::unavailable))
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, C> ConnectionCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
/// Python `RedisCache.test_connection`, or `RedisClusterCache.test_connection` for a
|
||||
/// cluster topology, which differs only in its messages.
|
||||
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
|
||||
let label = match self.topology {
|
||||
RedisTopology::Standalone => "Redis",
|
||||
RedisTopology::Cluster { .. } => "Redis Cluster",
|
||||
};
|
||||
let ping = self
|
||||
.run(|connection| Ok(connection.ping().map_err(Error::unavailable)))
|
||||
.await
|
||||
.and_then(|ping| ping)
|
||||
.map_err(|error| litellm_error::Report::new(&error).to_string());
|
||||
Ok(match ping {
|
||||
Ok(true) => CacheConnectionResult {
|
||||
status: CacheConnectionStatus::Success,
|
||||
message: format!("{label} connection test successful"),
|
||||
error: None,
|
||||
},
|
||||
Ok(false) => CacheConnectionResult {
|
||||
status: CacheConnectionStatus::Failed,
|
||||
message: format!("{label} ping returned False"),
|
||||
error: None,
|
||||
},
|
||||
Err(error) => CacheConnectionResult {
|
||||
status: CacheConnectionStatus::Failed,
|
||||
message: format!("{label} connection failed: {error}"),
|
||||
error: Some(error),
|
||||
},
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, C> DisconnectCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
async fn disconnect(&self) -> Result<(), Error> {
|
||||
self.connections.disconnect();
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, C> ClientInfoCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
type ClientList = String;
|
||||
type Info = String;
|
||||
|
||||
fn client_list(&self) -> Result<Self::ClientList, Error> {
|
||||
self.execute(|connection| connection.node_text(redis::cmd("CLIENT").arg("LIST")))
|
||||
}
|
||||
|
||||
fn info(&self) -> Result<Self::Info, Error> {
|
||||
self.execute(|connection| connection.node_text(&redis::cmd("INFO")))
|
||||
}
|
||||
}
|
||||
|
|
@ -1,226 +0,0 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use litellm_cache::{CacheCodec, Error, PopOperation, PushOperation, QueueCache, SetCache};
|
||||
|
||||
use crate::{cache::RedisCache, script::RedisArg};
|
||||
|
||||
pub type RedisRpushOperation = PushOperation<RedisArg>;
|
||||
pub type RedisLpopOperation = PopOperation;
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub enum RedisLpopResult {
|
||||
Missing,
|
||||
Value(Vec<u8>),
|
||||
Values(Vec<Vec<u8>>),
|
||||
}
|
||||
|
||||
impl<S, C> SetCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
type SetValue = RedisArg;
|
||||
type SetResult = usize;
|
||||
|
||||
async fn async_set_cache_sadd(
|
||||
&self,
|
||||
key: &str,
|
||||
values: Vec<Self::SetValue>,
|
||||
ttl: Option<Duration>,
|
||||
) -> Result<Self::SetResult, Error> {
|
||||
if values.is_empty() {
|
||||
return Err(Error::InvalidEntry(None));
|
||||
}
|
||||
let key = self.namespaced_key(key);
|
||||
let ttl = self.ttl_or_default(ttl);
|
||||
let mut pipeline = redis::pipe();
|
||||
pipeline
|
||||
.cmd("SADD")
|
||||
.arg(&key)
|
||||
.arg(values)
|
||||
.cmd("EXPIRE")
|
||||
.arg(&key)
|
||||
.arg(ttl)
|
||||
.ignore();
|
||||
self.run(move |connection| connection.query_pipeline(&pipeline))
|
||||
.await
|
||||
.map(|(added,)| added)
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, C> QueueCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
type QueueValue = RedisArg;
|
||||
type PopResult = RedisLpopResult;
|
||||
|
||||
async fn async_rpush(&self, key: &str, values: Vec<Self::QueueValue>) -> Result<usize, Error> {
|
||||
if values.is_empty() {
|
||||
return Err(Error::InvalidEntry(None));
|
||||
}
|
||||
let key = self.namespaced_key(key);
|
||||
self.run(move |connection| {
|
||||
redis::cmd("RPUSH")
|
||||
.arg(key)
|
||||
.arg(values)
|
||||
.query(connection)
|
||||
.map_err(Error::unavailable)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
async fn async_rpush_and_trim(
|
||||
&self,
|
||||
key: &str,
|
||||
values: Vec<Self::QueueValue>,
|
||||
max_len: usize,
|
||||
) -> Result<usize, Error> {
|
||||
if values.is_empty() {
|
||||
return Err(Error::InvalidEntry(None));
|
||||
}
|
||||
let key = self.namespaced_key(key);
|
||||
let start = i64::try_from(max_len).map_or(i64::MIN, |max_len| -max_len);
|
||||
let mut pipeline = redis::pipe();
|
||||
pipeline
|
||||
.atomic()
|
||||
.cmd("RPUSH")
|
||||
.arg(&key)
|
||||
.arg(values)
|
||||
.cmd("LTRIM")
|
||||
.arg(&key)
|
||||
.arg(start)
|
||||
.arg(-1)
|
||||
.ignore();
|
||||
self.run(move |connection| connection.query_pipeline(&pipeline))
|
||||
.await
|
||||
.map(|(length,)| length)
|
||||
}
|
||||
|
||||
async fn async_rpush_pipeline(
|
||||
&self,
|
||||
operations: Vec<RedisRpushOperation>,
|
||||
) -> Result<Vec<usize>, Error> {
|
||||
if operations.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let mut pipeline = redis::pipe();
|
||||
for operation in operations {
|
||||
if operation.values.is_empty() {
|
||||
return Err(Error::InvalidEntry(None));
|
||||
}
|
||||
pipeline
|
||||
.cmd("RPUSH")
|
||||
.arg(self.namespaced_key(&operation.key))
|
||||
.arg(operation.values);
|
||||
}
|
||||
self.run(move |connection| connection.query_pipeline(&pipeline))
|
||||
.await
|
||||
}
|
||||
|
||||
async fn async_lpop(&self, key: &str, count: Option<usize>) -> Result<Self::PopResult, Error> {
|
||||
if let Some(count) = count
|
||||
&& self.major_version().await < 7
|
||||
{
|
||||
return self.lpop_one_at_a_time(key, count).await;
|
||||
}
|
||||
let command = lpop(self.namespaced_key(key), count);
|
||||
let value = self
|
||||
.run(move |connection| {
|
||||
command
|
||||
.query::<redis::Value>(connection)
|
||||
.map_err(Error::unavailable)
|
||||
})
|
||||
.await?;
|
||||
lpop_result(value, count.is_some())
|
||||
}
|
||||
|
||||
async fn async_lpop_pipeline(
|
||||
&self,
|
||||
operations: Vec<RedisLpopOperation>,
|
||||
) -> Result<Vec<Self::PopResult>, Error> {
|
||||
if operations.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
if operations.iter().any(|operation| operation.count.is_some())
|
||||
&& self.major_version().await < 7
|
||||
{
|
||||
let mut results = Vec::with_capacity(operations.len());
|
||||
for operation in &operations {
|
||||
results.push(self.async_lpop(&operation.key, operation.count).await?);
|
||||
}
|
||||
return Ok(results);
|
||||
}
|
||||
let multiple = operations
|
||||
.iter()
|
||||
.map(|operation| operation.count.is_some())
|
||||
.collect::<Vec<_>>();
|
||||
let mut pipeline = redis::pipe();
|
||||
for operation in operations {
|
||||
pipeline.add_command(lpop(self.namespaced_key(&operation.key), operation.count));
|
||||
}
|
||||
self.run(move |connection| connection.query_pipeline::<Vec<redis::Value>>(&pipeline))
|
||||
.await?
|
||||
.into_iter()
|
||||
.zip(multiple)
|
||||
.map(|(value, multiple)| lpop_result(value, multiple))
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
fn lpop(key: String, count: Option<usize>) -> redis::Cmd {
|
||||
let mut command = redis::cmd("LPOP");
|
||||
command.arg(key);
|
||||
if let Some(count) = count {
|
||||
command.arg(count);
|
||||
}
|
||||
command
|
||||
}
|
||||
|
||||
fn redis_bytes(value: redis::Value) -> Result<Vec<u8>, Error> {
|
||||
match value {
|
||||
redis::Value::BulkString(bytes) => Ok(bytes),
|
||||
redis::Value::SimpleString(text) => Ok(text.into_bytes()),
|
||||
_ => Err(Error::InvalidEntry(None)),
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, C> RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
/// `handle_lpop_count_for_older_redis_versions`: `count` single-`LPOP` pipelines, keeping
|
||||
/// only the values actually popped.
|
||||
async fn lpop_one_at_a_time(&self, key: &str, count: usize) -> Result<RedisLpopResult, Error> {
|
||||
let key = self.namespaced_key(key);
|
||||
let mut values = Vec::new();
|
||||
for _ in 0..count {
|
||||
let mut pipeline = redis::pipe();
|
||||
pipeline.add_command(lpop(key.clone(), None));
|
||||
let replies = self
|
||||
.run(move |connection| connection.query_pipeline::<Vec<redis::Value>>(&pipeline))
|
||||
.await?;
|
||||
for reply in replies {
|
||||
if reply != redis::Value::Nil {
|
||||
values.push(redis_bytes(reply)?);
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(RedisLpopResult::Values(values))
|
||||
}
|
||||
}
|
||||
|
||||
fn lpop_result(value: redis::Value, multiple: bool) -> Result<RedisLpopResult, Error> {
|
||||
match value {
|
||||
redis::Value::Nil => Ok(RedisLpopResult::Missing),
|
||||
redis::Value::Array(values) if multiple => values
|
||||
.into_iter()
|
||||
.map(redis_bytes)
|
||||
.collect::<Result<Vec<_>, _>>()
|
||||
.map(RedisLpopResult::Values),
|
||||
value if !multiple => redis_bytes(value).map(RedisLpopResult::Value),
|
||||
_ => Err(Error::InvalidEntry(None)),
|
||||
}
|
||||
}
|
||||
|
|
@ -1,136 +0,0 @@
|
|||
use std::sync::Arc;
|
||||
|
||||
use litellm_cache::{CacheCodec, CacheScript, Error, ScriptCache};
|
||||
|
||||
use crate::{
|
||||
cache::{RedisCache, namespaced_key},
|
||||
connection::Connections,
|
||||
};
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub enum RedisArg {
|
||||
Bytes(Vec<u8>),
|
||||
Integer(i64),
|
||||
Float(f64),
|
||||
}
|
||||
|
||||
impl From<&str> for RedisArg {
|
||||
fn from(value: &str) -> Self {
|
||||
Self::Bytes(value.as_bytes().to_vec())
|
||||
}
|
||||
}
|
||||
|
||||
impl From<String> for RedisArg {
|
||||
fn from(value: String) -> Self {
|
||||
Self::Bytes(value.into_bytes())
|
||||
}
|
||||
}
|
||||
|
||||
impl From<Vec<u8>> for RedisArg {
|
||||
fn from(value: Vec<u8>) -> Self {
|
||||
Self::Bytes(value)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<i64> for RedisArg {
|
||||
fn from(value: i64) -> Self {
|
||||
Self::Integer(value)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<f64> for RedisArg {
|
||||
fn from(value: f64) -> Self {
|
||||
Self::Float(value)
|
||||
}
|
||||
}
|
||||
|
||||
impl redis::ToRedisArgs for RedisArg {
|
||||
fn write_redis_args<W>(&self, out: &mut W)
|
||||
where
|
||||
W: ?Sized + redis::RedisWrite,
|
||||
{
|
||||
match self {
|
||||
Self::Bytes(value) => value.write_redis_args(out),
|
||||
Self::Integer(value) => value.write_redis_args(out),
|
||||
Self::Float(value) => value.write_redis_args(out),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct RedisScript<C> {
|
||||
connections: Arc<Connections<C>>,
|
||||
namespace: Option<String>,
|
||||
source: String,
|
||||
}
|
||||
|
||||
impl<C> CacheScript for RedisScript<C>
|
||||
where
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
type Argument = RedisArg;
|
||||
type Output = redis::Value;
|
||||
|
||||
async fn invoke(
|
||||
&self,
|
||||
keys: Vec<String>,
|
||||
arguments: Vec<Self::Argument>,
|
||||
) -> Result<Self::Output, Error> {
|
||||
let keys = keys
|
||||
.into_iter()
|
||||
.map(|key| namespaced_key(self.namespace.as_deref(), &key))
|
||||
.collect::<Vec<_>>();
|
||||
let source = self.source.clone();
|
||||
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
|
||||
eval(connection, &source, keys, arguments)
|
||||
})
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, C> RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
pub async fn async_eval(
|
||||
&self,
|
||||
script: String,
|
||||
keys: Vec<String>,
|
||||
arguments: Vec<RedisArg>,
|
||||
) -> Result<redis::Value, Error> {
|
||||
let keys = self.namespaced_keys(&keys);
|
||||
self.run(move |connection| eval(connection, &script, keys, arguments))
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
fn eval(
|
||||
connection: &mut impl redis::ConnectionLike,
|
||||
script: &str,
|
||||
keys: Vec<String>,
|
||||
arguments: Vec<RedisArg>,
|
||||
) -> Result<redis::Value, Error> {
|
||||
redis::cmd("EVAL")
|
||||
.arg(script)
|
||||
.arg(keys.len())
|
||||
.arg(keys)
|
||||
.arg(arguments)
|
||||
.query(connection)
|
||||
.map_err(Error::unavailable)
|
||||
}
|
||||
|
||||
impl<S, C> ScriptCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
type Script = RedisScript<C>;
|
||||
|
||||
fn async_register_script(&self, source: String) -> Self::Script {
|
||||
RedisScript {
|
||||
connections: Arc::clone(&self.connections),
|
||||
namespace: self.namespace.clone(),
|
||||
source,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -1,232 +0,0 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use litellm_cache::{
|
||||
BaseCache, BatchCache, BatchEntry, BulkDeleteCache, CacheCodec, DeleteCache, Error,
|
||||
ExactCacheContext, FlushAllCache, FlushCache, TtlPipelineCache,
|
||||
};
|
||||
use redis::Commands;
|
||||
|
||||
use crate::{cache::RedisCache, connection::ConnectionRef};
|
||||
|
||||
impl<S, C> BaseCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
type Value = S::Value;
|
||||
type Context = ExactCacheContext;
|
||||
|
||||
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
|
||||
context.ttl.or(Some(self.default_ttl))
|
||||
}
|
||||
|
||||
fn set_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
value: Self::Value,
|
||||
context: &ExactCacheContext,
|
||||
) -> Result<(), Error> {
|
||||
let payload = self.codec.encode(&value)?;
|
||||
let ttl = self.ttl_or_default(context.ttl);
|
||||
let key = self.namespaced_key(key);
|
||||
self.execute(|connection| {
|
||||
connection
|
||||
.set_ex::<_, _, ()>(key, payload, ttl)
|
||||
.map_err(Error::unavailable)
|
||||
})
|
||||
}
|
||||
|
||||
fn get_cache(&self, key: &str, _: &ExactCacheContext) -> Result<Option<Self::Value>, Error> {
|
||||
let key = self.namespaced_key(key);
|
||||
let value = self.execute(|connection| {
|
||||
connection
|
||||
.get::<_, redis::Value>(key)
|
||||
.map_err(Error::unavailable)
|
||||
})?;
|
||||
self.decode_response(value)
|
||||
}
|
||||
|
||||
async fn async_set_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
value: Self::Value,
|
||||
context: ExactCacheContext,
|
||||
) -> Result<(), Error> {
|
||||
let payload = self.codec.encode(&value)?;
|
||||
let key = self.namespaced_key(key);
|
||||
let ttl = self.ttl_or_default(context.ttl);
|
||||
self.run(move |connection| {
|
||||
connection
|
||||
.set_ex::<_, _, ()>(key, payload, ttl)
|
||||
.map_err(Error::unavailable)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
async fn async_get_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
_: &ExactCacheContext,
|
||||
) -> Result<Option<Self::Value>, Error> {
|
||||
let key = self.namespaced_key(key);
|
||||
let value = self
|
||||
.run(move |connection| {
|
||||
connection
|
||||
.get::<_, redis::Value>(key)
|
||||
.map_err(Error::unavailable)
|
||||
})
|
||||
.await?;
|
||||
self.decode_response(value)
|
||||
}
|
||||
|
||||
async fn async_set_cache_pipeline(
|
||||
&self,
|
||||
cache_list: Vec<(String, Self::Value)>,
|
||||
context: ExactCacheContext,
|
||||
) -> Result<(), Error> {
|
||||
self.async_set_cache_pipeline_with_ttls(
|
||||
cache_list
|
||||
.into_iter()
|
||||
.map(|(key, value)| (key, value, context.ttl))
|
||||
.collect(),
|
||||
)
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, C> TtlPipelineCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
async fn async_set_cache_pipeline_with_ttls(
|
||||
&self,
|
||||
entries: Vec<(String, Self::Value, Option<Duration>)>,
|
||||
) -> Result<(), Error> {
|
||||
if entries.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
let mut pipeline = redis::pipe();
|
||||
for (key, value, ttl) in entries {
|
||||
pipeline
|
||||
.cmd("SETEX")
|
||||
.arg(self.namespaced_key(&key))
|
||||
.arg(self.ttl_or_default(ttl))
|
||||
.arg(self.codec.encode(&value)?)
|
||||
.ignore();
|
||||
}
|
||||
self.run(move |connection| connection.query_pipeline(&pipeline))
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, C> BatchCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
fn batch_get_cache(
|
||||
&self,
|
||||
keys: &[String],
|
||||
_: &ExactCacheContext,
|
||||
) -> Result<Vec<BatchEntry<Self::Value>>, Error> {
|
||||
let keys = self.namespaced_keys(keys);
|
||||
self.execute(|connection| mget(connection, keys))?
|
||||
.into_iter()
|
||||
.map(|value| self.decode_batch_response(value))
|
||||
.collect()
|
||||
}
|
||||
|
||||
async fn async_batch_get_cache(
|
||||
&self,
|
||||
keys: Vec<String>,
|
||||
_: ExactCacheContext,
|
||||
) -> Result<Vec<BatchEntry<Self::Value>>, Error> {
|
||||
let keys = self.namespaced_keys(&keys);
|
||||
self.run(move |connection| mget(connection, keys))
|
||||
.await?
|
||||
.into_iter()
|
||||
.map(|value| self.decode_batch_response(value))
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn mget(
|
||||
connection: &mut ConnectionRef<'_>,
|
||||
keys: Vec<String>,
|
||||
) -> Result<Vec<redis::Value>, Error> {
|
||||
redis::cmd("MGET")
|
||||
.arg(keys)
|
||||
.query(connection)
|
||||
.map_err(Error::unavailable)
|
||||
}
|
||||
|
||||
impl<S, C> DeleteCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
fn delete_cache(&self, key: &str) -> Result<(), Error> {
|
||||
let key = self.namespaced_key(key);
|
||||
self.execute(|connection| connection.del::<_, ()>(key).map_err(Error::unavailable))
|
||||
}
|
||||
|
||||
async fn async_delete_cache(&self, key: &str) -> Result<(), Error> {
|
||||
let key = self.namespaced_key(key);
|
||||
self.run(move |connection| connection.del::<_, ()>(key).map_err(Error::unavailable))
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, C> BulkDeleteCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
async fn delete_cache_keys(&self, keys: Vec<String>) -> Result<usize, Error> {
|
||||
if keys.is_empty() {
|
||||
return Ok(0);
|
||||
}
|
||||
let keys = self.namespaced_keys(&keys);
|
||||
self.run(move |connection| connection.del(keys).map_err(Error::unavailable))
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, C> FlushCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
fn flush_cache(&self) -> Result<(), Error> {
|
||||
let pattern = self.namespaced_pattern()?;
|
||||
self.execute(|connection| flush_matching(connection, &pattern))
|
||||
}
|
||||
|
||||
async fn async_flush_cache(&self) -> Result<(), Error> {
|
||||
let pattern = self.namespaced_pattern()?;
|
||||
self.run(move |connection| flush_matching(connection, &pattern))
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, C> FlushAllCache for RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
fn flushall(&self) -> Result<(), Error> {
|
||||
self.execute(|connection| connection.flushall())
|
||||
}
|
||||
}
|
||||
|
||||
fn flush_matching(connection: &mut ConnectionRef<'_>, pattern: &str) -> Result<(), Error> {
|
||||
connection.scan(pattern, 1000, |connection, keys| {
|
||||
if !keys.is_empty() {
|
||||
connection
|
||||
.del::<_, usize>(keys)
|
||||
.map_err(Error::unavailable)?;
|
||||
}
|
||||
Ok(true)
|
||||
})
|
||||
}
|
||||
|
|
@ -1,14 +0,0 @@
|
|||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct RedisNode {
|
||||
pub host: String,
|
||||
pub port: u16,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default, PartialEq, Eq)]
|
||||
pub enum RedisTopology {
|
||||
#[default]
|
||||
Standalone,
|
||||
Cluster {
|
||||
startup_nodes: Vec<RedisNode>,
|
||||
},
|
||||
}
|
||||
File diff suppressed because it is too large
Load diff
|
|
@ -1,563 +0,0 @@
|
|||
//! Tests against a real Redis Cluster. Set `LITELLM_TEST_REDIS_CLUSTER_NODES` to a comma
|
||||
//! separated `host:port` list (for example `127.0.0.1:7000,127.0.0.1:7001`) to run them.
|
||||
|
||||
mod support;
|
||||
|
||||
use std::{collections::HashSet, time::Duration};
|
||||
|
||||
use litellm_cache::{
|
||||
BaseCache, BatchCache, BatchEntry, BoundedCounterCache, BulkDeleteCache, CacheConnectionStatus,
|
||||
CacheScript, ClaimCache, ClientInfoCache, ConnectionCache, CounterCache, DeleteCache,
|
||||
DisconnectCache, Error, ExactCacheContext, FlushCache, IncrementOperation, JsonCodec,
|
||||
PingCache, QueueCache, RefreshTtlCache, ScanCache, ScriptCache, SetCache, TtlCache,
|
||||
TtlPipelineCache,
|
||||
};
|
||||
use litellm_cache_redis::{
|
||||
RedisArg, RedisCache, RedisLpopOperation, RedisLpopResult, RedisNode, RedisRpushOperation,
|
||||
RedisTopology,
|
||||
};
|
||||
use redis::cluster_routing::Slot;
|
||||
use rstest::{fixture, rstest};
|
||||
use serde_json::json;
|
||||
use support::{JsonCache, cluster_cache, cluster_url};
|
||||
|
||||
type Counter = RedisCache<JsonCodec<f64>>;
|
||||
|
||||
#[fixture]
|
||||
fn cache(#[default("cache")] label: &str) -> Option<JsonCache> {
|
||||
cluster_cache(label, Duration::from_secs(120), JsonCodec::new())
|
||||
}
|
||||
|
||||
#[fixture]
|
||||
fn counter(#[default("counter")] label: &str) -> Option<Counter> {
|
||||
cluster_cache(label, Duration::from_secs(60), JsonCodec::new())
|
||||
}
|
||||
|
||||
#[fixture]
|
||||
fn context() -> ExactCacheContext {
|
||||
ExactCacheContext::default()
|
||||
}
|
||||
|
||||
fn multi_slot_keys(count: usize) -> Vec<String> {
|
||||
let keys: Vec<String> = (0..count).map(|index| format!("key-{index}")).collect();
|
||||
let slots: HashSet<Slot> = keys.iter().map(Slot::for_key).collect();
|
||||
assert!(slots.len() > 1, "keys must span multiple slots");
|
||||
keys
|
||||
}
|
||||
|
||||
fn seconds(seconds: u64) -> Option<Duration> {
|
||||
Some(Duration::from_secs(seconds))
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::no_startup_nodes("redis://127.0.0.1:7000", Vec::new())]
|
||||
#[case::unix_socket_url(
|
||||
"redis+unix:///tmp/redis.sock",
|
||||
vec![RedisNode { host: "127.0.0.1".into(), port: 7000 }]
|
||||
)]
|
||||
fn constructor_rejects_unusable_cluster_configs(
|
||||
#[case] url: &str,
|
||||
#[case] startup_nodes: Vec<RedisNode>,
|
||||
) {
|
||||
let error = JsonCache::connect(
|
||||
url,
|
||||
&RedisTopology::Cluster { startup_nodes },
|
||||
None,
|
||||
JsonCodec::new(),
|
||||
)
|
||||
.err();
|
||||
assert!(matches!(error, Some(Error::Unavailable(_))));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn single_key_operations_round_trip_with_ttl_rounding(
|
||||
#[with("single")] cache: Option<JsonCache>,
|
||||
) {
|
||||
let Some(cache) = cache else { return };
|
||||
let context = ExactCacheContext {
|
||||
ttl: Some(Duration::from_millis(1500)),
|
||||
};
|
||||
let keys = multi_slot_keys(12);
|
||||
for (index, key) in keys.iter().enumerate() {
|
||||
cache
|
||||
.set_cache(key, json!({ "index": index }), &context)
|
||||
.unwrap();
|
||||
}
|
||||
for (index, key) in keys.iter().enumerate() {
|
||||
assert_eq!(
|
||||
cache.get_cache(key, &context).unwrap(),
|
||||
Some(json!({ "index": index }))
|
||||
);
|
||||
}
|
||||
assert_eq!(cache.async_get_ttl(&keys[0]).await.unwrap(), seconds(2));
|
||||
assert!(
|
||||
cache
|
||||
.async_refresh_ttl(&keys[0], seconds(40))
|
||||
.await
|
||||
.unwrap()
|
||||
);
|
||||
assert_eq!(cache.async_get_ttl(&keys[0]).await.unwrap(), seconds(40));
|
||||
cache.delete_cache(&keys[0]).unwrap();
|
||||
assert_eq!(cache.get_cache(&keys[0], &context).unwrap(), None);
|
||||
assert!(!cache.async_refresh_ttl(&keys[0], None).await.unwrap());
|
||||
assert!(cache.sync_ping().unwrap());
|
||||
cache.async_flush_cache().await.unwrap();
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn batch_reads_span_slots_and_preserve_order_with_malformed_entries(
|
||||
#[with("batch")] cache: Option<JsonCache>,
|
||||
context: ExactCacheContext,
|
||||
) {
|
||||
let Some(cache) = cache else { return };
|
||||
let keys = multi_slot_keys(40);
|
||||
for (index, key) in keys.iter().enumerate() {
|
||||
if index % 5 == 0 {
|
||||
continue;
|
||||
}
|
||||
cache
|
||||
.async_set_cache(key, json!(index), context.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
let mut raw = redis::cluster::ClusterClient::new(vec![cluster_url()])
|
||||
.unwrap()
|
||||
.get_connection()
|
||||
.unwrap();
|
||||
let malformed = format!("{}:{}", cache.namespace().unwrap(), keys[1]);
|
||||
redis::cmd("SET")
|
||||
.arg(&malformed)
|
||||
.arg("not json")
|
||||
.exec(&mut raw)
|
||||
.unwrap();
|
||||
|
||||
let entries = cache
|
||||
.async_batch_get_cache(keys.clone(), context.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(entries.len(), keys.len());
|
||||
for (index, entry) in entries.iter().enumerate() {
|
||||
let expected = if index == 1 {
|
||||
BatchEntry::Invalid
|
||||
} else if index % 5 == 0 {
|
||||
BatchEntry::Miss
|
||||
} else {
|
||||
BatchEntry::Hit(json!(index))
|
||||
};
|
||||
assert_eq!(*entry, expected, "entry {index}");
|
||||
}
|
||||
assert_eq!(cache.batch_get_cache(&keys, &context).unwrap(), entries);
|
||||
|
||||
cache.delete_cache_keys(keys.clone()).await.unwrap();
|
||||
let entries = cache.async_batch_get_cache(keys, context).await.unwrap();
|
||||
assert!(entries.iter().all(|entry| *entry == BatchEntry::Miss));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn pipelines_group_by_slot_and_return_results_in_submission_order(
|
||||
#[with("pipeline")] cache: Option<JsonCache>,
|
||||
counter: Option<Counter>,
|
||||
context: ExactCacheContext,
|
||||
) {
|
||||
let (Some(cache), Some(counter)) = (cache, counter) else {
|
||||
return;
|
||||
};
|
||||
let keys = multi_slot_keys(30);
|
||||
cache
|
||||
.async_set_cache_pipeline(
|
||||
keys.iter()
|
||||
.enumerate()
|
||||
.map(|(index, key)| (key.clone(), json!(index)))
|
||||
.collect(),
|
||||
context.clone(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let hits = cache
|
||||
.async_batch_get_cache(keys.clone(), context.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(
|
||||
hits.iter()
|
||||
.enumerate()
|
||||
.all(|(index, entry)| *entry == BatchEntry::Hit(json!(index)))
|
||||
);
|
||||
|
||||
cache
|
||||
.async_set_cache_pipeline_with_ttls(
|
||||
keys.iter()
|
||||
.enumerate()
|
||||
.map(|(index, key)| (key.clone(), json!(index), seconds(index as u64 + 10)))
|
||||
.collect(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
for (index, key) in keys.iter().enumerate() {
|
||||
assert_eq!(
|
||||
cache.async_get_ttl(key).await.unwrap(),
|
||||
seconds(index as u64 + 10),
|
||||
"{key}"
|
||||
);
|
||||
}
|
||||
|
||||
let queues: Vec<String> = keys.iter().map(|key| format!("queue:{key}")).collect();
|
||||
let pushed = cache
|
||||
.async_rpush_pipeline(
|
||||
queues
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(index, key)| RedisRpushOperation {
|
||||
key: key.clone(),
|
||||
values: (0..=index)
|
||||
.map(|value| RedisArg::Integer(value as i64))
|
||||
.collect(),
|
||||
})
|
||||
.collect(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(pushed, (1..=keys.len()).collect::<Vec<_>>());
|
||||
let popped = cache
|
||||
.async_lpop_pipeline(
|
||||
queues
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(index, key)| RedisLpopOperation {
|
||||
key: key.clone(),
|
||||
count: (index % 2 == 0).then_some(2),
|
||||
})
|
||||
.collect(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
for (index, result) in popped.into_iter().enumerate() {
|
||||
match result {
|
||||
RedisLpopResult::Value(value) => {
|
||||
assert_eq!(index % 2, 1, "queue {index}");
|
||||
assert_eq!(value, b"0");
|
||||
}
|
||||
RedisLpopResult::Values(values) => {
|
||||
assert_eq!(index % 2, 0, "queue {index}");
|
||||
let expected: Vec<Vec<u8>> = (0..=index)
|
||||
.take(2)
|
||||
.map(|value| value.to_string().into_bytes())
|
||||
.collect();
|
||||
assert_eq!(values, expected);
|
||||
}
|
||||
other => panic!("queue {index}: {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
let counters: Vec<String> = keys.iter().map(|key| format!("counter:{key}")).collect();
|
||||
let totals = counter
|
||||
.async_increment_pipeline(
|
||||
counters
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(index, key)| IncrementOperation {
|
||||
key: key.clone(),
|
||||
amount: index as f64 + 0.5,
|
||||
ttl: (index % 3 == 0).then_some(Duration::from_secs(30)),
|
||||
})
|
||||
.collect(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let expected: Vec<f64> = (0..keys.len()).map(|index| index as f64 + 0.5).collect();
|
||||
assert_eq!(totals, expected);
|
||||
assert_eq!(
|
||||
counter.async_get_ttl(&counters[0]).await.unwrap(),
|
||||
seconds(30)
|
||||
);
|
||||
assert_eq!(counter.async_get_ttl(&counters[1]).await.unwrap(), None);
|
||||
counter.async_flush_cache().await.unwrap();
|
||||
cache.async_flush_cache().await.unwrap();
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn rpush_and_trim_is_one_transaction_on_the_key_slot(
|
||||
#[with("trim")] cache: Option<JsonCache>,
|
||||
) {
|
||||
let Some(cache) = cache else { return };
|
||||
let values = |values: &[&str]| values.iter().map(|value| RedisArg::from(*value)).collect();
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_rpush_and_trim("buf", values(&["a", "b"]), 3)
|
||||
.await
|
||||
.unwrap(),
|
||||
2
|
||||
);
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_rpush_and_trim("buf", values(&["c", "d"]), 3)
|
||||
.await
|
||||
.unwrap(),
|
||||
4
|
||||
);
|
||||
assert_eq!(
|
||||
cache.async_lpop("buf", Some(10)).await.unwrap(),
|
||||
RedisLpopResult::Values(vec![b"b".to_vec(), b"c".to_vec(), b"d".to_vec()])
|
||||
);
|
||||
cache.async_flush_cache().await.unwrap();
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn scan_and_scoped_flush_cover_every_primary(
|
||||
#[with("flush")] cache: Option<JsonCache>,
|
||||
#[from(cache)]
|
||||
#[with("other")]
|
||||
other: Option<JsonCache>,
|
||||
context: ExactCacheContext,
|
||||
) {
|
||||
let (Some(cache), Some(other)) = (cache, other) else {
|
||||
return;
|
||||
};
|
||||
let keys = multi_slot_keys(60);
|
||||
for key in &keys {
|
||||
cache
|
||||
.async_set_cache(key, json!(true), context.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
other
|
||||
.async_set_cache(key, json!(true), context.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
let mut scanned = cache.async_scan_iter("key-", 1000).await.unwrap();
|
||||
scanned.sort();
|
||||
let mut expected: Vec<String> = keys
|
||||
.iter()
|
||||
.map(|key| format!("{}:{key}", cache.namespace().unwrap()))
|
||||
.collect();
|
||||
expected.sort();
|
||||
assert_eq!(scanned, expected);
|
||||
assert_eq!(cache.async_scan_iter("key-", 7).await.unwrap().len(), 7);
|
||||
|
||||
cache.flush_cache().unwrap();
|
||||
let flushed = cache
|
||||
.async_batch_get_cache(keys.clone(), context.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(flushed.iter().all(|entry| *entry == BatchEntry::Miss));
|
||||
let kept = other.async_batch_get_cache(keys, context).await.unwrap();
|
||||
assert!(
|
||||
kept.iter()
|
||||
.all(|entry| *entry == BatchEntry::Hit(json!(true)))
|
||||
);
|
||||
other.async_flush_cache().await.unwrap();
|
||||
}
|
||||
|
||||
fn ping_calls_per_node(startup: &redis::Client) -> Vec<(String, u64)> {
|
||||
let mut connection = startup.get_connection().unwrap();
|
||||
let nodes: String = redis::cmd("CLUSTER")
|
||||
.arg("NODES")
|
||||
.query(&mut connection)
|
||||
.unwrap();
|
||||
let mut counts: Vec<(String, u64)> = nodes
|
||||
.lines()
|
||||
.map(|line| {
|
||||
let address = line.split_whitespace().nth(1).unwrap();
|
||||
let address = address.split('@').next().unwrap();
|
||||
let mut node = redis::Client::open(format!("redis://{address}"))
|
||||
.unwrap()
|
||||
.get_connection()
|
||||
.unwrap();
|
||||
let stats: String = redis::cmd("INFO")
|
||||
.arg("commandstats")
|
||||
.query(&mut node)
|
||||
.unwrap();
|
||||
let calls = stats
|
||||
.lines()
|
||||
.find_map(|stat| stat.strip_prefix("cmdstat_ping:calls="))
|
||||
.and_then(|rest| rest.split(',').next())
|
||||
.map_or(0, |calls| calls.parse().unwrap());
|
||||
(address.to_string(), calls)
|
||||
})
|
||||
.collect();
|
||||
counts.sort();
|
||||
counts
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn ping_reaches_every_node(#[with("ping")] cache: Option<JsonCache>) {
|
||||
let Some(cache) = cache else { return };
|
||||
let startup = redis::Client::open(cluster_url()).unwrap();
|
||||
let before = ping_calls_per_node(&startup);
|
||||
assert!(before.len() >= 2, "{before:?}");
|
||||
assert!(cache.ping().await.unwrap());
|
||||
let after = ping_calls_per_node(&startup);
|
||||
for ((node, calls_before), (_, calls_after)) in before.iter().zip(&after) {
|
||||
assert!(calls_after > calls_before, "{node} was not pinged");
|
||||
}
|
||||
assert!(cache.sync_ping().unwrap());
|
||||
let result = cache.test_connection().await.unwrap();
|
||||
assert_eq!(result.status, CacheConnectionStatus::Success);
|
||||
assert_eq!(result.message, "Redis Cluster connection test successful");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn disconnect_closes_idle_connections_and_reconnects_on_demand(
|
||||
#[with("disconnect")] cache: Option<JsonCache>,
|
||||
) {
|
||||
let Some(cache) = cache else { return };
|
||||
assert!(cache.ping().await.unwrap());
|
||||
cache.disconnect().await.unwrap();
|
||||
assert!(cache.ping().await.unwrap());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::keep_existing_ttl(false)]
|
||||
#[case::refresh_ttl(true)]
|
||||
#[tokio::test]
|
||||
async fn increments_refresh_the_ttl_only_when_asked(
|
||||
counter: Option<Counter>,
|
||||
#[case] refresh_ttl: bool,
|
||||
) {
|
||||
let Some(counter) = counter else { return };
|
||||
let context = ExactCacheContext { ttl: seconds(60) };
|
||||
counter
|
||||
.async_set_cache("spend", 0.0, ExactCacheContext { ttl: seconds(600) })
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
counter
|
||||
.async_increment("spend", 1.5, context.clone(), refresh_ttl)
|
||||
.await
|
||||
.unwrap(),
|
||||
1.5
|
||||
);
|
||||
assert_eq!(
|
||||
counter
|
||||
.async_increment("spend", 2.0, context, refresh_ttl)
|
||||
.await
|
||||
.unwrap(),
|
||||
3.5
|
||||
);
|
||||
let ttl = counter.async_get_ttl("spend").await.unwrap().unwrap();
|
||||
assert_eq!(ttl <= Duration::from_secs(60), refresh_ttl, "{ttl:?}");
|
||||
counter.async_flush_cache().await.unwrap();
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn counters_claims_scripts_and_sets_work_on_the_cluster(
|
||||
counter: Option<Counter>,
|
||||
#[with("claim")] cache: Option<JsonCache>,
|
||||
context: ExactCacheContext,
|
||||
) {
|
||||
let (Some(counter), Some(cache)) = (counter, cache) else {
|
||||
return;
|
||||
};
|
||||
assert_eq!(
|
||||
counter
|
||||
.increment_cache("spend", 1.5, context.clone())
|
||||
.unwrap(),
|
||||
1.5
|
||||
);
|
||||
assert_eq!(
|
||||
counter
|
||||
.async_increment("spend", 2.0, context.clone(), false)
|
||||
.await
|
||||
.unwrap(),
|
||||
3.5
|
||||
);
|
||||
assert_eq!(
|
||||
counter
|
||||
.increment_with_floor("budget", -3, Duration::from_secs(30))
|
||||
.unwrap(),
|
||||
0
|
||||
);
|
||||
assert_eq!(
|
||||
counter
|
||||
.async_increment_with_floor("budget", 7, Duration::from_secs(30))
|
||||
.await
|
||||
.unwrap(),
|
||||
7
|
||||
);
|
||||
assert_eq!(counter.async_set_max("peak", 4.0, None).await.unwrap(), 4.0);
|
||||
assert_eq!(counter.async_set_max("peak", 2.0, None).await.unwrap(), 4.0);
|
||||
counter.flush_cache().unwrap();
|
||||
|
||||
let owner = json!("owner-a");
|
||||
let rival = json!("owner-b");
|
||||
assert_eq!(
|
||||
cache
|
||||
.claim_cache("lock", owner.clone(), &[], context.clone())
|
||||
.unwrap(),
|
||||
owner
|
||||
);
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_claim_cache("lock", rival.clone(), vec![owner.clone()], context.clone())
|
||||
.await
|
||||
.unwrap(),
|
||||
owner
|
||||
);
|
||||
assert_eq!(
|
||||
cache
|
||||
.claim_cache("lock", rival.clone(), &[], context.clone())
|
||||
.unwrap(),
|
||||
owner
|
||||
);
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_claim_cache("lock", rival.clone(), vec![rival.clone()], context.clone())
|
||||
.await
|
||||
.unwrap(),
|
||||
rival
|
||||
);
|
||||
|
||||
let script = cache
|
||||
.async_register_script("return redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2])".into());
|
||||
let reply = script
|
||||
.invoke(
|
||||
vec!["scripted".into()],
|
||||
vec![RedisArg::Bytes(b"payload".to_vec()), RedisArg::Integer(5)],
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(reply, redis::Value::Okay);
|
||||
assert_eq!(cache.async_get_ttl("scripted").await.unwrap(), seconds(5));
|
||||
let evaluated = cache
|
||||
.async_eval(
|
||||
"return redis.call('GET', KEYS[1])".into(),
|
||||
vec!["scripted".into()],
|
||||
Vec::new(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(evaluated, redis::Value::BulkString(b"payload".to_vec()));
|
||||
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_set_cache_sadd(
|
||||
"members",
|
||||
vec![
|
||||
RedisArg::Bytes(b"a".to_vec()),
|
||||
RedisArg::Bytes(b"b".to_vec())
|
||||
],
|
||||
seconds(9),
|
||||
)
|
||||
.await
|
||||
.unwrap(),
|
||||
2
|
||||
);
|
||||
assert_eq!(cache.async_get_ttl("members").await.unwrap(), seconds(9));
|
||||
|
||||
let result = cache.test_connection().await.unwrap();
|
||||
assert_eq!(result.status, CacheConnectionStatus::Success);
|
||||
assert!(cache.ping().await.unwrap());
|
||||
let info = cache.info().unwrap();
|
||||
assert!(info.matches("redis_version").count() > 1, "{info}");
|
||||
assert!(cache.client_list().unwrap().contains("id="));
|
||||
cache.async_flush_cache().await.unwrap();
|
||||
assert_eq!(cache.async_get_ttl("members").await.unwrap(), None);
|
||||
assert_eq!(cache.get_cache("lock", &context).unwrap(), None);
|
||||
}
|
||||
|
|
@ -1,95 +0,0 @@
|
|||
//! The shared cache contracts, run against the in-process fake connection and, when
|
||||
//! `LITELLM_TEST_REDIS_CLUSTER_NODES` is set, against a live Redis Cluster.
|
||||
|
||||
mod support;
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use litellm_cache::{ExactCacheContext, JsonCodec};
|
||||
use litellm_cache_testing as contract;
|
||||
use rstest::rstest;
|
||||
use serde_json::json;
|
||||
use support::{JsonCache, cluster_cache, fake_cache};
|
||||
|
||||
const PREFIX: &str = "contract:";
|
||||
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
enum Contract {
|
||||
HitAndMiss,
|
||||
SyncAsyncEquivalence,
|
||||
OverwriteReplaces,
|
||||
PipelineWritesEveryEntry,
|
||||
BatchPreservesOrder,
|
||||
DeleteRemovesKey,
|
||||
FlushClears,
|
||||
CounterAccumulates,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
enum Server {
|
||||
Fake,
|
||||
Cluster,
|
||||
}
|
||||
|
||||
async fn check<C>(contract: Contract, cache: &JsonCache<C>)
|
||||
where
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
let context = ExactCacheContext::default();
|
||||
match contract {
|
||||
Contract::HitAndMiss => {
|
||||
contract::hit_and_miss(cache, context, PREFIX, json!({"answer": 42})).await
|
||||
}
|
||||
Contract::SyncAsyncEquivalence => {
|
||||
contract::sync_async_equivalence(cache, context, PREFIX, json!("first"), json!([2]))
|
||||
.await
|
||||
}
|
||||
Contract::OverwriteReplaces => {
|
||||
contract::overwrite_replaces(cache, context, PREFIX, json!(1), json!({"b": 2})).await
|
||||
}
|
||||
Contract::PipelineWritesEveryEntry => {
|
||||
contract::pipeline_writes_every_entry(
|
||||
cache,
|
||||
context,
|
||||
PREFIX,
|
||||
vec![json!("a"), json!(2), json!({"c": true})],
|
||||
)
|
||||
.await
|
||||
}
|
||||
Contract::BatchPreservesOrder => {
|
||||
contract::batch_preserves_order(cache, context, PREFIX, json!("first"), json!(2)).await
|
||||
}
|
||||
Contract::DeleteRemovesKey => {
|
||||
contract::delete_removes_key(cache, context, PREFIX, json!("value")).await
|
||||
}
|
||||
Contract::FlushClears => {
|
||||
contract::flush_clears(cache, context, PREFIX, json!("value")).await
|
||||
}
|
||||
Contract::CounterAccumulates => contract::counter_accumulates(cache, context, PREFIX).await,
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::hit_and_miss(Contract::HitAndMiss)]
|
||||
#[case::sync_async_equivalence(Contract::SyncAsyncEquivalence)]
|
||||
#[case::overwrite_replaces(Contract::OverwriteReplaces)]
|
||||
#[case::pipeline_writes_every_entry(Contract::PipelineWritesEveryEntry)]
|
||||
#[case::batch_preserves_order(Contract::BatchPreservesOrder)]
|
||||
#[case::delete_removes_key(Contract::DeleteRemovesKey)]
|
||||
#[case::flush_clears(Contract::FlushClears)]
|
||||
#[case::counter_accumulates(Contract::CounterAccumulates)]
|
||||
#[tokio::test]
|
||||
async fn redis_satisfies_the_cache_contract(
|
||||
#[case] contract: Contract,
|
||||
#[values(Server::Fake, Server::Cluster)] server: Server,
|
||||
) {
|
||||
match server {
|
||||
Server::Fake => check(contract, &fake_cache("contract")).await,
|
||||
Server::Cluster => {
|
||||
let label = format!("{contract:?}");
|
||||
if let Some(cache) = cluster_cache(&label, Duration::from_secs(120), JsonCodec::new()) {
|
||||
check(contract, &cache).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -1,231 +0,0 @@
|
|||
#![allow(dead_code)]
|
||||
|
||||
use std::{
|
||||
collections::BTreeMap,
|
||||
time::{Duration, SystemTime, UNIX_EPOCH},
|
||||
};
|
||||
|
||||
use litellm_cache::{CacheCodec, Error, JsonCodec};
|
||||
use litellm_cache_redis::{RedisCache, RedisNode, RedisTopology};
|
||||
|
||||
/// Encodes a byte behind a tag, so a value written with another tag decodes as invalid.
|
||||
pub struct TaggedByteCodec(pub u8);
|
||||
|
||||
impl CacheCodec for TaggedByteCodec {
|
||||
type Value = u8;
|
||||
|
||||
fn encode(&self, value: &u8) -> Result<Vec<u8>, Error> {
|
||||
if *value > 127 {
|
||||
return Err(Error::InvalidEntry(None));
|
||||
}
|
||||
Ok(vec![self.0, *value])
|
||||
}
|
||||
|
||||
fn decode(&self, bytes: &[u8]) -> Result<u8, Error> {
|
||||
match bytes {
|
||||
[tag, value] if *tag == self.0 => Ok(*value),
|
||||
_ => Err(Error::InvalidEntry(None)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A stateful in-process stand-in for a Redis server that understands the string commands the
|
||||
/// shared contracts exercise, so they run without a live server. TTLs are accepted and ignored.
|
||||
#[derive(Default)]
|
||||
pub struct FakeRedis {
|
||||
strings: BTreeMap<Vec<u8>, Vec<u8>>,
|
||||
}
|
||||
|
||||
impl FakeRedis {
|
||||
fn run(&mut self, command: Vec<Vec<u8>>) -> redis::RedisResult<redis::Value> {
|
||||
let name = String::from_utf8_lossy(&command[0]).to_ascii_uppercase();
|
||||
let args = &command[1..];
|
||||
Ok(match name.as_str() {
|
||||
"PING" => redis::Value::SimpleString("PONG".into()),
|
||||
"SET" | "SETEX" => {
|
||||
let value = if name == "SET" { &args[1] } else { &args[2] };
|
||||
self.strings.insert(args[0].clone(), value.clone());
|
||||
redis::Value::Okay
|
||||
}
|
||||
"GET" => self.get(&args[0]),
|
||||
"MGET" => redis::Value::Array(args.iter().map(|key| self.get(key)).collect()),
|
||||
"DEL" => {
|
||||
let removed = args
|
||||
.iter()
|
||||
.filter(|key| self.strings.remove(*key).is_some())
|
||||
.count();
|
||||
redis::Value::Int(removed as i64)
|
||||
}
|
||||
"SCAN" => {
|
||||
let pattern = &args[2];
|
||||
let keys = self
|
||||
.strings
|
||||
.keys()
|
||||
.filter(|key| glob(pattern, key))
|
||||
.map(|key| redis::Value::BulkString(key.clone()))
|
||||
.collect();
|
||||
redis::Value::Array(vec![
|
||||
redis::Value::BulkString(b"0".to_vec()),
|
||||
redis::Value::Array(keys),
|
||||
])
|
||||
}
|
||||
"EVAL" if args[0].windows(11).any(|window| window == b"INCRBYFLOAT") => {
|
||||
self.increment_by_float(&args[2], &args[3])
|
||||
}
|
||||
"INCRBYFLOAT" => self.increment_by_float(&args[0], &args[1]),
|
||||
_ => {
|
||||
return Err(redis::RedisError::from((
|
||||
redis::ErrorKind::Client,
|
||||
"unsupported command",
|
||||
name,
|
||||
)));
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
fn get(&self, key: &[u8]) -> redis::Value {
|
||||
self.strings.get(key).map_or(redis::Value::Nil, |value| {
|
||||
redis::Value::BulkString(value.clone())
|
||||
})
|
||||
}
|
||||
|
||||
fn increment_by_float(&mut self, key: &[u8], amount: &[u8]) -> redis::Value {
|
||||
let current = self
|
||||
.strings
|
||||
.get(key)
|
||||
.map_or(0.0, |value| parse_float(value));
|
||||
let total = format!("{}", current + parse_float(amount));
|
||||
self.strings
|
||||
.insert(key.to_vec(), total.clone().into_bytes());
|
||||
redis::Value::BulkString(total.into_bytes())
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_float(bytes: &[u8]) -> f64 {
|
||||
std::str::from_utf8(bytes).unwrap().parse().unwrap()
|
||||
}
|
||||
|
||||
/// Redis `MATCH` globbing for `*`, `?` and backslash escapes.
|
||||
fn glob(pattern: &[u8], key: &[u8]) -> bool {
|
||||
match pattern.split_first() {
|
||||
None => key.is_empty(),
|
||||
Some((b'*', rest)) => (0..=key.len()).any(|skip| glob(rest, &key[skip..])),
|
||||
Some((b'?', rest)) => !key.is_empty() && glob(rest, &key[1..]),
|
||||
Some((b'\\', [escaped, rest @ ..])) => {
|
||||
key.first() == Some(escaped) && glob(rest, &key[1..])
|
||||
}
|
||||
Some((literal, rest)) => key.first() == Some(literal) && glob(rest, &key[1..]),
|
||||
}
|
||||
}
|
||||
|
||||
/// Splits RESP request bytes into the commands they carry.
|
||||
fn commands(mut bytes: &[u8]) -> Vec<Vec<Vec<u8>>> {
|
||||
fn line<'a>(bytes: &mut &'a [u8]) -> &'a [u8] {
|
||||
let end = bytes
|
||||
.windows(2)
|
||||
.position(|window| window == b"\r\n")
|
||||
.unwrap();
|
||||
let (line, rest) = bytes.split_at(end);
|
||||
*bytes = &rest[2..];
|
||||
line
|
||||
}
|
||||
fn length(line: &[u8]) -> usize {
|
||||
std::str::from_utf8(&line[1..]).unwrap().parse().unwrap()
|
||||
}
|
||||
let mut commands = Vec::new();
|
||||
while !bytes.is_empty() {
|
||||
let count = length(line(&mut bytes));
|
||||
let command = (0..count)
|
||||
.map(|_| {
|
||||
let size = length(line(&mut bytes));
|
||||
let (argument, rest) = bytes.split_at(size);
|
||||
bytes = &rest[2..];
|
||||
argument.to_vec()
|
||||
})
|
||||
.collect();
|
||||
commands.push(command);
|
||||
}
|
||||
commands
|
||||
}
|
||||
|
||||
impl redis::ConnectionLike for FakeRedis {
|
||||
fn req_packed_command(&mut self, cmd: &[u8]) -> redis::RedisResult<redis::Value> {
|
||||
let command = commands(cmd).into_iter().next().unwrap();
|
||||
self.run(command)
|
||||
}
|
||||
|
||||
fn req_packed_commands(
|
||||
&mut self,
|
||||
cmd: &[u8],
|
||||
offset: usize,
|
||||
count: usize,
|
||||
) -> redis::RedisResult<Vec<redis::Value>> {
|
||||
let replies = commands(cmd)
|
||||
.into_iter()
|
||||
.map(|command| self.run(command))
|
||||
.collect::<redis::RedisResult<Vec<_>>>()?;
|
||||
Ok(replies.into_iter().skip(offset).take(count).collect())
|
||||
}
|
||||
|
||||
fn get_db(&self) -> i64 {
|
||||
0
|
||||
}
|
||||
|
||||
fn check_connection(&mut self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn is_open(&self) -> bool {
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
pub type JsonCache<C = redis::Connection> = RedisCache<JsonCodec<serde_json::Value>, C>;
|
||||
|
||||
pub fn fake_cache(namespace: &str) -> JsonCache<FakeRedis> {
|
||||
RedisCache::with_connection(FakeRedis::default(), None, JsonCodec::new())
|
||||
.with_namespace(Some(namespace.into()))
|
||||
}
|
||||
|
||||
/// Startup nodes from `LITELLM_TEST_REDIS_CLUSTER_NODES` (`host:port,host:port`); tests that
|
||||
/// need a live cluster skip when it is unset.
|
||||
pub fn cluster_topology() -> Option<RedisTopology> {
|
||||
let nodes = std::env::var("LITELLM_TEST_REDIS_CLUSTER_NODES").ok()?;
|
||||
let startup_nodes = nodes
|
||||
.split(',')
|
||||
.map(|node| {
|
||||
let (host, port) = node.trim().rsplit_once(':').expect("host:port");
|
||||
RedisNode {
|
||||
host: host.to_string(),
|
||||
port: port.parse().expect("port"),
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
Some(RedisTopology::Cluster { startup_nodes })
|
||||
}
|
||||
|
||||
pub fn cluster_url() -> String {
|
||||
std::env::var("LITELLM_TEST_REDIS_CLUSTER_URL")
|
||||
.unwrap_or_else(|_| "redis://127.0.0.1:7000".into())
|
||||
}
|
||||
|
||||
pub fn unique_namespace(label: &str) -> String {
|
||||
let nanos = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_nanos();
|
||||
format!("cluster-test:{label}:{nanos}")
|
||||
}
|
||||
|
||||
pub fn cluster_cache<S: CacheCodec>(
|
||||
label: &str,
|
||||
default_ttl: Duration,
|
||||
codec: S,
|
||||
) -> Option<RedisCache<S>> {
|
||||
let topology = cluster_topology()?;
|
||||
Some(
|
||||
RedisCache::connect(&cluster_url(), &topology, Some(default_ttl), codec)
|
||||
.expect("cluster connection")
|
||||
.with_namespace(Some(unique_namespace(label))),
|
||||
)
|
||||
}
|
||||
|
|
@ -16,9 +16,6 @@ sha2.workspace = true
|
|||
litellm-cache-gcs.workspace = true
|
||||
litellm-http = { workspace = true, features = ["test-support"] }
|
||||
litellm-cache-memory.workspace = true
|
||||
litellm-cache-redis.workspace = true
|
||||
redis.workspace = true
|
||||
redis-test.workspace = true
|
||||
rstest.workspace = true
|
||||
tokio.workspace = true
|
||||
wiremock.workspace = true
|
||||
|
|
|
|||
|
|
@ -2,79 +2,11 @@ mod support;
|
|||
|
||||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use litellm_cache::CacheConnectionStatus;
|
||||
use litellm_cache_memory::InMemoryCache;
|
||||
use litellm_cache_response::{
|
||||
CacheEntry, ConnectionProbe, ExactResponseCache, ResponseCache, ResponseCacheRequest,
|
||||
};
|
||||
use redis_test::MockCmd;
|
||||
use litellm_cache_response::{CacheEntry, ExactResponseCache, ResponseCache, ResponseCacheRequest};
|
||||
use rstest::rstest;
|
||||
use serde_json::json;
|
||||
use support::{keyed, memory, redis, request};
|
||||
|
||||
#[rstest]
|
||||
#[case::reachable(
|
||||
Ok("PONG"),
|
||||
CacheConnectionStatus::Success,
|
||||
"Redis connection test successful",
|
||||
false
|
||||
)]
|
||||
#[case::unexpected_reply(
|
||||
Ok("NOPE"),
|
||||
CacheConnectionStatus::Failed,
|
||||
"Redis ping returned False",
|
||||
false
|
||||
)]
|
||||
#[case::connection_refused(
|
||||
Err(redis::RedisError::from((redis::ErrorKind::Io, "connection refused"))),
|
||||
CacheConnectionStatus::Failed,
|
||||
"Redis connection failed:",
|
||||
true
|
||||
)]
|
||||
#[tokio::test]
|
||||
async fn connection_backends_are_reachable_as_a_probe(
|
||||
#[case] reply: redis::RedisResult<&'static str>,
|
||||
#[case] status: CacheConnectionStatus,
|
||||
#[case] message: &str,
|
||||
#[case] has_error: bool,
|
||||
) {
|
||||
let probe: Arc<dyn ConnectionProbe> =
|
||||
Arc::new(redis(vec![MockCmd::new(redis::cmd("PING"), reply)], None));
|
||||
|
||||
let result = probe.test_connection().await.unwrap();
|
||||
assert_eq!(result.status, status);
|
||||
assert!(result.message.starts_with(message), "{}", result.message);
|
||||
assert_eq!(result.error.is_some(), has_error);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn one_service_serves_both_the_exact_cache_and_its_probe(request: ResponseCacheRequest) {
|
||||
let service = Arc::new(redis(
|
||||
vec![
|
||||
MockCmd::new(redis::cmd("PING"), Ok("PONG")),
|
||||
MockCmd::new(
|
||||
redis::cmd("GET").arg("tenant:key"),
|
||||
Ok(br#"{"timestamp":100.0,"response":{"ok":true}}"#.to_vec()),
|
||||
),
|
||||
],
|
||||
Some("tenant"),
|
||||
));
|
||||
let probe: Arc<dyn ConnectionProbe> = service.clone();
|
||||
let exact: Arc<dyn ExactResponseCache> = service;
|
||||
|
||||
assert_eq!(
|
||||
probe.test_connection().await.unwrap().status,
|
||||
CacheConnectionStatus::Success
|
||||
);
|
||||
assert_eq!(
|
||||
exact
|
||||
.async_lookup(&request, Duration::from_secs(100))
|
||||
.await
|
||||
.unwrap(),
|
||||
Some(json!({"ok": true}))
|
||||
);
|
||||
}
|
||||
use support::{keyed, memory, request};
|
||||
|
||||
/// The in-memory backend has no `test_connection`, as in Python, and still serves every response
|
||||
/// operation.
|
||||
|
|
|
|||
|
|
@ -17,10 +17,9 @@ use litellm_cache_response::{
|
|||
CacheControls, CacheEntry, CacheKeyField, CacheKeyInput, ResponseCache, ResponseCacheRequest,
|
||||
WriteBuffer, cache_key,
|
||||
};
|
||||
use redis_test::MockCmd;
|
||||
use rstest::{fixture, rstest};
|
||||
use serde_json::{Value, json};
|
||||
use support::{keyed, memory, redis, request};
|
||||
use support::{keyed, memory, request};
|
||||
|
||||
type Memory = Arc<ResponseCache<InMemoryCache<CacheEntry>>>;
|
||||
|
||||
|
|
@ -331,79 +330,6 @@ async fn directives_skip_io_and_keep_reads_and_writes_independent(
|
|||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::python_sync_literal(
|
||||
br#"{'timestamp': 100.0, 'response': '{"ok": true, "text": "cached"}'}"#.as_slice()
|
||||
)]
|
||||
#[case::python_async_json(br#"{"timestamp":100.0,"response":{"ok":true,"text":"cached"}}"#.as_slice())]
|
||||
#[tokio::test]
|
||||
async fn redis_consumer_reads_python_sync_and_async_envelopes(
|
||||
request: ResponseCacheRequest,
|
||||
#[case] stored: &[u8],
|
||||
#[values(false, true)] asynchronous: bool,
|
||||
) {
|
||||
let cache = redis(
|
||||
vec![MockCmd::new(
|
||||
redis::cmd("GET").arg("tenant:key"),
|
||||
Ok(stored.to_vec()),
|
||||
)],
|
||||
Some("tenant"),
|
||||
);
|
||||
let now = Duration::from_secs(101);
|
||||
let found = if asynchronous {
|
||||
cache.async_lookup(&request, now).await.unwrap()
|
||||
} else {
|
||||
cache.lookup(&request, now).unwrap()
|
||||
};
|
||||
assert_eq!(found, Some(json!({"ok": true, "text": "cached"})));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::object(
|
||||
json!({"ok": true, "text": "cached"}),
|
||||
br#"{"timestamp":100.0,"response":{"ok":true,"text":"cached"}}"#.as_slice()
|
||||
)]
|
||||
#[case::array(json!([1, 2]), br#"{"timestamp":100.0,"response":"[1,2]"}"#.as_slice())]
|
||||
#[tokio::test]
|
||||
async fn redis_consumer_writes_python_compatible_json(
|
||||
request: ResponseCacheRequest,
|
||||
#[case] response: Value,
|
||||
#[case] wire: &[u8],
|
||||
) {
|
||||
let cache = redis(
|
||||
vec![MockCmd::new(
|
||||
redis::cmd("SETEX").arg("tenant:key").arg(600).arg(wire),
|
||||
Ok("OK"),
|
||||
)],
|
||||
Some("tenant"),
|
||||
);
|
||||
cache
|
||||
.async_store(&request, response, Duration::from_secs(100))
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn invalid_entries_are_misses_and_disabled_reads_do_not_touch_redis(
|
||||
mut request: ResponseCacheRequest,
|
||||
) {
|
||||
let cache = redis(
|
||||
vec![MockCmd::new(
|
||||
redis::cmd("GET").arg("tenant:key"),
|
||||
Ok(b"invalid".to_vec()),
|
||||
)],
|
||||
None,
|
||||
);
|
||||
request.controls.no_cache = true;
|
||||
assert_eq!(cache.lookup(&request, Duration::ZERO).unwrap(), None);
|
||||
request.controls.no_cache = false;
|
||||
assert_eq!(
|
||||
cache.async_lookup(&request, Duration::ZERO).await.unwrap(),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn captured_service_keeps_the_selected_backend_for_background_writes(
|
||||
|
|
@ -541,27 +467,6 @@ async fn batch_lookup_reports_partial_hits_and_batch_store_populates_misses(
|
|||
assert_eq!(memory.lookup(&requests[2], now).unwrap(), None);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn batch_lookup_with_no_readable_request_skips_the_backend(
|
||||
#[values(false, true)] asynchronous: bool,
|
||||
) {
|
||||
let cache = redis(Vec::new(), None);
|
||||
let mut request = keyed("key");
|
||||
request.controls.no_cache = true;
|
||||
let requests = [request.clone(), request];
|
||||
let partial = if asynchronous {
|
||||
cache
|
||||
.async_lookup_batch(&requests, Duration::ZERO)
|
||||
.await
|
||||
.unwrap()
|
||||
} else {
|
||||
cache.lookup_batch(&requests, Duration::ZERO).unwrap()
|
||||
};
|
||||
assert_eq!(partial.values, vec![None, None]);
|
||||
assert_eq!(partial.missing_indices, vec![0, 1]);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn deferred_entries_keep_the_time_they_were_produced(
|
||||
|
|
|
|||
|
|
@ -1,15 +1,9 @@
|
|||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use litellm_cache_memory::InMemoryCache;
|
||||
use litellm_cache_redis::RedisCache;
|
||||
use litellm_cache_response::{
|
||||
CacheEntry, CacheKeyInput, ResponseCache, ResponseCacheCodec, ResponseCacheRequest,
|
||||
};
|
||||
use redis_test::{MockCmd, MockRedisConnection};
|
||||
use litellm_cache_response::{CacheEntry, CacheKeyInput, ResponseCache, ResponseCacheRequest};
|
||||
use rstest::fixture;
|
||||
|
||||
pub type MockedRedis = RedisCache<ResponseCacheCodec, MockRedisConnection>;
|
||||
|
||||
#[fixture]
|
||||
pub fn memory() -> Arc<ResponseCache<InMemoryCache<CacheEntry>>> {
|
||||
Arc::new(ResponseCache::new(Arc::new(InMemoryCache::new(
|
||||
|
|
@ -29,12 +23,3 @@ pub fn keyed(key: &str) -> ResponseCacheRequest {
|
|||
..Default::default()
|
||||
})
|
||||
}
|
||||
|
||||
/// A Redis response cache that must receive exactly `commands`, in order.
|
||||
pub fn redis(commands: Vec<MockCmd>, namespace: Option<&str>) -> ResponseCache<MockedRedis> {
|
||||
let connection = MockRedisConnection::new(commands).assert_all_commands_consumed();
|
||||
ResponseCache::new(Arc::new(
|
||||
RedisCache::with_connection(connection, None, ResponseCacheCodec)
|
||||
.with_namespace(namespace.map(str::to_owned)),
|
||||
))
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,22 +0,0 @@
|
|||
[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
|
||||
redis = { workspace = true, features = ["tls-rustls"] }
|
||||
sha2.workspace = true
|
||||
thiserror.workspace = true
|
||||
uuid = { version = "1", features = ["v4"] }
|
||||
|
||||
[dev-dependencies]
|
||||
litellm-error.workspace = true
|
||||
litellm-cache-testing.workspace = true
|
||||
redis-test.workspace = true
|
||||
rstest.workspace = true
|
||||
serde_json.workspace = true
|
||||
tokio.workspace = true
|
||||
|
|
@ -1,245 +0,0 @@
|
|||
use std::{
|
||||
sync::{Arc, Mutex},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use litellm_cache::{
|
||||
BaseCache, CacheCodec, Error, SemanticCacheContext,
|
||||
semantic::{Embedder, SemanticCache, SemanticLookup, prompt_from_context},
|
||||
};
|
||||
use litellm_cache_redis::{RedisTopology, connection::Connections};
|
||||
|
||||
use crate::{
|
||||
ValkeySemanticConfig,
|
||||
index::IndexState,
|
||||
search::{embedding_bytes, scope_tag, search_document, write_document},
|
||||
};
|
||||
|
||||
/// `ValkeySemanticCache`: a semantic cache on valkey-search's TAG + VECTOR index. Values go
|
||||
/// through the injected codec, so the response layer decides what a cached entry is.
|
||||
pub struct ValkeySemanticCache<E, S, C = redis::Connection> {
|
||||
connections: Arc<Connections<C>>,
|
||||
embedder: E,
|
||||
codec: S,
|
||||
config: ValkeySemanticConfig,
|
||||
index_dimension: Arc<Mutex<Option<usize>>>,
|
||||
}
|
||||
|
||||
impl<E: Embedder, S: CacheCodec> ValkeySemanticCache<E, S> {
|
||||
pub fn new(
|
||||
url: &str,
|
||||
embedder: E,
|
||||
codec: S,
|
||||
config: ValkeySemanticConfig,
|
||||
) -> Result<Self, Error> {
|
||||
Ok(Self {
|
||||
connections: Arc::new(Connections::open(url, &RedisTopology::Standalone)?),
|
||||
embedder,
|
||||
codec,
|
||||
config,
|
||||
index_dimension: Arc::new(Mutex::new(None)),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl<E, S, C> ValkeySemanticCache<E, S, C>
|
||||
where
|
||||
E: Embedder,
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
pub fn with_connection(
|
||||
connection: C,
|
||||
embedder: E,
|
||||
codec: S,
|
||||
config: ValkeySemanticConfig,
|
||||
) -> Self {
|
||||
Self {
|
||||
connections: Arc::new(Connections::fixed(connection)),
|
||||
embedder,
|
||||
codec,
|
||||
config,
|
||||
index_dimension: Arc::new(Mutex::new(None)),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn similarity_threshold(&self) -> f64 {
|
||||
self.config.similarity_threshold
|
||||
}
|
||||
|
||||
pub fn index_name(&self) -> &str {
|
||||
&self.config.index_name
|
||||
}
|
||||
|
||||
fn index_state(&self) -> IndexState {
|
||||
IndexState {
|
||||
name: self.config.index_name.clone(),
|
||||
prefix: format!("{}:", self.config.index_name),
|
||||
dimension: Arc::clone(&self.index_dimension),
|
||||
similarity_threshold: self.config.similarity_threshold,
|
||||
}
|
||||
}
|
||||
|
||||
fn decode(&self, lookup: SemanticLookup<Vec<u8>>) -> Result<SemanticLookup<S::Value>, Error> {
|
||||
Ok(SemanticLookup {
|
||||
value: lookup
|
||||
.value
|
||||
.map(|bytes| self.codec.decode(&bytes))
|
||||
.transpose()?,
|
||||
similarity: lookup.similarity,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl<E, S, C> ValkeySemanticCache<E, S, C>
|
||||
where
|
||||
E: Embedder,
|
||||
S: CacheCodec + Clone,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
/// The same index and connections behind a different embedder.
|
||||
pub fn with_embedder<E2: Embedder>(&self, embedder: E2) -> ValkeySemanticCache<E2, S, C> {
|
||||
ValkeySemanticCache {
|
||||
connections: Arc::clone(&self.connections),
|
||||
embedder,
|
||||
codec: self.codec.clone(),
|
||||
config: self.config.clone(),
|
||||
index_dimension: Arc::clone(&self.index_dimension),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<E, S, C> BaseCache for ValkeySemanticCache<E, S, C>
|
||||
where
|
||||
E: Embedder,
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
type Value = S::Value;
|
||||
type Context = SemanticCacheContext;
|
||||
|
||||
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
|
||||
context.ttl
|
||||
}
|
||||
|
||||
fn set_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
value: Self::Value,
|
||||
context: &Self::Context,
|
||||
) -> Result<(), Error> {
|
||||
let Some(prompt) = prompt_from_context(context) else {
|
||||
return Ok(());
|
||||
};
|
||||
let embedding = self.embedder.embed(&prompt, context.metadata.as_ref())?;
|
||||
let response = self.codec.encode(&value)?;
|
||||
let index = self.index_state();
|
||||
self.connections.execute(|connection| {
|
||||
write_document(
|
||||
connection,
|
||||
&index,
|
||||
&scope_tag(key),
|
||||
&prompt,
|
||||
response,
|
||||
embedding_bytes(&embedding),
|
||||
self.get_ttl(context),
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
fn get_cache(&self, key: &str, context: &Self::Context) -> Result<Option<Self::Value>, Error> {
|
||||
self.get_cache_with_similarity(key, context)
|
||||
.map(|lookup| lookup.value)
|
||||
}
|
||||
|
||||
async fn async_set_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
value: Self::Value,
|
||||
context: Self::Context,
|
||||
) -> Result<(), Error> {
|
||||
let Some(prompt) = prompt_from_context(&context) else {
|
||||
return Ok(());
|
||||
};
|
||||
let embedding = self
|
||||
.embedder
|
||||
.async_embed(&prompt, context.metadata.as_ref())
|
||||
.await?;
|
||||
let response = self.codec.encode(&value)?;
|
||||
let index = self.index_state();
|
||||
let scope = scope_tag(key);
|
||||
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
|
||||
write_document(
|
||||
connection,
|
||||
&index,
|
||||
&scope,
|
||||
&prompt,
|
||||
response,
|
||||
embedding_bytes(&embedding),
|
||||
context.ttl,
|
||||
)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
async fn async_get_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
context: &Self::Context,
|
||||
) -> Result<Option<Self::Value>, Error> {
|
||||
self.async_get_cache_with_similarity(key, context)
|
||||
.await
|
||||
.map(|lookup| lookup.value)
|
||||
}
|
||||
}
|
||||
|
||||
/// Python stamps a similarity of `0.0` when there is no prompt or no document in the key's
|
||||
/// scope, and the closest document's similarity even when it misses the threshold.
|
||||
impl<E, S, C> SemanticCache for ValkeySemanticCache<E, S, C>
|
||||
where
|
||||
E: Embedder,
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
fn get_cache_with_similarity(
|
||||
&self,
|
||||
key: &str,
|
||||
context: &Self::Context,
|
||||
) -> Result<SemanticLookup<Self::Value>, Error> {
|
||||
let Some(prompt) = prompt_from_context(context) else {
|
||||
return Ok(SemanticLookup::miss(Some(0.0)));
|
||||
};
|
||||
let embedding = self.embedder.embed(&prompt, context.metadata.as_ref())?;
|
||||
let index = self.index_state();
|
||||
let lookup = self.connections.execute(|connection| {
|
||||
search_document(
|
||||
connection,
|
||||
&index,
|
||||
&scope_tag(key),
|
||||
embedding_bytes(&embedding),
|
||||
)
|
||||
})?;
|
||||
self.decode(lookup)
|
||||
}
|
||||
|
||||
async fn async_get_cache_with_similarity(
|
||||
&self,
|
||||
key: &str,
|
||||
context: &Self::Context,
|
||||
) -> Result<SemanticLookup<Self::Value>, Error> {
|
||||
let Some(prompt) = prompt_from_context(context) else {
|
||||
return Ok(SemanticLookup::miss(Some(0.0)));
|
||||
};
|
||||
let embedding = self
|
||||
.embedder
|
||||
.async_embed(&prompt, context.metadata.as_ref())
|
||||
.await?;
|
||||
let index = self.index_state();
|
||||
let scope = scope_tag(key);
|
||||
let lookup = Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
|
||||
search_document(connection, &index, &scope, embedding_bytes(&embedding))
|
||||
})
|
||||
.await?;
|
||||
self.decode(lookup)
|
||||
}
|
||||
}
|
||||
|
|
@ -1,8 +0,0 @@
|
|||
/// `ValkeySemanticCache.DEFAULT_VALKEY_INDEX_NAME`.
|
||||
pub const DEFAULT_INDEX_NAME: &str = "litellm_semantic_cache_index";
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct ValkeySemanticConfig {
|
||||
pub similarity_threshold: f64,
|
||||
pub index_name: String,
|
||||
}
|
||||
|
|
@ -1,117 +0,0 @@
|
|||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use litellm_cache::Error;
|
||||
use litellm_cache_redis::connection::ConnectionRef;
|
||||
|
||||
use crate::search::value_text;
|
||||
|
||||
/// The valkey-search index one cache writes to, with the dimension it was last ensured for.
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct IndexState {
|
||||
pub(crate) name: String,
|
||||
pub(crate) prefix: String,
|
||||
pub(crate) dimension: Arc<Mutex<Option<usize>>>,
|
||||
pub(crate) similarity_threshold: f64,
|
||||
}
|
||||
|
||||
/// `_ensure_index_sync` / `_ensure_index_async`: create the TAG + HNSW index once per dimension,
|
||||
/// and accept an existing index unless it reports a different dimension.
|
||||
pub(crate) fn ensure_index(
|
||||
connection: &mut ConnectionRef<'_>,
|
||||
index: &IndexState,
|
||||
dimension: usize,
|
||||
) -> Result<(), Error> {
|
||||
if index
|
||||
.dimension
|
||||
.lock()?
|
||||
.is_some_and(|existing| existing == dimension)
|
||||
{
|
||||
return Ok(());
|
||||
}
|
||||
let create = redis::cmd("FT.CREATE")
|
||||
.arg(&index.name)
|
||||
.arg("ON")
|
||||
.arg("HASH")
|
||||
.arg("PREFIX")
|
||||
.arg(1)
|
||||
.arg(&index.prefix)
|
||||
.arg("SCHEMA")
|
||||
.arg("litellm_cache_key")
|
||||
.arg("TAG")
|
||||
.arg("embedding")
|
||||
.arg("VECTOR")
|
||||
.arg("HNSW")
|
||||
.arg(6)
|
||||
.arg("TYPE")
|
||||
.arg("FLOAT32")
|
||||
.arg("DIM")
|
||||
.arg(dimension)
|
||||
.arg("DISTANCE_METRIC")
|
||||
.arg("COSINE")
|
||||
.query::<String>(connection)
|
||||
.map(|_| ());
|
||||
if let Err(error) = create {
|
||||
if !index_exists(&error) {
|
||||
return Err(Error::unavailable(error));
|
||||
}
|
||||
let info = redis::cmd("FT.INFO")
|
||||
.arg(&index.name)
|
||||
.query::<redis::Value>(connection)
|
||||
.map_err(Error::unavailable)?;
|
||||
if let Some(existing) =
|
||||
index_dimension_from_info(&info).filter(|existing| *existing != dimension)
|
||||
{
|
||||
return Err(Error::unavailable(DimensionMismatch {
|
||||
existing,
|
||||
dimension,
|
||||
}));
|
||||
}
|
||||
}
|
||||
*index.dimension.lock()? = Some(dimension);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn index_exists(error: &redis::RedisError) -> bool {
|
||||
error.code().is_some()
|
||||
&& error
|
||||
.detail()
|
||||
.is_some_and(|detail| detail.to_ascii_lowercase().contains("already exists"))
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
#[error("index exists with dimension {existing}, not {dimension}")]
|
||||
struct DimensionMismatch {
|
||||
existing: usize,
|
||||
dimension: usize,
|
||||
}
|
||||
|
||||
/// `_extract_index_dim`: flatten each attribute one level and read the value after
|
||||
/// `dimensions`.
|
||||
fn index_dimension_from_info(value: &redis::Value) -> Option<usize> {
|
||||
let redis::Value::Array(values) = value else {
|
||||
return None;
|
||||
};
|
||||
let attributes = values.windows(2).find_map(|pair| {
|
||||
(value_text(&pair[0]).as_deref() == Some("attributes")).then_some(&pair[1])
|
||||
})?;
|
||||
let redis::Value::Array(fields) = attributes else {
|
||||
return None;
|
||||
};
|
||||
fields.iter().find_map(|field| {
|
||||
let redis::Value::Array(values) = field else {
|
||||
return None;
|
||||
};
|
||||
let values = values
|
||||
.iter()
|
||||
.flat_map(|value| match value {
|
||||
redis::Value::Array(values) => values.as_slice(),
|
||||
_ => std::slice::from_ref(value),
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
values.windows(2).find_map(|pair| {
|
||||
(value_text(pair[0]).as_deref() == Some("dimensions"))
|
||||
.then(|| value_text(pair[1]).and_then(|value| value.parse().ok()))
|
||||
.flatten()
|
||||
})
|
||||
})
|
||||
}
|
||||
|
|
@ -1,7 +0,0 @@
|
|||
mod cache;
|
||||
mod config;
|
||||
mod index;
|
||||
mod search;
|
||||
|
||||
pub use cache::ValkeySemanticCache;
|
||||
pub use config::{DEFAULT_INDEX_NAME, ValkeySemanticConfig};
|
||||
|
|
@ -1,161 +0,0 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use litellm_cache::{Error, semantic::SemanticLookup};
|
||||
use litellm_cache_redis::connection::ConnectionRef;
|
||||
use sha2::{Digest, Sha256};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::index::{IndexState, ensure_index};
|
||||
|
||||
/// `_scope_tag`: valkey-search TAG fields cannot match arbitrary keys verbatim, so scopes are
|
||||
/// the key's lowercase SHA-256.
|
||||
pub(crate) fn scope_tag(key: &str) -> String {
|
||||
let digest = Sha256::digest(key.as_bytes());
|
||||
digest.iter().map(|byte| format!("{byte:02x}")).collect()
|
||||
}
|
||||
|
||||
pub(crate) fn embedding_bytes(embedding: &[f32]) -> Vec<u8> {
|
||||
embedding
|
||||
.iter()
|
||||
.flat_map(|value| value.to_le_bytes())
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// `HSET` a fresh `<prefix><scope>:<uuid4>` document, then `EXPIRE` it when a TTL is set.
|
||||
pub(crate) fn write_document(
|
||||
connection: &mut ConnectionRef<'_>,
|
||||
index: &IndexState,
|
||||
scope: &str,
|
||||
prompt: &str,
|
||||
response: Vec<u8>,
|
||||
vector: Vec<u8>,
|
||||
ttl: Option<Duration>,
|
||||
) -> Result<(), Error> {
|
||||
ensure_index(connection, index, vector.len() / size_of::<f32>())?;
|
||||
let document = format!("{}{scope}:{}", index.prefix, Uuid::new_v4());
|
||||
let mut pipeline = redis::pipe();
|
||||
pipeline
|
||||
.cmd("HSET")
|
||||
.arg(&document)
|
||||
.arg("litellm_cache_key")
|
||||
.arg(scope)
|
||||
.arg("prompt")
|
||||
.arg(prompt)
|
||||
.arg("response")
|
||||
.arg(response)
|
||||
.arg("embedding")
|
||||
.arg(vector)
|
||||
.ignore();
|
||||
if let Some(ttl) = ttl {
|
||||
pipeline
|
||||
.cmd("EXPIRE")
|
||||
.arg(&document)
|
||||
.arg(ttl.as_secs())
|
||||
.ignore();
|
||||
}
|
||||
pipeline.query::<()>(connection).map_err(Error::unavailable)
|
||||
}
|
||||
|
||||
/// The KNN-1 search within `scope`: the closest document's similarity, and its stored response
|
||||
/// when that similarity reaches the threshold. No document reads as a similarity of `0.0`.
|
||||
pub(crate) fn search_document(
|
||||
connection: &mut ConnectionRef<'_>,
|
||||
index: &IndexState,
|
||||
scope: &str,
|
||||
vector: Vec<u8>,
|
||||
) -> Result<SemanticLookup<Vec<u8>>, Error> {
|
||||
ensure_index(connection, index, vector.len() / size_of::<f32>())?;
|
||||
let query =
|
||||
format!("(@litellm_cache_key:{{{scope}}})=>[KNN 1 @embedding $vec AS vector_distance]");
|
||||
let response = redis::cmd("FT.SEARCH")
|
||||
.arg(&index.name)
|
||||
.arg(query)
|
||||
.arg("PARAMS")
|
||||
.arg(2)
|
||||
.arg("vec")
|
||||
.arg(vector)
|
||||
.arg("RETURN")
|
||||
.arg(2)
|
||||
.arg("response")
|
||||
.arg("vector_distance")
|
||||
.arg("DIALECT")
|
||||
.arg(2)
|
||||
.query::<redis::Value>(connection)
|
||||
.map_err(Error::unavailable)?;
|
||||
let Some(fields) = search_fields(response)? else {
|
||||
return Ok(SemanticLookup::miss(Some(0.0)));
|
||||
};
|
||||
let field = |name: &str| {
|
||||
fields
|
||||
.iter()
|
||||
.find_map(|(field, value)| (field == name).then(|| value.clone()))
|
||||
.ok_or(Error::InvalidEntry(None))
|
||||
};
|
||||
let response = field("response")?;
|
||||
let similarity = 1.0 - parse_f64(&field("vector_distance")?)?;
|
||||
Ok(SemanticLookup {
|
||||
value: (similarity >= index.similarity_threshold).then_some(response),
|
||||
similarity: Some(similarity),
|
||||
})
|
||||
}
|
||||
|
||||
type SearchFields = Vec<(String, Vec<u8>)>;
|
||||
|
||||
fn search_fields(value: redis::Value) -> Result<Option<SearchFields>, Error> {
|
||||
let redis::Value::Array(values) = value else {
|
||||
return Err(Error::InvalidEntry(None));
|
||||
};
|
||||
let total = parse_i64(values.first().ok_or(Error::InvalidEntry(None))?)?;
|
||||
if total <= 0 || values.len() < 3 {
|
||||
return Ok(None);
|
||||
}
|
||||
let redis::Value::Array(fields) = &values[2] else {
|
||||
return Err(Error::InvalidEntry(None));
|
||||
};
|
||||
let (pairs, remainder) = fields.as_chunks::<2>();
|
||||
if !remainder.is_empty() {
|
||||
return Err(Error::InvalidEntry(None));
|
||||
}
|
||||
pairs
|
||||
.iter()
|
||||
.map(|pair| {
|
||||
Ok((
|
||||
value_text(&pair[0]).ok_or(Error::InvalidEntry(None))?,
|
||||
value_bytes(&pair[1])?,
|
||||
))
|
||||
})
|
||||
.collect::<Result<Vec<_>, Error>>()
|
||||
.map(Some)
|
||||
}
|
||||
|
||||
fn parse_i64(value: &redis::Value) -> Result<i64, Error> {
|
||||
value_text(value)
|
||||
.ok_or(Error::InvalidEntry(None))?
|
||||
.parse()
|
||||
.map_err(Error::invalid_entry)
|
||||
}
|
||||
|
||||
fn parse_f64(value: &[u8]) -> Result<f64, Error> {
|
||||
std::str::from_utf8(value)
|
||||
.map_err(Error::invalid_entry)?
|
||||
.parse()
|
||||
.map_err(Error::invalid_entry)
|
||||
}
|
||||
|
||||
pub(crate) fn value_text(value: &redis::Value) -> Option<String> {
|
||||
match value {
|
||||
redis::Value::BulkString(bytes) => String::from_utf8(bytes.clone()).ok(),
|
||||
redis::Value::SimpleString(value) => Some(value.clone()),
|
||||
redis::Value::Int(value) => Some(value.to_string()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn value_bytes(value: &redis::Value) -> Result<Vec<u8>, Error> {
|
||||
match value {
|
||||
redis::Value::BulkString(bytes) => Ok(bytes.clone()),
|
||||
redis::Value::SimpleString(value) => Ok(value.as_bytes().to_vec()),
|
||||
redis::Value::Int(value) => Ok(value.to_string().into_bytes()),
|
||||
_ => Err(Error::InvalidEntry(None)),
|
||||
}
|
||||
}
|
||||
|
|
@ -1,430 +0,0 @@
|
|||
mod support;
|
||||
|
||||
use std::{
|
||||
sync::{Arc, Mutex},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use litellm_cache::{
|
||||
BaseCache, Error, JsonCodec, SemanticCacheContext,
|
||||
semantic::{PreparedEmbedding, SemanticCache, SemanticLookup},
|
||||
};
|
||||
use litellm_cache_valkey_semantic::{
|
||||
DEFAULT_INDEX_NAME, ValkeySemanticCache, ValkeySemanticConfig,
|
||||
};
|
||||
use redis_test::MockRedisConnection;
|
||||
use rstest::{fixture, rstest};
|
||||
use serde_json::{Value, json};
|
||||
use support::{EmbedCalls, FakeEmbedder, RecordingConnection};
|
||||
|
||||
type Requests = Arc<Mutex<Vec<Vec<u8>>>>;
|
||||
type RecordingCache = ValkeySemanticCache<FakeEmbedder, JsonCodec<Value>, RecordingConnection>;
|
||||
|
||||
const KEY_SCOPE: &str = "2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683";
|
||||
|
||||
struct Recording {
|
||||
cache: RecordingCache,
|
||||
requests: Requests,
|
||||
calls: EmbedCalls,
|
||||
}
|
||||
|
||||
impl Recording {
|
||||
fn text(&self) -> String {
|
||||
self.requests
|
||||
.lock()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.map(|request| String::from_utf8_lossy(request).into_owned())
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n")
|
||||
}
|
||||
}
|
||||
|
||||
fn config() -> ValkeySemanticConfig {
|
||||
ValkeySemanticConfig {
|
||||
similarity_threshold: 0.8,
|
||||
index_name: "test".into(),
|
||||
}
|
||||
}
|
||||
|
||||
fn recording(replies: impl IntoIterator<Item = redis::RedisResult<redis::Value>>) -> Recording {
|
||||
let connection = RecordingConnection::new(replies);
|
||||
let requests = connection.requests();
|
||||
let embedder = FakeEmbedder::new(&[]);
|
||||
let calls = Arc::clone(&embedder.calls);
|
||||
Recording {
|
||||
cache: ValkeySemanticCache::with_connection(
|
||||
connection,
|
||||
embedder,
|
||||
JsonCodec::new(),
|
||||
config(),
|
||||
),
|
||||
requests,
|
||||
calls,
|
||||
}
|
||||
}
|
||||
|
||||
#[fixture]
|
||||
fn entry() -> Value {
|
||||
json!({"timestamp": 1.0, "response": {"answer": "ok"}})
|
||||
}
|
||||
|
||||
#[fixture]
|
||||
fn context() -> SemanticCacheContext {
|
||||
SemanticCacheContext {
|
||||
messages: Some(json!([{"role": "user", "content": "hello"}])),
|
||||
metadata: Some(json!({"source": "test"})),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
fn ok() -> redis::RedisResult<redis::Value> {
|
||||
Ok(redis::Value::SimpleString("OK".into()))
|
||||
}
|
||||
|
||||
fn already_exists() -> redis::RedisResult<redis::Value> {
|
||||
Err(support::server_error("Index already exists"))
|
||||
}
|
||||
|
||||
fn bulk(value: &[u8]) -> redis::Value {
|
||||
redis::Value::BulkString(value.to_vec())
|
||||
}
|
||||
|
||||
fn search_hit(response: &[u8], distance: &str) -> redis::Value {
|
||||
redis::Value::Array(vec![
|
||||
redis::Value::Int(1),
|
||||
bulk(b"test:document"),
|
||||
redis::Value::Array(vec![
|
||||
bulk(b"response"),
|
||||
bulk(response),
|
||||
bulk(b"vector_distance"),
|
||||
bulk(distance.as_bytes()),
|
||||
]),
|
||||
])
|
||||
}
|
||||
|
||||
fn encoded(value: &Value) -> Vec<u8> {
|
||||
serde_json::to_vec(value).unwrap()
|
||||
}
|
||||
|
||||
/// `FT.INFO` with the vector field's dimension nested one level down, as valkey-search reports.
|
||||
fn nested_dimension_info(dimension: i64) -> redis::Value {
|
||||
redis::Value::Array(vec![
|
||||
redis::Value::SimpleString("attributes".into()),
|
||||
redis::Value::Array(vec![redis::Value::Array(vec![
|
||||
redis::Value::SimpleString("embedding".into()),
|
||||
redis::Value::Array(vec![
|
||||
redis::Value::SimpleString("dimensions".into()),
|
||||
redis::Value::Int(dimension),
|
||||
]),
|
||||
])]),
|
||||
])
|
||||
}
|
||||
|
||||
/// `FT.INFO` with `dimensions` as a sibling string of the identifier.
|
||||
fn flat_dimension_info(dimension: &str) -> redis::Value {
|
||||
redis::Value::Array(vec![
|
||||
redis::Value::SimpleString("attributes".into()),
|
||||
redis::Value::Array(vec![redis::Value::Array(vec![
|
||||
redis::Value::SimpleString("identifier".into()),
|
||||
redis::Value::SimpleString("embedding".into()),
|
||||
redis::Value::Array(vec![
|
||||
redis::Value::SimpleString("dimensions".into()),
|
||||
redis::Value::SimpleString(dimension.into()),
|
||||
]),
|
||||
])]),
|
||||
])
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::string_content(json!([{"content": "hello"}]), None, Some("hello"))]
|
||||
#[case::text_parts(json!([{"content": [{"text": "hello"}, {"text": " world"}]}]), None, Some("hello world"))]
|
||||
#[case::non_object_parts_are_skipped(json!([{"content": ["raw", {"text": "hello"}]}]), None, Some("hello"))]
|
||||
#[case::search_results(json!([{"search_results": [{"source": "s", "title": "t", "content": [{"text": "c"}], "citations": ["x"]}]}]), None, Some(r#"stc["x"]"#))]
|
||||
#[case::responses_string_input(json!([]), Some(json!(" hello ")), Some("hello"))]
|
||||
#[case::responses_item_input(json!([]), Some(json!([{"content": "first"}, {"text": "second"}])), Some("first\nsecond"))]
|
||||
#[case::blank_input(json!([]), Some(json!(" ")), None)]
|
||||
fn prompt_shapes_follow_redis_semantic_extraction(
|
||||
#[case] messages: Value,
|
||||
#[case] input: Option<Value>,
|
||||
#[case] expected: Option<&str>,
|
||||
) {
|
||||
let recording = recording([ok(), Ok(redis::Value::Array(vec![redis::Value::Int(0)]))]);
|
||||
let context = SemanticCacheContext {
|
||||
messages: Some(messages),
|
||||
input,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
assert_eq!(recording.cache.get_cache("key", &context).unwrap(), None);
|
||||
|
||||
let prompts = recording
|
||||
.calls
|
||||
.lock()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.map(|(prompt, _)| prompt.clone())
|
||||
.collect::<Vec<_>>();
|
||||
assert_eq!(
|
||||
prompts,
|
||||
expected.into_iter().map(str::to_owned).collect::<Vec<_>>()
|
||||
);
|
||||
assert_eq!(
|
||||
recording.requests.lock().unwrap().is_empty(),
|
||||
expected.is_none()
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::key("key", KEY_SCOPE)]
|
||||
#[case::empty_key("", "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855")]
|
||||
fn documents_are_scoped_by_the_keys_sha256(
|
||||
#[case] key: &str,
|
||||
#[case] scope: &str,
|
||||
entry: Value,
|
||||
context: SemanticCacheContext,
|
||||
) {
|
||||
let recording = recording([ok()]);
|
||||
|
||||
recording.cache.set_cache(key, entry, &context).unwrap();
|
||||
|
||||
let text = recording.text();
|
||||
assert!(text.contains(&format!("test:{scope}:")));
|
||||
assert!(text.contains(&format!("litellm_cache_key\r\n$64\r\n{scope}\r\n")));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::no_ttl(None, None)]
|
||||
#[case::whole_seconds(Some(Duration::from_secs(5)), Some("5"))]
|
||||
#[case::fractional_seconds_truncate(Some(Duration::from_millis(1900)), Some("1"))]
|
||||
fn set_writes_hset_and_expires_only_with_a_ttl(
|
||||
#[case] ttl: Option<Duration>,
|
||||
#[case] expire: Option<&str>,
|
||||
entry: Value,
|
||||
context: SemanticCacheContext,
|
||||
) {
|
||||
let recording = recording([ok()]);
|
||||
|
||||
recording
|
||||
.cache
|
||||
.set_cache(
|
||||
"key",
|
||||
entry,
|
||||
&SemanticCacheContext {
|
||||
ttl,
|
||||
..context.clone()
|
||||
},
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let text = recording.text();
|
||||
assert!(text.contains("FT.CREATE"));
|
||||
assert!(text.contains("HSET"));
|
||||
let expire_seconds = text
|
||||
.split_once("EXPIRE\r\n")
|
||||
.and_then(|(_, rest)| rest.split("\r\n").nth(3));
|
||||
assert_eq!(expire_seconds, expire);
|
||||
assert_eq!(
|
||||
*recording.calls.lock().unwrap(),
|
||||
vec![("hello".to_owned(), context.metadata)]
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn second_set_skips_create_after_dimension_is_cached(entry: Value, context: SemanticCacheContext) {
|
||||
let recording = recording([ok()]);
|
||||
|
||||
recording
|
||||
.cache
|
||||
.set_cache("key", entry.clone(), &context)
|
||||
.unwrap();
|
||||
recording.cache.set_cache("key", entry, &context).unwrap();
|
||||
|
||||
let text = recording.text();
|
||||
assert_eq!(text.matches("FT.CREATE").count(), 1);
|
||||
assert_eq!(text.matches("HSET").count(), 2);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::nested_matching(nested_dimension_info(3), Ok(()))]
|
||||
#[case::nested_mismatch(
|
||||
nested_dimension_info(2),
|
||||
Err("cache is unavailable: index exists with dimension 2, not 3")
|
||||
)]
|
||||
#[case::flat_matching(flat_dimension_info("3"), Ok(()))]
|
||||
#[case::flat_mismatch(
|
||||
flat_dimension_info("2"),
|
||||
Err("cache is unavailable: index exists with dimension 2, not 3")
|
||||
)]
|
||||
#[case::unreported_dimension_is_accepted(redis::Value::Array(vec![]), Ok(()))]
|
||||
fn existing_index_dimension_must_match_embedding(
|
||||
#[case] info: redis::Value,
|
||||
#[case] expected: Result<(), &str>,
|
||||
entry: Value,
|
||||
context: SemanticCacheContext,
|
||||
) {
|
||||
let recording = recording([already_exists(), Ok(info)]);
|
||||
|
||||
let result = recording.cache.set_cache("key", entry, &context);
|
||||
assert_eq!(
|
||||
result.map_err(|error| litellm_error::Report::new(&error).to_string()),
|
||||
expected.map_err(str::to_owned)
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::create_failure(Err(redis::RedisError::from((redis::ErrorKind::Io, "boom"))))]
|
||||
fn index_creation_failures_are_unavailable(
|
||||
#[case] reply: redis::RedisResult<redis::Value>,
|
||||
entry: Value,
|
||||
context: SemanticCacheContext,
|
||||
) {
|
||||
let recording = recording([reply]);
|
||||
|
||||
assert!(matches!(
|
||||
recording.cache.set_cache("key", entry, &context),
|
||||
Err(Error::Unavailable(Some(_)))
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::within_threshold(search_hit(&encoded(&entry()), "0.1"), Ok(Some(entry())))]
|
||||
#[case::at_threshold(search_hit(&encoded(&entry()), "0.2"), Ok(Some(entry())))]
|
||||
#[case::beyond_threshold(search_hit(&encoded(&entry()), "0.5"), Ok(None))]
|
||||
#[case::zero_documents(redis::Value::Array(vec![redis::Value::Int(0)]), Ok(None))]
|
||||
#[case::missing_response(
|
||||
redis::Value::Array(vec![
|
||||
redis::Value::Int(1),
|
||||
bulk(b"document"),
|
||||
redis::Value::Array(vec![bulk(b"vector_distance"), bulk(b"0.1")]),
|
||||
]),
|
||||
Err(Error::InvalidEntry(None))
|
||||
)]
|
||||
#[case::unparsable_distance(search_hit(b"not-json", "abc"), Err(Error::InvalidEntry(None)))]
|
||||
#[case::undecodable_response(search_hit(b"not-json", "0.1"), Err(Error::InvalidEntry(None)))]
|
||||
fn get_applies_threshold_and_decodes_entry(
|
||||
#[case] reply: redis::Value,
|
||||
#[case] expected: Result<Option<Value>, Error>,
|
||||
context: SemanticCacheContext,
|
||||
) {
|
||||
let recording = recording([ok(), Ok(reply)]);
|
||||
|
||||
let kind = |result: Result<Option<Value>, Error>| {
|
||||
result.map_err(|error| std::mem::discriminant(&error))
|
||||
};
|
||||
assert_eq!(
|
||||
kind(recording.cache.get_cache("key", &context)),
|
||||
kind(expected)
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::hit(context(), Some(search_hit(&encoded(&entry()), "0.1")), Some(entry()), Some(1.0 - 0.1))]
|
||||
#[case::below_threshold(context(), Some(search_hit(&encoded(&entry()), "0.5")), None, Some(1.0 - 0.5))]
|
||||
#[case::no_results(context(), Some(redis::Value::Array(vec![redis::Value::Int(0)])), None, Some(0.0))]
|
||||
#[case::no_prompt(SemanticCacheContext::default(), None, None, Some(0.0))]
|
||||
#[tokio::test]
|
||||
async fn lookup_reports_python_semantic_similarity(
|
||||
#[case] context: SemanticCacheContext,
|
||||
#[case] reply: Option<redis::Value>,
|
||||
#[case] value: Option<Value>,
|
||||
#[case] similarity: Option<f64>,
|
||||
#[values(false, true)] use_async: bool,
|
||||
) {
|
||||
let searched = reply.is_some();
|
||||
let recording = recording(reply.map_or_else(Vec::new, |reply| vec![ok(), Ok(reply)]));
|
||||
|
||||
let lookup = if use_async {
|
||||
recording
|
||||
.cache
|
||||
.async_get_cache_with_similarity("key", &context)
|
||||
.await
|
||||
} else {
|
||||
recording.cache.get_cache_with_similarity("key", &context)
|
||||
};
|
||||
|
||||
assert_eq!(lookup, Ok(SemanticLookup { value, similarity }));
|
||||
assert_eq!(recording.text().contains("FT.SEARCH"), searched);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn missing_prompt_does_not_touch_valkey(entry: Value) {
|
||||
let cache = ValkeySemanticCache::with_connection(
|
||||
MockRedisConnection::new([]).assert_all_commands_consumed(),
|
||||
FakeEmbedder::new(&[]),
|
||||
JsonCodec::new(),
|
||||
config(),
|
||||
);
|
||||
let context = SemanticCacheContext::default();
|
||||
|
||||
cache.set_cache("key", entry.clone(), &context).unwrap();
|
||||
assert_eq!(cache.get_cache("key", &context).unwrap(), None);
|
||||
cache
|
||||
.async_set_cache("key", entry, context.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(cache.async_get_cache("key", &context).await.unwrap(), None);
|
||||
assert_eq!(cache.get_ttl(&context), None);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn with_embedder_shares_index_state_and_connections(entry: Value, context: SemanticCacheContext) {
|
||||
let recording = recording([ok(), ok(), Ok(search_hit(&encoded(&entry), "0.1"))]);
|
||||
recording
|
||||
.cache
|
||||
.set_cache("key", entry.clone(), &context)
|
||||
.unwrap();
|
||||
|
||||
let prepared = recording
|
||||
.cache
|
||||
.with_embedder(PreparedEmbedding(vec![0.1, 0.2, 0.3]));
|
||||
|
||||
assert_eq!(prepared.get_cache("key", &context).unwrap(), Some(entry));
|
||||
assert_eq!(recording.text().matches("FT.CREATE").count(), 1);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn accessors_report_the_config() {
|
||||
let recording = recording([]);
|
||||
|
||||
assert_eq!(recording.cache.index_name(), "test");
|
||||
assert_eq!(recording.cache.similarity_threshold(), 0.8);
|
||||
assert_eq!(DEFAULT_INDEX_NAME, "litellm_semantic_cache_index");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn async_set_and_get_use_shared_document_helpers(
|
||||
entry: Value,
|
||||
context: SemanticCacheContext,
|
||||
) {
|
||||
let recording = recording([ok(), ok(), ok(), Ok(search_hit(&encoded(&entry), "0.1"))]);
|
||||
let context = SemanticCacheContext {
|
||||
ttl: Some(Duration::from_millis(1900)),
|
||||
..context
|
||||
};
|
||||
|
||||
recording
|
||||
.cache
|
||||
.async_set_cache("key", entry.clone(), context.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
recording
|
||||
.cache
|
||||
.async_get_cache("key", &context)
|
||||
.await
|
||||
.unwrap(),
|
||||
Some(entry)
|
||||
);
|
||||
|
||||
let text = recording.text();
|
||||
assert!(text.contains("FT.CREATE"));
|
||||
assert!(text.contains("HSET"));
|
||||
assert!(text.contains("EXPIRE"));
|
||||
assert_eq!(
|
||||
*recording.calls.lock().unwrap(),
|
||||
vec![("hello".to_owned(), context.metadata.clone()); 2]
|
||||
);
|
||||
}
|
||||
|
|
@ -1,62 +0,0 @@
|
|||
//! `overwrite_replaces` does not apply: like Python, every write is a new `<scope>:<uuid4>`
|
||||
//! document, so a second write with the same prompt adds a tie instead of replacing the first.
|
||||
|
||||
mod support;
|
||||
|
||||
use litellm_cache::{JsonCodec, SemanticCacheContext, semantic::PreparedEmbedding};
|
||||
use litellm_cache_testing as contract;
|
||||
use litellm_cache_valkey_semantic::{
|
||||
DEFAULT_INDEX_NAME, ValkeySemanticCache, ValkeySemanticConfig,
|
||||
};
|
||||
use rstest::{fixture, rstest};
|
||||
use serde_json::{Value, json};
|
||||
use support::FakeSearch;
|
||||
|
||||
type Cache = ValkeySemanticCache<PreparedEmbedding, JsonCodec<Value>, FakeSearch>;
|
||||
|
||||
const PREFIX: &str = "contract:";
|
||||
|
||||
#[fixture]
|
||||
fn cache() -> Cache {
|
||||
ValkeySemanticCache::with_connection(
|
||||
FakeSearch::default(),
|
||||
PreparedEmbedding(vec![0.6, 0.8]),
|
||||
JsonCodec::new(),
|
||||
ValkeySemanticConfig {
|
||||
similarity_threshold: 0.9,
|
||||
index_name: DEFAULT_INDEX_NAME.into(),
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
#[fixture]
|
||||
fn context() -> SemanticCacheContext {
|
||||
SemanticCacheContext {
|
||||
messages: Some(json!([{"role": "user", "content": "contract prompt"}])),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn hit_and_miss(cache: Cache, context: SemanticCacheContext) {
|
||||
contract::hit_and_miss(&cache, context, PREFIX, json!({"answer": 42})).await;
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn sync_async_equivalence(cache: Cache, context: SemanticCacheContext) {
|
||||
contract::sync_async_equivalence(&cache, context, PREFIX, json!("first"), json!([2])).await;
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn pipeline_writes_every_entry(cache: Cache, context: SemanticCacheContext) {
|
||||
contract::pipeline_writes_every_entry(
|
||||
&cache,
|
||||
context,
|
||||
PREFIX,
|
||||
vec![json!("a"), json!(2), json!({"c": true})],
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
|
@ -1,352 +0,0 @@
|
|||
#![allow(dead_code)]
|
||||
|
||||
use std::{
|
||||
collections::{BTreeMap, HashMap, VecDeque},
|
||||
sync::{Arc, Mutex},
|
||||
};
|
||||
|
||||
use litellm_cache::{Error, semantic::Embedder};
|
||||
use serde_json::Value;
|
||||
|
||||
pub type EmbedCalls = Arc<Mutex<Vec<(String, Option<Value>)>>>;
|
||||
|
||||
/// Embeds known prompts to fixed vectors, anything else to `[0.1, 0.2, 0.3]`, and records every
|
||||
/// prompt with its metadata.
|
||||
pub struct FakeEmbedder {
|
||||
vectors: HashMap<String, Vec<f32>>,
|
||||
pub calls: EmbedCalls,
|
||||
}
|
||||
|
||||
impl FakeEmbedder {
|
||||
pub fn new(vectors: &[(&str, &[f32])]) -> Self {
|
||||
Self {
|
||||
vectors: vectors
|
||||
.iter()
|
||||
.map(|(prompt, vector)| ((*prompt).to_owned(), vector.to_vec()))
|
||||
.collect(),
|
||||
calls: EmbedCalls::default(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Embedder for FakeEmbedder {
|
||||
fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result<Vec<f32>, Error> {
|
||||
self.calls
|
||||
.lock()
|
||||
.unwrap()
|
||||
.push((prompt.to_owned(), metadata.cloned()));
|
||||
Ok(self
|
||||
.vectors
|
||||
.get(prompt)
|
||||
.cloned()
|
||||
.unwrap_or_else(|| vec![0.1, 0.2, 0.3]))
|
||||
}
|
||||
|
||||
async fn async_embed(&self, prompt: &str, metadata: Option<&Value>) -> Result<Vec<f32>, Error> {
|
||||
self.embed(prompt, metadata)
|
||||
}
|
||||
}
|
||||
|
||||
struct FakeIndex {
|
||||
prefix: Vec<u8>,
|
||||
dims: usize,
|
||||
vector_field: String,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct SearchState {
|
||||
indexes: HashMap<String, FakeIndex>,
|
||||
hashes: BTreeMap<Vec<u8>, BTreeMap<String, Vec<u8>>>,
|
||||
}
|
||||
|
||||
/// An in-memory valkey-search speaking the `FT.*`, `HSET` and `EXPIRE` subset the semantic cache
|
||||
/// sends, with exact cosine KNN over the hashes under an index prefix.
|
||||
#[derive(Clone, Default)]
|
||||
pub struct FakeSearch {
|
||||
state: Arc<Mutex<SearchState>>,
|
||||
}
|
||||
|
||||
impl FakeSearch {
|
||||
fn run(&self, args: Vec<Vec<u8>>) -> redis::RedisResult<redis::Value> {
|
||||
let mut state = self.state.lock().unwrap();
|
||||
let text = |index: usize| String::from_utf8_lossy(&args[index]).into_owned();
|
||||
match text(0).to_uppercase().as_str() {
|
||||
"FT.CREATE" => {
|
||||
let name = text(1);
|
||||
if state.indexes.contains_key(&name) {
|
||||
return Err(server_error("Index already exists"));
|
||||
}
|
||||
let position = |token: &str| args.iter().position(|arg| arg == token.as_bytes());
|
||||
let prefix = args[position("PREFIX").unwrap() + 2].clone();
|
||||
let dims = text(position("DIM").unwrap() + 1).parse().unwrap();
|
||||
let vector_field = text(position("VECTOR").unwrap() - 1);
|
||||
state.indexes.insert(
|
||||
name,
|
||||
FakeIndex {
|
||||
prefix,
|
||||
dims,
|
||||
vector_field,
|
||||
},
|
||||
);
|
||||
Ok(redis::Value::Okay)
|
||||
}
|
||||
"FT.INFO" => {
|
||||
let index = state
|
||||
.indexes
|
||||
.get(&text(1))
|
||||
.ok_or_else(|| server_error("Unknown index name"))?;
|
||||
Ok(index_info(index))
|
||||
}
|
||||
"FT.DROPINDEX" => {
|
||||
state.indexes.remove(&text(1));
|
||||
Ok(redis::Value::Okay)
|
||||
}
|
||||
"HSET" => {
|
||||
let hash = state.hashes.entry(args[1].clone()).or_default();
|
||||
for pair in args[2..].chunks(2) {
|
||||
hash.insert(
|
||||
String::from_utf8_lossy(&pair[0]).into_owned(),
|
||||
pair[1].clone(),
|
||||
);
|
||||
}
|
||||
Ok(redis::Value::Int(((args.len() - 2) / 2) as i64))
|
||||
}
|
||||
"EXPIRE" => Ok(redis::Value::Int(i64::from(
|
||||
state.hashes.contains_key(&args[1]),
|
||||
))),
|
||||
"FT.SEARCH" => {
|
||||
let index = state
|
||||
.indexes
|
||||
.get(&text(1))
|
||||
.ok_or_else(|| server_error("no such index"))?;
|
||||
let query = text(2);
|
||||
let tag = query_tag(&query);
|
||||
let params = args.iter().position(|arg| arg == b"PARAMS").unwrap();
|
||||
let vector = floats(&args[params + 3]);
|
||||
let best = state
|
||||
.hashes
|
||||
.iter()
|
||||
.filter(|(key, _)| key.starts_with(&index.prefix))
|
||||
.filter(|(_, fields)| {
|
||||
fields.get("litellm_cache_key").map(Vec::as_slice) == Some(tag.as_bytes())
|
||||
})
|
||||
.filter_map(|(key, fields)| {
|
||||
let stored = floats(fields.get(&index.vector_field)?);
|
||||
(stored.len() == index.dims)
|
||||
.then(|| (key, fields, 1.0 - cosine(&vector, &stored)))
|
||||
})
|
||||
.min_by(|left, right| left.2.total_cmp(&right.2));
|
||||
let Some((key, fields, distance)) = best else {
|
||||
return Ok(redis::Value::Array(vec![redis::Value::Int(0)]));
|
||||
};
|
||||
let mut reply = fields
|
||||
.iter()
|
||||
.filter(|(name, _)| **name != index.vector_field)
|
||||
.flat_map(|(name, value)| [bulk(name.as_bytes()), bulk(value)])
|
||||
.collect::<Vec<_>>();
|
||||
reply.extend([
|
||||
bulk(b"vector_distance"),
|
||||
bulk(distance.to_string().as_bytes()),
|
||||
]);
|
||||
Ok(redis::Value::Array(vec![
|
||||
redis::Value::Int(1),
|
||||
bulk(key),
|
||||
redis::Value::Array(reply),
|
||||
]))
|
||||
}
|
||||
"PING" => Ok(redis::Value::SimpleString("PONG".into())),
|
||||
_ => Err(server_error("unsupported command")),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl redis::ConnectionLike for FakeSearch {
|
||||
fn req_packed_command(&mut self, command: &[u8]) -> redis::RedisResult<redis::Value> {
|
||||
let mut commands = parse_commands(command);
|
||||
self.run(commands.remove(0))
|
||||
}
|
||||
|
||||
fn req_packed_commands(
|
||||
&mut self,
|
||||
commands: &[u8],
|
||||
offset: usize,
|
||||
count: usize,
|
||||
) -> redis::RedisResult<Vec<redis::Value>> {
|
||||
let replies = parse_commands(commands)
|
||||
.into_iter()
|
||||
.map(|args| self.run(args))
|
||||
.collect::<redis::RedisResult<Vec<_>>>()?;
|
||||
Ok(replies.into_iter().skip(offset).take(count).collect())
|
||||
}
|
||||
|
||||
fn get_db(&self) -> i64 {
|
||||
0
|
||||
}
|
||||
|
||||
fn check_connection(&mut self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn is_open(&self) -> bool {
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
pub fn server_error(reply: &str) -> redis::RedisError {
|
||||
redis::parse_redis_value(format!("-{reply}\r\n").as_bytes())
|
||||
.expect("error reply")
|
||||
.extract_error()
|
||||
.expect_err("server error")
|
||||
}
|
||||
|
||||
fn bulk(bytes: &[u8]) -> redis::Value {
|
||||
redis::Value::BulkString(bytes.to_vec())
|
||||
}
|
||||
|
||||
fn index_info(index: &FakeIndex) -> redis::Value {
|
||||
redis::Value::Array(vec![
|
||||
bulk(b"index_name"),
|
||||
bulk(b"fake"),
|
||||
bulk(b"attributes"),
|
||||
redis::Value::Array(vec![
|
||||
redis::Value::Array(vec![
|
||||
bulk(b"identifier"),
|
||||
bulk(b"litellm_cache_key"),
|
||||
bulk(b"type"),
|
||||
bulk(b"TAG"),
|
||||
]),
|
||||
redis::Value::Array(vec![
|
||||
bulk(b"identifier"),
|
||||
bulk(index.vector_field.as_bytes()),
|
||||
bulk(b"type"),
|
||||
bulk(b"VECTOR"),
|
||||
bulk(b"index"),
|
||||
redis::Value::Array(vec![
|
||||
bulk(b"dimensions"),
|
||||
redis::Value::Int(index.dims as i64),
|
||||
]),
|
||||
]),
|
||||
]),
|
||||
])
|
||||
}
|
||||
|
||||
/// The tag inside `@litellm_cache_key:{...}`, with query escapes removed.
|
||||
fn query_tag(query: &str) -> String {
|
||||
let start = query.find("@litellm_cache_key:{").unwrap() + "@litellm_cache_key:{".len();
|
||||
let mut tag = String::new();
|
||||
let mut characters = query[start..].chars();
|
||||
while let Some(character) = characters.next() {
|
||||
match character {
|
||||
'\\' => tag.extend(characters.next()),
|
||||
'}' => break,
|
||||
character => tag.push(character),
|
||||
}
|
||||
}
|
||||
tag
|
||||
}
|
||||
|
||||
fn floats(bytes: &[u8]) -> Vec<f32> {
|
||||
bytes
|
||||
.as_chunks::<4>()
|
||||
.0
|
||||
.iter()
|
||||
.map(|chunk| f32::from_le_bytes(*chunk))
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn cosine(left: &[f32], right: &[f32]) -> f64 {
|
||||
let dot = left
|
||||
.iter()
|
||||
.zip(right)
|
||||
.map(|(left, right)| f64::from(*left) * f64::from(*right))
|
||||
.sum::<f64>();
|
||||
let norm = |vector: &[f32]| {
|
||||
vector
|
||||
.iter()
|
||||
.map(|value| f64::from(*value).powi(2))
|
||||
.sum::<f64>()
|
||||
.sqrt()
|
||||
};
|
||||
dot / (norm(left) * norm(right))
|
||||
}
|
||||
|
||||
/// Splits a packed RESP request into each command's arguments.
|
||||
fn parse_commands(mut bytes: &[u8]) -> Vec<Vec<Vec<u8>>> {
|
||||
let line = |bytes: &mut &[u8]| {
|
||||
let end = bytes
|
||||
.windows(2)
|
||||
.position(|window| window == b"\r\n")
|
||||
.unwrap();
|
||||
let text = String::from_utf8(bytes[1..end].to_vec()).unwrap();
|
||||
*bytes = &bytes[end + 2..];
|
||||
text.parse::<usize>().unwrap()
|
||||
};
|
||||
let mut commands = Vec::new();
|
||||
while !bytes.is_empty() {
|
||||
let count = line(&mut bytes);
|
||||
let mut args = Vec::with_capacity(count);
|
||||
for _ in 0..count {
|
||||
let length = line(&mut bytes);
|
||||
args.push(bytes[..length].to_vec());
|
||||
bytes = &bytes[length + 2..];
|
||||
}
|
||||
commands.push(args);
|
||||
}
|
||||
commands
|
||||
}
|
||||
|
||||
/// Records every packed request and answers from a script, `OK` once the script runs out.
|
||||
pub struct RecordingConnection {
|
||||
requests: Arc<Mutex<Vec<Vec<u8>>>>,
|
||||
replies: Mutex<VecDeque<redis::RedisResult<redis::Value>>>,
|
||||
}
|
||||
|
||||
impl RecordingConnection {
|
||||
pub fn new(replies: impl IntoIterator<Item = redis::RedisResult<redis::Value>>) -> Self {
|
||||
Self {
|
||||
requests: Arc::default(),
|
||||
replies: Mutex::new(replies.into_iter().collect()),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn requests(&self) -> Arc<Mutex<Vec<Vec<u8>>>> {
|
||||
Arc::clone(&self.requests)
|
||||
}
|
||||
|
||||
fn reply(&self) -> redis::RedisResult<redis::Value> {
|
||||
self.replies
|
||||
.lock()
|
||||
.unwrap()
|
||||
.pop_front()
|
||||
.unwrap_or_else(|| Ok(redis::Value::SimpleString("OK".into())))
|
||||
}
|
||||
}
|
||||
|
||||
impl redis::ConnectionLike for RecordingConnection {
|
||||
fn req_packed_command(&mut self, command: &[u8]) -> redis::RedisResult<redis::Value> {
|
||||
self.requests.lock().unwrap().push(command.to_vec());
|
||||
self.reply()
|
||||
}
|
||||
|
||||
fn req_packed_commands(
|
||||
&mut self,
|
||||
command: &[u8],
|
||||
_offset: usize,
|
||||
count: usize,
|
||||
) -> redis::RedisResult<Vec<redis::Value>> {
|
||||
self.requests.lock().unwrap().push(command.to_vec());
|
||||
(0..count).map(|_| self.reply()).collect()
|
||||
}
|
||||
|
||||
fn get_db(&self) -> i64 {
|
||||
0
|
||||
}
|
||||
|
||||
fn check_connection(&mut self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn is_open(&self) -> bool {
|
||||
true
|
||||
}
|
||||
}
|
||||
|
|
@ -33,6 +33,17 @@
|
|||
"api_key",
|
||||
"additional_args"
|
||||
],
|
||||
"record_response": [
|
||||
"logger",
|
||||
"status",
|
||||
"headers",
|
||||
"body"
|
||||
],
|
||||
"record_api_window": [
|
||||
"logger",
|
||||
"start",
|
||||
"end"
|
||||
],
|
||||
"defers_async_logging": [
|
||||
"logger"
|
||||
],
|
||||
|
|
@ -91,6 +102,13 @@
|
|||
"logger",
|
||||
"head"
|
||||
],
|
||||
"stamp_stream_head": [
|
||||
"head",
|
||||
"logger",
|
||||
"kwargs",
|
||||
"start_time",
|
||||
"end_time"
|
||||
],
|
||||
"stream_success": [
|
||||
"logger",
|
||||
"url_route",
|
||||
|
|
|
|||
|
|
@ -57,6 +57,8 @@ pub struct LegacyLogging {
|
|||
asynchronous: bool,
|
||||
internal: bool,
|
||||
cache_key: Option<String>,
|
||||
/// When the provider request went out, until its window is recorded (Messages only).
|
||||
api_started: Option<f64>,
|
||||
}
|
||||
|
||||
fn datetime(py: Python<'_>, epoch_seconds: f64) -> PyResult<Py<PyAny>> {
|
||||
|
|
@ -87,6 +89,60 @@ impl LegacyLogging {
|
|||
asynchronous,
|
||||
internal: false,
|
||||
cache_key: None,
|
||||
api_started: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// The kwargs `Logging.update_from_kwargs` sees. Python's Messages handler takes `metadata` as
|
||||
/// its own argument, so the logger never sees the caller's raw metadata, only the validated
|
||||
/// request metadata merged in through `litellm_params`. Python updates the logger once, before
|
||||
/// the call; the refresh once the result is ready leaves out `litellm_metadata` too, so the
|
||||
/// logged metadata is not rebuilt from what the call wrote into it since.
|
||||
fn logging_kwargs<'py>(&self, py: Python<'py>, refresh: bool) -> PyResult<Bound<'py, PyDict>> {
|
||||
let kwargs = self.call.kwargs().bind(py).copy()?;
|
||||
if self.operation != LoggingOperation::Messages {
|
||||
return Ok(kwargs);
|
||||
}
|
||||
let consumed: &[&str] = if refresh {
|
||||
&["metadata", "litellm_metadata"]
|
||||
} else {
|
||||
&["metadata"]
|
||||
};
|
||||
for name in consumed {
|
||||
if kwargs.contains(name)? {
|
||||
kwargs.del_item(name)?;
|
||||
}
|
||||
}
|
||||
Ok(kwargs)
|
||||
}
|
||||
|
||||
/// Python's Messages handler hands the logger its request params, validated `metadata`
|
||||
/// included, as `litellm_params`, beside the `aanthropic_messages` call flag.
|
||||
fn logged_request_params(
|
||||
&self,
|
||||
context: &RequestContext,
|
||||
) -> Option<serde_json::Map<String, Value>> {
|
||||
if self.operation != LoggingOperation::Messages {
|
||||
return None;
|
||||
}
|
||||
let mut params = context
|
||||
.optional_params
|
||||
.as_object()
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
params.insert(
|
||||
"aanthropic_messages".to_owned(),
|
||||
Value::Bool(self.asynchronous),
|
||||
);
|
||||
Some(params)
|
||||
}
|
||||
|
||||
/// Python times the provider exchange inside the async `AsyncHTTPHandler.post`; its sync
|
||||
/// client is untimed.
|
||||
fn close_api_window(&mut self, py: Python<'_>, end: f64) -> PyResult<()> {
|
||||
match self.api_started.take() {
|
||||
Some(start) => self.logger()?.record_api_window(py, start, end),
|
||||
None => Ok(()),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -111,6 +167,16 @@ impl LegacyLogging {
|
|||
}
|
||||
}
|
||||
|
||||
/// The key `pre_call` logs. Python's Messages handler logs an empty string, never the key.
|
||||
fn logged_api_key<'a>(&self, context: &'a RequestContext) -> Option<&'a str> {
|
||||
match self.operation {
|
||||
LoggingOperation::Messages => Some(""),
|
||||
LoggingOperation::Completion | LoggingOperation::Responses | LoggingOperation::Ocr => {
|
||||
context.api_key.as_ref().map(|api_key| api_key.expose())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn stream_billing(&self) -> Option<PassThroughStream> {
|
||||
match self.operation {
|
||||
LoggingOperation::Messages => Some(PassThroughStream {
|
||||
|
|
@ -377,7 +443,13 @@ impl LegacyLogging {
|
|||
context: &RequestContext,
|
||||
) -> PyResult<HookStep<Self, Box<WireRequest>>> {
|
||||
let logger = self.logger()?;
|
||||
logger.update_from_kwargs(py, self.call.kwargs(), &wire, context)?;
|
||||
logger.update_from_kwargs(
|
||||
py,
|
||||
&self.logging_kwargs(py, false)?,
|
||||
&wire,
|
||||
context,
|
||||
self.logged_request_params(context).as_ref(),
|
||||
)?;
|
||||
let body = to_py(py, &wire.body)?
|
||||
.into_bound(py)
|
||||
.cast_into::<PyDict>()?;
|
||||
|
|
@ -400,11 +472,14 @@ impl LegacyLogging {
|
|||
self.logger()?.pre_call(
|
||||
py,
|
||||
self.input_description(),
|
||||
context.api_key.as_ref().map(|api_key| api_key.expose()),
|
||||
self.logged_api_key(context),
|
||||
&body,
|
||||
&headers,
|
||||
&wire.url,
|
||||
)?;
|
||||
if self.operation == LoggingOperation::Messages && self.asynchronous {
|
||||
self.api_started = Some(epoch_seconds());
|
||||
}
|
||||
let headers = headers
|
||||
.iter()
|
||||
.map(|(name, value)| Ok((name.extract::<String>()?, value.extract::<String>()?)))
|
||||
|
|
@ -455,7 +530,7 @@ impl LegacyLogging {
|
|||
py,
|
||||
(
|
||||
&logger,
|
||||
self.call.kwargs(),
|
||||
self.logging_kwargs(py, true)?,
|
||||
&facts.provider.model,
|
||||
logger.getattr("optional_params")?,
|
||||
params,
|
||||
|
|
@ -483,6 +558,16 @@ impl LegacyLogging {
|
|||
py: Python<'_>,
|
||||
raw: &RawResponse,
|
||||
) -> PyResult<HookStep<Self, ()>> {
|
||||
if self.operation == LoggingOperation::Messages {
|
||||
// Python's Messages handler never calls `post_call` at receipt: it stores the raw
|
||||
// response, and logging replays it through `AnthropicConfig.transform_response`.
|
||||
self.close_api_window(py, epoch_seconds())?;
|
||||
if let Some(head) = &raw.head {
|
||||
self.logger()?
|
||||
.record_response(py, head.status, &head.headers, &raw.body)?;
|
||||
}
|
||||
return Ok(HookStep::Ready(()));
|
||||
}
|
||||
let api_key = self
|
||||
.request
|
||||
.as_ref()
|
||||
|
|
@ -522,6 +607,9 @@ impl LegacyLogging {
|
|||
) -> PyResult<HookStep<Self, ()>> {
|
||||
self.end = Some(datetime(py, timing.end_time)?);
|
||||
self.error = Some(error.clone_ref(py).into_value(py));
|
||||
if self.logger.is_some() {
|
||||
self.close_api_window(py, timing.end_time)?;
|
||||
}
|
||||
if self.stream.is_some() {
|
||||
return self.stream_failure(py);
|
||||
}
|
||||
|
|
@ -543,6 +631,20 @@ impl LegacyLogging {
|
|||
head.bind(py).set_item("cache_key", key)?;
|
||||
head.bind(py).set_item("cache_hit", true)?;
|
||||
}
|
||||
self.close_api_window(py, epoch_seconds())?;
|
||||
if self.operation == LoggingOperation::Messages {
|
||||
// Python's `@client` stamps the call metadata onto the stream it hands back.
|
||||
Streaming::StampHead.call(
|
||||
py,
|
||||
(
|
||||
head.bind(py),
|
||||
self.logger()?.object(py),
|
||||
self.call.kwargs(),
|
||||
&self.start,
|
||||
datetime(py, epoch_seconds())?,
|
||||
),
|
||||
)?;
|
||||
}
|
||||
Streaming::Opened.call(py, (self.logger()?.object(py), head.bind(py)))?;
|
||||
self.stream = Some(DeliveredStream {
|
||||
chunks: PyList::empty(py).unbind(),
|
||||
|
|
@ -1117,9 +1219,7 @@ assert logger.update['custom_llm_provider'] == 'edenai', logger.update
|
|||
let step = logging
|
||||
.before_provider_request(py, Box::new(wire), context)
|
||||
.unwrap();
|
||||
let raw = RawResponse {
|
||||
body: "raw response".into(),
|
||||
};
|
||||
let raw = RawResponse::new("raw response".into());
|
||||
assert!(matches!(
|
||||
logging
|
||||
.on_event(
|
||||
|
|
@ -1869,9 +1969,10 @@ assert logger.calls[1][1] is response
|
|||
py,
|
||||
&locals,
|
||||
c"
|
||||
assert logger.names() == ['stream_opened', 'stream_success'], logger.calls
|
||||
assert logger.names() == ['stamp_stream_head', 'stream_opened', 'stream_success'], logger.calls
|
||||
assert logger.calls[0][1] is head
|
||||
chunks = logger.calls[1][1]
|
||||
assert logger.calls[1][1] is head
|
||||
chunks = logger.calls[2][1]
|
||||
assert len(chunks) == 2
|
||||
assert chunks[0] is first
|
||||
assert chunks[1] is last
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@
|
|||
use litellm_host::interceptors::{RequestContext, WireRequest};
|
||||
use litellm_host_python::to_py;
|
||||
use pyo3::{exceptions::PyBaseException, prelude::*, types::PyDict};
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::logger::PythonLogger;
|
||||
use crate::python::{Logging, Wrapper};
|
||||
|
|
@ -15,9 +16,10 @@ pub trait LegacyCallbacks {
|
|||
fn update_from_kwargs(
|
||||
&self,
|
||||
py: Python<'_>,
|
||||
kwargs: &Py<PyDict>,
|
||||
kwargs: &Bound<'_, PyDict>,
|
||||
wire: &WireRequest,
|
||||
context: &RequestContext,
|
||||
request_params: Option<&serde_json::Map<String, Value>>,
|
||||
) -> PyResult<()>;
|
||||
|
||||
/// `Logging.pre_call`.
|
||||
|
|
@ -31,6 +33,20 @@ pub trait LegacyCallbacks {
|
|||
url: &str,
|
||||
) -> PyResult<()>;
|
||||
|
||||
/// Records the provider's raw HTTP response on the logger, as Python's Messages handler
|
||||
/// does with `httpx_response` and `response_headers`.
|
||||
fn record_response(
|
||||
&self,
|
||||
py: Python<'_>,
|
||||
status: u16,
|
||||
headers: &[(String, String)],
|
||||
body: &str,
|
||||
) -> PyResult<()>;
|
||||
|
||||
/// Records one provider HTTP exchange's window, as Python's timed `AsyncHTTPHandler.post`
|
||||
/// does with `llm_api_duration_ms` and the metadata's `llm_api_timing_windows`.
|
||||
fn record_api_window(&self, py: Python<'_>, start: f64, end: f64) -> PyResult<()>;
|
||||
|
||||
/// `Logging.post_call`.
|
||||
fn post_call(
|
||||
&self,
|
||||
|
|
@ -83,12 +99,13 @@ impl LegacyCallbacks for PythonLogger {
|
|||
fn update_from_kwargs(
|
||||
&self,
|
||||
py: Python<'_>,
|
||||
kwargs: &Py<PyDict>,
|
||||
kwargs: &Bound<'_, PyDict>,
|
||||
wire: &WireRequest,
|
||||
context: &RequestContext,
|
||||
request_params: Option<&serde_json::Map<String, Value>>,
|
||||
) -> PyResult<()> {
|
||||
let secret_fields: Vec<&str> = context.secret_fields.iter().map(String::as_str).collect();
|
||||
let redacted_kwargs = redact(py, kwargs.bind(py), &secret_fields)?;
|
||||
let redacted_kwargs = redact(py, kwargs, &secret_fields)?;
|
||||
let optional_params = redact(
|
||||
py,
|
||||
&to_py(py, &context.optional_params)?
|
||||
|
|
@ -97,23 +114,23 @@ impl LegacyCallbacks for PythonLogger {
|
|||
&secret_fields,
|
||||
)?;
|
||||
let params = PyDict::new(py);
|
||||
params.set_item(
|
||||
"litellm_call_id",
|
||||
kwargs.bind(py).get_item("litellm_call_id")?,
|
||||
)?;
|
||||
params.set_item("litellm_call_id", kwargs.get_item("litellm_call_id")?)?;
|
||||
params.set_item("api_base", &wire.url)?;
|
||||
for name in ["logger_fn", "litellm_request_debug"] {
|
||||
if let Some(value) = kwargs.bind(py).get_item(name)? {
|
||||
if let Some(value) = kwargs.get_item(name)? {
|
||||
params.set_item(name, value)?;
|
||||
}
|
||||
}
|
||||
for name in custom_pricing_fields(py)? {
|
||||
if let Some(value) = kwargs.bind(py).get_item(&name)?
|
||||
if let Some(value) = kwargs.get_item(&name)?
|
||||
&& !value.is_none()
|
||||
{
|
||||
params.set_item(name, value)?;
|
||||
}
|
||||
}
|
||||
for (name, value) in request_params.into_iter().flatten() {
|
||||
params.set_item(name, to_py(py, value)?)?;
|
||||
}
|
||||
Logging::Update.call(
|
||||
py,
|
||||
(
|
||||
|
|
@ -145,6 +162,22 @@ impl LegacyCallbacks for PythonLogger {
|
|||
Ok(())
|
||||
}
|
||||
|
||||
fn record_response(
|
||||
&self,
|
||||
py: Python<'_>,
|
||||
status: u16,
|
||||
headers: &[(String, String)],
|
||||
body: &str,
|
||||
) -> PyResult<()> {
|
||||
Logging::RecordResponse.call(py, (self.object(py), status, headers.to_vec(), body))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn record_api_window(&self, py: Python<'_>, start: f64, end: f64) -> PyResult<()> {
|
||||
Logging::RecordApiWindow.call(py, (self.object(py), start, end))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn post_call(
|
||||
&self,
|
||||
py: Python<'_>,
|
||||
|
|
|
|||
|
|
@ -44,6 +44,10 @@ pub(crate) enum Logging {
|
|||
PreCall,
|
||||
#[strum(serialize = "post_call")]
|
||||
PostCall,
|
||||
#[strum(serialize = "record_response")]
|
||||
RecordResponse,
|
||||
#[strum(serialize = "record_api_window")]
|
||||
RecordApiWindow,
|
||||
#[strum(serialize = "defers_async_logging")]
|
||||
DefersAsync,
|
||||
#[strum(serialize = "defer_success")]
|
||||
|
|
@ -77,6 +81,8 @@ pub(crate) enum DeploymentHooks {
|
|||
pub(crate) enum Streaming {
|
||||
#[strum(serialize = "stream_opened")]
|
||||
Opened,
|
||||
#[strum(serialize = "stamp_stream_head")]
|
||||
StampHead,
|
||||
#[strum(serialize = "stream_success")]
|
||||
Success,
|
||||
#[strum(serialize = "stream_failure")]
|
||||
|
|
|
|||
|
|
@ -65,6 +65,12 @@ FAKES = {
|
|||
'post_call': lambda logger, original_response, api_key, additional_args: logger.post_call(
|
||||
original_response, api_key, additional_args
|
||||
),
|
||||
'record_response': lambda logger, status, headers, body: logger.model_call_details.update(
|
||||
response_headers=dict(headers)
|
||||
),
|
||||
'record_api_window': lambda logger, start, end: logger.model_call_details.update(
|
||||
llm_api_duration_ms=(end - start) * 1000
|
||||
),
|
||||
'defers_async_logging': lambda logger: bool(getattr(logger, '_defer_async_logging', False)),
|
||||
'defer_success': lambda logger, pending: setattr(logger, '_native_pending_logging', pending),
|
||||
'sync_success_for_async_call': lambda logger, response, start, end: logger.handle_sync_success_callbacks_for_async_calls(
|
||||
|
|
@ -85,6 +91,7 @@ FAKES = {
|
|||
),
|
||||
'after_deployment_failure': lambda kwargs, error, call_type: kwargs['logger'].hook('failure', error, call_type),
|
||||
'stream_opened': lambda logger, head: logger.record('stream_opened', head),
|
||||
'stamp_stream_head': lambda head, logger, kwargs, start, end: logger.record('stamp_stream_head', head),
|
||||
'stream_success': lambda logger, url_route, endpoint_type, request_body, chunks, start, end, first_chunk: logger.record(
|
||||
'stream_success', list(chunks)
|
||||
),
|
||||
|
|
|
|||
|
|
@ -151,6 +151,8 @@ pub enum Uplift {
|
|||
Residency(Residency),
|
||||
/// Vertex non-global endpoint, from `regional_endpoint_uplift_multiplier`
|
||||
NonGlobalEndpoint,
|
||||
/// OpenAI FedRAMP endpoint (gov.api.openai.com), from `fedramp_uplift_multiplier`
|
||||
FedrampEndpoint,
|
||||
/// Anthropic `speed: fast`, from `provider_specific_entry.fast`
|
||||
Speed(Speed),
|
||||
/// Anthropic and Bedrock geo, from `provider_specific_entry.<geo>`, lowercased as on main
|
||||
|
|
|
|||
|
|
@ -215,9 +215,7 @@ async fn projection_custom_operations_and_hooks_feed_the_http_response(intercept
|
|||
)
|
||||
.await?;
|
||||
route_hooks
|
||||
.after_provider_response(RawResponse {
|
||||
body: wire.url.clone(),
|
||||
})
|
||||
.after_provider_response(RawResponse::new(wire.url.clone()))
|
||||
.await?;
|
||||
Ok(CallOutput::Complete(Bytes::from(wire.url)))
|
||||
},
|
||||
|
|
@ -501,9 +499,7 @@ async fn hook_rejection_stops_execution_and_is_reported_once(
|
|||
move |_, _, route_hooks, _observations| async move {
|
||||
if event {
|
||||
route_hooks
|
||||
.after_provider_response(RawResponse {
|
||||
body: "response".into(),
|
||||
})
|
||||
.after_provider_response(RawResponse::new("response".into()))
|
||||
.await?;
|
||||
} else {
|
||||
route_hooks
|
||||
|
|
@ -637,9 +633,7 @@ async fn unary_calls_use_into_response_after_hooks_and_before_success(intercepto
|
|||
)
|
||||
.await?;
|
||||
route_hooks
|
||||
.after_provider_response(RawResponse {
|
||||
body: wire.url.clone(),
|
||||
})
|
||||
.after_provider_response(RawResponse::new(wire.url.clone()))
|
||||
.await?;
|
||||
Ok(CallOutput::Complete(json!({"url": wire.url})))
|
||||
},
|
||||
|
|
@ -689,7 +683,7 @@ async fn unary_failure_preserves_the_error_without_converting(
|
|||
None,
|
||||
|_, _, route_hooks, _observations| async move {
|
||||
route_hooks
|
||||
.after_provider_response(RawResponse { body: "raw".into() })
|
||||
.after_provider_response(RawResponse::new("raw".into()))
|
||||
.await?;
|
||||
Err(TestError::Provider)
|
||||
},
|
||||
|
|
|
|||
|
|
@ -170,9 +170,7 @@ fn dispatching_call() -> TestMachine {
|
|||
)
|
||||
.await?;
|
||||
route_hooks
|
||||
.after_provider_response(RawResponse {
|
||||
body: wire.url.clone(),
|
||||
})
|
||||
.after_provider_response(RawResponse::new(wire.url.clone()))
|
||||
.await?;
|
||||
Ok(CallOutput::Complete(wire.url))
|
||||
},
|
||||
|
|
@ -660,9 +658,7 @@ async fn response_interception_waits_and_can_reject_after_observation(#[case] re
|
|||
"input",
|
||||
Some(sender),
|
||||
|_, _, interceptors, observers| async move {
|
||||
let raw = RawResponse {
|
||||
body: "provider response".into(),
|
||||
};
|
||||
let raw = RawResponse::new("provider response".into());
|
||||
observers.unwrap().emit(CallEvent::Execution(
|
||||
ExecutionEvent::ProviderResponseReceived { raw: raw.clone() },
|
||||
));
|
||||
|
|
|
|||
|
|
@ -467,7 +467,7 @@ pub fn success_machine() -> MachineFactory {
|
|||
.before_provider_request(wire(), context())
|
||||
.await?;
|
||||
host.interceptors
|
||||
.after_provider_response(RawResponse { body: "raw".into() })
|
||||
.after_provider_response(RawResponse::new("raw".into()))
|
||||
.await?;
|
||||
Ok(format!("{projected}|{signed}|{}", wire.url))
|
||||
})
|
||||
|
|
|
|||
|
|
@ -28,6 +28,28 @@ pub struct RequestContext {
|
|||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct RawResponse {
|
||||
pub body: String,
|
||||
/// The provider's HTTP status and headers, when the route forwards them to logging.
|
||||
pub head: Option<ResponseHead>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct ResponseHead {
|
||||
pub status: u16,
|
||||
pub headers: Vec<(String, String)>,
|
||||
}
|
||||
|
||||
impl RawResponse {
|
||||
pub fn new(body: String) -> Self {
|
||||
Self { body, head: None }
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn with_head(self, head: ResponseHead) -> Self {
|
||||
Self {
|
||||
head: Some(head),
|
||||
..self
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
|
|
@ -186,7 +208,7 @@ mod tests {
|
|||
.await?;
|
||||
Interceptors::after_provider_response(
|
||||
&channel.interceptors,
|
||||
RawResponse { body: "raw".into() },
|
||||
RawResponse::new("raw".into()),
|
||||
)
|
||||
.await?;
|
||||
Ok((sent, ()))
|
||||
|
|
|
|||
|
|
@ -74,9 +74,7 @@ fn snapshots_do_not_retain_runtime_objects(capacity: NonZeroUsize, #[case] faile
|
|||
|
||||
#[rstest]
|
||||
fn provider_snapshots_own_the_response_body(capacity: NonZeroUsize) {
|
||||
let mut raw = RawResponse {
|
||||
body: "provider response".into(),
|
||||
};
|
||||
let mut raw = RawResponse::new("provider response".into());
|
||||
let event: CallEvent<(), (), &RawResponse> =
|
||||
CallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw: &raw });
|
||||
let expected = raw.clone();
|
||||
|
|
|
|||
|
|
@ -58,8 +58,20 @@ pub(crate) enum Repr {
|
|||
Network(#[source] reqwest::Error),
|
||||
#[error("upstream request timed out")]
|
||||
Timeout(#[source] reqwest::Error),
|
||||
/// No bytes arrived within the read timeout, which the caller enforces per read; `body` says
|
||||
/// the response head had already arrived.
|
||||
#[error("upstream request timed out")]
|
||||
ReadTimeout {
|
||||
#[source]
|
||||
elapsed: Elapsed,
|
||||
body: bool,
|
||||
},
|
||||
#[error("upstream request failed with status {status}: {body}")]
|
||||
Status { status: u16, body: UpstreamBody },
|
||||
Status {
|
||||
status: u16,
|
||||
body: UpstreamBody,
|
||||
headers: Vec<(String, String)>,
|
||||
},
|
||||
|
||||
#[error("invalid WebSocket URL")]
|
||||
WebSocketUrl(#[source] url::ParseError),
|
||||
|
|
@ -136,13 +148,59 @@ impl Repr {
|
|||
|
||||
impl Error {
|
||||
pub fn upstream(status: u16, body: String) -> Self {
|
||||
Self::upstream_with_headers(status, body, Vec::new())
|
||||
}
|
||||
|
||||
/// An upstream failure that keeps the provider's response headers (`retry-after`,
|
||||
/// `request-id`) for the caller.
|
||||
pub fn upstream_with_headers(
|
||||
status: u16,
|
||||
body: String,
|
||||
headers: Vec<(String, String)>,
|
||||
) -> Self {
|
||||
Repr::Status {
|
||||
status,
|
||||
body: UpstreamBody(body),
|
||||
headers,
|
||||
}
|
||||
.into()
|
||||
}
|
||||
|
||||
/// The provider's response headers on an upstream failure; empty for every other error.
|
||||
pub fn upstream_headers(&self) -> &[(String, String)] {
|
||||
match &*self.0 {
|
||||
Repr::Status { headers, .. } => headers,
|
||||
_ => &[],
|
||||
}
|
||||
}
|
||||
|
||||
/// A timeout reached before the connection was established, as opposed to while reading.
|
||||
pub fn is_connect_timeout(&self) -> bool {
|
||||
matches!(&*self.0, Repr::Timeout(error) if error.is_connect())
|
||||
}
|
||||
|
||||
/// A read the caller timed itself (httpx's per-read timeout) ran out.
|
||||
pub fn read_timeout(elapsed: Elapsed) -> Self {
|
||||
Repr::ReadTimeout {
|
||||
elapsed,
|
||||
body: false,
|
||||
}
|
||||
.into()
|
||||
}
|
||||
|
||||
/// A read of the response body the caller timed itself ran out, after the head arrived.
|
||||
pub fn body_read_timeout(elapsed: Elapsed) -> Self {
|
||||
Repr::ReadTimeout {
|
||||
elapsed,
|
||||
body: true,
|
||||
}
|
||||
.into()
|
||||
}
|
||||
|
||||
pub fn is_body_read_timeout(&self) -> bool {
|
||||
matches!(&*self.0, Repr::ReadTimeout { body: true, .. })
|
||||
}
|
||||
|
||||
pub fn computed_header(name: impl Into<String>) -> Self {
|
||||
Repr::ComputedHeader(name.into()).into()
|
||||
}
|
||||
|
|
@ -169,6 +227,7 @@ impl Error {
|
|||
match &*self.0 {
|
||||
Repr::Network(_)
|
||||
| Repr::Timeout(_)
|
||||
| Repr::ReadTimeout { .. }
|
||||
| Repr::Status { .. }
|
||||
| Repr::Handshake(_)
|
||||
| Repr::HandshakeTimeout(_)
|
||||
|
|
@ -227,6 +286,7 @@ impl Error {
|
|||
| Repr::Connect(_)
|
||||
| Repr::Network(_)
|
||||
| Repr::Timeout(_)
|
||||
| Repr::ReadTimeout { .. }
|
||||
| Repr::Status { .. }
|
||||
| Repr::Handshake(_)
|
||||
| Repr::HandshakeTimeout(_)
|
||||
|
|
@ -259,6 +319,7 @@ impl Error {
|
|||
| Repr::Connect(_)
|
||||
| Repr::Network(_)
|
||||
| Repr::Timeout(_)
|
||||
| Repr::ReadTimeout { .. }
|
||||
| Repr::Status { .. }
|
||||
| Repr::WebSocketUrl(_)
|
||||
| Repr::WebSocketScheme(_)
|
||||
|
|
@ -295,6 +356,7 @@ impl Error {
|
|||
| Repr::Connect(_)
|
||||
| Repr::Network(_)
|
||||
| Repr::Timeout(_)
|
||||
| Repr::ReadTimeout { .. }
|
||||
| Repr::Status { .. }
|
||||
| Repr::WebSocketUrl(_)
|
||||
| Repr::WebSocketScheme(_)
|
||||
|
|
@ -333,6 +395,7 @@ impl Error {
|
|||
| Repr::Signature(_)
|
||||
| Repr::Network(_)
|
||||
| Repr::Timeout(_)
|
||||
| Repr::ReadTimeout { .. }
|
||||
| Repr::Status { .. }
|
||||
| Repr::WebSocketUrl(_)
|
||||
| Repr::WebSocketScheme(_)
|
||||
|
|
@ -356,7 +419,10 @@ impl Error {
|
|||
|
||||
pub fn is_timeout(&self) -> bool {
|
||||
match &*self.0 {
|
||||
Repr::Timeout(_) | Repr::HandshakeTimeout(_) | Repr::DownloadTimeout(_) => true,
|
||||
Repr::Timeout(_)
|
||||
| Repr::ReadTimeout { .. }
|
||||
| Repr::HandshakeTimeout(_)
|
||||
| Repr::DownloadTimeout(_) => true,
|
||||
Repr::Read { .. }
|
||||
| Repr::InvalidPem { .. }
|
||||
| Repr::Client(_)
|
||||
|
|
@ -392,7 +458,7 @@ impl Error {
|
|||
|
||||
pub fn upstream_response(&self) -> Option<(u16, &str)> {
|
||||
match &*self.0 {
|
||||
Repr::Status { status, body } => Some((*status, &body.0)),
|
||||
Repr::Status { status, body, .. } => Some((*status, &body.0)),
|
||||
Repr::Read { .. }
|
||||
| Repr::InvalidPem { .. }
|
||||
| Repr::Client(_)
|
||||
|
|
@ -408,6 +474,7 @@ impl Error {
|
|||
| Repr::Connect(_)
|
||||
| Repr::Network(_)
|
||||
| Repr::Timeout(_)
|
||||
| Repr::ReadTimeout { .. }
|
||||
| Repr::WebSocketUrl(_)
|
||||
| Repr::WebSocketScheme(_)
|
||||
| Repr::WebSocketSchemeRewrite(_)
|
||||
|
|
@ -431,7 +498,7 @@ impl Error {
|
|||
pub fn status(&self) -> Option<u16> {
|
||||
match &*self.0 {
|
||||
Repr::Status { status, .. } | Repr::DownloadStatus(status) => Some(*status),
|
||||
Repr::Timeout(_) | Repr::DownloadTimeout(_) => Some(408),
|
||||
Repr::Timeout(_) | Repr::ReadTimeout { .. } | Repr::DownloadTimeout(_) => Some(408),
|
||||
Repr::Read { .. }
|
||||
| Repr::InvalidPem { .. }
|
||||
| Repr::Client(_)
|
||||
|
|
@ -482,6 +549,7 @@ impl Error {
|
|||
| Repr::Connect(_)
|
||||
| Repr::Network(_)
|
||||
| Repr::Timeout(_)
|
||||
| Repr::ReadTimeout { .. }
|
||||
| Repr::Status { .. }
|
||||
| Repr::WebSocketUrl(_)
|
||||
| Repr::WebSocketScheme(_)
|
||||
|
|
|
|||
|
|
@ -108,7 +108,7 @@ pub(super) async fn execute(
|
|||
if !status.is_success() {
|
||||
return Err(litellm_http::Error::upstream(status.as_u16(), text).into());
|
||||
}
|
||||
let raw = RawResponse { body: text.clone() };
|
||||
let raw = RawResponse::new(text.clone());
|
||||
if let Some(observers) = observers {
|
||||
observers.emit(litellm_host::lifecycle::CallEvent::Execution(
|
||||
ExecutionEvent::ProviderResponseReceived { raw: raw.clone() },
|
||||
|
|
|
|||
|
|
@ -24,6 +24,7 @@ litellm-tracing.workspace = true
|
|||
reqwest.workspace = true
|
||||
serde.workspace = true
|
||||
serde_json = { workspace = true, features = ["preserve_order"] }
|
||||
tokio.workspace = true
|
||||
tracing.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
|
|
@ -34,5 +35,4 @@ litellm-host-native.workspace = true
|
|||
litellm-http = { workspace = true, features = ["test-support"] }
|
||||
litellm-inference-testing.workspace = true
|
||||
rstest.workspace = true
|
||||
tokio.workspace = true
|
||||
wiremock.workspace = true
|
||||
|
|
|
|||
|
|
@ -35,6 +35,14 @@ impl MessagesProvider {
|
|||
pub(crate) fn config(self) -> &'static dyn BaseMessagesConfig {
|
||||
self.config
|
||||
}
|
||||
|
||||
/// Python's `MESSAGES_SERVICE_TIER_PROVIDERS`: everyone else has `service_tier` dropped.
|
||||
pub(crate) fn accepts_service_tier(self) -> bool {
|
||||
matches!(
|
||||
self.provider,
|
||||
LlmProviders::Anthropic | LlmProviders::AzureAi | LlmProviders::BedrockMantle
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/// Python's `get_provider_anthropic_messages_config`: Vertex AI serves only its Claude
|
||||
|
|
@ -119,3 +127,69 @@ mod tests {
|
|||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// Python's `_merge_messages_header_sources`: later sources win per header name, except
|
||||
/// `anthropic-beta` (any casing), whose values from every source are unioned, deduplicated and
|
||||
/// sorted into one lowercase `anthropic-beta` entry.
|
||||
pub fn merge_messages_header_sources(
|
||||
sources: impl IntoIterator<Item = Map<String, Value>>,
|
||||
) -> Map<String, Value> {
|
||||
let mut merged = Map::new();
|
||||
let mut betas = std::collections::BTreeSet::new();
|
||||
for (name, value) in sources.into_iter().flatten() {
|
||||
if !name.eq_ignore_ascii_case("anthropic-beta") {
|
||||
merged.insert(name, value);
|
||||
continue;
|
||||
}
|
||||
let values: Vec<&str> = match &value {
|
||||
Value::String(text) => vec![text.as_str()],
|
||||
Value::Array(entries) => entries.iter().filter_map(Value::as_str).collect(),
|
||||
_ => Vec::new(),
|
||||
};
|
||||
betas.extend(
|
||||
values
|
||||
.iter()
|
||||
.flat_map(|value| value.split(','))
|
||||
.map(str::trim)
|
||||
.filter(|beta| !beta.is_empty())
|
||||
.map(str::to_owned),
|
||||
);
|
||||
}
|
||||
if !betas.is_empty() {
|
||||
merged.insert(
|
||||
"anthropic-beta".to_owned(),
|
||||
Value::String(betas.into_iter().collect::<Vec<_>>().join(",")),
|
||||
);
|
||||
}
|
||||
merged
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod merge_tests {
|
||||
use serde_json::{Map, Value, json};
|
||||
|
||||
use super::merge_messages_header_sources;
|
||||
|
||||
fn map(value: Value) -> Map<String, Value> {
|
||||
value.as_object().cloned().unwrap()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn later_sources_win_except_anthropic_beta_which_is_unioned_and_sorted() {
|
||||
let merged = merge_messages_header_sources([
|
||||
map(json!({"x-a": "1", "anthropic-beta": "b-2, a-1"})),
|
||||
map(json!({"x-a": "2", "Anthropic-Beta": ["c-3", 7, "a-1"]})),
|
||||
map(json!({"anthropic-beta": 5})),
|
||||
]);
|
||||
assert_eq!(
|
||||
Value::Object(merged),
|
||||
json!({"x-a": "2", "anthropic-beta": "a-1,b-2,c-3"})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_beta_values_leave_no_beta_header() {
|
||||
let merged = merge_messages_header_sources([map(json!({"anthropic-beta": ""}))]);
|
||||
assert!(merged.is_empty());
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -10,9 +10,9 @@ use litellm_core_utils::{
|
|||
get_provider_specific_headers::get_provider_specific_headers, settings::Lookup,
|
||||
};
|
||||
use litellm_host::interceptors::{
|
||||
Cost, ExecutionFacts, Interceptors, ProviderIdentity, RequestContext, ResultSource, WireRequest,
|
||||
Cost, ExecutionFacts, Interceptors, ProviderIdentity, RawResponse, RequestContext,
|
||||
ResponseHead, ResultSource, WireRequest,
|
||||
};
|
||||
use litellm_http::request::with_default_headers;
|
||||
use litellm_llms::base_llm::{
|
||||
auth::{Authenticated, ValidatedEnvironment, resolve_auth},
|
||||
messages::{
|
||||
|
|
@ -21,14 +21,16 @@ use litellm_llms::base_llm::{
|
|||
transformation::BaseMessagesConfig,
|
||||
},
|
||||
};
|
||||
use litellm_llms_types::formats::messages::{MessagesRequest, MessagesResponse};
|
||||
use litellm_llms_types::formats::messages::{
|
||||
MessagesOptionalParams, MessagesRequest, MessagesResponse,
|
||||
};
|
||||
use litellm_secrets::source::Secrets;
|
||||
use litellm_tracing::ByteChunk;
|
||||
use serde_json::Value;
|
||||
|
||||
use super::{
|
||||
Error, MessagesCall, MessagesCallResponse, MessagesRoute,
|
||||
common_utils::{MessagesProvider, string_headers},
|
||||
common_utils::{MessagesProvider, merge_messages_header_sources, string_headers},
|
||||
constants::MESSAGES_TIMEOUT_SECS,
|
||||
route::Messages,
|
||||
types::invalid_request,
|
||||
|
|
@ -126,14 +128,18 @@ async fn provider_call(
|
|||
MessagesRequest { model, ..body },
|
||||
shaping.settings.reasoning_auto_summary,
|
||||
)?;
|
||||
let transform_context =
|
||||
MessagesTransformContext::new(shaping.capabilities, shaping.settings.drop_params);
|
||||
let filtered = requested_optional_params(sanitized, provider, &transform_context)?;
|
||||
let trimmed =
|
||||
without_additional_drop_params(sanitized, &shaping.settings.additional_drop_params)?;
|
||||
without_additional_drop_params(filtered, &shaping.settings.additional_drop_params)?;
|
||||
|
||||
let scoped =
|
||||
get_provider_specific_headers(provider_specific_header.as_ref(), provider.as_str());
|
||||
let forwarded = string_headers(Some(
|
||||
extra_headers.into_iter().flatten().chain(scoped).collect(),
|
||||
))?;
|
||||
let forwarded = string_headers(Some(merge_messages_header_sources([
|
||||
extra_headers.unwrap_or_default(),
|
||||
scoped,
|
||||
])))?;
|
||||
let validated = config.validate_environment(
|
||||
forwarded,
|
||||
api_key.as_deref(),
|
||||
|
|
@ -144,16 +150,13 @@ async fn provider_call(
|
|||
)?;
|
||||
|
||||
let body = config.request_body(
|
||||
config.transform_anthropic_messages_request(
|
||||
trimmed,
|
||||
&MessagesTransformContext::new(shaping.capabilities, shaping.settings.drop_params),
|
||||
)?,
|
||||
config.transform_anthropic_messages_request(trimmed, &transform_context)?,
|
||||
&validated.headers,
|
||||
shaping.settings.drop_params,
|
||||
)?;
|
||||
let environment = ValidatedEnvironment {
|
||||
headers: config.request_headers(
|
||||
with_default_headers(validated.headers, config.default_headers()),
|
||||
with_python_default_headers(validated.headers, config.default_headers()),
|
||||
&body,
|
||||
),
|
||||
auth: validated.auth,
|
||||
|
|
@ -184,7 +187,7 @@ async fn provider_call(
|
|||
.before_provider_request(
|
||||
WireRequest {
|
||||
url,
|
||||
headers: authenticated.headers,
|
||||
headers: joined_header_values(authenticated.headers),
|
||||
body: config.wire_body(serde_json::to_value(&body).map_err(serialize_failure)?),
|
||||
},
|
||||
request_context,
|
||||
|
|
@ -213,6 +216,64 @@ async fn provider_call(
|
|||
})
|
||||
}
|
||||
|
||||
/// Python sends its header dict through `httpx.Headers`, which folds names that differ only in
|
||||
/// case into one lowercase header whose values are joined with `", "`.
|
||||
fn joined_header_values(headers: Vec<(String, String)>) -> Vec<(String, String)> {
|
||||
let mut joined: Vec<(String, String)> = Vec::with_capacity(headers.len());
|
||||
let mut positions: std::collections::HashMap<String, usize> =
|
||||
std::collections::HashMap::with_capacity(headers.len());
|
||||
for (name, value) in headers {
|
||||
match positions.entry(name.to_ascii_lowercase()) {
|
||||
std::collections::hash_map::Entry::Occupied(position) => {
|
||||
let present = &mut joined[*position.get()].1;
|
||||
present.push_str(", ");
|
||||
present.push_str(&value);
|
||||
}
|
||||
std::collections::hash_map::Entry::Vacant(position) => {
|
||||
position.insert(joined.len());
|
||||
joined.push((name, value));
|
||||
}
|
||||
}
|
||||
}
|
||||
joined
|
||||
}
|
||||
|
||||
/// Python's `finalize_messages_headers` adds a default only when the exact lowercase name is
|
||||
/// absent, so a caller's `Anthropic-Version` is sent alongside the default, not instead of it.
|
||||
fn with_python_default_headers(
|
||||
headers: Vec<(String, String)>,
|
||||
defaults: &[(&str, &str)],
|
||||
) -> Vec<(String, String)> {
|
||||
let missing: Vec<(String, String)> = defaults
|
||||
.iter()
|
||||
.filter(|(name, _)| !headers.iter().any(|(present, _)| present == name))
|
||||
.map(|(name, value)| ((*name).to_string(), (*value).to_string()))
|
||||
.collect();
|
||||
headers.into_iter().chain(missing).collect()
|
||||
}
|
||||
|
||||
/// Python's `get_requested_anthropic_messages_optional_param`, which gates speed, service tier
|
||||
/// and sampling params before `additional_drop_params` and the environment check run.
|
||||
fn requested_optional_params(
|
||||
request: MessagesRequest,
|
||||
provider: MessagesProvider,
|
||||
context: &MessagesTransformContext,
|
||||
) -> Result<MessagesRequest, Error> {
|
||||
let request = litellm_llms::anthropic::messages::transformation::drop_unsupported_params(
|
||||
request, context,
|
||||
)?;
|
||||
if provider.accepts_service_tier() {
|
||||
return Ok(request);
|
||||
}
|
||||
Ok(MessagesRequest {
|
||||
params: MessagesOptionalParams {
|
||||
service_tier: None,
|
||||
..request.params
|
||||
},
|
||||
..request
|
||||
})
|
||||
}
|
||||
|
||||
fn without_additional_drop_params(
|
||||
request: MessagesRequest,
|
||||
paths: &[String],
|
||||
|
|
@ -245,6 +306,7 @@ async fn call_provider(
|
|||
} = request;
|
||||
let provider_name = provider.as_str();
|
||||
log_request_body(provider_name, stream, &wire.body);
|
||||
let timeout = timeout.unwrap_or(Duration::from_secs(MESSAGES_TIMEOUT_SECS));
|
||||
let response = send(
|
||||
http,
|
||||
Authenticated {
|
||||
|
|
@ -254,10 +316,11 @@ async fn call_provider(
|
|||
&wire.url,
|
||||
&wire.body,
|
||||
timeout,
|
||||
stream,
|
||||
)
|
||||
.await?;
|
||||
if !response.status().is_success() {
|
||||
return Err(provider_error(response).await);
|
||||
return Err(provider_error(response, timeout).await);
|
||||
}
|
||||
let config = provider.config();
|
||||
if stream {
|
||||
|
|
@ -265,11 +328,24 @@ async fn call_provider(
|
|||
response,
|
||||
config.stream_decoder(),
|
||||
provider_name,
|
||||
timeout,
|
||||
));
|
||||
}
|
||||
let head = ResponseHead {
|
||||
status: response.status().as_u16(),
|
||||
headers: litellm_http::request::header_pairs(response.headers()),
|
||||
};
|
||||
let status = head.status;
|
||||
let text = response.text().await.map_err(litellm_http::Error::from)?;
|
||||
log_response_body(&text);
|
||||
context.response_received(&text).await?;
|
||||
context
|
||||
.raw_response_received(RawResponse::new(text.clone()).with_head(head))
|
||||
.await?;
|
||||
if serde_json::from_str::<serde::de::IgnoredAny>(&text).is_err() {
|
||||
// Python raises `AnthropicError(status_code=<the response's>, message=<its text>)` for a
|
||||
// body it cannot parse, and the public mapping classifies it by that status.
|
||||
return Err(litellm_http::Error::upstream(status, text).into());
|
||||
}
|
||||
decode_response(config, &identity.model, &text)
|
||||
.map(|message| MessagesCallResponse::Complete(Box::new(message)))
|
||||
}
|
||||
|
|
@ -281,30 +357,53 @@ fn serialize_failure(err: serde_json::Error) -> Error {
|
|||
))
|
||||
}
|
||||
|
||||
/// Python's httpx timeout bounds each phase and each read, never the whole exchange. A complete
|
||||
/// response keeps the client's overall deadline; a stream is bounded until its head arrives and
|
||||
/// then per read (see [`timed_chunks`]), so a long stream that keeps sending is not cut off.
|
||||
async fn send(
|
||||
http: &litellm_http::Client,
|
||||
authenticated: Authenticated,
|
||||
url: &str,
|
||||
body: &Value,
|
||||
timeout: Option<Duration>,
|
||||
timeout: Duration,
|
||||
stream: bool,
|
||||
) -> Result<reqwest::Response, Error> {
|
||||
let request = outbound_request(
|
||||
authenticated,
|
||||
url,
|
||||
body,
|
||||
Some(timeout.unwrap_or(Duration::from_secs(MESSAGES_TIMEOUT_SECS))),
|
||||
)?;
|
||||
Ok(litellm_inference::outbound::send(request, http).await?)
|
||||
let request = outbound_request(authenticated, url, body, (!stream).then_some(timeout))?;
|
||||
if !stream {
|
||||
return Ok(litellm_inference::outbound::send(request, http).await?);
|
||||
}
|
||||
tokio::time::timeout(timeout, litellm_inference::outbound::send(request, http))
|
||||
.await
|
||||
.map_err(litellm_http::Error::read_timeout)?
|
||||
.map_err(Error::from)
|
||||
}
|
||||
|
||||
async fn provider_error(response: reqwest::Response) -> Error {
|
||||
/// The stream's body, each read bounded by `timeout` as httpx bounds each read.
|
||||
fn timed_chunks(
|
||||
response: reqwest::Response,
|
||||
timeout: Duration,
|
||||
) -> impl futures_util::Stream<Item = Result<Bytes, litellm_http::Error>> + Send + 'static {
|
||||
futures_util::stream::try_unfold(response, move |mut response| async move {
|
||||
let chunk = tokio::time::timeout(timeout, response.chunk())
|
||||
.await
|
||||
.map_err(litellm_http::Error::body_read_timeout)?
|
||||
.map_err(litellm_http::Error::from)?;
|
||||
Ok(chunk.map(|chunk| (chunk, response)))
|
||||
})
|
||||
}
|
||||
|
||||
/// The provider's error, its body read under `timeout`: a stream has no overall deadline once
|
||||
/// its head arrives, so an error body that stalls would otherwise hold the call forever.
|
||||
async fn provider_error(response: reqwest::Response, timeout: Duration) -> Error {
|
||||
let status = response.status().as_u16();
|
||||
match response.text().await {
|
||||
Ok(text) => {
|
||||
let headers = litellm_http::request::header_pairs(response.headers());
|
||||
match tokio::time::timeout(timeout, response.text()).await {
|
||||
Ok(Ok(text)) => {
|
||||
log_error_body(status, &text);
|
||||
litellm_http::Error::upstream(status, text).into()
|
||||
litellm_http::Error::upstream_with_headers(status, text, headers).into()
|
||||
}
|
||||
Err(error) => litellm_http::Error::from(error).into(),
|
||||
Ok(Err(error)) => litellm_http::Error::from(error).into(),
|
||||
Err(elapsed) => litellm_http::Error::read_timeout(elapsed).into(),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -329,22 +428,22 @@ fn streaming_response(
|
|||
response: reqwest::Response,
|
||||
decoder: Option<StreamDecoder>,
|
||||
provider: &'static str,
|
||||
timeout: Duration,
|
||||
) -> MessagesCallResponse {
|
||||
let headers = response
|
||||
.headers()
|
||||
.iter()
|
||||
.filter_map(|(name, value)| Some((name.to_string(), value.to_str().ok()?.to_string())))
|
||||
.collect();
|
||||
let chunks = timed_chunks(response, timeout)
|
||||
.inspect_ok(move |chunk| log_chunk(provider, "provider_response", chunk));
|
||||
let chunks = match decoder {
|
||||
None => futures_util::stream::try_unfold(response, move |mut response| async move {
|
||||
let chunk = response.chunk().await.map_err(litellm_http::Error::from)?;
|
||||
Ok(chunk.map(|chunk| {
|
||||
log_chunk(provider, "provider_response", &chunk);
|
||||
(chunk, response)
|
||||
}))
|
||||
})
|
||||
.boxed(),
|
||||
Some(decode) => decoded_chunks(response, decode, provider),
|
||||
None => chunks.map_err(Error::from).boxed(),
|
||||
Some(decode) => decoded_chunks(
|
||||
chunks.map_err(std::io::Error::other).boxed(),
|
||||
decode,
|
||||
provider,
|
||||
),
|
||||
};
|
||||
MessagesCallResponse::Stream {
|
||||
head: super::route::MessagesStreamHead { headers },
|
||||
|
|
@ -353,15 +452,10 @@ fn streaming_response(
|
|||
}
|
||||
|
||||
fn decoded_chunks(
|
||||
response: reqwest::Response,
|
||||
bytes: ByteStream,
|
||||
decode: StreamDecoder,
|
||||
provider: &'static str,
|
||||
) -> BoxStream<'static, Result<Bytes, Error>> {
|
||||
let bytes: ByteStream = response
|
||||
.bytes_stream()
|
||||
.inspect_ok(move |chunk| log_chunk(provider, "provider_response", chunk))
|
||||
.map_err(|error| std::io::Error::other(litellm_http::Error::from(error)))
|
||||
.boxed();
|
||||
futures_util::stream::try_unfold(decode(bytes), move |mut events| async move {
|
||||
let Some(event) = events.try_next().await? else {
|
||||
return Ok(None);
|
||||
|
|
@ -964,9 +1058,12 @@ mod tests {
|
|||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
let MessagesCallResponse::Stream { mut chunks, .. } =
|
||||
streaming_response(response, Some(streaming::event_stream), "test")
|
||||
else {
|
||||
let MessagesCallResponse::Stream { mut chunks, .. } = streaming_response(
|
||||
response,
|
||||
Some(streaming::event_stream),
|
||||
"test",
|
||||
Duration::from_secs(5),
|
||||
) else {
|
||||
panic!("a streaming response returns chunks");
|
||||
};
|
||||
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ use std::sync::Arc;
|
|||
|
||||
use crate::common_utils::messages_provider;
|
||||
|
||||
pub use common_utils::merge_messages_header_sources;
|
||||
pub use litellm_inference::Error;
|
||||
pub use types::{
|
||||
MessagesCall, MessagesCallResponse, MessagesSettings, MessagesShaping, litellm_params,
|
||||
|
|
|
|||
|
|
@ -440,7 +440,12 @@ async fn an_oauth_key_sends_the_browser_access_header_and_the_oauth_beta(call: M
|
|||
#[case::deepseek("deepseek")]
|
||||
#[case::minimax("minimax")]
|
||||
#[tokio::test]
|
||||
async fn caller_protocol_headers_win_over_the_defaults(call: MessagesCall, #[case] provider: &str) {
|
||||
/// Python adds a default unless that exact lowercase name is present, then folds the casings
|
||||
/// into one header, so a capitalised caller header is sent beside the default.
|
||||
async fn capitalised_caller_protocol_headers_are_joined_with_the_defaults(
|
||||
call: MessagesCall,
|
||||
#[case] provider: &str,
|
||||
) {
|
||||
let upstream = upstream([message_response()]).await;
|
||||
|
||||
run_message(MessagesCall {
|
||||
|
|
@ -457,10 +462,13 @@ async fn caller_protocol_headers_win_over_the_defaults(call: MessagesCall, #[cas
|
|||
.await;
|
||||
|
||||
let request = only_request(&upstream).await;
|
||||
assert_eq!(request.header_values("anthropic-version"), ["2024-01-01"]);
|
||||
assert_eq!(
|
||||
request.header_values("anthropic-version"),
|
||||
["2024-01-01, 2023-06-01"]
|
||||
);
|
||||
assert_eq!(
|
||||
request.header_values("content-type"),
|
||||
["application/json; charset=utf-8"]
|
||||
["application/json; charset=utf-8, application/json"]
|
||||
);
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -225,14 +225,25 @@ async fn an_upstream_error_keeps_its_status_and_body(call: MessagesCall, #[case]
|
|||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::not_json(ResponseTemplate::new(200).set_body_string("not json"))]
|
||||
#[case::not_a_message(json_response(json!({"unexpected": true})))]
|
||||
#[tokio::test]
|
||||
async fn an_unreadable_success_body_is_an_invalid_response(
|
||||
call: MessagesCall,
|
||||
#[case] response: ResponseTemplate,
|
||||
) {
|
||||
let upstream = upstream([response]).await;
|
||||
async fn an_unparseable_success_body_is_reported_as_the_upstream_answer(call: MessagesCall) {
|
||||
let upstream = upstream([ResponseTemplate::new(200).set_body_string("not json")]).await;
|
||||
|
||||
let error = run(MessagesCall {
|
||||
api_key: Some("sk".into()),
|
||||
api_base: Some(upstream.uri()),
|
||||
..call
|
||||
})
|
||||
.await
|
||||
.expect_err("an unparseable body fails");
|
||||
|
||||
assert_eq!(error.upstream_response(), Some((200, "not json")));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn an_unreadable_success_body_is_an_invalid_response(call: MessagesCall) {
|
||||
let upstream = upstream([json_response(json!({"unexpected": true}))]).await;
|
||||
|
||||
let error = run(MessagesCall {
|
||||
api_key: Some("sk".into()),
|
||||
|
|
|
|||
|
|
@ -432,6 +432,49 @@ async fn dropping_the_sdk_stream_closes_the_unfinished_upstream(
|
|||
.unwrap();
|
||||
}
|
||||
|
||||
/// Answers with an error head whose promised body never arrives.
|
||||
async fn stalling_error_upstream() -> (String, JoinHandle<()>) {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let base = format!("http://{}", listener.local_addr().unwrap());
|
||||
let connection = tokio::spawn(async move {
|
||||
let (mut socket, _) = listener.accept().await.unwrap();
|
||||
let mut request = vec![0; 4096];
|
||||
let _ = socket.read(&mut request).await;
|
||||
socket
|
||||
.write_all(b"HTTP/1.1 429 Too Many Requests\r\ncontent-type: application/json\r\ncontent-length: 100\r\n\r\n")
|
||||
.await
|
||||
.unwrap();
|
||||
let _ = socket.read_to_end(&mut Vec::new()).await;
|
||||
});
|
||||
(base, connection)
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn the_timeout_covers_a_stalled_stream_error_body(call: MessagesCall) {
|
||||
let (base, connection) = stalling_error_upstream().await;
|
||||
let host = RecordingStreamHost::new(
|
||||
MessagesCall {
|
||||
timeout: Some(Duration::from_millis(300)),
|
||||
..streaming(call, base)
|
||||
},
|
||||
usize::MAX,
|
||||
);
|
||||
|
||||
let error = tokio::time::timeout(Duration::from_secs(5), stream_through(&host))
|
||||
.await
|
||||
.expect("the stalled error body gives up within the timeout")
|
||||
.expect_err("a stalled error body fails the call");
|
||||
|
||||
assert!(error.dispatched(), "{error:?}");
|
||||
assert_eq!(error.upstream_response(), None);
|
||||
assert!(host.seen.into_inner().unwrap().is_empty());
|
||||
tokio::time::timeout(Duration::from_secs(5), connection)
|
||||
.await
|
||||
.expect("timing out closes the upstream connection")
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn the_sdk_yields_a_body_error_once_after_delivered_chunks(call: MessagesCall) {
|
||||
|
|
|
|||
|
|
@ -91,9 +91,7 @@ impl<H: Interceptors<Error>> CallHooks<Error> for OcrCallHooks<'_, H> {
|
|||
}
|
||||
|
||||
fn response_received<'a>(&'a self, body: &'a [u8]) -> BoxFuture<'a, Result<(), Error>> {
|
||||
let raw = RawResponse {
|
||||
body: String::from_utf8_lossy(body).into_owned(),
|
||||
};
|
||||
let raw = RawResponse::new(String::from_utf8_lossy(body).into_owned());
|
||||
if let Some(observers) = self.observers {
|
||||
observers.emit(litellm_host::lifecycle::CallEvent::Execution(
|
||||
ExecutionEvent::ProviderResponseReceived { raw: raw.clone() },
|
||||
|
|
|
|||
|
|
@ -114,7 +114,7 @@ pub(super) async fn execute(
|
|||
});
|
||||
}
|
||||
let body = response.text().await.map_err(litellm_http::Error::from)?;
|
||||
let raw = RawResponse { body: body.clone() };
|
||||
let raw = RawResponse::new(body.clone());
|
||||
if let Some(observers) = observers {
|
||||
observers.emit(litellm_host::lifecycle::CallEvent::Execution(
|
||||
ExecutionEvent::ProviderResponseReceived { raw: raw.clone() },
|
||||
|
|
|
|||
|
|
@ -42,9 +42,11 @@ impl<'a, I: Interceptors<Error>> CallContext<'a, I> {
|
|||
}
|
||||
|
||||
pub async fn response_received(&self, body: &str) -> Result<(), Error> {
|
||||
let raw = RawResponse {
|
||||
body: body.to_owned(),
|
||||
};
|
||||
self.raw_response_received(RawResponse::new(body.to_owned()))
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn raw_response_received(&self, raw: RawResponse) -> Result<(), Error> {
|
||||
if let Some(observers) = &self.observers {
|
||||
observers.emit(CallEvent::Execution(
|
||||
ExecutionEvent::ProviderResponseReceived { raw: raw.clone() },
|
||||
|
|
|
|||
|
|
@ -146,6 +146,34 @@ impl Error {
|
|||
}
|
||||
}
|
||||
|
||||
/// The provider's response headers on an upstream failure; empty for every other error.
|
||||
pub fn upstream_headers(&self) -> &[(String, String)] {
|
||||
match self {
|
||||
Self::Llm(error) => error.upstream_headers(),
|
||||
Self::Http(error) => error.upstream_headers(),
|
||||
_ => &[],
|
||||
}
|
||||
}
|
||||
|
||||
/// A read of a streamed body timed out after the response head arrived.
|
||||
pub fn body_read_timeout(&self) -> bool {
|
||||
match self {
|
||||
Self::Llm(error) => error.is_body_read_timeout(),
|
||||
Self::Http(error) => error.is_body_read_timeout(),
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
/// `Some(connect)` when the call timed out, where `connect` says the connection was never
|
||||
/// established; `None` for every other error.
|
||||
pub fn timeout(&self) -> Option<bool> {
|
||||
match self {
|
||||
Self::Llm(error) => error.is_timeout(),
|
||||
Self::Http(error) if error.is_timeout() => Some(error.is_connect_timeout()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn upstream_response(&self) -> Option<(u16, &str)> {
|
||||
match self {
|
||||
Self::Llm(error) => error.upstream_response(),
|
||||
|
|
|
|||
|
|
@ -264,7 +264,12 @@ pub enum ThinkingDisplay {
|
|||
#[macro_rules_attribute::apply(crate::wire_type)]
|
||||
#[derive(Default)]
|
||||
pub struct EnabledThinking {
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
/// An explicit `null` stays `Unrecognized(Null)`: Python rejects it, but defaults a missing key to 0.
|
||||
#[serde(
|
||||
default,
|
||||
skip_serializing_if = "Option::is_none",
|
||||
deserialize_with = "crate::recognized::present"
|
||||
)]
|
||||
pub budget_tokens: Option<Recognized<u64>>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub display: Option<Recognized<ThinkingDisplay>>,
|
||||
|
|
|
|||
|
|
@ -10,14 +10,29 @@ pub struct MessagesResponse {
|
|||
pub content: Vec<Value>,
|
||||
pub stop_reason: Option<String>,
|
||||
pub stop_sequence: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
/// `Some(Value::Null)` when the provider sent an explicit `null`, which Python keeps.
|
||||
#[serde(
|
||||
default,
|
||||
deserialize_with = "present_value",
|
||||
skip_serializing_if = "Option::is_none"
|
||||
)]
|
||||
pub usage: Option<Value>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
#[serde(
|
||||
default,
|
||||
deserialize_with = "present_value",
|
||||
skip_serializing_if = "Option::is_none"
|
||||
)]
|
||||
pub container: Option<Value>,
|
||||
#[serde(flatten)]
|
||||
pub extra: Map<String, Value>,
|
||||
}
|
||||
|
||||
fn present_value<'de, D: serde::Deserializer<'de>>(
|
||||
deserializer: D,
|
||||
) -> Result<Option<Value>, D::Error> {
|
||||
<Value as serde::Deserialize<'de>>::deserialize(deserializer).map(Some)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
|
@ -73,4 +88,21 @@ mod tests {
|
|||
assert_eq!(body.get("usage").cloned(), usage);
|
||||
assert_eq!(body.get("container").cloned(), container);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn explicit_null_usage_and_container_survive_a_round_trip() {
|
||||
let body = json!({"id": "msg_1", "type": "message", "role": "assistant", "model": "m",
|
||||
"content": [], "stop_reason": null, "stop_sequence": null,
|
||||
"usage": null, "container": null});
|
||||
let parsed: MessagesResponse = serde_json::from_value(body.clone()).expect("parses");
|
||||
assert_eq!(serde_json::to_value(parsed).expect("serializable"), body);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn absent_usage_and_container_stay_absent() {
|
||||
let body = json!({"id": "msg_1", "type": "message", "role": "assistant", "model": "m",
|
||||
"content": [], "stop_reason": null, "stop_sequence": null});
|
||||
let parsed: MessagesResponse = serde_json::from_value(body.clone()).expect("parses");
|
||||
assert_eq!(serde_json::to_value(parsed).expect("serializable"), body);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -124,12 +124,86 @@ pub enum AnthropicBeta {
|
|||
ToolSearchTool20251019,
|
||||
#[strum(serialize = "inline-tools-2026-09-15")]
|
||||
InlineTools20260915,
|
||||
#[strum(serialize = "afk-mode-2026-01-31")]
|
||||
AfkMode20260131,
|
||||
#[strum(serialize = "agent-memory-2026-07-22")]
|
||||
AgentMemory20260722,
|
||||
#[strum(serialize = "analytics-usage-cost-2026-08-24")]
|
||||
AnalyticsUsageCost20260824,
|
||||
#[strum(serialize = "cache-diagnosis-2026-04-07")]
|
||||
CacheDiagnosis20260407,
|
||||
#[strum(serialize = "ce-plugins-2026-09-01")]
|
||||
CePlugins20260901,
|
||||
#[strum(serialize = "ce-user-management-2026-07-13")]
|
||||
CeUserManagement20260713,
|
||||
#[strum(serialize = "code-execution-2025-05-22")]
|
||||
CodeExecution20250522,
|
||||
#[strum(serialize = "computer-use-2024-10-22")]
|
||||
ComputerUse20241022,
|
||||
#[strum(serialize = "dev-full-thinking-2025-05-14")]
|
||||
DevFullThinking20250514,
|
||||
#[strum(serialize = "dreaming-2026-04-21")]
|
||||
Dreaming20260421,
|
||||
#[strum(serialize = "experimental-claude-tag-2026-10-21")]
|
||||
ExperimentalClaudeTag20261021,
|
||||
#[strum(serialize = "extended-cache-ttl-2025-04-11")]
|
||||
ExtendedCacheTtl20250411,
|
||||
#[strum(serialize = "fallback-credit-2026-06-01")]
|
||||
FallbackCredit20260601,
|
||||
#[strum(serialize = "fallback-credit-2026-07-01")]
|
||||
FallbackCredit20260701,
|
||||
#[strum(serialize = "managed-agents-2026-04-01")]
|
||||
ManagedAgents20260401,
|
||||
#[strum(serialize = "mcp-client-2026-09-15")]
|
||||
McpClient20260915,
|
||||
#[strum(serialize = "mcp-tunnels-2026-06-22")]
|
||||
McpTunnels20260622,
|
||||
#[strum(serialize = "message-batches-2024-09-24")]
|
||||
MessageBatches20240924,
|
||||
#[strum(serialize = "message-threads-2026-08-12")]
|
||||
MessageThreads20260812,
|
||||
#[strum(serialize = "mid-conversation-system-clear-at-2026-08-21")]
|
||||
MidConversationSystemClearAt20260821,
|
||||
#[strum(serialize = "model-context-window-exceeded-2025-08-26")]
|
||||
ModelContextWindowExceeded20250826,
|
||||
#[strum(serialize = "output-300k-2026-03-24")]
|
||||
Output300k20260324,
|
||||
#[strum(serialize = "pdfs-2024-09-25")]
|
||||
Pdfs20240925,
|
||||
#[strum(serialize = "prompt-caching-2024-07-31")]
|
||||
PromptCaching20240731,
|
||||
#[strum(serialize = "redact-thinking-2026-02-12")]
|
||||
RedactThinking20260212,
|
||||
#[strum(serialize = "server-side-fallback-2026-06-01")]
|
||||
ServerSideFallback20260601,
|
||||
#[strum(serialize = "server-side-fallback-2026-07-01")]
|
||||
ServerSideFallback20260701,
|
||||
#[strum(serialize = "spend-limit-reads-2026-09-26")]
|
||||
SpendLimitReads20260926,
|
||||
#[strum(serialize = "structured-outputs-2025-12-15")]
|
||||
StructuredOutputs20251215,
|
||||
#[strum(serialize = "task-budgets-2026-03-13")]
|
||||
TaskBudgets20260313,
|
||||
#[strum(serialize = "telemetry-destinations-2026-08-11")]
|
||||
TelemetryDestinations20260811,
|
||||
#[strum(serialize = "thinking-token-count-2026-05-13")]
|
||||
ThinkingTokenCount20260513,
|
||||
#[strum(serialize = "token-counting-2024-11-01")]
|
||||
TokenCounting20241101,
|
||||
#[strum(serialize = "user-profiles-2026-03-24")]
|
||||
UserProfiles20260324,
|
||||
#[strum(serialize = "user-profiles-2026-08-18")]
|
||||
UserProfiles20260818,
|
||||
#[strum(serialize = "user-profiles-2026-09-04")]
|
||||
UserProfiles20260904,
|
||||
#[strum(serialize = "webhooks-2026-10-07")]
|
||||
Webhooks20261007,
|
||||
#[strum(default, transparent)]
|
||||
Other(String),
|
||||
}
|
||||
|
||||
impl AnthropicBeta {
|
||||
pub const KNOWN: [Self; 40] = [
|
||||
pub const KNOWN: [Self; 77] = [
|
||||
Self::Oauth20250420,
|
||||
Self::WebFetch20250910,
|
||||
Self::WebSearch20250305,
|
||||
|
|
@ -170,6 +244,43 @@ impl AnthropicBeta {
|
|||
Self::ToolExamples20251029,
|
||||
Self::ToolSearchTool20251019,
|
||||
Self::InlineTools20260915,
|
||||
Self::AfkMode20260131,
|
||||
Self::AgentMemory20260722,
|
||||
Self::AnalyticsUsageCost20260824,
|
||||
Self::CacheDiagnosis20260407,
|
||||
Self::CePlugins20260901,
|
||||
Self::CeUserManagement20260713,
|
||||
Self::CodeExecution20250522,
|
||||
Self::ComputerUse20241022,
|
||||
Self::DevFullThinking20250514,
|
||||
Self::Dreaming20260421,
|
||||
Self::ExperimentalClaudeTag20261021,
|
||||
Self::ExtendedCacheTtl20250411,
|
||||
Self::FallbackCredit20260601,
|
||||
Self::FallbackCredit20260701,
|
||||
Self::ManagedAgents20260401,
|
||||
Self::McpClient20260915,
|
||||
Self::McpTunnels20260622,
|
||||
Self::MessageBatches20240924,
|
||||
Self::MessageThreads20260812,
|
||||
Self::MidConversationSystemClearAt20260821,
|
||||
Self::ModelContextWindowExceeded20250826,
|
||||
Self::Output300k20260324,
|
||||
Self::Pdfs20240925,
|
||||
Self::PromptCaching20240731,
|
||||
Self::RedactThinking20260212,
|
||||
Self::ServerSideFallback20260601,
|
||||
Self::ServerSideFallback20260701,
|
||||
Self::SpendLimitReads20260926,
|
||||
Self::StructuredOutputs20251215,
|
||||
Self::TaskBudgets20260313,
|
||||
Self::TelemetryDestinations20260811,
|
||||
Self::ThinkingTokenCount20260513,
|
||||
Self::TokenCounting20241101,
|
||||
Self::UserProfiles20260324,
|
||||
Self::UserProfiles20260818,
|
||||
Self::UserProfiles20260904,
|
||||
Self::Webhooks20261007,
|
||||
];
|
||||
|
||||
pub fn as_str(&self) -> &str {
|
||||
|
|
@ -195,10 +306,9 @@ impl AnthropicBeta {
|
|||
|
||||
fn rejected_by(&self) -> &'static [BetaProvider] {
|
||||
match self {
|
||||
Self::AdvancedToolUse20251120 => {
|
||||
Self::AdvancedToolUse20251120 | Self::ToolSearchTool20251019 => {
|
||||
&[BetaProvider::BedrockConverse, BetaProvider::Databricks]
|
||||
}
|
||||
Self::ContextManagement20250627 => &[BetaProvider::BedrockConverse],
|
||||
Self::AdvisorTool20260301 => &[
|
||||
BetaProvider::AzureAi,
|
||||
BetaProvider::BedrockConverse,
|
||||
|
|
@ -207,98 +317,99 @@ impl AnthropicBeta {
|
|||
BetaProvider::VertexAi,
|
||||
BetaProvider::Databricks,
|
||||
],
|
||||
Self::Compact20260904 => &[
|
||||
Self::AfkMode20260131
|
||||
| Self::CePlugins20260901
|
||||
| Self::CeUserManagement20260713
|
||||
| Self::Compact20260904
|
||||
| Self::Dreaming20260421
|
||||
| Self::ExperimentalClaudeTag20261021
|
||||
| Self::ExtendedCacheTtl20250411
|
||||
| Self::FallbackCredit20260601
|
||||
| Self::FallbackCredit20260701
|
||||
| Self::InlineTools20260915
|
||||
| Self::ManagedAgents20260401
|
||||
| Self::McpTunnels20260622
|
||||
| Self::MidConversationSystemClearAt20260821
|
||||
| Self::ModelContextWindowExceeded20250826
|
||||
| Self::Oauth20250420
|
||||
| Self::Pdfs20240925
|
||||
| Self::PromptCaching20240731
|
||||
| Self::PromptCachingScope20260105
|
||||
| Self::RedactThinking20260212
|
||||
| Self::ServerSideFallback20260601
|
||||
| Self::SpendLimitReads20260926
|
||||
| Self::StructuredOutputs20251215
|
||||
| Self::TaskBudgets20260313
|
||||
| Self::TelemetryDestinations20260811
|
||||
| Self::ThinkingTokenCount20260513
|
||||
| Self::TokenCounting20241101
|
||||
| Self::UserProfiles20260324
|
||||
| Self::UserProfiles20260818
|
||||
| Self::UserProfiles20260904 => &[
|
||||
BetaProvider::BedrockConverse,
|
||||
BetaProvider::Bedrock,
|
||||
BetaProvider::BedrockMantle,
|
||||
BetaProvider::Databricks,
|
||||
],
|
||||
Self::Bash20241022
|
||||
Self::AgentMemory20260722
|
||||
| Self::Bash20241022
|
||||
| Self::Bash20250124
|
||||
| Self::DevFullThinking20250514
|
||||
| Self::McpServers20251204
|
||||
| Self::StructuredOutput20240301
|
||||
| Self::TextEditor20241022
|
||||
| Self::TextEditor20250124 => BetaProvider::VARIANTS,
|
||||
Self::ClaudeCode20250219 | Self::ToolExamples20251029 => &[
|
||||
BetaProvider::Anthropic,
|
||||
BetaProvider::AzureAi,
|
||||
BetaProvider::BedrockConverse,
|
||||
BetaProvider::Bedrock,
|
||||
BetaProvider::VertexAi,
|
||||
BetaProvider::Databricks,
|
||||
],
|
||||
Self::CodeExecution20250825
|
||||
Self::AnalyticsUsageCost20260824
|
||||
| Self::CacheDiagnosis20260407
|
||||
| Self::CodeExecution20250522
|
||||
| Self::CodeExecution20250825
|
||||
| Self::ComputerUse20241022
|
||||
| Self::FastMode20260201
|
||||
| Self::FilesApi20250414
|
||||
| Self::McpClient20250404
|
||||
| Self::McpClient20251120
|
||||
| Self::PromptCachingScope20260105
|
||||
| Self::McpClient20260915
|
||||
| Self::MessageBatches20240924
|
||||
| Self::MessageThreads20260812
|
||||
| Self::Output300k20260324
|
||||
| Self::ServerSideFallback20260701
|
||||
| Self::Skills20251002
|
||||
| Self::WebFetch20250910 => &[
|
||||
| Self::WebFetch20250910
|
||||
| Self::Webhooks20261007 => &[
|
||||
BetaProvider::BedrockConverse,
|
||||
BetaProvider::Bedrock,
|
||||
BetaProvider::BedrockMantle,
|
||||
BetaProvider::VertexAi,
|
||||
BetaProvider::Databricks,
|
||||
],
|
||||
Self::Compact20260112 => &[BetaProvider::AzureAi, BetaProvider::BedrockConverse],
|
||||
Self::ClaudeCode20250219
|
||||
| Self::PerTurnControl20260701
|
||||
| Self::ToolExamples20251029 => &[
|
||||
BetaProvider::BedrockConverse,
|
||||
BetaProvider::Bedrock,
|
||||
BetaProvider::Databricks,
|
||||
],
|
||||
Self::Compact20260112
|
||||
| Self::ContextManagement20250627
|
||||
| Self::DangerousToolUse20260903 => &[BetaProvider::BedrockConverse],
|
||||
Self::ComputerUse20250124
|
||||
| Self::ComputerUse20251124
|
||||
| Self::Context1m20250807
|
||||
| Self::Effort20251124
|
||||
| Self::FineGrainedToolStreaming20250514
|
||||
| Self::ThinkingBindingControls20260801
|
||||
| Self::ThinkingDisplayUpdates20260818 => &[],
|
||||
Self::DangerousToolUse20260903 => &[BetaProvider::BedrockConverse],
|
||||
Self::Effort20251124 => &[BetaProvider::VertexAi],
|
||||
Self::FastMode20260201 | Self::Oauth20250420 => &[
|
||||
BetaProvider::AzureAi,
|
||||
BetaProvider::BedrockConverse,
|
||||
BetaProvider::Bedrock,
|
||||
BetaProvider::BedrockMantle,
|
||||
BetaProvider::VertexAi,
|
||||
BetaProvider::Databricks,
|
||||
],
|
||||
Self::FineGrainedToolStreaming20250514 => {
|
||||
&[BetaProvider::AzureAi, BetaProvider::VertexAi]
|
||||
}
|
||||
Self::InterleavedThinking20250514 | Self::WebSearch20250305 => {
|
||||
&[BetaProvider::BedrockConverse, BetaProvider::Bedrock]
|
||||
}
|
||||
Self::MidConversationOutputConfig20260701 => &[
|
||||
BetaProvider::AzureAi,
|
||||
BetaProvider::BedrockConverse,
|
||||
BetaProvider::BedrockMantle,
|
||||
BetaProvider::VertexAi,
|
||||
BetaProvider::Databricks,
|
||||
],
|
||||
Self::MidConversationToolChanges20260701 => &[
|
||||
Self::InterleavedThinking20250514
|
||||
| Self::Output128k20250219
|
||||
| Self::TokenEfficientTools20250219
|
||||
| Self::WebSearch20250305 => &[BetaProvider::BedrockConverse, BetaProvider::Bedrock],
|
||||
Self::MidConversationOutputConfig20260701
|
||||
| Self::MidConversationToolChanges20260701 => &[
|
||||
BetaProvider::BedrockConverse,
|
||||
BetaProvider::BedrockMantle,
|
||||
BetaProvider::Databricks,
|
||||
],
|
||||
Self::Output128k20250219 | Self::TokenEfficientTools20250219 => &[
|
||||
BetaProvider::AzureAi,
|
||||
BetaProvider::BedrockConverse,
|
||||
BetaProvider::Bedrock,
|
||||
BetaProvider::VertexAi,
|
||||
],
|
||||
Self::PerTurnControl20260701 => &[
|
||||
BetaProvider::BedrockConverse,
|
||||
BetaProvider::Bedrock,
|
||||
BetaProvider::Databricks,
|
||||
],
|
||||
Self::StructuredOutputs20251113 => &[BetaProvider::Bedrock, BetaProvider::VertexAi],
|
||||
Self::ToolSearchTool20251019 => &[
|
||||
BetaProvider::Anthropic,
|
||||
BetaProvider::AzureAi,
|
||||
BetaProvider::BedrockConverse,
|
||||
BetaProvider::Databricks,
|
||||
],
|
||||
Self::InlineTools20260915 => &[
|
||||
BetaProvider::AzureAi,
|
||||
BetaProvider::BedrockConverse,
|
||||
BetaProvider::Bedrock,
|
||||
BetaProvider::BedrockMantle,
|
||||
BetaProvider::Databricks,
|
||||
],
|
||||
Self::StructuredOutputs20251113 => &[BetaProvider::Bedrock],
|
||||
Self::Other(_) => BetaProvider::VARIANTS,
|
||||
}
|
||||
}
|
||||
|
|
@ -451,7 +562,44 @@ mod tests {
|
|||
AnthropicBeta::TokenEfficientTools20250219,
|
||||
AnthropicBeta::ToolExamples20251029,
|
||||
AnthropicBeta::ToolSearchTool20251019,
|
||||
AnthropicBeta::InlineTools20260915
|
||||
AnthropicBeta::InlineTools20260915,
|
||||
AnthropicBeta::AfkMode20260131,
|
||||
AnthropicBeta::AgentMemory20260722,
|
||||
AnthropicBeta::AnalyticsUsageCost20260824,
|
||||
AnthropicBeta::CacheDiagnosis20260407,
|
||||
AnthropicBeta::CePlugins20260901,
|
||||
AnthropicBeta::CeUserManagement20260713,
|
||||
AnthropicBeta::CodeExecution20250522,
|
||||
AnthropicBeta::ComputerUse20241022,
|
||||
AnthropicBeta::DevFullThinking20250514,
|
||||
AnthropicBeta::Dreaming20260421,
|
||||
AnthropicBeta::ExperimentalClaudeTag20261021,
|
||||
AnthropicBeta::ExtendedCacheTtl20250411,
|
||||
AnthropicBeta::FallbackCredit20260601,
|
||||
AnthropicBeta::FallbackCredit20260701,
|
||||
AnthropicBeta::ManagedAgents20260401,
|
||||
AnthropicBeta::McpClient20260915,
|
||||
AnthropicBeta::McpTunnels20260622,
|
||||
AnthropicBeta::MessageBatches20240924,
|
||||
AnthropicBeta::MessageThreads20260812,
|
||||
AnthropicBeta::MidConversationSystemClearAt20260821,
|
||||
AnthropicBeta::ModelContextWindowExceeded20250826,
|
||||
AnthropicBeta::Output300k20260324,
|
||||
AnthropicBeta::Pdfs20240925,
|
||||
AnthropicBeta::PromptCaching20240731,
|
||||
AnthropicBeta::RedactThinking20260212,
|
||||
AnthropicBeta::ServerSideFallback20260601,
|
||||
AnthropicBeta::ServerSideFallback20260701,
|
||||
AnthropicBeta::SpendLimitReads20260926,
|
||||
AnthropicBeta::StructuredOutputs20251215,
|
||||
AnthropicBeta::TaskBudgets20260313,
|
||||
AnthropicBeta::TelemetryDestinations20260811,
|
||||
AnthropicBeta::ThinkingTokenCount20260513,
|
||||
AnthropicBeta::TokenCounting20241101,
|
||||
AnthropicBeta::UserProfiles20260324,
|
||||
AnthropicBeta::UserProfiles20260818,
|
||||
AnthropicBeta::UserProfiles20260904,
|
||||
AnthropicBeta::Webhooks20261007
|
||||
)]
|
||||
beta: AnthropicBeta,
|
||||
) {
|
||||
|
|
@ -551,7 +699,7 @@ mod tests {
|
|||
BetaProvider::Anthropic,
|
||||
Some(AnthropicBeta::Oauth20250420)
|
||||
)]
|
||||
#[case::rejected("effort-2025-11-24", BetaProvider::VertexAi, None)]
|
||||
#[case::rejected("skills-2025-10-02", BetaProvider::VertexAi, None)]
|
||||
#[case::unknown("example-beta-2099-01-01", BetaProvider::Anthropic, None)]
|
||||
fn on_resolves_known_spellings_held_as_other(
|
||||
#[case] raw: &str,
|
||||
|
|
@ -580,12 +728,11 @@ mod tests {
|
|||
|
||||
#[test]
|
||||
fn unknown_values_are_kept_verbatim() {
|
||||
let parsed: AnthropicBeta = "claude-code-20250219".parse().unwrap();
|
||||
assert_eq!(
|
||||
parsed,
|
||||
AnthropicBeta::Other("claude-code-20250219".to_string())
|
||||
let parsed: AnthropicBeta = "example-beta-2099-01-01".parse().unwrap();
|
||||
assert!(
|
||||
matches!(&parsed, AnthropicBeta::Other(value) if value == "example-beta-2099-01-01")
|
||||
);
|
||||
assert_eq!(parsed.to_string(), "claude-code-20250219");
|
||||
assert_eq!(parsed.to_string(), "example-beta-2099-01-01");
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
|
|
|||
|
|
@ -7,6 +7,16 @@ pub enum Recognized<T> {
|
|||
Unrecognized(Value),
|
||||
}
|
||||
|
||||
/// `deserialize_with` for an optional `Recognized` field whose explicit `null` must stay
|
||||
/// distinguishable from an absent key: `null` becomes `Some(Unrecognized(Null))`.
|
||||
pub fn present<'de, D, T>(deserializer: D) -> Result<Option<Recognized<T>>, D::Error>
|
||||
where
|
||||
D: serde::Deserializer<'de>,
|
||||
Recognized<T>: serde::Deserialize<'de>,
|
||||
{
|
||||
<Recognized<T> as serde::Deserialize<'de>>::deserialize(deserializer).map(Some)
|
||||
}
|
||||
|
||||
impl<T> Recognized<T> {
|
||||
pub fn known(&self) -> Option<&T> {
|
||||
match self {
|
||||
|
|
|
|||
|
|
@ -96,7 +96,10 @@ mod tests {
|
|||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::azure_ai(BetaProvider::AzureAi, "fast-mode-2026-02-01,example-beta-2099-01-01")]
|
||||
#[case::azure_ai(
|
||||
BetaProvider::AzureAi,
|
||||
"advisor-tool-2026-03-01,example-beta-2099-01-01"
|
||||
)]
|
||||
#[case::anthropic(BetaProvider::Anthropic, "bash_20241022")]
|
||||
#[case::vertex_ai(BetaProvider::VertexAi, " , ")]
|
||||
fn filter_removes_header_when_nothing_survives(
|
||||
|
|
|
|||
|
|
@ -154,11 +154,42 @@ pub fn complete_anthropic_url(
|
|||
) -> String {
|
||||
let api_base = resolve_anthropic_api_base(api_base, env_lookup);
|
||||
|
||||
let api_base = api_base.trim_end_matches('/');
|
||||
if api_base.ends_with(MESSAGES_PATH) {
|
||||
ensure_api_route(&api_base, MESSAGES_PATH)
|
||||
}
|
||||
|
||||
/// Python's `urlsplit(url).path`, for the URLs an `api_base` holds.
|
||||
fn url_path(url: &str) -> &str {
|
||||
let before_fragment = url.split('#').next().unwrap_or_default();
|
||||
let before_query = before_fragment.split('?').next().unwrap_or_default();
|
||||
match before_query.split_once("://") {
|
||||
Some((_, rest)) => rest.find('/').map_or("", |slash| &rest[slash..]),
|
||||
None => before_query,
|
||||
}
|
||||
}
|
||||
|
||||
/// Python's `_strip_trailing_slashes_keeping_the_scheme`.
|
||||
fn strip_trailing_slashes(api_base: &str) -> String {
|
||||
let stripped = api_base.trim_end_matches('/');
|
||||
if stripped.ends_with(':') && !stripped.contains('/') && !stripped.contains('?') {
|
||||
return api_base.to_string();
|
||||
}
|
||||
format!("{api_base}{MESSAGES_PATH}")
|
||||
match stripped.split_once('?') {
|
||||
Some((before, query)) if url_path(stripped).ends_with('/') => {
|
||||
format!("{}?{query}", before.trim_end_matches('/'))
|
||||
}
|
||||
_ => stripped.to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Python's `ensure_api_route`: an `api_base` whose path already ends in the route is kept as
|
||||
/// it is, query string included; anything else gets the route appended to the whole string.
|
||||
fn ensure_api_route(api_base: &str, route: &str) -> String {
|
||||
let base = strip_trailing_slashes(api_base);
|
||||
let route = route.trim_start_matches('/');
|
||||
if url_path(&base).ends_with(&format!("/{route}")) {
|
||||
return base;
|
||||
}
|
||||
format!("{}/{route}", strip_trailing_slashes(&base))
|
||||
}
|
||||
|
||||
pub fn existing_betas(headers: &[(String, String)]) -> BetaSet {
|
||||
|
|
@ -1922,4 +1953,24 @@ mod tests {
|
|||
expected
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::bare_host("https://api.example.com", "https://api.example.com/v1/messages")]
|
||||
#[case::trailing_slash("https://api.example.com/", "https://api.example.com/v1/messages")]
|
||||
#[case::route_present(
|
||||
"https://api.example.com/v1/messages/",
|
||||
"https://api.example.com/v1/messages"
|
||||
)]
|
||||
#[case::route_and_query(
|
||||
"https://api.example.com/v1/messages?beta=true",
|
||||
"https://api.example.com/v1/messages?beta=true"
|
||||
)]
|
||||
#[case::query_without_route(
|
||||
"https://api.example.com/?x=1",
|
||||
"https://api.example.com?x=1/v1/messages"
|
||||
)]
|
||||
#[case::scheme_only("https://", "https:///v1/messages")]
|
||||
fn the_messages_url_follows_python_ensure_api_route(#[case] api_base: &str, #[case] url: &str) {
|
||||
assert_eq!(complete_anthropic_url(Some(api_base), &|_| None), url);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -2,8 +2,8 @@ use litellm_llms_types::{
|
|||
formats::{
|
||||
chat_completions::ReasoningEffort,
|
||||
messages::{
|
||||
EffortLevel, MessagesOptionalParams, MessagesRequest, OutputConfig, ThinkingConfig,
|
||||
ThinkingDisplay,
|
||||
AdaptiveThinking, EffortLevel, MessagesOptionalParams, MessagesRequest, OutputConfig,
|
||||
ThinkingConfig, ThinkingDisplay,
|
||||
},
|
||||
},
|
||||
recognized::Recognized,
|
||||
|
|
@ -36,16 +36,17 @@ fn budget_for_effort(budgets: &ThinkingBudgets, effort: ReasoningEffort) -> Opti
|
|||
|
||||
fn effort_for_budget(
|
||||
budgets: &ThinkingBudgets,
|
||||
budget_tokens: u64,
|
||||
budget_tokens: i64,
|
||||
capabilities: &MessagesModelCapabilities,
|
||||
) -> EffortLevel {
|
||||
if budget_tokens >= budgets.xhigh && capabilities.effort_tiers.xhigh {
|
||||
let reaches = |tier: u64| i128::from(budget_tokens) >= i128::from(tier);
|
||||
if reaches(budgets.xhigh) && capabilities.effort_tiers.xhigh {
|
||||
return EffortLevel::Xhigh;
|
||||
}
|
||||
if budget_tokens >= budgets.high {
|
||||
if reaches(budgets.high) {
|
||||
return EffortLevel::High;
|
||||
}
|
||||
if budget_tokens >= budgets.medium {
|
||||
if reaches(budgets.medium) {
|
||||
return EffortLevel::Medium;
|
||||
}
|
||||
EffortLevel::Low
|
||||
|
|
@ -231,32 +232,129 @@ fn drop_disabled_thinking(request: MessagesRequest, context: &ThinkingContext) -
|
|||
}
|
||||
}
|
||||
|
||||
/// `_numeric_budget_tokens`: Python reads any int, float (truncated) or int-parseable string
|
||||
/// as a budget, and nothing else.
|
||||
fn numeric_budget_tokens(budget: &Recognized<u64>) -> Option<i64> {
|
||||
let value = match budget {
|
||||
Recognized::Known(budget) => return Some(i64::try_from(*budget).unwrap_or(i64::MAX)),
|
||||
Recognized::Unrecognized(value) => value,
|
||||
};
|
||||
match value {
|
||||
Value::Number(number) => number
|
||||
.as_i64()
|
||||
.or_else(|| number.as_f64().map(|float| float.trunc() as i64)),
|
||||
Value::String(text) => python_int(text),
|
||||
Value::Null | Value::Bool(_) | Value::Array(_) | Value::Object(_) => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Python's `int(str)`: surrounding whitespace, a sign, and single underscores between digits.
|
||||
fn python_int(text: &str) -> Option<i64> {
|
||||
let trimmed = text.trim();
|
||||
let digits = trimmed.trim_start_matches(['+', '-']);
|
||||
let well_formed = trimmed.len() - digits.len() <= 1
|
||||
&& !digits.is_empty()
|
||||
&& !digits.starts_with('_')
|
||||
&& !digits.ends_with('_')
|
||||
&& !digits.contains("__")
|
||||
&& digits.chars().all(|ch| ch.is_ascii_digit() || ch == '_');
|
||||
if !well_formed {
|
||||
return None;
|
||||
}
|
||||
let cleaned: String = trimmed.chars().filter(|ch| *ch != '_').collect();
|
||||
cleaned.parse::<i64>().ok().or_else(|| {
|
||||
Some(if cleaned.starts_with('-') {
|
||||
i64::MIN
|
||||
} else {
|
||||
i64::MAX
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
fn python_type_name(value: &Value) -> &'static str {
|
||||
match value {
|
||||
Value::Null => "NoneType",
|
||||
Value::Bool(_) => "bool",
|
||||
Value::Number(number) if number.is_f64() => "float",
|
||||
Value::Number(_) => "int",
|
||||
Value::String(_) => "str",
|
||||
Value::Array(_) => "list",
|
||||
Value::Object(_) => "dict",
|
||||
}
|
||||
}
|
||||
|
||||
fn kept_display(
|
||||
display: Option<&Recognized<ThinkingDisplay>>,
|
||||
context: &ThinkingContext,
|
||||
) -> Option<Recognized<ThinkingDisplay>> {
|
||||
match display? {
|
||||
Recognized::Known(ThinkingDisplay::Summarized | ThinkingDisplay::Omitted) => {
|
||||
display.cloned()
|
||||
}
|
||||
Recognized::Known(ThinkingDisplay::Updates) if context.keeps_display_updates => {
|
||||
display.cloned()
|
||||
}
|
||||
Recognized::Known(ThinkingDisplay::Updates) | Recognized::Unrecognized(_) => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Keys Python's `carried_thinking_keys` keeps on the rebuilt adaptive config: the adaptive
|
||||
/// keys always, and any other key unless `drop_params` is set.
|
||||
fn carried_thinking_keys(
|
||||
extra: &serde_json::Map<String, Value>,
|
||||
context: &ThinkingContext,
|
||||
) -> serde_json::Map<String, Value> {
|
||||
extra
|
||||
.iter()
|
||||
.filter(|(key, _)| key.as_str() == "block_binding" || !context.drop_params)
|
||||
.map(|(key, value)| (key.clone(), value.clone()))
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// `AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model_silently`.
|
||||
fn translate_legacy_thinking_for_adaptive_model(
|
||||
request: MessagesRequest,
|
||||
context: &ThinkingContext,
|
||||
) -> MessagesRequest {
|
||||
) -> Result<MessagesRequest, Error> {
|
||||
let capabilities = &context.capabilities;
|
||||
if !capabilities.supports_adaptive_thinking || capabilities.supports_legacy_thinking {
|
||||
return request;
|
||||
return Ok(request);
|
||||
}
|
||||
let Some(ThinkingConfig::Enabled(enabled)) = known_thinking(&request) else {
|
||||
return request;
|
||||
return Ok(request);
|
||||
};
|
||||
let budget = enabled
|
||||
.budget_tokens
|
||||
.as_ref()
|
||||
.and_then(Recognized::known)
|
||||
.copied()
|
||||
.unwrap_or(0);
|
||||
let level = effort_for_budget(&context.budgets, budget, capabilities);
|
||||
MessagesRequest {
|
||||
let budget = match &enabled.budget_tokens {
|
||||
None => Some(0),
|
||||
Some(budget) => numeric_budget_tokens(budget),
|
||||
};
|
||||
if budget.is_none() && !context.drop_params {
|
||||
let actual = match &enabled.budget_tokens {
|
||||
Some(Recognized::Unrecognized(value)) => python_type_name(value),
|
||||
Some(Recognized::Known(_)) | None => "int",
|
||||
};
|
||||
return Err(Error::InvalidRequest(
|
||||
crate::ErrorDetail::NonNumericThinkingBudget { actual },
|
||||
));
|
||||
}
|
||||
let thinking = ThinkingConfig::Adaptive(AdaptiveThinking {
|
||||
display: kept_display(enabled.display.as_ref(), context),
|
||||
extra: carried_thinking_keys(&enabled.extra, context),
|
||||
});
|
||||
let output_config = match budget {
|
||||
Some(budget) => with_default_effort(
|
||||
request.params.output_config,
|
||||
effort_for_budget(&context.budgets, budget, capabilities),
|
||||
),
|
||||
None => request.params.output_config,
|
||||
};
|
||||
Ok(MessagesRequest {
|
||||
params: MessagesOptionalParams {
|
||||
thinking: Some(Recognized::Known(ThinkingConfig::adaptive(None))),
|
||||
output_config: with_default_effort(request.params.output_config, level),
|
||||
thinking: Some(Recognized::Known(thinking)),
|
||||
output_config,
|
||||
..request.params
|
||||
},
|
||||
..request
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
fn translate_adaptive_effort_for_non_adaptive_model(
|
||||
|
|
@ -337,7 +435,7 @@ pub fn translate_thinking(
|
|||
) -> Result<MessagesRequest, Error> {
|
||||
let request = translate_reasoning_effort(request, context)?;
|
||||
let request = drop_disabled_thinking(request, context);
|
||||
let request = translate_legacy_thinking_for_adaptive_model(request, context);
|
||||
let request = translate_legacy_thinking_for_adaptive_model(request, context)?;
|
||||
let request = translate_adaptive_effort_for_non_adaptive_model(request, context)?;
|
||||
Ok(drop_incompatible_temperature_for_thinking(request, context))
|
||||
}
|
||||
|
|
@ -363,6 +461,7 @@ mod tests {
|
|||
ThinkingContext {
|
||||
capabilities,
|
||||
budgets: ThinkingBudgets::default(),
|
||||
..ThinkingContext::default()
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -1239,10 +1338,102 @@ mod tests {
|
|||
let context = ThinkingContext {
|
||||
capabilities,
|
||||
budgets: overridden_budgets(overrides),
|
||||
..ThinkingContext::default()
|
||||
};
|
||||
assert_eq!(
|
||||
translate_thinking(request(input), &context),
|
||||
Ok(request(expected))
|
||||
);
|
||||
}
|
||||
|
||||
use serde_json::json;
|
||||
|
||||
fn legacy_context(drop_params: bool, keeps_display_updates: bool) -> ThinkingContext {
|
||||
ThinkingContext {
|
||||
drop_params,
|
||||
keeps_display_updates,
|
||||
..context(opus_4_7())
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::float_budget(json!(8192.0), "xhigh")]
|
||||
#[case::float_below_a_tier(json!(1023.9), "low")]
|
||||
#[case::string_budget(json!(" 16_000 "), "xhigh")]
|
||||
#[case::negative_budget(json!(-5), "low")]
|
||||
fn legacy_budgets_read_like_python_numbers(#[case] budget: Value, #[case] effort: &str) {
|
||||
let translated = translate_thinking(
|
||||
request(
|
||||
json!({"max_tokens": 64000, "thinking": {"type": "enabled", "budget_tokens": budget}}),
|
||||
),
|
||||
&legacy_context(false, false),
|
||||
);
|
||||
assert_eq!(
|
||||
translated,
|
||||
Ok(request(json!({
|
||||
"max_tokens": 64000,
|
||||
"thinking": {"type": "adaptive"},
|
||||
"output_config": {"effort": effort},
|
||||
})))
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::text(json!("lots"), "str")]
|
||||
#[case::list(json!([1]), "list")]
|
||||
#[case::null(Value::Null, "NoneType")]
|
||||
#[case::boolean(json!(true), "bool")]
|
||||
fn a_non_numeric_legacy_budget_is_rejected_unless_dropped(
|
||||
#[case] budget: Value,
|
||||
#[case] actual: &'static str,
|
||||
) {
|
||||
let fields =
|
||||
json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": budget}});
|
||||
assert_eq!(
|
||||
translate_thinking(request(fields.clone()), &legacy_context(false, false)),
|
||||
Err(Error::InvalidRequest(
|
||||
crate::ErrorDetail::NonNumericThinkingBudget { actual }
|
||||
))
|
||||
);
|
||||
assert_eq!(
|
||||
translate_thinking(request(fields), &legacy_context(true, false)),
|
||||
Ok(request(
|
||||
json!({"max_tokens": 8192, "thinking": {"type": "adaptive"}})
|
||||
))
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::summarized("summarized", false, Some("summarized"))]
|
||||
#[case::updates_kept("updates", true, Some("updates"))]
|
||||
#[case::updates_dropped("updates", false, None)]
|
||||
#[case::unknown("verbose", true, None)]
|
||||
fn legacy_display_survives_only_where_python_keeps_it(
|
||||
#[case] display: &str,
|
||||
#[case] keeps_display_updates: bool,
|
||||
#[case] kept: Option<&str>,
|
||||
) {
|
||||
let translated = translate_thinking(
|
||||
request(json!({"max_tokens": 4096, "thinking": {"type": "enabled", "budget_tokens": 2048, "display": display}})),
|
||||
&legacy_context(false, keeps_display_updates),
|
||||
)
|
||||
.unwrap();
|
||||
let thinking = serde_json::to_value(translated.params.thinking).unwrap();
|
||||
assert_eq!(thinking.get("display").and_then(Value::as_str), kept);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::kept(false, json!({"type": "adaptive", "block_binding": "strict", "note": "x"}))]
|
||||
#[case::dropped(true, json!({"type": "adaptive", "block_binding": "strict"}))]
|
||||
fn legacy_extra_keys_follow_drop_params(#[case] drop_params: bool, #[case] expected: Value) {
|
||||
let translated = translate_thinking(
|
||||
request(json!({"max_tokens": 4096, "thinking": {"type": "enabled", "budget_tokens": 2048, "block_binding": "strict", "note": "x"}})),
|
||||
&legacy_context(drop_params, false),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
serde_json::to_value(translated.params.thinking).unwrap(),
|
||||
expected
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -56,7 +56,7 @@ impl BaseMessagesConfig for AnthropicMessagesConfig {
|
|||
request: MessagesRequest,
|
||||
context: &MessagesTransformContext,
|
||||
) -> Result<MessagesRequest, Error> {
|
||||
transform_messages_request(request, context)
|
||||
transform_messages_request(request, &context.keeping_display_updates())
|
||||
}
|
||||
|
||||
fn secret_names(&self) -> &'static [&'static str] {
|
||||
|
|
@ -206,6 +206,8 @@ fn feature_betas(request: &MessagesRequest) -> BetaSet {
|
|||
.then_some(AnthropicBeta::PerTurnControl20260701),
|
||||
uses_thinking_display_updates(params)
|
||||
.then_some(AnthropicBeta::ThinkingDisplayUpdates20260818),
|
||||
uses_thinking_block_binding(params)
|
||||
.then_some(AnthropicBeta::ThinkingBindingControls20260801),
|
||||
messages_carry_tool_changes(&request.messages)
|
||||
.then_some(AnthropicBeta::MidConversationToolChanges20260701),
|
||||
has_advisor_tool(tools).then_some(AnthropicBeta::AdvisorTool20260301),
|
||||
|
|
@ -267,6 +269,25 @@ fn uses_thinking_display_updates(params: &MessagesOptionalParams) -> bool {
|
|||
display.and_then(Recognized::known) == Some(&ThinkingDisplay::Updates)
|
||||
}
|
||||
|
||||
/// `AnthropicModelInfo.thinking_betas`: any thinking object carrying `block_binding`.
|
||||
fn uses_thinking_block_binding(params: &MessagesOptionalParams) -> bool {
|
||||
match params.thinking.as_ref() {
|
||||
Some(Recognized::Known(ThinkingConfig::Enabled(thinking))) => {
|
||||
thinking.extra.contains_key("block_binding")
|
||||
}
|
||||
Some(Recognized::Known(ThinkingConfig::Adaptive(thinking))) => {
|
||||
thinking.extra.contains_key("block_binding")
|
||||
}
|
||||
Some(Recognized::Known(ThinkingConfig::Disabled(thinking))) => {
|
||||
thinking.extra.contains_key("block_binding")
|
||||
}
|
||||
Some(Recognized::Unrecognized(Value::Object(thinking))) => {
|
||||
thinking.contains_key("block_binding")
|
||||
}
|
||||
Some(Recognized::Unrecognized(_)) | None => false,
|
||||
}
|
||||
}
|
||||
|
||||
fn is_tool_reference_change(block: &ContentBlock) -> bool {
|
||||
matches!(
|
||||
&block.block_type,
|
||||
|
|
@ -297,7 +318,8 @@ fn unsupported_param(model: &str, param: &str, value: &str, hint: &str) -> Error
|
|||
})
|
||||
}
|
||||
|
||||
fn drop_unsupported_params(
|
||||
/// The speed and sampling gates of Python's `get_requested_anthropic_messages_optional_param`.
|
||||
pub fn drop_unsupported_params(
|
||||
request: MessagesRequest,
|
||||
context: &MessagesTransformContext,
|
||||
) -> Result<MessagesRequest, Error> {
|
||||
|
|
@ -838,6 +860,8 @@ mod tests {
|
|||
low: 2000,
|
||||
..ThinkingBudgets::default()
|
||||
},
|
||||
drop_params: true,
|
||||
keeps_display_updates: false,
|
||||
},
|
||||
drop_params: true,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -68,7 +68,7 @@ impl BaseMessagesConfig for AzureAnthropicMessagesConfig {
|
|||
.capabilities
|
||||
.supports_mid_conversation_system,
|
||||
)),
|
||||
context,
|
||||
&context.keeping_display_updates(),
|
||||
)
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -110,6 +110,11 @@ impl ThinkingBudgets {
|
|||
pub struct ThinkingContext {
|
||||
pub capabilities: MessagesModelCapabilities,
|
||||
pub budgets: ThinkingBudgets,
|
||||
/// `drop_params` as Python's thinking translation reads it.
|
||||
pub drop_params: bool,
|
||||
/// The provider accepts `display: "updates"` on thinking translated from the legacy shape
|
||||
/// (Python's `LEGACY_THINKING_DISPLAY_UPDATES_PROVIDERS`).
|
||||
pub keeps_display_updates: bool,
|
||||
}
|
||||
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
|
||||
pub struct MessagesTransformContext {
|
||||
|
|
@ -130,8 +135,22 @@ impl MessagesTransformContext {
|
|||
thinking: ThinkingContext {
|
||||
capabilities,
|
||||
budgets: ThinkingBudgets::from_lookup(env),
|
||||
drop_params,
|
||||
keeps_display_updates: false,
|
||||
},
|
||||
drop_params,
|
||||
}
|
||||
}
|
||||
|
||||
/// The same context for a provider in Python's `LEGACY_THINKING_DISPLAY_UPDATES_PROVIDERS`.
|
||||
#[must_use]
|
||||
pub fn keeping_display_updates(&self) -> Self {
|
||||
Self {
|
||||
thinking: ThinkingContext {
|
||||
keeps_display_updates: true,
|
||||
..self.thinking
|
||||
},
|
||||
..*self
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -92,6 +92,24 @@ impl Error {
|
|||
}
|
||||
}
|
||||
|
||||
pub fn upstream_headers(&self) -> &[(String, String)] {
|
||||
match self {
|
||||
Self::Http(error) => error.upstream_headers(),
|
||||
_ => &[],
|
||||
}
|
||||
}
|
||||
|
||||
pub fn is_body_read_timeout(&self) -> bool {
|
||||
matches!(self, Self::Http(error) if error.is_body_read_timeout())
|
||||
}
|
||||
|
||||
pub fn is_timeout(&self) -> Option<bool> {
|
||||
match self {
|
||||
Self::Http(error) if error.is_timeout() => Some(error.is_connect_timeout()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn upstream_response(&self) -> Option<(u16, &str)> {
|
||||
match self {
|
||||
Self::Http(error) => error.upstream_response(),
|
||||
|
|
@ -158,6 +176,10 @@ pub enum ErrorDetail {
|
|||
value: String,
|
||||
hint: String,
|
||||
},
|
||||
#[error(
|
||||
"thinking.budget_tokens must be a number, got {actual}. To drop it, set `litellm.drop_params=True`, or `drop_params: true` under `litellm_settings` on the proxy"
|
||||
)]
|
||||
NonNumericThinkingBudget { actual: &'static str },
|
||||
#[error(
|
||||
"Missing API Base - A call is being made to LLM Provider but no api base is set either in the environment variables ({{LLM_PROVIDER}}_API_KEY) or via params"
|
||||
)]
|
||||
|
|
|
|||
|
|
@ -77,7 +77,7 @@ impl BaseMessagesConfig for VertexAiPartnerModelsAnthropicMessagesConfig {
|
|||
request: MessagesRequest,
|
||||
context: &MessagesTransformContext,
|
||||
) -> Result<MessagesRequest, Error> {
|
||||
let request = transform_messages_request(request, context)?;
|
||||
let request = transform_messages_request(request, &context.keeping_display_updates())?;
|
||||
let request = strip_cache_control_scope(normalize_system_role_messages(
|
||||
request,
|
||||
context
|
||||
|
|
|
|||
|
|
@ -104,6 +104,7 @@ impl From<&EntryWire> for Adjustments {
|
|||
(Uplift::NonGlobalEndpoint, wire.regional_endpoint_uplift),
|
||||
(Uplift::Residency(Residency::Eu), wire.residency_uplift_eu),
|
||||
(Uplift::Residency(Residency::Us), wire.residency_uplift_us),
|
||||
(Uplift::FedrampEndpoint, wire.fedramp_uplift),
|
||||
]
|
||||
.into_iter()
|
||||
.filter_map(|(uplift, factor)| Some((uplift, factor?))),
|
||||
|
|
@ -746,7 +747,8 @@ mod tests {
|
|||
fn uplifts_are_adjustments_not_prices() {
|
||||
let wire = wire(json!({
|
||||
"regional_endpoint_uplift_multiplier": 1.1,
|
||||
"regional_processing_uplift_multiplier_eu": "1.2"
|
||||
"regional_processing_uplift_multiplier_eu": "1.2",
|
||||
"fedramp_uplift_multiplier": 1.1
|
||||
}));
|
||||
let adjustments = Adjustments::from(&wire);
|
||||
|
||||
|
|
@ -764,6 +766,12 @@ mod tests {
|
|||
Some(1.2)
|
||||
);
|
||||
assert_eq!(adjustments.get(&Uplift::Residency(Residency::Us)), None);
|
||||
assert_eq!(
|
||||
adjustments
|
||||
.get(&Uplift::FedrampEndpoint)
|
||||
.map(Multiplier::get),
|
||||
Some(1.1)
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
|
|
|
|||
|
|
@ -96,6 +96,7 @@ pub struct ModelInfoFields {
|
|||
pub output_cost_per_token_ultrafast: Option<f64>,
|
||||
pub regional_processing_uplift_multiplier_eu: Option<f64>,
|
||||
pub regional_processing_uplift_multiplier_us: Option<f64>,
|
||||
pub fedramp_uplift_multiplier: Option<f64>,
|
||||
pub regional_endpoint_uplift_multiplier: Option<f64>,
|
||||
pub output_cost_per_audio_token: Option<f64>,
|
||||
pub output_cost_per_character: Option<f64>,
|
||||
|
|
@ -186,6 +187,9 @@ pub struct ModelInfoFields {
|
|||
pub supports_image_size: Option<bool>,
|
||||
pub supported_audio_formats: Option<Value>,
|
||||
pub token_count_framing: Option<i64>,
|
||||
pub token_count_tool_framing: Option<i64>,
|
||||
pub token_count_responses_framing: Option<i64>,
|
||||
pub token_count_responses_tool_framing: Option<i64>,
|
||||
pub token_count_model: Option<Value>,
|
||||
pub token_count_tokenizer: Option<Value>,
|
||||
pub vertex_ai_audio_api: Option<Value>,
|
||||
|
|
|
|||
|
|
@ -29,6 +29,8 @@ struct PriceBlocks {
|
|||
regional_processing_uplift_multiplier_eu: Option<Multiplier>,
|
||||
/// Multiplier applied to all token costs for US data residency (e.g. 1.10 = +10%).
|
||||
regional_processing_uplift_multiplier_us: Option<Multiplier>,
|
||||
/// Multiplier applied to all token costs when served from the OpenAI FedRAMP endpoint (e.g. 1.10 = +10%).
|
||||
fedramp_uplift_multiplier: Option<Multiplier>,
|
||||
/// USD cost per web search query, keyed by search context size.
|
||||
search_context_cost_per_query: Option<SearchContextCostPerQuery>,
|
||||
/// Context-length or result-count tiered rates; each tier's costs apply within its range.
|
||||
|
|
|
|||
|
|
@ -93,6 +93,12 @@ pub enum Endpoint {
|
|||
Realtime,
|
||||
#[serde(rename = "/v1/realtime/transcription_sessions")]
|
||||
RealtimeTranscriptionSessions,
|
||||
#[serde(rename = "/v1/realtime/translations")]
|
||||
RealtimeTranslations,
|
||||
#[serde(rename = "/v1/realtime/translations/calls")]
|
||||
RealtimeTranslationCalls,
|
||||
#[serde(rename = "/v1/realtime/translations/client_secrets")]
|
||||
RealtimeTranslationClientSecrets,
|
||||
#[serde(rename = "/v1/videos")]
|
||||
Videos,
|
||||
#[serde(rename = "/v1/ocr")]
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@ use crate::error::{DateCause, Error, Repr};
|
|||
use crate::vocabulary::{
|
||||
AudioFormat, EffortLevel, Endpoint, InputModality, Mode, OutputModality, VertexAiAudioApi,
|
||||
};
|
||||
use crate::wire::lax::Count;
|
||||
use crate::wire::lax::{Count, NumericCount};
|
||||
|
||||
macro_rules! attributes {
|
||||
($($(#[$doc:meta])* $name:ident: $ty:ty $(as $via:ty)?),* $(,)?) => {
|
||||
|
|
@ -93,6 +93,10 @@ attributes! {
|
|||
search_free_requests: SearchFreeRequests,
|
||||
/// URL of the provider pricing/model page this entry was taken from.
|
||||
source: String,
|
||||
/// Model the provider recommends migrating to before `deprecation_date`, as a catalog key when one exists.
|
||||
successor_model: String,
|
||||
/// URL of the provider page or feed `deprecation_date` was read from.
|
||||
deprecation_source: String,
|
||||
/// Audio container formats the model can return.
|
||||
supported_audio_formats: Vec<AudioFormat>,
|
||||
/// OpenAI-style API routes this model can be called through, e.g. /v1/chat/completions.
|
||||
|
|
@ -104,9 +108,14 @@ attributes! {
|
|||
/// Cloud regions the model is available in ('global' or region ids).
|
||||
supported_regions: Vec<String>,
|
||||
thinking_always_on: bool,
|
||||
token_count_framing: u64 as Count,
|
||||
/// Whether the model thinks when the request carries no thinking param.
|
||||
thinking_on_by_default: bool,
|
||||
token_count_framing: u64 as NumericCount,
|
||||
token_count_model: TokenCountModel,
|
||||
token_count_responses_framing: u64 as NumericCount,
|
||||
token_count_responses_tool_framing: u64 as NumericCount,
|
||||
token_count_tokenizer: TokenCountTokenizer,
|
||||
token_count_tool_framing: u64 as NumericCount,
|
||||
/// Provider default tokens-per-minute limit.
|
||||
tpm: u64 as Count,
|
||||
use_openai_responses_path: bool,
|
||||
|
|
@ -298,10 +307,28 @@ mod tests {
|
|||
#[case::known_mode("mode", json!("chat"), true)]
|
||||
#[case::unknown_mode("mode", json!("chatt"), false)]
|
||||
#[case::negative_dbu("input_dbu_cost_per_token", json!(-1), false)]
|
||||
#[case::successor_model("successor_model", json!("gpt-6-luna"), true)]
|
||||
#[case::successor_model_not_a_name("successor_model", json!(["gpt-6-luna"]), false)]
|
||||
#[case::deprecation_source("deprecation_source", json!("https://developers.openai.com/api/docs/deprecations"), true)]
|
||||
#[case::deprecation_source_not_a_url_string("deprecation_source", json!(7), false)]
|
||||
#[case::thinking_on_by_default("thinking_on_by_default", json!(true), true)]
|
||||
#[case::thinking_on_by_default_not_a_flag("thinking_on_by_default", json!("yes"), false)]
|
||||
#[case::reasoning_contexts("reasoning_contexts", json!(["reasoning", "non_reasoning"]), true)]
|
||||
#[case::reasoning_modes("reasoning_modes", json!(["none", "adaptive"]), true)]
|
||||
#[case::reasoning_context_not_a_list("reasoning_contexts", json!("reasoning"), false)]
|
||||
#[case::reasoning_mode_not_a_string("reasoning_modes", json!([1]), false)]
|
||||
#[case::framing("token_count_framing", json!(12), true)]
|
||||
#[case::framing_written_as_text("token_count_framing", json!("12"), false)]
|
||||
#[case::framing_written_as_whole_float("token_count_framing", json!(12.0), true)]
|
||||
#[case::fractional_tool_framing("token_count_tool_framing", json!(1.5), false)]
|
||||
#[case::tool_framing("token_count_tool_framing", json!(214), true)]
|
||||
#[case::tool_framing_written_as_text("token_count_tool_framing", json!("214"), false)]
|
||||
#[case::tool_framing_as_a_flag("token_count_tool_framing", json!(true), false)]
|
||||
#[case::negative_tool_framing("token_count_tool_framing", json!(-1), false)]
|
||||
#[case::responses_framing("token_count_responses_framing", json!(38), true)]
|
||||
#[case::responses_framing_as_a_flag("token_count_responses_framing", json!(false), false)]
|
||||
#[case::responses_tool_framing("token_count_responses_tool_framing", json!(267), true)]
|
||||
#[case::responses_tool_framing_written_as_text("token_count_responses_tool_framing", json!("267"), false)]
|
||||
fn each_attribute_reads_its_own_type(
|
||||
#[case] key: &str,
|
||||
#[case] value: Value,
|
||||
|
|
|
|||
|
|
@ -24,6 +24,7 @@ pub struct EntryWire {
|
|||
pub regional_endpoint_uplift: Option<Multiplier>,
|
||||
pub residency_uplift_eu: Option<Multiplier>,
|
||||
pub residency_uplift_us: Option<Multiplier>,
|
||||
pub fedramp_uplift: Option<Multiplier>,
|
||||
pub unknown: Fields,
|
||||
pub problems: Vec<Error>,
|
||||
}
|
||||
|
|
@ -57,6 +58,7 @@ impl EntryWire {
|
|||
"regional_processing_uplift_multiplier_us" => {
|
||||
self.residency_uplift_us = reader.value(path(), value)
|
||||
}
|
||||
"fedramp_uplift_multiplier" => self.fedramp_uplift = reader.value(path(), value),
|
||||
_ => self.read_open_key(reader, key, value),
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -8,6 +8,8 @@ pub(crate) struct LaxI64;
|
|||
|
||||
pub(crate) struct Count;
|
||||
|
||||
pub(crate) struct NumericCount;
|
||||
|
||||
impl<'de> DeserializeAs<'de, i64> for LaxI64 {
|
||||
fn deserialize_as<D: Deserializer<'de>>(deserializer: D) -> Result<i64, D::Error> {
|
||||
deserializer.deserialize_any(Self)
|
||||
|
|
@ -52,6 +54,34 @@ impl<'de> DeserializeAs<'de, u64> for Count {
|
|||
}
|
||||
}
|
||||
|
||||
impl<'de> DeserializeAs<'de, u64> for NumericCount {
|
||||
fn deserialize_as<D: Deserializer<'de>>(deserializer: D) -> Result<u64, D::Error> {
|
||||
deserializer.deserialize_any(Self)
|
||||
}
|
||||
}
|
||||
|
||||
impl Visitor<'_> for NumericCount {
|
||||
type Value = u64;
|
||||
|
||||
fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
formatter.write_str("a whole number of zero or more")
|
||||
}
|
||||
|
||||
fn visit_u64<E: de::Error>(self, value: u64) -> Result<u64, E> {
|
||||
Ok(value)
|
||||
}
|
||||
|
||||
fn visit_i64<E: de::Error>(self, value: i64) -> Result<u64, E> {
|
||||
u64::try_from(value).map_err(|_| E::invalid_value(de::Unexpected::Signed(value), &self))
|
||||
}
|
||||
|
||||
fn visit_f64<E: de::Error>(self, value: f64) -> Result<u64, E> {
|
||||
integral_float(value)
|
||||
.and_then(|whole| u64::try_from(whole).ok())
|
||||
.ok_or_else(|| E::invalid_value(de::Unexpected::Float(value), &self))
|
||||
}
|
||||
}
|
||||
|
||||
fn integer_string(value: &str) -> Option<i64> {
|
||||
let integer = match value.split_once('.') {
|
||||
Some((integer, fraction)) => {
|
||||
|
|
@ -99,6 +129,12 @@ mod tests {
|
|||
.map(DeserializeAsWrap::into_inner)
|
||||
}
|
||||
|
||||
fn numeric_count(value: Value) -> Option<u64> {
|
||||
DeserializeAsWrap::<u64, NumericCount>::deserialize(value)
|
||||
.ok()
|
||||
.map(DeserializeAsWrap::into_inner)
|
||||
}
|
||||
|
||||
use serde::Deserialize;
|
||||
|
||||
#[rstest]
|
||||
|
|
@ -125,4 +161,23 @@ mod tests {
|
|||
) {
|
||||
assert_eq!(count(written), expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::integer(json!(214), Some(214))]
|
||||
#[case::zero(json!(0), Some(0))]
|
||||
#[case::whole_float(json!(214.0), Some(214))]
|
||||
#[case::negative_zero(json!(-0.0), Some(0))]
|
||||
#[case::above_i64(json!(u64::MAX), Some(u64::MAX))]
|
||||
#[case::negative(json!(-1), None)]
|
||||
#[case::negative_whole_float(json!(-1.0), None)]
|
||||
#[case::fraction(json!(1.5), None)]
|
||||
#[case::string(json!("214"), None)]
|
||||
#[case::boolean(json!(true), None)]
|
||||
#[case::list(json!([214]), None)]
|
||||
fn numeric_counts_accept_only_whole_json_numbers(
|
||||
#[case] written: Value,
|
||||
#[case] expected: Option<u64>,
|
||||
) {
|
||||
assert_eq!(numeric_count(written), expected);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -29,8 +29,6 @@ bytes.workspace = true
|
|||
futures-util.workspace = true
|
||||
litellm-cache.workspace = true
|
||||
litellm-error.workspace = true
|
||||
litellm-cache-memory.workspace = true
|
||||
litellm-cache-redis.workspace = true
|
||||
litellm-cache-response.workspace = true
|
||||
serde.workspace = true
|
||||
litellm-auth.workspace = true
|
||||
|
|
@ -56,7 +54,6 @@ pyo3.workspace = true
|
|||
prost.workspace = true
|
||||
pyo3-async-runtimes.workspace = true
|
||||
reqwest.workspace = true
|
||||
redis = { workspace = true, features = ["tls-rustls"] }
|
||||
serde_json.workspace = true
|
||||
strum.workspace = true
|
||||
veil.workspace = true
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
# Cache boundary
|
||||
|
||||
This folder owns how Rust inference reaches the selected cache: global cache selection, route admission, delegation to a Python cache and the experimental V2 native handles. Cache algorithms, storage protocols and response-cache semantics belong to their cache crates
|
||||
This folder owns how Rust inference reaches the selected Python cache. Cache algorithms, storage protocols and response-cache semantics belong to their cache crates
|
||||
|
||||
`mod.rs` exposes the cache boundary to routes and module registration. Adapter directories remain private. `selection.rs` asks `litellm.rust_bridge.host.cache` for the live configured cache, and owns route admission and inference protocol composition for both adapters
|
||||
`mod.rs` exposes the cache boundary to routes. `selection.rs` asks `litellm.rust_bridge.host.cache` for the live configured cache and owns route admission and inference protocol composition
|
||||
|
||||
`python/` delegates operations to the selected Python cache without discovering configuration. `native/` owns the experimental V2 native cache handles. Neither adapter depends on shared selection or the other adapter. Shared composition depends on the adapters, and routes use only the parent module's exports
|
||||
`python/` delegates operations to the selected Python cache without discovering configuration. Routes use only the parent module's exports
|
||||
|
|
|
|||
|
|
@ -1,32 +1,15 @@
|
|||
mod native;
|
||||
mod python;
|
||||
mod selection;
|
||||
|
||||
pub(crate) use native::NativeCacheHandle;
|
||||
pub(crate) use python::{CacheCall, PythonCache};
|
||||
pub(crate) use selection::{Cached, Selection, configure, configured_native};
|
||||
pub(crate) use selection::{Cached, Selection, configure, ensure_cache_disabled};
|
||||
|
||||
use litellm_cache::Error;
|
||||
use pyo3::{
|
||||
exceptions::{PyNotImplementedError, PyRuntimeError, PyValueError},
|
||||
prelude::*,
|
||||
};
|
||||
use pyo3::{exceptions::PyRuntimeError, prelude::*};
|
||||
|
||||
/// A cache call the bridge rejects itself, before or after the cache ran.
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub(crate) enum HandleError {
|
||||
#[error("recreate the v2 cache after fork")]
|
||||
Forked,
|
||||
#[error("cache limits must be positive")]
|
||||
Limits,
|
||||
#[error("namespace and a positive cache limit are required")]
|
||||
RedisLimits,
|
||||
#[error("cache durations must be finite and positive")]
|
||||
Duration(#[source] Option<std::time::TryFromFloatSecsError>),
|
||||
#[error("unsupported v2 cache control: {0}")]
|
||||
Control(String),
|
||||
#[error("native inference needs a v2 cache, the catalog keeps Python caches on Python")]
|
||||
NotNative,
|
||||
#[error("this inference route does not support configured caches")]
|
||||
Unsupported,
|
||||
#[error("cache operation without configured cache")]
|
||||
Unconfigured,
|
||||
#[error("cache arguments unavailable")]
|
||||
|
|
@ -39,66 +22,10 @@ impl From<HandleError> for PyErr {
|
|||
fn from(error: HandleError) -> Self {
|
||||
let message = litellm_error::Report::new(&error).to_string();
|
||||
match error {
|
||||
HandleError::Limits
|
||||
| HandleError::RedisLimits
|
||||
| HandleError::Duration(_)
|
||||
| HandleError::Control(_) => PyValueError::new_err(message),
|
||||
HandleError::Forked
|
||||
| HandleError::NotNative
|
||||
HandleError::Unsupported
|
||||
| HandleError::Unconfigured
|
||||
| HandleError::MissingArguments
|
||||
| HandleError::NothingPending => PyRuntimeError::new_err(message),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn cache_error(error: Error) -> PyErr {
|
||||
let message = litellm_error::Report::new(&error).to_string();
|
||||
match error {
|
||||
Error::InvalidEntry(_) => PyValueError::new_err(message),
|
||||
Error::UnsupportedOperation => PyNotImplementedError::new_err(message),
|
||||
Error::Unavailable(_) | Error::UnscopedFlush | Error::MissingPrompt => {
|
||||
PyRuntimeError::new_err(message)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use pyo3::{
|
||||
exceptions::{PyRuntimeError, PyValueError},
|
||||
prelude::*,
|
||||
};
|
||||
use rstest::rstest;
|
||||
|
||||
use super::{Error, cache_error};
|
||||
|
||||
fn refused() -> std::io::Error {
|
||||
std::io::Error::other("connection refused")
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::unavailable(
|
||||
Error::unavailable(refused()),
|
||||
false,
|
||||
"cache is unavailable: connection refused"
|
||||
)]
|
||||
#[case::invalid_entry(
|
||||
Error::invalid_entry(refused()),
|
||||
true,
|
||||
"invalid cache entry: connection refused"
|
||||
)]
|
||||
fn python_sees_the_class_and_the_whole_chain(
|
||||
#[case] error: Error,
|
||||
#[case] caller_fault: bool,
|
||||
#[case] message: &str,
|
||||
) {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let failure = cache_error(error);
|
||||
assert_eq!(failure.is_instance_of::<PyValueError>(py), caller_fault);
|
||||
assert_eq!(failure.is_instance_of::<PyRuntimeError>(py), !caller_fault);
|
||||
assert_eq!(failure.value(py).to_string(), message);
|
||||
});
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,9 +0,0 @@
|
|||
# Native cache bindings
|
||||
|
||||
This directory exposes the experimental V2 native cache handles (`NativeCacheHandle`) to Python. Cache algorithms and storage protocols remain in their cache crates
|
||||
|
||||
Accept the cache object selected by the parent module. Do not read global `litellm.cache`, decide route admission, or select the Python cache adapter here
|
||||
|
||||
Keep Python-facing cache classes and method signatures stable when reorganizing modules. Native internals stay private to this directory unless the shared cache boundary or Python module registration needs them
|
||||
|
||||
Verify changes with the V2 cache tests using a freshly built extension. Test cache behavior, not module paths or file structure
|
||||
|
|
@ -1,3 +0,0 @@
|
|||
pub(super) mod v2;
|
||||
|
||||
pub(crate) use v2::NativeCacheHandle;
|
||||
|
|
@ -1,336 +0,0 @@
|
|||
use crate::cache::{HandleError, cache_error};
|
||||
use std::{
|
||||
sync::Arc,
|
||||
time::{Duration, SystemTime, UNIX_EPOCH},
|
||||
};
|
||||
|
||||
use litellm_cache::{DeleteCache, DisconnectCache, PingCache};
|
||||
use litellm_host_python::{from_py, release_gil, to_py};
|
||||
use serde_json::Value;
|
||||
|
||||
use litellm_cache_memory::InMemoryCache;
|
||||
use litellm_cache_redis::{RedisCache, RedisTopology};
|
||||
use litellm_cache_response::{
|
||||
CacheEntry, CacheKeyInput, ExactResponseCache, ResponseCache, ResponseCacheCodec,
|
||||
ResponseCacheConfig, ResponseCacheRequest, ResponseCacheService,
|
||||
};
|
||||
use pyo3::{prelude::*, types::PyDict};
|
||||
|
||||
#[pyclass(
|
||||
frozen,
|
||||
name = "NativeCacheHandle",
|
||||
module = "litellm.rust_bridge._native"
|
||||
)]
|
||||
pub(crate) struct NativeCacheHandle {
|
||||
service: Arc<dyn ResponseCacheService>,
|
||||
backend: Arc<dyn ExactResponseCache>,
|
||||
storage: Storage,
|
||||
pid: u32,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
enum Storage {
|
||||
Memory(Arc<InMemoryCache<CacheEntry>>),
|
||||
Redis(Arc<RedisCache<ResponseCacheCodec>>),
|
||||
}
|
||||
|
||||
impl NativeCacheHandle {
|
||||
fn check_process(&self) -> PyResult<()> {
|
||||
if self.pid != std::process::id() {
|
||||
return Err(HandleError::Forked.into());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn request(key: String, ttl: Option<f64>) -> PyResult<ResponseCacheRequest> {
|
||||
let mut request: ResponseCacheRequest = ResponseCacheRequest::new(CacheKeyInput {
|
||||
preset: Some(key),
|
||||
..Default::default()
|
||||
});
|
||||
request.context.ttl = ttl.map(duration).transpose()?;
|
||||
Ok(request)
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl NativeCacheHandle {
|
||||
#[staticmethod]
|
||||
#[pyo3(signature = (*, ttl=600.0, capacity=200, max_entry_bytes=4194304))]
|
||||
fn memory(ttl: f64, capacity: usize, max_entry_bytes: usize) -> PyResult<Self> {
|
||||
let ttl = duration(ttl)?;
|
||||
if capacity == 0 || max_entry_bytes == 0 {
|
||||
return Err(HandleError::Limits.into());
|
||||
}
|
||||
let storage = Arc::new(InMemoryCache::new(Some(capacity), Some(ttl)));
|
||||
let backend = Arc::new(ResponseCache::new(storage.clone()).with_config(
|
||||
ResponseCacheConfig {
|
||||
namespace: "sdk".into(),
|
||||
max_entry_bytes,
|
||||
},
|
||||
));
|
||||
Ok(Self {
|
||||
service: backend.clone(),
|
||||
backend,
|
||||
storage: Storage::Memory(storage),
|
||||
pid: std::process::id(),
|
||||
})
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
#[pyo3(signature = (url, *, namespace, ttl=600.0, max_entry_bytes=4194304))]
|
||||
fn redis(
|
||||
py: Python<'_>,
|
||||
url: &str,
|
||||
namespace: String,
|
||||
ttl: f64,
|
||||
max_entry_bytes: usize,
|
||||
) -> PyResult<Self> {
|
||||
let ttl = duration(ttl)?;
|
||||
if namespace.is_empty() || max_entry_bytes == 0 {
|
||||
return Err(HandleError::RedisLimits.into());
|
||||
}
|
||||
let storage = Arc::new(
|
||||
release_gil(py, || {
|
||||
RedisCache::connect(
|
||||
url,
|
||||
&RedisTopology::Standalone,
|
||||
Some(ttl),
|
||||
ResponseCacheCodec,
|
||||
)
|
||||
})
|
||||
.map_err(cache_error)?
|
||||
.with_namespace(Some(namespace.clone())),
|
||||
);
|
||||
let backend = Arc::new(ResponseCache::new(storage.clone()).with_config(
|
||||
ResponseCacheConfig {
|
||||
namespace,
|
||||
max_entry_bytes,
|
||||
},
|
||||
));
|
||||
Ok(Self {
|
||||
service: backend.clone(),
|
||||
backend,
|
||||
storage: Storage::Redis(storage),
|
||||
pid: std::process::id(),
|
||||
})
|
||||
}
|
||||
fn get(&self, py: Python<'_>, key: String) -> PyResult<Py<PyAny>> {
|
||||
self.check_process()?;
|
||||
let request = request(key, None)?;
|
||||
let value =
|
||||
release_gil(py, || self.backend.lookup(&request, now())).map_err(cache_error)?;
|
||||
to_py(py, &value)
|
||||
}
|
||||
|
||||
#[pyo3(signature = (key, value, *, ttl=None))]
|
||||
fn set(
|
||||
&self,
|
||||
py: Python<'_>,
|
||||
key: String,
|
||||
value: &Bound<'_, PyAny>,
|
||||
ttl: Option<f64>,
|
||||
) -> PyResult<()> {
|
||||
self.check_process()?;
|
||||
let request = request(key, ttl)?;
|
||||
let value: Value = from_py(value)?;
|
||||
release_gil(py, || self.backend.store(&request, value, now())).map_err(cache_error)
|
||||
}
|
||||
|
||||
fn async_get<'py>(&self, py: Python<'py>, key: String) -> PyResult<Bound<'py, PyAny>> {
|
||||
self.check_process()?;
|
||||
let request = request(key, None)?;
|
||||
let backend = self.backend.clone();
|
||||
crate::execution::run_async(
|
||||
py,
|
||||
async move { backend.async_lookup(&request, now()).await },
|
||||
cache_error,
|
||||
)
|
||||
}
|
||||
|
||||
#[pyo3(signature = (key, value, *, ttl=None))]
|
||||
fn async_set<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
key: String,
|
||||
value: &Bound<'_, PyAny>,
|
||||
ttl: Option<f64>,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
self.check_process()?;
|
||||
let request = request(key, ttl)?;
|
||||
let value: Value = from_py(value)?;
|
||||
let backend = self.backend.clone();
|
||||
crate::execution::run_async(
|
||||
py,
|
||||
async move { backend.async_store(&request, value, now()).await },
|
||||
cache_error,
|
||||
)
|
||||
}
|
||||
|
||||
#[pyo3(signature = (entries, *, ttl=None))]
|
||||
fn async_set_many<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
entries: &Bound<'_, PyAny>,
|
||||
ttl: Option<f64>,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
self.check_process()?;
|
||||
let entries: Vec<(String, Value)> = from_py(entries)?;
|
||||
let entries = entries
|
||||
.into_iter()
|
||||
.map(|(key, value)| Ok((request(key, ttl)?, value)))
|
||||
.collect::<PyResult<Vec<_>>>()?;
|
||||
let backend = self.backend.clone();
|
||||
crate::execution::run_async(
|
||||
py,
|
||||
async move { backend.async_store_batch(entries, now()).await },
|
||||
cache_error,
|
||||
)
|
||||
}
|
||||
|
||||
fn flush(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
|
||||
self.check_process()?;
|
||||
let backend = self.backend.clone();
|
||||
crate::execution::run_sync(py, async move { backend.async_flush().await }, cache_error)
|
||||
}
|
||||
|
||||
fn async_flush<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
|
||||
self.check_process()?;
|
||||
let backend = self.backend.clone();
|
||||
crate::execution::run_async(py, async move { backend.async_flush().await }, cache_error)
|
||||
}
|
||||
|
||||
fn ping<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
|
||||
self.check_process()?;
|
||||
let storage = self.storage.clone();
|
||||
crate::execution::run_async(
|
||||
py,
|
||||
async move {
|
||||
match storage {
|
||||
Storage::Memory(_) => Ok(true),
|
||||
Storage::Redis(cache) => cache.ping().await,
|
||||
}
|
||||
},
|
||||
cache_error,
|
||||
)
|
||||
}
|
||||
|
||||
fn disconnect<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
|
||||
self.check_process()?;
|
||||
let storage = self.storage.clone();
|
||||
crate::execution::run_async(
|
||||
py,
|
||||
async move {
|
||||
match storage {
|
||||
Storage::Memory(cache) => cache.disconnect().await,
|
||||
Storage::Redis(cache) => cache.disconnect().await,
|
||||
}
|
||||
},
|
||||
cache_error,
|
||||
)
|
||||
}
|
||||
|
||||
fn delete<'py>(&self, py: Python<'py>, keys: Vec<String>) -> PyResult<Bound<'py, PyAny>> {
|
||||
self.check_process()?;
|
||||
let storage = self.storage.clone();
|
||||
crate::execution::run_async(
|
||||
py,
|
||||
async move {
|
||||
for key in keys {
|
||||
match &storage {
|
||||
Storage::Memory(cache) => cache.async_delete_cache(&key).await?,
|
||||
Storage::Redis(cache) => cache.async_delete_cache(&key).await?,
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
},
|
||||
cache_error,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
fn duration(seconds: f64) -> Result<Duration, HandleError> {
|
||||
match Duration::try_from_secs_f64(seconds) {
|
||||
Ok(duration) if !duration.is_zero() => Ok(duration),
|
||||
Ok(_) => Err(HandleError::Duration(None)),
|
||||
Err(error) => Err(HandleError::Duration(Some(error))),
|
||||
}
|
||||
}
|
||||
|
||||
pub(in crate::cache) fn native_handle<'py>(
|
||||
configured: &Bound<'py, PyAny>,
|
||||
) -> PyResult<Option<Bound<'py, PyAny>>> {
|
||||
Ok(configured
|
||||
.getattr_opt("cache")?
|
||||
.map(|backend| backend.getattr_opt("native_handle"))
|
||||
.transpose()?
|
||||
.flatten()
|
||||
.filter(|handle| handle.is_instance_of::<NativeCacheHandle>()))
|
||||
}
|
||||
|
||||
pub(in crate::cache) fn configured(
|
||||
configured: &Bound<'_, PyAny>,
|
||||
kwargs: &Bound<'_, PyDict>,
|
||||
) -> PyResult<(
|
||||
Option<Arc<dyn ResponseCacheService>>,
|
||||
litellm_cache_response::CacheOptions,
|
||||
)> {
|
||||
let handle = native_handle(configured)?.ok_or(HandleError::NotNative)?;
|
||||
let cache = handle.extract::<PyRef<'_, NativeCacheHandle>>()?;
|
||||
cache.check_process()?;
|
||||
let controls = kwargs.get_item("cache")?.filter(|value| !value.is_none());
|
||||
let controls = controls
|
||||
.as_ref()
|
||||
.map(|value| value.cast::<PyDict>())
|
||||
.transpose()?;
|
||||
if let Some(controls) = controls {
|
||||
for name in controls.keys() {
|
||||
let name = name.extract::<String>()?;
|
||||
if !matches!(
|
||||
name.as_str(),
|
||||
"no-cache" | "no-store" | "ttl" | "s-maxage" | "s-max-age" | "use-cache"
|
||||
) {
|
||||
return Err(HandleError::Control(name).into());
|
||||
}
|
||||
}
|
||||
}
|
||||
let boolean = |name: &str| -> PyResult<bool> {
|
||||
controls
|
||||
.map(|values| values.get_item(name))
|
||||
.transpose()?
|
||||
.flatten()
|
||||
.map(|value| value.extract())
|
||||
.transpose()
|
||||
.map(|value| value.unwrap_or(false))
|
||||
};
|
||||
let seconds = |name: &str| -> PyResult<Option<Duration>> {
|
||||
controls
|
||||
.map(|values| values.get_item(name))
|
||||
.transpose()?
|
||||
.flatten()
|
||||
.map(|value| Ok(duration(value.extract()?)?))
|
||||
.transpose()
|
||||
};
|
||||
Ok((
|
||||
Some(cache.service.clone()),
|
||||
litellm_cache_response::CacheOptions {
|
||||
policy: litellm_cache_response::CachePolicy {
|
||||
caching: kwargs
|
||||
.get_item("caching")?
|
||||
.filter(|value| !value.is_none())
|
||||
.map(|value| value.extract())
|
||||
.transpose()?,
|
||||
no_cache: boolean("no-cache")?,
|
||||
no_store: boolean("no-store")?,
|
||||
ttl: seconds("ttl")?,
|
||||
max_age: seconds("s-max-age")?.or(seconds("s-maxage")?),
|
||||
},
|
||||
scope: litellm_cache_response::CacheScope::Shared,
|
||||
},
|
||||
))
|
||||
}
|
||||
|
||||
fn now() -> Duration {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue