mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-26 01:12:21 +00:00
Merge branch 'main' into litellm_durable_background_interaction_settlement
Some checks failed
LiteLLM Rust / rust-lint (push) Has been cancelled
LiteLLM Rust / rust-test (push) Has been cancelled
LiteLLM Rust / rust-wheel (push) Has been cancelled
Terraform Provider / gofmt, vet, build, test (push) Has been cancelled
Terraform Provider / Provider endpoints vs proxy OpenAPI schema (push) Has been cancelled
Some checks failed
LiteLLM Rust / rust-lint (push) Has been cancelled
LiteLLM Rust / rust-test (push) Has been cancelled
LiteLLM Rust / rust-wheel (push) Has been cancelled
Terraform Provider / gofmt, vet, build, test (push) Has been cancelled
Terraform Provider / Provider endpoints vs proxy OpenAPI schema (push) Has been cancelled
This commit is contained in:
commit
5573b3ed35
491 changed files with 31411 additions and 4446 deletions
10
.github/actions/cache-cargo-build/action.yml
vendored
10
.github/actions/cache-cargo-build/action.yml
vendored
|
|
@ -15,6 +15,12 @@ description: >-
|
|||
cache the same directory for different workloads, and a shared key would let
|
||||
whichever ran first deny the others a save.
|
||||
|
||||
inputs:
|
||||
profile:
|
||||
description: "Cargo profile the build uses (dev or release)"
|
||||
required: false
|
||||
default: "dev"
|
||||
|
||||
runs:
|
||||
using: composite
|
||||
steps:
|
||||
|
|
@ -25,6 +31,6 @@ runs:
|
|||
~/.cargo/registry
|
||||
~/.cargo/git
|
||||
litellm-rust/target
|
||||
key: ${{ runner.os }}-maturin-dev-${{ hashFiles('litellm-rust/Cargo.lock') }}
|
||||
key: ${{ runner.os }}-maturin-${{ inputs.profile }}-${{ hashFiles('litellm-rust/Cargo.lock') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-maturin-dev-
|
||||
${{ runner.os }}-maturin-${{ inputs.profile }}-
|
||||
|
|
|
|||
1
.github/e2e-stack/select_tests.py
vendored
1
.github/e2e-stack/select_tests.py
vendored
|
|
@ -9,6 +9,7 @@ UNSUPPORTED: Final = re.compile(
|
|||
r"|^tests/e2e/llm_translation/realtime/test_realtime_pipecat_audio_e2e\.py$"
|
||||
r"|^tests/e2e/batches/test_managed_files_enforcement_e2e\.py$"
|
||||
r"|^tests/e2e/guardrails/test_presidio_masking_e2e\.py$"
|
||||
r"|^tests/e2e/logging/test_otel_v2_langfuse_generation_output_e2e\.py$"
|
||||
)
|
||||
HARNESS: Final = re.compile(
|
||||
r"^tests/e2e/[A-Za-z0-9_.-]+\.(py|ini)$"
|
||||
|
|
|
|||
14
.github/scripts/verify_linux_native_wheel.py
vendored
14
.github/scripts/verify_linux_native_wheel.py
vendored
|
|
@ -134,7 +134,16 @@ def main(
|
|||
uncompressed_wheel_size: Final = sum(member.file_size for member in wheel_members)
|
||||
native_path: Final = wheel.parent / "native" / Path(native_member.filename).name
|
||||
native_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
native_path.write_bytes(archive.read(native_member))
|
||||
native_bytes: Final = archive.read(native_member)
|
||||
native_path.write_bytes(native_bytes)
|
||||
duplicated_vocabularies: Final = tuple(
|
||||
member.filename
|
||||
for member in wheel_members
|
||||
if member.filename.startswith("litellm/litellm_core_utils/tokenizers/")
|
||||
and re.fullmatch(r"[0-9a-f]{40}", PurePosixPath(member.filename).name)
|
||||
and member.file_size > 0
|
||||
and archive.read(member) in native_bytes
|
||||
)
|
||||
|
||||
wheel_metadata_tags_match: Final = (
|
||||
len(wheel_metadata_tags) == len(expanded_filename_tags)
|
||||
|
|
@ -205,7 +214,7 @@ def main(
|
|||
native_module: Final = load_native_module(native_path)
|
||||
native_module_loads: Final = native_module is not None
|
||||
panic_test_hook_absent: Final = native_module is not None and not hasattr(native_module, "_panic_for_test")
|
||||
native_size_limit: Final = 40_000_000
|
||||
native_size_limit: Final = 35_000_000
|
||||
native_size_within_limit: Final = native_member.file_size <= native_size_limit
|
||||
validations: Final = (
|
||||
(f"Python tag is {EXPECTED_PYTHON_TAG}", python_tag == EXPECTED_PYTHON_TAG),
|
||||
|
|
@ -223,6 +232,7 @@ def main(
|
|||
("Native module loads", native_module_loads),
|
||||
("Production module omits the panic test hook", panic_test_hook_absent),
|
||||
(f"Native extension does not exceed {native_size_limit / 1_000_000:.0f} MB", native_size_within_limit),
|
||||
("Tokenizer vocabularies are not duplicated in the native extension", not duplicated_vocabularies),
|
||||
("Wheel contents are valid", not unexpected_members),
|
||||
)
|
||||
|
||||
|
|
|
|||
36
.github/workflows/codspeed.yml
vendored
36
.github/workflows/codspeed.yml
vendored
|
|
@ -13,6 +13,7 @@ on:
|
|||
- ".github/workflows/codspeed.yml"
|
||||
- ".github/actions/setup-uv-with-retries/**"
|
||||
- ".github/actions/cache-cargo-build/**"
|
||||
- ".github/scripts/uv_sync_with_retries.sh"
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
|
|
@ -25,6 +26,7 @@ on:
|
|||
- ".github/workflows/codspeed.yml"
|
||||
- ".github/actions/setup-uv-with-retries/**"
|
||||
- ".github/actions/cache-cargo-build/**"
|
||||
- ".github/scripts/uv_sync_with_retries.sh"
|
||||
# Allow CodSpeed to trigger backtest performance analysis
|
||||
# in order to generate initial data
|
||||
workflow_dispatch:
|
||||
|
|
@ -59,19 +61,27 @@ jobs:
|
|||
|
||||
- name: Cache the Rust build
|
||||
uses: ./.github/actions/cache-cargo-build
|
||||
with:
|
||||
profile: release
|
||||
|
||||
# Build the wheel and resolve every dependency outside the CodSpeed
|
||||
# runner: the same maturin build took 42 minutes inside `codspeed run`
|
||||
# versus under 3 minutes as a plain step (LIT-6183)
|
||||
- name: Build environment
|
||||
- name: Build the release wheel
|
||||
run: uv build --wheel --out-dir dist
|
||||
|
||||
- name: Install the wheel into the benchmark environment
|
||||
run: |
|
||||
UV_PROJECT_ENVIRONMENT="${RUNNER_TEMP}/benchmark-venv" .github/scripts/uv_sync_with_retries.sh --frozen --no-default-groups --group benchmarks --no-install-project --python 3.12
|
||||
uv pip install --python "${RUNNER_TEMP}/benchmark-venv/bin/python" --no-deps dist/*.whl
|
||||
|
||||
- name: Collect benchmarks
|
||||
env:
|
||||
PYTEST_DISABLE_PLUGIN_AUTOLOAD: "1"
|
||||
LITELLM_REQUIRE_INSTALLED_WHEEL: "1"
|
||||
run: >
|
||||
env PYTEST_DISABLE_PLUGIN_AUTOLOAD=1
|
||||
uv run --frozen --no-default-groups
|
||||
--with pytest==8.3.5
|
||||
--with pytest-codspeed==4.3.0
|
||||
--with "mcp>=2.2.0,<3.0"
|
||||
--with "a2a-sdk>=1.1.0,<2.0"
|
||||
pytest
|
||||
"${RUNNER_TEMP}/benchmark-venv/bin/python" -I -m pytest
|
||||
--import-mode=importlib
|
||||
-p pytest_codspeed.plugin
|
||||
tests/benchmarks/
|
||||
--codspeed
|
||||
|
|
@ -82,13 +92,9 @@ jobs:
|
|||
with:
|
||||
mode: simulation
|
||||
run: >
|
||||
env PYTEST_DISABLE_PLUGIN_AUTOLOAD=1
|
||||
uv run --frozen --no-default-groups
|
||||
--with pytest==8.3.5
|
||||
--with pytest-codspeed==4.3.0
|
||||
--with "mcp>=2.2.0,<3.0"
|
||||
--with "a2a-sdk>=1.1.0,<2.0"
|
||||
pytest
|
||||
env PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 LITELLM_REQUIRE_INSTALLED_WHEEL=1
|
||||
"${RUNNER_TEMP}/benchmark-venv/bin/python" -I -m pytest
|
||||
--import-mode=importlib
|
||||
-p pytest_codspeed.plugin
|
||||
tests/benchmarks/
|
||||
--codspeed
|
||||
|
|
|
|||
|
|
@ -61,6 +61,8 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
|||
--extra saml \
|
||||
--python python3.13
|
||||
|
||||
RUN cp "$(python -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')"/litellm/rust_bridge/_native*.so litellm/rust_bridge/
|
||||
|
||||
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
|
||||
npm_config_cache=/root/.npm \
|
||||
prisma generate --schema=./schema.prisma
|
||||
|
|
|
|||
|
|
@ -1,8 +1,11 @@
|
|||
"""Guard the cost map on pull requests.
|
||||
|
||||
Every pull request gets the file checks: the three cost map files parse, the backup copy matches the root file,
|
||||
and the JSON schema is in sync and validates the map. Pull requests from the cost map sync bot (branches named
|
||||
litellm_cost_map_sync_*) additionally may only touch those three files and may only add or update models.
|
||||
Every pull request whose diff against its merge base touches one of the three cost map files gets the file
|
||||
checks: the files parse, the backup copy matches the root file, and the JSON schema is in sync and validates the
|
||||
map. A pull request that leaves all three untouched skips them, since merging it keeps the base branch's copies
|
||||
and its head tree only carries whatever state the branch was cut from. Pull requests from the cost map sync bot
|
||||
(branches named litellm_cost_map_sync_*) always get the file checks and additionally may only touch those three
|
||||
files and may only add or update models.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
|
@ -108,20 +111,37 @@ def _bot_failures(base: Snapshot, head_map: CostMap, changed_files: Sequence[str
|
|||
)
|
||||
|
||||
|
||||
def touches_cost_map(changed_files: Sequence[str]) -> bool:
|
||||
return any(path in GUARDED_PATHS for path in changed_files)
|
||||
|
||||
|
||||
def contract_for(bot: bool, changed_files: Sequence[str]) -> str:
|
||||
if bot:
|
||||
return "bot contract enforced"
|
||||
return "human PR, file checks only" if touches_cost_map(changed_files) else "human PR, cost map untouched"
|
||||
|
||||
|
||||
def guard_failures(base: Snapshot, head: Snapshot, changed_files: Sequence[str], bot: bool) -> tuple[str, ...]:
|
||||
if not bot and not touches_cost_map(changed_files):
|
||||
return ()
|
||||
head_map: Final = _parse_object(head.cost_map, COST_MAP_PATH)
|
||||
if isinstance(head_map, str):
|
||||
return (head_map,)
|
||||
return (*_file_failures(head, head_map), *(_bot_failures(base, head_map, changed_files) if bot else ()))
|
||||
|
||||
|
||||
def _git(*args: str) -> str:
|
||||
def _git(*args: str) -> str | None:
|
||||
result: Final = subprocess.run(("git", *args), check=False, capture_output=True, text=True)
|
||||
return result.stdout if result.returncode == 0 else ""
|
||||
return result.stdout if result.returncode == 0 else None
|
||||
|
||||
|
||||
def snapshot(revision: str) -> Snapshot:
|
||||
return Snapshot(*(_git("show", f"{revision}:{path}") for path in GUARDED_PATHS))
|
||||
return Snapshot(*(_git("show", f"{revision}:{path}") or "" for path in GUARDED_PATHS))
|
||||
|
||||
|
||||
def changed_files(base: str, head: str) -> tuple[str, ...] | None:
|
||||
diff: Final = _git("diff", "--name-only", "--no-renames", base, head)
|
||||
return None if diff is None else tuple(diff.splitlines())
|
||||
|
||||
|
||||
def main(argv: Sequence[str]) -> int:
|
||||
|
|
@ -131,9 +151,12 @@ def main(argv: Sequence[str]) -> int:
|
|||
parser.add_argument("--head-ref", required=True, help="head branch name of the pull request")
|
||||
args: Final = parser.parse_args(argv)
|
||||
bot: Final = args.head_ref.startswith(BOT_BRANCH_PREFIX)
|
||||
changed_files: Final = tuple(_git("diff", "--name-only", args.base, args.head).splitlines())
|
||||
failures: Final = guard_failures(snapshot(args.base), snapshot(args.head), changed_files, bot)
|
||||
contract: Final = "bot contract enforced" if bot else "human PR, file checks only"
|
||||
changed: Final = changed_files(args.base, args.head)
|
||||
if changed is None:
|
||||
print(f"cost map guard failed: git diff {args.base} {args.head} failed, so the changed files are unknown")
|
||||
return 1
|
||||
failures: Final = guard_failures(snapshot(args.base), snapshot(args.head), changed, bot)
|
||||
contract: Final = contract_for(bot, changed)
|
||||
if failures:
|
||||
print(f"cost map guard failed ({contract}):")
|
||||
print("\n".join(f"- {failure}" for failure in failures))
|
||||
|
|
|
|||
|
|
@ -6267,6 +6267,63 @@
|
|||
],
|
||||
"title": "Spend update queue sizes (litellm_<queue>_size)",
|
||||
"type": "timeseries"
|
||||
},
|
||||
{
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "${DS_PROMETHEUS}"
|
||||
},
|
||||
"description": "Requests that carried usage but were logged at $0 on a model whose pricing entry has a non-zero rate, by requested model and reason",
|
||||
"fieldConfig": {
|
||||
"defaults": {
|
||||
"color": {
|
||||
"mode": "palette-classic"
|
||||
},
|
||||
"custom": {
|
||||
"drawStyle": "line",
|
||||
"fillOpacity": 10,
|
||||
"lineWidth": 1,
|
||||
"showPoints": "never",
|
||||
"spanNulls": false
|
||||
},
|
||||
"unit": "short"
|
||||
},
|
||||
"overrides": []
|
||||
},
|
||||
"gridPos": {
|
||||
"h": 8,
|
||||
"w": 12,
|
||||
"x": 0,
|
||||
"y": 430
|
||||
},
|
||||
"id": 110,
|
||||
"options": {
|
||||
"legend": {
|
||||
"calcs": [],
|
||||
"displayMode": "list",
|
||||
"placement": "bottom",
|
||||
"showLegend": true
|
||||
},
|
||||
"tooltip": {
|
||||
"mode": "multi",
|
||||
"sort": "desc"
|
||||
}
|
||||
},
|
||||
"targets": [
|
||||
{
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "${DS_PROMETHEUS}"
|
||||
},
|
||||
"editorMode": "code",
|
||||
"expr": "sum(rate(litellm_zero_cost_requests_total[$__rate_interval])) by (requested_model, reason)",
|
||||
"legendFormat": "{{requested_model}} / {{reason}}",
|
||||
"range": true,
|
||||
"refId": "A"
|
||||
}
|
||||
],
|
||||
"title": "litellm_zero_cost_requests rate",
|
||||
"type": "timeseries"
|
||||
}
|
||||
],
|
||||
"preload": false,
|
||||
|
|
|
|||
|
|
@ -96,16 +96,19 @@ GATEWAY_PATH_PREFIXES: tuple[str, ...] = (
|
|||
"/assemblyai/",
|
||||
"/eu.assemblyai/",
|
||||
"/deepgram/",
|
||||
"/fal_ai/",
|
||||
"/langfuse/",
|
||||
"/vllm/",
|
||||
"/mistral/",
|
||||
"/typesafe/",
|
||||
"/openrouter/",
|
||||
"/nvidia_nim/",
|
||||
"/groq/",
|
||||
"/voyage/",
|
||||
"/cursor/",
|
||||
"/milvus/",
|
||||
"/openai_passthrough/",
|
||||
"/tinyfish/",
|
||||
# Dynamic provider / toolset passthrough (path templates)
|
||||
"/{provider}/",
|
||||
"/toolset/",
|
||||
|
|
|
|||
|
|
@ -0,0 +1,2 @@
|
|||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN IF NOT EXISTS "access_group_ids" TEXT[] DEFAULT ARRAY[]::TEXT[];
|
||||
|
|
@ -73,6 +73,7 @@ model LiteLLM_AgentsTable {
|
|||
static_headers Json? @default("{}")
|
||||
extra_headers String[] @default([])
|
||||
agent_access_groups String[] @default([])
|
||||
access_group_ids String[] @default([])
|
||||
object_permission_id String?
|
||||
object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id])
|
||||
spend Float @default(0.0)
|
||||
|
|
|
|||
146
litellm-rust/Cargo.lock
generated
146
litellm-rust/Cargo.lock
generated
|
|
@ -650,6 +650,49 @@ dependencies = [
|
|||
"tracing",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "axum"
|
||||
version = "0.8.9"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "31b698c5f9a010f6573133b09e0de5408834d0c82f8d7475a89fc1867a71cd90"
|
||||
dependencies = [
|
||||
"axum-core",
|
||||
"bytes",
|
||||
"futures-util",
|
||||
"http 1.4.2",
|
||||
"http-body 1.1.0",
|
||||
"http-body-util",
|
||||
"itoa",
|
||||
"matchit",
|
||||
"memchr",
|
||||
"mime",
|
||||
"percent-encoding",
|
||||
"pin-project-lite",
|
||||
"serde_core",
|
||||
"sync_wrapper",
|
||||
"tower",
|
||||
"tower-layer",
|
||||
"tower-service",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "axum-core"
|
||||
version = "0.5.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "08c78f31d7b1291f7ee735c1c6780ccde7785daae9a9206026862dab7d8792d1"
|
||||
dependencies = [
|
||||
"bytes",
|
||||
"futures-core",
|
||||
"http 1.4.2",
|
||||
"http-body 1.1.0",
|
||||
"http-body-util",
|
||||
"mime",
|
||||
"pin-project-lite",
|
||||
"sync_wrapper",
|
||||
"tower-layer",
|
||||
"tower-service",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "azure_core"
|
||||
version = "1.1.0"
|
||||
|
|
@ -1439,7 +1482,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
|||
checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb"
|
||||
dependencies = [
|
||||
"libc",
|
||||
"windows-sys 0.52.0",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -2241,7 +2284,7 @@ dependencies = [
|
|||
"libc",
|
||||
"percent-encoding",
|
||||
"pin-project-lite",
|
||||
"socket2 0.5.10",
|
||||
"socket2 0.6.5",
|
||||
"tokio",
|
||||
"tower-service",
|
||||
"tracing",
|
||||
|
|
@ -2729,6 +2772,26 @@ dependencies = [
|
|||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-cache-qdrant-semantic"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"futures-util",
|
||||
"litellm-cache",
|
||||
"litellm-cache-response",
|
||||
"qdrant-client",
|
||||
"reqwest 0.12.28",
|
||||
"rstest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
"tokio-stream",
|
||||
"tonic",
|
||||
"tonic-prost",
|
||||
"uuid",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-cache-redis"
|
||||
version = "0.1.0"
|
||||
|
|
@ -2741,6 +2804,21 @@ dependencies = [
|
|||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-cache-redis-semantic"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"litellm-cache",
|
||||
"litellm-cache-redis",
|
||||
"litellm-cache-response",
|
||||
"r2d2",
|
||||
"redis",
|
||||
"redis-test",
|
||||
"serde_json",
|
||||
"sha2 0.10.9",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-cache-response"
|
||||
version = "0.1.0"
|
||||
|
|
@ -2816,6 +2894,7 @@ dependencies = [
|
|||
"litellm-host",
|
||||
"litellm-http",
|
||||
"litellm-llms",
|
||||
"litellm-secrets",
|
||||
"litellm-types",
|
||||
"mime_guess",
|
||||
"moka",
|
||||
|
|
@ -2928,6 +3007,7 @@ dependencies = [
|
|||
"litellm-framing",
|
||||
"litellm-host",
|
||||
"litellm-http",
|
||||
"litellm-secrets",
|
||||
"litellm-types",
|
||||
"reqwest 0.12.28",
|
||||
"rstest",
|
||||
|
|
@ -2946,6 +3026,7 @@ dependencies = [
|
|||
name = "litellm-python-bridge"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"aws-sdk-secretsmanager",
|
||||
"bytes",
|
||||
"criterion",
|
||||
"futures-util",
|
||||
|
|
@ -2957,7 +3038,9 @@ dependencies = [
|
|||
"litellm-cache-disk",
|
||||
"litellm-cache-gcs",
|
||||
"litellm-cache-memory",
|
||||
"litellm-cache-qdrant-semantic",
|
||||
"litellm-cache-redis",
|
||||
"litellm-cache-redis-semantic",
|
||||
"litellm-cache-response",
|
||||
"litellm-cache-s3",
|
||||
"litellm-cache-valkey-semantic",
|
||||
|
|
@ -2967,11 +3050,16 @@ dependencies = [
|
|||
"litellm-host-python",
|
||||
"litellm-http",
|
||||
"litellm-llms",
|
||||
"litellm-secrets",
|
||||
"litellm-secrets-aws",
|
||||
"litellm-secrets-types",
|
||||
"litellm-token-counter",
|
||||
"litellm-types",
|
||||
"pyo3",
|
||||
"pyo3-async-runtimes",
|
||||
"qdrant-client",
|
||||
"redis",
|
||||
"reqwest 0.12.28",
|
||||
"rstest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
|
|
@ -2979,6 +3067,8 @@ dependencies = [
|
|||
"sha2 0.10.9",
|
||||
"tokio",
|
||||
"tokio-tungstenite",
|
||||
"url",
|
||||
"wiremock",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -3163,6 +3253,7 @@ dependencies = [
|
|||
name = "litellm-token-counter-huggingface"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"serde_json",
|
||||
"thiserror 2.0.19",
|
||||
"tokenizers",
|
||||
]
|
||||
|
|
@ -3171,6 +3262,9 @@ dependencies = [
|
|||
name = "litellm-token-counter-tiktoken"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"base64 0.22.1",
|
||||
"once_cell",
|
||||
"rustc-hash",
|
||||
"thiserror 2.0.19",
|
||||
"tiktoken-rs",
|
||||
]
|
||||
|
|
@ -3235,6 +3329,12 @@ version = "0.2.3"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "fc04a4c58212d57930a24bf47d3fa87485264a3a054e9c10e042eb373573ad3c"
|
||||
|
||||
[[package]]
|
||||
name = "matchit"
|
||||
version = "0.8.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "47e1ffaa40ddd1f3ed91f717a33c8c0ee23fff369e3aa8772b9605cc1d22f4c3"
|
||||
|
||||
[[package]]
|
||||
name = "md-5"
|
||||
version = "0.11.0"
|
||||
|
|
@ -3858,6 +3958,27 @@ dependencies = [
|
|||
"serde",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "qdrant-client"
|
||||
version = "1.19.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "dddc19df129bad7346ebd027288621ab1ac7e52678371f906b9a8622d7aaf87e"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"derive_builder",
|
||||
"futures",
|
||||
"parking_lot",
|
||||
"prost",
|
||||
"prost-types",
|
||||
"semver",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
"tonic",
|
||||
"tonic-prost",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "quick-error"
|
||||
version = "1.2.3"
|
||||
|
|
@ -3887,7 +4008,7 @@ dependencies = [
|
|||
"quinn-udp",
|
||||
"rustc-hash",
|
||||
"rustls 0.23.42",
|
||||
"socket2 0.5.10",
|
||||
"socket2 0.6.5",
|
||||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
"tracing",
|
||||
|
|
@ -3926,9 +4047,9 @@ dependencies = [
|
|||
"cfg_aliases",
|
||||
"libc",
|
||||
"once_cell",
|
||||
"socket2 0.5.10",
|
||||
"socket2 0.6.5",
|
||||
"tracing",
|
||||
"windows-sys 0.52.0",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -4414,7 +4535,7 @@ dependencies = [
|
|||
"errno",
|
||||
"libc",
|
||||
"linux-raw-sys",
|
||||
"windows-sys 0.52.0",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -4485,7 +4606,7 @@ dependencies = [
|
|||
"security-framework",
|
||||
"security-framework-sys",
|
||||
"webpki-root-certs",
|
||||
"windows-sys 0.52.0",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -5057,10 +5178,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
|||
checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd"
|
||||
dependencies = [
|
||||
"fastrand",
|
||||
"getrandom 0.3.4",
|
||||
"getrandom 0.4.3",
|
||||
"once_cell",
|
||||
"rustix",
|
||||
"windows-sys 0.52.0",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -5349,8 +5470,12 @@ version = "0.14.6"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ac2a5518c70fa84342385732db33fb3f44bc4cc748936eb5833d2df34d6445ef"
|
||||
dependencies = [
|
||||
"async-trait",
|
||||
"axum",
|
||||
"base64 0.22.1",
|
||||
"bytes",
|
||||
"flate2",
|
||||
"h2 0.4.15",
|
||||
"http 1.4.2",
|
||||
"http-body 1.1.0",
|
||||
"http-body-util",
|
||||
|
|
@ -5360,6 +5485,7 @@ dependencies = [
|
|||
"percent-encoding",
|
||||
"pin-project",
|
||||
"rustls-native-certs",
|
||||
"socket2 0.6.5",
|
||||
"sync_wrapper",
|
||||
"tokio",
|
||||
"tokio-rustls 0.26.4",
|
||||
|
|
@ -5933,7 +6059,7 @@ version = "0.1.11"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22"
|
||||
dependencies = [
|
||||
"windows-sys 0.52.0",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
|
|||
|
|
@ -36,7 +36,9 @@ 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-token-counter = { path = "crates/token-counter" }
|
||||
litellm-token-counter-fast = { path = "crates/token-counter-fast" }
|
||||
litellm-token-counter-huggingface = { path = "crates/token-counter-huggingface" }
|
||||
|
|
@ -54,6 +56,8 @@ pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] }
|
|||
pythonize = "0.29.0"
|
||||
rand = "0.8"
|
||||
reqwest = { version = "0.12", default-features = false, features = ["json", "multipart", "rustls-tls", "http2", "stream"] }
|
||||
qdrant-client = { version = "1.19.0", default-features = false }
|
||||
uuid = { version = "1", features = ["v4"] }
|
||||
rstest = "0.26.1"
|
||||
rstest_reuse = "0.7.0"
|
||||
rustls = { version = "0.23", default-features = false, features = ["ring", "std", "tls12"] }
|
||||
|
|
@ -85,7 +89,7 @@ veil = "0.3.0"
|
|||
|
||||
[profile.release]
|
||||
opt-level = 3
|
||||
lto = "thin"
|
||||
lto = "fat"
|
||||
codegen-units = 1
|
||||
panic = "unwind"
|
||||
debug = false
|
||||
|
|
|
|||
|
|
@ -621,6 +621,26 @@ mod tests {
|
|||
None
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn secret_names_cover_environment_reads() {
|
||||
let seen = std::sync::Arc::new(std::sync::Mutex::new(
|
||||
std::collections::BTreeSet::<String>::new(),
|
||||
));
|
||||
let recorded = seen.clone();
|
||||
let env = |name: &str| {
|
||||
recorded.lock().unwrap().insert(name.to_string());
|
||||
None
|
||||
};
|
||||
resolve_aws_region(None, &Map::new(), &env);
|
||||
aws_auth_config(&Map::new(), &env);
|
||||
assert!(
|
||||
seen.lock()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.all(|name| crate::constants::SECRET_NAMES.contains(&name.as_str()))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_region_comes_from_the_call_then_the_model_then_the_environment() {
|
||||
let params = Map::from_iter([("aws_region_name".to_string(), Value::from("eu-west-1"))]);
|
||||
|
|
|
|||
|
|
@ -14,6 +14,19 @@ pub const AWS_WEB_IDENTITY_TOKEN_FILE: &str = "AWS_WEB_IDENTITY_TOKEN_FILE";
|
|||
pub const AWS_STS_ENDPOINT: &str = "AWS_STS_ENDPOINT";
|
||||
pub const AWS_EXTERNAL_ID: &str = "AWS_EXTERNAL_ID";
|
||||
pub const AWS_BEARER_TOKEN_BEDROCK: &str = "AWS_BEARER_TOKEN_BEDROCK";
|
||||
pub const SECRET_NAMES: &[&str] = &[
|
||||
AWS_ACCESS_KEY_ID,
|
||||
AWS_SECRET_ACCESS_KEY,
|
||||
AWS_SESSION_TOKEN,
|
||||
AWS_REGION_NAME,
|
||||
AWS_REGION,
|
||||
AWS_SESSION_NAME,
|
||||
AWS_PROFILE_NAME,
|
||||
AWS_ROLE_NAME,
|
||||
AWS_WEB_IDENTITY_TOKEN,
|
||||
AWS_STS_ENDPOINT,
|
||||
AWS_EXTERNAL_ID,
|
||||
];
|
||||
|
||||
/// Headers SigV4 covers, beyond the `x-amz-` / `x-amzn-` prefixes. Mirrors
|
||||
/// Python's `_filter_headers_for_aws_signature` allowlist.
|
||||
|
|
|
|||
|
|
@ -3,5 +3,5 @@ mod native;
|
|||
mod resolve;
|
||||
mod types;
|
||||
|
||||
pub use resolve::AzureAuthService;
|
||||
pub use resolve::{AzureAuthService, SECRET_NAMES};
|
||||
pub use types::{AzureAuthInputs, ConfigValue};
|
||||
|
|
|
|||
|
|
@ -19,6 +19,17 @@ const AZURE_AUTHORITY_HOST_ENV: &str = "AZURE_AUTHORITY_HOST";
|
|||
const AZURE_CREDENTIAL_ENV: &str = "AZURE_CREDENTIAL";
|
||||
const AZURE_FEDERATED_TOKEN_FILE_ENV: &str = "AZURE_FEDERATED_TOKEN_FILE";
|
||||
|
||||
pub const SECRET_NAMES: &[&str] = &[
|
||||
AZURE_AD_TOKEN_ENV,
|
||||
AZURE_TENANT_ID_ENV,
|
||||
AZURE_CLIENT_ID_ENV,
|
||||
AZURE_CLIENT_SECRET_ENV,
|
||||
AZURE_SCOPE_ENV,
|
||||
AZURE_AUTHORITY_HOST_ENV,
|
||||
AZURE_CREDENTIAL_ENV,
|
||||
AZURE_FEDERATED_TOKEN_FILE_ENV,
|
||||
];
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub(crate) enum AzureCredentialPlan {
|
||||
Supplied(Sourced<ResolvedCredential>),
|
||||
|
|
@ -440,13 +451,14 @@ fn non_empty_reference(value: &str, kind: &str) -> Result<String, Error> {
|
|||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::BTreeSet;
|
||||
use std::future::Future;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use serde_json::json;
|
||||
|
||||
use super::{
|
||||
AzureAuthService, AzureCredentialPlan, AzureTokenAcquirer, oidc_reference,
|
||||
AzureAuthService, AzureCredentialPlan, AzureTokenAcquirer, SECRET_NAMES, oidc_reference,
|
||||
resolve_reference, select_auth_plan,
|
||||
};
|
||||
use crate::native::ValidatedAzureRequest;
|
||||
|
|
@ -517,6 +529,24 @@ mod tests {
|
|||
assert!(matches!(plan, AzureCredentialPlan::Native(_)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn secret_names_cover_environment_reads() {
|
||||
let seen = std::sync::Arc::new(std::sync::Mutex::new(BTreeSet::<String>::new()));
|
||||
let recorded = seen.clone();
|
||||
let inputs = AzureAuthInputs::default();
|
||||
select_auth_plan(&inputs, &|name| {
|
||||
recorded.lock().unwrap().insert(name.to_string());
|
||||
None
|
||||
})
|
||||
.unwrap();
|
||||
assert!(
|
||||
seen.lock()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.all(|name| SECRET_NAMES.contains(&name.as_str()))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn supplied_token_does_not_require_refresh() {
|
||||
let params = json!({"azure_ad_token": "token"});
|
||||
|
|
|
|||
|
|
@ -23,6 +23,16 @@ const VERTEXAI_PROJECT_ENV: &str = "VERTEXAI_PROJECT";
|
|||
const VERTEXAI_LOCATION_ENV: &str = "VERTEXAI_LOCATION";
|
||||
const VERTEX_LOCATION_ENV: &str = "VERTEX_LOCATION";
|
||||
|
||||
pub const SECRET_NAMES: &[&str] = &[
|
||||
VERTEX_AI_API_KEY_ENV,
|
||||
VERTEXAI_API_KEY_ENV,
|
||||
VERTEXAI_CREDENTIALS_ENV,
|
||||
GOOGLE_APPLICATION_CREDENTIALS_ENV,
|
||||
VERTEXAI_PROJECT_ENV,
|
||||
VERTEXAI_LOCATION_ENV,
|
||||
VERTEX_LOCATION_ENV,
|
||||
];
|
||||
|
||||
#[derive(Clone, Debug, Default)]
|
||||
pub struct VertexConfig {
|
||||
credentials: Option<Sourced<SecretValue>>,
|
||||
|
|
@ -406,6 +416,7 @@ fn auth_acquisition_error(error: gcp_auth::Error) -> Error {
|
|||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::BTreeSet;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
|
||||
use serde_json::json;
|
||||
|
|
@ -476,6 +487,27 @@ mod tests {
|
|||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn secret_names_cover_environment_reads() {
|
||||
let seen = Arc::new(std::sync::Mutex::new(BTreeSet::<String>::new()));
|
||||
let recorded = seen.clone();
|
||||
let env = |name: &str| {
|
||||
recorded.lock().unwrap().insert(name.to_string());
|
||||
None
|
||||
};
|
||||
let auth = auth(Arc::new(AtomicUsize::new(0)), Arc::new(AtomicUsize::new(0)));
|
||||
auth.validate_environment(Vec::new(), None, &VertexConfig::default(), &env)
|
||||
.await
|
||||
.unwrap();
|
||||
get_vertex_ai_location(&VertexConfig::default(), &env);
|
||||
assert!(
|
||||
seen.lock()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.all(|name| SECRET_NAMES.contains(&name.as_str()))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_primary_values_fall_back_to_python_aliases() {
|
||||
let config = config(json!({
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
use serde::Deserialize;
|
||||
use std::hash::{Hash, Hasher};
|
||||
use veil::Redact;
|
||||
|
||||
#[derive(Redact, Clone, Deserialize)]
|
||||
|
|
@ -23,6 +24,12 @@ impl PartialEq for SecretValue {
|
|||
|
||||
impl Eq for SecretValue {}
|
||||
|
||||
impl Hash for SecretValue {
|
||||
fn hash<H: Hasher>(&self, state: &mut H) {
|
||||
self.0.hash(state);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::SecretValue;
|
||||
|
|
|
|||
24
litellm-rust/crates/cache-qdrant-semantic/Cargo.toml
Normal file
24
litellm-rust/crates/cache-qdrant-semantic/Cargo.toml
Normal file
|
|
@ -0,0 +1,24 @@
|
|||
[package]
|
||||
name = "litellm-cache-qdrant-semantic"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
futures-util.workspace = true
|
||||
litellm-cache.workspace = true
|
||||
qdrant-client = { workspace = true, features = ["serde"] }
|
||||
reqwest.workspace = true
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
thiserror.workspace = true
|
||||
tokio.workspace = true
|
||||
uuid.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
litellm-cache-response.workspace = true
|
||||
rstest.workspace = true
|
||||
tonic = "0.14"
|
||||
tonic-prost = "0.14"
|
||||
tokio-stream = "0.1"
|
||||
75
litellm-rust/crates/cache-qdrant-semantic/src/embedder.rs
Normal file
75
litellm-rust/crates/cache-qdrant-semantic/src/embedder.rs
Normal file
|
|
@ -0,0 +1,75 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use litellm_cache::Error;
|
||||
use reqwest::Client;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::Embedder;
|
||||
|
||||
pub struct OpenAiEmbedder {
|
||||
client: Client,
|
||||
api_base: String,
|
||||
api_key: String,
|
||||
model: String,
|
||||
timeout: Option<Duration>,
|
||||
}
|
||||
|
||||
pub struct OpenAiEmbedderConfig {
|
||||
pub api_base: String,
|
||||
pub api_key: String,
|
||||
pub model: String,
|
||||
pub timeout: Option<Duration>,
|
||||
}
|
||||
|
||||
impl OpenAiEmbedder {
|
||||
pub fn new(client: Client, config: OpenAiEmbedderConfig) -> Self {
|
||||
Self {
|
||||
client,
|
||||
api_base: config.api_base.trim_end_matches('/').to_owned(),
|
||||
api_key: config.api_key,
|
||||
model: config.model,
|
||||
timeout: config.timeout,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Embedder for OpenAiEmbedder {
|
||||
fn model(&self) -> &str {
|
||||
&self.model
|
||||
}
|
||||
|
||||
async fn embed(&self, input: &str) -> Result<Vec<f32>, Error> {
|
||||
let request = self
|
||||
.client
|
||||
.post(format!("{}/embeddings", self.api_base))
|
||||
.bearer_auth(&self.api_key)
|
||||
.json(&serde_json::json!({
|
||||
"model": self.model,
|
||||
"input": input,
|
||||
"encoding_format": "float",
|
||||
}));
|
||||
let response = if let Some(timeout) = self.timeout {
|
||||
request.timeout(timeout)
|
||||
} else {
|
||||
request
|
||||
}
|
||||
.send()
|
||||
.await
|
||||
.map_err(|_| Error::Unavailable)?
|
||||
.error_for_status()
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
let body: Value = response.json().await.map_err(|_| Error::Unavailable)?;
|
||||
body.get("data")
|
||||
.and_then(Value::as_array)
|
||||
.and_then(|data| data.first())
|
||||
.and_then(|item| item.get("embedding"))
|
||||
.and_then(Value::as_array)
|
||||
.and_then(|embedding| {
|
||||
embedding
|
||||
.iter()
|
||||
.map(|value| value.as_f64().map(|value| value as f32))
|
||||
.collect::<Option<Vec<_>>>()
|
||||
})
|
||||
.ok_or(Error::Unavailable)
|
||||
}
|
||||
}
|
||||
7
litellm-rust/crates/cache-qdrant-semantic/src/lib.rs
Normal file
7
litellm-rust/crates/cache-qdrant-semantic/src/lib.rs
Normal file
|
|
@ -0,0 +1,7 @@
|
|||
mod embedder;
|
||||
mod prompt;
|
||||
mod semantic;
|
||||
|
||||
pub use embedder::{OpenAiEmbedder, OpenAiEmbedderConfig};
|
||||
pub use prompt::prompt_from_messages;
|
||||
pub use semantic::{Embedder, QdrantSemanticCache, QdrantSemanticConfig, Quantization};
|
||||
59
litellm-rust/crates/cache-qdrant-semantic/src/prompt.rs
Normal file
59
litellm-rust/crates/cache-qdrant-semantic/src/prompt.rs
Normal file
|
|
@ -0,0 +1,59 @@
|
|||
use serde_json::Value;
|
||||
|
||||
fn search_results_text(search_results: Option<&Value>) -> String {
|
||||
let Some(Value::Array(results)) = search_results else {
|
||||
return String::new();
|
||||
};
|
||||
results
|
||||
.iter()
|
||||
.filter_map(Value::as_object)
|
||||
.flat_map(|result| {
|
||||
let source = result
|
||||
.get("source")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::to_owned);
|
||||
let title = result
|
||||
.get("title")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::to_owned);
|
||||
let content = result
|
||||
.get("content")
|
||||
.and_then(Value::as_array)
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.filter_map(Value::as_object)
|
||||
.filter_map(|block| block.get("text").and_then(Value::as_str).map(str::to_owned));
|
||||
let citations = result
|
||||
.get("citations")
|
||||
.filter(|value| !value.is_null())
|
||||
.map(|value| serde_json::to_string(value).unwrap_or_default());
|
||||
source
|
||||
.into_iter()
|
||||
.chain(title)
|
||||
.chain(content)
|
||||
.chain(citations)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub fn prompt_from_messages(messages: &[Value]) -> String {
|
||||
messages
|
||||
.iter()
|
||||
.filter_map(Value::as_object)
|
||||
.map(|message| {
|
||||
let content = match message.get("content") {
|
||||
Some(Value::String(content)) => content.clone(),
|
||||
Some(Value::Array(parts)) => parts
|
||||
.iter()
|
||||
.filter_map(Value::as_object)
|
||||
.filter_map(|part| part.get("text").and_then(Value::as_str))
|
||||
.collect(),
|
||||
_ => String::new(),
|
||||
};
|
||||
format!(
|
||||
"{content}{}",
|
||||
search_results_text(message.get("search_results"))
|
||||
)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
262
litellm-rust/crates/cache-qdrant-semantic/src/semantic.rs
Normal file
262
litellm-rust/crates/cache-qdrant-semantic/src/semantic.rs
Normal file
|
|
@ -0,0 +1,262 @@
|
|||
use std::future::Future;
|
||||
|
||||
use futures_util::future::try_join_all;
|
||||
use litellm_cache::{BaseCache, CacheCodec, CacheConnectionResult, Error, SemanticCacheContext};
|
||||
use qdrant_client::{
|
||||
Payload, Qdrant,
|
||||
qdrant::{
|
||||
BinaryQuantizationBuilder, CompressionRatio, Condition, CreateCollectionBuilder,
|
||||
CreateFieldIndexCollectionBuilder, Distance, FieldType, Filter, PointStruct,
|
||||
ProductQuantizationBuilder, QuantizationSearchParamsBuilder, ScalarQuantizationBuilder,
|
||||
SearchParamsBuilder, SearchPointsBuilder, UpsertPointsBuilder, VectorParamsBuilder,
|
||||
},
|
||||
};
|
||||
use serde_json::{Map, Value, json};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::prompt_from_messages;
|
||||
|
||||
pub trait Embedder: Send + Sync + 'static {
|
||||
fn model(&self) -> &str;
|
||||
fn embed(&self, input: &str) -> impl Future<Output = Result<Vec<f32>, Error>> + Send;
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub enum Quantization {
|
||||
Binary,
|
||||
Scalar,
|
||||
Product,
|
||||
}
|
||||
|
||||
pub struct QdrantSemanticConfig {
|
||||
pub collection_name: String,
|
||||
pub similarity_threshold: f64,
|
||||
pub vector_size: u64,
|
||||
pub quantization: Quantization,
|
||||
}
|
||||
|
||||
pub struct QdrantSemanticCache<E: Embedder, C: CacheCodec> {
|
||||
client: Qdrant,
|
||||
embedder: E,
|
||||
codec: C,
|
||||
config: QdrantSemanticConfig,
|
||||
runtime: tokio::runtime::Handle,
|
||||
}
|
||||
|
||||
impl<E: Embedder, C: CacheCodec> QdrantSemanticCache<E, C> {
|
||||
pub async fn connect(
|
||||
client: Qdrant,
|
||||
embedder: E,
|
||||
codec: C,
|
||||
config: QdrantSemanticConfig,
|
||||
runtime: tokio::runtime::Handle,
|
||||
) -> Result<Self, Error> {
|
||||
let exists = client
|
||||
.collection_exists(config.collection_name.clone())
|
||||
.await
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
if !exists {
|
||||
client
|
||||
.create_collection(
|
||||
CreateCollectionBuilder::new(config.collection_name.clone())
|
||||
.vectors_config(VectorParamsBuilder::new(
|
||||
config.vector_size,
|
||||
Distance::Cosine,
|
||||
))
|
||||
.quantization_config(quantization(&config.quantization)),
|
||||
)
|
||||
.await
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
}
|
||||
let _ = client
|
||||
.create_field_index(CreateFieldIndexCollectionBuilder::new(
|
||||
config.collection_name.clone(),
|
||||
"litellm_cache_key".to_owned(),
|
||||
FieldType::Keyword,
|
||||
))
|
||||
.await;
|
||||
Ok(Self {
|
||||
client,
|
||||
embedder,
|
||||
codec,
|
||||
config,
|
||||
runtime,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn collection_name(&self) -> &str {
|
||||
&self.config.collection_name
|
||||
}
|
||||
|
||||
pub fn similarity_threshold(&self) -> f64 {
|
||||
self.config.similarity_threshold
|
||||
}
|
||||
|
||||
pub fn vector_size(&self) -> u64 {
|
||||
self.config.vector_size
|
||||
}
|
||||
|
||||
pub fn embedder(&self) -> &E {
|
||||
&self.embedder
|
||||
}
|
||||
|
||||
fn prompt(context: &SemanticCacheContext) -> Result<String, Error> {
|
||||
let Some(messages) = context.messages.as_ref().and_then(Value::as_array) else {
|
||||
return Err(Error::MissingPrompt);
|
||||
};
|
||||
if messages.is_empty() {
|
||||
return Err(Error::MissingPrompt);
|
||||
}
|
||||
Ok(prompt_from_messages(messages))
|
||||
}
|
||||
|
||||
async fn set(
|
||||
&self,
|
||||
key: &str,
|
||||
value: C::Value,
|
||||
context: &SemanticCacheContext,
|
||||
) -> Result<(), Error> {
|
||||
let prompt = Self::prompt(context)?;
|
||||
let vector = self.embedder.embed(&prompt).await?;
|
||||
let response =
|
||||
String::from_utf8(self.codec.encode(&value)?).map_err(|_| Error::InvalidEntry)?;
|
||||
let payload = Payload::try_from(json!({
|
||||
"litellm_cache_key": key,
|
||||
"text": prompt,
|
||||
"response": response,
|
||||
}))
|
||||
.map_err(|_| Error::InvalidEntry)?;
|
||||
self.client
|
||||
.upsert_points(
|
||||
UpsertPointsBuilder::new(
|
||||
self.collection_name(),
|
||||
vec![PointStruct::new(
|
||||
Uuid::new_v4().to_string(),
|
||||
vector,
|
||||
payload,
|
||||
)],
|
||||
)
|
||||
.wait(true),
|
||||
)
|
||||
.await
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn get(
|
||||
&self,
|
||||
key: &str,
|
||||
context: &SemanticCacheContext,
|
||||
) -> Result<Option<C::Value>, Error> {
|
||||
let prompt = Self::prompt(context)?;
|
||||
let vector = self.embedder.embed(&prompt).await?;
|
||||
let result = self
|
||||
.client
|
||||
.search_points(
|
||||
SearchPointsBuilder::new(self.collection_name(), vector, 1)
|
||||
.with_payload(true)
|
||||
.filter(Filter::must([Condition::matches(
|
||||
"litellm_cache_key",
|
||||
key.to_owned(),
|
||||
)]))
|
||||
.params(
|
||||
SearchParamsBuilder::default().quantization(
|
||||
QuantizationSearchParamsBuilder::default()
|
||||
.ignore(false)
|
||||
.rescore(true)
|
||||
.oversampling(3.0),
|
||||
),
|
||||
),
|
||||
)
|
||||
.await
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
let Some(point) = result.result.into_iter().next() else {
|
||||
return Ok(None);
|
||||
};
|
||||
let payload: Map<String, Value> = Payload::from(point.payload).into();
|
||||
if payload.get("litellm_cache_key").and_then(Value::as_str) != Some(key) {
|
||||
return Ok(None);
|
||||
}
|
||||
if f64::from(point.score) < self.config.similarity_threshold {
|
||||
return Ok(None);
|
||||
}
|
||||
let response = payload
|
||||
.get("response")
|
||||
.and_then(Value::as_str)
|
||||
.ok_or(Error::InvalidEntry)?;
|
||||
self.codec.decode(response.as_bytes()).map(Some)
|
||||
}
|
||||
}
|
||||
|
||||
fn quantization(value: &Quantization) -> qdrant_client::qdrant::quantization_config::Quantization {
|
||||
match value {
|
||||
Quantization::Binary => BinaryQuantizationBuilder::new(false).into(),
|
||||
Quantization::Scalar => ScalarQuantizationBuilder::default()
|
||||
.quantile(0.99)
|
||||
.always_ram(false)
|
||||
.into(),
|
||||
Quantization::Product => ProductQuantizationBuilder::new(CompressionRatio::X16.into())
|
||||
.always_ram(false)
|
||||
.into(),
|
||||
}
|
||||
}
|
||||
|
||||
impl<E: Embedder, C: CacheCodec> BaseCache for QdrantSemanticCache<E, C> {
|
||||
type Value = C::Value;
|
||||
type Context = SemanticCacheContext;
|
||||
|
||||
fn get_ttl(&self, _: &Self::Context) -> Option<std::time::Duration> {
|
||||
None
|
||||
}
|
||||
|
||||
fn set_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
value: Self::Value,
|
||||
context: &Self::Context,
|
||||
) -> Result<(), Error> {
|
||||
self.runtime.block_on(self.set(key, value, context))
|
||||
}
|
||||
|
||||
fn get_cache(&self, key: &str, context: &Self::Context) -> Result<Option<Self::Value>, Error> {
|
||||
self.runtime.block_on(self.get(key, context))
|
||||
}
|
||||
|
||||
async fn async_set_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
value: Self::Value,
|
||||
context: Self::Context,
|
||||
) -> Result<(), Error> {
|
||||
self.set(key, value, &context).await
|
||||
}
|
||||
|
||||
async fn async_get_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
context: &Self::Context,
|
||||
) -> Result<Option<Self::Value>, Error> {
|
||||
self.get(key, context).await
|
||||
}
|
||||
|
||||
async fn async_set_cache_pipeline(
|
||||
&self,
|
||||
entries: Vec<(String, Self::Value)>,
|
||||
context: Self::Context,
|
||||
) -> Result<(), Error> {
|
||||
try_join_all(entries.into_iter().map(|(key, value)| {
|
||||
let context = context.clone();
|
||||
async move { self.async_set_cache(&key, value, context).await }
|
||||
}))
|
||||
.await
|
||||
.map(|_| ())
|
||||
}
|
||||
|
||||
async fn disconnect(&self) -> Result<(), Error> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
|
||||
Err(Error::UnsupportedOperation)
|
||||
}
|
||||
}
|
||||
166
litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs
Normal file
166
litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs
Normal file
|
|
@ -0,0 +1,166 @@
|
|||
use std::{
|
||||
sync::{Arc, Mutex},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use litellm_cache::Error;
|
||||
use litellm_cache_qdrant_semantic::{Embedder, OpenAiEmbedder, OpenAiEmbedderConfig};
|
||||
use serde_json::Value;
|
||||
use tokio::{
|
||||
io::{AsyncReadExt, AsyncWriteExt},
|
||||
net::TcpListener,
|
||||
};
|
||||
|
||||
struct TestHttpServer {
|
||||
address: std::net::SocketAddr,
|
||||
request: Arc<Mutex<Option<Vec<u8>>>>,
|
||||
task: tokio::task::JoinHandle<()>,
|
||||
}
|
||||
|
||||
impl TestHttpServer {
|
||||
async fn response(status: &str, body: &str) -> Self {
|
||||
Self::response_after(status, body, Duration::ZERO).await
|
||||
}
|
||||
|
||||
async fn response_after(status: &str, body: &str, delay: Duration) -> Self {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let address = listener.local_addr().unwrap();
|
||||
let request = Arc::new(Mutex::new(None));
|
||||
let captured = request.clone();
|
||||
let status = status.to_owned();
|
||||
let body = body.to_owned();
|
||||
let task = tokio::spawn(async move {
|
||||
let (mut stream, _) = listener.accept().await.unwrap();
|
||||
let request_bytes = read_request(&mut stream).await;
|
||||
*captured.lock().unwrap() = Some(request_bytes);
|
||||
tokio::time::sleep(delay).await;
|
||||
let response = format!(
|
||||
"HTTP/1.1 {status}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}",
|
||||
body.len()
|
||||
);
|
||||
stream.write_all(response.as_bytes()).await.unwrap();
|
||||
});
|
||||
Self {
|
||||
address,
|
||||
request,
|
||||
task,
|
||||
}
|
||||
}
|
||||
|
||||
fn base_url(&self) -> String {
|
||||
format!("http://{}", self.address)
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for TestHttpServer {
|
||||
fn drop(&mut self) {
|
||||
self.task.abort();
|
||||
}
|
||||
}
|
||||
|
||||
async fn read_request(stream: &mut tokio::net::TcpStream) -> Vec<u8> {
|
||||
let mut bytes = Vec::new();
|
||||
let header_end = loop {
|
||||
let mut chunk = [0_u8; 1024];
|
||||
let count = stream.read(&mut chunk).await.unwrap();
|
||||
assert_ne!(count, 0);
|
||||
bytes.extend_from_slice(&chunk[..count]);
|
||||
if let Some(end) = bytes.windows(4).position(|window| window == b"\r\n\r\n") {
|
||||
break end + 4;
|
||||
}
|
||||
};
|
||||
let headers = String::from_utf8_lossy(&bytes[..header_end]);
|
||||
let content_length = headers
|
||||
.lines()
|
||||
.find_map(|line| {
|
||||
line.split_once(':')
|
||||
.filter(|(name, _)| name.eq_ignore_ascii_case("content-length"))
|
||||
.map(|(_, value)| value.trim())
|
||||
})
|
||||
.unwrap()
|
||||
.parse::<usize>()
|
||||
.unwrap();
|
||||
while bytes.len() < header_end + content_length {
|
||||
let mut chunk = [0_u8; 1024];
|
||||
let count = stream.read(&mut chunk).await.unwrap();
|
||||
assert_ne!(count, 0);
|
||||
bytes.extend_from_slice(&chunk[..count]);
|
||||
}
|
||||
bytes
|
||||
}
|
||||
|
||||
fn config(base: String, timeout: Option<Duration>) -> OpenAiEmbedderConfig {
|
||||
OpenAiEmbedderConfig {
|
||||
api_base: base,
|
||||
api_key: "test-key".to_owned(),
|
||||
model: "test-model".to_owned(),
|
||||
timeout,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn posts_embeddings_request_and_parses_vector() {
|
||||
let server = TestHttpServer::response("200 OK", r#"{"data":[{"embedding":[0.1,0.2]}]}"#).await;
|
||||
let embedder = OpenAiEmbedder::new(
|
||||
reqwest::Client::new(),
|
||||
config(
|
||||
format!("{}/", server.base_url()),
|
||||
Some(Duration::from_secs(1)),
|
||||
),
|
||||
);
|
||||
assert_eq!(embedder.embed("hello").await.unwrap(), vec![0.1, 0.2]);
|
||||
let request = server.request.lock().unwrap().clone().unwrap();
|
||||
let request_text = String::from_utf8(request).unwrap();
|
||||
assert!(request_text.starts_with("POST /embeddings HTTP/1.1\r\n"));
|
||||
assert!(request_text.contains("\r\nauthorization: Bearer test-key\r\n"));
|
||||
let body = request_text.split("\r\n\r\n").nth(1).unwrap();
|
||||
let body: Value = serde_json::from_str(body).unwrap();
|
||||
assert_eq!(body["model"], "test-model");
|
||||
assert_eq!(body["input"], "hello");
|
||||
assert_eq!(body["encoding_format"], "float");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn status_and_timeout_errors_are_unavailable() {
|
||||
let server = TestHttpServer::response("500 Internal Server Error", "{}").await;
|
||||
let embedder = OpenAiEmbedder::new(reqwest::Client::new(), config(server.base_url(), None));
|
||||
assert_eq!(embedder.embed("hello").await, Err(Error::Unavailable));
|
||||
|
||||
let server = TestHttpServer::response_after(
|
||||
"200 OK",
|
||||
r#"{"data":[{"embedding":[0.1,0.2]}]}"#,
|
||||
Duration::from_millis(500),
|
||||
)
|
||||
.await;
|
||||
let embedder = OpenAiEmbedder::new(
|
||||
reqwest::Client::new(),
|
||||
config(server.base_url(), Some(Duration::from_millis(200))),
|
||||
);
|
||||
assert_eq!(embedder.embed("hello").await, Err(Error::Unavailable));
|
||||
|
||||
let server = TestHttpServer::response_after(
|
||||
"200 OK",
|
||||
r#"{"data":[{"embedding":[0.1,0.2]}]}"#,
|
||||
Duration::from_millis(100),
|
||||
)
|
||||
.await;
|
||||
let embedder = OpenAiEmbedder::new(
|
||||
reqwest::Client::new(),
|
||||
config(server.base_url(), Some(Duration::from_secs(1))),
|
||||
);
|
||||
assert_eq!(embedder.embed("hello").await.unwrap(), vec![0.1, 0.2]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn uses_the_injected_client() {
|
||||
let server = TestHttpServer::response("200 OK", r#"{"data":[{"embedding":[0.1,0.2]}]}"#).await;
|
||||
let client = reqwest::Client::builder()
|
||||
.user_agent("litellm-embedder-test")
|
||||
.build()
|
||||
.unwrap();
|
||||
let embedder = OpenAiEmbedder::new(client, config(server.base_url(), None));
|
||||
assert_eq!(embedder.embed("hello").await.unwrap(), vec![0.1, 0.2]);
|
||||
let request = server.request.lock().unwrap().clone().unwrap();
|
||||
let request_text = String::from_utf8(request).unwrap();
|
||||
assert!(request_text.contains("\r\nuser-agent: litellm-embedder-test\r\n"));
|
||||
}
|
||||
38
litellm-rust/crates/cache-qdrant-semantic/tests/prompt.rs
Normal file
38
litellm-rust/crates/cache-qdrant-semantic/tests/prompt.rs
Normal file
|
|
@ -0,0 +1,38 @@
|
|||
use litellm_cache_qdrant_semantic::prompt_from_messages;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn prompt_matches_python_message_content_rules() {
|
||||
let messages = vec![
|
||||
json!({"role": "user", "content": "hello"}),
|
||||
json!({
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "world"},
|
||||
{"type": "image_url", "image_url": {"url": "ignored"}},
|
||||
{"type": "text", "text": "!"},
|
||||
],
|
||||
}),
|
||||
];
|
||||
|
||||
assert_eq!(prompt_from_messages(&messages), "helloworld!");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prompt_includes_search_result_text_and_compact_citations() {
|
||||
let messages = vec![json!({
|
||||
"role": "tool",
|
||||
"content": null,
|
||||
"search_results": [{
|
||||
"source": "source",
|
||||
"title": "title",
|
||||
"content": [{"text": "body"}],
|
||||
"citations": {"page": 1, "section": "intro"},
|
||||
}],
|
||||
})];
|
||||
|
||||
assert_eq!(
|
||||
prompt_from_messages(&messages),
|
||||
r#"sourcetitlebody{"page":1,"section":"intro"}"#
|
||||
);
|
||||
}
|
||||
422
litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs
Normal file
422
litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs
Normal file
|
|
@ -0,0 +1,422 @@
|
|||
#[path = "support/mod.rs"]
|
||||
mod support;
|
||||
|
||||
use std::{collections::HashMap, sync::Arc, time::Duration};
|
||||
|
||||
use litellm_cache::{BaseCache, CacheCodec, CacheContext, Error, SemanticCacheContext};
|
||||
use litellm_cache_qdrant_semantic::{
|
||||
Embedder, QdrantSemanticCache, QdrantSemanticConfig, Quantization,
|
||||
};
|
||||
use litellm_cache_response::{
|
||||
CacheEntry, CacheKeyInput, ResponseCache, ResponseCacheCodec, ResponseCacheRequest,
|
||||
};
|
||||
use qdrant_client::Payload;
|
||||
use qdrant_client::{
|
||||
Qdrant,
|
||||
qdrant::{self, CompressionRatio, Distance, PointId, QuantizationType, Value, VectorParams},
|
||||
};
|
||||
use serde_json::{Value as JsonValue, json};
|
||||
|
||||
use support::{FakeQdrant, FakeState, StoredPoint};
|
||||
|
||||
#[derive(Clone)]
|
||||
struct FixedEmbedder {
|
||||
vectors: Arc<HashMap<String, Vec<f32>>>,
|
||||
}
|
||||
|
||||
impl FixedEmbedder {
|
||||
fn new(vectors: impl IntoIterator<Item = (&'static str, Vec<f32>)>) -> Self {
|
||||
Self {
|
||||
vectors: Arc::new(
|
||||
vectors
|
||||
.into_iter()
|
||||
.map(|(prompt, vector)| (prompt.to_owned(), vector))
|
||||
.collect(),
|
||||
),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Embedder for FixedEmbedder {
|
||||
fn model(&self) -> &str {
|
||||
"fixed"
|
||||
}
|
||||
|
||||
async fn embed(&self, input: &str) -> Result<Vec<f32>, Error> {
|
||||
self.vectors.get(input).cloned().ok_or(Error::Unavailable)
|
||||
}
|
||||
}
|
||||
|
||||
fn config(quantization: Quantization) -> QdrantSemanticConfig {
|
||||
QdrantSemanticConfig {
|
||||
collection_name: "semantic".to_owned(),
|
||||
similarity_threshold: 0.9,
|
||||
vector_size: 2,
|
||||
quantization,
|
||||
}
|
||||
}
|
||||
|
||||
fn context(prompt: &str) -> SemanticCacheContext {
|
||||
SemanticCacheContext {
|
||||
messages: Some(json!([{"role": "user", "content": prompt}])),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
fn value(response: JsonValue) -> CacheEntry {
|
||||
CacheEntry {
|
||||
timestamp: Some(1.0),
|
||||
response,
|
||||
}
|
||||
}
|
||||
|
||||
async fn connect(
|
||||
server: &FakeQdrant,
|
||||
vectors: impl IntoIterator<Item = (&'static str, Vec<f32>)>,
|
||||
) -> QdrantSemanticCache<FixedEmbedder, ResponseCacheCodec> {
|
||||
let client = Qdrant::from_url(&server.url()).build().unwrap();
|
||||
QdrantSemanticCache::connect(
|
||||
client,
|
||||
FixedEmbedder::new(vectors),
|
||||
ResponseCacheCodec,
|
||||
config(Quantization::Binary),
|
||||
tokio::runtime::Handle::current(),
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
#[expect(
|
||||
deprecated,
|
||||
reason = "the test verifies Qdrant's legacy always_ram quantization contract"
|
||||
)]
|
||||
async fn connect_sets_collection_quantization_and_index() {
|
||||
for (quantization, expected) in [
|
||||
(Quantization::Binary, 0),
|
||||
(Quantization::Scalar, 1),
|
||||
(Quantization::Product, 2),
|
||||
] {
|
||||
let server = FakeQdrant::start(FakeState::default()).await;
|
||||
let client = Qdrant::from_url(&server.url()).build().unwrap();
|
||||
QdrantSemanticCache::connect(
|
||||
client,
|
||||
FixedEmbedder::new([]),
|
||||
ResponseCacheCodec,
|
||||
config(quantization),
|
||||
tokio::runtime::Handle::current(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let state = server.state.lock().unwrap();
|
||||
let request = &state.created_collections[0];
|
||||
let Some(qdrant::vectors_config::Config::Params(VectorParams { size, distance, .. })) =
|
||||
request
|
||||
.vectors_config
|
||||
.as_ref()
|
||||
.and_then(|config| config.config.clone())
|
||||
else {
|
||||
panic!("missing vector params");
|
||||
};
|
||||
assert_eq!(size, 2);
|
||||
assert_eq!(distance, Distance::Cosine as i32);
|
||||
let quantization_config = request
|
||||
.quantization_config
|
||||
.as_ref()
|
||||
.unwrap()
|
||||
.quantization
|
||||
.unwrap();
|
||||
match (expected, quantization_config) {
|
||||
(0, qdrant::quantization_config::Quantization::Binary(binary)) => {
|
||||
assert_eq!(binary.always_ram, Some(false));
|
||||
}
|
||||
(1, qdrant::quantization_config::Quantization::Scalar(scalar)) => {
|
||||
assert_eq!(scalar.r#type, QuantizationType::Int8 as i32);
|
||||
assert_eq!(scalar.quantile, Some(0.99));
|
||||
assert_eq!(scalar.always_ram, Some(false));
|
||||
}
|
||||
(2, qdrant::quantization_config::Quantization::Product(product)) => {
|
||||
assert_eq!(product.compression, CompressionRatio::X16 as i32);
|
||||
assert_eq!(product.always_ram, Some(false));
|
||||
}
|
||||
_ => panic!("unexpected quantization"),
|
||||
}
|
||||
assert!(state.index_creations >= 1);
|
||||
assert_eq!(state.field_indexes[0].collection_name, "semantic");
|
||||
assert_eq!(state.field_indexes[0].field_name, "litellm_cache_key");
|
||||
assert_eq!(
|
||||
state.field_indexes[0].field_type,
|
||||
Some(qdrant::FieldType::Keyword as i32)
|
||||
);
|
||||
server.stop();
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn existing_collection_skips_create_and_index_failure_is_non_fatal() {
|
||||
let server = FakeQdrant::start(FakeState {
|
||||
collections: ["semantic".to_owned()].into_iter().collect(),
|
||||
fail_field_index: true,
|
||||
..Default::default()
|
||||
})
|
||||
.await;
|
||||
let _cache = connect(&server, [("hello", vec![1.0, 0.0])]).await;
|
||||
let state = server.state.lock().unwrap();
|
||||
assert!(state.created_collections.is_empty());
|
||||
assert!(state.index_creations >= 1);
|
||||
server.stop();
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn async_and_sync_set_get_store_exact_payload() {
|
||||
let server = FakeQdrant::start(FakeState::default()).await;
|
||||
let cache = Arc::new(connect(&server, [("hello", vec![1.0, 0.0])]).await);
|
||||
let ctx = context("hello");
|
||||
let entry = value(json!({"answer": 42}));
|
||||
cache
|
||||
.async_set_cache("key", entry.clone(), ctx.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
cache.async_get_cache("key", &ctx).await.unwrap().as_ref(),
|
||||
Some(&entry)
|
||||
);
|
||||
{
|
||||
let state = server.state.lock().unwrap();
|
||||
let payload = &state.points[0].payload;
|
||||
let mut payload_keys = payload.keys().cloned().collect::<Vec<_>>();
|
||||
payload_keys.sort();
|
||||
assert_eq!(payload_keys, ["litellm_cache_key", "response", "text"]);
|
||||
assert_eq!(payload["litellm_cache_key"], Value::from("key"));
|
||||
assert_eq!(
|
||||
payload["response"],
|
||||
Value::from(String::from_utf8(ResponseCacheCodec.encode(&entry).unwrap()).unwrap())
|
||||
);
|
||||
}
|
||||
let sync_entry = entry.clone();
|
||||
let sync_cache = cache.clone();
|
||||
let sync_ctx = ctx.clone();
|
||||
tokio::task::spawn_blocking(move || {
|
||||
sync_cache
|
||||
.set_cache("sync", sync_entry.clone(), &sync_ctx)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
sync_cache.get_cache("sync", &sync_ctx).unwrap(),
|
||||
Some(sync_entry)
|
||||
);
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
server.stop();
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn misses_and_payload_validation_are_safe() {
|
||||
let server = FakeQdrant::start(FakeState::default()).await;
|
||||
let cache = connect(
|
||||
&server,
|
||||
[("hello", vec![1.0, 0.0]), ("near", vec![0.7, 0.71414286])],
|
||||
)
|
||||
.await;
|
||||
let entry = value(json!({"answer": 1}));
|
||||
cache
|
||||
.async_set_cache("key", entry, context("hello"))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_get_cache("other", &context("hello"))
|
||||
.await
|
||||
.unwrap(),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_get_cache("key", &context("near"))
|
||||
.await
|
||||
.unwrap(),
|
||||
None
|
||||
);
|
||||
server.insert_point(StoredPoint {
|
||||
id: Some(PointId::from(99_u64)),
|
||||
vector: vec![1.0, 0.0],
|
||||
payload: Payload::try_from(json!({
|
||||
"litellm_cache_key": 99,
|
||||
"response": "{}",
|
||||
}))
|
||||
.unwrap()
|
||||
.into(),
|
||||
});
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_get_cache("99", &context("hello"))
|
||||
.await
|
||||
.unwrap(),
|
||||
None
|
||||
);
|
||||
server.stop();
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn decoding_errors_missing_prompt_pipeline_and_ttl_behave_as_required() {
|
||||
let server = FakeQdrant::start(FakeState::default()).await;
|
||||
let cache = connect(&server, [("one", vec![1.0, 0.0]), ("two", vec![0.0, 1.0])]).await;
|
||||
let empty = SemanticCacheContext::default();
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_set_cache("key", value(json!({})), empty.clone())
|
||||
.await,
|
||||
Err(Error::MissingPrompt)
|
||||
);
|
||||
assert_eq!(
|
||||
cache.async_get_cache("key", &empty).await,
|
||||
Err(Error::MissingPrompt)
|
||||
);
|
||||
assert_eq!(
|
||||
cache.async_get_cache("key", &context("unknown")).await,
|
||||
Err(Error::Unavailable)
|
||||
);
|
||||
cache
|
||||
.async_set_cache(
|
||||
"ttl",
|
||||
value(json!({"ttl": true})),
|
||||
context("one").with_ttl(Some(Duration::from_secs(1))),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
tokio::time::sleep(Duration::from_millis(1_100)).await;
|
||||
assert!(
|
||||
cache
|
||||
.async_get_cache(
|
||||
"ttl",
|
||||
&context("one").with_ttl(Some(Duration::from_secs(1))),
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.is_some()
|
||||
);
|
||||
cache
|
||||
.async_set_cache_pipeline(
|
||||
vec![
|
||||
("one".to_owned(), value(json!({"n": 1}))),
|
||||
("two".to_owned(), value(json!({"n": 2}))),
|
||||
],
|
||||
context("one"),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(
|
||||
cache
|
||||
.async_get_cache("one", &context("one"))
|
||||
.await
|
||||
.unwrap()
|
||||
.is_some()
|
||||
);
|
||||
assert!(
|
||||
cache
|
||||
.async_get_cache("two", &context("one"))
|
||||
.await
|
||||
.unwrap()
|
||||
.is_some()
|
||||
);
|
||||
assert_eq!(
|
||||
server.state.lock().unwrap().upsert_waits,
|
||||
vec![Some(true), Some(true), Some(true)]
|
||||
);
|
||||
assert_eq!(cache.get_ttl(&context("one")), None);
|
||||
assert_eq!(
|
||||
cache.test_connection().await,
|
||||
Err(Error::UnsupportedOperation)
|
||||
);
|
||||
server.stop();
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn response_payloads_decode_and_invalid_entries_fail() {
|
||||
let server = FakeQdrant::start(FakeState::default()).await;
|
||||
let cache = connect(&server, [("hello", vec![1.0, 0.0])]).await;
|
||||
for (key, response) in [
|
||||
("python", json!("{'timestamp': 1.0, 'response': {'a': 1}}")),
|
||||
("garbage", json!("not json")),
|
||||
("missing", json!("unused")),
|
||||
] {
|
||||
let mut payload = serde_json::Map::new();
|
||||
payload.insert("litellm_cache_key".to_owned(), json!(key));
|
||||
if key != "missing" {
|
||||
payload.insert("response".to_owned(), response);
|
||||
}
|
||||
server.insert_point(StoredPoint {
|
||||
id: Some(PointId::from(key.len() as u64)),
|
||||
vector: vec![1.0, 0.0],
|
||||
payload: Payload::try_from(JsonValue::Object(payload))
|
||||
.unwrap()
|
||||
.into(),
|
||||
});
|
||||
}
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_get_cache("python", &context("hello"))
|
||||
.await
|
||||
.unwrap(),
|
||||
Some(value(json!({"a": 1})))
|
||||
);
|
||||
assert_eq!(
|
||||
cache.async_get_cache("garbage", &context("hello")).await,
|
||||
Err(Error::InvalidEntry)
|
||||
);
|
||||
assert_eq!(
|
||||
cache.async_get_cache("missing", &context("hello")).await,
|
||||
Err(Error::InvalidEntry)
|
||||
);
|
||||
server.stop();
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn response_cache_facade_turns_invalid_entry_into_miss() {
|
||||
let server = FakeQdrant::start(FakeState::default()).await;
|
||||
let cache = Arc::new(connect(&server, [("hello", vec![1.0, 0.0])]).await);
|
||||
let request = ResponseCacheRequest::<SemanticCacheContext>::new(CacheKeyInput {
|
||||
preset: Some("key".to_owned()),
|
||||
..Default::default()
|
||||
})
|
||||
.with_context(context("hello"));
|
||||
let response = json!({"answer": 42});
|
||||
let facade = ResponseCache::new(cache.clone());
|
||||
facade
|
||||
.async_store(&request, response.clone(), Duration::from_secs(1))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
facade
|
||||
.async_lookup(&request, Duration::from_secs(1))
|
||||
.await
|
||||
.unwrap(),
|
||||
Some(response)
|
||||
);
|
||||
{
|
||||
let mut state = server.state.lock().unwrap();
|
||||
state.points[0]
|
||||
.payload
|
||||
.insert("response".to_owned(), Value::from("not json"));
|
||||
}
|
||||
assert_eq!(
|
||||
facade
|
||||
.async_lookup(&request, Duration::from_secs(1))
|
||||
.await
|
||||
.unwrap(),
|
||||
None
|
||||
);
|
||||
server.stop();
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn stopped_qdrant_server_maps_to_unavailable() {
|
||||
let server = FakeQdrant::start(FakeState::default()).await;
|
||||
let cache = connect(&server, [("hello", vec![1.0, 0.0])]).await;
|
||||
server.stop();
|
||||
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||
assert_eq!(
|
||||
cache.async_get_cache("key", &context("hello")).await,
|
||||
Err(Error::Unavailable)
|
||||
);
|
||||
}
|
||||
342
litellm-rust/crates/cache-qdrant-semantic/tests/support/mod.rs
Normal file
342
litellm-rust/crates/cache-qdrant-semantic/tests/support/mod.rs
Normal file
|
|
@ -0,0 +1,342 @@
|
|||
use std::{
|
||||
collections::{HashMap, HashSet},
|
||||
net::SocketAddr,
|
||||
sync::{Arc, Mutex},
|
||||
};
|
||||
|
||||
use qdrant_client::qdrant::collections_server::CollectionsServer;
|
||||
use qdrant_client::qdrant::{
|
||||
self, CollectionExists, CollectionExistsRequest, CollectionExistsResponse,
|
||||
CollectionOperationResponse, CreateCollection, CreateFieldIndexCollection, Filter, PointId,
|
||||
PointsOperationResponse, ScoredPoint, SearchPoints, SearchResponse, Value, Vector, Vectors,
|
||||
collections_server::Collections,
|
||||
points_server::{Points, PointsServer},
|
||||
};
|
||||
use tokio::sync::oneshot;
|
||||
use tokio_stream::wrappers::TcpListenerStream;
|
||||
use tonic::{Request, Response, Status, transport::Server};
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct StoredPoint {
|
||||
pub id: Option<PointId>,
|
||||
pub vector: Vec<f32>,
|
||||
pub payload: HashMap<String, Value>,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
pub struct FakeState {
|
||||
pub collections: HashSet<String>,
|
||||
pub created_collections: Vec<CreateCollection>,
|
||||
pub field_indexes: Vec<CreateFieldIndexCollection>,
|
||||
pub points: Vec<StoredPoint>,
|
||||
pub upsert_waits: Vec<Option<bool>>,
|
||||
pub index_creations: usize,
|
||||
pub fail_field_index: bool,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct FakeQdrant {
|
||||
pub state: Arc<Mutex<FakeState>>,
|
||||
pub address: SocketAddr,
|
||||
shutdown: Arc<Mutex<Option<oneshot::Sender<()>>>>,
|
||||
}
|
||||
|
||||
impl FakeQdrant {
|
||||
pub async fn start(state: FakeState) -> Self {
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let address = listener.local_addr().unwrap();
|
||||
let state = Arc::new(Mutex::new(state));
|
||||
let service = FakeService {
|
||||
state: state.clone(),
|
||||
};
|
||||
let (shutdown_tx, shutdown_rx) = oneshot::channel();
|
||||
tokio::spawn(async move {
|
||||
Server::builder()
|
||||
.add_service(CollectionsServer::new(service.clone()))
|
||||
.add_service(PointsServer::new(service))
|
||||
.serve_with_incoming_shutdown(TcpListenerStream::new(listener), async {
|
||||
let _ = shutdown_rx.await;
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
});
|
||||
Self {
|
||||
state,
|
||||
address,
|
||||
shutdown: Arc::new(Mutex::new(Some(shutdown_tx))),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn url(&self) -> String {
|
||||
format!("http://{}", self.address)
|
||||
}
|
||||
|
||||
pub fn stop(&self) {
|
||||
self.shutdown
|
||||
.lock()
|
||||
.unwrap()
|
||||
.take()
|
||||
.unwrap()
|
||||
.send(())
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
pub fn insert_point(&self, point: StoredPoint) {
|
||||
self.state.lock().unwrap().points.push(point);
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct FakeService {
|
||||
state: Arc<Mutex<FakeState>>,
|
||||
}
|
||||
|
||||
macro_rules! unimplemented_collections {
|
||||
($($name:ident, $request:ty, $response:ty);* $(;)?) => {
|
||||
$(
|
||||
fn $name<'life0, 'async_trait>(
|
||||
&'life0 self,
|
||||
_: Request<$request>,
|
||||
) -> std::pin::Pin<
|
||||
Box<
|
||||
dyn std::future::Future<
|
||||
Output = Result<Response<$response>, Status>,
|
||||
> + Send
|
||||
+ 'async_trait,
|
||||
>,
|
||||
>
|
||||
where
|
||||
'life0: 'async_trait,
|
||||
Self: 'async_trait,
|
||||
{
|
||||
Box::pin(async { Err(Status::unimplemented(stringify!($name))) })
|
||||
}
|
||||
)*
|
||||
};
|
||||
}
|
||||
|
||||
macro_rules! unimplemented_points {
|
||||
($($name:ident, $request:ty, $response:ty);* $(;)?) => {
|
||||
$(
|
||||
fn $name<'life0, 'async_trait>(
|
||||
&'life0 self,
|
||||
_: Request<$request>,
|
||||
) -> std::pin::Pin<
|
||||
Box<
|
||||
dyn std::future::Future<
|
||||
Output = Result<Response<$response>, Status>,
|
||||
> + Send
|
||||
+ 'async_trait,
|
||||
>,
|
||||
>
|
||||
where
|
||||
'life0: 'async_trait,
|
||||
Self: 'async_trait,
|
||||
{
|
||||
Box::pin(async { Err(Status::unimplemented(stringify!($name))) })
|
||||
}
|
||||
)*
|
||||
};
|
||||
}
|
||||
|
||||
#[tonic::async_trait]
|
||||
impl Collections for FakeService {
|
||||
async fn create(
|
||||
&self,
|
||||
request: Request<CreateCollection>,
|
||||
) -> Result<Response<CollectionOperationResponse>, Status> {
|
||||
let request = request.into_inner();
|
||||
let mut state = self.state.lock().unwrap();
|
||||
state.collections.insert(request.collection_name.clone());
|
||||
state.created_collections.push(request);
|
||||
Ok(Response::new(CollectionOperationResponse {
|
||||
result: true,
|
||||
..Default::default()
|
||||
}))
|
||||
}
|
||||
|
||||
async fn collection_exists(
|
||||
&self,
|
||||
request: Request<CollectionExistsRequest>,
|
||||
) -> Result<Response<CollectionExistsResponse>, Status> {
|
||||
let exists = self
|
||||
.state
|
||||
.lock()
|
||||
.unwrap()
|
||||
.collections
|
||||
.contains(&request.into_inner().collection_name);
|
||||
Ok(Response::new(CollectionExistsResponse {
|
||||
result: Some(CollectionExists { exists }),
|
||||
..Default::default()
|
||||
}))
|
||||
}
|
||||
|
||||
unimplemented_collections!(
|
||||
get, qdrant::GetCollectionInfoRequest, qdrant::GetCollectionInfoResponse;
|
||||
list, qdrant::ListCollectionsRequest, qdrant::ListCollectionsResponse;
|
||||
update, qdrant::UpdateCollection, qdrant::CollectionOperationResponse;
|
||||
delete, qdrant::DeleteCollection, qdrant::CollectionOperationResponse;
|
||||
update_aliases, qdrant::ChangeAliases, qdrant::CollectionOperationResponse;
|
||||
list_collection_aliases, qdrant::ListCollectionAliasesRequest, qdrant::ListAliasesResponse;
|
||||
list_aliases, qdrant::ListAliasesRequest, qdrant::ListAliasesResponse;
|
||||
collection_cluster_info, qdrant::CollectionClusterInfoRequest, qdrant::CollectionClusterInfoResponse;
|
||||
update_collection_cluster_setup, qdrant::UpdateCollectionClusterSetupRequest, qdrant::UpdateCollectionClusterSetupResponse;
|
||||
create_shard_key, qdrant::CreateShardKeyRequest, qdrant::CreateShardKeyResponse;
|
||||
delete_shard_key, qdrant::DeleteShardKeyRequest, qdrant::DeleteShardKeyResponse;
|
||||
list_shard_keys, qdrant::ListShardKeysRequest, qdrant::ListShardKeysResponse;
|
||||
);
|
||||
}
|
||||
|
||||
#[tonic::async_trait]
|
||||
impl Points for FakeService {
|
||||
async fn create_field_index(
|
||||
&self,
|
||||
request: Request<CreateFieldIndexCollection>,
|
||||
) -> Result<Response<PointsOperationResponse>, Status> {
|
||||
let mut state = self.state.lock().unwrap();
|
||||
state.index_creations += 1;
|
||||
state.field_indexes.push(request.into_inner());
|
||||
if state.fail_field_index {
|
||||
return Err(Status::internal("field index failure"));
|
||||
}
|
||||
Ok(Response::new(PointsOperationResponse::default()))
|
||||
}
|
||||
|
||||
async fn upsert(
|
||||
&self,
|
||||
request: Request<qdrant::UpsertPoints>,
|
||||
) -> Result<Response<PointsOperationResponse>, Status> {
|
||||
let request = request.into_inner();
|
||||
let mut state = self.state.lock().unwrap();
|
||||
state.upsert_waits.push(request.wait);
|
||||
for point in request.points {
|
||||
let stored = StoredPoint {
|
||||
id: point.id.clone(),
|
||||
vector: dense_vector(point.vectors)?,
|
||||
payload: point.payload,
|
||||
};
|
||||
if let Some(existing) = state
|
||||
.points
|
||||
.iter_mut()
|
||||
.find(|existing| existing.id == stored.id)
|
||||
{
|
||||
*existing = stored;
|
||||
} else {
|
||||
state.points.push(stored);
|
||||
}
|
||||
}
|
||||
Ok(Response::new(PointsOperationResponse::default()))
|
||||
}
|
||||
|
||||
async fn search(
|
||||
&self,
|
||||
request: Request<SearchPoints>,
|
||||
) -> Result<Response<SearchResponse>, Status> {
|
||||
let request = request.into_inner();
|
||||
let key_filter = keyword_filter(request.filter.as_ref());
|
||||
let state = self.state.lock().unwrap();
|
||||
let mut results = state
|
||||
.points
|
||||
.iter()
|
||||
.filter(|point| {
|
||||
key_filter.as_ref().is_none_or(|(field, expected)| {
|
||||
point
|
||||
.payload
|
||||
.get(field)
|
||||
.and_then(|value| {
|
||||
let value: serde_json::Value = value.clone().into();
|
||||
value
|
||||
.as_str()
|
||||
.map(str::to_owned)
|
||||
.or_else(|| value.as_i64().map(|value| value.to_string()))
|
||||
})
|
||||
.is_some_and(|value| value == *expected)
|
||||
})
|
||||
})
|
||||
.map(|point| ScoredPoint {
|
||||
id: point.id.clone(),
|
||||
payload: point.payload.clone(),
|
||||
score: cosine(&request.vector, &point.vector),
|
||||
..Default::default()
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
results.sort_by(|left, right| right.score.total_cmp(&left.score));
|
||||
results.truncate(request.limit as usize);
|
||||
Ok(Response::new(SearchResponse {
|
||||
result: results,
|
||||
..Default::default()
|
||||
}))
|
||||
}
|
||||
|
||||
unimplemented_points!(
|
||||
delete, qdrant::DeletePoints, qdrant::PointsOperationResponse;
|
||||
get, qdrant::GetPoints, qdrant::GetResponse;
|
||||
update_vectors, qdrant::UpdatePointVectors, qdrant::PointsOperationResponse;
|
||||
delete_vectors, qdrant::DeletePointVectors, qdrant::PointsOperationResponse;
|
||||
set_payload, qdrant::SetPayloadPoints, qdrant::PointsOperationResponse;
|
||||
overwrite_payload, qdrant::SetPayloadPoints, qdrant::PointsOperationResponse;
|
||||
delete_payload, qdrant::DeletePayloadPoints, qdrant::PointsOperationResponse;
|
||||
clear_payload, qdrant::ClearPayloadPoints, qdrant::PointsOperationResponse;
|
||||
delete_field_index, qdrant::DeleteFieldIndexCollection, qdrant::PointsOperationResponse;
|
||||
create_vector_name, qdrant::CreateVectorNameRequest, qdrant::PointsOperationResponse;
|
||||
delete_vector_name, qdrant::DeleteVectorNameRequest, qdrant::PointsOperationResponse;
|
||||
search_batch, qdrant::SearchBatchPoints, qdrant::SearchBatchResponse;
|
||||
search_groups, qdrant::SearchPointGroups, qdrant::SearchGroupsResponse;
|
||||
scroll, qdrant::ScrollPoints, qdrant::ScrollResponse;
|
||||
recommend, qdrant::RecommendPoints, qdrant::RecommendResponse;
|
||||
recommend_batch, qdrant::RecommendBatchPoints, qdrant::RecommendBatchResponse;
|
||||
recommend_groups, qdrant::RecommendPointGroups, qdrant::RecommendGroupsResponse;
|
||||
discover, qdrant::DiscoverPoints, qdrant::DiscoverResponse;
|
||||
discover_batch, qdrant::DiscoverBatchPoints, qdrant::DiscoverBatchResponse;
|
||||
count, qdrant::CountPoints, qdrant::CountResponse;
|
||||
update_batch, qdrant::UpdateBatchPoints, qdrant::UpdateBatchResponse;
|
||||
query, qdrant::QueryPoints, qdrant::QueryResponse;
|
||||
query_batch, qdrant::QueryBatchPoints, qdrant::QueryBatchResponse;
|
||||
query_groups, qdrant::QueryPointGroups, qdrant::QueryGroupsResponse;
|
||||
facet, qdrant::FacetCounts, qdrant::FacetResponse;
|
||||
search_matrix_pairs, qdrant::SearchMatrixPoints, qdrant::SearchMatrixPairsResponse;
|
||||
search_matrix_offsets, qdrant::SearchMatrixPoints, qdrant::SearchMatrixOffsetsResponse;
|
||||
);
|
||||
}
|
||||
|
||||
fn dense_vector(vectors: Option<Vectors>) -> Result<Vec<f32>, Status> {
|
||||
let Some(Vectors {
|
||||
vectors_options:
|
||||
Some(qdrant::vectors::VectorsOptions::Vector(Vector {
|
||||
vector: Some(qdrant::vector::Vector::Dense(qdrant::DenseVector { data })),
|
||||
..
|
||||
})),
|
||||
}) = vectors
|
||||
else {
|
||||
return Err(Status::invalid_argument("expected dense vector"));
|
||||
};
|
||||
Ok(data)
|
||||
}
|
||||
|
||||
fn keyword_filter(filter: Option<&Filter>) -> Option<(String, String)> {
|
||||
filter?
|
||||
.must
|
||||
.iter()
|
||||
.find_map(|condition| match condition.condition_one_of.as_ref()? {
|
||||
qdrant::condition::ConditionOneOf::Field(field) => {
|
||||
let qdrant::r#match::MatchValue::Keyword(value) =
|
||||
field.r#match.as_ref()?.match_value.as_ref()?
|
||||
else {
|
||||
return None;
|
||||
};
|
||||
Some((field.key.clone(), value.clone()))
|
||||
}
|
||||
_ => None,
|
||||
})
|
||||
}
|
||||
|
||||
fn cosine(left: &[f32], right: &[f32]) -> f32 {
|
||||
let dot = left
|
||||
.iter()
|
||||
.zip(right)
|
||||
.map(|(left, right)| left * right)
|
||||
.sum::<f32>();
|
||||
let left_norm = left.iter().map(|value| value * value).sum::<f32>().sqrt();
|
||||
let right_norm = right.iter().map(|value| value * value).sum::<f32>().sqrt();
|
||||
dot / (left_norm * right_norm)
|
||||
}
|
||||
21
litellm-rust/crates/cache-redis-semantic/Cargo.toml
Normal file
21
litellm-rust/crates/cache-redis-semantic/Cargo.toml
Normal file
|
|
@ -0,0 +1,21 @@
|
|||
[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
|
||||
litellm-cache-response.workspace = true
|
||||
redis = { version = "1.7.0", features = ["tls-rustls"] }
|
||||
r2d2 = "0.8.10"
|
||||
serde_json.workspace = true
|
||||
sha2.workspace = true
|
||||
tokio.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
redis-test = "1.0.4"
|
||||
serde_json.workspace = true
|
||||
tokio.workspace = true
|
||||
618
litellm-rust/crates/cache-redis-semantic/src/cache.rs
Normal file
618
litellm-rust/crates/cache-redis-semantic/src/cache.rs
Normal file
|
|
@ -0,0 +1,618 @@
|
|||
use std::{
|
||||
future::Future,
|
||||
sync::{Arc, OnceLock},
|
||||
time::{Duration, SystemTime, UNIX_EPOCH},
|
||||
};
|
||||
|
||||
use litellm_cache::{
|
||||
BaseCache, CacheCodec, CacheConnectionResult, CacheConnectionStatus, Error,
|
||||
SemanticCacheContext,
|
||||
};
|
||||
use litellm_cache_redis::{
|
||||
RedisTopology,
|
||||
connection::{ConnectionRef, Connections},
|
||||
};
|
||||
use litellm_cache_response::{CacheEntry, ResponseCacheCodec};
|
||||
use serde_json::Value;
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
use crate::prompt::prompt_from_context;
|
||||
|
||||
const CACHE_KEY_FIELD: &str = "litellm_cache_key";
|
||||
const VECTOR_FIELD: &str = "prompt_vector";
|
||||
|
||||
pub trait Embedder: Send + Sync + 'static {
|
||||
fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result<Vec<f32>, Error>;
|
||||
|
||||
fn async_embed(
|
||||
&self,
|
||||
prompt: &str,
|
||||
metadata: Option<&Value>,
|
||||
) -> impl Future<Output = Result<Vec<f32>, Error>> + Send;
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct RedisSemanticConfig {
|
||||
pub index_name: String,
|
||||
pub similarity_threshold: f32,
|
||||
}
|
||||
|
||||
struct Inner {
|
||||
index_name: String,
|
||||
distance_threshold: f64,
|
||||
resolved_index: OnceLock<String>,
|
||||
codec: ResponseCacheCodec,
|
||||
clock: fn() -> f64,
|
||||
}
|
||||
|
||||
impl Inner {
|
||||
fn new(config: RedisSemanticConfig) -> Self {
|
||||
Self {
|
||||
index_name: config.index_name,
|
||||
distance_threshold: 1.0 - f64::from(config.similarity_threshold),
|
||||
resolved_index: OnceLock::new(),
|
||||
codec: ResponseCacheCodec,
|
||||
clock: timestamp,
|
||||
}
|
||||
}
|
||||
|
||||
fn ensure_index(
|
||||
&self,
|
||||
connection: &mut ConnectionRef<'_>,
|
||||
dims: usize,
|
||||
) -> Result<String, Error> {
|
||||
if let Some(name) = self.resolved_index.get() {
|
||||
return Ok(name.clone());
|
||||
}
|
||||
let name = match index_compatible(connection, &self.index_name, dims)? {
|
||||
Some(true) => self.index_name.clone(),
|
||||
Some(false) => self.isolated_index(connection, dims)?,
|
||||
None => match create_index(connection, &self.index_name, dims) {
|
||||
Ok(()) => self.index_name.clone(),
|
||||
Err(_) => match index_compatible(connection, &self.index_name, dims)? {
|
||||
Some(true) => self.index_name.clone(),
|
||||
Some(false) => self.isolated_index(connection, dims)?,
|
||||
None => return Err(Error::Unavailable),
|
||||
},
|
||||
},
|
||||
};
|
||||
let _ = self.resolved_index.set(name.clone());
|
||||
Ok(name)
|
||||
}
|
||||
|
||||
fn isolated_index(
|
||||
&self,
|
||||
connection: &mut ConnectionRef<'_>,
|
||||
dims: usize,
|
||||
) -> Result<String, Error> {
|
||||
let name = format!("{}_isolated", self.index_name);
|
||||
match index_compatible(connection, &name, dims)? {
|
||||
Some(true) => Ok(name),
|
||||
Some(false) => {
|
||||
redis::cmd("FT.DROPINDEX")
|
||||
.arg(&name)
|
||||
.query::<()>(connection)
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
create_index(connection, &name, dims)?;
|
||||
Ok(name)
|
||||
}
|
||||
None => {
|
||||
create_index(connection, &name, dims)?;
|
||||
Ok(name)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn store(
|
||||
&self,
|
||||
connection: &mut ConnectionRef<'_>,
|
||||
tag: &str,
|
||||
value: &CacheEntry,
|
||||
prompt: &str,
|
||||
vector: &[f32],
|
||||
ttl: Option<Duration>,
|
||||
) -> Result<(), Error> {
|
||||
let index = self.ensure_index(connection, vector.len())?;
|
||||
let entry_id = entry_id(prompt, tag);
|
||||
let hash_key = format!("{index}:{entry_id}");
|
||||
let response = self.codec.encode(value)?;
|
||||
redis::cmd("HSET")
|
||||
.arg(&hash_key)
|
||||
.arg("entry_id")
|
||||
.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<Option<CacheEntry>, Error> {
|
||||
let index = self.ensure_index(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(None);
|
||||
};
|
||||
if string_field(fields, CACHE_KEY_FIELD).as_deref() != Some(tag) {
|
||||
return Ok(None);
|
||||
}
|
||||
if number_field(fields, "vector_distance")
|
||||
.is_none_or(|distance| distance > self.distance_threshold)
|
||||
{
|
||||
return Ok(None);
|
||||
}
|
||||
let Some(response) = bytes_field(fields, "response") else {
|
||||
return Ok(None);
|
||||
};
|
||||
self.codec.decode(&response).map(Some)
|
||||
}
|
||||
}
|
||||
|
||||
pub struct RedisSemanticCache<E: Embedder, C = redis::Connection> {
|
||||
connections: Arc<Connections<C>>,
|
||||
embedder: E,
|
||||
inner: Arc<Inner>,
|
||||
}
|
||||
|
||||
impl<E: Embedder> RedisSemanticCache<E> {
|
||||
pub fn new(url: &str, embedder: E, config: RedisSemanticConfig) -> Result<Self, Error> {
|
||||
Ok(Self {
|
||||
connections: Arc::new(Connections::open(url, &RedisTopology::Standalone)?),
|
||||
embedder,
|
||||
inner: Arc::new(Inner::new(config)),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl<E: Embedder, C: redis::ConnectionLike + Send + 'static> RedisSemanticCache<E, C> {
|
||||
pub fn with_connection(connection: C, embedder: E, config: RedisSemanticConfig) -> Self {
|
||||
Self {
|
||||
connections: Arc::new(Connections::fixed(connection)),
|
||||
embedder,
|
||||
inner: Arc::new(Inner::new(config)),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_clock(self, clock: fn() -> f64) -> Self {
|
||||
Self {
|
||||
inner: Arc::new(Inner {
|
||||
index_name: self.inner.index_name.clone(),
|
||||
distance_threshold: self.inner.distance_threshold,
|
||||
resolved_index: OnceLock::new(),
|
||||
codec: self.inner.codec,
|
||||
clock,
|
||||
}),
|
||||
..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)
|
||||
}
|
||||
}
|
||||
|
||||
impl<E: Embedder, C: redis::ConnectionLike + Send + 'static> BaseCache
|
||||
for RedisSemanticCache<E, C>
|
||||
{
|
||||
type Value = CacheEntry;
|
||||
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 vector = self.embedder.embed(&prompt, context.metadata.as_ref())?;
|
||||
let tag = Self::tag(key, context).to_string();
|
||||
self.connections.execute(|connection| {
|
||||
self.inner
|
||||
.store(connection, &tag, &value, &prompt, &vector, context.ttl)
|
||||
})
|
||||
}
|
||||
|
||||
fn get_cache(&self, key: &str, context: &Self::Context) -> Result<Option<Self::Value>, Error> {
|
||||
let Some(prompt) = prompt_from_context(context) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let vector = self.embedder.embed(&prompt, context.metadata.as_ref())?;
|
||||
let tag = Self::tag(key, context).to_string();
|
||||
self.connections
|
||||
.execute(|connection| self.inner.lookup(connection, &tag, &vector))
|
||||
}
|
||||
|
||||
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 vector = self
|
||||
.embedder
|
||||
.async_embed(&prompt, context.metadata.as_ref())
|
||||
.await?;
|
||||
let tag = Self::tag(key, &context).to_string();
|
||||
let inner = Arc::clone(&self.inner);
|
||||
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
|
||||
inner.store(connection, &tag, &value, &prompt, &vector, context.ttl)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
async fn async_get_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
context: &Self::Context,
|
||||
) -> Result<Option<Self::Value>, Error> {
|
||||
let Some(prompt) = prompt_from_context(context) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let vector = self
|
||||
.embedder
|
||||
.async_embed(&prompt, context.metadata.as_ref())
|
||||
.await?;
|
||||
let tag = Self::tag(key, context).to_string();
|
||||
let inner = Arc::clone(&self.inner);
|
||||
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
|
||||
inner.lookup(connection, &tag, &vector)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
async fn disconnect(&self) -> Result<(), Error> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
|
||||
match Connections::run_blocking(Arc::clone(&self.connections), |connection| {
|
||||
Ok(match redis::cmd("PING").query::<String>(connection) {
|
||||
Ok(_) => CacheConnectionResult {
|
||||
status: CacheConnectionStatus::Success,
|
||||
message: "Redis cache connection test successful".into(),
|
||||
error: None,
|
||||
},
|
||||
Err(error) => CacheConnectionResult {
|
||||
status: CacheConnectionStatus::Failed,
|
||||
message: format!("Redis connection failed: {error}"),
|
||||
error: Some(error.to_string()),
|
||||
},
|
||||
})
|
||||
})
|
||||
.await
|
||||
{
|
||||
Ok(result) => Ok(result),
|
||||
Err(error) => Ok(CacheConnectionResult {
|
||||
status: CacheConnectionStatus::Failed,
|
||||
message: format!("Redis connection failed: {error}"),
|
||||
error: Some(error.to_string()),
|
||||
}),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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 {
|
||||
value
|
||||
.chars()
|
||||
.flat_map(|ch| {
|
||||
if matches!(
|
||||
ch,
|
||||
',' | '.'
|
||||
| '<'
|
||||
| '>'
|
||||
| '{'
|
||||
| '}'
|
||||
| '['
|
||||
| ']'
|
||||
| '\\'
|
||||
| '"'
|
||||
| '\''
|
||||
| ':'
|
||||
| ';'
|
||||
| '!'
|
||||
| '@'
|
||||
| '#'
|
||||
| '$'
|
||||
| '%'
|
||||
| '^'
|
||||
| '&'
|
||||
| '*'
|
||||
| '('
|
||||
| ')'
|
||||
| '-'
|
||||
| '+'
|
||||
| '='
|
||||
| '~'
|
||||
| '|'
|
||||
| '/'
|
||||
| ' '
|
||||
| '?'
|
||||
) {
|
||||
vec!['\\', ch]
|
||||
} else {
|
||||
vec![ch]
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn create_index(connection: &mut ConnectionRef<'_>, name: &str, dims: usize) -> Result<(), Error> {
|
||||
redis::cmd("FT.CREATE")
|
||||
.arg(name)
|
||||
.arg("ON")
|
||||
.arg("HASH")
|
||||
.arg("PREFIX")
|
||||
.arg(1)
|
||||
.arg(name)
|
||||
.arg("SCORE")
|
||||
.arg(1.0)
|
||||
.arg("SCHEMA")
|
||||
.arg("prompt")
|
||||
.arg("TEXT")
|
||||
.arg("WEIGHT")
|
||||
.arg(1)
|
||||
.arg("response")
|
||||
.arg("TEXT")
|
||||
.arg("WEIGHT")
|
||||
.arg(1)
|
||||
.arg("inserted_at")
|
||||
.arg("NUMERIC")
|
||||
.arg("updated_at")
|
||||
.arg("NUMERIC")
|
||||
.arg(VECTOR_FIELD)
|
||||
.arg("VECTOR")
|
||||
.arg("FLAT")
|
||||
.arg(6)
|
||||
.arg("TYPE")
|
||||
.arg("FLOAT32")
|
||||
.arg("DIM")
|
||||
.arg(dims)
|
||||
.arg("DISTANCE_METRIC")
|
||||
.arg("COSINE")
|
||||
.arg(CACHE_KEY_FIELD)
|
||||
.arg("TAG")
|
||||
.arg("SEPARATOR")
|
||||
.arg(",")
|
||||
.query::<()>(connection)
|
||||
.map_err(|_| Error::Unavailable)
|
||||
}
|
||||
|
||||
fn index_compatible(
|
||||
connection: &mut ConnectionRef<'_>,
|
||||
name: &str,
|
||||
dims: usize,
|
||||
) -> Result<Option<bool>, Error> {
|
||||
let info = match redis::cmd("FT.INFO")
|
||||
.arg(name)
|
||||
.query::<redis::Value>(connection)
|
||||
{
|
||||
Ok(info) => info,
|
||||
Err(error) if unknown_index(&error) => return Ok(None),
|
||||
Err(_) => return Err(Error::Unavailable),
|
||||
};
|
||||
Ok(Some(schema_compatible(&info, dims)))
|
||||
}
|
||||
|
||||
fn unknown_index(error: &redis::RedisError) -> bool {
|
||||
let message = error.to_string().to_lowercase();
|
||||
message.contains("unknown") && message.contains("index")
|
||||
}
|
||||
|
||||
fn schema_compatible(info: &redis::Value, dims: usize) -> bool {
|
||||
let redis::Value::Array(entries) = info else {
|
||||
return false;
|
||||
};
|
||||
let attributes = entries
|
||||
.as_chunks::<2>()
|
||||
.0
|
||||
.iter()
|
||||
.find(|pair| string_value(&pair[0]).as_deref() == Some("attributes"))
|
||||
.map(|pair| &pair[1]);
|
||||
let Some(redis::Value::Array(attributes)) = attributes else {
|
||||
return false;
|
||||
};
|
||||
let fields = attributes
|
||||
.iter()
|
||||
.map(|attribute| {
|
||||
let redis::Value::Array(attribute) = attribute else {
|
||||
return (None, None, None, None, None);
|
||||
};
|
||||
let mut name = None;
|
||||
let mut field_type = None;
|
||||
let mut dim = None;
|
||||
let mut data_type = None;
|
||||
let mut distance_metric = None;
|
||||
for pair in attribute.as_chunks::<2>().0 {
|
||||
match string_value(&pair[0]).as_deref() {
|
||||
Some("identifier") => name = string_value(&pair[1]),
|
||||
Some("type") => field_type = string_value(&pair[1]),
|
||||
Some("dim") => dim = number_value(&pair[1]),
|
||||
Some("data_type") => data_type = string_value(&pair[1]),
|
||||
Some("distance_metric") => distance_metric = string_value(&pair[1]),
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
(name, field_type, dim, data_type, distance_metric)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let has_field = |name: &str, field_type: &str| {
|
||||
fields
|
||||
.iter()
|
||||
.any(|(n, t, ..)| n.as_deref() == Some(name) && t.as_deref() == Some(field_type))
|
||||
};
|
||||
has_field("prompt", "TEXT")
|
||||
&& has_field("response", "TEXT")
|
||||
&& has_field("inserted_at", "NUMERIC")
|
||||
&& has_field("updated_at", "NUMERIC")
|
||||
&& has_field(CACHE_KEY_FIELD, "TAG")
|
||||
&& fields.iter().any(|(n, t, d, data, metric)| {
|
||||
n.as_deref() == Some(VECTOR_FIELD)
|
||||
&& t.as_deref() == Some("VECTOR")
|
||||
&& *d == Some(dims as f64)
|
||||
&& data
|
||||
.as_deref()
|
||||
.is_some_and(|data| data.eq_ignore_ascii_case("float32"))
|
||||
&& metric
|
||||
.as_deref()
|
||||
.is_some_and(|metric| metric.eq_ignore_ascii_case("cosine"))
|
||||
})
|
||||
}
|
||||
|
||||
fn string_value(value: &redis::Value) -> Option<String> {
|
||||
match value {
|
||||
redis::Value::BulkString(bytes) => String::from_utf8(bytes.clone()).ok(),
|
||||
redis::Value::SimpleString(text) => Some(text.clone()),
|
||||
redis::Value::VerbatimString { text, .. } => Some(text.clone()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn number_value(value: &redis::Value) -> Option<f64> {
|
||||
match value {
|
||||
redis::Value::Int(number) => Some(*number as f64),
|
||||
redis::Value::Double(number) => Some(*number),
|
||||
_ => string_value(value).and_then(|text| text.parse().ok()),
|
||||
}
|
||||
}
|
||||
|
||||
fn first_document(result: &redis::Value) -> Option<&[redis::Value]> {
|
||||
let redis::Value::Array(items) = result else {
|
||||
return None;
|
||||
};
|
||||
let [count, _document_id, fields, ..] = items.as_slice() else {
|
||||
return None;
|
||||
};
|
||||
if !matches!(count, redis::Value::Int(count) if *count > 0) {
|
||||
return None;
|
||||
}
|
||||
match fields {
|
||||
redis::Value::Array(fields) => Some(fields.as_slice()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn field_value<'a>(fields: &'a [redis::Value], name: &str) -> Option<&'a redis::Value> {
|
||||
fields
|
||||
.as_chunks::<2>()
|
||||
.0
|
||||
.iter()
|
||||
.find(|pair| string_value(&pair[0]).as_deref() == Some(name))
|
||||
.map(|pair| &pair[1])
|
||||
}
|
||||
|
||||
fn string_field(fields: &[redis::Value], name: &str) -> Option<String> {
|
||||
field_value(fields, name).and_then(string_value)
|
||||
}
|
||||
|
||||
fn number_field(fields: &[redis::Value], name: &str) -> Option<f64> {
|
||||
field_value(fields, name).and_then(number_value)
|
||||
}
|
||||
|
||||
fn bytes_field(fields: &[redis::Value], name: &str) -> Option<Vec<u8>> {
|
||||
match field_value(fields, name)? {
|
||||
redis::Value::BulkString(bytes) => Some(bytes.clone()),
|
||||
redis::Value::SimpleString(text) => Some(text.clone().into_bytes()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn ttl_seconds(ttl: Duration) -> u64 {
|
||||
ttl.as_secs()
|
||||
.saturating_add(u64::from(ttl.subsec_nanos() > 0))
|
||||
.max(1)
|
||||
}
|
||||
5
litellm-rust/crates/cache-redis-semantic/src/lib.rs
Normal file
5
litellm-rust/crates/cache-redis-semantic/src/lib.rs
Normal file
|
|
@ -0,0 +1,5 @@
|
|||
mod cache;
|
||||
mod prompt;
|
||||
|
||||
pub use cache::{Embedder, RedisSemanticCache, RedisSemanticConfig};
|
||||
pub use prompt::prompt_from_context;
|
||||
97
litellm-rust/crates/cache-redis-semantic/src/prompt.rs
Normal file
97
litellm-rust/crates/cache-redis-semantic/src/prompt.rs
Normal file
|
|
@ -0,0 +1,97 @@
|
|||
use litellm_cache::SemanticCacheContext;
|
||||
use serde_json::Value;
|
||||
|
||||
pub fn prompt_from_context(context: &SemanticCacheContext) -> Option<String> {
|
||||
if let Some(messages) = context.messages.as_ref().and_then(Value::as_array)
|
||||
&& !messages.is_empty()
|
||||
{
|
||||
return Some(messages_text(messages));
|
||||
}
|
||||
let input = context.input.as_ref()?;
|
||||
let mut parts = Vec::new();
|
||||
collect_input_text(input, &mut parts);
|
||||
let prompt = parts.join("\n").trim().to_string();
|
||||
(!prompt.is_empty()).then_some(prompt)
|
||||
}
|
||||
|
||||
fn messages_text(messages: &[Value]) -> String {
|
||||
let mut text = String::new();
|
||||
for message in messages {
|
||||
let Some(message) = message.as_object() else {
|
||||
continue;
|
||||
};
|
||||
match message.get("content") {
|
||||
Some(Value::String(content)) => text.push_str(content),
|
||||
Some(Value::Array(parts)) => {
|
||||
for part in parts {
|
||||
if let Some(text_content) = part.get("text").and_then(Value::as_str) {
|
||||
text.push_str(text_content);
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
text.push_str(&search_results_text(message.get("search_results")));
|
||||
}
|
||||
text
|
||||
}
|
||||
|
||||
fn search_results_text(search_results: Option<&Value>) -> String {
|
||||
let Some(Value::Array(results)) = search_results else {
|
||||
return String::new();
|
||||
};
|
||||
let mut text = String::new();
|
||||
for result in results {
|
||||
let Some(result) = result.as_object() else {
|
||||
continue;
|
||||
};
|
||||
for key in ["source", "title"] {
|
||||
if let Some(value) = result.get(key).and_then(Value::as_str) {
|
||||
text.push_str(value);
|
||||
}
|
||||
}
|
||||
if let Some(Value::Array(content)) = result.get("content") {
|
||||
for block in content {
|
||||
if let Some(value) = block.get("text").and_then(Value::as_str) {
|
||||
text.push_str(value);
|
||||
}
|
||||
}
|
||||
}
|
||||
if let Some(citations) = result.get("citations") {
|
||||
text.push_str(&citations.to_string());
|
||||
}
|
||||
}
|
||||
text
|
||||
}
|
||||
|
||||
fn collect_input_text(value: &Value, parts: &mut Vec<String>) {
|
||||
match value {
|
||||
Value::String(text) => {
|
||||
let trimmed = text.trim();
|
||||
if !trimmed.is_empty() {
|
||||
parts.push(trimmed.to_string());
|
||||
}
|
||||
}
|
||||
Value::Array(items) => {
|
||||
for item in items {
|
||||
collect_input_text(item, parts);
|
||||
}
|
||||
}
|
||||
Value::Object(map) => {
|
||||
if let Some(content) = map.get("content").filter(|content| !content.is_null()) {
|
||||
collect_input_text(content, parts);
|
||||
return;
|
||||
}
|
||||
for key in ["text", "output", "input_text", "output_text"] {
|
||||
if let Some(Value::String(text)) = map.get(key) {
|
||||
let trimmed = text.trim();
|
||||
if !trimmed.is_empty() {
|
||||
parts.push(trimmed.to_string());
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
1003
litellm-rust/crates/cache-redis-semantic/tests/cache.rs
Normal file
1003
litellm-rust/crates/cache-redis-semantic/tests/cache.rs
Normal file
File diff suppressed because it is too large
Load diff
|
|
@ -12,7 +12,7 @@ use redis::{
|
|||
use super::REDIS_TIMEOUT;
|
||||
use crate::topology::RedisNode;
|
||||
|
||||
pub(super) struct PooledConnection<C> {
|
||||
pub struct PooledConnection<C> {
|
||||
pub(super) connection: C,
|
||||
pub(super) failed: bool,
|
||||
}
|
||||
|
|
@ -20,7 +20,7 @@ pub(super) struct PooledConnection<C> {
|
|||
/// Pools connections without a checkout PING, which would double every operation's round trips.
|
||||
/// A timed-out command leaves its reply on the socket while redis still reports the connection
|
||||
/// open, so any connection whose operation failed is discarded instead of being reused.
|
||||
pub(super) struct ConnectionManager(redis::Client);
|
||||
pub struct ConnectionManager(redis::Client);
|
||||
|
||||
impl ConnectionManager {
|
||||
pub(super) fn open(url: &str) -> Result<Self, Error> {
|
||||
|
|
@ -54,7 +54,7 @@ impl r2d2::ManageConnection for ConnectionManager {
|
|||
}
|
||||
}
|
||||
|
||||
pub(super) struct ClusterConnectionManager(ClusterClient);
|
||||
pub struct ClusterConnectionManager(ClusterClient);
|
||||
|
||||
impl ClusterConnectionManager {
|
||||
pub(super) fn open(url: &str, startup_nodes: &[RedisNode]) -> Result<Self, Error> {
|
||||
|
|
|
|||
|
|
@ -36,6 +36,8 @@ Callers supply Unix time for response freshness. Backend TTL uses its own clock.
|
|||
|
||||
The extension keeps a private test harness for memory and Redis single and batch response lookup and storage. Batch lookup returns ordered values plus missing indices for embedding partial-hit wiring. No bridge-only cache type is part of the public API
|
||||
|
||||
The bridge also exposes a production-shaped response cache runtime selected through the Rust catalog. Its shipped rule set is empty, so current SDK, Router, and proxy calls stay on Python and do not construct native cache resources. Tests can inject a rule and build the native memory runtime from an ordinary Python `Cache` configuration without changing the legacy cache classes
|
||||
|
||||
Object responses are written as they are, and every other response shape is written as a serialized string, which is the pair of shapes Python reads. A string on the wire is therefore always a serialized response, so string-valued responses round trip. Typed backends such as memory never pass through the codec
|
||||
|
||||
The resolver reads the namespace's `cache` attribute each time it resolves. A captured binding retains the selected service for its operation, including background writes. `None` disables caching. Custom Python cache objects keep their original methods, arguments, returned awaitables, exceptions, and caller-task execution
|
||||
|
|
|
|||
|
|
@ -1,9 +1,9 @@
|
|||
use std::{sync::Mutex, time::Duration};
|
||||
|
||||
use litellm_cache::{BaseCache, Error, ExactCacheContext};
|
||||
use litellm_cache::Error;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{CacheEntry, ResponseCache, ResponseCacheRequest};
|
||||
use crate::{ExactResponseCache, ResponseCacheRequest};
|
||||
|
||||
pub struct WriteBuffer {
|
||||
flush_size: usize,
|
||||
|
|
@ -18,9 +18,9 @@ impl WriteBuffer {
|
|||
}
|
||||
}
|
||||
|
||||
pub async fn async_store<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>>(
|
||||
pub async fn async_store(
|
||||
&self,
|
||||
cache: &ResponseCache<B>,
|
||||
cache: &dyn ExactResponseCache,
|
||||
request: &ResponseCacheRequest,
|
||||
response: Value,
|
||||
now: Duration,
|
||||
|
|
|
|||
148
litellm-rust/crates/cache-response/src/exact.rs
Normal file
148
litellm-rust/crates/cache-response/src/exact.rs
Normal file
|
|
@ -0,0 +1,148 @@
|
|||
use std::{future::Future, pin::Pin, time::Duration};
|
||||
|
||||
use litellm_cache::{
|
||||
BaseCache, BatchCache, CacheConnectionResult, Error, ExactCacheContext, FlushCache,
|
||||
};
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{CacheEntry, PartialHits, ResponseCache, ResponseCacheRequest};
|
||||
|
||||
type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
|
||||
|
||||
/// Object-safe view of a `ResponseCache` over an exact-match backend, so hosts can hold every
|
||||
/// exact backend behind one pointer without erasing which backend it is elsewhere.
|
||||
pub trait ExactResponseCache: Send + Sync {
|
||||
fn default_ttl(&self) -> Option<Duration>;
|
||||
|
||||
fn lookup(&self, request: &ResponseCacheRequest, now: Duration)
|
||||
-> Result<Option<Value>, Error>;
|
||||
|
||||
fn store(
|
||||
&self,
|
||||
request: &ResponseCacheRequest,
|
||||
response: Value,
|
||||
now: Duration,
|
||||
) -> Result<(), Error>;
|
||||
|
||||
fn lookup_batch(
|
||||
&self,
|
||||
requests: &[ResponseCacheRequest],
|
||||
now: Duration,
|
||||
) -> Result<PartialHits, Error>;
|
||||
|
||||
fn async_lookup<'a>(
|
||||
&'a self,
|
||||
request: &'a ResponseCacheRequest,
|
||||
now: Duration,
|
||||
) -> BoxFuture<'a, Result<Option<Value>, Error>>;
|
||||
|
||||
fn async_store<'a>(
|
||||
&'a self,
|
||||
request: &'a ResponseCacheRequest,
|
||||
response: Value,
|
||||
now: Duration,
|
||||
) -> BoxFuture<'a, Result<(), Error>>;
|
||||
|
||||
fn async_lookup_batch<'a>(
|
||||
&'a self,
|
||||
requests: &'a [ResponseCacheRequest],
|
||||
now: Duration,
|
||||
) -> BoxFuture<'a, Result<PartialHits, Error>>;
|
||||
|
||||
fn async_store_batch<'a>(
|
||||
&'a self,
|
||||
entries: Vec<(ResponseCacheRequest, Value)>,
|
||||
now: Duration,
|
||||
) -> BoxFuture<'a, Result<(), Error>>;
|
||||
|
||||
fn async_store_entries<'a>(
|
||||
&'a self,
|
||||
entries: Vec<(ResponseCacheRequest, Value, Duration)>,
|
||||
) -> BoxFuture<'a, Result<(), Error>>;
|
||||
|
||||
fn async_flush<'a>(&'a self) -> BoxFuture<'a, Result<(), Error>>;
|
||||
|
||||
fn test_connection<'a>(&'a self) -> BoxFuture<'a, Result<CacheConnectionResult, Error>>;
|
||||
}
|
||||
|
||||
impl<B> ExactResponseCache for ResponseCache<B>
|
||||
where
|
||||
B: BaseCache<Value = CacheEntry, Context = ExactCacheContext> + BatchCache + FlushCache,
|
||||
{
|
||||
fn default_ttl(&self) -> Option<Duration> {
|
||||
ResponseCache::default_ttl(self)
|
||||
}
|
||||
|
||||
fn lookup(
|
||||
&self,
|
||||
request: &ResponseCacheRequest,
|
||||
now: Duration,
|
||||
) -> Result<Option<Value>, Error> {
|
||||
ResponseCache::lookup(self, request, now)
|
||||
}
|
||||
|
||||
fn store(
|
||||
&self,
|
||||
request: &ResponseCacheRequest,
|
||||
response: Value,
|
||||
now: Duration,
|
||||
) -> Result<(), Error> {
|
||||
ResponseCache::store(self, request, response, now)
|
||||
}
|
||||
|
||||
fn lookup_batch(
|
||||
&self,
|
||||
requests: &[ResponseCacheRequest],
|
||||
now: Duration,
|
||||
) -> Result<PartialHits, Error> {
|
||||
ResponseCache::lookup_batch(self, requests, now)
|
||||
}
|
||||
|
||||
fn async_lookup<'a>(
|
||||
&'a self,
|
||||
request: &'a ResponseCacheRequest,
|
||||
now: Duration,
|
||||
) -> BoxFuture<'a, Result<Option<Value>, Error>> {
|
||||
Box::pin(ResponseCache::async_lookup(self, request, now))
|
||||
}
|
||||
|
||||
fn async_store<'a>(
|
||||
&'a self,
|
||||
request: &'a ResponseCacheRequest,
|
||||
response: Value,
|
||||
now: Duration,
|
||||
) -> BoxFuture<'a, Result<(), Error>> {
|
||||
Box::pin(ResponseCache::async_store(self, request, response, now))
|
||||
}
|
||||
|
||||
fn async_lookup_batch<'a>(
|
||||
&'a self,
|
||||
requests: &'a [ResponseCacheRequest],
|
||||
now: Duration,
|
||||
) -> BoxFuture<'a, Result<PartialHits, Error>> {
|
||||
Box::pin(ResponseCache::async_lookup_batch(self, requests, now))
|
||||
}
|
||||
|
||||
fn async_store_batch<'a>(
|
||||
&'a self,
|
||||
entries: Vec<(ResponseCacheRequest, Value)>,
|
||||
now: Duration,
|
||||
) -> BoxFuture<'a, Result<(), Error>> {
|
||||
Box::pin(ResponseCache::async_store_batch(self, entries, now))
|
||||
}
|
||||
|
||||
fn async_store_entries<'a>(
|
||||
&'a self,
|
||||
entries: Vec<(ResponseCacheRequest, Value, Duration)>,
|
||||
) -> BoxFuture<'a, Result<(), Error>> {
|
||||
Box::pin(ResponseCache::async_store_entries(self, entries))
|
||||
}
|
||||
|
||||
fn async_flush<'a>(&'a self) -> BoxFuture<'a, Result<(), Error>> {
|
||||
Box::pin(ResponseCache::async_flush(self))
|
||||
}
|
||||
|
||||
fn test_connection<'a>(&'a self) -> BoxFuture<'a, Result<CacheConnectionResult, Error>> {
|
||||
Box::pin(ResponseCache::test_connection(self))
|
||||
}
|
||||
}
|
||||
|
|
@ -2,6 +2,7 @@ mod buffer;
|
|||
mod caching;
|
||||
mod codec;
|
||||
mod embedding;
|
||||
mod exact;
|
||||
mod response;
|
||||
|
||||
pub use buffer::WriteBuffer;
|
||||
|
|
@ -11,4 +12,5 @@ pub use caching::{
|
|||
};
|
||||
pub use codec::ResponseCacheCodec;
|
||||
pub use embedding::PartialHits;
|
||||
pub use exact::ExactResponseCache;
|
||||
pub use response::{ResponseCache, ResponseCacheRequest};
|
||||
|
|
|
|||
|
|
@ -32,6 +32,17 @@ impl<C: CacheContext + Default> ResponseCacheRequest<C> {
|
|||
}
|
||||
}
|
||||
|
||||
impl<C: CacheContext> ResponseCacheRequest<C> {
|
||||
pub fn with_context<D: CacheContext>(self, context: D) -> ResponseCacheRequest<D> {
|
||||
ResponseCacheRequest {
|
||||
key: self.key,
|
||||
controls: self.controls,
|
||||
context,
|
||||
max_age: self.max_age,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct ResponseCache<B: BaseCache<Value = CacheEntry>>
|
||||
where
|
||||
B::Context: Default + PartialEq,
|
||||
|
|
|
|||
|
|
@ -1,12 +1,15 @@
|
|||
use std::{
|
||||
sync::{
|
||||
Arc,
|
||||
Arc, Mutex,
|
||||
atomic::{AtomicU64, Ordering},
|
||||
},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use litellm_cache::{BaseCache, CacheCodec, Error};
|
||||
use litellm_cache::{
|
||||
BaseCache, CacheCodec, CacheConnectionResult, CacheConnectionStatus, Error,
|
||||
SemanticCacheContext,
|
||||
};
|
||||
use litellm_cache_memory::InMemoryCache;
|
||||
use litellm_cache_redis::RedisCache;
|
||||
use litellm_cache_response::{
|
||||
|
|
@ -30,6 +33,82 @@ fn request() -> ResponseCacheRequest {
|
|||
})
|
||||
}
|
||||
|
||||
struct SemanticBackend {
|
||||
entries: Mutex<Vec<(String, CacheEntry)>>,
|
||||
contexts: Mutex<Vec<SemanticCacheContext>>,
|
||||
}
|
||||
|
||||
impl BaseCache for SemanticBackend {
|
||||
type Value = CacheEntry;
|
||||
type Context = SemanticCacheContext;
|
||||
|
||||
fn get_ttl(&self, _: &Self::Context) -> Option<Duration> {
|
||||
None
|
||||
}
|
||||
|
||||
fn set_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
value: Self::Value,
|
||||
context: &Self::Context,
|
||||
) -> Result<(), Error> {
|
||||
self.contexts.lock().unwrap().push(context.clone());
|
||||
self.entries.lock().unwrap().push((key.to_owned(), value));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn get_cache(&self, key: &str, context: &Self::Context) -> Result<Option<Self::Value>, Error> {
|
||||
self.contexts.lock().unwrap().push(context.clone());
|
||||
Ok(self
|
||||
.entries
|
||||
.lock()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.find(|(entry_key, _)| entry_key == key)
|
||||
.map(|(_, entry)| entry.clone()))
|
||||
}
|
||||
|
||||
async fn disconnect(&self) -> Result<(), Error> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
|
||||
Ok(CacheConnectionResult {
|
||||
status: CacheConnectionStatus::Success,
|
||||
message: "ok".into(),
|
||||
error: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn semantic_context_reaches_backend_for_store_and_lookup() {
|
||||
let backend = Arc::new(SemanticBackend {
|
||||
entries: Mutex::new(Vec::new()),
|
||||
contexts: Mutex::new(Vec::new()),
|
||||
});
|
||||
let cache = ResponseCache::new(backend.clone());
|
||||
let context = SemanticCacheContext {
|
||||
messages: Some(json!([{"role": "user", "content": "hello"}])),
|
||||
..Default::default()
|
||||
};
|
||||
let request = request().with_context(context.clone());
|
||||
let response = json!({"answer": 42});
|
||||
|
||||
cache
|
||||
.store(&request, response.clone(), Duration::from_secs(100))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
cache.lookup(&request, Duration::from_secs(100)).unwrap(),
|
||||
Some(response)
|
||||
);
|
||||
assert_eq!(
|
||||
backend.contexts.lock().unwrap().as_slice(),
|
||||
&[context.clone(), context]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sync_and_async_consumers_share_keys_ttls_and_freshness() {
|
||||
let clock = Arc::new(AtomicU64::new(100));
|
||||
|
|
|
|||
2
litellm-rust/crates/cache/src/error.rs
vendored
2
litellm-rust/crates/cache/src/error.rs
vendored
|
|
@ -8,4 +8,6 @@ pub enum Error {
|
|||
UnscopedFlush,
|
||||
#[error("operation is not supported by this cache")]
|
||||
UnsupportedOperation,
|
||||
#[error("semantic cache requires request messages")]
|
||||
MissingPrompt,
|
||||
}
|
||||
|
|
|
|||
21
litellm-rust/crates/cache/tests/caching.rs
vendored
21
litellm-rust/crates/cache/tests/caching.rs
vendored
|
|
@ -1,7 +1,8 @@
|
|||
use std::{sync::Mutex, time::Duration};
|
||||
|
||||
use litellm_cache::{
|
||||
BaseCache, CacheConnectionResult, CacheContext, Error, ExactCacheContext, get_cache,
|
||||
BaseCache, CacheConnectionResult, CacheContext, Error, ExactCacheContext, SemanticCacheContext,
|
||||
get_cache,
|
||||
};
|
||||
|
||||
struct TestCache {
|
||||
|
|
@ -126,6 +127,24 @@ fn associated_context_preserves_backend_specific_lookup_inputs() {
|
|||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn semantic_context_with_ttl_preserves_lookup_inputs() {
|
||||
let context = SemanticCacheContext {
|
||||
input: Some(serde_json::json!("text")),
|
||||
messages: Some(serde_json::json!([{"role": "user", "content": "hi"}])),
|
||||
metadata: Some(serde_json::json!({"key": "value"})),
|
||||
scope: Some("scope".into()),
|
||||
ttl: None,
|
||||
};
|
||||
let updated = context.with_ttl(Some(Duration::from_secs(30)));
|
||||
assert_eq!(updated.ttl(), Some(Duration::from_secs(30)));
|
||||
assert_eq!(updated.input, context.input);
|
||||
assert_eq!(updated.messages, context.messages);
|
||||
assert_eq!(updated.metadata, context.metadata);
|
||||
assert_eq!(updated.scope, context.scope);
|
||||
assert_eq!(context.with_ttl(None).ttl(), None);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn default_batch_operations_use_async_writes_and_stop_on_failure() {
|
||||
let cache = TestCache {
|
||||
|
|
|
|||
|
|
@ -17,6 +17,11 @@ pub fn parse_str_bool(value: &str) -> Option<bool> {
|
|||
token.eq_ignore_ascii_case("false").then_some(false)
|
||||
}
|
||||
|
||||
/// `redis-py` string Booleans: only `1`, `true`, and `yes` (case-insensitive) are true.
|
||||
pub fn parse_redis_bool(value: &str) -> bool {
|
||||
value == "1" || value.eq_ignore_ascii_case("true") || value.eq_ignore_ascii_case("yes")
|
||||
}
|
||||
|
||||
impl<'de> DeserializeAs<'de, i64> for LaxI64 {
|
||||
fn deserialize_as<D: Deserializer<'de>>(deserializer: D) -> Result<i64, D::Error> {
|
||||
deserializer.deserialize_any(Self)
|
||||
|
|
|
|||
|
|
@ -36,6 +36,7 @@ url.workspace = true
|
|||
veil.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
litellm-secrets.workspace = true
|
||||
litellm-auth-gcp.workspace = true
|
||||
litellm-llms = { workspace = true, features = ["test-support"] }
|
||||
rstest.workspace = true
|
||||
|
|
|
|||
|
|
@ -22,7 +22,12 @@ pub(crate) async fn perform_ocr_request(
|
|||
) -> Result<LiteLLMOcrResponse, Error> {
|
||||
request.response_format()?;
|
||||
let config = request.config;
|
||||
let request = prepare_request(request, caller_document, client);
|
||||
let secrets = client
|
||||
.secret_source()
|
||||
.resolve(&config.secret_names())
|
||||
.await
|
||||
.map_err(|error| Error::Secret(std::sync::Arc::new(error)))?;
|
||||
let request = prepare_request(request, caller_document, client, secrets);
|
||||
let hooks = OcrCallHooks::new(host.clone(), &request, config);
|
||||
config.ocr(client, &request, &hooks).await
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,7 +1,10 @@
|
|||
use litellm_auth::{InputSource, SecretValue, Sourced};
|
||||
use litellm_llms::base_llm::ocr::{
|
||||
handler::OcrClient,
|
||||
transformation::{OcrConnection, OcrCredentialInputs, PreparedOcrRequest},
|
||||
use litellm_llms::base_llm::{
|
||||
inference::secrets::Secrets,
|
||||
ocr::{
|
||||
handler::OcrClient,
|
||||
transformation::{OcrConnection, OcrCredentialInputs, PreparedOcrRequest},
|
||||
},
|
||||
};
|
||||
|
||||
use super::provider_config::OcrProvider;
|
||||
|
|
@ -11,6 +14,7 @@ pub(crate) fn prepare_request(
|
|||
request: ResolvedOcrRequest,
|
||||
caller_document: bool,
|
||||
client: &OcrClient,
|
||||
secrets: Secrets,
|
||||
) -> PreparedOcrRequest {
|
||||
let credentials = request.credentials.clone();
|
||||
let (preferred_api_key_env, api_base_env) = match request.config.provider() {
|
||||
|
|
@ -24,7 +28,7 @@ pub(crate) fn prepare_request(
|
|||
| OcrProvider::Reducto
|
||||
| OcrProvider::VertexAi => (None, None),
|
||||
};
|
||||
let secret = |name: &str| client.secrets().truthy(name);
|
||||
let secret = |name: &str| secrets.truthy(name);
|
||||
let dynamic_api_key = credentials.dynamic_api_key.or_else(|| {
|
||||
credentials.api_key.clone().or_else(|| {
|
||||
preferred_api_key_env
|
||||
|
|
@ -60,12 +64,7 @@ pub(crate) fn prepare_request(
|
|||
PreparedOcrRequest {
|
||||
model,
|
||||
document,
|
||||
connection: OcrConnection::new(
|
||||
resolved,
|
||||
transport,
|
||||
client.settings().clone(),
|
||||
client.secrets().clone(),
|
||||
),
|
||||
connection: OcrConnection::new(resolved, transport, client.settings().clone(), secrets),
|
||||
caller_document,
|
||||
optional_params,
|
||||
input_sources,
|
||||
|
|
@ -79,6 +78,7 @@ pub(crate) fn prepare_request_for_test(request: ResolvedOcrRequest) -> PreparedO
|
|||
request,
|
||||
true,
|
||||
&OcrClient::for_test(reqwest::Client::new(), reqwest::Client::new()),
|
||||
std::sync::Arc::new(litellm_core_utils::settings::ProcessEnvironment),
|
||||
)
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -114,6 +114,10 @@ impl OcrConfigKind {
|
|||
with_config!(self, config => config.get_api_key_env_var())
|
||||
}
|
||||
|
||||
pub(crate) fn secret_names(self) -> Vec<&'static str> {
|
||||
with_config!(self, config => config.secret_names())
|
||||
}
|
||||
|
||||
pub(crate) fn get_health_check_document(self) -> OcrDocument {
|
||||
with_config!(self, config => config.get_health_check_document())
|
||||
}
|
||||
|
|
@ -213,6 +217,8 @@ fn is_document_intelligence_model(model: &str) -> bool {
|
|||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::HashSet;
|
||||
|
||||
use litellm_auth::{InputSource, Sourced};
|
||||
use litellm_llms::{
|
||||
base_llm::ocr::document::InlineDocument, cohere::ocr::transformation::validate_document,
|
||||
|
|
@ -221,6 +227,27 @@ mod tests {
|
|||
|
||||
use super::*;
|
||||
|
||||
#[rstest]
|
||||
#[case(OcrConfigKind::AwsTextract)]
|
||||
#[case(OcrConfigKind::AwsTextractAnalyze)]
|
||||
#[case(OcrConfigKind::Cohere)]
|
||||
#[case(OcrConfigKind::Mistral)]
|
||||
#[case(OcrConfigKind::AzureAi)]
|
||||
#[case(OcrConfigKind::AzureCohere)]
|
||||
#[case(OcrConfigKind::AzureDocumentIntelligence)]
|
||||
#[case(OcrConfigKind::ReductoLegacy)]
|
||||
#[case(OcrConfigKind::ReductoV3)]
|
||||
#[case(OcrConfigKind::VertexAi)]
|
||||
#[case(OcrConfigKind::VertexDeepSeek)]
|
||||
fn secret_names_include_api_keys_without_duplicates(#[case] config: OcrConfigKind) {
|
||||
let names = config.secret_names();
|
||||
let unique = names.iter().collect::<HashSet<_>>();
|
||||
assert_eq!(names.len(), unique.len());
|
||||
if let Some(api_key) = config.get_api_key_env_var() {
|
||||
assert!(names.contains(&api_key));
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case("cohere")]
|
||||
#[case("mistral")]
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use futures_util::future::BoxFuture;
|
||||
use litellm_auth_gcp::VertexAuth;
|
||||
use litellm_host::{
|
||||
event::{CallEvent, MachineEvent, WireRequest},
|
||||
|
|
@ -10,11 +11,14 @@ use litellm_http::{
|
|||
HttpClientPool, HttpSettings, Resolution,
|
||||
media::{PublicDnsResolver, UrlPolicy},
|
||||
};
|
||||
use litellm_llms::base_llm::inference::secrets::{SecretSource, Secrets};
|
||||
use litellm_llms::base_llm::ocr::{
|
||||
error::Error as OcrError,
|
||||
handler::OcrClient,
|
||||
settings::OcrSettings,
|
||||
transformation::{LiteLLMOcrResponse, OCR_RESPONSE_MAX_BYTES, OcrTransportConfig},
|
||||
transformation::{
|
||||
BaseOcrConfig, LiteLLMOcrResponse, OCR_RESPONSE_MAX_BYTES, OcrTransportConfig,
|
||||
},
|
||||
};
|
||||
use rstest::rstest;
|
||||
use serde_json::{Value, json};
|
||||
|
|
@ -27,6 +31,32 @@ use super::{
|
|||
};
|
||||
use crate::ocr::route::{LocalOcrHost, OcrOp, OcrOpResult, ocr_machine};
|
||||
|
||||
struct RecordingSecretSource {
|
||||
names: Arc<Mutex<Vec<&'static str>>>,
|
||||
values: &'static [(&'static str, &'static str)],
|
||||
api_base: String,
|
||||
}
|
||||
|
||||
impl SecretSource for RecordingSecretSource {
|
||||
fn resolve<'a>(
|
||||
&'a self,
|
||||
names: &'a [&'static str],
|
||||
) -> BoxFuture<'a, Result<Secrets, litellm_secrets::Error>> {
|
||||
*self.names.lock().unwrap() = names.to_vec();
|
||||
let values = self.values;
|
||||
let api_base = self.api_base.clone();
|
||||
Box::pin(async move {
|
||||
Ok(Arc::new(move |name: &str| match name {
|
||||
"MISTRAL_AZURE_API_BASE" => Some(api_base.clone()),
|
||||
_ => values
|
||||
.iter()
|
||||
.find(|(key, _)| *key == name)
|
||||
.map(|(_, value)| value.to_string()),
|
||||
}) as Secrets)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::mistral("mistral/model", json!({}))]
|
||||
#[case::vertex("vertex_ai/mistral-ocr-latest", json!({"vertex_project":"test-project", "vertex_location":"us-central1"}))]
|
||||
|
|
@ -184,14 +214,11 @@ async fn mistral_env_fallbacks_follow_python_through_the_injected_secret_source(
|
|||
#[case] expected_key: &str,
|
||||
) {
|
||||
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await;
|
||||
let secret_base = base.clone();
|
||||
let client = ocr_client().with_secrets(Arc::new(move |name: &str| match name {
|
||||
"MISTRAL_AZURE_API_BASE" => Some(secret_base.clone()),
|
||||
"MISTRAL_API_BASE" => Some("http://127.0.0.1:9/never-read".into()),
|
||||
_ => secrets
|
||||
.iter()
|
||||
.find(|(key, _)| *key == name)
|
||||
.map(|(_, value)| value.to_string()),
|
||||
let names = Arc::new(Mutex::new(Vec::new()));
|
||||
let client = ocr_client().with_secrets(Arc::new(RecordingSecretSource {
|
||||
names: names.clone(),
|
||||
values: secrets,
|
||||
api_base: base.clone(),
|
||||
}));
|
||||
let request = decode_request(OcrWireRequest {
|
||||
model: "mistral/model".into(),
|
||||
|
|
@ -208,9 +235,47 @@ async fn mistral_env_fallbacks_follow_python_through_the_injected_secret_source(
|
|||
|
||||
crate::ocr::client::perform(&client, request).await.unwrap();
|
||||
server.await.unwrap();
|
||||
assert_eq!(
|
||||
*names.lock().unwrap(),
|
||||
litellm_llms::mistral::ocr::transformation::MistralOcrConfig.secret_names()
|
||||
);
|
||||
assert!(seen.lock().unwrap()[0].contains(&format!("authorization: Bearer {expected_key}")));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn mistral_ocr_resolves_provider_secrets_before_transformation() {
|
||||
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await;
|
||||
let names = Arc::new(Mutex::new(Vec::new()));
|
||||
let client = ocr_client().with_secrets(Arc::new(RecordingSecretSource {
|
||||
names: names.clone(),
|
||||
values: &[("MISTRAL_API_KEY", "source-key")],
|
||||
api_base: base.clone(),
|
||||
}));
|
||||
let request = decode_request(OcrWireRequest {
|
||||
model: "mistral/mistral-ocr-latest".into(),
|
||||
document: json!({
|
||||
"type":"document_url",
|
||||
"document_url":"data:application/pdf;base64,YWJj"
|
||||
}),
|
||||
api_key: None,
|
||||
api_base: None,
|
||||
custom_llm_provider: None,
|
||||
extra_headers: None,
|
||||
optional_params: Default::default(),
|
||||
input_sources: Default::default(),
|
||||
timeout_seconds: Some(2.0),
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
crate::ocr::client::perform(&client, request).await.unwrap();
|
||||
server.await.unwrap();
|
||||
assert_eq!(
|
||||
*names.lock().unwrap(),
|
||||
litellm_llms::mistral::ocr::transformation::MistralOcrConfig.secret_names()
|
||||
);
|
||||
assert!(seen.lock().unwrap()[0].contains("authorization: Bearer source-key"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn ocr_client_uses_the_injected_http_pool_configuration() {
|
||||
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await;
|
||||
|
|
@ -224,7 +289,7 @@ async fn ocr_client_uses_the_injected_http_pool_configuration() {
|
|||
UrlPolicy::default(),
|
||||
VertexAuth::default(),
|
||||
OcrSettings::default(),
|
||||
Arc::new(litellm_core_utils::settings::ProcessEnvironment),
|
||||
Arc::new(litellm_llms::base_llm::inference::secrets::EnvironmentSecrets),
|
||||
)
|
||||
.unwrap();
|
||||
crate::ocr::client::perform(&client, wire_request("mistral/model", &base, json!({})))
|
||||
|
|
|
|||
|
|
@ -29,7 +29,7 @@ pyo3::create_exception!(
|
|||
|
||||
static FORK_GATE: ForkGate = ForkGate::new();
|
||||
|
||||
/// Whether this process has started the Tokio runtime.
|
||||
/// Whether this process has entered process-bound native execution.
|
||||
pub fn runtime_started() -> bool {
|
||||
FORK_GATE.started(std::process::id())
|
||||
}
|
||||
|
|
@ -40,9 +40,8 @@ pub fn reserve_process_for_forking() -> Result<(), RuntimeAlreadyStarted> {
|
|||
FORK_GATE.reserve(std::process::id())
|
||||
}
|
||||
|
||||
/// The only door to the Tokio runtime: every route reaches it through this module, which is
|
||||
/// what lets the gate speak for the whole extension. `clippy.toml` disallows going around it.
|
||||
fn enter_runtime() -> PyResult<()> {
|
||||
/// Claims process-bound native state before runtime startup or tokenizer execution.
|
||||
pub fn enter_native() -> PyResult<()> {
|
||||
FORK_GATE
|
||||
.enter(std::process::id())
|
||||
.map_err(|refused| match refused {
|
||||
|
|
@ -60,7 +59,7 @@ fn enter_runtime() -> PyResult<()> {
|
|||
|
||||
#[expect(clippy::disallowed_methods, reason = "this is the gated door")]
|
||||
fn runtime() -> PyResult<&'static Runtime> {
|
||||
enter_runtime()?;
|
||||
enter_native()?;
|
||||
Ok(pyo3_async_runtimes::tokio::get_runtime())
|
||||
}
|
||||
|
||||
|
|
@ -70,7 +69,7 @@ where
|
|||
F: Future<Output = PyResult<T>> + Send + 'static,
|
||||
T: for<'py> IntoPyObject<'py> + Send + 'static,
|
||||
{
|
||||
enter_runtime()?;
|
||||
enter_native()?;
|
||||
pyo3_async_runtimes::tokio::future_into_py(py, future)
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -20,7 +20,7 @@ pub use argument::lookup;
|
|||
pub use callable::wrap_failure;
|
||||
pub use driver::run_call;
|
||||
pub use execution::{
|
||||
ForkedAfterNativeRuntimeStarted, ProcessReservedForForking, poll_async_value,
|
||||
ForkedAfterNativeRuntimeStarted, ProcessReservedForForking, enter_native, poll_async_value,
|
||||
reserve_process_for_forking, run_async, run_async_value, run_sync, run_sync_value,
|
||||
runtime_started,
|
||||
};
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@ litellm-auth-gcp.workspace = true
|
|||
litellm-host.workspace = true
|
||||
litellm-framing.workspace = true
|
||||
litellm-http.workspace = true
|
||||
litellm-secrets.workspace = true
|
||||
base64.workspace = true
|
||||
bytes.workspace = true
|
||||
data-url = "0.3.2"
|
||||
|
|
|
|||
|
|
@ -40,6 +40,10 @@ impl BaseOcrConfig for TextractAnalyzeDocumentConfig {
|
|||
type ProviderRequest = AnalyzeDocumentRequest;
|
||||
type Environment = TextractEnvironment;
|
||||
|
||||
fn secret_names(&self) -> Vec<&'static str> {
|
||||
litellm_auth_aws::constants::SECRET_NAMES.to_vec()
|
||||
}
|
||||
|
||||
fn get_supported_ocr_params(&self, _model: &str) -> &'static [&'static str] {
|
||||
&["feature_types"]
|
||||
}
|
||||
|
|
|
|||
|
|
@ -29,6 +29,10 @@ impl BaseOcrConfig for TextractDetectTextConfig {
|
|||
type ProviderRequest = DetectDocumentTextRequest;
|
||||
type Environment = TextractEnvironment;
|
||||
|
||||
fn secret_names(&self) -> Vec<&'static str> {
|
||||
litellm_auth_aws::constants::SECRET_NAMES.to_vec()
|
||||
}
|
||||
|
||||
fn get_health_check_document(&self) -> OcrDocument {
|
||||
health_check_document()
|
||||
}
|
||||
|
|
|
|||
|
|
@ -28,6 +28,10 @@ impl BaseOcrConfig for AzureAICohereParseConfig {
|
|||
super::transformation::AzureAiOcrConfig.get_api_key_env_var()
|
||||
}
|
||||
|
||||
fn secret_names(&self) -> Vec<&'static str> {
|
||||
super::transformation::AzureAiOcrConfig.secret_names()
|
||||
}
|
||||
|
||||
fn get_health_check_document(&self) -> OcrDocument {
|
||||
CohereParseConfig.get_health_check_document()
|
||||
}
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@ use std::{collections::BTreeSet, time::Duration};
|
|||
|
||||
use base64::{Engine, engine::general_purpose::STANDARD};
|
||||
use litellm_auth::{InputSource, Sourced};
|
||||
use litellm_auth_azure::AzureAuthInputs;
|
||||
use litellm_auth_azure::{AzureAuthInputs, SECRET_NAMES as AZURE_AUTH_SECRET_NAMES};
|
||||
use litellm_core_utils::{
|
||||
call_arguments::CallArguments,
|
||||
serde_compat::{FiniteF64, LaxI64},
|
||||
|
|
@ -141,6 +141,17 @@ impl BaseOcrConfig for AzureDocumentIntelligenceOcrConfig {
|
|||
Some(AZURE_DI_API_KEY_ENV)
|
||||
}
|
||||
|
||||
fn secret_names(&self) -> Vec<&'static str> {
|
||||
[
|
||||
[AZURE_DI_API_KEY_ENV, AZURE_DI_ENDPOINT_ENV].as_slice(),
|
||||
AZURE_AUTH_SECRET_NAMES,
|
||||
]
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.copied()
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn resolve_connection_params(&self, inputs: OcrCredentialInputs) -> ResolvedOcrCredentials {
|
||||
ResolvedOcrCredentials {
|
||||
api_key: inputs.api_key.and_then(|key| {
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
use litellm_auth::{InputSource, Sourced};
|
||||
use litellm_auth_azure::AzureAuthInputs;
|
||||
use litellm_auth_azure::SECRET_NAMES as AZURE_AUTH_SECRET_NAMES;
|
||||
use litellm_core_utils::{call_arguments::CallArguments, params::OpaqueParams, url_utils::ApiUrl};
|
||||
use serde_json::Value;
|
||||
|
||||
|
|
@ -37,6 +38,17 @@ impl BaseOcrConfig for AzureAiOcrConfig {
|
|||
Some(AZURE_AI_API_KEY_ENV)
|
||||
}
|
||||
|
||||
fn secret_names(&self) -> Vec<&'static str> {
|
||||
[
|
||||
[AZURE_AI_API_KEY_ENV, AZURE_AI_API_BASE_ENV].as_slice(),
|
||||
AZURE_AUTH_SECRET_NAMES,
|
||||
]
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.copied()
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn map_ocr_params(
|
||||
&self,
|
||||
non_default_params: &CallArguments,
|
||||
|
|
|
|||
1
litellm-rust/crates/llms/src/base_llm/inference/mod.rs
Normal file
1
litellm-rust/crates/llms/src/base_llm/inference/mod.rs
Normal file
|
|
@ -0,0 +1 @@
|
|||
pub mod secrets;
|
||||
19
litellm-rust/crates/llms/src/base_llm/inference/secrets.rs
Normal file
19
litellm-rust/crates/llms/src/base_llm/inference/secrets.rs
Normal file
|
|
@ -0,0 +1,19 @@
|
|||
use std::sync::Arc;
|
||||
|
||||
use futures_util::future::BoxFuture;
|
||||
use litellm_core_utils::settings::{Lookup, ProcessEnvironment};
|
||||
use litellm_secrets::Error;
|
||||
|
||||
pub type Secrets = Arc<dyn Lookup + Send + Sync>;
|
||||
|
||||
pub trait SecretSource: Send + Sync {
|
||||
fn resolve<'a>(&'a self, names: &'a [&'static str]) -> BoxFuture<'a, Result<Secrets, Error>>;
|
||||
}
|
||||
|
||||
pub struct EnvironmentSecrets;
|
||||
|
||||
impl SecretSource for EnvironmentSecrets {
|
||||
fn resolve<'a>(&'a self, _names: &'a [&'static str]) -> BoxFuture<'a, Result<Secrets, Error>> {
|
||||
Box::pin(async { Ok(Arc::new(ProcessEnvironment) as Secrets) })
|
||||
}
|
||||
}
|
||||
|
|
@ -2,5 +2,6 @@ pub mod anthropic_messages;
|
|||
pub mod audio_transcription;
|
||||
pub mod base_model_iterator;
|
||||
pub mod chat;
|
||||
pub mod inference;
|
||||
pub mod ocr;
|
||||
pub mod responses;
|
||||
|
|
|
|||
|
|
@ -98,6 +98,8 @@ pub enum Error {
|
|||
"Missing REDUCTO_API_KEY - set it in the environment or pass api_key to litellm.ocr()/litellm.aocr()"
|
||||
)]
|
||||
MissingReductoApiKey,
|
||||
#[error("secret resolution failed: {0}")]
|
||||
Secret(#[source] std::sync::Arc<litellm_secrets::Error>),
|
||||
#[error(transparent)]
|
||||
Auth(#[from] litellm_auth::Error),
|
||||
#[error(transparent)]
|
||||
|
|
|
|||
|
|
@ -1,3 +1,5 @@
|
|||
use std::sync::Arc;
|
||||
|
||||
use bytes::{Bytes, BytesMut};
|
||||
use futures_util::future::BoxFuture;
|
||||
use litellm_auth_gcp::VertexAuth;
|
||||
|
|
@ -11,9 +13,10 @@ use litellm_http::{
|
|||
use serde::{Serialize, de::DeserializeOwned};
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::base_llm::inference::secrets::SecretSource;
|
||||
use crate::base_llm::ocr::{
|
||||
error::Error,
|
||||
settings::{OcrSettings, Secrets},
|
||||
settings::OcrSettings,
|
||||
transformation::{
|
||||
BaseOcrConfig, DecodedOcrResponse, LiteLLMOcrResponse, OcrDocument, OcrResponseContext,
|
||||
PreparedOcrRequest, decode_request_value, decode_response,
|
||||
|
|
@ -35,7 +38,7 @@ pub struct OcrClient {
|
|||
document_fetcher: MediaFetcher,
|
||||
vertex_auth: VertexAuth,
|
||||
settings: OcrSettings,
|
||||
secrets: Secrets,
|
||||
secrets: Arc<dyn SecretSource>,
|
||||
}
|
||||
|
||||
impl OcrClient {
|
||||
|
|
@ -45,7 +48,7 @@ impl OcrClient {
|
|||
url_policy: UrlPolicy,
|
||||
vertex_auth: VertexAuth,
|
||||
settings: OcrSettings,
|
||||
secrets: Secrets,
|
||||
secrets: Arc<dyn SecretSource>,
|
||||
) -> Result<Self, litellm_http::Error> {
|
||||
Ok(Self {
|
||||
provider_http: pool.client(config, ClientVariant::Provider)?,
|
||||
|
|
@ -77,7 +80,7 @@ impl OcrClient {
|
|||
&self.settings
|
||||
}
|
||||
|
||||
pub fn secrets(&self) -> &Secrets {
|
||||
pub fn secret_source(&self) -> &Arc<dyn SecretSource> {
|
||||
&self.secrets
|
||||
}
|
||||
|
||||
|
|
@ -92,7 +95,7 @@ impl OcrClient {
|
|||
document_fetcher: MediaFetcher::for_test(document_http),
|
||||
vertex_auth: VertexAuth::default(),
|
||||
settings: OcrSettings::default(),
|
||||
secrets: std::sync::Arc::new(litellm_core_utils::settings::ProcessEnvironment),
|
||||
secrets: Arc::new(crate::base_llm::inference::secrets::EnvironmentSecrets),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -102,7 +105,7 @@ impl OcrClient {
|
|||
}
|
||||
|
||||
#[cfg(any(test, feature = "test-support"))]
|
||||
pub fn with_secrets(self, secrets: Secrets) -> Self {
|
||||
pub fn with_secrets(self, secrets: Arc<dyn SecretSource>) -> Self {
|
||||
Self { secrets, ..self }
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,9 +1,7 @@
|
|||
use std::{sync::Arc, time::Duration};
|
||||
use std::time::Duration;
|
||||
|
||||
use litellm_core_utils::settings::Lookup;
|
||||
|
||||
pub type Secrets = Arc<dyn Lookup + Send + Sync>;
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct OcrSettings {
|
||||
pub request_timeout: Duration,
|
||||
|
|
|
|||
|
|
@ -14,10 +14,13 @@ use serde::{
|
|||
use serde_json::{Map, Value};
|
||||
use serde_with::serde_as;
|
||||
|
||||
use crate::base_llm::ocr::{
|
||||
error::Error,
|
||||
handler::{CallHooks, OcrClient, read_response_bytes, transform_request_body},
|
||||
settings::{OcrSettings, Secrets},
|
||||
use crate::base_llm::{
|
||||
inference::secrets::Secrets,
|
||||
ocr::{
|
||||
error::Error,
|
||||
handler::{CallHooks, OcrClient, read_response_bytes, transform_request_body},
|
||||
settings::OcrSettings,
|
||||
},
|
||||
};
|
||||
|
||||
pub const OCR_RESPONSE_MAX_BYTES: usize = 64 * 1024 * 1024;
|
||||
|
|
@ -436,6 +439,8 @@ pub trait BaseOcrConfig: Send + Sync + Sized + 'static {
|
|||
None
|
||||
}
|
||||
|
||||
fn secret_names(&self) -> Vec<&'static str>;
|
||||
|
||||
fn resolve_connection_params(&self, inputs: OcrCredentialInputs) -> ResolvedOcrCredentials {
|
||||
ResolvedOcrCredentials {
|
||||
api_key: inputs
|
||||
|
|
|
|||
|
|
@ -102,6 +102,10 @@ impl BaseOcrConfig for CohereParseConfig {
|
|||
Some(COHERE_API_KEY_ENV)
|
||||
}
|
||||
|
||||
fn secret_names(&self) -> Vec<&'static str> {
|
||||
vec![COHERE_API_KEY_ENV]
|
||||
}
|
||||
|
||||
fn get_health_check_document(&self) -> OcrDocument {
|
||||
OcrDocument::ImageUrl {
|
||||
image_url: COHERE_PARSE_HEALTH_CHECK_IMAGE_DATA_URI.into(),
|
||||
|
|
|
|||
|
|
@ -69,6 +69,14 @@ impl BaseOcrConfig for MistralOcrConfig {
|
|||
Some(MISTRAL_OCR_API_KEY_ENV_VAR)
|
||||
}
|
||||
|
||||
fn secret_names(&self) -> Vec<&'static str> {
|
||||
vec![
|
||||
MISTRAL_OCR_API_KEY_ENV_VAR,
|
||||
"MISTRAL_AZURE_API_KEY",
|
||||
"MISTRAL_AZURE_API_BASE",
|
||||
]
|
||||
}
|
||||
|
||||
fn map_ocr_params(
|
||||
&self,
|
||||
non_default_params: &CallArguments,
|
||||
|
|
|
|||
|
|
@ -92,6 +92,10 @@ impl BaseOcrConfig for ReductoParseV3Config {
|
|||
type ProviderRequest = ReductoV3Request;
|
||||
type Environment = Vec<(String, String)>;
|
||||
|
||||
fn secret_names(&self) -> Vec<&'static str> {
|
||||
vec![REDUCTO_API_KEY_ENV]
|
||||
}
|
||||
|
||||
fn get_supported_ocr_params(&self, _model: &str) -> &'static [&'static str] {
|
||||
&["formatting", "retrieval", "settings"]
|
||||
}
|
||||
|
|
@ -180,6 +184,10 @@ impl BaseOcrConfig for ReductoParseLegacyConfig {
|
|||
type ProviderRequest = ReductoLegacyRequest;
|
||||
type Environment = Vec<(String, String)>;
|
||||
|
||||
fn secret_names(&self) -> Vec<&'static str> {
|
||||
vec![REDUCTO_API_KEY_ENV]
|
||||
}
|
||||
|
||||
fn get_supported_ocr_params(&self, _model: &str) -> &'static [&'static str] {
|
||||
&["enhance"]
|
||||
}
|
||||
|
|
|
|||
|
|
@ -105,6 +105,10 @@ impl BaseOcrConfig for VertexAIDeepSeekOCRConfig {
|
|||
VertexAiOcrConfig.get_api_key_env_var()
|
||||
}
|
||||
|
||||
fn secret_names(&self) -> Vec<&'static str> {
|
||||
VertexAiOcrConfig.secret_names()
|
||||
}
|
||||
|
||||
fn map_ocr_params(
|
||||
&self,
|
||||
_arguments: &CallArguments,
|
||||
|
|
|
|||
|
|
@ -34,6 +34,10 @@ impl BaseOcrConfig for VertexAiOcrConfig {
|
|||
Some("VERTEX_AI_API_KEY")
|
||||
}
|
||||
|
||||
fn secret_names(&self) -> Vec<&'static str> {
|
||||
litellm_auth_gcp::SECRET_NAMES.to_vec()
|
||||
}
|
||||
|
||||
fn map_ocr_params(
|
||||
&self,
|
||||
non_default_params: &CallArguments,
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@ name = "_native"
|
|||
crate-type = ["cdylib"]
|
||||
|
||||
[features]
|
||||
default = ["abi3", "fast"]
|
||||
default = ["abi3", "fast", "huggingface", "tiktoken"]
|
||||
abi3 = ["pyo3/abi3-py310"]
|
||||
extension-module = ["pyo3/extension-module"]
|
||||
panic-test = []
|
||||
|
|
@ -20,6 +20,7 @@ tiktoken = ["litellm-token-counter/tiktoken"]
|
|||
|
||||
[dependencies]
|
||||
bytes.workspace = true
|
||||
futures-util.workspace = true
|
||||
litellm-cache.workspace = true
|
||||
litellm-cache-azure-blob.workspace = true
|
||||
litellm-cache-memory.workspace = true
|
||||
|
|
@ -27,7 +28,10 @@ litellm-cache-redis.workspace = true
|
|||
litellm-cache-s3.workspace = true
|
||||
litellm-cache-gcs.workspace = true
|
||||
litellm-cache-disk.workspace = true
|
||||
litellm-cache-redis-semantic.workspace = true
|
||||
litellm-cache-response.workspace = true
|
||||
litellm-cache-qdrant-semantic.workspace = true
|
||||
qdrant-client.workspace = true
|
||||
litellm-cache-valkey-semantic = { path = "../cache-valkey-semantic" }
|
||||
serde.workspace = true
|
||||
litellm-auth.workspace = true
|
||||
|
|
@ -38,16 +42,21 @@ litellm-core-utils.workspace = true
|
|||
litellm-auth-gcp.workspace = true
|
||||
litellm-http.workspace = true
|
||||
litellm-llms.workspace = true
|
||||
litellm-secrets = { workspace = true, features = ["aws"] }
|
||||
litellm-secrets-types.workspace = true
|
||||
litellm-types.workspace = true
|
||||
litellm-host-python.workspace = true
|
||||
litellm-token-counter = { path = "../token-counter", default-features = false }
|
||||
pyo3.workspace = true
|
||||
pyo3-async-runtimes.workspace = true
|
||||
reqwest.workspace = true
|
||||
redis = { version = "1.7.0", features = ["tls-rustls"] }
|
||||
serde_json.workspace = true
|
||||
tokio = { workspace = true, features = ["sync"] }
|
||||
url.workspace = true
|
||||
tokio = { workspace = true, features = ["rt", "sync"] }
|
||||
|
||||
[dev-dependencies]
|
||||
litellm-secrets-aws.workspace = true
|
||||
serde.workspace = true
|
||||
serde_with.workspace = true
|
||||
criterion.workspace = true
|
||||
|
|
@ -55,6 +64,8 @@ futures-util.workspace = true
|
|||
rstest.workspace = true
|
||||
sha2.workspace = true
|
||||
tokio-tungstenite.workspace = true
|
||||
wiremock = "0.6.5"
|
||||
aws-sdk-secretsmanager = "1.117.0"
|
||||
|
||||
[[bench]]
|
||||
name = "serialization"
|
||||
|
|
|
|||
5
litellm-rust/crates/python-bridge/README.md
Normal file
5
litellm-rust/crates/python-bridge/README.md
Normal file
|
|
@ -0,0 +1,5 @@
|
|||
Native OCR uses `SecretSource` with `EnvironmentSecrets`, preserving process-environment reads. Readable Python secret managers still make OCR decline to the existing Python implementation. `ResolvedSecrets` and the separate `secret_manager_binding()` snapshot are inactive foundations for a later rollout
|
||||
|
||||
Cache and secret-manager catalog entries remain Python-only, including when `LITELLM_RUST=1`. The new cache runtime is not connected to SDK or gateway caching
|
||||
|
||||
OCR provider requests use the shared `litellm-http` pool. AWS and Google secret-manager SDK clients keep their SDK transports, which do not yet inherit the pool's proxy, TLS, certificate, timeout, or observability configuration. Preserve those SDK transports and configure them equivalently instead of forcing them through reqwest
|
||||
|
|
@ -1,154 +0,0 @@
|
|||
{
|
||||
"http_settings": {
|
||||
"version": 1,
|
||||
"fields": {
|
||||
"ssl_verify": {
|
||||
"adapter": "SslVerifyInput",
|
||||
"required": true,
|
||||
"precedence": "module_global",
|
||||
"sensitive": false,
|
||||
"shapes": [
|
||||
"none",
|
||||
"bool",
|
||||
"str"
|
||||
],
|
||||
"unsupported_live": "configuration_error"
|
||||
},
|
||||
"ssl_certificate": {
|
||||
"adapter": "OptionalStrictString",
|
||||
"required": true,
|
||||
"precedence": "module_global",
|
||||
"sensitive": false,
|
||||
"shapes": [],
|
||||
"unsupported_live": null
|
||||
},
|
||||
"ssl_security_level": {
|
||||
"adapter": "TuningString",
|
||||
"required": true,
|
||||
"precedence": "module_global",
|
||||
"sensitive": false,
|
||||
"shapes": [],
|
||||
"unsupported_live": null
|
||||
},
|
||||
"ssl_ecdh_curve": {
|
||||
"adapter": "TuningString",
|
||||
"required": true,
|
||||
"precedence": "module_global",
|
||||
"sensitive": false,
|
||||
"shapes": [],
|
||||
"unsupported_live": null
|
||||
},
|
||||
"force_ipv4": {
|
||||
"adapter": "Truthy",
|
||||
"required": true,
|
||||
"precedence": "module_global",
|
||||
"sensitive": false,
|
||||
"shapes": [],
|
||||
"unsupported_live": null
|
||||
},
|
||||
"http2": {
|
||||
"adapter": "ExactTrue",
|
||||
"required": true,
|
||||
"precedence": "module_global",
|
||||
"sensitive": false,
|
||||
"shapes": [],
|
||||
"unsupported_live": null
|
||||
},
|
||||
"aiohttp_trust_env": {
|
||||
"adapter": "Truthy",
|
||||
"required": true,
|
||||
"precedence": "module_global",
|
||||
"sensitive": false,
|
||||
"shapes": [],
|
||||
"unsupported_live": null
|
||||
},
|
||||
"disable_aiohttp_trust_env": {
|
||||
"adapter": "Truthy",
|
||||
"required": true,
|
||||
"precedence": "module_global",
|
||||
"sensitive": false,
|
||||
"shapes": [],
|
||||
"unsupported_live": null
|
||||
},
|
||||
"disable_aiohttp_transport": {
|
||||
"adapter": "ExactTrue",
|
||||
"required": true,
|
||||
"precedence": "module_global",
|
||||
"sensitive": false,
|
||||
"shapes": [],
|
||||
"unsupported_live": null
|
||||
},
|
||||
"user_agent": {
|
||||
"adapter": "StrictString",
|
||||
"required": true,
|
||||
"precedence": "accessor",
|
||||
"sensitive": false,
|
||||
"shapes": [],
|
||||
"unsupported_live": null
|
||||
}
|
||||
}
|
||||
},
|
||||
"url_policy": {
|
||||
"version": 1,
|
||||
"fields": {
|
||||
"user_url_validation": {
|
||||
"adapter": "Truthy",
|
||||
"required": true,
|
||||
"precedence": "module_global",
|
||||
"sensitive": false,
|
||||
"shapes": [],
|
||||
"unsupported_live": null
|
||||
},
|
||||
"user_url_allowed_hosts": {
|
||||
"adapter": "HostCollection",
|
||||
"required": true,
|
||||
"precedence": "module_global",
|
||||
"sensitive": false,
|
||||
"shapes": [],
|
||||
"unsupported_live": null
|
||||
}
|
||||
}
|
||||
},
|
||||
"provider_defaults": {
|
||||
"version": 1,
|
||||
"fields": {
|
||||
"vertex_project": {
|
||||
"adapter": "FalsyOptionalString",
|
||||
"required": true,
|
||||
"precedence": "module_global",
|
||||
"sensitive": true,
|
||||
"shapes": [],
|
||||
"unsupported_live": null
|
||||
},
|
||||
"vertex_location": {
|
||||
"adapter": "FalsyOptionalString",
|
||||
"required": true,
|
||||
"precedence": "module_global",
|
||||
"sensitive": true,
|
||||
"shapes": [],
|
||||
"unsupported_live": null
|
||||
},
|
||||
"enable_azure_ad_token_refresh": {
|
||||
"adapter": "ExactTrue",
|
||||
"required": true,
|
||||
"precedence": "module_global",
|
||||
"sensitive": false,
|
||||
"shapes": [],
|
||||
"unsupported_live": null
|
||||
}
|
||||
}
|
||||
},
|
||||
"secret_manager": {
|
||||
"version": 1,
|
||||
"fields": {
|
||||
"readable": {
|
||||
"adapter": "StrictBool",
|
||||
"required": true,
|
||||
"precedence": "accessor",
|
||||
"sensitive": false,
|
||||
"shapes": [],
|
||||
"unsupported_live": null
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -11,10 +11,12 @@ use serde_json::Value;
|
|||
use super::{
|
||||
cache_error,
|
||||
callback::PythonCallback,
|
||||
config::{CacheBackendConfig, CacheConfigProjection, NativeCacheConfig},
|
||||
future::{ready_none, ready_value},
|
||||
native::NativeResponseCache,
|
||||
request::{now, request, requests},
|
||||
};
|
||||
use crate::errors::RustBridgeDeclined;
|
||||
|
||||
pub(super) enum CacheBinding {
|
||||
Disabled,
|
||||
|
|
@ -22,7 +24,7 @@ pub(super) enum CacheBinding {
|
|||
PythonCallback(PythonCallback),
|
||||
}
|
||||
|
||||
#[pyclass(frozen, name = "_CacheTestBinding")]
|
||||
#[pyclass(frozen, name = "_ResponseCacheRuntime")]
|
||||
pub(crate) struct ResolvedCache {
|
||||
binding: CacheBinding,
|
||||
pid: u32,
|
||||
|
|
@ -66,6 +68,33 @@ impl ResolvedCache {
|
|||
|
||||
#[pymethods]
|
||||
impl ResolvedCache {
|
||||
#[staticmethod]
|
||||
fn from_cache(cache: &Bound<'_, PyAny>) -> PyResult<Self> {
|
||||
let config = match NativeCacheConfig::project(cache)? {
|
||||
CacheConfigProjection::Native(config) => *config,
|
||||
CacheConfigProjection::Unsupported(reason) => {
|
||||
return Err(RustBridgeDeclined::new_err(reason.message()));
|
||||
}
|
||||
};
|
||||
let service = match config.backend {
|
||||
CacheBackendConfig::Memory(memory) => NativeResponseCache::memory(
|
||||
memory.capacity,
|
||||
memory.default_ttl,
|
||||
memory.max_entry_bytes,
|
||||
),
|
||||
_ => {
|
||||
return Err(RustBridgeDeclined::new_err(
|
||||
"native response cache activation is not implemented for this backend",
|
||||
));
|
||||
}
|
||||
};
|
||||
Ok(Self::new(CacheBinding::Native(
|
||||
service
|
||||
.with_scope(config.policy.semantic_cache_scope)
|
||||
.with_redis_flush_size(config.policy.redis_flush_size),
|
||||
)))
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn kind(&self) -> &'static str {
|
||||
match self.binding {
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@ use std::{path::PathBuf, time::Duration};
|
|||
|
||||
use litellm_auth_aws::AwsAuthConfig;
|
||||
use litellm_cache::CacheType;
|
||||
use litellm_cache_qdrant_semantic::{OpenAiEmbedderConfig, QdrantSemanticConfig, Quantization};
|
||||
use litellm_cache_redis::{RedisNode, RedisTopology};
|
||||
use litellm_cache_s3::{S3CacheConfig, S3Endpoint};
|
||||
use pyo3::{
|
||||
|
|
@ -10,7 +11,7 @@ use pyo3::{
|
|||
types::{PyAny, PyBool, PyDict, PyList, PyString},
|
||||
};
|
||||
|
||||
use super::{native::NativeResponseCache, request::duration};
|
||||
use super::{identity::BackendIdentity, native::NativeResponseCache, request::duration};
|
||||
|
||||
#[allow(dead_code, reason = "consumed by the cache activation follow-up")]
|
||||
pub(super) struct CachePolicy {
|
||||
|
|
@ -88,6 +89,24 @@ pub(super) struct GcsCacheConfig {
|
|||
pub(super) path_service_account: Option<String>,
|
||||
}
|
||||
|
||||
pub(super) struct AzureBlobCacheConfig {
|
||||
pub(super) account_url: String,
|
||||
pub(super) container: String,
|
||||
}
|
||||
|
||||
#[allow(
|
||||
dead_code,
|
||||
reason = "embedding settings are projected so drift falls back to Python"
|
||||
)]
|
||||
pub(super) struct RedisSemanticCacheConfig {
|
||||
pub(super) redis_url: String,
|
||||
pub(super) index_name: String,
|
||||
pub(super) similarity_threshold: f64,
|
||||
pub(super) embedding_model: String,
|
||||
pub(super) embedding_max_input_tokens: Option<u64>,
|
||||
pub(super) embedding_timeout: Option<f64>,
|
||||
}
|
||||
|
||||
struct RedisClientProjection<'py> {
|
||||
topology: RedisTopology,
|
||||
host: String,
|
||||
|
|
@ -107,9 +126,25 @@ pub(super) struct ValkeySemanticCacheConfig {
|
|||
pub(super) connection: RedisConnectionConfig,
|
||||
}
|
||||
|
||||
pub(super) struct AzureBlobCacheConfig {
|
||||
pub(super) account_url: String,
|
||||
pub(super) container: String,
|
||||
pub(super) struct QdrantSemanticCacheConfig {
|
||||
pub(super) grpc_url: String,
|
||||
pub(super) api_key: Option<String>,
|
||||
pub(super) collection_name: String,
|
||||
pub(super) similarity_threshold: f64,
|
||||
pub(super) vector_size: u64,
|
||||
pub(super) embedding: OpenAiEmbedderConfig,
|
||||
pub(super) quantization: Quantization,
|
||||
}
|
||||
|
||||
impl QdrantSemanticCacheConfig {
|
||||
pub(super) fn to_qdrant_config(&self) -> QdrantSemanticConfig {
|
||||
QdrantSemanticConfig {
|
||||
collection_name: self.collection_name.clone(),
|
||||
similarity_threshold: self.similarity_threshold,
|
||||
vector_size: self.vector_size,
|
||||
quantization: self.quantization.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) enum CacheBackendConfig {
|
||||
|
|
@ -120,6 +155,8 @@ pub(super) enum CacheBackendConfig {
|
|||
ValkeySemantic(Box<ValkeySemanticCacheConfig>),
|
||||
Disk(DiskCacheConfig),
|
||||
AzureBlob(AzureBlobCacheConfig),
|
||||
RedisSemantic(Box<RedisSemanticCacheConfig>),
|
||||
QdrantSemantic(Box<QdrantSemanticCacheConfig>),
|
||||
}
|
||||
|
||||
#[allow(dead_code, reason = "consumed by the cache activation follow-up")]
|
||||
|
|
@ -139,6 +176,8 @@ pub(super) enum UnsupportedCacheConfig {
|
|||
S3Option,
|
||||
GcsBucket,
|
||||
DiskStore,
|
||||
QdrantEndpoint,
|
||||
SemanticEmbedding,
|
||||
}
|
||||
|
||||
impl UnsupportedCacheConfig {
|
||||
|
|
@ -154,6 +193,10 @@ impl UnsupportedCacheConfig {
|
|||
Self::S3Option => "native S3 configuration requires Python",
|
||||
Self::GcsBucket => "native GCS cache requires a configured bucket name",
|
||||
Self::DiskStore => "native disk cache requires the built-in diskcache store",
|
||||
Self::QdrantEndpoint => {
|
||||
"native Qdrant requires the default REST port so the gRPC port can be derived"
|
||||
}
|
||||
Self::SemanticEmbedding => "native semantic embedding requires Python",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -224,138 +267,188 @@ impl NativeCacheConfig {
|
|||
}))),
|
||||
Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)),
|
||||
},
|
||||
Some(CacheType::QdrantSemantic) => match project_qdrant_semantic(&backend)? {
|
||||
Ok(backend) => Ok(CacheConfigProjection::Native(Box::new(Self {
|
||||
policy,
|
||||
backend: CacheBackendConfig::QdrantSemantic(Box::new(backend)),
|
||||
}))),
|
||||
Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)),
|
||||
},
|
||||
Some(CacheType::AzureBlob) => project_azure_blob(&backend).map(|backend| {
|
||||
CacheConfigProjection::Native(Box::new(Self {
|
||||
policy,
|
||||
backend: CacheBackendConfig::AzureBlob(backend),
|
||||
}))
|
||||
}),
|
||||
Some(CacheType::RedisSemantic | CacheType::QdrantSemantic) | None => Ok(
|
||||
CacheConfigProjection::Unsupported(UnsupportedCacheConfig::Backend),
|
||||
),
|
||||
Some(CacheType::RedisSemantic) => project_redis_semantic(&backend).map(|backend| {
|
||||
CacheConfigProjection::Native(Box::new(Self {
|
||||
policy,
|
||||
backend: CacheBackendConfig::RedisSemantic(Box::new(backend)),
|
||||
}))
|
||||
}),
|
||||
None => Ok(CacheConfigProjection::Unsupported(
|
||||
UnsupportedCacheConfig::Backend,
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn service_mismatch(&self, service: &NativeResponseCache) -> Option<&'static str> {
|
||||
let default_ttl = match &self.backend {
|
||||
CacheBackendConfig::Memory(config) => Some(config.default_ttl),
|
||||
CacheBackendConfig::Redis(config) => Some(config.default_ttl),
|
||||
CacheBackendConfig::S3(_) => None,
|
||||
CacheBackendConfig::ValkeySemantic(_) => Some(Duration::ZERO),
|
||||
CacheBackendConfig::Disk(_)
|
||||
| CacheBackendConfig::AzureBlob(_)
|
||||
| CacheBackendConfig::Gcs(_) => None,
|
||||
};
|
||||
if !matches!(self.backend, CacheBackendConfig::ValkeySemantic(_))
|
||||
&& service.default_ttl() != default_ttl
|
||||
{
|
||||
return Some("facade and native backend default TTLs must match");
|
||||
}
|
||||
match &self.backend {
|
||||
CacheBackendConfig::Memory(config) if service.kind() != "memory" => {
|
||||
Some("facade and native backend types must match")
|
||||
}
|
||||
CacheBackendConfig::Memory(config) if service.capacity() != Some(config.capacity) => {
|
||||
Some("facade and native backend capacities must match")
|
||||
}
|
||||
CacheBackendConfig::Memory(config)
|
||||
if service.max_entry_bytes() != Some(config.max_entry_bytes) =>
|
||||
{
|
||||
Some("facade and native backend item limits must match")
|
||||
}
|
||||
CacheBackendConfig::Memory(_) => None,
|
||||
CacheBackendConfig::Redis(_) if service.kind() != "redis" => {
|
||||
Some("facade and native backend types must match")
|
||||
}
|
||||
CacheBackendConfig::Redis(config) if service.topology() != Some(&config.topology) => {
|
||||
Some("facade and native backend topologies must match")
|
||||
}
|
||||
CacheBackendConfig::Redis(config) => (service.namespace()
|
||||
!= config.namespace.as_deref())
|
||||
.then_some("facade and native backend namespaces must match"),
|
||||
CacheBackendConfig::S3(_) if service.kind() != "s3" => {
|
||||
Some("facade and native backend types must match")
|
||||
}
|
||||
CacheBackendConfig::S3(config) if service.bucket() != Some(config.bucket.as_str()) => {
|
||||
Some("facade and native backend buckets must match")
|
||||
}
|
||||
CacheBackendConfig::S3(config)
|
||||
if service.key_prefix() != Some(config.key_prefix.as_str()) =>
|
||||
{
|
||||
Some("facade and native backend key prefixes must match")
|
||||
}
|
||||
CacheBackendConfig::S3(config) if service.region() != Some(config.region.as_str()) => {
|
||||
Some("facade and native backend regions must match")
|
||||
}
|
||||
CacheBackendConfig::S3(config)
|
||||
if service.endpoint()
|
||||
!= config
|
||||
.endpoint
|
||||
.as_ref()
|
||||
.map(|endpoint| endpoint.url.as_str()) =>
|
||||
{
|
||||
Some("facade and native backend endpoints must match")
|
||||
}
|
||||
CacheBackendConfig::S3(_) => None,
|
||||
CacheBackendConfig::Gcs(_) if service.kind() != "gcs" => {
|
||||
Some("facade and native backend types must match")
|
||||
}
|
||||
CacheBackendConfig::Gcs(config)
|
||||
if service
|
||||
.gcs_backend()
|
||||
.is_none_or(|backend| backend.bucket_name() != config.bucket_name) =>
|
||||
{
|
||||
Some("facade and native backend buckets must match")
|
||||
}
|
||||
CacheBackendConfig::Gcs(config)
|
||||
if service
|
||||
.gcs_backend()
|
||||
.is_none_or(|backend| backend.key_prefix() != config.key_prefix) =>
|
||||
{
|
||||
Some("facade and native backend key prefixes must match")
|
||||
}
|
||||
CacheBackendConfig::Gcs(config)
|
||||
if service.gcs_backend().is_none_or(|backend| {
|
||||
backend.path_service_account() != config.path_service_account.as_deref()
|
||||
}) =>
|
||||
{
|
||||
Some("facade and native backend credentials must match")
|
||||
}
|
||||
CacheBackendConfig::Gcs(_) => None,
|
||||
CacheBackendConfig::ValkeySemantic(config) => {
|
||||
if service.kind() != "valkey-semantic" {
|
||||
return Some("facade and native backend types must match");
|
||||
}
|
||||
let Some((threshold, index_name)) = service.semantic_config() else {
|
||||
return Some("facade and native backend types must match");
|
||||
};
|
||||
(threshold != config.similarity_threshold || index_name != config.index_name)
|
||||
.then_some("facade and native semantic settings must match")
|
||||
}
|
||||
CacheBackendConfig::Disk(_) if service.kind() != "disk" => {
|
||||
Some("facade and native backend types must match")
|
||||
}
|
||||
CacheBackendConfig::Disk(config) => {
|
||||
let Some(directory) = service.directory() else {
|
||||
return Some("facade and native backend types must match");
|
||||
};
|
||||
let native = std::fs::canonicalize(directory).ok();
|
||||
let facade = std::fs::canonicalize(&config.directory).ok();
|
||||
(native != facade).then_some("facade and native backend directories must match")
|
||||
}
|
||||
CacheBackendConfig::AzureBlob(config) => match service.azure_blob_identity() {
|
||||
None => Some("facade and native backend types must match"),
|
||||
Some((account_url, container))
|
||||
if account_url != config.account_url || container != config.container =>
|
||||
{
|
||||
Some("facade and native backend containers must match")
|
||||
}
|
||||
Some(_) => None,
|
||||
self.backend.identity().mismatch(&service.identity())
|
||||
}
|
||||
}
|
||||
|
||||
impl CacheBackendConfig {
|
||||
/// The identity a native backend must have for this facade configuration to describe it.
|
||||
pub(super) fn identity(&self) -> BackendIdentity {
|
||||
match self {
|
||||
Self::Memory(config) => BackendIdentity::Memory {
|
||||
capacity: config.capacity,
|
||||
max_entry_bytes: Some(config.max_entry_bytes),
|
||||
default_ttl: Some(config.default_ttl),
|
||||
},
|
||||
Self::Redis(config) => BackendIdentity::Redis {
|
||||
topology: config.topology.clone(),
|
||||
namespace: config.namespace.clone(),
|
||||
default_ttl: Some(config.default_ttl),
|
||||
},
|
||||
Self::S3(config) => BackendIdentity::S3 {
|
||||
bucket: config.bucket.clone(),
|
||||
key_prefix: config.key_prefix.clone(),
|
||||
region: config.region.clone(),
|
||||
endpoint: config
|
||||
.endpoint
|
||||
.as_ref()
|
||||
.map(|endpoint| endpoint.url.clone()),
|
||||
},
|
||||
Self::Gcs(config) => BackendIdentity::Gcs {
|
||||
bucket_name: config.bucket_name.clone(),
|
||||
key_prefix: config.key_prefix.clone(),
|
||||
path_service_account: config.path_service_account.clone(),
|
||||
},
|
||||
Self::ValkeySemantic(config) => BackendIdentity::ValkeySemantic {
|
||||
index_name: config.index_name.clone(),
|
||||
similarity_threshold: config.similarity_threshold,
|
||||
},
|
||||
Self::Disk(config) => BackendIdentity::Disk {
|
||||
directory: config.directory.clone(),
|
||||
},
|
||||
Self::AzureBlob(config) => BackendIdentity::AzureBlob {
|
||||
account_url: config.account_url.clone(),
|
||||
container: config.container.clone(),
|
||||
},
|
||||
Self::RedisSemantic(config) => BackendIdentity::RedisSemantic {
|
||||
index_name: config.index_name.clone(),
|
||||
similarity_threshold: config.similarity_threshold as f32,
|
||||
},
|
||||
Self::QdrantSemantic(config) => BackendIdentity::QdrantSemantic {
|
||||
collection_name: config.collection_name.clone(),
|
||||
similarity_threshold: config.similarity_threshold,
|
||||
vector_size: config.vector_size,
|
||||
embedding_model: config.embedding.model.clone(),
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[inline(never)]
|
||||
fn project_qdrant_semantic(
|
||||
backend: &Bound<'_, PyAny>,
|
||||
) -> PyResult<Result<QdrantSemanticCacheConfig, UnsupportedCacheConfig>> {
|
||||
let rest_url = backend.getattr("qdrant_api_base")?.extract::<String>()?;
|
||||
let parsed = match url::Url::parse(&rest_url) {
|
||||
Ok(value) => value,
|
||||
Err(_) => return Ok(Err(UnsupportedCacheConfig::QdrantEndpoint)),
|
||||
};
|
||||
if !matches!(parsed.scheme(), "http" | "https")
|
||||
|| (!parsed.path().is_empty() && parsed.path() != "/")
|
||||
|| parsed.query().is_some()
|
||||
|| parsed.host_str().is_none()
|
||||
|| parsed.port() != Some(6333)
|
||||
{
|
||||
return Ok(Err(UnsupportedCacheConfig::QdrantEndpoint));
|
||||
}
|
||||
let mut grpc_url = parsed;
|
||||
if grpc_url.set_port(Some(6334)).is_err() {
|
||||
return Ok(Err(UnsupportedCacheConfig::QdrantEndpoint));
|
||||
}
|
||||
grpc_url.set_path("");
|
||||
grpc_url.set_query(None);
|
||||
|
||||
if optional_attribute(backend, "embedding_max_input_tokens")?
|
||||
.is_some_and(|value| !value.is_none())
|
||||
{
|
||||
return Ok(Err(UnsupportedCacheConfig::SemanticEmbedding));
|
||||
}
|
||||
let configured_model = backend.getattr("embedding_model")?.extract::<String>()?;
|
||||
let embedding_model = configured_model
|
||||
.strip_prefix("openai/")
|
||||
.unwrap_or(&configured_model)
|
||||
.to_owned();
|
||||
if !embedding_model.starts_with("text-embedding-") {
|
||||
return Ok(Err(UnsupportedCacheConfig::SemanticEmbedding));
|
||||
}
|
||||
let proxy_server = py_sys_module(backend.py())?;
|
||||
if let Some(proxy_server) = proxy_server {
|
||||
let router = proxy_server.getattr("llm_router")?;
|
||||
let model_list = proxy_server.getattr("llm_model_list")?;
|
||||
let embedding_router = backend.py().import("litellm.caching._embedding_router")?;
|
||||
if !embedding_router
|
||||
.getattr("resolve_embedding_router")?
|
||||
.call1((configured_model.as_str(), router, model_list))?
|
||||
.is_none()
|
||||
{
|
||||
return Ok(Err(UnsupportedCacheConfig::SemanticEmbedding));
|
||||
}
|
||||
}
|
||||
let litellm = backend.py().import("litellm")?;
|
||||
for name in ["api_key", "openai_key", "api_base"] {
|
||||
if !litellm.getattr(name)?.is_none() {
|
||||
return Ok(Err(UnsupportedCacheConfig::SemanticEmbedding));
|
||||
}
|
||||
}
|
||||
let Ok(embedding_api_key) = std::env::var("OPENAI_API_KEY") else {
|
||||
return Ok(Err(UnsupportedCacheConfig::SemanticEmbedding));
|
||||
};
|
||||
if embedding_api_key.is_empty() {
|
||||
return Ok(Err(UnsupportedCacheConfig::SemanticEmbedding));
|
||||
}
|
||||
let embedding_api_base = std::env::var("OPENAI_BASE_URL")
|
||||
.or_else(|_| std::env::var("OPENAI_API_BASE"))
|
||||
.unwrap_or_else(|_| "https://api.openai.com/v1".to_owned());
|
||||
let timeout = optional_attribute(backend, "embedding_timeout")?
|
||||
.map(|value| value.extract::<Option<f64>>())
|
||||
.transpose()?
|
||||
.flatten()
|
||||
.map(duration)
|
||||
.transpose()?;
|
||||
Ok(Ok(QdrantSemanticCacheConfig {
|
||||
grpc_url: grpc_url.to_string().trim_end_matches('/').to_owned(),
|
||||
api_key: optional_string(backend.getattr("qdrant_api_key")?)?,
|
||||
collection_name: backend.getattr("collection_name")?.extract()?,
|
||||
similarity_threshold: backend.getattr("similarity_threshold")?.extract()?,
|
||||
vector_size: backend.getattr("vector_size")?.extract::<u64>()?,
|
||||
embedding: OpenAiEmbedderConfig {
|
||||
api_base: embedding_api_base,
|
||||
api_key: embedding_api_key,
|
||||
model: embedding_model,
|
||||
timeout,
|
||||
},
|
||||
quantization: Quantization::Binary,
|
||||
}))
|
||||
}
|
||||
|
||||
fn py_sys_module(py: Python<'_>) -> PyResult<Option<Bound<'_, PyAny>>> {
|
||||
match py
|
||||
.import("sys")?
|
||||
.getattr("modules")?
|
||||
.get_item("litellm.proxy.proxy_server")
|
||||
{
|
||||
Ok(module) => Ok(Some(module)),
|
||||
Err(error) if error.is_instance_of::<pyo3::exceptions::PyKeyError>(py) => Ok(None),
|
||||
Err(error) => Err(error),
|
||||
}
|
||||
}
|
||||
|
||||
#[inline(never)]
|
||||
fn project_azure_blob(backend: &Bound<'_, PyAny>) -> PyResult<AzureBlobCacheConfig> {
|
||||
let client = backend.getattr("container_client")?;
|
||||
|
|
@ -371,6 +464,27 @@ fn project_azure_blob(backend: &Bound<'_, PyAny>) -> PyResult<AzureBlobCacheConf
|
|||
})
|
||||
}
|
||||
|
||||
#[inline(never)]
|
||||
pub(super) fn project_redis_semantic(
|
||||
backend: &Bound<'_, PyAny>,
|
||||
) -> PyResult<RedisSemanticCacheConfig> {
|
||||
Ok(RedisSemanticCacheConfig {
|
||||
redis_url: backend.getattr("_redis_url")?.extract::<String>()?,
|
||||
index_name: backend
|
||||
.getattr("_index_name")?
|
||||
.extract::<Option<String>>()?
|
||||
.unwrap_or_else(|| "litellm_semantic_cache_index".into()),
|
||||
similarity_threshold: backend.getattr("similarity_threshold")?.extract::<f64>()?,
|
||||
embedding_model: backend.getattr("embedding_model")?.extract::<String>()?,
|
||||
embedding_max_input_tokens: backend
|
||||
.getattr("embedding_max_input_tokens")?
|
||||
.extract::<Option<u64>>()?,
|
||||
embedding_timeout: backend
|
||||
.getattr("embedding_timeout")?
|
||||
.extract::<Option<f64>>()?,
|
||||
})
|
||||
}
|
||||
|
||||
#[inline(never)]
|
||||
fn project_memory(backend: &Bound<'_, PyAny>) -> PyResult<MemoryCacheConfig> {
|
||||
let max_size_kib = backend.getattr("max_size_per_item")?.extract::<usize>()?;
|
||||
|
|
@ -961,12 +1075,13 @@ mod tests {
|
|||
use pyo3::{prelude::*, types::PyDict};
|
||||
|
||||
use litellm_cache_redis::{RedisNode, RedisTopology};
|
||||
use litellm_cache_redis_semantic::RedisSemanticConfig;
|
||||
|
||||
use super::{
|
||||
CacheBackendConfig, CacheConfigProjection, CertificateRequirement, GcsCacheConfig,
|
||||
NativeCacheConfig, RedisProtocol, UnsupportedCacheConfig,
|
||||
CacheBackendConfig, CacheConfigProjection, CachePolicy, CertificateRequirement,
|
||||
GcsCacheConfig, NativeCacheConfig, RedisProtocol, UnsupportedCacheConfig,
|
||||
};
|
||||
use crate::cache::native::NativeResponseCache;
|
||||
use crate::cache::{embedder::PythonEmbedder, native::NativeResponseCache};
|
||||
|
||||
fn cluster_facade<'py>(py: Python<'py>, startup_nodes: &str, hook: &str) -> Bound<'py, PyAny> {
|
||||
facade(
|
||||
|
|
@ -1039,6 +1154,49 @@ mod tests {
|
|||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn redis_semantic_service_mismatch_accepts_backend_precision_threshold() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let facade = facade(
|
||||
py,
|
||||
"backend = SimpleNamespace(_redis_url='redis://127.0.0.1/', _index_name='semantic_idx', similarity_threshold=0.8, embedding_model='text-embedding-3-small', embedding_max_input_tokens=None, embedding_timeout=None)\n\
|
||||
facade = SimpleNamespace(type='redis-semantic', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)",
|
||||
);
|
||||
let backend = facade.getattr("cache").unwrap();
|
||||
let embedder = PythonEmbedder::new(backend.clone().unbind());
|
||||
let CacheConfigProjection::Native(config) =
|
||||
NativeCacheConfig::project(&facade).unwrap()
|
||||
else {
|
||||
panic!("Redis semantic cache should be supported");
|
||||
};
|
||||
let CacheBackendConfig::RedisSemantic(config) = config.backend else {
|
||||
panic!("expected Redis semantic configuration");
|
||||
};
|
||||
let service = NativeResponseCache::redis_semantic(
|
||||
&config.redis_url,
|
||||
embedder,
|
||||
RedisSemanticConfig {
|
||||
index_name: config.index_name.clone(),
|
||||
similarity_threshold: config.similarity_threshold as f32,
|
||||
},
|
||||
)
|
||||
.unwrap();
|
||||
let matching_config = NativeCacheConfig {
|
||||
policy: CachePolicy {
|
||||
mode: "default-on".into(),
|
||||
ttl: None,
|
||||
namespace: None,
|
||||
supported_call_types: None,
|
||||
redis_flush_size: None,
|
||||
semantic_cache_scope: "key".into(),
|
||||
},
|
||||
backend: CacheBackendConfig::RedisSemantic(config),
|
||||
};
|
||||
assert_eq!(matching_config.service_mismatch(&service), None);
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn projects_resolved_redis_tls_configuration() {
|
||||
Python::initialize();
|
||||
|
|
|
|||
|
|
@ -1,63 +1,144 @@
|
|||
use std::{future::Future, sync::Arc};
|
||||
use std::future::Future;
|
||||
|
||||
use litellm_cache::Error;
|
||||
use litellm_cache_valkey_semantic::Embedder;
|
||||
use litellm_host_python::to_py;
|
||||
use pyo3::{PyTraverseError, PyVisit, prelude::*};
|
||||
use pyo3::{PyTraverseError, PyVisit, prelude::*, types::PyDict};
|
||||
use serde_json::Value;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(super) struct PythonEmbedder {
|
||||
sync_embed: Arc<Py<PyAny>>,
|
||||
async_embed_callable: Arc<Py<PyAny>>,
|
||||
tokio::task_local! {
|
||||
static PREPARED_EMBEDDING: Result<Vec<f32>, Error>;
|
||||
}
|
||||
|
||||
/// Runs `future` with the vector the Python embedder already produced, so the backend's
|
||||
/// `async_embed` never has to call back into Python from the runtime.
|
||||
pub(super) fn with_prepared_embedding<F: Future>(
|
||||
vector: Result<Vec<f32>, Error>,
|
||||
future: F,
|
||||
) -> impl Future<Output = F::Output> {
|
||||
PREPARED_EMBEDDING.scope(vector, future)
|
||||
}
|
||||
|
||||
/// The Python object that owns embedding for a semantic backend.
|
||||
pub(super) struct PythonEmbedder(Py<PyAny>);
|
||||
|
||||
impl Clone for PythonEmbedder {
|
||||
fn clone(&self) -> Self {
|
||||
Python::attach(|py| Self(self.0.clone_ref(py)))
|
||||
}
|
||||
}
|
||||
|
||||
impl PythonEmbedder {
|
||||
pub(super) fn from_backend(backend: &Bound<'_, PyAny>) -> PyResult<Self> {
|
||||
Ok(Self {
|
||||
sync_embed: Arc::new(backend.getattr("_get_embedding")?.unbind()),
|
||||
async_embed_callable: Arc::new(backend.getattr("_get_async_embedding")?.unbind()),
|
||||
})
|
||||
pub(super) fn new(object: Py<PyAny>) -> Self {
|
||||
Self(object)
|
||||
}
|
||||
|
||||
pub(super) fn async_embed_awaitable<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
prompt: &str,
|
||||
metadata: &Option<Value>,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let metadata = to_py(py, metadata)?;
|
||||
self.async_embed_callable.bind(py).call1((prompt, metadata))
|
||||
pub(super) fn object(&self) -> &Py<PyAny> {
|
||||
&self.0
|
||||
}
|
||||
|
||||
pub(super) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
|
||||
visit.call(&*self.sync_embed)?;
|
||||
visit.call(&*self.async_embed_callable)
|
||||
visit.call(&self.0)
|
||||
}
|
||||
|
||||
fn metadata_kwargs<'py>(
|
||||
py: Python<'py>,
|
||||
metadata: Option<&Value>,
|
||||
) -> PyResult<Bound<'py, PyDict>> {
|
||||
let kwargs = PyDict::new(py);
|
||||
kwargs.set_item("metadata", to_py(py, &metadata)?)?;
|
||||
Ok(kwargs)
|
||||
}
|
||||
|
||||
/// The awaitable of `_get_async_embedding(prompt, metadata=...)`, to run in the caller's loop.
|
||||
pub(super) fn async_embedding(
|
||||
&self,
|
||||
py: Python<'_>,
|
||||
prompt: &str,
|
||||
metadata: Option<&Value>,
|
||||
) -> PyResult<Py<PyAny>> {
|
||||
let kwargs = Self::metadata_kwargs(py, metadata)?;
|
||||
self.0
|
||||
.bind(py)
|
||||
.call_method("_get_async_embedding", (prompt,), Some(&kwargs))
|
||||
.map(Bound::unbind)
|
||||
}
|
||||
|
||||
pub(super) fn extract(vector: Bound<'_, PyAny>) -> PyResult<Vec<f32>> {
|
||||
Ok(vector
|
||||
.extract::<Vec<f64>>()?
|
||||
.into_iter()
|
||||
.map(|value| value as f32)
|
||||
.collect())
|
||||
}
|
||||
|
||||
fn embed_sync(&self, prompt: &str, metadata: Option<&Value>) -> Result<Vec<f32>, Error> {
|
||||
Python::attach(|py| {
|
||||
let kwargs = Self::metadata_kwargs(py, metadata)?;
|
||||
Self::extract(self.0.bind(py).call_method(
|
||||
"_get_embedding",
|
||||
(prompt,),
|
||||
Some(&kwargs),
|
||||
)?)
|
||||
})
|
||||
.map_err(|_| Error::Unavailable)
|
||||
}
|
||||
|
||||
fn seeded_embedding() -> Result<Vec<f32>, Error> {
|
||||
PREPARED_EMBEDDING
|
||||
.try_with(Clone::clone)
|
||||
.unwrap_or(Err(Error::Unavailable))
|
||||
}
|
||||
}
|
||||
|
||||
impl Embedder for PythonEmbedder {
|
||||
impl litellm_cache_valkey_semantic::Embedder for PythonEmbedder {
|
||||
fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result<Vec<f32>, Error> {
|
||||
let result = Python::attach(|py| -> PyResult<Vec<f64>> {
|
||||
let metadata = to_py(py, &metadata)?;
|
||||
self.sync_embed
|
||||
.bind(py)
|
||||
.call1((prompt, metadata))?
|
||||
.extract()
|
||||
})
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
Ok(result.into_iter().map(|value| value as f32).collect())
|
||||
self.embed_sync(prompt, metadata)
|
||||
}
|
||||
|
||||
#[expect(
|
||||
clippy::manual_async_fn,
|
||||
reason = "the shared Embedder trait uses an impl Future return"
|
||||
)]
|
||||
fn async_embed(
|
||||
&self,
|
||||
_prompt: &str,
|
||||
_metadata: Option<&Value>,
|
||||
) -> impl Future<Output = Result<Vec<f32>, Error>> + Send {
|
||||
async { Err(Error::Unavailable) }
|
||||
std::future::ready(Self::seeded_embedding())
|
||||
}
|
||||
}
|
||||
|
||||
impl litellm_cache_redis_semantic::Embedder for PythonEmbedder {
|
||||
fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result<Vec<f32>, Error> {
|
||||
self.embed_sync(prompt, metadata)
|
||||
}
|
||||
|
||||
fn async_embed(
|
||||
&self,
|
||||
_prompt: &str,
|
||||
_metadata: Option<&Value>,
|
||||
) -> impl Future<Output = Result<Vec<f32>, Error>> + Send {
|
||||
std::future::ready(Self::seeded_embedding())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn async_embed_returns_the_seeded_vector_or_unavailable() {
|
||||
Python::initialize();
|
||||
let object = Python::attach(|py| py.None());
|
||||
let embedder = PythonEmbedder::new(object);
|
||||
let scoped_embedder = embedder.clone();
|
||||
let scoped = with_prepared_embedding(Ok(vec![0.25]), async move {
|
||||
litellm_cache_redis_semantic::Embedder::async_embed(&scoped_embedder, "prompt", None)
|
||||
.await
|
||||
});
|
||||
assert_eq!(scoped.await, Ok(vec![0.25]));
|
||||
let unscoped =
|
||||
litellm_cache_redis_semantic::Embedder::async_embed(&embedder, "prompt", None).await;
|
||||
assert_eq!(unscoped, Err(Error::Unavailable));
|
||||
let valkey = with_prepared_embedding(Ok(vec![0.5]), async move {
|
||||
litellm_cache_valkey_semantic::Embedder::async_embed(&embedder, "prompt", None).await
|
||||
});
|
||||
assert_eq!(valkey.await, Ok(vec![0.5]));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ use serde_json::Value;
|
|||
use super::{
|
||||
config::{CacheConfigProjection, NativeCacheConfig},
|
||||
handle::CacheTestHandle,
|
||||
identity::BackendIdentity,
|
||||
native::NativeResponseCache,
|
||||
};
|
||||
|
||||
|
|
@ -352,37 +353,41 @@ impl FacadeGuard {
|
|||
facade: &Bound<'_, PyAny>,
|
||||
service: &NativeResponseCache,
|
||||
) -> PyResult<Self> {
|
||||
let kind = service.kind();
|
||||
let identity = service.identity();
|
||||
let kind = identity.kind();
|
||||
let cache_type = py.import("litellm.caching.caching")?.getattr("Cache")?;
|
||||
if !facade.get_type().is(&cache_type) {
|
||||
return Err(PyTypeError::new_err(
|
||||
"only exact built-in Cache facades can be registered",
|
||||
));
|
||||
}
|
||||
let cluster = matches!(service.topology(), Some(RedisTopology::Cluster { .. }));
|
||||
let (module, name, cache_kind) = match (kind, cluster) {
|
||||
("memory", _) => ("litellm.caching.in_memory_cache", "InMemoryCache", "local"),
|
||||
("redis", false) => ("litellm.caching.redis_cache", "RedisCache", "redis"),
|
||||
("redis", true) => (
|
||||
"litellm.caching.redis_cluster_cache",
|
||||
"RedisClusterCache",
|
||||
"redis",
|
||||
let cluster = matches!(
|
||||
identity,
|
||||
BackendIdentity::Redis {
|
||||
topology: RedisTopology::Cluster { .. },
|
||||
..
|
||||
}
|
||||
);
|
||||
let (module, name) = match (kind, cluster) {
|
||||
("memory", _) => ("litellm.caching.in_memory_cache", "InMemoryCache"),
|
||||
("redis", false) => ("litellm.caching.redis_cache", "RedisCache"),
|
||||
("redis", true) => ("litellm.caching.redis_cluster_cache", "RedisClusterCache"),
|
||||
("redis_semantic", _) => ("litellm.caching.redis_semantic_cache", "RedisSemanticCache"),
|
||||
("qdrant_semantic", _) => (
|
||||
"litellm.caching.qdrant_semantic_cache",
|
||||
"QdrantSemanticCache",
|
||||
),
|
||||
("gcs", _) => ("litellm.caching.gcs_cache", "GCSCache", "gcs"),
|
||||
("valkey-semantic", false) => (
|
||||
("gcs", _) => ("litellm.caching.gcs_cache", "GCSCache"),
|
||||
("valkey-semantic", _) => (
|
||||
"litellm.caching.valkey_semantic_cache",
|
||||
"ValkeySemanticCache",
|
||||
"valkey-semantic",
|
||||
),
|
||||
("disk", _) => ("litellm.caching.disk_cache", "DiskCache", "disk"),
|
||||
("azure-blob", _) => (
|
||||
"litellm.caching.azure_blob_cache",
|
||||
"AzureBlobCache",
|
||||
"azure-blob",
|
||||
),
|
||||
("s3", _) => ("litellm.caching.s3_cache", "S3Cache", "s3"),
|
||||
("disk", _) => ("litellm.caching.disk_cache", "DiskCache"),
|
||||
("azure-blob", _) => ("litellm.caching.azure_blob_cache", "AzureBlobCache"),
|
||||
("s3", _) => ("litellm.caching.s3_cache", "S3Cache"),
|
||||
_ => unreachable!(),
|
||||
};
|
||||
let cache_kind = identity.cache_type();
|
||||
let backend = facade.getattr("cache")?;
|
||||
if facade.getattr("type")?.extract::<String>()? != cache_kind
|
||||
|| !backend.get_type().is(&py.import(module)?.getattr(name)?)
|
||||
|
|
@ -400,6 +405,15 @@ impl FacadeGuard {
|
|||
if let Some(message) = config.service_mismatch(service) {
|
||||
return Err(PyTypeError::new_err(message));
|
||||
}
|
||||
if kind == "redis_semantic"
|
||||
&& service
|
||||
.embedder_object()
|
||||
.is_none_or(|embedder| !backend.is(embedder.bind(py)))
|
||||
{
|
||||
return Err(PyTypeError::new_err(
|
||||
"facade backend must be the native embedder",
|
||||
));
|
||||
}
|
||||
Ok(Self {
|
||||
outer: ObjectGuard::capture(
|
||||
py,
|
||||
|
|
@ -425,6 +439,17 @@ impl FacadeGuard {
|
|||
"redis_kwargs",
|
||||
"redis_flush_size",
|
||||
"similarity_threshold",
|
||||
"distance_threshold",
|
||||
"embedding_model",
|
||||
"embedding_max_input_tokens",
|
||||
"embedding_timeout",
|
||||
"qdrant_api_base",
|
||||
"qdrant_api_key",
|
||||
"collection_name",
|
||||
"vector_size",
|
||||
"_index_name",
|
||||
"_redis_url",
|
||||
"similarity_threshold",
|
||||
"embedding_model",
|
||||
"index_name",
|
||||
"embedding_max_input_tokens",
|
||||
|
|
|
|||
|
|
@ -1,12 +1,25 @@
|
|||
use litellm_auth_aws::AwsAuthConfig;
|
||||
use litellm_cache_gcs::{DEFAULT_ENDPOINT, GcsConfig};
|
||||
use litellm_cache_qdrant_semantic::{OpenAiEmbedderConfig, Quantization};
|
||||
use litellm_cache_redis::{RedisNode, RedisTopology};
|
||||
use litellm_cache_redis_semantic::RedisSemanticConfig;
|
||||
use litellm_cache_s3::{S3CacheConfig, S3Endpoint};
|
||||
use litellm_host_python::{release_gil, run_sync_value};
|
||||
use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*};
|
||||
use litellm_http::ClientVariant;
|
||||
use pyo3::{
|
||||
PyTraverseError, PyVisit,
|
||||
exceptions::{PyRuntimeError, PyTypeError},
|
||||
prelude::*,
|
||||
types::PyDict,
|
||||
};
|
||||
use url::Url;
|
||||
|
||||
use super::{
|
||||
cache_error, embedder::PythonEmbedder, facade::FacadeGuard, native::NativeResponseCache,
|
||||
cache_error,
|
||||
config::{QdrantSemanticCacheConfig, project_redis_semantic},
|
||||
embedder::PythonEmbedder,
|
||||
facade::FacadeGuard,
|
||||
native::NativeResponseCache,
|
||||
request::duration,
|
||||
};
|
||||
|
||||
|
|
@ -141,6 +154,105 @@ impl CacheTestHandle {
|
|||
})
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
#[pyo3(signature = (url, *, collection_name, similarity_threshold, vector_size, embedding_model="text-embedding-3-small", api_key=None, embedding_api_key=None, embedding_api_base=None, embedding_timeout_seconds=None, quantization="binary"))]
|
||||
#[expect(
|
||||
clippy::too_many_arguments,
|
||||
reason = "the test handle exposes the complete Qdrant constructor"
|
||||
)]
|
||||
fn qdrant_semantic(
|
||||
py: Python<'_>,
|
||||
url: String,
|
||||
collection_name: String,
|
||||
similarity_threshold: f64,
|
||||
vector_size: u64,
|
||||
embedding_model: &str,
|
||||
api_key: Option<String>,
|
||||
embedding_api_key: Option<String>,
|
||||
embedding_api_base: Option<String>,
|
||||
embedding_timeout_seconds: Option<f64>,
|
||||
quantization: &str,
|
||||
) -> PyResult<Self> {
|
||||
let parsed = Url::parse(&url).map_err(|_| {
|
||||
pyo3::exceptions::PyValueError::new_err(
|
||||
"native Qdrant requires the default REST port so the gRPC port can be derived",
|
||||
)
|
||||
})?;
|
||||
if !matches!(parsed.scheme(), "http" | "https")
|
||||
|| (!parsed.path().is_empty() && parsed.path() != "/")
|
||||
|| parsed.query().is_some()
|
||||
|| parsed.host_str().is_none()
|
||||
|| parsed.port() != Some(6333)
|
||||
{
|
||||
return Err(pyo3::exceptions::PyValueError::new_err(
|
||||
"native Qdrant requires the default REST port so the gRPC port can be derived",
|
||||
));
|
||||
}
|
||||
let mut grpc_url = parsed;
|
||||
grpc_url.set_port(Some(6334)).map_err(|_| {
|
||||
pyo3::exceptions::PyValueError::new_err(
|
||||
"native Qdrant requires the default REST port so the gRPC port can be derived",
|
||||
)
|
||||
})?;
|
||||
grpc_url.set_path("");
|
||||
grpc_url.set_query(None);
|
||||
let embedding_api_key = embedding_api_key
|
||||
.or_else(|| {
|
||||
std::env::var("OPENAI_API_KEY")
|
||||
.ok()
|
||||
.filter(|value| !value.is_empty())
|
||||
})
|
||||
.ok_or_else(|| {
|
||||
pyo3::exceptions::PyValueError::new_err(
|
||||
"native semantic embedding requires an OpenAI API key",
|
||||
)
|
||||
})?;
|
||||
let embedding_api_base = embedding_api_base.unwrap_or_else(|| {
|
||||
std::env::var("OPENAI_BASE_URL")
|
||||
.or_else(|_| std::env::var("OPENAI_API_BASE"))
|
||||
.unwrap_or_else(|_| "https://api.openai.com/v1".to_owned())
|
||||
});
|
||||
let quantization = match quantization {
|
||||
"binary" => Quantization::Binary,
|
||||
"scalar" => Quantization::Scalar,
|
||||
"product" => Quantization::Product,
|
||||
_ => {
|
||||
return Err(pyo3::exceptions::PyValueError::new_err(
|
||||
"unsupported Qdrant quantization",
|
||||
));
|
||||
}
|
||||
};
|
||||
let config = QdrantSemanticCacheConfig {
|
||||
grpc_url: grpc_url.to_string().trim_end_matches('/').to_owned(),
|
||||
api_key,
|
||||
collection_name,
|
||||
similarity_threshold,
|
||||
vector_size,
|
||||
embedding: OpenAiEmbedderConfig {
|
||||
api_base: embedding_api_base,
|
||||
api_key: embedding_api_key,
|
||||
model: embedding_model.to_owned(),
|
||||
timeout: embedding_timeout_seconds.map(duration).transpose()?,
|
||||
},
|
||||
quantization,
|
||||
};
|
||||
let http_config = crate::http::call_config(py, &PyDict::new(py), true)?;
|
||||
let client = crate::http::pool()
|
||||
.client(&http_config, ClientVariant::Provider)
|
||||
.map_err(crate::http::client_error)?;
|
||||
let service = run_sync_value(py, async move {
|
||||
let handle = tokio::runtime::Handle::current();
|
||||
NativeResponseCache::qdrant_semantic(config, client, handle)
|
||||
.await
|
||||
.map_err(cache_error)
|
||||
})?;
|
||||
Ok(Self {
|
||||
service,
|
||||
guard: None,
|
||||
pid: std::process::id(),
|
||||
})
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
#[pyo3(signature = (url, similarity_threshold, index_name, embedder))]
|
||||
fn valkey_semantic(
|
||||
|
|
@ -149,7 +261,7 @@ impl CacheTestHandle {
|
|||
index_name: String,
|
||||
embedder: &Bound<'_, PyAny>,
|
||||
) -> PyResult<Self> {
|
||||
let python_embedder = PythonEmbedder::from_backend(embedder)?;
|
||||
let python_embedder = PythonEmbedder::new(embedder.clone().unbind());
|
||||
let service = NativeResponseCache::valkey_semantic(
|
||||
&url,
|
||||
similarity_threshold,
|
||||
|
|
@ -179,6 +291,36 @@ impl CacheTestHandle {
|
|||
})
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
fn redis_semantic(py: Python<'_>, backend: Bound<'_, PyAny>) -> PyResult<Self> {
|
||||
let class = py
|
||||
.import("litellm.caching.redis_semantic_cache")?
|
||||
.getattr("RedisSemanticCache")?;
|
||||
if !backend.get_type().is(&class) {
|
||||
return Err(PyTypeError::new_err(
|
||||
"native redis-semantic handles require the built-in RedisSemanticCache",
|
||||
));
|
||||
}
|
||||
let config = project_redis_semantic(&backend)?;
|
||||
let embedder = PythonEmbedder::new(backend.unbind());
|
||||
let service = release_gil(py, move || {
|
||||
NativeResponseCache::redis_semantic(
|
||||
&config.redis_url,
|
||||
embedder,
|
||||
RedisSemanticConfig {
|
||||
index_name: config.index_name,
|
||||
similarity_threshold: config.similarity_threshold as f32,
|
||||
},
|
||||
)
|
||||
})
|
||||
.map_err(cache_error)?;
|
||||
Ok(Self {
|
||||
service,
|
||||
guard: None,
|
||||
pid: std::process::id(),
|
||||
})
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn backend(&self) -> &'static str {
|
||||
self.service.kind()
|
||||
|
|
@ -210,6 +352,7 @@ impl CacheTestHandle {
|
|||
}
|
||||
|
||||
fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
|
||||
self.service.traverse(&visit)?;
|
||||
if let Some(guard) = &self.guard {
|
||||
guard.traverse(visit)?;
|
||||
}
|
||||
|
|
|
|||
511
litellm-rust/crates/python-bridge/src/cache/identity.rs
vendored
Normal file
511
litellm-rust/crates/python-bridge/src/cache/identity.rs
vendored
Normal file
|
|
@ -0,0 +1,511 @@
|
|||
use std::{path::PathBuf, time::Duration};
|
||||
|
||||
use litellm_cache_redis::RedisTopology;
|
||||
|
||||
/// What makes a native backend the one a Python facade describes: the configuration a user can
|
||||
/// observe on the Python object, captured once so facade projection and native construction
|
||||
/// compare plain data instead of reaching into each backend type.
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub(super) enum BackendIdentity {
|
||||
Memory {
|
||||
capacity: usize,
|
||||
max_entry_bytes: Option<usize>,
|
||||
default_ttl: Option<Duration>,
|
||||
},
|
||||
Redis {
|
||||
topology: RedisTopology,
|
||||
namespace: Option<String>,
|
||||
default_ttl: Option<Duration>,
|
||||
},
|
||||
S3 {
|
||||
bucket: String,
|
||||
key_prefix: String,
|
||||
region: String,
|
||||
endpoint: Option<String>,
|
||||
},
|
||||
Gcs {
|
||||
bucket_name: String,
|
||||
key_prefix: String,
|
||||
path_service_account: Option<String>,
|
||||
},
|
||||
Disk {
|
||||
directory: PathBuf,
|
||||
},
|
||||
AzureBlob {
|
||||
account_url: String,
|
||||
container: String,
|
||||
},
|
||||
RedisSemantic {
|
||||
index_name: String,
|
||||
/// The backend stores the threshold as `f32`; a facade's `f64` is compared at that width.
|
||||
similarity_threshold: f32,
|
||||
},
|
||||
ValkeySemantic {
|
||||
index_name: String,
|
||||
similarity_threshold: f64,
|
||||
},
|
||||
QdrantSemantic {
|
||||
collection_name: String,
|
||||
similarity_threshold: f64,
|
||||
vector_size: u64,
|
||||
embedding_model: String,
|
||||
},
|
||||
}
|
||||
|
||||
const TYPES: &str = "facade and native backend types must match";
|
||||
|
||||
impl BackendIdentity {
|
||||
/// The native backend name reported to Python through `_CacheTestHandle.backend`.
|
||||
pub(super) fn kind(&self) -> &'static str {
|
||||
match self {
|
||||
Self::Memory { .. } => "memory",
|
||||
Self::Redis { .. } => "redis",
|
||||
Self::S3 { .. } => "s3",
|
||||
Self::Gcs { .. } => "gcs",
|
||||
Self::ValkeySemantic { .. } => "valkey-semantic",
|
||||
Self::RedisSemantic { .. } => "redis_semantic",
|
||||
Self::QdrantSemantic { .. } => "qdrant_semantic",
|
||||
Self::Disk { .. } => "disk",
|
||||
Self::AzureBlob { .. } => "azure-blob",
|
||||
}
|
||||
}
|
||||
|
||||
/// The `LiteLLMCacheType` value a facade of this backend carries in `Cache.type`.
|
||||
pub(super) fn cache_type(&self) -> &'static str {
|
||||
match self {
|
||||
Self::Memory { .. } => "local",
|
||||
Self::Redis { .. } => "redis",
|
||||
Self::S3 { .. } => "s3",
|
||||
Self::Gcs { .. } => "gcs",
|
||||
Self::ValkeySemantic { .. } => "valkey-semantic",
|
||||
Self::RedisSemantic { .. } => "redis-semantic",
|
||||
Self::QdrantSemantic { .. } => "qdrant-semantic",
|
||||
Self::Disk { .. } => "disk",
|
||||
Self::AzureBlob { .. } => "azure-blob",
|
||||
}
|
||||
}
|
||||
|
||||
/// The first difference between the facade's configuration (`self`) and the native
|
||||
/// backend (`native`), in the order Python users see the attributes.
|
||||
pub(super) fn mismatch(&self, native: &Self) -> Option<&'static str> {
|
||||
let mut differences: Vec<(bool, &'static str)> = Vec::new();
|
||||
let mut differs = |condition: bool, message: &'static str| {
|
||||
differences.push((condition, message));
|
||||
};
|
||||
match (self, native) {
|
||||
(
|
||||
Self::Memory {
|
||||
capacity,
|
||||
max_entry_bytes,
|
||||
default_ttl,
|
||||
},
|
||||
Self::Memory {
|
||||
capacity: native_capacity,
|
||||
max_entry_bytes: native_max_entry_bytes,
|
||||
default_ttl: native_default_ttl,
|
||||
},
|
||||
) => {
|
||||
differs(
|
||||
default_ttl != native_default_ttl,
|
||||
"facade and native backend default TTLs must match",
|
||||
);
|
||||
differs(
|
||||
capacity != native_capacity,
|
||||
"facade and native backend capacities must match",
|
||||
);
|
||||
differs(
|
||||
max_entry_bytes != native_max_entry_bytes,
|
||||
"facade and native backend item limits must match",
|
||||
);
|
||||
}
|
||||
(
|
||||
Self::Redis {
|
||||
topology,
|
||||
namespace,
|
||||
default_ttl,
|
||||
},
|
||||
Self::Redis {
|
||||
topology: native_topology,
|
||||
namespace: native_namespace,
|
||||
default_ttl: native_default_ttl,
|
||||
},
|
||||
) => {
|
||||
differs(
|
||||
default_ttl != native_default_ttl,
|
||||
"facade and native backend default TTLs must match",
|
||||
);
|
||||
differs(
|
||||
topology != native_topology,
|
||||
"facade and native backend topologies must match",
|
||||
);
|
||||
differs(
|
||||
namespace != native_namespace,
|
||||
"facade and native backend namespaces must match",
|
||||
);
|
||||
}
|
||||
(
|
||||
Self::S3 {
|
||||
bucket,
|
||||
key_prefix,
|
||||
region,
|
||||
endpoint,
|
||||
},
|
||||
Self::S3 {
|
||||
bucket: native_bucket,
|
||||
key_prefix: native_key_prefix,
|
||||
region: native_region,
|
||||
endpoint: native_endpoint,
|
||||
},
|
||||
) => {
|
||||
differs(
|
||||
bucket != native_bucket,
|
||||
"facade and native backend buckets must match",
|
||||
);
|
||||
differs(
|
||||
key_prefix != native_key_prefix,
|
||||
"facade and native backend key prefixes must match",
|
||||
);
|
||||
differs(
|
||||
region != native_region,
|
||||
"facade and native backend regions must match",
|
||||
);
|
||||
differs(
|
||||
endpoint != native_endpoint,
|
||||
"facade and native backend endpoints must match",
|
||||
);
|
||||
}
|
||||
(
|
||||
Self::Gcs {
|
||||
bucket_name,
|
||||
key_prefix,
|
||||
path_service_account,
|
||||
},
|
||||
Self::Gcs {
|
||||
bucket_name: native_bucket_name,
|
||||
key_prefix: native_key_prefix,
|
||||
path_service_account: native_path_service_account,
|
||||
},
|
||||
) => {
|
||||
differs(
|
||||
bucket_name != native_bucket_name,
|
||||
"facade and native backend buckets must match",
|
||||
);
|
||||
differs(
|
||||
key_prefix != native_key_prefix,
|
||||
"facade and native backend key prefixes must match",
|
||||
);
|
||||
differs(
|
||||
path_service_account != native_path_service_account,
|
||||
"facade and native backend credentials must match",
|
||||
);
|
||||
}
|
||||
(
|
||||
Self::Disk { directory },
|
||||
Self::Disk {
|
||||
directory: native_directory,
|
||||
},
|
||||
) => {
|
||||
let canonical = |path: &PathBuf| std::fs::canonicalize(path).ok();
|
||||
differs(
|
||||
canonical(directory) != canonical(native_directory),
|
||||
"facade and native backend directories must match",
|
||||
);
|
||||
}
|
||||
(
|
||||
Self::AzureBlob {
|
||||
account_url,
|
||||
container,
|
||||
},
|
||||
Self::AzureBlob {
|
||||
account_url: native_account_url,
|
||||
container: native_container,
|
||||
},
|
||||
) => {
|
||||
differs(
|
||||
account_url != native_account_url || container != native_container,
|
||||
"facade and native backend containers must match",
|
||||
);
|
||||
}
|
||||
(
|
||||
Self::RedisSemantic {
|
||||
index_name,
|
||||
similarity_threshold,
|
||||
},
|
||||
Self::RedisSemantic {
|
||||
index_name: native_index_name,
|
||||
similarity_threshold: native_similarity_threshold,
|
||||
},
|
||||
) => {
|
||||
differs(
|
||||
index_name != native_index_name,
|
||||
"facade and native backend index names must match",
|
||||
);
|
||||
differs(
|
||||
similarity_threshold != native_similarity_threshold,
|
||||
"facade and native backend similarity thresholds must match",
|
||||
);
|
||||
}
|
||||
(
|
||||
Self::ValkeySemantic {
|
||||
index_name,
|
||||
similarity_threshold,
|
||||
},
|
||||
Self::ValkeySemantic {
|
||||
index_name: native_index_name,
|
||||
similarity_threshold: native_similarity_threshold,
|
||||
},
|
||||
) => {
|
||||
differs(
|
||||
index_name != native_index_name
|
||||
|| similarity_threshold != native_similarity_threshold,
|
||||
"facade and native semantic settings must match",
|
||||
);
|
||||
}
|
||||
(
|
||||
Self::QdrantSemantic {
|
||||
collection_name,
|
||||
similarity_threshold,
|
||||
vector_size,
|
||||
embedding_model,
|
||||
},
|
||||
Self::QdrantSemantic {
|
||||
collection_name: native_collection_name,
|
||||
similarity_threshold: native_similarity_threshold,
|
||||
vector_size: native_vector_size,
|
||||
embedding_model: native_embedding_model,
|
||||
},
|
||||
) => {
|
||||
differs(
|
||||
collection_name != native_collection_name,
|
||||
"facade and native backend collections must match",
|
||||
);
|
||||
differs(
|
||||
similarity_threshold != native_similarity_threshold,
|
||||
"facade and native backend similarity thresholds must match",
|
||||
);
|
||||
differs(
|
||||
vector_size != native_vector_size,
|
||||
"facade and native backend vector sizes must match",
|
||||
);
|
||||
differs(
|
||||
embedding_model != native_embedding_model,
|
||||
"facade and native backend embedding models must match",
|
||||
);
|
||||
}
|
||||
_ => return Some(TYPES),
|
||||
}
|
||||
differences
|
||||
.into_iter()
|
||||
.find_map(|(condition, message)| condition.then_some(message))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::time::Duration;
|
||||
|
||||
use litellm_cache_redis::{RedisNode, RedisTopology};
|
||||
|
||||
use super::BackendIdentity;
|
||||
|
||||
fn memory() -> BackendIdentity {
|
||||
BackendIdentity::Memory {
|
||||
capacity: 200,
|
||||
max_entry_bytes: Some(1024),
|
||||
default_ttl: Some(Duration::from_secs(60)),
|
||||
}
|
||||
}
|
||||
|
||||
fn redis() -> BackendIdentity {
|
||||
BackendIdentity::Redis {
|
||||
topology: RedisTopology::Standalone,
|
||||
namespace: Some("team".into()),
|
||||
default_ttl: Some(Duration::from_secs(60)),
|
||||
}
|
||||
}
|
||||
|
||||
fn s3() -> BackendIdentity {
|
||||
BackendIdentity::S3 {
|
||||
bucket: "bucket".into(),
|
||||
key_prefix: "cache/".into(),
|
||||
region: "us-east-1".into(),
|
||||
endpoint: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn gcs() -> BackendIdentity {
|
||||
BackendIdentity::Gcs {
|
||||
bucket_name: "bucket".into(),
|
||||
key_prefix: "cache/".into(),
|
||||
path_service_account: Some("credentials.json".into()),
|
||||
}
|
||||
}
|
||||
|
||||
fn azure() -> BackendIdentity {
|
||||
BackendIdentity::AzureBlob {
|
||||
account_url: "https://account.blob.core.windows.net".into(),
|
||||
container: "cache".into(),
|
||||
}
|
||||
}
|
||||
|
||||
fn redis_semantic() -> BackendIdentity {
|
||||
BackendIdentity::RedisSemantic {
|
||||
index_name: "idx".into(),
|
||||
similarity_threshold: 0.8,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn redis_semantic_thresholds_compare_at_backend_precision() {
|
||||
let facade = BackendIdentity::RedisSemantic {
|
||||
index_name: "idx".into(),
|
||||
similarity_threshold: 0.8_f64 as f32,
|
||||
};
|
||||
assert_eq!(facade.mismatch(&redis_semantic()), None);
|
||||
}
|
||||
|
||||
fn valkey_semantic() -> BackendIdentity {
|
||||
BackendIdentity::ValkeySemantic {
|
||||
index_name: "idx".into(),
|
||||
similarity_threshold: 0.8,
|
||||
}
|
||||
}
|
||||
|
||||
fn qdrant() -> BackendIdentity {
|
||||
BackendIdentity::QdrantSemantic {
|
||||
collection_name: "collection".into(),
|
||||
similarity_threshold: 0.8,
|
||||
vector_size: 1536,
|
||||
embedding_model: "text-embedding-3-small".into(),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn identical_identities_have_no_mismatch() {
|
||||
for identity in [
|
||||
memory(),
|
||||
redis(),
|
||||
s3(),
|
||||
gcs(),
|
||||
azure(),
|
||||
redis_semantic(),
|
||||
valkey_semantic(),
|
||||
qdrant(),
|
||||
BackendIdentity::Disk {
|
||||
directory: std::env::temp_dir(),
|
||||
},
|
||||
] {
|
||||
assert_eq!(identity.mismatch(&identity), None, "{identity:?}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn different_kinds_report_a_type_mismatch() {
|
||||
assert_eq!(
|
||||
memory().mismatch(&redis()),
|
||||
Some("facade and native backend types must match")
|
||||
);
|
||||
assert_eq!(
|
||||
redis_semantic().mismatch(&valkey_semantic()),
|
||||
Some("facade and native backend types must match")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_first_differing_field_names_the_mismatch() {
|
||||
let BackendIdentity::Memory { capacity, .. } = memory() else {
|
||||
unreachable!()
|
||||
};
|
||||
assert_eq!(
|
||||
memory().mismatch(&BackendIdentity::Memory {
|
||||
capacity: capacity + 1,
|
||||
max_entry_bytes: Some(1),
|
||||
default_ttl: Some(Duration::from_secs(60)),
|
||||
}),
|
||||
Some("facade and native backend capacities must match")
|
||||
);
|
||||
assert_eq!(
|
||||
memory().mismatch(&BackendIdentity::Memory {
|
||||
capacity,
|
||||
max_entry_bytes: Some(1),
|
||||
default_ttl: Some(Duration::from_secs(61)),
|
||||
}),
|
||||
Some("facade and native backend default TTLs must match")
|
||||
);
|
||||
assert_eq!(
|
||||
redis().mismatch(&BackendIdentity::Redis {
|
||||
topology: RedisTopology::Cluster {
|
||||
startup_nodes: vec![RedisNode {
|
||||
host: "node".into(),
|
||||
port: 7000,
|
||||
}],
|
||||
},
|
||||
namespace: None,
|
||||
default_ttl: Some(Duration::from_secs(60)),
|
||||
}),
|
||||
Some("facade and native backend topologies must match")
|
||||
);
|
||||
assert_eq!(
|
||||
s3().mismatch(&BackendIdentity::S3 {
|
||||
bucket: "bucket".into(),
|
||||
key_prefix: "cache/".into(),
|
||||
region: "us-east-1".into(),
|
||||
endpoint: Some("http://localhost:9000".into()),
|
||||
}),
|
||||
Some("facade and native backend endpoints must match")
|
||||
);
|
||||
assert_eq!(
|
||||
gcs().mismatch(&BackendIdentity::Gcs {
|
||||
bucket_name: "bucket".into(),
|
||||
key_prefix: "cache/".into(),
|
||||
path_service_account: None,
|
||||
}),
|
||||
Some("facade and native backend credentials must match")
|
||||
);
|
||||
assert_eq!(
|
||||
azure().mismatch(&BackendIdentity::AzureBlob {
|
||||
account_url: "https://account.blob.core.windows.net".into(),
|
||||
container: "other".into(),
|
||||
}),
|
||||
Some("facade and native backend containers must match")
|
||||
);
|
||||
assert_eq!(
|
||||
valkey_semantic().mismatch(&BackendIdentity::ValkeySemantic {
|
||||
index_name: "idx".into(),
|
||||
similarity_threshold: 0.9,
|
||||
}),
|
||||
Some("facade and native semantic settings must match")
|
||||
);
|
||||
assert_eq!(
|
||||
qdrant().mismatch(&BackendIdentity::QdrantSemantic {
|
||||
collection_name: "collection".into(),
|
||||
similarity_threshold: 0.8,
|
||||
vector_size: 1536,
|
||||
embedding_model: "text-embedding-3-large".into(),
|
||||
}),
|
||||
Some("facade and native backend embedding models must match")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn disk_directories_compare_canonically() {
|
||||
let directory = std::env::temp_dir();
|
||||
let mut indirect = directory.clone();
|
||||
indirect.push(".");
|
||||
assert_eq!(
|
||||
BackendIdentity::Disk {
|
||||
directory: directory.clone()
|
||||
}
|
||||
.mismatch(&BackendIdentity::Disk {
|
||||
directory: indirect
|
||||
}),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
BackendIdentity::Disk { directory }.mismatch(&BackendIdentity::Disk {
|
||||
directory: "/definitely/missing".into()
|
||||
}),
|
||||
Some("facade and native backend directories must match")
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
@ -5,10 +5,11 @@ mod embedder;
|
|||
mod facade;
|
||||
mod future;
|
||||
mod handle;
|
||||
mod identity;
|
||||
mod native;
|
||||
mod request;
|
||||
mod resolver;
|
||||
mod semantic_step;
|
||||
mod semantic;
|
||||
|
||||
use litellm_cache::Error;
|
||||
use pyo3::{
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -1,7 +1,7 @@
|
|||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||
|
||||
use litellm_cache::ExactCacheContext;
|
||||
use litellm_cache_response::{CacheControls, CacheKeyInput, ResponseCacheRequest};
|
||||
use litellm_cache::{ExactCacheContext, SemanticCacheContext};
|
||||
use litellm_cache_response::{CacheControls, CacheKeyField, CacheKeyInput, ResponseCacheRequest};
|
||||
use litellm_host_python::from_py;
|
||||
use pyo3::{exceptions::PyValueError, prelude::*};
|
||||
use serde::Deserialize;
|
||||
|
|
@ -19,8 +19,10 @@ struct RequestInput {
|
|||
metadata: Option<Value>,
|
||||
litellm_metadata: Option<Value>,
|
||||
litellm_params: Option<Value>,
|
||||
scope: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(super) struct NativeRequest {
|
||||
pub(super) key: CacheKeyInput,
|
||||
pub(super) controls: CacheControls,
|
||||
|
|
@ -31,6 +33,100 @@ pub(super) struct NativeRequest {
|
|||
pub(super) metadata: Option<Value>,
|
||||
pub(super) litellm_metadata: Option<Value>,
|
||||
pub(super) litellm_params: Option<Value>,
|
||||
pub(super) scope: Option<String>,
|
||||
}
|
||||
|
||||
impl NativeRequest {
|
||||
pub(super) fn exact(&self) -> ResponseCacheRequest<ExactCacheContext> {
|
||||
ResponseCacheRequest {
|
||||
key: self.key.clone(),
|
||||
controls: self.controls,
|
||||
context: ExactCacheContext { ttl: self.ttl },
|
||||
max_age: self.max_age,
|
||||
}
|
||||
}
|
||||
|
||||
/// The request as a semantic backend that keys on the caller's scope sees it.
|
||||
pub(super) fn semantic(&self) -> ResponseCacheRequest<SemanticCacheContext> {
|
||||
self.semantic_with(self.key.clone(), self.scope.clone())
|
||||
}
|
||||
|
||||
/// The request keyed the way Python's Valkey semantic cache keys it: prompt fields drop out
|
||||
/// and the tenant identifiers for `scope` join the key.
|
||||
pub(super) fn scoped_semantic(
|
||||
&self,
|
||||
scope: &str,
|
||||
) -> ResponseCacheRequest<SemanticCacheContext> {
|
||||
self.semantic_with(semantic_key(self, scope), Some(scope.to_owned()))
|
||||
}
|
||||
|
||||
fn semantic_with(
|
||||
&self,
|
||||
key: CacheKeyInput,
|
||||
scope: Option<String>,
|
||||
) -> ResponseCacheRequest<SemanticCacheContext> {
|
||||
ResponseCacheRequest {
|
||||
key,
|
||||
controls: self.controls,
|
||||
context: SemanticCacheContext {
|
||||
input: self.input.clone(),
|
||||
messages: self.messages.clone(),
|
||||
metadata: self.metadata.clone(),
|
||||
scope,
|
||||
ttl: self.ttl,
|
||||
},
|
||||
max_age: self.max_age,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn semantic_key(request: &NativeRequest, scope: &str) -> CacheKeyInput {
|
||||
let mut key = request.key.clone();
|
||||
if key.preset.is_some() {
|
||||
return key;
|
||||
}
|
||||
key.fields
|
||||
.retain(|field| !matches!(field.name.as_str(), "messages" | "prompt" | "input"));
|
||||
const TENANT: [&str; 3] = [
|
||||
"user_api_key",
|
||||
"user_api_key_team_id",
|
||||
"user_api_key_org_id",
|
||||
];
|
||||
let end_user = (scope == "end_user").then_some("user_api_key_end_user_id");
|
||||
for name in TENANT.into_iter().chain(end_user) {
|
||||
let sources = [
|
||||
request.metadata.as_ref(),
|
||||
request.litellm_metadata.as_ref(),
|
||||
request
|
||||
.litellm_params
|
||||
.as_ref()
|
||||
.and_then(|params| params.get("metadata")),
|
||||
request
|
||||
.litellm_params
|
||||
.as_ref()
|
||||
.and_then(|params| params.get("litellm_metadata")),
|
||||
];
|
||||
let Some(value) = sources.into_iter().flatten().find_map(|source| {
|
||||
source
|
||||
.as_object()
|
||||
.and_then(|values| values.get(name))
|
||||
.filter(|value| !value.is_null())
|
||||
}) else {
|
||||
continue;
|
||||
};
|
||||
let value = match value {
|
||||
Value::Null => continue,
|
||||
Value::String(text) => text.clone(),
|
||||
other => other.to_string(),
|
||||
};
|
||||
key.fields.push(CacheKeyField {
|
||||
name: name.to_owned(),
|
||||
value: Some(value),
|
||||
api_parameter: true,
|
||||
internal_parameter: false,
|
||||
});
|
||||
}
|
||||
key
|
||||
}
|
||||
|
||||
pub(super) fn request(value: &Bound<'_, PyAny>) -> PyResult<NativeRequest> {
|
||||
|
|
@ -52,6 +148,7 @@ fn request_input(input: RequestInput) -> PyResult<NativeRequest> {
|
|||
metadata: input.metadata,
|
||||
litellm_metadata: input.litellm_metadata,
|
||||
litellm_params: input.litellm_params,
|
||||
scope: input.scope,
|
||||
})
|
||||
}
|
||||
|
||||
|
|
@ -72,3 +169,90 @@ pub(super) fn now() -> Duration {
|
|||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use litellm_cache_response::{CacheControls, CacheKeyInput, cache_key};
|
||||
use serde_json::json;
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
use super::*;
|
||||
|
||||
fn native_request(key: CacheKeyInput, metadata: Value) -> NativeRequest {
|
||||
NativeRequest {
|
||||
key,
|
||||
controls: CacheControls::default(),
|
||||
ttl: None,
|
||||
max_age: None,
|
||||
messages: Some(json!([{"role": "user", "content": "prompt"}])),
|
||||
input: None,
|
||||
metadata: Some(metadata),
|
||||
litellm_metadata: None,
|
||||
litellm_params: None,
|
||||
scope: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn semantic_key_matches_python_scope_material() {
|
||||
let key = CacheKeyInput {
|
||||
fields: vec![
|
||||
CacheKeyField {
|
||||
name: "model".to_owned(),
|
||||
value: Some("gpt-4.1".to_owned()),
|
||||
api_parameter: true,
|
||||
internal_parameter: false,
|
||||
},
|
||||
CacheKeyField {
|
||||
name: "messages".to_owned(),
|
||||
value: Some("prompt".to_owned()),
|
||||
api_parameter: true,
|
||||
internal_parameter: false,
|
||||
},
|
||||
],
|
||||
..Default::default()
|
||||
};
|
||||
let request = native_request(
|
||||
key,
|
||||
json!({"user_api_key": "k1", "user_api_key_team_id": null}),
|
||||
);
|
||||
let expected = format!("{:x}", Sha256::digest(b"model: gpt-4.1user_api_key: k1"));
|
||||
assert_eq!(cache_key(&semantic_key(&request, "key")), expected);
|
||||
assert_eq!(cache_key(&request.scoped_semantic("key").key), expected);
|
||||
|
||||
let end_user_request = native_request(
|
||||
request.key.clone(),
|
||||
json!({"user_api_key": "k1", "user_api_key_end_user_id": "u1"}),
|
||||
);
|
||||
let expected = format!(
|
||||
"{:x}",
|
||||
Sha256::digest(b"model: gpt-4.1user_api_key: k1user_api_key_end_user_id: u1")
|
||||
);
|
||||
assert_eq!(
|
||||
cache_key(&semantic_key(&end_user_request, "end_user")),
|
||||
expected
|
||||
);
|
||||
|
||||
let preset_request = native_request(
|
||||
CacheKeyInput {
|
||||
preset: Some("preset-key".to_owned()),
|
||||
..Default::default()
|
||||
},
|
||||
json!({"user_api_key": "k1"}),
|
||||
);
|
||||
assert_eq!(
|
||||
semantic_key(&preset_request, "end_user").preset.as_deref(),
|
||||
Some("preset-key")
|
||||
);
|
||||
assert!(semantic_key(&preset_request, "end_user").fields.is_empty());
|
||||
assert_eq!(preset_request.semantic().context.scope, None);
|
||||
assert_eq!(
|
||||
preset_request
|
||||
.scoped_semantic("end_user")
|
||||
.context
|
||||
.scope
|
||||
.as_deref(),
|
||||
Some("end_user")
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
189
litellm-rust/crates/python-bridge/src/cache/semantic.rs
vendored
Normal file
189
litellm-rust/crates/python-bridge/src/cache/semantic.rs
vendored
Normal file
|
|
@ -0,0 +1,189 @@
|
|||
use std::{collections::VecDeque, time::Duration};
|
||||
|
||||
use litellm_cache::Error;
|
||||
use litellm_host_python::{Execution, ExecutionBody, ExecutionStep, run_async};
|
||||
use pyo3::{
|
||||
PyTraverseError, PyVisit,
|
||||
exceptions::{PyException, PyRuntimeError},
|
||||
prelude::*,
|
||||
};
|
||||
use serde_json::Value;
|
||||
|
||||
use super::{
|
||||
cache_error,
|
||||
embedder::{PythonEmbedder, with_prepared_embedding},
|
||||
native::NativeResponseCache,
|
||||
request::{NativeRequest, now},
|
||||
};
|
||||
|
||||
pub(super) enum SemanticOperation {
|
||||
Lookup(NativeRequest),
|
||||
Store(NativeRequest, Value),
|
||||
StoreBatch(VecDeque<(NativeRequest, Value)>),
|
||||
}
|
||||
|
||||
/// What an exception from the Python embedder means for the operation.
|
||||
#[derive(Clone, Copy)]
|
||||
pub(super) enum EmbeddingFailure {
|
||||
/// Raise the Python exception unchanged.
|
||||
Propagate,
|
||||
/// Treat the embedding as unavailable and let the backend report that.
|
||||
Unavailable,
|
||||
}
|
||||
|
||||
enum Phase {
|
||||
Start,
|
||||
AwaitingEmbedding,
|
||||
AwaitingBackend,
|
||||
}
|
||||
|
||||
/// Runs a semantic cache operation whose embedding comes from Python: await the Python
|
||||
/// embedder in the caller's event loop, seed the native backend with the vector, await the
|
||||
/// backend, and repeat for each entry of a batch.
|
||||
pub(super) struct SemanticExecution {
|
||||
service: NativeResponseCache,
|
||||
embedder: PythonEmbedder,
|
||||
failure: EmbeddingFailure,
|
||||
operation: SemanticOperation,
|
||||
pending: Option<(NativeRequest, Option<Value>)>,
|
||||
phase: Phase,
|
||||
now: Duration,
|
||||
}
|
||||
|
||||
impl SemanticExecution {
|
||||
pub(super) fn new(
|
||||
service: NativeResponseCache,
|
||||
embedder: PythonEmbedder,
|
||||
failure: EmbeddingFailure,
|
||||
operation: SemanticOperation,
|
||||
) -> Self {
|
||||
Self {
|
||||
service,
|
||||
embedder,
|
||||
failure,
|
||||
operation,
|
||||
pending: None,
|
||||
phase: Phase::Start,
|
||||
now: now(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Takes the next entry of the operation; `None` once a batch is exhausted.
|
||||
fn next_pending(&mut self) -> Option<(NativeRequest, Option<Value>)> {
|
||||
match &mut self.operation {
|
||||
SemanticOperation::Lookup(request) => Some((request.clone(), None)),
|
||||
SemanticOperation::Store(request, response) => {
|
||||
Some((request.clone(), Some(std::mem::take(response))))
|
||||
}
|
||||
SemanticOperation::StoreBatch(queue) => queue
|
||||
.pop_front()
|
||||
.map(|(request, response)| (request, Some(response))),
|
||||
}
|
||||
}
|
||||
|
||||
fn start(&mut self, py: Python<'_>) -> PyResult<ExecutionStep> {
|
||||
let Some(pending) = self.next_pending() else {
|
||||
return Ok(ExecutionStep::Return(py.None()));
|
||||
};
|
||||
let (request, response) = &pending;
|
||||
let enabled = match response {
|
||||
None => request.controls.reads(),
|
||||
Some(_) => request.controls.writes(),
|
||||
};
|
||||
let input = enabled
|
||||
.then(|| self.service.embedding_input(request))
|
||||
.flatten();
|
||||
self.pending = Some(pending);
|
||||
let Some(input) = input else {
|
||||
return self.backend_step(py, Err(Error::Unavailable));
|
||||
};
|
||||
let awaitable =
|
||||
self.embedder
|
||||
.async_embedding(py, &input.prompt, input.metadata.as_ref())?;
|
||||
self.phase = Phase::AwaitingEmbedding;
|
||||
Ok(ExecutionStep::Await(awaitable))
|
||||
}
|
||||
|
||||
fn embedded(&mut self, py: Python<'_>, result: PyResult<Py<PyAny>>) -> PyResult<ExecutionStep> {
|
||||
let seed = match result {
|
||||
Ok(vector) => {
|
||||
PythonEmbedder::extract(vector.into_bound(py)).map_err(|_| Error::Unavailable)
|
||||
}
|
||||
Err(error) => match self.failure {
|
||||
EmbeddingFailure::Propagate => return Err(error),
|
||||
EmbeddingFailure::Unavailable if error.is_instance_of::<PyException>(py) => {
|
||||
Err(Error::Unavailable)
|
||||
}
|
||||
EmbeddingFailure::Unavailable => return Err(error),
|
||||
},
|
||||
};
|
||||
self.backend_step(py, seed)
|
||||
}
|
||||
|
||||
fn backend_step(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
seed: Result<Vec<f32>, Error>,
|
||||
) -> PyResult<ExecutionStep> {
|
||||
self.phase = Phase::AwaitingBackend;
|
||||
let (request, response) = self.pending.take().ok_or_else(|| {
|
||||
PyRuntimeError::new_err("semantic execution resumed without a pending operation")
|
||||
})?;
|
||||
let service = self.service.clone();
|
||||
let now = self.now;
|
||||
let future = async move {
|
||||
match response {
|
||||
None => service.async_lookup(&request, now).await,
|
||||
Some(response) => service
|
||||
.async_store(&request, response, now)
|
||||
.await
|
||||
.map(|_| None),
|
||||
}
|
||||
};
|
||||
let awaitable = run_async(py, with_prepared_embedding(seed, future), cache_error)?;
|
||||
Ok(ExecutionStep::Await(awaitable.unbind()))
|
||||
}
|
||||
|
||||
fn resume_py(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
result: Option<PyResult<Py<PyAny>>>,
|
||||
) -> PyResult<ExecutionStep> {
|
||||
match (&self.phase, result) {
|
||||
(Phase::Start, None) => self.start(py),
|
||||
(Phase::AwaitingEmbedding, Some(result)) => self.embedded(py, result),
|
||||
(Phase::AwaitingBackend, Some(Err(error))) => Err(error),
|
||||
(Phase::AwaitingBackend, Some(Ok(value))) => {
|
||||
let more = matches!(
|
||||
&self.operation,
|
||||
SemanticOperation::StoreBatch(queue) if !queue.is_empty()
|
||||
);
|
||||
if more {
|
||||
self.phase = Phase::Start;
|
||||
return self.start(py);
|
||||
}
|
||||
Ok(ExecutionStep::Return(value))
|
||||
}
|
||||
_ => Err(PyRuntimeError::new_err(
|
||||
"invalid semantic cache execution state",
|
||||
)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ExecutionBody for SemanticExecution {
|
||||
fn resume(&mut self, result: Option<PyResult<Py<PyAny>>>) -> PyResult<ExecutionStep> {
|
||||
Python::attach(|py| self.resume_py(py, result))
|
||||
}
|
||||
|
||||
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
|
||||
self.embedder.traverse(visit)
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn drive(py: Python<'_>, body: SemanticExecution) -> PyResult<Bound<'_, PyAny>> {
|
||||
let execution = Py::new(py, Execution::new(body))?;
|
||||
py.import("litellm.rust_bridge.lifecycle")?
|
||||
.getattr("drive")?
|
||||
.call1((execution,))
|
||||
}
|
||||
|
|
@ -1,249 +0,0 @@
|
|||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use litellm_cache::SemanticCacheContext;
|
||||
use litellm_cache_response::{ResponseCache, ResponseCacheCodec, ResponseCacheRequest};
|
||||
use litellm_cache_valkey_semantic::{PreparedEmbedding, ValkeySemanticCache, prompt_from_context};
|
||||
use litellm_host_python::{Execution, ExecutionBody, ExecutionStep, run_async};
|
||||
use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*};
|
||||
use serde_json::Value;
|
||||
|
||||
use super::{cache_error, embedder::PythonEmbedder};
|
||||
|
||||
pub(super) enum Op {
|
||||
Lookup,
|
||||
Store(Value),
|
||||
StoreBatch(Vec<Value>),
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
enum State {
|
||||
Start,
|
||||
AwaitingEmbedding,
|
||||
AwaitingStorage,
|
||||
Done,
|
||||
}
|
||||
|
||||
pub(super) struct SemanticEmbedExecution {
|
||||
backend: Arc<ValkeySemanticCache<PythonEmbedder, ResponseCacheCodec>>,
|
||||
embedder: PythonEmbedder,
|
||||
requests: Vec<ResponseCacheRequest<SemanticCacheContext>>,
|
||||
op: Op,
|
||||
now: Option<Duration>,
|
||||
prepared: Vec<Option<Vec<f32>>>,
|
||||
index: usize,
|
||||
state: State,
|
||||
}
|
||||
|
||||
impl SemanticEmbedExecution {
|
||||
pub(super) fn lookup(
|
||||
backend: Arc<ValkeySemanticCache<PythonEmbedder, ResponseCacheCodec>>,
|
||||
embedder: PythonEmbedder,
|
||||
request: ResponseCacheRequest<SemanticCacheContext>,
|
||||
) -> Self {
|
||||
Self {
|
||||
backend,
|
||||
embedder,
|
||||
requests: vec![request],
|
||||
op: Op::Lookup,
|
||||
now: None,
|
||||
prepared: vec![None],
|
||||
index: 0,
|
||||
state: State::Start,
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn store(
|
||||
backend: Arc<ValkeySemanticCache<PythonEmbedder, ResponseCacheCodec>>,
|
||||
embedder: PythonEmbedder,
|
||||
request: ResponseCacheRequest<SemanticCacheContext>,
|
||||
response: Value,
|
||||
) -> Self {
|
||||
Self {
|
||||
backend,
|
||||
embedder,
|
||||
requests: vec![request],
|
||||
op: Op::Store(response),
|
||||
now: None,
|
||||
prepared: vec![None],
|
||||
index: 0,
|
||||
state: State::Start,
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn store_batch(
|
||||
backend: Arc<ValkeySemanticCache<PythonEmbedder, ResponseCacheCodec>>,
|
||||
embedder: PythonEmbedder,
|
||||
requests: Vec<ResponseCacheRequest<SemanticCacheContext>>,
|
||||
responses: Vec<Value>,
|
||||
) -> Self {
|
||||
Self {
|
||||
backend,
|
||||
embedder,
|
||||
prepared: vec![None; requests.len()],
|
||||
requests,
|
||||
op: Op::StoreBatch(responses),
|
||||
now: None,
|
||||
index: 0,
|
||||
state: State::Start,
|
||||
}
|
||||
}
|
||||
|
||||
fn start(&mut self, py: Python<'_>) -> PyResult<ExecutionStep> {
|
||||
if self.now.is_none() {
|
||||
self.now = Some(super::request::now());
|
||||
}
|
||||
while self.index < self.requests.len() {
|
||||
let request = &self.requests[self.index];
|
||||
let enabled = match &self.op {
|
||||
Op::Lookup => request.controls.reads(),
|
||||
Op::Store(_) | Op::StoreBatch(_) => request.controls.writes(),
|
||||
};
|
||||
if !enabled {
|
||||
self.index += 1;
|
||||
continue;
|
||||
}
|
||||
let Some(prompt) = prompt_from_context(&request.context) else {
|
||||
self.index += 1;
|
||||
continue;
|
||||
};
|
||||
let metadata = request.context.metadata.clone();
|
||||
let awaitable = self
|
||||
.embedder
|
||||
.async_embed_awaitable(py, &prompt, &metadata)?;
|
||||
self.state = State::AwaitingEmbedding;
|
||||
return Ok(ExecutionStep::Await(awaitable.unbind()));
|
||||
}
|
||||
self.state = State::AwaitingStorage;
|
||||
self.storage_step(py)
|
||||
}
|
||||
|
||||
fn storage_step(&self, py: Python<'_>) -> PyResult<ExecutionStep> {
|
||||
let requests = self.requests.clone();
|
||||
let prepared = self.prepared.clone();
|
||||
let backend = Arc::clone(&self.backend);
|
||||
let now = self
|
||||
.now
|
||||
.ok_or_else(|| PyRuntimeError::new_err("semantic cache timestamp is unavailable"))?;
|
||||
let awaitable = match &self.op {
|
||||
Op::Lookup => {
|
||||
let Some(request) = requests.into_iter().next() else {
|
||||
return Err(PyRuntimeError::new_err(
|
||||
"semantic lookup requires one request",
|
||||
));
|
||||
};
|
||||
match prepared.into_iter().next().flatten() {
|
||||
Some(values) => {
|
||||
let backend = backend.with_embedder(PreparedEmbedding(values));
|
||||
let cache = Arc::new(ResponseCache::new(Arc::new(backend)));
|
||||
run_async(
|
||||
py,
|
||||
async move { cache.async_lookup(&request, now).await },
|
||||
cache_error,
|
||||
)?
|
||||
}
|
||||
None => {
|
||||
let cache = Arc::new(ResponseCache::new(backend));
|
||||
run_async(
|
||||
py,
|
||||
async move { cache.async_lookup(&request, now).await },
|
||||
cache_error,
|
||||
)?
|
||||
}
|
||||
}
|
||||
}
|
||||
Op::Store(response) => {
|
||||
let Some(request) = requests.into_iter().next() else {
|
||||
return Err(PyRuntimeError::new_err(
|
||||
"semantic store requires one request",
|
||||
));
|
||||
};
|
||||
let response = response.clone();
|
||||
match prepared.into_iter().next().flatten() {
|
||||
Some(values) => {
|
||||
let backend = backend.with_embedder(PreparedEmbedding(values));
|
||||
let cache = Arc::new(ResponseCache::new(Arc::new(backend)));
|
||||
run_async(
|
||||
py,
|
||||
async move { cache.async_store(&request, response, now).await },
|
||||
cache_error,
|
||||
)?
|
||||
}
|
||||
None => {
|
||||
let cache = Arc::new(ResponseCache::new(backend));
|
||||
run_async(
|
||||
py,
|
||||
async move { cache.async_store(&request, response, now).await },
|
||||
cache_error,
|
||||
)?
|
||||
}
|
||||
}
|
||||
}
|
||||
Op::StoreBatch(responses) => {
|
||||
let responses = responses.clone();
|
||||
run_async(
|
||||
py,
|
||||
async move {
|
||||
for ((request, response), prepared) in
|
||||
requests.into_iter().zip(responses).zip(prepared)
|
||||
{
|
||||
let Some(values) = prepared else {
|
||||
continue;
|
||||
};
|
||||
let backend = backend.with_embedder(PreparedEmbedding(values));
|
||||
let cache = ResponseCache::new(Arc::new(backend));
|
||||
cache.async_store(&request, response, now).await?;
|
||||
}
|
||||
Ok(())
|
||||
},
|
||||
cache_error,
|
||||
)?
|
||||
}
|
||||
};
|
||||
Ok(ExecutionStep::Await(awaitable.unbind()))
|
||||
}
|
||||
|
||||
fn resume_py(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
result: Option<PyResult<Py<PyAny>>>,
|
||||
) -> PyResult<ExecutionStep> {
|
||||
match (self.state, result) {
|
||||
(State::Start, None) => self.start(py),
|
||||
(State::AwaitingEmbedding, Some(Ok(value))) => {
|
||||
let values = value.bind(py).extract::<Vec<f64>>()?;
|
||||
self.prepared[self.index] =
|
||||
Some(values.into_iter().map(|value| value as f32).collect());
|
||||
self.index += 1;
|
||||
self.start(py)
|
||||
}
|
||||
(State::AwaitingStorage, Some(Ok(value))) => {
|
||||
self.state = State::Done;
|
||||
Ok(ExecutionStep::Return(value))
|
||||
}
|
||||
(_, Some(Err(error))) => Err(error),
|
||||
_ => Err(PyRuntimeError::new_err(
|
||||
"invalid semantic cache execution state",
|
||||
)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ExecutionBody for SemanticEmbedExecution {
|
||||
fn resume(&mut self, result: Option<PyResult<Py<PyAny>>>) -> PyResult<ExecutionStep> {
|
||||
Python::attach(|py| self.resume_py(py, result))
|
||||
}
|
||||
|
||||
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
|
||||
self.embedder.traverse(visit)
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn drive_semantic<'py>(
|
||||
py: Python<'py>,
|
||||
body: SemanticEmbedExecution,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let execution = Py::new(py, Execution::new(body))?;
|
||||
py.import("litellm.rust_bridge.lifecycle")?
|
||||
.getattr("drive")?
|
||||
.call1((execution,))
|
||||
}
|
||||
|
|
@ -1,7 +1,4 @@
|
|||
use std::collections::BTreeSet;
|
||||
|
||||
use litellm_core_utils::serde_compat::parse_str_bool;
|
||||
use litellm_http::SslVerify;
|
||||
use pyo3::{
|
||||
exceptions::{PyAttributeError, PyRuntimeError, PyValueError},
|
||||
prelude::*,
|
||||
|
|
@ -33,36 +30,52 @@ impl From<ProjectionError> for PyErr {
|
|||
}
|
||||
}
|
||||
|
||||
pub(crate) struct Truthy(pub bool);
|
||||
pub(crate) struct ExactTrue(pub bool);
|
||||
pub(crate) struct StrBool(pub Option<bool>);
|
||||
pub(crate) struct OptionalStrictString(pub Option<String>);
|
||||
pub(crate) struct FalsyOptionalString(pub Option<String>);
|
||||
pub(crate) struct TuningString(pub Option<String>);
|
||||
pub(crate) struct StringCollection(pub Vec<String>);
|
||||
pub(crate) struct SslVerifyInput(pub Option<SslVerify>);
|
||||
pub(crate) struct FieldSpec<T> {
|
||||
name: &'static str,
|
||||
decode: fn(&Field<'_>) -> Result<T, ProjectionError>,
|
||||
}
|
||||
|
||||
impl<T> FieldSpec<T> {
|
||||
pub(crate) const fn new(
|
||||
name: &'static str,
|
||||
decode: fn(&Field<'_>) -> Result<T, ProjectionError>,
|
||||
) -> Self {
|
||||
Self { name, decode }
|
||||
}
|
||||
|
||||
pub(crate) fn read(
|
||||
&self,
|
||||
snapshot: &Bound<'_, PyAny>,
|
||||
group: &'static str,
|
||||
) -> Result<T, ProjectionError> {
|
||||
(self.decode)(&Field::read(snapshot, group, self.name)?)
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) struct Field<'py> {
|
||||
path: &'static str,
|
||||
group: &'static str,
|
||||
name: &'static str,
|
||||
value: Bound<'py, PyAny>,
|
||||
}
|
||||
|
||||
impl<'py> Field<'py> {
|
||||
pub(crate) fn new(path: &'static str, value: Bound<'py, PyAny>) -> Self {
|
||||
Self { path, value }
|
||||
pub(crate) fn new(group: &'static str, name: &'static str, value: Bound<'py, PyAny>) -> Self {
|
||||
Self { group, name, value }
|
||||
}
|
||||
|
||||
/// Reads `snapshot.<name>`, distinguishing a field the accessor never declared from a
|
||||
/// descriptor that raised `AttributeError`.
|
||||
pub(crate) fn read(
|
||||
snapshot: &Bound<'py, PyAny>,
|
||||
path: &'static str,
|
||||
group: &'static str,
|
||||
name: &'static str,
|
||||
) -> Result<Self, ProjectionError> {
|
||||
let name = path.rsplit('.').next().unwrap_or(path);
|
||||
match snapshot.getattr(name) {
|
||||
Ok(value) => Ok(Self::new(path, value)),
|
||||
Ok(value) => Ok(Self::new(group, name, value)),
|
||||
Err(error) if error.is_instance_of::<PyAttributeError>(snapshot.py()) => {
|
||||
match Self::missing_field(snapshot, name) {
|
||||
Ok(true) => Err(ProjectionError::InternalSchemaFailure(format!(
|
||||
"{path}: missing snapshot field"
|
||||
"{group}.{name}: missing snapshot field"
|
||||
))),
|
||||
_ => Err(error.into()),
|
||||
}
|
||||
|
|
@ -84,27 +97,49 @@ impl<'py> Field<'py> {
|
|||
&& getter.is(object.getattr("__getattribute__")?))
|
||||
}
|
||||
|
||||
fn expected(&self, expected: &'static str) -> Result<String, ProjectionError> {
|
||||
pub(crate) fn path(&self) -> String {
|
||||
format!("{}.{}", self.group, self.name)
|
||||
}
|
||||
|
||||
/// A member of this field's collection, reported under the same path.
|
||||
pub(crate) fn member(&self, value: Bound<'py, PyAny>) -> Self {
|
||||
Self::new(self.group, self.name, value)
|
||||
}
|
||||
|
||||
pub(crate) fn expected(&self, expected: &str) -> Result<String, ProjectionError> {
|
||||
Ok(format!(
|
||||
"{}: expected {expected}, got {}",
|
||||
self.path,
|
||||
self.path(),
|
||||
self.value.get_type().name()?
|
||||
))
|
||||
}
|
||||
|
||||
fn invalid(&self, expected: &'static str) -> ProjectionError {
|
||||
pub(crate) fn invalid(&self, expected: &str) -> ProjectionError {
|
||||
match self.expected(expected) {
|
||||
Ok(message) => ProjectionError::InvalidConfiguration(message),
|
||||
Err(error) => error,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn truthy(&self) -> Result<Truthy, ProjectionError> {
|
||||
Ok(Truthy(self.value.is_truthy()?))
|
||||
pub(crate) fn value(&self) -> &Bound<'py, PyAny> {
|
||||
&self.value
|
||||
}
|
||||
|
||||
pub(crate) fn exact_true(&self) -> ExactTrue {
|
||||
ExactTrue(self.value.is(PyBool::new(self.value.py(), true)))
|
||||
pub(crate) fn truthy(&self) -> Result<bool, ProjectionError> {
|
||||
Ok(self.value.is_truthy()?)
|
||||
}
|
||||
|
||||
pub(crate) fn exact_true(&self) -> bool {
|
||||
self.value.is(PyBool::new(self.value.py(), true))
|
||||
}
|
||||
|
||||
pub(crate) fn schema_bool(&self) -> Result<bool, ProjectionError> {
|
||||
if !self.value.is_instance_of::<PyBool>() {
|
||||
return Err(ProjectionError::InternalSchemaFailure(
|
||||
self.expected("a Boolean")?,
|
||||
));
|
||||
}
|
||||
Ok(self.exact_true())
|
||||
}
|
||||
|
||||
pub(crate) fn strict_string(&self) -> Result<String, ProjectionError> {
|
||||
|
|
@ -124,108 +159,366 @@ impl<'py> Field<'py> {
|
|||
self.strict_string()
|
||||
}
|
||||
|
||||
pub(crate) fn schema_bool(&self) -> Result<bool, ProjectionError> {
|
||||
if !self.value.is_instance_of::<PyBool>() {
|
||||
return Err(ProjectionError::InternalSchemaFailure(
|
||||
self.expected("a Boolean")?,
|
||||
));
|
||||
}
|
||||
Ok(self.exact_true().0)
|
||||
}
|
||||
|
||||
pub(crate) fn str_bool(&self) -> Result<StrBool, ProjectionError> {
|
||||
pub(crate) fn str_bool(&self) -> Result<Option<bool>, ProjectionError> {
|
||||
if self.value.is_none() {
|
||||
return Ok(StrBool(None));
|
||||
return Ok(None);
|
||||
}
|
||||
Ok(StrBool(parse_str_bool(&self.strict_string()?)))
|
||||
Ok(parse_str_bool(&self.strict_string()?))
|
||||
}
|
||||
|
||||
pub(crate) fn optional_strict_string(&self) -> Result<OptionalStrictString, ProjectionError> {
|
||||
pub(crate) fn optional_strict_string(&self) -> Result<Option<String>, ProjectionError> {
|
||||
if self.value.is_none() {
|
||||
return Ok(OptionalStrictString(None));
|
||||
return Ok(None);
|
||||
}
|
||||
self.strict_string().map(Some).map(OptionalStrictString)
|
||||
self.strict_string().map(Some)
|
||||
}
|
||||
|
||||
pub(crate) fn falsy_optional_string(&self) -> Result<FalsyOptionalString, ProjectionError> {
|
||||
if !self.truthy()?.0 {
|
||||
return Ok(FalsyOptionalString(None));
|
||||
pub(crate) fn falsy_optional_string(&self) -> Result<Option<String>, ProjectionError> {
|
||||
if !self.truthy()? {
|
||||
return Ok(None);
|
||||
}
|
||||
self.strict_string().map(Some).map(FalsyOptionalString)
|
||||
self.strict_string().map(Some)
|
||||
}
|
||||
|
||||
pub(crate) fn tuning_string(&self) -> Result<TuningString, ProjectionError> {
|
||||
if !self.truthy()?.0 || !self.value.is_instance_of::<PyString>() {
|
||||
return Ok(TuningString(None));
|
||||
pub(crate) fn tuning_string(&self) -> Result<Option<String>, ProjectionError> {
|
||||
if !self.truthy()? || !self.value.is_instance_of::<PyString>() {
|
||||
return Ok(None);
|
||||
}
|
||||
self.strict_string().map(Some).map(TuningString)
|
||||
self.strict_string().map(Some)
|
||||
}
|
||||
|
||||
pub(crate) fn string_collection(&self) -> Result<StringCollection, ProjectionError> {
|
||||
if !self.truthy()?.0 {
|
||||
return Ok(StringCollection(Vec::new()));
|
||||
pub(crate) fn string_collection(&self) -> Result<Vec<String>, ProjectionError> {
|
||||
if !self.truthy()? {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
if self.value.is_instance_of::<PyString>() {
|
||||
return self
|
||||
.strict_string()
|
||||
.map(|value| StringCollection(vec![value]));
|
||||
return self.strict_string().map(|value| vec![value]);
|
||||
}
|
||||
let values = self
|
||||
.value
|
||||
self.value
|
||||
.try_iter()?
|
||||
.filter_map(|item| {
|
||||
let member = match item {
|
||||
Ok(value) => Self::new(self.path, value),
|
||||
Ok(value) => self.member(value),
|
||||
Err(error) => return Some(Err(error.into())),
|
||||
};
|
||||
match member.truthy() {
|
||||
Ok(Truthy(false)) => None,
|
||||
Ok(Truthy(true)) => Some(member.strict_string()),
|
||||
Ok(false) => None,
|
||||
Ok(true) => Some(member.strict_string()),
|
||||
Err(error) => Some(Err(error)),
|
||||
}
|
||||
})
|
||||
.collect::<Result<Vec<_>, ProjectionError>>()?;
|
||||
Ok(StringCollection(values))
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub(crate) fn host_collection(&self) -> Result<StringCollection, ProjectionError> {
|
||||
let values = self
|
||||
.string_collection()?
|
||||
.0
|
||||
.into_iter()
|
||||
.map(|host| litellm_http::media::normalize_host(&host))
|
||||
.collect::<BTreeSet<_>>();
|
||||
Ok(StringCollection(values.into_iter().collect()))
|
||||
}
|
||||
|
||||
pub(crate) fn ssl_verify(&self) -> Result<SslVerifyInput, ProjectionError> {
|
||||
pub(crate) fn optional_string_collection(
|
||||
&self,
|
||||
) -> Result<Option<Vec<String>>, ProjectionError> {
|
||||
if self.value.is_none() {
|
||||
return Ok(SslVerifyInput(None));
|
||||
return Ok(None);
|
||||
}
|
||||
if self.value.is_instance_of::<PyBool>() {
|
||||
return Ok(SslVerifyInput(Some(if self.exact_true().0 {
|
||||
SslVerify::Enabled
|
||||
} else {
|
||||
SslVerify::Disabled
|
||||
})));
|
||||
}
|
||||
if self.value.is_instance_of::<PyString>() {
|
||||
let parsed = match self.str_bool()?.0 {
|
||||
Some(true) => SslVerify::Enabled,
|
||||
Some(false) => SslVerify::Disabled,
|
||||
None => SslVerify::CaBundle(self.strict_string()?.into()),
|
||||
};
|
||||
return Ok(SslVerifyInput(Some(parsed)));
|
||||
}
|
||||
let context = self.value.py().import("ssl")?.getattr("SSLContext")?;
|
||||
if self.value.is_instance(&context)? {
|
||||
return Err(ProjectionError::UnsupportedLiveObject(self.expected(
|
||||
"a Boolean, Boolean string, CA path, or None; live SSLContext is unsupported",
|
||||
)?));
|
||||
}
|
||||
Err(self.invalid("a Boolean, Boolean string, CA path, or None"))
|
||||
self.string_collection().map(Some)
|
||||
}
|
||||
|
||||
pub(crate) fn python_binding(&self) -> Option<Py<PyAny>> {
|
||||
(!self.value.is_none()).then(|| self.value.clone().unbind())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
mod tests {
|
||||
use std::ffi::CString;
|
||||
|
||||
use pyo3::{
|
||||
exceptions::{PyLookupError, PyRuntimeError, PyValueError},
|
||||
types::PyDict,
|
||||
};
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
|
||||
fn evaluate<'py>(py: Python<'py>, source: &str) -> Bound<'py, PyAny> {
|
||||
py.eval(&CString::new(source).unwrap(), None, None).unwrap()
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case("None", false, false)]
|
||||
#[case("False", false, false)]
|
||||
#[case("True", true, true)]
|
||||
#[case("0", false, false)]
|
||||
#[case("1", true, false)]
|
||||
#[case("''", false, false)]
|
||||
#[case("'false'", true, false)]
|
||||
#[case("[]", false, false)]
|
||||
#[case("[0]", true, false)]
|
||||
#[case("{}", false, false)]
|
||||
#[case("object()", true, false)]
|
||||
fn boolean_operations_have_distinct_python_semantics(
|
||||
#[case] source: &str,
|
||||
#[case] truth: bool,
|
||||
#[case] exact: bool,
|
||||
) {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let value = evaluate(py, source);
|
||||
let field = Field::new("test", "flag", value.clone());
|
||||
assert_eq!(field.truthy().unwrap(), truth);
|
||||
assert_eq!(field.exact_true(), exact);
|
||||
assert_eq!(
|
||||
field.truthy().unwrap(),
|
||||
py.import("builtins")
|
||||
.unwrap()
|
||||
.getattr("bool")
|
||||
.unwrap()
|
||||
.call1((value,))
|
||||
.unwrap()
|
||||
.extract::<bool>()
|
||||
.unwrap()
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case("None", Ok(None), Ok(None), Ok(None))]
|
||||
#[case("''", Ok(Some("")), Ok(None), Ok(None))]
|
||||
#[case(
|
||||
"' value '",
|
||||
Ok(Some(" value ")),
|
||||
Ok(Some(" value ")),
|
||||
Ok(Some(" value "))
|
||||
)]
|
||||
#[case("[]", Err(()), Ok(None), Ok(None))]
|
||||
#[case("0", Err(()), Ok(None), Ok(None))]
|
||||
#[case("1", Err(()), Err(()), Ok(None))]
|
||||
#[case("object()", Err(()), Err(()), Ok(None))]
|
||||
fn string_operations_do_not_conflate_absence_and_type_checks(
|
||||
#[case] source: &str,
|
||||
#[case] strict: Result<Option<&str>, ()>,
|
||||
#[case] fallback: Result<Option<&str>, ()>,
|
||||
#[case] tuning: Result<Option<&str>, ()>,
|
||||
) {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let field = Field::new("test", "string", evaluate(py, source));
|
||||
let owned =
|
||||
|expected: Result<Option<&str>, ()>| expected.map(|value| value.map(str::to_owned));
|
||||
assert_eq!(
|
||||
field.optional_strict_string().map_err(|_| ()),
|
||||
owned(strict)
|
||||
);
|
||||
assert_eq!(
|
||||
field.falsy_optional_string().map_err(|_| ()),
|
||||
owned(fallback)
|
||||
);
|
||||
assert_eq!(field.tuning_string().map_err(|_| ()), owned(tuning));
|
||||
});
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case("None", None)]
|
||||
#[case("' True '", Some(true))]
|
||||
#[case("' fAlSe '", Some(false))]
|
||||
#[case("'yes'", None)]
|
||||
#[case("'1'", None)]
|
||||
#[case("'unknown'", None)]
|
||||
fn string_boolean_tokens_remain_separate_from_truthiness(
|
||||
#[case] source: &str,
|
||||
#[case] expected: Option<bool>,
|
||||
) {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
assert_eq!(
|
||||
Field::new("test", "flag", evaluate(py, source))
|
||||
.str_bool()
|
||||
.unwrap(),
|
||||
expected
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn protocol_errors_preserve_exception_identity_traceback_cause_and_context() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let locals = PyDict::new(py);
|
||||
py.run(
|
||||
c"
|
||||
failure = LookupError('protocol failed')
|
||||
cause = ValueError('cause')
|
||||
context = RuntimeError('context')
|
||||
def fail():
|
||||
try:
|
||||
raise context
|
||||
except RuntimeError:
|
||||
raise failure from cause
|
||||
class Bool:
|
||||
def __bool__(self): return fail()
|
||||
class Length:
|
||||
def __len__(self): return fail()
|
||||
class Iter:
|
||||
def __iter__(self): return fail()
|
||||
class Next:
|
||||
def __iter__(self): return self
|
||||
def __next__(self): return fail()
|
||||
class Descriptor:
|
||||
@property
|
||||
def flag(self): return fail()
|
||||
values = (Bool(), Length(), Iter(), Next(), [Bool()])
|
||||
descriptor = Descriptor()
|
||||
",
|
||||
Some(&locals),
|
||||
Some(&locals),
|
||||
)
|
||||
.unwrap();
|
||||
let values = locals.get_item("values").unwrap().unwrap();
|
||||
for value in values.try_iter().unwrap() {
|
||||
let error = Field::new("test", "flag", value.unwrap())
|
||||
.string_collection()
|
||||
.err()
|
||||
.unwrap();
|
||||
let error = PyErr::from(error);
|
||||
assert!(
|
||||
error
|
||||
.value(py)
|
||||
.is(locals.get_item("failure").unwrap().unwrap())
|
||||
);
|
||||
assert!(error.is_instance_of::<PyLookupError>(py));
|
||||
assert!(error.traceback(py).is_some());
|
||||
assert!(
|
||||
error
|
||||
.value(py)
|
||||
.getattr("__cause__")
|
||||
.unwrap()
|
||||
.is(locals.get_item("cause").unwrap().unwrap())
|
||||
);
|
||||
assert!(
|
||||
error
|
||||
.value(py)
|
||||
.getattr("__context__")
|
||||
.unwrap()
|
||||
.is(locals.get_item("context").unwrap().unwrap())
|
||||
);
|
||||
}
|
||||
let error = Field::read(
|
||||
&locals.get_item("descriptor").unwrap().unwrap(),
|
||||
"test",
|
||||
"flag",
|
||||
)
|
||||
.err()
|
||||
.unwrap();
|
||||
assert!(
|
||||
PyErr::from(error)
|
||||
.value(py)
|
||||
.is(locals.get_item("failure").unwrap().unwrap())
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn identity_and_string_contents_do_not_invoke_unrelated_protocols() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let locals = PyDict::new(py);
|
||||
py.run(
|
||||
c"
|
||||
class Hostile:
|
||||
def __bool__(self): raise AssertionError('bool called')
|
||||
def __eq__(self, other): raise AssertionError('eq called')
|
||||
def __str__(self): raise AssertionError('str called')
|
||||
class Text(str):
|
||||
def __str__(self): raise AssertionError('str called')
|
||||
def strip(self): raise AssertionError('strip called')
|
||||
def lower(self): raise AssertionError('lower called')
|
||||
hostile = Hostile()
|
||||
text = Text(' False ')
|
||||
",
|
||||
Some(&locals),
|
||||
Some(&locals),
|
||||
)
|
||||
.unwrap();
|
||||
let hostile = Field::new("test", "flag", locals.get_item("hostile").unwrap().unwrap());
|
||||
assert!(!hostile.exact_true());
|
||||
assert!(matches!(
|
||||
hostile.strict_string(),
|
||||
Err(ProjectionError::InvalidConfiguration(_))
|
||||
));
|
||||
let text = Field::new("test", "flag", locals.get_item("text").unwrap().unwrap());
|
||||
assert_eq!(text.strict_string().unwrap(), " False ");
|
||||
assert_eq!(text.str_bool().unwrap(), Some(false));
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn missing_snapshot_fields_and_descriptor_attribute_errors_are_distinct() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let locals = PyDict::new(py);
|
||||
py.run(
|
||||
c"
|
||||
failure = AttributeError('descriptor failed')
|
||||
class Snapshot:
|
||||
@property
|
||||
def flag(self): raise failure
|
||||
snapshot = Snapshot()
|
||||
class Dynamic:
|
||||
def __getattr__(self, name): raise failure
|
||||
class Intercepted:
|
||||
def __getattribute__(self, name): raise failure
|
||||
dynamic = Dynamic()
|
||||
intercepted = Intercepted()
|
||||
",
|
||||
Some(&locals),
|
||||
Some(&locals),
|
||||
)
|
||||
.unwrap();
|
||||
let snapshot = locals.get_item("snapshot").unwrap().unwrap();
|
||||
let descriptor = PyErr::from(Field::read(&snapshot, "test", "flag").err().unwrap());
|
||||
assert!(
|
||||
descriptor
|
||||
.value(py)
|
||||
.is(locals.get_item("failure").unwrap().unwrap())
|
||||
);
|
||||
for name in ["dynamic", "intercepted"] {
|
||||
let value = locals.get_item(name).unwrap().unwrap();
|
||||
let error = PyErr::from(Field::read(&value, "test", "flag").err().unwrap());
|
||||
assert!(
|
||||
error
|
||||
.value(py)
|
||||
.is(locals.get_item("failure").unwrap().unwrap())
|
||||
);
|
||||
}
|
||||
let missing = PyErr::from(Field::read(&snapshot, "test", "missing").err().unwrap());
|
||||
assert!(missing.is_instance_of::<PyRuntimeError>(py));
|
||||
assert!(missing.to_string().contains("test.missing"));
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn configuration_errors_name_fields_without_exposing_values() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
for source in [
|
||||
"{'secret': 'do-not-print'}",
|
||||
"['host.test', {'secret': 'do-not-print'}]",
|
||||
] {
|
||||
let field = Field::new("test", "setting", evaluate(py, source));
|
||||
let error = PyErr::from(field.falsy_optional_string().err().unwrap());
|
||||
assert!(error.is_instance_of::<PyValueError>(py));
|
||||
assert!(error.to_string().contains("test.setting"));
|
||||
assert!(!error.to_string().contains("do-not-print"));
|
||||
}
|
||||
let hosts = Field::new(
|
||||
"url_policy",
|
||||
"user_url_allowed_hosts",
|
||||
evaluate(py, "['host.test', 1]"),
|
||||
);
|
||||
assert!(matches!(
|
||||
hosts.string_collection(),
|
||||
Err(ProjectionError::InvalidConfiguration(_))
|
||||
));
|
||||
assert!(matches!(
|
||||
Field::new("test", "flag", evaluate(py, "1")).str_bool(),
|
||||
Err(ProjectionError::InvalidConfiguration(_))
|
||||
));
|
||||
});
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,372 +0,0 @@
|
|||
use std::ffi::CString;
|
||||
|
||||
use pyo3::{
|
||||
exceptions::{PyLookupError, PyRuntimeError, PyValueError},
|
||||
types::PyDict,
|
||||
};
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
|
||||
fn evaluate<'py>(py: Python<'py>, source: &str) -> Bound<'py, PyAny> {
|
||||
py.eval(&CString::new(source).unwrap(), None, None).unwrap()
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case("None", false, false)]
|
||||
#[case("False", false, false)]
|
||||
#[case("True", true, true)]
|
||||
#[case("0", false, false)]
|
||||
#[case("1", true, false)]
|
||||
#[case("''", false, false)]
|
||||
#[case("'false'", true, false)]
|
||||
#[case("[]", false, false)]
|
||||
#[case("[0]", true, false)]
|
||||
#[case("{}", false, false)]
|
||||
#[case("object()", true, false)]
|
||||
fn boolean_operations_have_distinct_python_semantics(
|
||||
#[case] source: &str,
|
||||
#[case] truth: bool,
|
||||
#[case] exact: bool,
|
||||
) {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let value = evaluate(py, source);
|
||||
let field = Field::new("test.flag", value.clone());
|
||||
assert_eq!(field.truthy().unwrap().0, truth);
|
||||
assert_eq!(field.exact_true().0, exact);
|
||||
assert_eq!(
|
||||
field.truthy().unwrap().0,
|
||||
py.import("builtins")
|
||||
.unwrap()
|
||||
.getattr("bool")
|
||||
.unwrap()
|
||||
.call1((value,))
|
||||
.unwrap()
|
||||
.extract::<bool>()
|
||||
.unwrap()
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case("None", Ok(None), Ok(None), Ok(None))]
|
||||
#[case("''", Ok(Some("")), Ok(None), Ok(None))]
|
||||
#[case(
|
||||
"' value '",
|
||||
Ok(Some(" value ")),
|
||||
Ok(Some(" value ")),
|
||||
Ok(Some(" value "))
|
||||
)]
|
||||
#[case("[]", Err(()), Ok(None), Ok(None))]
|
||||
#[case("0", Err(()), Ok(None), Ok(None))]
|
||||
#[case("1", Err(()), Err(()), Ok(None))]
|
||||
#[case("object()", Err(()), Err(()), Ok(None))]
|
||||
fn string_operations_do_not_conflate_absence_and_type_checks(
|
||||
#[case] source: &str,
|
||||
#[case] strict: Result<Option<&str>, ()>,
|
||||
#[case] fallback: Result<Option<&str>, ()>,
|
||||
#[case] tuning: Result<Option<&str>, ()>,
|
||||
) {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let field = Field::new("test.string", evaluate(py, source));
|
||||
let owned =
|
||||
|expected: Result<Option<&str>, ()>| expected.map(|value| value.map(str::to_owned));
|
||||
assert_eq!(
|
||||
field
|
||||
.optional_strict_string()
|
||||
.map(|value| value.0)
|
||||
.map_err(|_| ()),
|
||||
owned(strict)
|
||||
);
|
||||
assert_eq!(
|
||||
field
|
||||
.falsy_optional_string()
|
||||
.map(|value| value.0)
|
||||
.map_err(|_| ()),
|
||||
owned(fallback)
|
||||
);
|
||||
assert_eq!(
|
||||
field.tuning_string().map(|value| value.0).map_err(|_| ()),
|
||||
owned(tuning)
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case("None", None)]
|
||||
#[case("' True '", Some(true))]
|
||||
#[case("' fAlSe '", Some(false))]
|
||||
#[case("'yes'", None)]
|
||||
#[case("'1'", None)]
|
||||
#[case("'unknown'", None)]
|
||||
fn string_boolean_tokens_remain_separate_from_truthiness(
|
||||
#[case] source: &str,
|
||||
#[case] expected: Option<bool>,
|
||||
) {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
assert_eq!(
|
||||
Field::new("test.flag", evaluate(py, source))
|
||||
.str_bool()
|
||||
.unwrap()
|
||||
.0,
|
||||
expected
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case("'EXAMPLE.TEST.'", vec!["example.test"])]
|
||||
#[case("['B.test', '', None, 0, [], 'A.test.', 'b.test']", vec!["a.test", "b.test"])]
|
||||
#[case("('B.test', 'a.test')", vec!["a.test", "b.test"])]
|
||||
#[case("{'B.test', 'a.test'}", vec!["a.test", "b.test"])]
|
||||
#[case("(host for host in ['B.test', 'a.test'])", vec!["a.test", "b.test"])]
|
||||
#[case("None", vec![])]
|
||||
#[case("False", vec![])]
|
||||
fn host_collection_is_owned_normalized_and_deterministic(
|
||||
#[case] source: &str,
|
||||
#[case] expected: Vec<&str>,
|
||||
) {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
assert_eq!(
|
||||
Field::new("url_policy.user_url_allowed_hosts", evaluate(py, source))
|
||||
.host_collection()
|
||||
.unwrap()
|
||||
.0,
|
||||
expected
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn protocol_errors_preserve_exception_identity_traceback_cause_and_context() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let locals = PyDict::new(py);
|
||||
py.run(
|
||||
c"
|
||||
failure = LookupError('protocol failed')
|
||||
cause = ValueError('cause')
|
||||
context = RuntimeError('context')
|
||||
def fail():
|
||||
try:
|
||||
raise context
|
||||
except RuntimeError:
|
||||
raise failure from cause
|
||||
class Bool:
|
||||
def __bool__(self): return fail()
|
||||
class Length:
|
||||
def __len__(self): return fail()
|
||||
class Iter:
|
||||
def __iter__(self): return fail()
|
||||
class Next:
|
||||
def __iter__(self): return self
|
||||
def __next__(self): return fail()
|
||||
class Descriptor:
|
||||
@property
|
||||
def flag(self): return fail()
|
||||
values = (Bool(), Length(), Iter(), Next(), [Bool()])
|
||||
descriptor = Descriptor()
|
||||
",
|
||||
Some(&locals),
|
||||
Some(&locals),
|
||||
)
|
||||
.unwrap();
|
||||
let values = locals.get_item("values").unwrap().unwrap();
|
||||
for value in values.try_iter().unwrap() {
|
||||
let error = Field::new("test.flag", value.unwrap())
|
||||
.host_collection()
|
||||
.err()
|
||||
.unwrap();
|
||||
let error = PyErr::from(error);
|
||||
assert!(
|
||||
error
|
||||
.value(py)
|
||||
.is(locals.get_item("failure").unwrap().unwrap())
|
||||
);
|
||||
assert!(error.is_instance_of::<PyLookupError>(py));
|
||||
assert!(error.traceback(py).is_some());
|
||||
assert!(
|
||||
error
|
||||
.value(py)
|
||||
.getattr("__cause__")
|
||||
.unwrap()
|
||||
.is(locals.get_item("cause").unwrap().unwrap())
|
||||
);
|
||||
assert!(
|
||||
error
|
||||
.value(py)
|
||||
.getattr("__context__")
|
||||
.unwrap()
|
||||
.is(locals.get_item("context").unwrap().unwrap())
|
||||
);
|
||||
}
|
||||
let error = Field::read(
|
||||
&locals.get_item("descriptor").unwrap().unwrap(),
|
||||
"test.flag",
|
||||
)
|
||||
.err()
|
||||
.unwrap();
|
||||
assert!(
|
||||
PyErr::from(error)
|
||||
.value(py)
|
||||
.is(locals.get_item("failure").unwrap().unwrap())
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn identity_and_string_contents_do_not_invoke_unrelated_protocols() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let locals = PyDict::new(py);
|
||||
py.run(
|
||||
c"
|
||||
class Hostile:
|
||||
def __bool__(self): raise AssertionError('bool called')
|
||||
def __eq__(self, other): raise AssertionError('eq called')
|
||||
def __str__(self): raise AssertionError('str called')
|
||||
class Text(str):
|
||||
def __str__(self): raise AssertionError('str called')
|
||||
def strip(self): raise AssertionError('strip called')
|
||||
def lower(self): raise AssertionError('lower called')
|
||||
hostile = Hostile()
|
||||
text = Text(' False ')
|
||||
",
|
||||
Some(&locals),
|
||||
Some(&locals),
|
||||
)
|
||||
.unwrap();
|
||||
let hostile = Field::new("test.flag", locals.get_item("hostile").unwrap().unwrap());
|
||||
assert!(!hostile.exact_true().0);
|
||||
assert!(matches!(
|
||||
hostile.strict_string(),
|
||||
Err(ProjectionError::InvalidConfiguration(_))
|
||||
));
|
||||
let text = Field::new("test.flag", locals.get_item("text").unwrap().unwrap());
|
||||
assert_eq!(text.strict_string().unwrap(), " False ");
|
||||
assert_eq!(text.str_bool().unwrap().0, Some(false));
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn missing_snapshot_fields_and_descriptor_attribute_errors_are_distinct() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let locals = PyDict::new(py);
|
||||
py.run(
|
||||
c"
|
||||
failure = AttributeError('descriptor failed')
|
||||
class Snapshot:
|
||||
@property
|
||||
def flag(self): raise failure
|
||||
snapshot = Snapshot()
|
||||
class Dynamic:
|
||||
def __getattr__(self, name): raise failure
|
||||
class Intercepted:
|
||||
def __getattribute__(self, name): raise failure
|
||||
dynamic = Dynamic()
|
||||
intercepted = Intercepted()
|
||||
",
|
||||
Some(&locals),
|
||||
Some(&locals),
|
||||
)
|
||||
.unwrap();
|
||||
let snapshot = locals.get_item("snapshot").unwrap().unwrap();
|
||||
let descriptor = PyErr::from(Field::read(&snapshot, "test.flag").err().unwrap());
|
||||
assert!(
|
||||
descriptor
|
||||
.value(py)
|
||||
.is(locals.get_item("failure").unwrap().unwrap())
|
||||
);
|
||||
for name in ["dynamic", "intercepted"] {
|
||||
let value = locals.get_item(name).unwrap().unwrap();
|
||||
let error = PyErr::from(Field::read(&value, "test.flag").err().unwrap());
|
||||
assert!(
|
||||
error
|
||||
.value(py)
|
||||
.is(locals.get_item("failure").unwrap().unwrap())
|
||||
);
|
||||
}
|
||||
let missing = PyErr::from(Field::read(&snapshot, "test.missing").err().unwrap());
|
||||
assert!(missing.is_instance_of::<PyRuntimeError>(py));
|
||||
assert!(missing.to_string().contains("test.missing"));
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn configuration_errors_name_fields_without_exposing_values() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
for source in [
|
||||
"{'secret': 'do-not-print'}",
|
||||
"['host.test', {'secret': 'do-not-print'}]",
|
||||
] {
|
||||
let field = Field::new("test.setting", evaluate(py, source));
|
||||
let error = PyErr::from(field.falsy_optional_string().err().unwrap());
|
||||
assert!(error.is_instance_of::<PyValueError>(py));
|
||||
assert!(error.to_string().contains("test.setting"));
|
||||
assert!(!error.to_string().contains("do-not-print"));
|
||||
}
|
||||
let hosts = Field::new(
|
||||
"url_policy.user_url_allowed_hosts",
|
||||
evaluate(py, "['host.test', 1]"),
|
||||
);
|
||||
assert!(matches!(
|
||||
hosts.host_collection(),
|
||||
Err(ProjectionError::InvalidConfiguration(_))
|
||||
));
|
||||
assert!(matches!(
|
||||
Field::new("test.flag", evaluate(py, "1")).str_bool(),
|
||||
Err(ProjectionError::InvalidConfiguration(_))
|
||||
));
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn projection_releases_the_source_collection() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let source = evaluate(py, "['A.test']");
|
||||
let projected = Field::new("test.hosts", source.clone())
|
||||
.host_collection()
|
||||
.unwrap()
|
||||
.0;
|
||||
source.call_method1("append", ("b.test",)).unwrap();
|
||||
assert_eq!(projected, ["a.test"]);
|
||||
assert_eq!(
|
||||
Field::new("test.hosts", source)
|
||||
.host_collection()
|
||||
.unwrap()
|
||||
.0,
|
||||
["a.test", "b.test"]
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case("True", Some(true))]
|
||||
#[case("False", Some(false))]
|
||||
#[case("1", None)]
|
||||
#[case("None", None)]
|
||||
#[case("[]", None)]
|
||||
fn accessor_booleans_are_strict_schema_values(
|
||||
#[case] source: &str,
|
||||
#[case] expected: Option<bool>,
|
||||
) {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let result = Field::new("secret_manager.readable", evaluate(py, source)).schema_bool();
|
||||
match expected {
|
||||
Some(expected) => assert_eq!(result.unwrap(), expected),
|
||||
None => {
|
||||
let error = PyErr::from(result.unwrap_err());
|
||||
assert!(error.is_instance_of::<PyRuntimeError>(py));
|
||||
assert!(error.to_string().contains("secret_manager.readable"));
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
|
@ -1,5 +1,5 @@
|
|||
use std::{
|
||||
collections::HashSet,
|
||||
collections::{BTreeSet, HashSet},
|
||||
path::{Path, PathBuf},
|
||||
sync::{Arc, LazyLock, Mutex, PoisonError},
|
||||
};
|
||||
|
|
@ -10,9 +10,75 @@ use litellm_http::{
|
|||
TlsSource, Unsupported,
|
||||
media::{PublicDnsResolver, UrlPolicy},
|
||||
};
|
||||
use pyo3::{exceptions::PyValueError, prelude::*, types::PyDict};
|
||||
use pyo3::{
|
||||
exceptions::PyValueError,
|
||||
prelude::*,
|
||||
types::{PyBool, PyDict, PyString},
|
||||
};
|
||||
|
||||
use crate::{coercion::Field, python_settings::PythonSettings};
|
||||
use crate::{
|
||||
coercion::{Field, FieldSpec, ProjectionError},
|
||||
python_settings::{PythonSettings, Snapshot},
|
||||
};
|
||||
|
||||
const SSL_VERIFY: FieldSpec<Option<SslVerify>> = FieldSpec::new("ssl_verify", decode_ssl_verify);
|
||||
const SSL_CERTIFICATE: FieldSpec<Option<String>> =
|
||||
FieldSpec::new("ssl_certificate", |field| field.optional_strict_string());
|
||||
const SSL_SECURITY_LEVEL: FieldSpec<Option<String>> =
|
||||
FieldSpec::new("ssl_security_level", |field| field.tuning_string());
|
||||
const SSL_ECDH_CURVE: FieldSpec<Option<String>> =
|
||||
FieldSpec::new("ssl_ecdh_curve", |field| field.tuning_string());
|
||||
const FORCE_IPV4: FieldSpec<bool> = FieldSpec::new("force_ipv4", |field| field.truthy());
|
||||
const HTTP2: FieldSpec<bool> = FieldSpec::new("http2", |field| Ok(field.exact_true()));
|
||||
const AIOHTTP_TRUST_ENV: FieldSpec<bool> =
|
||||
FieldSpec::new("aiohttp_trust_env", |field| field.truthy());
|
||||
const DISABLE_AIOHTTP_TRUST_ENV: FieldSpec<bool> =
|
||||
FieldSpec::new("disable_aiohttp_trust_env", |field| field.truthy());
|
||||
const DISABLE_AIOHTTP_TRANSPORT: FieldSpec<bool> =
|
||||
FieldSpec::new("disable_aiohttp_transport", |field| Ok(field.exact_true()));
|
||||
const USER_AGENT: FieldSpec<String> = FieldSpec::new("user_agent", |field| field.schema_string());
|
||||
const USER_URL_VALIDATION: FieldSpec<bool> =
|
||||
FieldSpec::new("user_url_validation", |field| field.truthy());
|
||||
const USER_URL_ALLOWED_HOSTS: FieldSpec<Vec<String>> =
|
||||
FieldSpec::new("user_url_allowed_hosts", decode_hosts);
|
||||
|
||||
fn decode_hosts(field: &Field<'_>) -> Result<Vec<String>, ProjectionError> {
|
||||
Ok(field
|
||||
.string_collection()?
|
||||
.into_iter()
|
||||
.map(|host| litellm_http::media::normalize_host(&host))
|
||||
.collect::<BTreeSet<_>>()
|
||||
.into_iter()
|
||||
.collect())
|
||||
}
|
||||
|
||||
fn decode_ssl_verify(field: &Field<'_>) -> Result<Option<SslVerify>, ProjectionError> {
|
||||
let value = field.value();
|
||||
if value.is_none() {
|
||||
return Ok(None);
|
||||
}
|
||||
if value.is_instance_of::<PyBool>() {
|
||||
return Ok(Some(if field.exact_true() {
|
||||
SslVerify::Enabled
|
||||
} else {
|
||||
SslVerify::Disabled
|
||||
}));
|
||||
}
|
||||
if value.is_instance_of::<PyString>() {
|
||||
return Ok(Some(match field.str_bool()? {
|
||||
Some(true) => SslVerify::Enabled,
|
||||
Some(false) => SslVerify::Disabled,
|
||||
None => SslVerify::CaBundle(field.strict_string()?.into()),
|
||||
}));
|
||||
}
|
||||
let context = value.py().import("ssl")?.getattr("SSLContext")?;
|
||||
if value.is_instance(&context)? {
|
||||
return Err(ProjectionError::UnsupportedLiveObject(field.expected(
|
||||
"a Boolean, Boolean string, CA path, or None; live SSLContext is unsupported",
|
||||
)?));
|
||||
}
|
||||
Err(field.invalid("a Boolean, Boolean string, CA path, or None"))
|
||||
}
|
||||
|
||||
static POOL: LazyLock<HttpClientPool> =
|
||||
LazyLock::new(|| HttpClientPool::new(Arc::new(PublicDnsResolver)));
|
||||
|
|
@ -80,20 +146,20 @@ pub(crate) fn url_policy(py: Python<'_>) -> PyResult<UrlPolicy> {
|
|||
project_url_policy(&PythonSettings::UrlPolicy.read(py)?)
|
||||
}
|
||||
|
||||
fn project_url_policy(value: &Bound<'_, PyAny>) -> PyResult<UrlPolicy> {
|
||||
fn project_url_policy(snapshot: &Snapshot<'_>) -> PyResult<UrlPolicy> {
|
||||
Ok(UrlPolicy {
|
||||
validate: Field::read(value, "url_policy.user_url_validation")?
|
||||
.truthy()?
|
||||
.0,
|
||||
allowed_hosts: Field::read(value, "url_policy.user_url_allowed_hosts")?
|
||||
.host_collection()?
|
||||
.0,
|
||||
validate: snapshot.read(&USER_URL_VALIDATION)?,
|
||||
allowed_hosts: snapshot.read(&USER_URL_ALLOWED_HOSTS)?,
|
||||
})
|
||||
}
|
||||
|
||||
fn call_ssl_verify(kwargs: &Bound<'_, PyDict>) -> PyResult<Option<SslVerify>> {
|
||||
match kwargs.get_item("ssl_verify")? {
|
||||
Some(value) => Ok(Field::new("request.ssl_verify", value).ssl_verify()?.0),
|
||||
Some(value) => Ok(decode_ssl_verify(&Field::new(
|
||||
"request",
|
||||
"ssl_verify",
|
||||
value,
|
||||
))?),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
|
@ -106,39 +172,18 @@ fn for_call(call_ssl_verify: Option<SslVerify>, asynchronous: bool) -> HttpSetti
|
|||
}
|
||||
}
|
||||
|
||||
fn configured(value: &Bound<'_, PyAny>) -> PyResult<HttpSettingsLayer> {
|
||||
fn configured(snapshot: &Snapshot<'_>) -> PyResult<HttpSettingsLayer> {
|
||||
Ok(HttpSettingsLayer {
|
||||
ssl_verify: Field::read(value, "http_settings.ssl_verify")?
|
||||
.ssl_verify()?
|
||||
.0,
|
||||
ssl_certificate: Field::read(value, "http_settings.ssl_certificate")?
|
||||
.optional_strict_string()?
|
||||
.0
|
||||
.map(PathBuf::from),
|
||||
ssl_security_level: Field::read(value, "http_settings.ssl_security_level")?
|
||||
.tuning_string()?
|
||||
.0,
|
||||
ssl_ecdh_curve: Field::read(value, "http_settings.ssl_ecdh_curve")?
|
||||
.tuning_string()?
|
||||
.0,
|
||||
force_ipv4: Some(Field::read(value, "http_settings.force_ipv4")?.truthy()?.0),
|
||||
http2: Some(Field::read(value, "http_settings.http2")?.exact_true().0),
|
||||
aiohttp_trust_env: Some(
|
||||
Field::read(value, "http_settings.aiohttp_trust_env")?
|
||||
.truthy()?
|
||||
.0,
|
||||
),
|
||||
disable_aiohttp_trust_env: Some(
|
||||
Field::read(value, "http_settings.disable_aiohttp_trust_env")?
|
||||
.truthy()?
|
||||
.0,
|
||||
),
|
||||
disable_aiohttp_transport: Some(
|
||||
Field::read(value, "http_settings.disable_aiohttp_transport")?
|
||||
.exact_true()
|
||||
.0,
|
||||
),
|
||||
user_agent: Some(Field::read(value, "http_settings.user_agent")?.schema_string()?),
|
||||
ssl_verify: snapshot.read(&SSL_VERIFY)?,
|
||||
ssl_certificate: snapshot.read(&SSL_CERTIFICATE)?.map(PathBuf::from),
|
||||
ssl_security_level: snapshot.read(&SSL_SECURITY_LEVEL)?,
|
||||
ssl_ecdh_curve: snapshot.read(&SSL_ECDH_CURVE)?,
|
||||
force_ipv4: Some(snapshot.read(&FORCE_IPV4)?),
|
||||
http2: Some(snapshot.read(&HTTP2)?),
|
||||
aiohttp_trust_env: Some(snapshot.read(&AIOHTTP_TRUST_ENV)?),
|
||||
disable_aiohttp_trust_env: Some(snapshot.read(&DISABLE_AIOHTTP_TRUST_ENV)?),
|
||||
disable_aiohttp_transport: Some(snapshot.read(&DISABLE_AIOHTTP_TRANSPORT)?),
|
||||
user_agent: Some(snapshot.read(&USER_AGENT)?),
|
||||
..HttpSettingsLayer::default()
|
||||
})
|
||||
}
|
||||
|
|
@ -150,12 +195,15 @@ mod tests {
|
|||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
use crate::python_settings::CONTRACT;
|
||||
|
||||
fn python_settings<'py>(py: Python<'py>, overrides: &str) -> Bound<'py, PyAny> {
|
||||
fn evaluate<'py>(py: Python<'py>, source: &str) -> Bound<'py, PyAny> {
|
||||
py.eval(&std::ffi::CString::new(source).unwrap(), None, None)
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
fn python_settings<'py>(py: Python<'py>, overrides: &str) -> Snapshot<'py> {
|
||||
let source = format!(
|
||||
"
|
||||
import json
|
||||
import types
|
||||
defaults = dict(
|
||||
ssl_verify=True,
|
||||
|
|
@ -170,14 +218,13 @@ defaults = dict(
|
|||
user_agent='litellm/test',
|
||||
)
|
||||
defaults.update(dict({overrides}))
|
||||
settings = types.SimpleNamespace(**{{name: defaults[name] for name in json.loads(contract)['http_settings']['fields']}})
|
||||
settings = types.SimpleNamespace(**defaults)
|
||||
"
|
||||
);
|
||||
let locals = PyDict::new(py);
|
||||
locals.set_item("contract", CONTRACT).unwrap();
|
||||
let source = std::ffi::CString::new(source).unwrap();
|
||||
py.run(&source, Some(&locals), Some(&locals)).unwrap();
|
||||
locals.get_item("settings").unwrap().unwrap()
|
||||
PythonSettings::Http.snapshot(locals.get_item("settings").unwrap().unwrap())
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
|
@ -395,7 +442,7 @@ user_agent='litellm/9.9.9',
|
|||
Python::attach(|py| {
|
||||
let value = py.eval(c"__import__('types').SimpleNamespace(user_url_validation=[], user_url_allowed_hosts=['B.test', 'a.test.', 'b.test'])", None, None).unwrap();
|
||||
assert_eq!(
|
||||
project_url_policy(&value).unwrap(),
|
||||
project_url_policy(&PythonSettings::UrlPolicy.snapshot(value)).unwrap(),
|
||||
UrlPolicy {
|
||||
validate: false,
|
||||
allowed_hosts: vec!["a.test".into(), "b.test".into()],
|
||||
|
|
@ -419,4 +466,44 @@ user_agent='litellm/9.9.9',
|
|||
let settings = HttpSettings::from_layers([for_call(None, asynchronous), opted_out]);
|
||||
assert_eq!(settings.trust_proxy_env, expected);
|
||||
}
|
||||
#[rstest]
|
||||
#[case("'EXAMPLE.TEST.'", vec!["example.test"])]
|
||||
#[case("['B.test', '', None, 0, [], 'A.test.', 'b.test']", vec!["a.test", "b.test"])]
|
||||
#[case("('B.test', 'a.test')", vec!["a.test", "b.test"])]
|
||||
#[case("{'B.test', 'a.test'}", vec!["a.test", "b.test"])]
|
||||
#[case("(host for host in ['B.test', 'a.test'])", vec!["a.test", "b.test"])]
|
||||
#[case("None", vec![])]
|
||||
#[case("False", vec![])]
|
||||
fn host_collection_is_owned_normalized_and_deterministic(
|
||||
#[case] source: &str,
|
||||
#[case] expected: Vec<&str>,
|
||||
) {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
assert_eq!(
|
||||
decode_hosts(&Field::new(
|
||||
"url_policy",
|
||||
"user_url_allowed_hosts",
|
||||
evaluate(py, source)
|
||||
))
|
||||
.unwrap(),
|
||||
expected
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn projection_releases_the_source_collection() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let source = evaluate(py, "['A.test']");
|
||||
let projected = decode_hosts(&Field::new("test", "hosts", source.clone())).unwrap();
|
||||
source.call_method1("append", ("b.test",)).unwrap();
|
||||
assert_eq!(projected, ["a.test"]);
|
||||
assert_eq!(
|
||||
decode_hosts(&Field::new("test", "hosts", source)).unwrap(),
|
||||
["a.test", "b.test"]
|
||||
);
|
||||
});
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -7,7 +7,12 @@ mod http;
|
|||
mod marshal;
|
||||
mod python_settings;
|
||||
mod routes;
|
||||
mod token_counter;
|
||||
#[allow(
|
||||
dead_code,
|
||||
reason = "secret-manager foundations await rollout activation"
|
||||
)]
|
||||
mod secrets;
|
||||
mod tokenizer;
|
||||
|
||||
#[pymodule(gil_used = true)]
|
||||
mod _native {
|
||||
|
|
@ -32,7 +37,12 @@ mod _native {
|
|||
#[pymodule_export]
|
||||
use crate::routes::responses::ResponsesWebSocketConnection;
|
||||
#[pymodule_export]
|
||||
use crate::token_counter::TokenCounter;
|
||||
use crate::routes::token_counter::TokenCounter;
|
||||
#[cfg(feature = "huggingface")]
|
||||
#[pymodule_export]
|
||||
use crate::tokenizer::HuggingFaceEncoding;
|
||||
#[pymodule_export]
|
||||
use crate::tokenizer::Tokenizer;
|
||||
#[pymodule_export]
|
||||
use litellm_host_python::{ForkedAfterNativeRuntimeStarted, ProcessReservedForForking};
|
||||
use pyo3::{prelude::*, types::PyModule};
|
||||
|
|
@ -43,7 +53,7 @@ mod _native {
|
|||
let dict = module.dict();
|
||||
dict.set_item("_CacheTestHandle", py.get_type::<CacheTestHandle>())?;
|
||||
dict.set_item("_CacheTestResolver", py.get_type::<CacheTestResolver>())?;
|
||||
dict.set_item("_CacheTestBinding", py.get_type::<ResolvedCache>())
|
||||
dict.set_item("_ResponseCacheRuntime", py.get_type::<ResolvedCache>())
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -78,10 +88,13 @@ mod tests {
|
|||
"achat_completions",
|
||||
"ResponsesWebSocketConnection",
|
||||
"TokenCounter",
|
||||
"Tokenizer",
|
||||
"gil_stats",
|
||||
"process_state_started",
|
||||
"reserve_process_for_forking",
|
||||
];
|
||||
#[cfg(feature = "huggingface")]
|
||||
expected.push("HuggingFaceEncoding");
|
||||
expected.sort_unstable();
|
||||
|
||||
let mut public_names: Vec<String> = native_module(py)
|
||||
|
|
|
|||
|
|
@ -1,5 +1,7 @@
|
|||
use pyo3::prelude::*;
|
||||
|
||||
use crate::coercion::{FieldSpec, ProjectionError};
|
||||
|
||||
const MODULE: &str = "litellm.rust_bridge.settings";
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
|
|
@ -8,28 +10,39 @@ pub(crate) enum PythonSettings {
|
|||
UrlPolicy,
|
||||
ProviderDefaults,
|
||||
SecretManager,
|
||||
SecretManagerBinding,
|
||||
}
|
||||
|
||||
pub(crate) struct Snapshot<'py> {
|
||||
group: PythonSettings,
|
||||
value: Bound<'py, PyAny>,
|
||||
}
|
||||
|
||||
impl Snapshot<'_> {
|
||||
pub(crate) fn read<T>(&self, spec: &FieldSpec<T>) -> Result<T, ProjectionError> {
|
||||
spec.read(&self.value, self.group.name())
|
||||
}
|
||||
}
|
||||
|
||||
impl PythonSettings {
|
||||
#[cfg(test)]
|
||||
pub(crate) const ALL: [Self; 4] = [
|
||||
Self::Http,
|
||||
Self::UrlPolicy,
|
||||
Self::ProviderDefaults,
|
||||
Self::SecretManager,
|
||||
];
|
||||
|
||||
pub(crate) fn name(self) -> &'static str {
|
||||
match self {
|
||||
Self::Http => "http_settings",
|
||||
Self::UrlPolicy => "url_policy",
|
||||
Self::ProviderDefaults => "provider_defaults",
|
||||
Self::SecretManager => "secret_manager",
|
||||
Self::SecretManagerBinding => "secret_manager_binding",
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn read(self, py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
|
||||
py.import(MODULE)?.getattr(self.name())?.call0()
|
||||
pub(crate) fn read(self, py: Python<'_>) -> PyResult<Snapshot<'_>> {
|
||||
let value = py.import(MODULE)?.getattr(self.name())?.call0()?;
|
||||
Ok(Snapshot { group: self, value })
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) fn snapshot(self, value: Bound<'_, PyAny>) -> Snapshot<'_> {
|
||||
Snapshot { group: self, value }
|
||||
}
|
||||
|
||||
pub(crate) fn warn(py: Python<'_>, message: &str) -> PyResult<()> {
|
||||
|
|
@ -38,209 +51,98 @@ impl PythonSettings {
|
|||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) const CONTRACT: &str = include_str!("../python_settings.json");
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{CONTRACT, PythonSettings};
|
||||
use pyo3::prelude::*;
|
||||
use serde_json::{Value, json};
|
||||
use pyo3::{exceptions::PyRuntimeError, prelude::*, types::PyDict};
|
||||
|
||||
struct SettingSpec {
|
||||
group: &'static str,
|
||||
name: &'static str,
|
||||
adapter: &'static str,
|
||||
precedence: &'static str,
|
||||
sensitive: bool,
|
||||
shapes: &'static [&'static str],
|
||||
unsupported_live: Option<&'static str>,
|
||||
}
|
||||
|
||||
const SETTINGS: &[SettingSpec] = &[
|
||||
SettingSpec {
|
||||
group: "http_settings",
|
||||
name: "ssl_verify",
|
||||
adapter: "SslVerifyInput",
|
||||
precedence: "module_global",
|
||||
sensitive: false,
|
||||
shapes: &["none", "bool", "str"],
|
||||
unsupported_live: Some("configuration_error"),
|
||||
},
|
||||
SettingSpec {
|
||||
group: "http_settings",
|
||||
name: "ssl_certificate",
|
||||
adapter: "OptionalStrictString",
|
||||
precedence: "module_global",
|
||||
sensitive: false,
|
||||
shapes: &[],
|
||||
unsupported_live: None,
|
||||
},
|
||||
SettingSpec {
|
||||
group: "http_settings",
|
||||
name: "ssl_security_level",
|
||||
adapter: "TuningString",
|
||||
precedence: "module_global",
|
||||
sensitive: false,
|
||||
shapes: &[],
|
||||
unsupported_live: None,
|
||||
},
|
||||
SettingSpec {
|
||||
group: "http_settings",
|
||||
name: "ssl_ecdh_curve",
|
||||
adapter: "TuningString",
|
||||
precedence: "module_global",
|
||||
sensitive: false,
|
||||
shapes: &[],
|
||||
unsupported_live: None,
|
||||
},
|
||||
SettingSpec {
|
||||
group: "http_settings",
|
||||
name: "force_ipv4",
|
||||
adapter: "Truthy",
|
||||
precedence: "module_global",
|
||||
sensitive: false,
|
||||
shapes: &[],
|
||||
unsupported_live: None,
|
||||
},
|
||||
SettingSpec {
|
||||
group: "http_settings",
|
||||
name: "http2",
|
||||
adapter: "ExactTrue",
|
||||
precedence: "module_global",
|
||||
sensitive: false,
|
||||
shapes: &[],
|
||||
unsupported_live: None,
|
||||
},
|
||||
SettingSpec {
|
||||
group: "http_settings",
|
||||
name: "aiohttp_trust_env",
|
||||
adapter: "Truthy",
|
||||
precedence: "module_global",
|
||||
sensitive: false,
|
||||
shapes: &[],
|
||||
unsupported_live: None,
|
||||
},
|
||||
SettingSpec {
|
||||
group: "http_settings",
|
||||
name: "disable_aiohttp_trust_env",
|
||||
adapter: "Truthy",
|
||||
precedence: "module_global",
|
||||
sensitive: false,
|
||||
shapes: &[],
|
||||
unsupported_live: None,
|
||||
},
|
||||
SettingSpec {
|
||||
group: "http_settings",
|
||||
name: "disable_aiohttp_transport",
|
||||
adapter: "ExactTrue",
|
||||
precedence: "module_global",
|
||||
sensitive: false,
|
||||
shapes: &[],
|
||||
unsupported_live: None,
|
||||
},
|
||||
SettingSpec {
|
||||
group: "http_settings",
|
||||
name: "user_agent",
|
||||
adapter: "StrictString",
|
||||
precedence: "accessor",
|
||||
sensitive: false,
|
||||
shapes: &[],
|
||||
unsupported_live: None,
|
||||
},
|
||||
SettingSpec {
|
||||
group: "url_policy",
|
||||
name: "user_url_validation",
|
||||
adapter: "Truthy",
|
||||
precedence: "module_global",
|
||||
sensitive: false,
|
||||
shapes: &[],
|
||||
unsupported_live: None,
|
||||
},
|
||||
SettingSpec {
|
||||
group: "url_policy",
|
||||
name: "user_url_allowed_hosts",
|
||||
adapter: "HostCollection",
|
||||
precedence: "module_global",
|
||||
sensitive: false,
|
||||
shapes: &[],
|
||||
unsupported_live: None,
|
||||
},
|
||||
SettingSpec {
|
||||
group: "provider_defaults",
|
||||
name: "vertex_project",
|
||||
adapter: "FalsyOptionalString",
|
||||
precedence: "module_global",
|
||||
sensitive: true,
|
||||
shapes: &[],
|
||||
unsupported_live: None,
|
||||
},
|
||||
SettingSpec {
|
||||
group: "provider_defaults",
|
||||
name: "vertex_location",
|
||||
adapter: "FalsyOptionalString",
|
||||
precedence: "module_global",
|
||||
sensitive: true,
|
||||
shapes: &[],
|
||||
unsupported_live: None,
|
||||
},
|
||||
SettingSpec {
|
||||
group: "provider_defaults",
|
||||
name: "enable_azure_ad_token_refresh",
|
||||
adapter: "ExactTrue",
|
||||
precedence: "module_global",
|
||||
sensitive: false,
|
||||
shapes: &[],
|
||||
unsupported_live: None,
|
||||
},
|
||||
SettingSpec {
|
||||
group: "secret_manager",
|
||||
name: "readable",
|
||||
adapter: "StrictBool",
|
||||
precedence: "accessor",
|
||||
sensitive: false,
|
||||
shapes: &[],
|
||||
unsupported_live: None,
|
||||
},
|
||||
];
|
||||
use super::PythonSettings;
|
||||
use crate::coercion::FieldSpec;
|
||||
|
||||
#[test]
|
||||
fn settings_manifest_matches_the_semantic_contract() {
|
||||
pyo3::Python::initialize();
|
||||
let manifest: Value = pyo3::Python::attach(|py| {
|
||||
let value = py
|
||||
.import("json")
|
||||
.unwrap()
|
||||
.call_method1("loads", (CONTRACT,))
|
||||
.unwrap();
|
||||
litellm_host_python::from_py(&value).unwrap()
|
||||
fn declarations_select_the_decoder_and_read_only_the_requested_field() {
|
||||
const TRUTHY: FieldSpec<bool> = FieldSpec::new("flag", |field| field.truthy());
|
||||
const EXACT: FieldSpec<bool> = FieldSpec::new("flag", |field| Ok(field.exact_true()));
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let locals = PyDict::new(py);
|
||||
py.run(
|
||||
c"
|
||||
reads = []
|
||||
class Settings:
|
||||
value = 1
|
||||
@property
|
||||
def flag(self):
|
||||
reads.append('flag')
|
||||
return self.value
|
||||
@property
|
||||
def unrelated(self):
|
||||
raise AssertionError('unrequested field')
|
||||
settings = Settings()
|
||||
",
|
||||
Some(&locals),
|
||||
Some(&locals),
|
||||
)
|
||||
.unwrap();
|
||||
let value = locals.get_item("settings").unwrap().unwrap();
|
||||
let snapshot = PythonSettings::Http.snapshot(value.clone());
|
||||
assert!(snapshot.read(&TRUTHY).unwrap());
|
||||
assert!(!snapshot.read(&EXACT).unwrap());
|
||||
value.setattr("value", true).unwrap();
|
||||
assert!(snapshot.read(&EXACT).unwrap());
|
||||
assert_eq!(
|
||||
locals
|
||||
.get_item("reads")
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.extract::<Vec<String>>()
|
||||
.unwrap(),
|
||||
["flag", "flag", "flag"]
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn declared_reads_preserve_descriptor_and_decoder_failures_and_name_missing_fields() {
|
||||
const FLAG: FieldSpec<bool> = FieldSpec::new("flag", |field| field.truthy());
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let locals = PyDict::new(py);
|
||||
py.run(
|
||||
c"
|
||||
from types import SimpleNamespace
|
||||
failure = AttributeError('read failed')
|
||||
class Descriptor:
|
||||
@property
|
||||
def flag(self): raise failure
|
||||
class Truth:
|
||||
def __bool__(self): raise failure
|
||||
values = (Descriptor(), SimpleNamespace(flag=Truth()))
|
||||
",
|
||||
Some(&locals),
|
||||
Some(&locals),
|
||||
)
|
||||
.unwrap();
|
||||
let failure = locals.get_item("failure").unwrap().unwrap();
|
||||
for value in locals
|
||||
.get_item("values")
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.try_iter()
|
||||
.unwrap()
|
||||
{
|
||||
let snapshot = PythonSettings::Http.snapshot(value.unwrap());
|
||||
let error = PyErr::from(snapshot.read(&FLAG).unwrap_err());
|
||||
assert!(error.value(py).is(&failure));
|
||||
assert!(error.traceback(py).is_some());
|
||||
}
|
||||
let missing = PythonSettings::Http.snapshot(py.eval(c"object()", None, None).unwrap());
|
||||
let error = PyErr::from(missing.read(&FLAG).unwrap_err());
|
||||
assert!(error.is_instance_of::<PyRuntimeError>(py));
|
||||
assert!(
|
||||
error
|
||||
.to_string()
|
||||
.contains("http_settings.flag: missing snapshot field")
|
||||
);
|
||||
});
|
||||
let expected: serde_json::Map<String, Value> = PythonSettings::ALL
|
||||
.into_iter()
|
||||
.map(|group| {
|
||||
let fields: serde_json::Map<String, Value> = SETTINGS
|
||||
.iter()
|
||||
.filter(|spec| spec.group == group.name())
|
||||
.map(|spec| {
|
||||
(
|
||||
spec.name.to_owned(),
|
||||
json!({
|
||||
"adapter": spec.adapter,
|
||||
"required": true,
|
||||
"precedence": spec.precedence,
|
||||
"sensitive": spec.sensitive,
|
||||
"shapes": spec.shapes,
|
||||
"unsupported_live": spec.unsupported_live,
|
||||
}),
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
(
|
||||
group.name().to_owned(),
|
||||
json!({"version": 1, "fields": fields}),
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
assert_eq!(manifest, Value::Object(expected));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ pub(crate) mod chat_completions;
|
|||
pub(crate) mod messages;
|
||||
pub(crate) mod ocr;
|
||||
pub(crate) mod responses;
|
||||
pub(crate) mod token_counter;
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
|
|
|
|||
|
|
@ -130,6 +130,11 @@ impl RouteHost for OcrRouteHost {
|
|||
}
|
||||
|
||||
fn classify(&self, py: Python<'_>, error: Error) -> PyResult<PyErr> {
|
||||
if let Error::Secret(source) = &error
|
||||
&& let Some(original) = crate::secrets::callback::python_error(py, source)
|
||||
{
|
||||
return Ok(original);
|
||||
}
|
||||
Ok(self.map_failure(py, ocr_error_to_pyerr(error)))
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -10,16 +10,33 @@ use litellm_auth_gcp::VertexAuth;
|
|||
use litellm_callbacks_legacy_python::{LegacySurface, PublicCall, run_legacy_call};
|
||||
use litellm_core::ocr::route::ocr_machine;
|
||||
use litellm_core_utils::settings::ProcessEnvironment;
|
||||
use litellm_llms::base_llm::ocr::{
|
||||
handler::OcrClient,
|
||||
settings::{OcrSettings, Secrets},
|
||||
use litellm_llms::base_llm::{
|
||||
inference::secrets::{EnvironmentSecrets, SecretSource},
|
||||
ocr::{handler::OcrClient, settings::OcrSettings},
|
||||
};
|
||||
use pyo3::{
|
||||
prelude::*,
|
||||
types::{PyDict, PyTuple},
|
||||
};
|
||||
|
||||
use crate::{coercion::Field, errors::RustBridgeDeclined, http, python_settings::PythonSettings};
|
||||
use crate::{
|
||||
coercion::FieldSpec,
|
||||
errors::RustBridgeDeclined,
|
||||
http,
|
||||
python_settings::{PythonSettings, Snapshot},
|
||||
};
|
||||
|
||||
const SECRET_MANAGER_READABLE: FieldSpec<bool> =
|
||||
FieldSpec::new("readable", |field| field.schema_bool());
|
||||
|
||||
const VERTEX_PROJECT: FieldSpec<Option<String>> =
|
||||
FieldSpec::new("vertex_project", |field| field.falsy_optional_string());
|
||||
const VERTEX_LOCATION: FieldSpec<Option<String>> =
|
||||
FieldSpec::new("vertex_location", |field| field.falsy_optional_string());
|
||||
const ENABLE_AZURE_AD_TOKEN_REFRESH: FieldSpec<bool> =
|
||||
FieldSpec::new("enable_azure_ad_token_refresh", |field| {
|
||||
Ok(field.exact_true())
|
||||
});
|
||||
|
||||
const SURFACE: LegacySurface = LegacySurface {
|
||||
call_type: "ocr",
|
||||
|
|
@ -62,33 +79,24 @@ fn run_ocr(
|
|||
)
|
||||
}
|
||||
|
||||
fn process_environment_secrets(secret_manager: &Bound<'_, PyAny>) -> PyResult<Secrets> {
|
||||
if Field::read(secret_manager, "secret_manager.readable")?.schema_bool()? {
|
||||
fn process_environment_secrets(snapshot: &Snapshot<'_>) -> PyResult<Arc<dyn SecretSource>> {
|
||||
if snapshot.read(&SECRET_MANAGER_READABLE)? {
|
||||
return Err(RustBridgeDeclined::new_err(
|
||||
"a readable secret manager is configured and the Rust route only reads the process environment",
|
||||
));
|
||||
}
|
||||
Ok(Arc::new(ProcessEnvironment))
|
||||
Ok(Arc::new(EnvironmentSecrets))
|
||||
}
|
||||
|
||||
fn ocr_settings(py: Python<'_>) -> PyResult<OcrSettings> {
|
||||
project_provider_defaults(&PythonSettings::ProviderDefaults.read(py)?)
|
||||
}
|
||||
|
||||
fn project_provider_defaults(value: &Bound<'_, PyAny>) -> PyResult<OcrSettings> {
|
||||
fn project_provider_defaults(snapshot: &Snapshot<'_>) -> PyResult<OcrSettings> {
|
||||
Ok(OcrSettings {
|
||||
vertex_project: Field::read(value, "provider_defaults.vertex_project")?
|
||||
.falsy_optional_string()?
|
||||
.0,
|
||||
vertex_location: Field::read(value, "provider_defaults.vertex_location")?
|
||||
.falsy_optional_string()?
|
||||
.0,
|
||||
enable_azure_ad_token_refresh: Field::read(
|
||||
value,
|
||||
"provider_defaults.enable_azure_ad_token_refresh",
|
||||
)?
|
||||
.exact_true()
|
||||
.0,
|
||||
vertex_project: snapshot.read(&VERTEX_PROJECT)?,
|
||||
vertex_location: snapshot.read(&VERTEX_LOCATION)?,
|
||||
enable_azure_ad_token_refresh: snapshot.read(&ENABLE_AZURE_AD_TOKEN_REFRESH)?,
|
||||
..OcrSettings::from_environment(&ProcessEnvironment)
|
||||
})
|
||||
}
|
||||
|
|
@ -120,6 +128,8 @@ mod tests {
|
|||
use super::process_environment_secrets;
|
||||
use crate::errors::RustBridgeDeclined;
|
||||
|
||||
use crate::python_settings::PythonSettings;
|
||||
|
||||
fn secret_manager<'py>(py: Python<'py>, readable: bool) -> Bound<'py, PyAny> {
|
||||
let locals = PyDict::new(py);
|
||||
locals.set_item("readable", readable).unwrap();
|
||||
|
|
@ -132,12 +142,26 @@ mod tests {
|
|||
locals.get_item("manager").unwrap().unwrap()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_readable_secret_manager_sends_the_call_back_to_python() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let declined = process_environment_secrets(
|
||||
&PythonSettings::SecretManager.snapshot(secret_manager(py, true)),
|
||||
)
|
||||
.err()
|
||||
.expect("the Rust route declines");
|
||||
assert!(declined.is_instance_of::<RustBridgeDeclined>(py));
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_defaults_distinguish_falsey_values_and_exact_true() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let value = py.eval(c"__import__('types').SimpleNamespace(vertex_project=[], vertex_location=0, enable_azure_ad_token_refresh=1)", None, None).unwrap();
|
||||
let projected = super::project_provider_defaults(&value).unwrap();
|
||||
let snapshot = PythonSettings::ProviderDefaults.snapshot(value.clone());
|
||||
let projected = super::project_provider_defaults(&snapshot).unwrap();
|
||||
assert_eq!(projected.vertex_project, None);
|
||||
assert_eq!(projected.vertex_location, None);
|
||||
assert!(!projected.enable_azure_ad_token_refresh);
|
||||
|
|
@ -146,12 +170,12 @@ mod tests {
|
|||
value
|
||||
.setattr("enable_azure_ad_token_refresh", true)
|
||||
.unwrap();
|
||||
let next = super::project_provider_defaults(&value).unwrap();
|
||||
let next = super::project_provider_defaults(&snapshot).unwrap();
|
||||
assert_eq!(next.vertex_project.as_deref(), Some("project"));
|
||||
assert_eq!(next.vertex_location.as_deref(), Some("region"));
|
||||
assert!(next.enable_azure_ad_token_refresh);
|
||||
value.setattr("vertex_project", 1).unwrap();
|
||||
let error = super::project_provider_defaults(&value).err().unwrap();
|
||||
let error = super::project_provider_defaults(&snapshot).err().unwrap();
|
||||
assert!(error.is_instance_of::<pyo3::exceptions::PyValueError>(py));
|
||||
assert!(
|
||||
error
|
||||
|
|
@ -160,28 +184,4 @@ mod tests {
|
|||
);
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_readable_secret_manager_sends_the_call_back_to_python() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let declined = process_environment_secrets(&secret_manager(py, true))
|
||||
.err()
|
||||
.expect("the Rust route declines");
|
||||
assert!(declined.is_instance_of::<RustBridgeDeclined>(py));
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn without_a_readable_secret_manager_secrets_are_the_process_environment() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let secrets = process_environment_secrets(&secret_manager(py, false)).unwrap();
|
||||
assert_eq!(
|
||||
secrets.get("LITELLM_RUST_BRIDGE_UNSET_VARIABLE_FOR_TEST"),
|
||||
None
|
||||
);
|
||||
assert_eq!(secrets.get("PATH"), std::env::var("PATH").ok());
|
||||
});
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,87 @@
|
|||
use std::sync::Arc;
|
||||
use std::{num::NonZero, thread::available_parallelism};
|
||||
|
||||
use litellm_host_python::{enter_native, run_async};
|
||||
use litellm_token_counter::{
|
||||
CountableRequest, Error, InputTokenCount, TokenCounter as CoreTokenCounter,
|
||||
};
|
||||
use pyo3::{
|
||||
exceptions::{PyRuntimeError, PyValueError},
|
||||
prelude::*,
|
||||
types::PyAny,
|
||||
};
|
||||
use tokio::sync::Semaphore;
|
||||
|
||||
use crate::errors::RustBridgeDeclined;
|
||||
use crate::tokenizer::Tokenizer;
|
||||
|
||||
/// Counts the input tokens of a raw request body off the Python event loop with
|
||||
/// the GIL released. Python owns which requests get here and what to do with
|
||||
/// the count. At most one encode per core runs at a time; the rest wait in the
|
||||
/// async task, where a cancelled Python awaiter drops them before any blocking
|
||||
/// work is scheduled.
|
||||
#[pyclass(frozen)]
|
||||
pub(crate) struct TokenCounter {
|
||||
inner: Arc<CoreTokenCounter>,
|
||||
encode_slots: Arc<Semaphore>,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl TokenCounter {
|
||||
#[staticmethod]
|
||||
#[pyo3(signature = (tokenizer, fast = false))]
|
||||
fn from_tokenizer(py: Python<'_>, tokenizer: &Tokenizer, fast: bool) -> PyResult<Self> {
|
||||
enter_native()?;
|
||||
let inner = CoreTokenCounter::new(tokenizer.counter(py, fast));
|
||||
Ok(Self {
|
||||
inner: Arc::new(inner),
|
||||
encode_slots: Arc::new(Semaphore::new(encode_parallelism())),
|
||||
})
|
||||
}
|
||||
|
||||
fn acount_request<'py>(&self, py: Python<'py>, body: &[u8]) -> PyResult<Bound<'py, PyAny>> {
|
||||
let counter = Arc::clone(&self.inner);
|
||||
let encode_slots = Arc::clone(&self.encode_slots);
|
||||
let body = body.to_vec();
|
||||
run_async(
|
||||
py,
|
||||
async move {
|
||||
let _slot = encode_slots
|
||||
.acquire_owned()
|
||||
.await
|
||||
.map_err(|error| Error::Task(error.to_string()))?;
|
||||
tokio::task::spawn_blocking(move || count_body(&counter, &body))
|
||||
.await
|
||||
.map_err(|error| Error::Task(error.to_string()))?
|
||||
},
|
||||
token_count_error_to_pyerr,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
fn encode_parallelism() -> usize {
|
||||
available_parallelism().map_or(1, NonZero::get)
|
||||
}
|
||||
|
||||
fn count_body(counter: &CoreTokenCounter, body: &[u8]) -> Result<InputTokenCount, Error> {
|
||||
let request = CountableRequest::parse(body)?;
|
||||
counter.count_request(&request)
|
||||
}
|
||||
|
||||
pub(crate) fn token_count_error_to_pyerr(error: Error) -> PyErr {
|
||||
let message = error.to_string();
|
||||
match error {
|
||||
Error::Load(_)
|
||||
| Error::Ranks(_)
|
||||
| Error::UnicodeClasses
|
||||
| Error::UnsupportedTokenizer(_) => PyValueError::new_err(message),
|
||||
Error::RequestParse(_)
|
||||
| Error::MissingInput
|
||||
| Error::FloatText
|
||||
| Error::ContentBlock
|
||||
| Error::ArrayItems
|
||||
| Error::JsonSerialization(_)
|
||||
| Error::JsonUtf8(_) => RustBridgeDeclined::new_err(message),
|
||||
Error::Encode(_) | Error::Decode(_) | Error::Task(_) => PyRuntimeError::new_err(message),
|
||||
}
|
||||
}
|
||||
349
litellm-rust/crates/python-bridge/src/secrets/callback.rs
Normal file
349
litellm-rust/crates/python-bridge/src/secrets/callback.rs
Normal file
|
|
@ -0,0 +1,349 @@
|
|||
use std::{fmt, future::Future, pin::Pin};
|
||||
|
||||
use litellm_core_utils::settings::Lookup;
|
||||
use litellm_secrets::{
|
||||
Error, ExternalSecretManager, KeyManagementSettings, KeyManagementSystem, Secret, SecretValue,
|
||||
};
|
||||
use pyo3::{exceptions::PyBaseException, prelude::*, types::PyDict};
|
||||
|
||||
const HANDLER_MODULE: &str = "litellm.secret_managers.secret_manager_handler";
|
||||
|
||||
struct PythonSecretError(Py<PyBaseException>);
|
||||
|
||||
impl fmt::Debug for PythonSecretError {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
formatter.write_str("PythonSecretError")
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Display for PythonSecretError {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
formatter.write_str("Python secret manager failed")
|
||||
}
|
||||
}
|
||||
|
||||
impl std::error::Error for PythonSecretError {}
|
||||
|
||||
pub(crate) fn python_error(py: Python<'_>, error: &Error) -> Option<PyErr> {
|
||||
let Error::ExternalManager(source) = error else {
|
||||
return None;
|
||||
};
|
||||
source
|
||||
.downcast_ref::<PythonSecretError>()
|
||||
.map(|error| PyErr::from_value(error.0.clone_ref(py).into_bound(py).into_any()))
|
||||
}
|
||||
|
||||
/// A secret manager whose reads execute in Python: a custom manager, a legacy compatible
|
||||
/// client, or a manually assigned SDK client.
|
||||
pub(crate) struct PythonSecretManager {
|
||||
client: Py<PyAny>,
|
||||
system: Option<KeyManagementSystem>,
|
||||
/// The `key_manager` name Python's handler dispatches on.
|
||||
key_manager: &'static str,
|
||||
settings: Option<Py<PyAny>>,
|
||||
}
|
||||
|
||||
impl PythonSecretManager {
|
||||
pub(crate) fn new(
|
||||
client: Py<PyAny>,
|
||||
system: Option<KeyManagementSystem>,
|
||||
settings: Option<Py<PyAny>>,
|
||||
) -> Self {
|
||||
Self {
|
||||
client,
|
||||
system,
|
||||
key_manager: system.map_or("local", python_name),
|
||||
settings,
|
||||
}
|
||||
}
|
||||
|
||||
fn read(&self, py: Python<'_>, name: &str) -> PyResult<Option<String>> {
|
||||
let client = self.client.bind(py);
|
||||
if self.system == Some(KeyManagementSystem::Custom)
|
||||
|| (self.system.is_none() && client.hasattr("sync_read_secret")?)
|
||||
{
|
||||
let kwargs = PyDict::new(py);
|
||||
kwargs.set_item("secret_name", name)?;
|
||||
if self.system == Some(KeyManagementSystem::Custom) {
|
||||
let optional_params = self
|
||||
.settings
|
||||
.as_ref()
|
||||
.map(|settings| settings.bind(py).call_method0("model_dump"))
|
||||
.transpose()?;
|
||||
kwargs.set_item("optional_params", optional_params)?;
|
||||
}
|
||||
return client
|
||||
.call_method("sync_read_secret", (), Some(&kwargs))?
|
||||
.extract();
|
||||
}
|
||||
let kwargs = PyDict::new(py);
|
||||
kwargs.set_item("client", client)?;
|
||||
kwargs.set_item("key_manager", self.key_manager)?;
|
||||
kwargs.set_item("secret_name", name)?;
|
||||
kwargs.set_item(
|
||||
"key_management_settings",
|
||||
self.settings
|
||||
.as_ref()
|
||||
.map_or_else(|| py.None(), |settings| settings.clone_ref(py)),
|
||||
)?;
|
||||
py.import(HANDLER_MODULE)?
|
||||
.getattr("get_secret_from_manager")?
|
||||
.call((), Some(&kwargs))?
|
||||
.extract()
|
||||
}
|
||||
}
|
||||
|
||||
/// The `KeyManagementSystem` value as Python spells it.
|
||||
fn python_name(system: KeyManagementSystem) -> &'static str {
|
||||
match system {
|
||||
KeyManagementSystem::GoogleKms => "google_kms",
|
||||
KeyManagementSystem::AzureKeyVault => "azure_key_vault",
|
||||
KeyManagementSystem::AwsSecretManager => "aws_secret_manager",
|
||||
KeyManagementSystem::GoogleSecretManager => "google_secret_manager",
|
||||
KeyManagementSystem::HashicorpVault => "hashicorp_vault",
|
||||
KeyManagementSystem::Cyberark => "cyberark",
|
||||
KeyManagementSystem::Local => "local",
|
||||
KeyManagementSystem::AwsKms => "aws_kms",
|
||||
KeyManagementSystem::Custom => "custom",
|
||||
}
|
||||
}
|
||||
|
||||
impl ExternalSecretManager for PythonSecretManager {
|
||||
fn system(&self) -> KeyManagementSystem {
|
||||
self.system.unwrap_or(KeyManagementSystem::Custom)
|
||||
}
|
||||
|
||||
fn read_secret<'a>(
|
||||
&'a self,
|
||||
name: &'a str,
|
||||
_settings: &'a KeyManagementSettings,
|
||||
_environment: &'a (dyn Lookup + Send + Sync),
|
||||
) -> Pin<Box<dyn Future<Output = Result<Option<Secret>, Error>> + Send + 'a>> {
|
||||
Box::pin(async move {
|
||||
Python::attach(|py| {
|
||||
self.read(py, name)
|
||||
.map(|value| value.map(SecretValue::new).map(Secret::String))
|
||||
.map_err(|error| {
|
||||
Error::ExternalManager(Box::new(PythonSecretError(error.into_value(py))))
|
||||
})
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use litellm_secrets::{
|
||||
FailurePolicy, KeyManagementSettings, KeyManagementSystem, OidcResolver, SecretManager,
|
||||
SecretManagerState, SecretResolver,
|
||||
};
|
||||
use pyo3::{prelude::*, types::PyDict};
|
||||
|
||||
use super::{HANDLER_MODULE, PythonSecretManager, python_error, python_name};
|
||||
|
||||
#[tokio::test]
|
||||
async fn callback_failures_preserve_python_exceptions_even_with_environment_fallback() {
|
||||
Python::initialize();
|
||||
for failure_type in ["ValueError", "asyncio.CancelledError"] {
|
||||
for fallback in [None, Some("environment-key")] {
|
||||
let (reader, locals) = Python::attach(|py| {
|
||||
let locals = PyDict::new(py);
|
||||
locals.set_item("failure_type", failure_type).unwrap();
|
||||
py.run(
|
||||
c"
|
||||
import asyncio
|
||||
failure = eval(failure_type)('secret manager failed')
|
||||
cause = RuntimeError('original cause')
|
||||
context = RuntimeError('original context')
|
||||
failure.__cause__ = cause
|
||||
failure.__context__ = context
|
||||
class Manager:
|
||||
def sync_read_secret(self, secret_name):
|
||||
raise failure
|
||||
manager = Manager()
|
||||
",
|
||||
Some(&locals),
|
||||
Some(&locals),
|
||||
)
|
||||
.unwrap();
|
||||
let reader = PythonSecretManager::new(
|
||||
locals.get_item("manager").unwrap().unwrap().unbind(),
|
||||
None,
|
||||
None,
|
||||
);
|
||||
(reader, locals.unbind())
|
||||
});
|
||||
let resolver = SecretResolver::new(
|
||||
Arc::new(SecretManagerState::new(
|
||||
SecretManager::External(Arc::new(reader)),
|
||||
KeyManagementSettings::default(),
|
||||
)),
|
||||
Arc::new(move |_: &str| fallback.map(str::to_owned)),
|
||||
OidcResolver::default(),
|
||||
)
|
||||
.with_failure_policy(FailurePolicy::EnvironmentFallback);
|
||||
let error = resolver.get_secret("API_KEY", None).await.unwrap_err();
|
||||
Python::attach(|py| {
|
||||
let original = python_error(py, &error).unwrap();
|
||||
let locals = locals.bind(py);
|
||||
assert!(
|
||||
original
|
||||
.value(py)
|
||||
.is(locals.get_item("failure").unwrap().unwrap())
|
||||
);
|
||||
for (attribute, name) in [("__cause__", "cause"), ("__context__", "context")] {
|
||||
assert!(
|
||||
original
|
||||
.value(py)
|
||||
.getattr(attribute)
|
||||
.unwrap()
|
||||
.is(locals.get_item(name).unwrap().unwrap())
|
||||
);
|
||||
}
|
||||
assert!(original.traceback(py).is_some());
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Installs a fake `get_secret_from_manager` that records its kwargs, runs `body`, and
|
||||
/// removes the fake modules again.
|
||||
fn with_fake_handler<'py>(py: Python<'py>, body: impl FnOnce(&Bound<'py, PyDict>)) {
|
||||
let locals = PyDict::new(py);
|
||||
py.run(
|
||||
c"
|
||||
import sys, types
|
||||
calls = []
|
||||
def get_secret_from_manager(**kwargs):
|
||||
calls.append(kwargs)
|
||||
return 'handled-' + kwargs['secret_name']
|
||||
handler = types.ModuleType('litellm.secret_managers.secret_manager_handler')
|
||||
handler.get_secret_from_manager = get_secret_from_manager
|
||||
installed = {}
|
||||
for name in ('litellm', 'litellm.secret_managers'):
|
||||
if name not in sys.modules:
|
||||
sys.modules[name] = types.ModuleType(name)
|
||||
installed[name] = True
|
||||
sys.modules['litellm.secret_managers.secret_manager_handler'] = handler
|
||||
",
|
||||
Some(&locals),
|
||||
Some(&locals),
|
||||
)
|
||||
.unwrap();
|
||||
body(&locals);
|
||||
py.run(
|
||||
c"
|
||||
sys.modules.pop('litellm.secret_managers.secret_manager_handler', None)
|
||||
for name in installed:
|
||||
sys.modules.pop(name, None)
|
||||
",
|
||||
Some(&locals),
|
||||
Some(&locals),
|
||||
)
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn python_names_round_trip_through_serde() {
|
||||
for system in [
|
||||
KeyManagementSystem::GoogleKms,
|
||||
KeyManagementSystem::AzureKeyVault,
|
||||
KeyManagementSystem::AwsSecretManager,
|
||||
KeyManagementSystem::GoogleSecretManager,
|
||||
KeyManagementSystem::HashicorpVault,
|
||||
KeyManagementSystem::Cyberark,
|
||||
KeyManagementSystem::Local,
|
||||
KeyManagementSystem::AwsKms,
|
||||
KeyManagementSystem::Custom,
|
||||
] {
|
||||
assert_eq!(
|
||||
serde_json::to_value(system).unwrap(),
|
||||
serde_json::Value::String(python_name(system).to_owned())
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn custom_readers_without_a_system_are_called_directly() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let locals = PyDict::new(py);
|
||||
py.run(
|
||||
c"
|
||||
class Manager:
|
||||
def __init__(self):
|
||||
self.names = []
|
||||
def sync_read_secret(self, secret_name, optional_params=None, timeout=None):
|
||||
self.names.append(secret_name)
|
||||
return 'direct-' + secret_name
|
||||
manager = Manager()
|
||||
",
|
||||
Some(&locals),
|
||||
Some(&locals),
|
||||
)
|
||||
.unwrap();
|
||||
let manager = locals.get_item("manager").unwrap().unwrap();
|
||||
let reader = PythonSecretManager::new(manager.clone().unbind(), None, None);
|
||||
assert_eq!(
|
||||
reader.read(py, "API_KEY").unwrap().as_deref(),
|
||||
Some("direct-API_KEY")
|
||||
);
|
||||
assert_eq!(
|
||||
manager
|
||||
.getattr("names")
|
||||
.unwrap()
|
||||
.extract::<Vec<String>>()
|
||||
.unwrap(),
|
||||
["API_KEY"]
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn configured_systems_dispatch_through_the_python_handler_with_the_original_settings() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
with_fake_handler(py, |locals| {
|
||||
let client = py.eval(c"object()", None, None).unwrap();
|
||||
let settings = py.eval(c"object()", None, None).unwrap();
|
||||
let reader = PythonSecretManager::new(
|
||||
client.clone().unbind(),
|
||||
Some(KeyManagementSystem::AzureKeyVault),
|
||||
Some(settings.clone().unbind()),
|
||||
);
|
||||
assert_eq!(
|
||||
reader.read(py, "API_KEY").unwrap().as_deref(),
|
||||
Some("handled-API_KEY")
|
||||
);
|
||||
assert!(py.import(HANDLER_MODULE).is_ok());
|
||||
let calls = locals.get_item("calls").unwrap().unwrap();
|
||||
let call = calls.get_item(0).unwrap().cast_into::<PyDict>().unwrap();
|
||||
assert!(call.get_item("client").unwrap().unwrap().is(&client));
|
||||
assert!(
|
||||
call.get_item("key_management_settings")
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.is(&settings)
|
||||
);
|
||||
assert_eq!(
|
||||
call.get_item("key_manager")
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.extract::<String>()
|
||||
.unwrap(),
|
||||
"azure_key_vault"
|
||||
);
|
||||
assert_eq!(
|
||||
call.get_item("secret_name")
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.extract::<String>()
|
||||
.unwrap(),
|
||||
"API_KEY"
|
||||
);
|
||||
});
|
||||
});
|
||||
}
|
||||
}
|
||||
338
litellm-rust/crates/python-bridge/src/secrets/config.rs
Normal file
338
litellm-rust/crates/python-bridge/src/secrets/config.rs
Normal file
|
|
@ -0,0 +1,338 @@
|
|||
use std::sync::Arc;
|
||||
|
||||
use litellm_secrets::{SecretManager, SecretManagerState};
|
||||
use litellm_secrets_types::{AccessMode, KeyManagementSettings, KeyManagementSystem, SecretValue};
|
||||
use pyo3::prelude::*;
|
||||
use serde_json::Value;
|
||||
|
||||
use super::callback::PythonSecretManager;
|
||||
use crate::{
|
||||
coercion::{Field, FieldSpec, ProjectionError},
|
||||
python_settings::{PythonSettings, Snapshot},
|
||||
};
|
||||
|
||||
const SYSTEM: FieldSpec<Option<KeyManagementSystem>> =
|
||||
FieldSpec::new("system", parse_optional_system);
|
||||
const ACCESS_MODE: FieldSpec<AccessMode> = FieldSpec::new("access_mode", parse_access_mode);
|
||||
const HOSTED_KEYS: FieldSpec<Option<Vec<String>>> =
|
||||
FieldSpec::new("hosted_keys", |field| field.optional_string_collection());
|
||||
const STORE_VIRTUAL_KEYS: FieldSpec<bool> =
|
||||
FieldSpec::new("store_virtual_keys", |field| field.truthy());
|
||||
const PREFIX_FOR_STORED_VIRTUAL_KEYS: FieldSpec<String> =
|
||||
FieldSpec::new("prefix_for_stored_virtual_keys", |field| {
|
||||
field.strict_string()
|
||||
});
|
||||
const PRIMARY_SECRET_NAME: FieldSpec<Option<String>> =
|
||||
FieldSpec::new("primary_secret_name", |field| field.falsy_optional_string());
|
||||
const KMS_KEY_ID: FieldSpec<Option<String>> =
|
||||
FieldSpec::new("kms_key_id", |field| field.falsy_optional_string());
|
||||
const CUSTOM_SECRET_MANAGER: FieldSpec<Option<String>> =
|
||||
FieldSpec::new("custom_secret_manager", |field| {
|
||||
field.falsy_optional_string()
|
||||
});
|
||||
const AWS_REGION_NAME: FieldSpec<Option<String>> =
|
||||
FieldSpec::new("aws_region_name", |field| field.falsy_optional_string());
|
||||
const AWS_ROLE_NAME: FieldSpec<Option<String>> =
|
||||
FieldSpec::new("aws_role_name", |field| field.falsy_optional_string());
|
||||
const AWS_SESSION_NAME: FieldSpec<Option<String>> =
|
||||
FieldSpec::new("aws_session_name", |field| field.falsy_optional_string());
|
||||
const AWS_EXTERNAL_ID: FieldSpec<Option<String>> =
|
||||
FieldSpec::new("aws_external_id", |field| field.falsy_optional_string());
|
||||
const AWS_PROFILE_NAME: FieldSpec<Option<String>> =
|
||||
FieldSpec::new("aws_profile_name", |field| field.falsy_optional_string());
|
||||
const AWS_WEB_IDENTITY_TOKEN: FieldSpec<Option<String>> =
|
||||
FieldSpec::new("aws_web_identity_token", |field| {
|
||||
field.falsy_optional_string()
|
||||
});
|
||||
const AWS_STS_ENDPOINT: FieldSpec<Option<String>> =
|
||||
FieldSpec::new("aws_sts_endpoint", |field| field.falsy_optional_string());
|
||||
const REPLICA_REGIONS: FieldSpec<Option<Vec<String>>> =
|
||||
FieldSpec::new("replica_regions", |field| {
|
||||
field.optional_string_collection()
|
||||
});
|
||||
const CLIENT: FieldSpec<Option<Py<PyAny>>> =
|
||||
FieldSpec::new("client", |field| Ok(field.python_binding()));
|
||||
const SETTINGS_OBJECT: FieldSpec<Option<Py<PyAny>>> =
|
||||
FieldSpec::new("settings_object", |field| Ok(field.python_binding()));
|
||||
|
||||
/// `litellm.secret_manager_client` as the bridge classifies it.
|
||||
#[derive(Debug)]
|
||||
pub(crate) enum SecretManagerClient {
|
||||
/// `None`: reads come from the process environment.
|
||||
Local,
|
||||
/// A custom manager, legacy compatible client, or manually assigned SDK client that keeps
|
||||
/// executing in Python.
|
||||
PythonCallback(Py<PyAny>),
|
||||
}
|
||||
|
||||
/// One operation-local capture of the secret manager globals, taken while attached to Python.
|
||||
#[derive(Debug)]
|
||||
pub(crate) struct SecretManagerSnapshot {
|
||||
pub(crate) client: SecretManagerClient,
|
||||
pub(crate) system: Option<KeyManagementSystem>,
|
||||
/// Typed settings that drive native routing: access mode and hosted keys.
|
||||
pub(crate) settings: KeyManagementSettings,
|
||||
/// The original `KeyManagementSettings` object, handed back to Python callbacks unchanged.
|
||||
pub(crate) settings_object: Option<Py<PyAny>>,
|
||||
}
|
||||
|
||||
impl SecretManagerSnapshot {
|
||||
pub(crate) fn into_state(self) -> Arc<SecretManagerState> {
|
||||
match self.client {
|
||||
SecretManagerClient::Local => Arc::new(SecretManagerState::default()),
|
||||
SecretManagerClient::PythonCallback(client) => Arc::new(SecretManagerState::new(
|
||||
SecretManager::External(Arc::new(PythonSecretManager::new(
|
||||
client,
|
||||
self.system,
|
||||
self.settings_object,
|
||||
))),
|
||||
self.settings,
|
||||
)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Reads and projects the secret manager settings group in one attached operation.
|
||||
pub(crate) fn read(py: Python<'_>) -> PyResult<SecretManagerSnapshot> {
|
||||
Ok(project(&PythonSettings::SecretManagerBinding.read(py)?)?)
|
||||
}
|
||||
|
||||
pub(crate) fn project(snapshot: &Snapshot<'_>) -> Result<SecretManagerSnapshot, ProjectionError> {
|
||||
let system = snapshot.read(&SYSTEM)?;
|
||||
let access_mode = snapshot.read(&ACCESS_MODE)?;
|
||||
let settings = KeyManagementSettings {
|
||||
hosted_keys: snapshot.read(&HOSTED_KEYS)?,
|
||||
store_virtual_keys: Some(snapshot.read(&STORE_VIRTUAL_KEYS)?),
|
||||
prefix_for_stored_virtual_keys: snapshot.read(&PREFIX_FOR_STORED_VIRTUAL_KEYS)?,
|
||||
access_mode,
|
||||
primary_secret_name: snapshot.read(&PRIMARY_SECRET_NAME)?,
|
||||
kms_key_id: snapshot.read(&KMS_KEY_ID)?,
|
||||
custom_secret_manager: snapshot.read(&CUSTOM_SECRET_MANAGER)?,
|
||||
aws_region_name: snapshot.read(&AWS_REGION_NAME)?,
|
||||
aws_role_name: snapshot.read(&AWS_ROLE_NAME)?,
|
||||
aws_session_name: snapshot.read(&AWS_SESSION_NAME)?,
|
||||
aws_external_id: snapshot.read(&AWS_EXTERNAL_ID)?.map(SecretValue::new),
|
||||
aws_profile_name: snapshot.read(&AWS_PROFILE_NAME)?,
|
||||
aws_web_identity_token: snapshot
|
||||
.read(&AWS_WEB_IDENTITY_TOKEN)?
|
||||
.map(SecretValue::new),
|
||||
aws_sts_endpoint: snapshot.read(&AWS_STS_ENDPOINT)?,
|
||||
replica_regions: snapshot.read(&REPLICA_REGIONS)?,
|
||||
..KeyManagementSettings::default()
|
||||
};
|
||||
let client = match snapshot.read(&CLIENT)? {
|
||||
None => SecretManagerClient::Local,
|
||||
Some(client) => SecretManagerClient::PythonCallback(client),
|
||||
};
|
||||
Ok(SecretManagerSnapshot {
|
||||
client,
|
||||
system,
|
||||
settings,
|
||||
settings_object: snapshot.read(&SETTINGS_OBJECT)?,
|
||||
})
|
||||
}
|
||||
|
||||
fn parse_optional_system(
|
||||
field: &Field<'_>,
|
||||
) -> Result<Option<KeyManagementSystem>, ProjectionError> {
|
||||
let Some(value) = field.falsy_optional_string()? else {
|
||||
return Ok(None);
|
||||
};
|
||||
serde_json::from_value(Value::String(value))
|
||||
.map(Some)
|
||||
.map_err(|error| {
|
||||
ProjectionError::InvalidConfiguration(format!("secret manager system: {error}"))
|
||||
})
|
||||
}
|
||||
|
||||
fn parse_access_mode(field: &Field<'_>) -> Result<AccessMode, ProjectionError> {
|
||||
let value = field.strict_string()?;
|
||||
serde_json::from_value(Value::String(value)).map_err(|error| {
|
||||
ProjectionError::InvalidConfiguration(format!("secret manager access mode: {error}"))
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use pyo3::{
|
||||
prelude::*,
|
||||
types::{PyDict, PyTuple},
|
||||
};
|
||||
|
||||
use super::{SecretManagerClient, project};
|
||||
use crate::python_settings::PythonSettings;
|
||||
|
||||
fn snapshot<'py>(
|
||||
py: Python<'py>,
|
||||
system: &str,
|
||||
access_mode: &str,
|
||||
store_virtual_keys: Bound<'py, PyAny>,
|
||||
hosted_keys: Bound<'py, PyAny>,
|
||||
) -> crate::python_settings::Snapshot<'py> {
|
||||
snapshot_with_client(
|
||||
py,
|
||||
system,
|
||||
access_mode,
|
||||
store_virtual_keys,
|
||||
hosted_keys,
|
||||
py.None().into_bound(py),
|
||||
)
|
||||
}
|
||||
|
||||
fn snapshot_with_client<'py>(
|
||||
py: Python<'py>,
|
||||
system: &str,
|
||||
access_mode: &str,
|
||||
store_virtual_keys: Bound<'py, PyAny>,
|
||||
hosted_keys: Bound<'py, PyAny>,
|
||||
client: Bound<'py, PyAny>,
|
||||
) -> crate::python_settings::Snapshot<'py> {
|
||||
let locals = PyDict::new(py);
|
||||
locals.set_item("client", client).unwrap();
|
||||
locals.set_item("system", system).unwrap();
|
||||
locals.set_item("access_mode", access_mode).unwrap();
|
||||
locals
|
||||
.set_item("store_virtual_keys", store_virtual_keys)
|
||||
.unwrap();
|
||||
locals.set_item("hosted_keys", hosted_keys).unwrap();
|
||||
py.run(
|
||||
cr#"
|
||||
from dataclasses import dataclass
|
||||
from types import SimpleNamespace
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class SecretManager:
|
||||
system: object
|
||||
access_mode: object
|
||||
hosted_keys: object
|
||||
primary_secret_name: object
|
||||
store_virtual_keys: object
|
||||
prefix_for_stored_virtual_keys: object
|
||||
kms_key_id: object
|
||||
custom_secret_manager: object
|
||||
aws_region_name: object
|
||||
aws_role_name: object
|
||||
aws_session_name: object
|
||||
aws_external_id: object
|
||||
aws_profile_name: object
|
||||
aws_web_identity_token: object
|
||||
aws_sts_endpoint: object
|
||||
replica_regions: object
|
||||
client: object
|
||||
settings_object: object
|
||||
|
||||
root = SimpleNamespace(secret_manager=SecretManager(
|
||||
system=system,
|
||||
access_mode=access_mode,
|
||||
hosted_keys=hosted_keys,
|
||||
primary_secret_name=None,
|
||||
store_virtual_keys=store_virtual_keys,
|
||||
prefix_for_stored_virtual_keys="litellm/",
|
||||
kms_key_id=None,
|
||||
custom_secret_manager=None,
|
||||
aws_region_name=None,
|
||||
aws_role_name=None,
|
||||
aws_session_name=None,
|
||||
aws_external_id=None,
|
||||
aws_profile_name=None,
|
||||
aws_web_identity_token=None,
|
||||
aws_sts_endpoint=None,
|
||||
replica_regions=None,
|
||||
client=client,
|
||||
settings_object=None,
|
||||
))
|
||||
"#,
|
||||
Some(&locals),
|
||||
Some(&locals),
|
||||
)
|
||||
.unwrap();
|
||||
PythonSettings::SecretManagerBinding.snapshot(
|
||||
locals
|
||||
.get_item("root")
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.getattr("secret_manager")
|
||||
.unwrap(),
|
||||
)
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::string_true(Some("true"), false, true)]
|
||||
#[case::string_one(Some("1"), false, true)]
|
||||
#[case::true_value(None, true, true)]
|
||||
#[case::false_value(None, false, false)]
|
||||
#[case::string_false(Some("false"), false, true)]
|
||||
fn python_compatible_boolean_coercion(
|
||||
#[case] string_value: Option<&str>,
|
||||
#[case] bool_value: bool,
|
||||
#[case] expected: bool,
|
||||
) {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let store_virtual_keys = match string_value {
|
||||
Some(value) => value.into_pyobject(py).unwrap().into_any(),
|
||||
None => bool_value.into_pyobject(py).unwrap().to_owned().into_any(),
|
||||
};
|
||||
let hosted_keys = PyTuple::new(py, ["ONE"]).unwrap().into_any();
|
||||
let projected = project(&snapshot(
|
||||
py,
|
||||
"local",
|
||||
"read_only",
|
||||
store_virtual_keys,
|
||||
hosted_keys,
|
||||
))
|
||||
.unwrap();
|
||||
assert_eq!(projected.settings.store_virtual_keys, Some(expected));
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unknown_system_is_rejected() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let error = project(&snapshot(
|
||||
py,
|
||||
"unknown",
|
||||
"read_only",
|
||||
false.into_pyobject(py).unwrap().to_owned().into_any(),
|
||||
PyTuple::empty(py).into_any(),
|
||||
))
|
||||
.unwrap_err();
|
||||
let error: PyErr = error.into();
|
||||
assert!(error.is_instance_of::<pyo3::exceptions::PyValueError>(py));
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn client_identity_selects_local_or_python_callback() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let falsy = false.into_pyobject(py).unwrap().to_owned().into_any();
|
||||
let local = project(&snapshot(
|
||||
py,
|
||||
"local",
|
||||
"read_only",
|
||||
falsy.clone(),
|
||||
PyTuple::empty(py).into_any(),
|
||||
))
|
||||
.unwrap();
|
||||
assert!(matches!(local.client, SecretManagerClient::Local));
|
||||
assert!(local.settings_object.is_none());
|
||||
|
||||
let manager = py.eval(c"object()", None, None).unwrap();
|
||||
let custom = project(&snapshot_with_client(
|
||||
py,
|
||||
"custom",
|
||||
"read_only",
|
||||
falsy,
|
||||
PyTuple::empty(py).into_any(),
|
||||
manager.clone(),
|
||||
))
|
||||
.unwrap();
|
||||
let SecretManagerClient::PythonCallback(client) = custom.client else {
|
||||
panic!("a live client must stay a Python callback");
|
||||
};
|
||||
assert!(client.bind(py).is(&manager));
|
||||
});
|
||||
}
|
||||
}
|
||||
3
litellm-rust/crates/python-bridge/src/secrets/mod.rs
Normal file
3
litellm-rust/crates/python-bridge/src/secrets/mod.rs
Normal file
|
|
@ -0,0 +1,3 @@
|
|||
pub(crate) mod callback;
|
||||
pub(crate) mod config;
|
||||
pub(crate) mod resolved;
|
||||
252
litellm-rust/crates/python-bridge/src/secrets/resolved.rs
Normal file
252
litellm-rust/crates/python-bridge/src/secrets/resolved.rs
Normal file
|
|
@ -0,0 +1,252 @@
|
|||
use std::{collections::HashMap, sync::Arc};
|
||||
|
||||
use futures_util::{future::BoxFuture, future::try_join_all};
|
||||
use litellm_core_utils::settings::{Lookup, ProcessEnvironment};
|
||||
use litellm_llms::base_llm::inference::secrets::{SecretSource, Secrets};
|
||||
use litellm_secrets::{
|
||||
Error, FailurePolicy, OidcResolver, Secret, SecretManagerState, SecretResolver,
|
||||
};
|
||||
|
||||
use super::config::SecretManagerSnapshot;
|
||||
|
||||
pub(crate) struct ResolvedSecrets {
|
||||
resolver: SecretResolver,
|
||||
}
|
||||
|
||||
impl ResolvedSecrets {
|
||||
pub(crate) fn new(snapshot: SecretManagerSnapshot) -> Self {
|
||||
Self::from_state(snapshot.into_state())
|
||||
}
|
||||
|
||||
fn from_state(state: Arc<SecretManagerState>) -> Self {
|
||||
Self {
|
||||
resolver: SecretResolver::new(
|
||||
state,
|
||||
Arc::new(ProcessEnvironment),
|
||||
OidcResolver::default(),
|
||||
)
|
||||
.with_failure_policy(FailurePolicy::EnvironmentFallback),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl SecretSource for ResolvedSecrets {
|
||||
fn resolve<'a>(&'a self, names: &'a [&'static str]) -> BoxFuture<'a, Result<Secrets, Error>> {
|
||||
Box::pin(async move {
|
||||
let values = try_join_all(names.iter().map(|name| async move {
|
||||
self.resolver
|
||||
.get_secret(name, None)
|
||||
.await
|
||||
.map(|secret| secret.map(|secret| ((*name).to_owned(), secret_value(secret))))
|
||||
}))
|
||||
.await?
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.collect::<HashMap<_, _>>();
|
||||
Ok(Arc::new(ResolvedLookup { values }) as Secrets)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
struct ResolvedLookup {
|
||||
values: HashMap<String, String>,
|
||||
}
|
||||
|
||||
impl Lookup for ResolvedLookup {
|
||||
fn get(&self, name: &str) -> Option<String> {
|
||||
self.values
|
||||
.get(name)
|
||||
.cloned()
|
||||
.or_else(|| ProcessEnvironment.get(name))
|
||||
}
|
||||
}
|
||||
|
||||
fn secret_value(secret: Secret) -> String {
|
||||
match secret {
|
||||
Secret::String(value) => value.expose().to_owned(),
|
||||
Secret::Bool(value) => if value { "True" } else { "False" }.to_owned(),
|
||||
Secret::Json(value) => value.to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use aws_sdk_secretsmanager::Client;
|
||||
use aws_sdk_secretsmanager::config::{
|
||||
BehaviorVersion, Credentials, Region, retry::RetryConfig,
|
||||
};
|
||||
use litellm_secrets::{AccessMode, KeyManagementSettings, SecretManager, SecretManagerState};
|
||||
use litellm_secrets_aws::AwsSecretsManagerV2;
|
||||
use serde_json::json;
|
||||
use wiremock::{
|
||||
Mock, MockServer, ResponseTemplate,
|
||||
matchers::{body_partial_json, header},
|
||||
};
|
||||
|
||||
use super::ResolvedSecrets;
|
||||
use litellm_llms::base_llm::inference::secrets::SecretSource;
|
||||
|
||||
fn state(server: &MockServer, settings: KeyManagementSettings) -> Arc<SecretManagerState> {
|
||||
let client = Client::from_conf(
|
||||
aws_sdk_secretsmanager::Config::builder()
|
||||
.behavior_version(BehaviorVersion::latest())
|
||||
.region(Region::new("us-east-1"))
|
||||
.credentials_provider(Credentials::new("test", "test", None, None, "test"))
|
||||
.endpoint_url(server.uri())
|
||||
.retry_config(RetryConfig::disabled())
|
||||
.build(),
|
||||
);
|
||||
Arc::new(SecretManagerState::new(
|
||||
SecretManager::AwsSecretsManagerV2(AwsSecretsManagerV2::new(
|
||||
client,
|
||||
(&settings).into(),
|
||||
)),
|
||||
settings,
|
||||
))
|
||||
}
|
||||
|
||||
async fn resolve(state: Arc<SecretManagerState>, name: &'static str) -> Option<String> {
|
||||
ResolvedSecrets::from_state(state)
|
||||
.resolve(&[name])
|
||||
.await
|
||||
.unwrap()
|
||||
.get(name)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn hosted_key_miss_falls_back_to_environment() {
|
||||
let name = "LITELLM_RUST_BRIDGE_HOSTED_KEY_MISS";
|
||||
unsafe { std::env::set_var(name, "env-key") };
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(header("x-amz-target", "secretsmanager.GetSecretValue"))
|
||||
.and(body_partial_json(json!({"SecretId": name})))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200).set_body_json(json!({"SecretString": "manager-key"})),
|
||||
)
|
||||
.expect(0)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let result = resolve(
|
||||
state(
|
||||
&server,
|
||||
KeyManagementSettings {
|
||||
hosted_keys: Some(vec!["OTHER".into()]),
|
||||
..Default::default()
|
||||
},
|
||||
),
|
||||
name,
|
||||
)
|
||||
.await;
|
||||
unsafe { std::env::remove_var(name) };
|
||||
assert_eq!(result.as_deref(), Some("env-key"));
|
||||
assert_eq!(server.received_requests().await.unwrap().len(), 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn manager_failure_falls_back_to_environment() {
|
||||
let name = "LITELLM_RUST_BRIDGE_MANAGER_FAILURE";
|
||||
unsafe { std::env::set_var(name, "env-key") };
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(header("x-amz-target", "secretsmanager.GetSecretValue"))
|
||||
.respond_with(ResponseTemplate::new(500))
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let result = resolve(state(&server, KeyManagementSettings::default()), name).await;
|
||||
unsafe { std::env::remove_var(name) };
|
||||
assert_eq!(result.as_deref(), Some("env-key"));
|
||||
assert_eq!(server.received_requests().await.unwrap().len(), 1);
|
||||
|
||||
let missing_server = MockServer::start().await;
|
||||
Mock::given(header("x-amz-target", "secretsmanager.GetSecretValue"))
|
||||
.respond_with(ResponseTemplate::new(500))
|
||||
.expect(1)
|
||||
.mount(&missing_server)
|
||||
.await;
|
||||
let missing =
|
||||
ResolvedSecrets::from_state(state(&missing_server, KeyManagementSettings::default()))
|
||||
.resolve(&["LITELLM_RUST_BRIDGE_MANAGER_FAILURE_MISSING"])
|
||||
.await;
|
||||
assert!(matches!(missing, Err(litellm_secrets::Error::Aws(_))));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn write_only_mode_never_consults_the_manager() {
|
||||
let name = "LITELLM_RUST_BRIDGE_WRITE_ONLY";
|
||||
unsafe { std::env::set_var(name, "env-key") };
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(header("x-amz-target", "secretsmanager.GetSecretValue"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200).set_body_json(json!({"SecretString": "manager-key"})),
|
||||
)
|
||||
.expect(0)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let result = resolve(
|
||||
state(
|
||||
&server,
|
||||
KeyManagementSettings {
|
||||
access_mode: AccessMode::WriteOnly,
|
||||
..Default::default()
|
||||
},
|
||||
),
|
||||
name,
|
||||
)
|
||||
.await;
|
||||
unsafe { std::env::remove_var(name) };
|
||||
assert_eq!(result.as_deref(), Some("env-key"));
|
||||
assert_eq!(server.received_requests().await.unwrap().len(), 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn read_only_mode_resolves_from_the_manager() {
|
||||
let name = "LITELLM_RUST_BRIDGE_READ_ONLY";
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(header("x-amz-target", "secretsmanager.GetSecretValue"))
|
||||
.and(body_partial_json(json!({"SecretId": name})))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200).set_body_json(json!({"SecretString": "manager-key"})),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
assert_eq!(
|
||||
resolve(state(&server, KeyManagementSettings::default()), name)
|
||||
.await
|
||||
.as_deref(),
|
||||
Some("manager-key")
|
||||
);
|
||||
assert_eq!(server.received_requests().await.unwrap().len(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn oidc_failures_are_not_converted_to_missing_secrets() {
|
||||
let result = ResolvedSecrets::from_state(Arc::new(SecretManagerState::default()))
|
||||
.resolve(&["oidc/"])
|
||||
.await;
|
||||
assert!(matches!(result, Err(litellm_secrets::Error::InvalidOidc)));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn undeclared_names_still_read_the_process_environment() {
|
||||
let name = "LITELLM_RUST_BRIDGE_UNDECLARED";
|
||||
unsafe { std::env::set_var(name, "env-key") };
|
||||
let server = MockServer::start().await;
|
||||
let result = resolve(
|
||||
state(
|
||||
&server,
|
||||
KeyManagementSettings {
|
||||
hosted_keys: Some(vec!["OTHER".into()]),
|
||||
..Default::default()
|
||||
},
|
||||
),
|
||||
name,
|
||||
)
|
||||
.await;
|
||||
unsafe { std::env::remove_var(name) };
|
||||
assert_eq!(result.as_deref(), Some("env-key"));
|
||||
assert_eq!(server.received_requests().await.unwrap().len(), 0);
|
||||
}
|
||||
}
|
||||
|
|
@ -1,158 +0,0 @@
|
|||
use std::sync::Arc;
|
||||
|
||||
#[cfg(any(feature = "fast", feature = "huggingface", feature = "tiktoken"))]
|
||||
use std::{num::NonZero, thread::available_parallelism};
|
||||
|
||||
#[cfg(any(feature = "fast", feature = "huggingface", feature = "tiktoken"))]
|
||||
use litellm_host_python::release_gil;
|
||||
use litellm_host_python::run_async;
|
||||
use litellm_token_counter::{
|
||||
CountableRequest, Error, InputTokenCount, TokenCounter as CoreTokenCounter,
|
||||
};
|
||||
use pyo3::{
|
||||
exceptions::{PyRuntimeError, PyValueError},
|
||||
prelude::*,
|
||||
types::PyAny,
|
||||
};
|
||||
use tokio::sync::Semaphore;
|
||||
|
||||
use crate::errors::RustBridgeDeclined;
|
||||
|
||||
/// Counts the input tokens of a raw request body off the Python event loop with
|
||||
/// the GIL released. Python owns which requests get here and what to do with
|
||||
/// the count. At most one encode per core runs at a time; the rest wait in the
|
||||
/// async task, where a cancelled Python awaiter drops them before any blocking
|
||||
/// work is scheduled.
|
||||
#[pyclass(frozen)]
|
||||
pub(crate) struct TokenCounter {
|
||||
inner: Arc<CoreTokenCounter>,
|
||||
encode_slots: Arc<Semaphore>,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl TokenCounter {
|
||||
#[new]
|
||||
fn new(py: Python<'_>, tokenizer_json: &str) -> PyResult<Self> {
|
||||
#[cfg(feature = "fast")]
|
||||
{
|
||||
Self::load(py, || CoreTokenCounter::from_json_fast(tokenizer_json))
|
||||
}
|
||||
#[cfg(all(not(feature = "fast"), feature = "huggingface"))]
|
||||
{
|
||||
Self::load(py, || CoreTokenCounter::from_json(tokenizer_json))
|
||||
}
|
||||
#[cfg(not(any(feature = "fast", feature = "huggingface")))]
|
||||
{
|
||||
let _ = (py, tokenizer_json);
|
||||
Err(RustBridgeDeclined::new_err(
|
||||
"tokenizer backend requires the fast or huggingface feature",
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
fn from_cl100k_ranks(py: Python<'_>, rank_file: &str) -> PyResult<Self> {
|
||||
#[cfg(feature = "fast")]
|
||||
{
|
||||
Self::load(py, || CoreTokenCounter::from_cl100k_ranks(rank_file))
|
||||
}
|
||||
#[cfg(not(feature = "fast"))]
|
||||
{
|
||||
let _ = (py, rank_file);
|
||||
Err(RustBridgeDeclined::new_err(
|
||||
"tokenizer backend requires the fast feature",
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
fn from_o200k_ranks(py: Python<'_>, rank_file: &str) -> PyResult<Self> {
|
||||
#[cfg(feature = "fast")]
|
||||
{
|
||||
Self::load(py, || CoreTokenCounter::from_o200k_ranks(rank_file))
|
||||
}
|
||||
#[cfg(not(feature = "fast"))]
|
||||
{
|
||||
let _ = (py, rank_file);
|
||||
Err(RustBridgeDeclined::new_err(
|
||||
"tokenizer backend requires the fast feature",
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
fn from_tiktoken(py: Python<'_>, encoding: &str) -> PyResult<Self> {
|
||||
#[cfg(feature = "tiktoken")]
|
||||
{
|
||||
Self::load(py, || CoreTokenCounter::from_tiktoken(encoding))
|
||||
}
|
||||
#[cfg(not(feature = "tiktoken"))]
|
||||
{
|
||||
let _ = (py, encoding);
|
||||
Err(RustBridgeDeclined::new_err(
|
||||
"tokenizer backend requires the tiktoken feature",
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
fn acount_request<'py>(&self, py: Python<'py>, body: &[u8]) -> PyResult<Bound<'py, PyAny>> {
|
||||
let counter = Arc::clone(&self.inner);
|
||||
let encode_slots = Arc::clone(&self.encode_slots);
|
||||
let body = body.to_vec();
|
||||
run_async(
|
||||
py,
|
||||
async move {
|
||||
let _slot = encode_slots
|
||||
.acquire_owned()
|
||||
.await
|
||||
.map_err(|error| Error::Task(error.to_string()))?;
|
||||
tokio::task::spawn_blocking(move || count_body(&counter, &body))
|
||||
.await
|
||||
.map_err(|error| Error::Task(error.to_string()))?
|
||||
},
|
||||
token_count_error_to_pyerr,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
impl TokenCounter {
|
||||
#[cfg(any(feature = "fast", feature = "huggingface", feature = "tiktoken"))]
|
||||
fn load(
|
||||
py: Python<'_>,
|
||||
load: impl FnOnce() -> Result<CoreTokenCounter, Error> + Send,
|
||||
) -> PyResult<Self> {
|
||||
let inner = release_gil(py, load).map_err(token_count_error_to_pyerr)?;
|
||||
Ok(Self {
|
||||
inner: Arc::new(inner),
|
||||
encode_slots: Arc::new(Semaphore::new(encode_parallelism())),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(any(feature = "fast", feature = "huggingface", feature = "tiktoken"))]
|
||||
fn encode_parallelism() -> usize {
|
||||
available_parallelism().map_or(1, NonZero::get)
|
||||
}
|
||||
|
||||
fn count_body(counter: &CoreTokenCounter, body: &[u8]) -> Result<InputTokenCount, Error> {
|
||||
let request = CountableRequest::parse(body)?;
|
||||
counter.count_request(&request)
|
||||
}
|
||||
|
||||
fn token_count_error_to_pyerr(error: Error) -> PyErr {
|
||||
let message = error.to_string();
|
||||
match error {
|
||||
Error::Load(_)
|
||||
| Error::Ranks(_)
|
||||
| Error::UnicodeClasses
|
||||
| Error::UnsupportedTokenizer(_) => PyValueError::new_err(message),
|
||||
Error::RequestParse(_)
|
||||
| Error::MissingInput
|
||||
| Error::FloatText
|
||||
| Error::ContentBlock
|
||||
| Error::ArrayItems
|
||||
| Error::JsonSerialization(_)
|
||||
| Error::JsonUtf8(_) => RustBridgeDeclined::new_err(message),
|
||||
Error::Encode(_) | Error::Task(_) => PyRuntimeError::new_err(message),
|
||||
}
|
||||
}
|
||||
713
litellm-rust/crates/python-bridge/src/tokenizer.rs
Normal file
713
litellm-rust/crates/python-bridge/src/tokenizer.rs
Normal file
|
|
@ -0,0 +1,713 @@
|
|||
//! The Python face of the text codecs: one `Tokenizer` class over the tiktoken and Hugging
|
||||
//! Face backends, carrying the read-only surface of `tiktoken.Encoding` and
|
||||
//! `tokenizers.Tokenizer` that `litellm/litellm_core_utils/tokenizer.py` wraps.
|
||||
use std::borrow::Cow;
|
||||
#[cfg(any(feature = "tiktoken", feature = "huggingface"))]
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
#[cfg(feature = "fast")]
|
||||
use std::sync::OnceLock;
|
||||
|
||||
use litellm_host_python::{enter_native, release_gil};
|
||||
#[cfg(feature = "fast")]
|
||||
use litellm_token_counter::fast::{FastCounter, FastTokenizer};
|
||||
use litellm_token_counter::{Error, TextCodec};
|
||||
use pyo3::{exceptions::PyUnicodeEncodeError, prelude::*, types::PyString};
|
||||
|
||||
#[cfg(any(feature = "tiktoken", feature = "huggingface"))]
|
||||
use pyo3::exceptions::PyValueError;
|
||||
#[cfg(feature = "huggingface")]
|
||||
use pyo3::{exceptions::PyIOError, types::PyDict};
|
||||
#[cfg(feature = "tiktoken")]
|
||||
use pyo3::{
|
||||
exceptions::{PyKeyError, PyRuntimeError},
|
||||
types::PyBytes,
|
||||
};
|
||||
|
||||
#[cfg(not(all(feature = "tiktoken", feature = "huggingface")))]
|
||||
use crate::errors::RustBridgeDeclined;
|
||||
use crate::routes::token_counter::token_count_error_to_pyerr;
|
||||
|
||||
#[cfg(feature = "huggingface")]
|
||||
use litellm_token_counter::huggingface::{
|
||||
EncodeInput, Encoding, HuggingFaceTokenizer, InputSequence, PaddingDirection, PaddingStrategy,
|
||||
TruncationDirection, encoding_from_json, encoding_to_json,
|
||||
};
|
||||
#[cfg(feature = "tiktoken")]
|
||||
use litellm_token_counter::tiktoken::{TiktokenTokenizer, Vocabulary};
|
||||
|
||||
#[cfg(feature = "tiktoken")]
|
||||
pub(crate) fn load_tiktoken(py: Python<'_>, encoding: &str) -> PyResult<TiktokenTokenizer> {
|
||||
enter_native()?;
|
||||
let resource: std::path::PathBuf =
|
||||
PyModule::import(py, "litellm.litellm_core_utils.tokenizers")?
|
||||
.getattr("__file__")?
|
||||
.extract()?;
|
||||
release_gil(py, || {
|
||||
TiktokenTokenizer::from_cached_ranks(encoding, |file| {
|
||||
std::fs::read_to_string(resource.with_file_name(file))
|
||||
})
|
||||
})
|
||||
.map_err(|error| token_count_error_to_pyerr(error.into()))
|
||||
}
|
||||
|
||||
pub(crate) enum Codec {
|
||||
#[cfg(feature = "tiktoken")]
|
||||
Tiktoken(TiktokenTokenizer),
|
||||
#[cfg(feature = "huggingface")]
|
||||
HuggingFace(HuggingFaceTokenizer),
|
||||
}
|
||||
|
||||
impl Codec {
|
||||
pub(crate) fn codec(&self) -> &dyn TextCodec {
|
||||
match *self {
|
||||
#[cfg(feature = "tiktoken")]
|
||||
Self::Tiktoken(ref tokenizer) => tokenizer,
|
||||
#[cfg(feature = "huggingface")]
|
||||
Self::HuggingFace(ref tokenizer) => tokenizer,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "fast")]
|
||||
fn fast_counter(&self) -> Option<FastTokenizer> {
|
||||
match *self {
|
||||
#[cfg(feature = "tiktoken")]
|
||||
Self::Tiktoken(ref tokenizer) => tokenizer.fast_counter(),
|
||||
#[cfg(feature = "huggingface")]
|
||||
Self::HuggingFace(ref tokenizer) => tokenizer.fast_counter(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The loaded model is shared: `TokenCounter::from_tokenizer` counts with the same parse,
|
||||
/// and the opt-in count-only counter is derived from it once, on first use.
|
||||
#[pyclass(frozen, module = "litellm.rust_bridge._native")]
|
||||
pub(crate) struct Tokenizer {
|
||||
inner: Arc<Codec>,
|
||||
#[cfg(feature = "fast")]
|
||||
fast: OnceLock<Option<Arc<FastTokenizer>>>,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl Tokenizer {
|
||||
#[staticmethod]
|
||||
fn from_tiktoken(py: Python<'_>, encoding: &str) -> PyResult<Self> {
|
||||
#[cfg(feature = "tiktoken")]
|
||||
{
|
||||
let tokenizer = load_tiktoken(py, encoding)?;
|
||||
Ok(Self::new(Codec::Tiktoken(tokenizer)))
|
||||
}
|
||||
#[cfg(not(feature = "tiktoken"))]
|
||||
{
|
||||
let _ = (py, encoding);
|
||||
Err(RustBridgeDeclined::new_err(
|
||||
"tokenizer backend requires the tiktoken feature",
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
fn from_json(py: Python<'_>, tokenizer_json: &str) -> PyResult<Self> {
|
||||
#[cfg(feature = "huggingface")]
|
||||
{
|
||||
enter_native()?;
|
||||
let tokenizer = release_gil(py, || HuggingFaceTokenizer::from_json(tokenizer_json))
|
||||
.map_err(|error| token_count_error_to_pyerr(error.into()))?;
|
||||
Ok(Self::new(Codec::HuggingFace(tokenizer)))
|
||||
}
|
||||
#[cfg(not(feature = "huggingface"))]
|
||||
{
|
||||
let _ = (py, tokenizer_json);
|
||||
Err(RustBridgeDeclined::new_err(
|
||||
"tokenizer backend requires the huggingface feature",
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
#[pyo3(signature = (identifier, revision = "main", token = None))]
|
||||
fn from_pretrained(
|
||||
py: Python<'_>,
|
||||
identifier: &str,
|
||||
revision: &str,
|
||||
token: Option<&str>,
|
||||
) -> PyResult<Self> {
|
||||
#[cfg(feature = "huggingface")]
|
||||
{
|
||||
enter_native()?;
|
||||
let kwargs = PyDict::new(py);
|
||||
kwargs.set_item("repo_id", identifier)?;
|
||||
kwargs.set_item("filename", "tokenizer.json")?;
|
||||
kwargs.set_item("revision", revision)?;
|
||||
kwargs.set_item("token", token)?;
|
||||
let path: String = PyModule::import(py, "huggingface_hub")?
|
||||
.getattr("hf_hub_download")?
|
||||
.call((), Some(&kwargs))?
|
||||
.extract()?;
|
||||
let json =
|
||||
release_gil(py, || std::fs::read_to_string(path)).map_err(PyIOError::new_err)?;
|
||||
Self::from_json(py, &json)
|
||||
}
|
||||
#[cfg(not(feature = "huggingface"))]
|
||||
{
|
||||
let _ = (py, identifier, revision, token);
|
||||
Err(RustBridgeDeclined::new_err(
|
||||
"tokenizer backend requires the huggingface feature",
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
fn encode(&self, py: Python<'_>, text: &Bound<'_, PyString>) -> PyResult<Vec<u32>> {
|
||||
enter_native()?;
|
||||
let text = self.text(text)?;
|
||||
release_gil(py, || self.inner.codec().encode(&text)).map_err(token_count_error_to_pyerr)
|
||||
}
|
||||
|
||||
#[pyo3(signature = (ids, skip_special_tokens = true))]
|
||||
fn decode(&self, py: Python<'_>, ids: Vec<u32>, skip_special_tokens: bool) -> PyResult<String> {
|
||||
enter_native()?;
|
||||
release_gil(py, || self.inner.codec().decode(&ids, skip_special_tokens))
|
||||
.map_err(token_count_error_to_pyerr)
|
||||
}
|
||||
|
||||
#[pyo3(signature = (text, fast = false))]
|
||||
fn count(&self, py: Python<'_>, text: &Bound<'_, PyString>, fast: bool) -> PyResult<usize> {
|
||||
enter_native()?;
|
||||
let text = self.text(text)?;
|
||||
let counter = self.counter(py, fast);
|
||||
release_gil(py, || {
|
||||
litellm_token_counter::Tokenizer::count_tokens(&counter, &text)
|
||||
})
|
||||
.map_err(token_count_error_to_pyerr)
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn name(&self) -> &str {
|
||||
self.inner.codec().name()
|
||||
}
|
||||
|
||||
// ---- tiktoken: the `tiktoken.Encoding` surface ------------------------------------------
|
||||
|
||||
#[cfg(feature = "tiktoken")]
|
||||
fn encode_special(
|
||||
&self,
|
||||
py: Python<'_>,
|
||||
text: &Bound<'_, PyString>,
|
||||
allowed: Vec<String>,
|
||||
) -> PyResult<Vec<u32>> {
|
||||
enter_native()?;
|
||||
let tokenizer = self.tiktoken()?;
|
||||
let text = self.text(text)?;
|
||||
release_gil(py, || tokenizer.encode_special(&text, &allowed))
|
||||
.map_err(PyRuntimeError::new_err)
|
||||
}
|
||||
|
||||
/// tiktoken's `encode_with_unstable`: `(stable_tokens, completions)`.
|
||||
#[cfg(feature = "tiktoken")]
|
||||
fn encode_with_unstable(
|
||||
&self,
|
||||
py: Python<'_>,
|
||||
text: &Bound<'_, PyString>,
|
||||
allowed: Vec<String>,
|
||||
) -> PyResult<(Vec<u32>, Vec<Vec<u32>>)> {
|
||||
enter_native()?;
|
||||
let tokenizer = self.tiktoken()?;
|
||||
let text = self.text(text)?;
|
||||
Ok(release_gil(py, || {
|
||||
tokenizer.encode_with_unstable(&text, &allowed)
|
||||
}))
|
||||
}
|
||||
|
||||
/// The special tokens by text: tiktoken's `_special_tokens`.
|
||||
#[cfg(feature = "tiktoken")]
|
||||
fn special_tokens(&self) -> PyResult<HashMap<String, u32>> {
|
||||
Ok(self
|
||||
.vocabulary()?
|
||||
.special_tokens()
|
||||
.map(|(token, rank)| (token.to_owned(), rank))
|
||||
.collect())
|
||||
}
|
||||
|
||||
#[cfg(feature = "tiktoken")]
|
||||
fn max_token_value(&self) -> PyResult<u32> {
|
||||
Ok(self.vocabulary()?.max_token_value())
|
||||
}
|
||||
|
||||
#[cfg(feature = "tiktoken")]
|
||||
fn is_special_token(&self, token: u32) -> PyResult<bool> {
|
||||
Ok(self.vocabulary()?.is_special_token(token))
|
||||
}
|
||||
|
||||
/// Every mergeable token's bytes, sorted bytewise like tiktoken's `token_byte_values`.
|
||||
#[cfg(feature = "tiktoken")]
|
||||
fn token_byte_values<'py>(&self, py: Python<'py>) -> PyResult<Vec<Bound<'py, PyBytes>>> {
|
||||
let vocabulary = self.vocabulary()?;
|
||||
let values = release_gil(py, || vocabulary.token_byte_values());
|
||||
Ok(values.iter().map(|value| PyBytes::new(py, value)).collect())
|
||||
}
|
||||
|
||||
/// The token of one whole piece; `KeyError` when it is not in the vocabulary.
|
||||
#[cfg(feature = "tiktoken")]
|
||||
fn encode_single_token(&self, py: Python<'_>, piece: Vec<u8>) -> PyResult<u32> {
|
||||
self.vocabulary()?
|
||||
.encode_single_token(&piece)
|
||||
.ok_or_else(|| PyKeyError::new_err(PyBytes::new(py, &piece).unbind()))
|
||||
}
|
||||
|
||||
#[cfg(feature = "tiktoken")]
|
||||
fn decode_bytes<'py>(&self, py: Python<'py>, ids: Vec<u32>) -> PyResult<Bound<'py, PyBytes>> {
|
||||
enter_native()?;
|
||||
let tokenizer = self.tiktoken()?;
|
||||
let bytes =
|
||||
release_gil(py, || tokenizer.decode_bytes(&ids)).map_err(PyKeyError::new_err)?;
|
||||
Ok(PyBytes::new(py, &bytes))
|
||||
}
|
||||
|
||||
// ---- Hugging Face: the `tokenizers.Tokenizer` surface -----------------------------------
|
||||
|
||||
#[cfg(feature = "huggingface")]
|
||||
#[pyo3(signature = (sequence, pair = None, is_pretokenized = false, add_special_tokens = true, fast = false))]
|
||||
fn encode_huggingface(
|
||||
&self,
|
||||
py: Python<'_>,
|
||||
sequence: Sequence,
|
||||
pair: Option<Sequence>,
|
||||
is_pretokenized: bool,
|
||||
add_special_tokens: bool,
|
||||
fast: bool,
|
||||
) -> PyResult<HuggingFaceEncoding> {
|
||||
enter_native()?;
|
||||
let tokenizer = self.huggingface()?;
|
||||
let sequence = sequence.input(is_pretokenized)?;
|
||||
let input = match pair {
|
||||
Some(pair) => EncodeInput::Dual(sequence, pair.input(is_pretokenized)?),
|
||||
None => EncodeInput::Single(sequence),
|
||||
};
|
||||
release_gil(py, || {
|
||||
tokenizer.encode_result(input, add_special_tokens, fast)
|
||||
})
|
||||
.map(|inner| HuggingFaceEncoding { inner })
|
||||
.map_err(|error| token_count_error_to_pyerr(Error::from(error)))
|
||||
}
|
||||
|
||||
#[cfg(feature = "huggingface")]
|
||||
#[pyo3(signature = (inputs, is_pretokenized = false, add_special_tokens = true, fast = false))]
|
||||
fn encode_batch_huggingface(
|
||||
&self,
|
||||
py: Python<'_>,
|
||||
inputs: Vec<(Sequence, Option<Sequence>)>,
|
||||
is_pretokenized: bool,
|
||||
add_special_tokens: bool,
|
||||
fast: bool,
|
||||
) -> PyResult<Vec<HuggingFaceEncoding>> {
|
||||
enter_native()?;
|
||||
let tokenizer = self.huggingface()?;
|
||||
let inputs = inputs
|
||||
.into_iter()
|
||||
.map(|(sequence, pair)| {
|
||||
let sequence = sequence.input(is_pretokenized)?;
|
||||
match pair {
|
||||
Some(pair) => Ok(EncodeInput::Dual(sequence, pair.input(is_pretokenized)?)),
|
||||
None => Ok(EncodeInput::Single(sequence)),
|
||||
}
|
||||
})
|
||||
.collect::<PyResult<Vec<_>>>()?;
|
||||
release_gil(py, || {
|
||||
tokenizer.encode_batch_result(inputs, add_special_tokens, fast)
|
||||
})
|
||||
.map(|encodings| {
|
||||
encodings
|
||||
.into_iter()
|
||||
.map(|inner| HuggingFaceEncoding { inner })
|
||||
.collect()
|
||||
})
|
||||
.map_err(|error| token_count_error_to_pyerr(Error::from(error)))
|
||||
}
|
||||
|
||||
#[cfg(feature = "huggingface")]
|
||||
#[pyo3(signature = (pretty = false))]
|
||||
fn to_json(&self, py: Python<'_>, pretty: bool) -> PyResult<String> {
|
||||
enter_native()?;
|
||||
let tokenizer = self.huggingface()?;
|
||||
release_gil(py, || tokenizer.to_json(pretty))
|
||||
.map_err(|error| token_count_error_to_pyerr(Error::from(error)))
|
||||
}
|
||||
|
||||
#[cfg(feature = "huggingface")]
|
||||
fn token_to_id(&self, token: &str) -> PyResult<Option<u32>> {
|
||||
Ok(self.huggingface()?.token_to_id(token))
|
||||
}
|
||||
|
||||
#[cfg(feature = "huggingface")]
|
||||
fn id_to_token(&self, id: u32) -> PyResult<Option<String>> {
|
||||
Ok(self.huggingface()?.id_to_token(id))
|
||||
}
|
||||
|
||||
#[cfg(feature = "huggingface")]
|
||||
#[pyo3(signature = (with_added_tokens = true))]
|
||||
fn get_vocab(&self, py: Python<'_>, with_added_tokens: bool) -> PyResult<HashMap<String, u32>> {
|
||||
let tokenizer = self.huggingface()?;
|
||||
Ok(release_gil(py, || tokenizer.vocab(with_added_tokens)))
|
||||
}
|
||||
|
||||
#[cfg(feature = "huggingface")]
|
||||
#[pyo3(signature = (with_added_tokens = true))]
|
||||
fn get_vocab_size(&self, with_added_tokens: bool) -> PyResult<usize> {
|
||||
Ok(self.huggingface()?.vocab_size(with_added_tokens))
|
||||
}
|
||||
|
||||
/// The added tokens by id as `(id, (content, single_word, lstrip, rstrip, normalized,
|
||||
/// special))`, for Python to rebuild as `tokenizers.AddedToken`.
|
||||
#[cfg(feature = "huggingface")]
|
||||
fn added_tokens_decoder(&self) -> PyResult<Vec<(u32, AddedTokenFields)>> {
|
||||
Ok(self
|
||||
.huggingface()?
|
||||
.added_tokens_decoder()
|
||||
.into_iter()
|
||||
.map(|(id, token)| {
|
||||
(
|
||||
id,
|
||||
(
|
||||
token.content,
|
||||
token.single_word,
|
||||
token.lstrip,
|
||||
token.rstrip,
|
||||
token.normalized,
|
||||
token.special,
|
||||
),
|
||||
)
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
/// The padding parameters as `tokenizers.Tokenizer.padding` reports them.
|
||||
#[cfg(feature = "huggingface")]
|
||||
fn padding<'py>(&self, py: Python<'py>) -> PyResult<Option<Bound<'py, PyDict>>> {
|
||||
let Some(params) = self.huggingface()?.padding() else {
|
||||
return Ok(None);
|
||||
};
|
||||
let padding = PyDict::new(py);
|
||||
padding.set_item(
|
||||
"length",
|
||||
match params.strategy {
|
||||
PaddingStrategy::BatchLongest => None,
|
||||
PaddingStrategy::Fixed(length) => Some(length),
|
||||
},
|
||||
)?;
|
||||
padding.set_item("pad_to_multiple_of", params.pad_to_multiple_of)?;
|
||||
padding.set_item("pad_id", params.pad_id)?;
|
||||
padding.set_item("pad_type_id", params.pad_type_id)?;
|
||||
padding.set_item("pad_token", ¶ms.pad_token)?;
|
||||
padding.set_item("direction", params.direction.as_ref())?;
|
||||
Ok(Some(padding))
|
||||
}
|
||||
|
||||
/// The truncation parameters as `tokenizers.Tokenizer.truncation` reports them.
|
||||
#[cfg(feature = "huggingface")]
|
||||
fn truncation<'py>(&self, py: Python<'py>) -> PyResult<Option<Bound<'py, PyDict>>> {
|
||||
let Some(params) = self.huggingface()?.truncation() else {
|
||||
return Ok(None);
|
||||
};
|
||||
let truncation = PyDict::new(py);
|
||||
truncation.set_item("max_length", params.max_length)?;
|
||||
truncation.set_item("stride", params.stride)?;
|
||||
truncation.set_item("strategy", params.strategy.as_ref())?;
|
||||
truncation.set_item("direction", params.direction.as_ref())?;
|
||||
Ok(Some(truncation))
|
||||
}
|
||||
|
||||
#[cfg(feature = "huggingface")]
|
||||
fn num_special_tokens_to_add(&self, is_pair: bool) -> PyResult<usize> {
|
||||
Ok(self.huggingface()?.num_special_tokens_to_add(is_pair))
|
||||
}
|
||||
|
||||
#[cfg(feature = "huggingface")]
|
||||
fn encode_special_tokens(&self) -> PyResult<bool> {
|
||||
Ok(self.huggingface()?.encode_special_tokens())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "huggingface")]
|
||||
type AddedTokenFields = (String, bool, bool, bool, bool, bool);
|
||||
|
||||
impl Tokenizer {
|
||||
fn new(inner: Codec) -> Self {
|
||||
Self {
|
||||
inner: Arc::new(inner),
|
||||
#[cfg(feature = "fast")]
|
||||
fast: OnceLock::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn counter(&self, py: Python<'_>, fast: bool) -> SharedCounter {
|
||||
#[cfg(feature = "fast")]
|
||||
if fast {
|
||||
let counter = self.fast.get().unwrap_or_else(|| {
|
||||
release_gil(py, || {
|
||||
self.fast
|
||||
.get_or_init(|| self.inner.fast_counter().map(Arc::new))
|
||||
})
|
||||
});
|
||||
if let Some(counter) = counter {
|
||||
return SharedCounter::Fast(Arc::clone(counter));
|
||||
}
|
||||
}
|
||||
#[cfg(not(feature = "fast"))]
|
||||
let _ = (py, fast);
|
||||
SharedCounter::Codec(Arc::clone(&self.inner))
|
||||
}
|
||||
|
||||
/// A Python `str` as UTF-8. tiktoken replaces lone surrogates the way its Python `encode`
|
||||
/// does; `tokenizers` rejects them, so that backend keeps the encode error.
|
||||
fn text<'a>(&self, text: &'a Bound<'_, PyString>) -> PyResult<Cow<'a, str>> {
|
||||
match text.to_cow() {
|
||||
Ok(text) => Ok(text),
|
||||
Err(error) => match *self.inner {
|
||||
#[cfg(feature = "tiktoken")]
|
||||
Codec::Tiktoken(_) if error.is_instance_of::<PyUnicodeEncodeError>(text.py()) => {
|
||||
text.call_method1("encode", ("utf-16", "surrogatepass"))?
|
||||
.call_method1("decode", ("utf-16", "replace"))?
|
||||
.extract::<String>()
|
||||
.map(Cow::Owned)
|
||||
}
|
||||
_ => Err(error),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "tiktoken")]
|
||||
fn tiktoken(&self) -> PyResult<&TiktokenTokenizer> {
|
||||
match *self.inner {
|
||||
Codec::Tiktoken(ref tokenizer) => Ok(tokenizer),
|
||||
#[cfg(feature = "huggingface")]
|
||||
Codec::HuggingFace(_) => Err(PyValueError::new_err("requires a tiktoken encoding")),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "tiktoken")]
|
||||
fn vocabulary(&self) -> PyResult<&Vocabulary> {
|
||||
self.tiktoken()?.vocabulary().ok_or_else(|| {
|
||||
PyRuntimeError::new_err("this encoding was built without its vocabulary")
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(feature = "huggingface")]
|
||||
fn huggingface(&self) -> PyResult<&HuggingFaceTokenizer> {
|
||||
match *self.inner {
|
||||
Codec::HuggingFace(ref tokenizer) => Ok(tokenizer),
|
||||
#[cfg(feature = "tiktoken")]
|
||||
Codec::Tiktoken(_) => Err(PyValueError::new_err("requires a Hugging Face tokenizer")),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) enum SharedCounter {
|
||||
Codec(Arc<Codec>),
|
||||
#[cfg(feature = "fast")]
|
||||
Fast(Arc<FastTokenizer>),
|
||||
}
|
||||
|
||||
impl litellm_token_counter::Tokenizer for SharedCounter {
|
||||
fn count_tokens(&self, text: &str) -> Result<usize, Error> {
|
||||
match self {
|
||||
Self::Codec(codec) => codec.codec().count_tokens(text),
|
||||
#[cfg(feature = "fast")]
|
||||
Self::Fast(counter) => counter.count_tokens(text).map_err(Error::from),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "huggingface")]
|
||||
#[derive(FromPyObject)]
|
||||
pub(crate) enum Sequence {
|
||||
Text(String),
|
||||
Words(Vec<String>),
|
||||
}
|
||||
|
||||
#[cfg(feature = "huggingface")]
|
||||
impl Sequence {
|
||||
fn input(self, is_pretokenized: bool) -> PyResult<InputSequence<'static>> {
|
||||
match (self, is_pretokenized) {
|
||||
(Self::Text(text), false) => Ok(text.into()),
|
||||
(Self::Words(words), true) => Ok(words.into()),
|
||||
_ => Err(pyo3::exceptions::PyTypeError::new_err(
|
||||
"input must match is_pretokenized",
|
||||
)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "huggingface")]
|
||||
fn direction<T>(value: &str, left: T, right: T, what: &str) -> PyResult<T> {
|
||||
match value {
|
||||
"left" => Ok(left),
|
||||
"right" => Ok(right),
|
||||
other => Err(PyValueError::new_err(format!(
|
||||
"invalid {what} direction {other:?}: expected 'left' or 'right'"
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
/// `tokenizers.Encoding`, mutable like the original: `pad`, `truncate` and `set_sequence_id`
|
||||
/// change it in place.
|
||||
#[cfg(feature = "huggingface")]
|
||||
#[pyclass(module = "litellm.rust_bridge._native")]
|
||||
pub(crate) struct HuggingFaceEncoding {
|
||||
inner: Encoding,
|
||||
}
|
||||
|
||||
#[cfg(feature = "huggingface")]
|
||||
#[pymethods]
|
||||
impl HuggingFaceEncoding {
|
||||
#[new]
|
||||
#[pyo3(signature = (json = None))]
|
||||
fn new(json: Option<&str>) -> PyResult<Self> {
|
||||
let inner = match json {
|
||||
Some(json) => encoding_from_json(json)
|
||||
.map_err(|error| PyValueError::new_err(error.to_string()))?,
|
||||
None => Encoding::default(),
|
||||
};
|
||||
Ok(Self { inner })
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
#[pyo3(signature = (encodings, growing_offsets = true))]
|
||||
fn merge(encodings: Vec<PyRef<'_, Self>>, growing_offsets: bool) -> Self {
|
||||
Self {
|
||||
inner: Encoding::merge(
|
||||
encodings.iter().map(|encoding| encoding.inner.clone()),
|
||||
growing_offsets,
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
fn __reduce__<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
) -> PyResult<(Bound<'py, pyo3::types::PyType>, (String,))> {
|
||||
let json = encoding_to_json(&self.inner)
|
||||
.map_err(|error| PyValueError::new_err(error.to_string()))?;
|
||||
Ok((py.get_type::<Self>(), (json,)))
|
||||
}
|
||||
|
||||
fn __repr__(&self) -> String {
|
||||
format!(
|
||||
"Encoding(num_tokens={}, attributes=[ids, type_ids, tokens, offsets, \
|
||||
attention_mask, special_tokens_mask, overflowing])",
|
||||
self.inner.len()
|
||||
)
|
||||
}
|
||||
|
||||
fn __len__(&self) -> usize {
|
||||
self.inner.len()
|
||||
}
|
||||
#[getter]
|
||||
fn ids(&self) -> Vec<u32> {
|
||||
self.inner.get_ids().to_vec()
|
||||
}
|
||||
#[getter]
|
||||
fn tokens(&self) -> Vec<String> {
|
||||
self.inner.get_tokens().to_vec()
|
||||
}
|
||||
#[getter]
|
||||
fn offsets(&self) -> Vec<(usize, usize)> {
|
||||
self.inner.get_offsets().to_vec()
|
||||
}
|
||||
#[getter]
|
||||
fn type_ids(&self) -> Vec<u32> {
|
||||
self.inner.get_type_ids().to_vec()
|
||||
}
|
||||
#[getter]
|
||||
fn attention_mask(&self) -> Vec<u32> {
|
||||
self.inner.get_attention_mask().to_vec()
|
||||
}
|
||||
#[getter]
|
||||
fn special_tokens_mask(&self) -> Vec<u32> {
|
||||
self.inner.get_special_tokens_mask().to_vec()
|
||||
}
|
||||
#[getter]
|
||||
fn word_ids(&self) -> Vec<Option<u32>> {
|
||||
self.inner.get_word_ids().to_vec()
|
||||
}
|
||||
#[getter]
|
||||
fn sequence_ids(&self) -> Vec<Option<usize>> {
|
||||
self.inner.get_sequence_ids()
|
||||
}
|
||||
#[getter]
|
||||
fn overflowing(&self) -> Vec<Self> {
|
||||
self.inner
|
||||
.get_overflowing()
|
||||
.iter()
|
||||
.cloned()
|
||||
.map(|inner| Self { inner })
|
||||
.collect()
|
||||
}
|
||||
#[getter]
|
||||
fn n_sequences(&self) -> usize {
|
||||
self.inner.n_sequences()
|
||||
}
|
||||
|
||||
#[pyo3(signature = (word_index, sequence_index = 0))]
|
||||
fn word_to_tokens(&self, word_index: u32, sequence_index: usize) -> Option<(usize, usize)> {
|
||||
self.inner.word_to_tokens(word_index, sequence_index)
|
||||
}
|
||||
#[pyo3(signature = (word_index, sequence_index = 0))]
|
||||
fn word_to_chars(&self, word_index: u32, sequence_index: usize) -> Option<(usize, usize)> {
|
||||
self.inner.word_to_chars(word_index, sequence_index)
|
||||
}
|
||||
fn token_to_sequence(&self, token_index: usize) -> Option<usize> {
|
||||
self.inner.token_to_sequence(token_index)
|
||||
}
|
||||
fn token_to_chars(&self, token_index: usize) -> Option<(usize, usize)> {
|
||||
self.inner
|
||||
.token_to_chars(token_index)
|
||||
.map(|(_, offsets)| offsets)
|
||||
}
|
||||
fn token_to_word(&self, token_index: usize) -> Option<u32> {
|
||||
self.inner.token_to_word(token_index).map(|(_, word)| word)
|
||||
}
|
||||
#[pyo3(signature = (char_pos, sequence_index = 0))]
|
||||
fn char_to_token(&self, char_pos: usize, sequence_index: usize) -> Option<usize> {
|
||||
self.inner.char_to_token(char_pos, sequence_index)
|
||||
}
|
||||
#[pyo3(signature = (char_pos, sequence_index = 0))]
|
||||
fn char_to_word(&self, char_pos: usize, sequence_index: usize) -> Option<u32> {
|
||||
self.inner.char_to_word(char_pos, sequence_index)
|
||||
}
|
||||
|
||||
fn set_sequence_id(&mut self, sequence_id: usize) {
|
||||
self.inner.set_sequence_id(sequence_id);
|
||||
}
|
||||
|
||||
#[pyo3(signature = (length, direction = "right", pad_id = 0, pad_type_id = 0, pad_token = "[PAD]"))]
|
||||
fn pad(
|
||||
&mut self,
|
||||
length: usize,
|
||||
direction: &str,
|
||||
pad_id: u32,
|
||||
pad_type_id: u32,
|
||||
pad_token: &str,
|
||||
) -> PyResult<()> {
|
||||
let direction = self::direction(
|
||||
direction,
|
||||
PaddingDirection::Left,
|
||||
PaddingDirection::Right,
|
||||
"padding",
|
||||
)?;
|
||||
self.inner
|
||||
.pad(length, pad_id, pad_type_id, pad_token, direction);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[pyo3(signature = (max_length, stride = 0, direction = "right"))]
|
||||
fn truncate(&mut self, max_length: usize, stride: usize, direction: &str) -> PyResult<()> {
|
||||
let direction = self::direction(
|
||||
direction,
|
||||
TruncationDirection::Left,
|
||||
TruncationDirection::Right,
|
||||
"truncation",
|
||||
)?;
|
||||
self.inner.truncate(max_length, stride, direction);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
|
@ -4,7 +4,7 @@ use serde::{Deserialize, Serialize};
|
|||
|
||||
use crate::SecretValue;
|
||||
|
||||
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
|
||||
#[derive(Clone, Copy, Debug, Deserialize, Eq, Hash, PartialEq, Serialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum KeyManagementSystem {
|
||||
GoogleKms,
|
||||
|
|
@ -18,7 +18,7 @@ pub enum KeyManagementSystem {
|
|||
Custom,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
|
||||
#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, Hash, PartialEq, Serialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum AccessMode {
|
||||
#[default]
|
||||
|
|
@ -33,7 +33,7 @@ impl AccessMode {
|
|||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Deserialize, Serialize, PartialEq, Eq)]
|
||||
#[derive(Clone, Debug, Deserialize, Serialize, PartialEq, Eq, Hash)]
|
||||
#[serde(default)]
|
||||
pub struct KeyManagementSettings {
|
||||
pub hosted_keys: Option<Vec<String>>,
|
||||
|
|
|
|||
|
|
@ -24,6 +24,8 @@ pub enum Error {
|
|||
OidcFile,
|
||||
#[error("secret cannot be converted to {expected}")]
|
||||
TypeMismatch { expected: &'static str },
|
||||
#[error("external secret manager failed")]
|
||||
ExternalManager(#[source] Box<dyn std::error::Error + Send + Sync>),
|
||||
#[cfg(feature = "aws")]
|
||||
#[error(transparent)]
|
||||
Aws(#[from] litellm_secrets_aws::Error),
|
||||
|
|
|
|||
|
|
@ -1,10 +1,24 @@
|
|||
use std::{future::Future, pin::Pin, sync::Arc};
|
||||
|
||||
use litellm_core_utils::settings::Lookup;
|
||||
|
||||
use crate::{Error, KeyManagementSettings, KeyManagementSystem, Secret, SecretValue};
|
||||
|
||||
pub trait ExternalSecretManager: Send + Sync {
|
||||
fn system(&self) -> KeyManagementSystem;
|
||||
|
||||
fn read_secret<'a>(
|
||||
&'a self,
|
||||
name: &'a str,
|
||||
settings: &'a KeyManagementSettings,
|
||||
environment: &'a (dyn Lookup + Send + Sync),
|
||||
) -> Pin<Box<dyn Future<Output = Result<Option<Secret>, Error>> + Send + 'a>>;
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub enum SecretManager {
|
||||
Local,
|
||||
External(Arc<dyn ExternalSecretManager>),
|
||||
#[cfg(feature = "aws")]
|
||||
AwsKms(crate::aws::AwsKms),
|
||||
#[cfg(feature = "aws")]
|
||||
|
|
@ -25,6 +39,7 @@ impl SecretManager {
|
|||
pub fn system(&self) -> KeyManagementSystem {
|
||||
match self {
|
||||
Self::Local => KeyManagementSystem::Local,
|
||||
Self::External(manager) => manager.system(),
|
||||
#[cfg(feature = "aws")]
|
||||
Self::AwsKms(_) => KeyManagementSystem::AwsKms,
|
||||
#[cfg(feature = "aws")]
|
||||
|
|
@ -54,6 +69,11 @@ pub async fn get_secret_from_manager(
|
|||
.get(secret_name)
|
||||
.map(SecretValue::new)
|
||||
.map(Secret::String)),
|
||||
SecretManager::External(manager) => {
|
||||
manager
|
||||
.read_secret(secret_name, _settings, environment)
|
||||
.await
|
||||
}
|
||||
#[cfg(feature = "aws")]
|
||||
SecretManager::AwsKms(client) => {
|
||||
let ciphertext = environment
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ mod resolver;
|
|||
mod state;
|
||||
|
||||
pub use error::Error;
|
||||
pub use handler::{SecretManager, get_secret_from_manager};
|
||||
pub use handler::{ExternalSecretManager, SecretManager, get_secret_from_manager};
|
||||
pub use litellm_secrets_types::{
|
||||
AccessMode, KeyManagementSettings, KeyManagementSystem, Secret, SecretValue,
|
||||
};
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue