scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-12 08:42:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 207

Merge remote-tracking branch 'origin/main' into HEAD
Some checks failed
Cost map pins recheck / rerun-cost-map-guard (push) Has been cancelled
Details
LiteLLM Rust DB / db-queries (push) Has been cancelled
Details
LiteLLM Rust / rust-lint (push) Has been cancelled
Details
LiteLLM Rust / rust-test (push) Has been cancelled
Details
LiteLLM Rust / rust-wheel (push) Has been cancelled
Details
Terraform Provider / gofmt, vet, build, test (push) Has been cancelled
Details
Terraform Provider / Provider endpoints vs proxy OpenAPI schema (push) Has been cancelled
Details
Terraform Modules / fmt, validate, test (aws) (push) Has been cancelled
Details
Terraform Modules / fmt, validate, test (gcp) (push) Has been cancelled
Details

Browse source 
# Conflicts:
#	tests/unit/proxy/vector_store_files_endpoints/test_endpoints.py
This commit is contained in:
mateo-berri 2026-11-10 10:28:11 -08:00
commit c6174e0c6f
1799 changed files with 140403 additions and 34203 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
.github/workflows/test-unit.yml vendored
View file

@ -182,6 +182,7 @@ jobs:
tests/unit/google_genai
tests/unit/router_strategy
tests/unit/router_utils
tests/unit/redis
tests/unit/proxy/common_utils/test_cache_aware_routing.py
tests/unit/enterprise/enterprise_callbacks/send_emails
tests/unit/enterprise/proxy/test_afile_retrieve_returns_unified_id.py
@ -1034,9 +1035,9 @@ jobs:
run: |
diff -u model_prices_and_context_window.json litellm/model_prices_and_context_window_backup.json
if [ -z "$RUST_BRIDGE_ARTIFACT" ]; then
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra extra_proxy --extra proxy-runtime --extra utils
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra switchyard --extra saml --extra extra_proxy --extra proxy-runtime --extra utils
else
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra extra_proxy --extra proxy-runtime --extra utils --no-install-project --no-install-package litellm-proxy
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra switchyard --extra saml --extra extra_proxy --extra proxy-runtime --extra utils --no-install-project --no-install-package litellm-proxy
uv pip install --no-deps --python .venv/bin/python rust-bridge-dist/*.whl
cp rust-bridge-dist/litellm/rust_bridge/_native.abi3.so litellm/rust_bridge/_native.abi3.so
uv run --no-sync python -c "import importlib.metadata; import litellm.rust_bridge._native; print(importlib.metadata.version('litellm'))"
@ -1293,7 +1294,7 @@ jobs:
- name: Install dependencies
if: steps.changes.outputs.decision != 'skip'
run: |
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra switchyard
- name: Cache Prisma binaries
if: steps.changes.outputs.decision != 'skip'
@ -1496,9 +1497,9 @@ jobs:
RUST_BRIDGE_ARTIFACT: ${{ needs.rust-bridge.outputs.artifact }}
run: |
if [ -z "$RUST_BRIDGE_ARTIFACT" ]; then
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra extra_proxy --extra proxy-runtime --extra utils
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra switchyard --extra saml --extra extra_proxy --extra proxy-runtime --extra utils
else
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra extra_proxy --extra proxy-runtime --extra utils --no-install-project --no-install-package litellm-proxy
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra switchyard --extra saml --extra extra_proxy --extra proxy-runtime --extra utils --no-install-project --no-install-package litellm-proxy
uv pip install --no-deps --python .venv/bin/python rust-bridge-dist/*.whl
cp rust-bridge-dist/litellm/rust_bridge/_native.abi3.so litellm/rust_bridge/_native.abi3.so
fi

1
AGENTS.md
View file

@ -52,7 +52,6 @@ If you ever write any human-facing text (pull requests, issues, commit messages,
- don't use "—". Don't use ";" or ":" unless it would be nonsensical not to. Instead, reach for ",", ".", conjunction words, etc. in descending order of preference: vary among them, weighted toward the front of the list, and skip "," where it would cause a comma splice or the sentence is getting long. Overusing any one of them also feels AI-y. A word cap does not penalize you for adding more sentences
- don't use the pattern "It's not X, it's Y", "You're not X, you're Y", etc.
- unless explicitly asked, don't use bulleted or numbered lists unless it would be nonsensical not to. Instead, prefer prose
- don't add a trailing "." at the end of paragraphs (just like this file). That means every paragraph, not just the last one (of the markdown file, PR description, GitHub comment, etc.). Rule of thumb: if you're adding new line(s) before the next sentence, don't add a ".". Exception is that docs (i.e., litellm-docs) should have trailing "."
- don't use →. Instead, prefer not to use arrows, and if need be, use -> instead
- use plain, simple, everyday engineering language: the common phrase engineers actually say over rare compact phrasing, in grammatically complete sentences. When explicitly asked to use bullets or ordered lists and structure legitimately helps the reader, prefer nested bullets (any depth is fine) over dense lines in a flat structure

40
cookbook/ai_coding_tool_guides/chatgpt_subscription_passthrough.md
View file

@ -46,7 +46,43 @@ For a managed workspace rollout, use OpenAI's [gateway configuration guide](http
## Native compatibility and Pi
For a native Codex model catalog and request format, change Codex's `base_url` to `https://litellm.example.com/backend-api/codex`. This route uses the same credentials but does not need a model deployment or resolve model aliases. Restricted keys must allow the upstream model name.
For a native Codex model catalog and request format, change Codex's `base_url` to `https://litellm.example.com/backend-api/codex`. This route uses the same credentials. A ChatGPT model keeps its upstream name and needs no deployment of its own.
### Mix ChatGPT and LiteLLM models
On the native route, Codex's model picker shows OpenAI's catalog first and then the other chat and responses models the virtual key can use. Pick one of those and Codex's requests go through LiteLLM's `/v1/responses` with the key's usual routing, budgets, guardrails, and logging. The proxy removes the ChatGPT bearer and account header from that request, so the other provider and your callbacks never receive them. A name is routed this way when another provider serves it as an exact `model_name`, through a router `model_group_alias`, or as a team's public model name. A ChatGPT model, a name that only a wildcard deployment matches, or a name OpenAI's catalog lists for the caller's ChatGPT account keeps the passthrough, so a model OpenAI adds later still uses the subscription even when a deployment of another provider has the same name. The proxy learns the account's list from the catalog Codex fetches, or asks OpenAI with the caller's login. When it cannot get the list, for example because OpenAI is down, the request times out, or the client is not Codex, every request keeps the passthrough for the next minute, as it did before LiteLLM models were added to the picker. When the key can use many models, the picker lists the ones that fit in Codex's catalog size limit, and the proxy logs a warning with how many it left out.
Codex's picker sends a `none` reasoning effort for a LiteLLM model whose catalog entry lists no effort levels. With `litellm_settings: drop_params: true`, the proxy drops it for a model that does not support `none`, and the provider applies its own default. A model whose cost map entry or `model_info` sets `supports_none_reasoning_effort: true` keeps it. Without `drop_params`, the effort reaches the provider unchanged.
Codex also sends its ChatGPT app tools with every turn, which on a LiteLLM model adds a few hundred thousand prompt tokens to each turn. Set `apps = false` under `[features]` in Codex's `config.toml` to turn them off.
Only a key that may call `/v1/responses` gets these models. A key whose `allowed_routes` covers only the native routes sees OpenAI's catalog unchanged, and every request it sends uses the subscription.
Codex sends earlier reasoning items back with each turn. The proxy tags the reasoning that a LiteLLM model returns on this route, OpenAI and Azure deployments included. When a session switches to a ChatGPT model, the proxy drops those items before forwarding the request to OpenAI, which would reject them. Messages, tool calls, and ChatGPT's own reasoning are kept, and a LiteLLM model still gets its own reasoning back.
When the catalog adds LiteLLM models for a key, its `ETag` is a digest of OpenAI's ETag and those models, and that key's responses carry the same value in `X-Models-Etag`. Codex keeps its cached catalog while the two match and fetches it again when OpenAI's catalog or the key's models change. A key whose catalog adds no LiteLLM models keeps OpenAI's ETags.
### Restrict native models with an access group
A key with a model allowlist needs a rule for ChatGPT names that OpenAI adds over time. Put the subscription wildcards in an access group under `chatgpt/` names:
```yaml
model_list:
- model_name: chatgpt/gpt-*
litellm_params:
model: chatgpt/gpt-*
model_info:
access_groups: ["chatgpt-native"]
- model_name: chatgpt/codex-*
litellm_params:
model: chatgpt/codex-*
model_info:
access_groups: ["chatgpt-native"]
```
Then add `chatgpt-native` to the key's models, next to any other model names it should reach. On the native route, a key that holds only the `chatgpt/` form of a name is checked as that form, so a request for `gpt-5.6-sol` is allowed when the key may call `chatgpt/gpt-5.6-sol`. The key can call every ChatGPT model that matches the wildcards, including ones the proxy configuration does not name.
On `/v1` these wildcards match only a request that asks for a `chatgpt/` name, so a request for a plain `gpt-` or `codex-` name still goes where it went before, such as a `*` catch-all deployment. A `/v1` request for a `chatgpt/` name without a ChatGPT login falls back to the proxy's own saved login if it has one, so keep that login unset as described above.
### Configure Pi
@ -71,7 +107,7 @@ Keep Pi's subscription login as the provider credential. The header value names
Send a prompt and check the completed request in LiteLLM's Logs page. Its gateway user/key and token counts should match the caller. Give each person their own LiteLLM virtual key for separate attribution. The gateway credential and ChatGPT login are independently authenticated; LiteLLM does not establish that both identities belong to the same person.
The native routes are `GET /backend-api/codex/models` and `POST /backend-api/codex/responses`. Model requests use the upstream model name, so a restricted LiteLLM key must allow that name. OpenAI's model catalog is returned in its native format; seeing a model there does not override gateway permissions.
The native routes are `GET /backend-api/codex/models` and `POST /backend-api/codex/responses`. Model requests use the upstream model name, so a restricted LiteLLM key needs the `chatgpt-native` access group above or that exact name. OpenAI's model catalog is returned in its native format; seeing a model there does not override gateway permissions.
The proxy forwards the ChatGPT bearer and account header to the fixed OpenAI Codex endpoint. It removes gateway credentials and cookies from the upstream request and redacts credential headers in logging callbacks. The client remains responsible for refresh and reauthentication. Native passthrough returns authentication and quota errors without switching subscriptions. Standard model routes retain the caller login across ChatGPT retries; explicitly configured cross-provider fallbacks use their own credentials. An explicit workspace header is preserved for OpenAI to validate.

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

@ -8,7 +8,6 @@ import re
from collections.abc import Callable, Mapping, Sequence
from dataclasses import replace as dataclasses_replace
from datetime import datetime, timedelta
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, List, Literal, Optional, Protocol, Tuple, cast
from pydantic import TypeAdapter, ValidationError
@ -766,7 +765,6 @@ class CheckBatchCost:
credentials: Final = self.llm_router.registry().credentials(model_id, form="with_provider") or {}
_file_content = await afile_content(
file_id=cast(str, raw_output_file_id),
_litellm_internal_model_credentials=MappingProxyType(dict(credentials)),
**cast(FileContentCallOptions, credentials),
)
@ -867,10 +865,7 @@ class CheckBatchCost:
error_file_failed_requests: Final = await count_error_file_failed_requests(
response,
custom_llm_provider=batch_file_provider,
litellm_params={
**credentials,
"_litellm_internal_model_credentials": MappingProxyType(dict(credentials)),
},
litellm_params=credentials,
)
batch_result: Final = (
output_file_result
@ -1029,11 +1024,7 @@ class CheckBatchCost:
custom_llm_provider=deployment_info.litellm_params.custom_llm_provider,
)
cost_provider: Final = cast(BatchFileProvider, llm_provider)
credentials: Final = self.llm_router.registry().credentials(model_id, form="with_provider") or {}
file_access: Final = {
**credentials,
"_litellm_internal_model_credentials": MappingProxyType(dict(credentials)),
}
file_access: Final = self.llm_router.registry().credentials(model_id, form="with_provider") or {}
output_content: Final = (
await fetch_batch_managed_file_content(
response.output_file_id, custom_llm_provider=cost_provider, litellm_params=file_access

11
enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py
View file

@ -182,7 +182,7 @@ class ManagedObjectCreatorAttribution:
team_id: Final = getattr(job, "team_id", None)
rows: Final = _CreatorRows(
key=await self._live_key_row(api_key, job.created_by),
team=await self._live_team_row(team_id),
team=await self._live_team_row(team_id, job.created_at),
user=await self._user_row(object_id, _creator_user_id(job.created_by)),
)
return ResponseCreator(
@ -269,10 +269,15 @@ class ManagedObjectCreatorAttribution:
verbose_proxy_logger.error(f"{self._poller_name}: could not look up team {team_id} for {object_id}: {e}")
return None
async def _live_team_row(self, team_id: str | None) -> prisma_models.LiteLLM_TeamTable | None:
async def _live_team_row(
self, team_id: str | None, created_before: datetime | None = None
) -> prisma_models.LiteLLM_TeamTable | None:
if not team_id:
return None
return await _team_table(self._prisma_client).find_unique(where={"team_id": team_id})
team: Final = await _team_table(self._prisma_client).find_unique(where={"team_id": team_id})
if team is not None and created_before is not None and team.created_at > created_before:
return None
return team
def _response_creator_metadata(job: ManagedObjectRow, rows: _CreatorRows) -> Mapping[str, object]:

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

@ -78,6 +78,7 @@ from litellm.proxy.litellm_pre_call_utils import (
from litellm.proxy.openai_files_endpoints.common_utils import (
BATCH_CREATE_HIDDEN_PARAM,
FILE_LIST_CONTINUATION_CHUNK_SIZE,
TeamProviderDeployment,
_is_base64_encoded_unified_file_id,
apply_unified_file_ids,
decode_model_from_file_id,
@ -127,7 +128,6 @@ class _ManagedFileRetrieve(Protocol):
self,
*,
file_id: str,
_litellm_internal_model_credentials: Mapping[str, object] | None = None,
**kwargs: Unpack[FileRetrieveCallOptions],
) -> OpenAIFileObject: ...
@ -400,7 +400,7 @@ class _RouterFileCallKwargs(TypedDict, total=False):
_MODEL_FILE_ID_MAPPING: Final = TypeAdapter(dict[str, str])
_FILE_UPLOAD_ONLY_FIELDS: Final = frozenset({"file", "purpose", "expires_after", "extra_body"})
_FILE_DELETE_RESERVED_FIELDS: Final = frozenset({"model", "file_id", "_litellm_internal_model_credentials"})
_FILE_DELETE_RESERVED_FIELDS: Final = frozenset({"model", "file_id"})
@dataclass(frozen=True, slots=True)
@ -459,15 +459,7 @@ async def _delete_provider_file(
provider_file_id: str,
request_data: Mapping[str, object],
) -> None:
credentials: Final = llm_router.registry().credentials(model_id, form="with_provider")
delete_data: Final = {
**{k: v for k, v in request_data.items() if k not in _FILE_DELETE_RESERVED_FIELDS},
**(
{"_litellm_internal_model_credentials": MappingProxyType(dict(credentials))}
if credentials is not None
else {}
),
}
delete_data: Final = {k: v for k, v in request_data.items() if k not in _FILE_DELETE_RESERVED_FIELDS}
await llm_router.afile_delete(
model=model_id,
file_id=provider_file_id,
@ -1160,7 +1152,6 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
*,
provider_file_id: str,
call_options: Mapping[str, object],
internal_model_credentials: Mapping[str, object] | None = None,
) -> OpenAIFileObject:
"""Retrieve provider file details with transient retries and SDK retries disabled."""
retrieve_file: Final = cast(_ManagedFileRetrieve, litellm.afile_retrieve)
@ -1172,14 +1163,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
if attempt > 0:
await self._sleep(_PROVIDER_FILE_RETRIEVE_RETRY_DELAYS_SECONDS[attempt - 1])
try:
if internal_model_credentials is None:
return await retrieve_file(
file_id=provider_file_id,
**retrieve_options,
)
return await retrieve_file(
file_id=provider_file_id,
_litellm_internal_model_credentials=MappingProxyType(dict(internal_model_credentials)),
**retrieve_options,
)
except Exception as error:
@ -1221,7 +1206,6 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
self._afile_retrieve_with_retries(
provider_file_id=provider_file_id,
call_options=router_credentials,
internal_model_credentials=router_credentials,
),
timeout=BATCH_OUTPUT_FILE_LOOKUP_TIMEOUT_SECONDS,
)
@ -1443,6 +1427,43 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
user_api_key_dict=owner,
)
async def register_provider_file(
self,
*,
provider_file_id: str,
content_type: str,
deployment: TeamProviderDeployment,
owner: UserAPIKeyAuth,
llm_router: Router | None,
) -> str | None:
unified_file_id: Final = self.get_unified_output_file_id(
provider_file_id, deployment.deployment_id, deployment.public_model_name, file_type=content_type
)
provider_file_object, _ = await self._fetch_provider_file_object(
unified_file_id=unified_file_id,
provider_file_id=provider_file_id,
model_id=deployment.deployment_id,
model_name=None,
llm_router=llm_router,
)
if provider_file_object is None:
return None
try:
await self.store_unified_file_id(
file_id=unified_file_id,
file_object=provider_file_object,
litellm_parent_otel_span=None,
model_mappings={deployment.deployment_id: provider_file_id},
user_api_key_dict=owner,
)
except Exception as error: # noqa: BLE001 # the upload succeeded, so the caller keeps the provider id
verbose_logger.warning(
f"Failed to register provider_file_id={sanitize_for_log(provider_file_id)} as a managed file: "
f"{type(error).__name__} {sanitize_for_log(error)}"
)
return None
return unified_file_id
async def list_user_batches(
self,
user_api_key_dict: UserAPIKeyAuth,
@ -2270,10 +2291,12 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
unified_batch_id = SpecialEnums.LITELLM_MANAGED_BATCH_COMPLETE_STR.value.format(model_id, batch_id)
return base64.urlsafe_b64encode(unified_batch_id.encode()).decode().rstrip("=")
def get_unified_output_file_id(self, output_file_id: str, model_id: str, model_name: Optional[str]) -> str:
def get_unified_output_file_id(
self, output_file_id: str, model_id: str, model_name: Optional[str], file_type: str = "application/json"
) -> str:
deterministic_uuid: Final = uuid5(uuid5(NAMESPACE_URL, model_id), output_file_id)
unified_output_file_id = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format(
"application/json",
file_type,
str(deterministic_uuid),
model_name or "",
output_file_id,
@ -2866,15 +2889,6 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
exception_dict = {}
for model_id, provider_file_id in specific_model_file_id_mapping.items():
try:
# Cloud-storage providers (e.g. Bedrock S3) validate file ids
# against the deployment's configured bucket, which they only
# trust from this immutable server-side snapshot, never from
# request params.
credentials = llm_router.registry().credentials(model_id, form="with_provider")
if credentials is not None:
data["_litellm_internal_model_credentials"] = cast(Dict, MappingProxyType(dict(credentials)))
else:
data.pop("_litellm_internal_model_credentials", None)
return cast(
HttpxBinaryResponseContent,
await llm_router.afile_content(

4
enterprise/pyproject.toml
View file

@ -1,6 +1,6 @@
[project]
name = "litellm-enterprise"
version = "0.1.81"
version = "0.1.83"
description = "Package for LiteLLM Enterprise features"
readme = "README.md"
requires-python = ">=3.9"
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.1.81"
version = "0.1.83"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-enterprise==",

9
litellm-proxy-extras/litellm_proxy_extras/migrations/20261029000000_add_background_response_settlement/migration.sql Normal file
View file

@ -0,0 +1,9 @@
SET LOCAL lock_timeout = '10s';
-- CreateTable
CREATE TABLE IF NOT EXISTS "LiteLLM_BackgroundResponseSettlement" (
"response_id" TEXT NOT NULL,
"settled_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
CONSTRAINT "LiteLLM_BackgroundResponseSettlement_pkey" PRIMARY KEY ("response_id")
);

7
litellm-proxy-extras/litellm_proxy_extras/migrations/20261110000000_add_budget_model_tpm_rpm_limit/migration.sql Normal file
View file

@ -0,0 +1,7 @@
SET LOCAL lock_timeout = '10s';
-- AlterTable
ALTER TABLE "LiteLLM_BudgetTable" ADD COLUMN IF NOT EXISTS "model_tpm_limit" JSONB;
-- AlterTable
ALTER TABLE "LiteLLM_BudgetTable" ADD COLUMN IF NOT EXISTS "model_rpm_limit" JSONB;

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

@ -19,6 +19,8 @@ model LiteLLM_BudgetTable {
rpm_limit BigInt?
tpd_limit BigInt?
model_max_budget Json?
model_tpm_limit Json?
model_rpm_limit Json?
budget_duration String?
budget_reset_at DateTime?
allowed_models String[] @default([]) // per-member model scope; empty = inherit team models
@ -2053,6 +2055,11 @@ model LiteLLM_BackgroundInteractionSettlement {
@@index([settled_at], map: "idx_background_interaction_settlement_settled_at")
}
model LiteLLM_BackgroundResponseSettlement {
response_id String @id
settled_at DateTime @default(now())
}
model LiteLLM_Lens {
reviews LiteLLM_LensReview[]
id String @id

4
litellm-proxy-extras/pyproject.toml
View file

@ -1,6 +1,6 @@
[project]
name = "litellm-proxy-extras"
version = "0.4.112"
version = "0.4.114"
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
readme = "README.md"
requires-python = ">=3.9"
@ -30,7 +30,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.4.112"
version = "0.4.114"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-proxy-extras==",

53
litellm-rust/.agents/skills/rust-python-boundary/SKILL.md Normal file
View file

@ -0,0 +1,53 @@
---
name: rust-python-boundary
description: Decide where Python-shaped input becomes typed Rust in litellm-rust. Use when reading config.yaml values, os.environ/ references, REDIS_* or other environment layering, get_secret fallbacks, or kwargs and objects handed over from Python, when coercing text to a bool, int or float the way Python does, or when adding a dependency on litellm-python-compat
metadata:
internal: true
---
# Python boundary
Python-shaped input becomes typed Rust in exactly two crates. Every crate past them handles typed values, so no crate re-guesses how Python would have read a value
## The two boundaries
- `crates/config` reads config.yaml, which is Python-shaped: `os.environ/` references, scalars that Helm or the environment turned into text, raw kwargs maps. It also owns the environment layering and `get_secret` fallbacks Python applies while it loads that config. It has no Python runtime, so it reproduces Python value semantics with `litellm-python-compat`
- `crates/python-bridge` receives live Python objects, such as `RedisCache.redis_kwargs`. It reads them through PyO3 with the real protocol (`is_truthy()`, `str()`, `extract`), so it never imitates Python
- Both produce the same typed input, defined in the data crate the consumer already depends on (`redis-types` for Redis), so a downstream crate cannot tell which host it runs under
## Python value semantics
Reading a value the way Python itself reads it. This lives only at the two boundaries
- Truthiness and `x or fallback`, `str(x)`, `int(text)` and `float(text)`
- Each surface's true and false text: `_str_to_bool`, pydantic's lax bool, `get_secret`'s `true` and `false`
- `os.environ/` references and `REDIS_<KWARG>` style environment layering
A boundary knows which surface a value came from. `general_settings.coordination_redis` is a pydantic model and `cache_params` goes through `_str_to_bool`, so the boundary applies that surface's rule exactly instead of refusing the text two surfaces disagree on
Each typed field keeps where its value came from (the key, `REDIS_<KWARG>` or the fallback name), so an error or warning downstream still names it
## What stays downstream
A crate that stands in for a Python library keeps that library's semantics, because they are its compatibility contract and not value coercion. `redis-resolve` stands in for redis-py, so it keeps
- Which target wins, and the defaults redis-py applies
- How `parse_url` reads a url, including `int()` on `?db=` and redis-py's `to_bool` on the TLS query parameters
- Which kwargs redis-py accepts, and whether the Rust client reads, warns about or refuses each one
Its input is typed. A rule it applies reads typed fields, never a `serde_json::Value` it has to coerce first. A value it only carries back to Python untouched, like `passthrough`, may stay untyped
## litellm-python-compat
The rule is about where config and kwargs get coerced, not about who may depend on the crate
- Its data formats (`pickle`, `json.dumps`, `repr`) are fine anywhere Rust reads or writes state Python also reads or writes, such as cache values in a shared Redis
- Its coercions (`truthy`, `number`) belong at the two boundaries, and in a crate that stands in for a Python library when that library calls them itself
## Existing code
Some code still coerces Python-shaped values past the boundary, for example the `REDIS_*` merge in `gateway/src/redis/layers.rs` and the flag and number readers in `redis-resolve/src/options.rs`. Do not add more. When you change one, move the coercion to its boundary instead of extending it in place
## Tests
Pin a boundary's coercion with named cases next to the Python function it mirrors, in the Rust test and in its Python twin, like the Redis contract cases in `redis-resolve/tests/interpret.rs` and `tests/unit/test_redis.py`. A downstream crate's tests build its typed input directly and never feed it raw text to coerce

4
litellm-rust/AGENTS.md
View file

@ -8,6 +8,10 @@ For string-valued enums and their Serde conversions, follow [.agents/skills/rust
For error types, error variants and conversions between crate errors, follow [.agents/skills/rust-error/SKILL.md](.agents/skills/rust-error/SKILL.md)
For where config.yaml values, environment layering and kwargs from Python become typed Rust, follow [.agents/skills/rust-python-boundary/SKILL.md](.agents/skills/rust-python-boundary/SKILL.md)
The process-wide rustls `CryptoProvider` (ring) is installed once by the host. rust-proxy does it at the top of `main` through `litellm_gateway::install_crypto_provider`, and python-bridge must do the same in its module init before it builds any rustls config. Library crates never install it. The AWS SDK links aws-lc-rs next to ring, so without an installed default the first rustls config a library builds, such as redis-rs TLS, panics
For fieldless enums, derive `strum::VariantArray` and use `VARIANTS` instead of a hand-listed `ALL` array. Derive Strum string conversions instead of hand-written variant-to-string matches
Use the derived conversions directly (`<&'static str>::from(x)` / `.into()`, `str::parse`) with no `as_str`/`parse` wrapper that only delegates, and spell each variant with explicit `#[strum(serialize = "...")]` instead of `serialize_all`

245
litellm-rust/Cargo.lock generated
View file

@ -8,6 +8,17 @@ version = "2.0.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa"
[[package]]
name = "aes"
version = "0.8.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b169f7a6d4742236a0a00c541b845991d0ac43e546831af1249753ab4c3aa3a0"
dependencies = [
"cfg-if",
"cipher",
"cpufeatures 0.2.17",
]
[[package]]
name = "ahash"
version = "0.8.12"
@ -883,6 +894,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "cffb0e931875b666fc4fcb20fee52e9bbd1ef836fd9e9e04ec21555f9f85f7ef"
dependencies = [
"fastrand",
"tokio",
]
[[package]]
@ -982,6 +994,15 @@ dependencies = [
"hybrid-array",
]
[[package]]
name = "block-padding"
version = "0.3.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a8894febbff9f758034a5b8e12d87918f56dfc64a8e1fe757d65e29041538d93"
dependencies = [
"generic-array",
]
[[package]]
name = "bollard"
version = "0.20.2"
@ -1134,6 +1155,15 @@ dependencies = [
"rustversion",
]
[[package]]
name = "cbc"
version = "0.1.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "26b52a9543ae338f279b96b0b9fed9c8093744685043739079ce85cd58f289a6"
dependencies = [
"cipher",
]
[[package]]
name = "cc"
version = "1.3.0"
@ -1210,6 +1240,16 @@ dependencies = [
"half",
]
[[package]]
name = "cipher"
version = "0.4.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "773f3b9af64447d2ce9850330c473515014aa235e6a783b02db81ff39e4a3dad"
dependencies = [
"crypto-common 0.1.7",
"inout",
]
[[package]]
name = "clap"
version = "4.6.7"
@ -1774,6 +1814,15 @@ dependencies = [
"syn 2.0.119",
]
[[package]]
name = "des"
version = "0.8.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ffdd80ce8ce993de27e9f063a444a4d53ce8e8db4c1f00cc03af5ad5a9867a1e"
dependencies = [
"cipher",
]
[[package]]
name = "digest"
version = "0.10.7"
@ -3105,6 +3154,16 @@ dependencies = [
"serde_core",
]
[[package]]
name = "inout"
version = "0.1.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "879f10e63c20629ecabbb64a8010319738c66a5cd0c29b02d63d272b03751d01"
dependencies = [
"block-padding",
"generic-array",
]
[[package]]
name = "ipnet"
version = "2.12.0"
@ -3599,19 +3658,15 @@ dependencies = [
name = "litellm-cache-s3"
version = "0.1.0"
dependencies = [
"aws-credential-types",
"aws-sdk-s3",
"aws-smithy-runtime-api",
"aws-smithy-types",
"aws-types",
"futures-util",
"http 1.4.2",
"litellm-auth-aws",
"litellm-cache",
"litellm-cache-testing",
"litellm-error",
"litellm-http",
"reqwest",
"litellm-s3",
"rstest",
"serde_json",
"tokio",
@ -3663,8 +3718,7 @@ dependencies = [
"strum",
"tempfile",
"thiserror 2.0.19",
"tracing",
"tracing-subscriber",
"veil",
]
[[package]]
@ -3767,6 +3821,7 @@ dependencies = [
"http-body-util",
"litellm-circuit-breaker",
"litellm-config",
"litellm-core-utils",
"litellm-error",
"litellm-gateway-auth",
"litellm-gateway-forward",
@ -3777,11 +3832,15 @@ dependencies = [
"litellm-inference",
"litellm-llms",
"litellm-redis-client",
"litellm-redis-connect",
"litellm-redis-resolve",
"litellm-redis-testing",
"litellm-redis-types",
"litellm-secret-value",
"litellm-secrets",
"litellm-tracing",
"rstest",
"rustls 0.23.42",
"serde",
"serde_json",
"tempfile",
@ -3792,6 +3851,7 @@ dependencies = [
"tower-sessions-moka-store",
"tracing",
"uuid",
"veil",
]
[[package]]
@ -4283,12 +4343,14 @@ dependencies = [
"jsonschema",
"litellm-cost",
"litellm-error",
"litellm-tracing",
"macro_rules_attribute",
"rstest",
"schemars 1.2.2",
"serde",
"serde_json",
"serde_with",
"sha1 0.10.7",
"strum",
"thiserror 2.0.19",
"time",
@ -4298,9 +4360,13 @@ dependencies = [
name = "litellm-model-catalog-refresh"
version = "0.1.0"
dependencies = [
"aws-credential-types",
"backon",
"bytes",
"litellm-error",
"litellm-http",
"litellm-model-catalog",
"litellm-s3",
"litellm-tracing",
"reqwest",
"rstest",
"thiserror 2.0.19",
@ -4318,11 +4384,13 @@ dependencies = [
"criterion",
"fancy-regex 0.19.2",
"futures-util",
"indexmap 2.14.0",
"litellm-auth",
"litellm-auth-aws",
"litellm-cache",
"litellm-cache-response",
"litellm-callbacks-legacy-python",
"litellm-circuit-breaker",
"litellm-core-utils",
"litellm-error",
"litellm-host",
@ -4337,6 +4405,13 @@ dependencies = [
"litellm-llms",
"litellm-llms-types",
"litellm-model-catalog",
"litellm-model-catalog-refresh",
"litellm-python-compat",
"litellm-redis-client",
"litellm-redis-connect",
"litellm-redis-resolve",
"litellm-redis-store",
"litellm-redis-types",
"litellm-router-core",
"litellm-router-types",
"litellm-secret-value",
@ -4352,6 +4427,7 @@ dependencies = [
"pyo3-async-runtimes",
"reqwest",
"rstest",
"rustls 0.23.42",
"serde",
"serde_json",
"serde_with",
@ -4387,23 +4463,110 @@ version = "0.1.0"
dependencies = [
"futures-util",
"litellm-circuit-breaker",
"litellm-python-compat",
"litellm-redis-testing",
"litellm-redis-types",
"litellm-secret-value",
"pem",
"pkcs8",
"redis",
"rstest",
"rustls 0.23.42",
"rustls-native-certs",
"socket2 0.6.5",
"thiserror 2.0.19",
"tokio",
"tracing",
]
[[package]]
name = "litellm-redis-connect"
version = "0.1.0"
dependencies = [
"futures-util",
"litellm-auth-aws",
"litellm-auth-azure",
"litellm-auth-gcp",
"litellm-auth-types",
"litellm-circuit-breaker",
"litellm-core-utils",
"litellm-error",
"litellm-http",
"litellm-redis-client",
"litellm-redis-resolve",
"litellm-redis-testing",
"litellm-redis-types",
"litellm-secret-value",
"rstest",
"rustls 0.23.42",
"serde",
"serde_json",
"strum",
"thiserror 2.0.19",
"time",
"tokio",
"tracing",
"url",
"veil",
]
[[package]]
name = "litellm-redis-resolve"
version = "0.1.0"
dependencies = [
"litellm-python-compat",
"litellm-redis-types",
"litellm-secret-value",
"rstest",
"serde",
"serde_json",
"strum",
"thiserror 2.0.19",
"tracing",
"url",
]
[[package]]
name = "litellm-redis-store"
version = "0.1.0"
dependencies = [
"futures-util",
"litellm-cache",
"litellm-cache-memory",
"litellm-circuit-breaker",
"litellm-python-compat",
"litellm-redis-client",
"litellm-redis-testing",
"litellm-secret-value",
"litellm-tracing",
"rstest",
"strum",
"thiserror 2.0.19",
"tokio",
"tracing",
]
[[package]]
name = "litellm-redis-testing"
version = "0.1.0"
dependencies = [
"pkcs8",
"rcgen",
"tempfile",
"testcontainers-modules",
"thiserror 2.0.19",
"tokio",
]
[[package]]
name = "litellm-redis-types"
version = "0.1.0"
dependencies = [
"litellm-secret-value",
"rstest",
"url",
]
[[package]]
name = "litellm-router"
version = "0.1.0"
@ -4450,6 +4613,23 @@ dependencies = [
"serde_with",
]
[[package]]
name = "litellm-s3"
version = "0.1.0"
dependencies = [
"aws-config",
"aws-credential-types",
"aws-sdk-s3",
"aws-smithy-runtime-api",
"aws-smithy-types",
"aws-types",
"http 1.4.2",
"litellm-http",
"reqwest",
"rstest",
"thiserror 2.0.19",
]
[[package]]
name = "litellm-secret-value"
version = "0.1.0"
@ -5286,6 +5466,16 @@ version = "0.2.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2ee67f1008b1ba2321834326597b8e186293b049a023cdef258527550b9935b4"
[[package]]
name = "pbkdf2"
version = "0.12.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f8ed6a7761f76e3b9f92dfb0a60a6a6477c61024b775147ff0973a02653abaf2"
dependencies = [
"digest 0.10.7",
"hmac 0.12.1",
]
[[package]]
name = "pem"
version = "4.0.0"
@ -5396,6 +5586,22 @@ dependencies = [
"spki",
]
[[package]]
name = "pkcs5"
version = "0.7.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e847e2c91a18bfa887dd028ec33f2fe6f25db77db3619024764914affe8b69a6"
dependencies = [
"aes",
"cbc",
"der",
"des",
"pbkdf2",
"scrypt",
"sha2 0.10.9",
"spki",
]
[[package]]
name = "pkcs8"
version = "0.10.2"
@ -5403,6 +5609,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f950b2377845cebe5cf8b5165cb3cc1a5e0fa5cfa3e1f7f55707d8fd82e0a7b7"
dependencies = [
"der",
"pkcs5",
"rand_core 0.6.4",
"spki",
]
@ -5590,6 +5798,7 @@ version = "0.29.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4688ddedf473e32662b9b067670129a8afb8c18e351482c70d62ba4a88171e8b"
dependencies = [
"indexmap 2.14.0",
"libc",
"once_cell",
"portable-atomic",
@ -6411,6 +6620,15 @@ version = "1.0.23"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f"
[[package]]
name = "salsa20"
version = "0.10.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "97a22f5af31f73a954c10289c93e8a50cc23d971e80ee446f1f6f7137a088213"
dependencies = [
"cipher",
]
[[package]]
name = "same-file"
version = "1.0.6"
@ -6474,6 +6692,17 @@ version = "1.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49"
[[package]]
name = "scrypt"
version = "0.11.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0516a385866c09368f0b5bcd1caff3366aace790fcd46e2bb032697bb172fd1f"
dependencies = [
"pbkdf2",
"salsa20",
"sha2 0.10.9",
]
[[package]]
name = "sct"
version = "0.7.1"

7
litellm-rust/Cargo.toml
View file

@ -60,6 +60,7 @@ litellm-cache = { path = "crates/cache" }
litellm-cache-azure-blob = { path = "crates/cache-azure-blob" }
litellm-cache-memory = { path = "crates/cache-memory" }
litellm-cache-s3 = { path = "crates/cache-s3" }
litellm-s3 = { path = "crates/s3" }
litellm-cache-gcs = { path = "crates/cache-gcs" }
litellm-cache-disk = { path = "crates/cache-disk" }
litellm-cache-response = { path = "crates/cache-response" }
@ -73,7 +74,11 @@ litellm-host-python = { path = "crates/host-python" }
litellm-python-compat = { path = "crates/python-compat" }
litellm-circuit-breaker = { path = "crates/circuit-breaker" }
litellm-redis-client = { path = "crates/redis-client" }
litellm-redis-connect = { path = "crates/redis-connect" }
litellm-redis-store = { path = "crates/redis-store" }
litellm-redis-testing = { path = "crates/redis-testing" }
litellm-redis-resolve = { path = "crates/redis-resolve" }
litellm-redis-types = { path = "crates/redis-types" }
litellm-cost = { path = "crates/cost" }
litellm-model-catalog = { path = "crates/model-catalog" }
litellm-model-catalog-refresh = { path = "crates/model-catalog-refresh" }
@ -86,6 +91,8 @@ axum = { version = "0.8.9", default-features = false, features = ["http1", "toki
axum-login = "0.18.0"
tower-sessions = { version = "0.14.0", features = ["memory-store"] }
bytes = "1"
backon = { version = "1.6.0", default-features = false, features = ["std", "tokio-sleep"] }
aws-sdk-s3 = { version = "1.146.1", default-features = false, features = ["rt-tokio"] }
http = "1"
google-cloud-auth = { version = "1.16.0", default-features = false }
jsonwebtoken = { version = "11.1.0", default-features = false }

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

@ -1,9 +1,11 @@
mod aws;
pub mod constants;
mod error;
mod provider;
mod signer;
pub use aws::*;
pub use aws_credential_types::Credentials;
pub use error::{CredentialFlow, Error, SigningStage};
pub use provider::AwsCredentialsProvider;
pub use signer::SigV4Signer;

135
litellm-rust/crates/auth-aws/src/provider.rs Normal file
View file

@ -0,0 +1,135 @@
use std::sync::Arc;
use aws_credential_types::Credentials;
use aws_credential_types::provider::{ProvideCredentials, error::CredentialsError, future};
use crate::{AwsAuthConfig, AwsAuthService};
type Environment = Arc<dyn Fn(&str) -> Option<String> + Send + Sync>;
#[derive(Clone)]
pub struct AwsCredentialsProvider {
auth: AwsAuthService,
config: Box<AwsAuthConfig>,
environment: Environment,
}
impl AwsCredentialsProvider {
/// Explicit keys in `config` as given, otherwise LiteLLM's `AWS_*` settings read through `environment`
pub fn configured(
auth: AwsAuthService,
config: AwsAuthConfig,
environment: impl Fn(&str) -> Option<String> + Send + Sync + 'static,
) -> Self {
Self {
auth,
config: Box::new(config),
environment: Arc::new(environment),
}
}
async fn credentials(&self) -> Result<Credentials, crate::Error> {
match (&self.config.access_key_id, &self.config.secret_access_key) {
(Some(access_key_id), Some(secret_access_key)) => Ok(Credentials::new(
access_key_id.expose(),
secret_access_key.expose(),
self.config
.session_token
.as_ref()
.map(|token| token.expose().to_owned()),
None,
"litellm-configured",
)),
_ => {
self.auth
.resolve_credentials(
AwsAuthConfig::clone(&self.config),
self.environment.as_ref(),
)
.await
}
}
}
}
impl ProvideCredentials for AwsCredentialsProvider {
fn provide_credentials<'a>(&'a self) -> future::ProvideCredentials<'a>
where
Self: 'a,
{
future::ProvideCredentials::new(async {
self.credentials()
.await
.map_err(CredentialsError::provider_error)
})
}
}
impl std::fmt::Debug for AwsCredentialsProvider {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("AwsCredentialsProvider")
.finish_non_exhaustive()
}
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use super::*;
fn explicit(session_token: Option<&str>) -> AwsAuthConfig {
AwsAuthConfig {
access_key_id: Some("key".into()),
secret_access_key: Some("secret".into()),
session_token: session_token.map(Into::into),
region_name: Some("us-east-1".to_string()),
..Default::default()
}
}
fn ambient_token(name: &str) -> Option<String> {
(name == "AWS_SESSION_TOKEN").then(|| "ambient".to_string())
}
fn environment_keys(name: &str) -> Option<String> {
match name {
"AWS_ACCESS_KEY_ID" => Some("env-key".to_string()),
"AWS_SECRET_ACCESS_KEY" => Some("env-secret".to_string()),
"AWS_SESSION_TOKEN" => Some("env-token".to_string()),
_ => None,
}
}
#[rstest]
#[case::explicit_keys_ignore_an_ambient_session_token(
explicit(None), ambient_token, ("key", "secret", None)
)]
#[case::explicit_keys_keep_their_session_token(
explicit(Some("t")), ambient_token, ("key", "secret", Some("t"))
)]
#[case::environment_keys_resolve_with_their_session_token(
AwsAuthConfig::default(), environment_keys, ("env-key", "env-secret", Some("env-token"))
)]
#[tokio::test]
async fn configured_credentials_read_the_environment_they_are_given(
#[case] config: AwsAuthConfig,
#[case] environment: fn(&str) -> Option<String>,
#[case] expected: (&str, &str, Option<&str>),
) {
let credentials =
AwsCredentialsProvider::configured(AwsAuthService::default(), config, environment)
.provide_credentials()
.await
.unwrap();
assert_eq!(
(
credentials.access_key_id(),
credentials.secret_access_key(),
credentials.session_token(),
),
expected
);
}
}

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

@ -2,7 +2,10 @@ use std::{
cmp::Reverse,
collections::{BinaryHeap, HashMap, HashSet},
hash::Hash,
sync::{Arc, Mutex, MutexGuard},
sync::{
Arc, Mutex, MutexGuard,
atomic::{AtomicUsize, Ordering},
},
time::{Duration, SystemTime, UNIX_EPOCH},
};
@ -31,7 +34,7 @@ struct CacheState<V> {
pub struct InMemoryCache<V: Clone> {
state: Mutex<CacheState<V>>,
max_size_in_memory: usize,
max_size_in_memory: AtomicUsize,
default_ttl: Duration,
max_entry_bytes: Option<usize>,
measure_value: Option<ValueMeasure<V>>,
@ -74,7 +77,9 @@ impl<V: Clone> InMemoryCache<V> {
expirations: HashMap::new(),
expiration_heap: BinaryHeap::new(),
}),
max_size_in_memory: max_size_in_memory.unwrap_or(DEFAULT_MAX_SIZE_IN_MEMORY),
max_size_in_memory: AtomicUsize::new(
max_size_in_memory.unwrap_or(DEFAULT_MAX_SIZE_IN_MEMORY),
),
default_ttl: default_ttl
.filter(|ttl| !ttl.is_zero())
.unwrap_or(DEFAULT_TTL),
@ -94,7 +99,7 @@ impl<V: Clone> InMemoryCache<V> {
value: V,
ttl: Option<Duration>,
) -> Result<CacheWrite, Error> {
if self.max_size_in_memory == 0 {
if self.max_size_in_memory() == 0 {
return Ok(CacheWrite::Disabled);
}
let now = (self.now)();
@ -129,7 +134,7 @@ impl<V: Clone> InMemoryCache<V> {
pub fn evict_cache(&self) -> Result<(), Error> {
let now = (self.now)();
let mut state = self.lock()?;
Self::evict(&mut state, self.max_size_in_memory, now, None);
Self::evict(&mut state, self.max_size_in_memory(), now, None);
Ok(())
}
@ -170,7 +175,16 @@ impl<V: Clone> InMemoryCache<V> {
}
pub fn max_size_in_memory(&self) -> usize {
self.max_size_in_memory
self.max_size_in_memory.load(Ordering::Relaxed)
}
/// `DualCache.update_in_memory_max_size`: `None` restores the default. Entries already held stay
/// until the next write evicts down to the new size.
pub fn update_max_size_in_memory(&self, max_size_in_memory: Option<usize>) {
self.max_size_in_memory.store(
max_size_in_memory.unwrap_or(DEFAULT_MAX_SIZE_IN_MEMORY),
Ordering::Relaxed,
);
}
pub fn max_entry_bytes(&self) -> Option<usize> {
@ -200,6 +214,23 @@ impl<V: Clone> InMemoryCache<V> {
.collect())
}
/// Every live entry, for callers that inspect the whole tier.
pub fn entries(&self) -> Result<Vec<(String, V)>, Error> {
let now = (self.now)();
let state = self.lock()?;
Ok(state
.values
.iter()
.filter(|(key, _)| {
state
.expirations
.get(*key)
.is_none_or(|expiration| *expiration >= now)
})
.map(|(key, value)| (key.clone(), value.clone()))
.collect())
}
pub fn delete_cache(&self, key: &str) -> Result<(), Error> {
let mut state = self.lock()?;
Self::remove(&mut state, key);
@ -276,7 +307,7 @@ impl<V: Clone> InMemoryCache<V> {
ttl: Option<Duration>,
now: Duration,
) -> Result<CacheWrite, Error> {
Self::evict(state, self.max_size_in_memory, now, None);
Self::evict(state, self.max_size_in_memory(), now, None);
if !self.check_value_size(&value)? {
return Ok(CacheWrite::TooLarge);
}
@ -300,12 +331,12 @@ where
eligible: &[V],
context: ExactCacheContext,
) -> Result<V, Error> {
if self.max_size_in_memory == 0 {
if self.max_size_in_memory() == 0 {
return Ok(candidate);
}
let now = (self.now)();
let mut state = self.lock()?;
Self::evict(&mut state, self.max_size_in_memory, now, Some(key));
Self::evict(&mut state, self.max_size_in_memory(), now, Some(key));
let existing = state
.values
.get(key)
@ -335,7 +366,7 @@ impl CounterCache for InMemoryCache<f64> {
amount: f64,
context: ExactCacheContext,
) -> Result<f64, Error> {
if self.max_size_in_memory == 0 {
if self.max_size_in_memory() == 0 {
return Ok(amount);
}
let now = (self.now)();
@ -408,7 +439,7 @@ where
values: Vec<Self::SetValue>,
ttl: Option<Duration>,
) -> Result<Self::SetResult, Error> {
if self.max_size_in_memory == 0 {
if self.max_size_in_memory() == 0 {
return Ok(values);
}
let now = (self.now)();

10
litellm-rust/crates/cache-s3/Cargo.toml
View file

@ -9,14 +9,10 @@ repository.workspace = true
litellm-http.workspace = true
litellm-cache.workspace = true
litellm-auth-aws.workspace = true
aws-sdk-s3 = { version = "1.146.1", default-features = false, features = ["rustls", "rt-tokio"] }
aws-credential-types = "1.3.0"
aws-smithy-runtime-api = { version = "1.16.2", features = ["client", "http-1x"] }
aws-smithy-types = { version = "1.6.0", features = ["http-body-1-x"] }
aws-types = "1.6.0"
litellm-s3.workspace = true
aws-sdk-s3.workspace = true
aws-smithy-types = "1.6.0"
futures-util.workspace = true
http.workspace = true
reqwest.workspace = true
tokio.workspace = true
[dev-dependencies]

61
litellm-rust/crates/cache-s3/src/auth.rs
View file

@ -1,61 +0,0 @@
use aws_credential_types::{
Credentials as AwsCredentials,
provider::{ProvideCredentials, error::CredentialsError, future},
};
use litellm_auth_aws::{AwsAuthConfig, AwsAuthService};
#[derive(Clone)]
pub struct S3Credentials {
auth: AwsAuthService,
config: AwsAuthConfig,
env: fn(&str) -> Option<String>,
}
impl S3Credentials {
pub fn new(config: AwsAuthConfig) -> Self {
Self::with_env(config, |name| std::env::var(name).ok())
}
pub fn with_env(config: AwsAuthConfig, env: fn(&str) -> Option<String>) -> Self {
Self {
auth: AwsAuthService::default(),
config,
env,
}
}
}
impl ProvideCredentials for S3Credentials {
fn provide_credentials<'a>(&'a self) -> future::ProvideCredentials<'a>
where
Self: 'a,
{
future::ProvideCredentials::new(async {
if let (Some(access_key_id), Some(secret_access_key)) = (
self.config.access_key_id.clone(),
self.config.secret_access_key.clone(),
) {
return Ok(AwsCredentials::new(
access_key_id.expose(),
secret_access_key.expose(),
self.config
.session_token
.as_ref()
.map(|value| value.expose().to_owned()),
None,
"litellm-s3-cache",
));
}
self.auth
.resolve_credentials(self.config.clone(), &self.env)
.await
.map_err(CredentialsError::provider_error)
})
}
}
impl std::fmt::Debug for S3Credentials {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("S3Credentials").finish_non_exhaustive()
}
}

27
litellm-rust/crates/cache-s3/src/cache.rs
View file

@ -5,20 +5,18 @@ use std::{
};
use aws_sdk_s3::{
config::{BehaviorVersion, Region, RequestChecksumCalculation, ResponseChecksumValidation},
config::{RequestChecksumCalculation, ResponseChecksumValidation},
error::SdkError,
primitives::ByteStream,
};
use aws_smithy_types::{DateTime, date_time::Format};
use futures_util::future::try_join_all;
use litellm_auth_aws::AwsAuthConfig;
use litellm_auth_aws::AwsCredentialsProvider;
use litellm_cache::{
BaseCache, BatchCache, CacheCodec, DisconnectCache, Error, ExactCacheContext, FlushCache,
};
use tokio::runtime::Handle;
use crate::{auth::S3Credentials, transport::ReqwestHttpClient};
pub struct S3Endpoint {
pub url: String,
}
@ -28,7 +26,7 @@ pub struct S3CacheConfig {
pub key_prefix: String,
pub region: String,
pub endpoint: Option<S3Endpoint>,
pub auth: AwsAuthConfig,
pub credentials: AwsCredentialsProvider,
}
pub struct S3Cache<C: CacheCodec> {
@ -49,17 +47,14 @@ impl<C: CacheCodec> S3Cache<C> {
runtime: Handle,
) -> Self {
let endpoint_url: Option<String> = config.endpoint.map(|endpoint| endpoint.url);
let base = aws_sdk_s3::Config::builder()
.behavior_version(BehaviorVersion::latest())
.region(Region::new(config.region.clone()))
.http_client(ReqwestHttpClient(http))
.credentials_provider(S3Credentials::new(config.auth))
.request_checksum_calculation(RequestChecksumCalculation::WhenRequired)
.response_checksum_validation(ResponseChecksumValidation::WhenRequired);
let builder = match &endpoint_url {
Some(url) => base.endpoint_url(url).force_path_style(true),
None => base,
};
let builder = litellm_s3::config(
config.region.clone(),
endpoint_url.as_deref(),
http,
config.credentials,
)
.request_checksum_calculation(RequestChecksumCalculation::WhenRequired)
.response_checksum_validation(ResponseChecksumValidation::WhenRequired);
Self {
client: aws_sdk_s3::Client::from_conf(builder.build()),
codec,

3
litellm-rust/crates/cache-s3/src/lib.rs
View file

@ -1,6 +1,3 @@
mod auth;
mod cache;
mod transport;
pub use auth::S3Credentials;
pub use cache::{S3Cache, S3CacheConfig, S3Endpoint};

57
litellm-rust/crates/cache-s3/tests/auth.rs
View file

@ -1,57 +0,0 @@
use aws_credential_types::provider::ProvideCredentials;
use litellm_auth_aws::AwsAuthConfig;
use litellm_cache_s3::S3Credentials;
use rstest::rstest;
fn explicit(session_token: Option<&str>) -> AwsAuthConfig {
AwsAuthConfig {
access_key_id: Some("key".into()),
secret_access_key: Some("secret".into()),
session_token: session_token.map(Into::into),
region_name: Some("us-east-1".to_string()),
..Default::default()
}
}
fn ambient_token(name: &str) -> Option<String> {
(name == "AWS_SESSION_TOKEN").then(|| "ambient".to_string())
}
fn environment_keys(name: &str) -> Option<String> {
match name {
"AWS_ACCESS_KEY_ID" => Some("env-key".to_string()),
"AWS_SECRET_ACCESS_KEY" => Some("env-secret".to_string()),
"AWS_SESSION_TOKEN" => Some("env-token".to_string()),
_ => None,
}
}
#[rstest]
#[case::explicit_keys_ignore_an_ambient_session_token(
explicit(None), ambient_token, ("key", "secret", None)
)]
#[case::explicit_keys_keep_their_session_token(
explicit(Some("t")), ambient_token, ("key", "secret", Some("t"))
)]
#[case::environment_keys_resolve_with_their_session_token(
AwsAuthConfig::default(), environment_keys, ("env-key", "env-secret", Some("env-token"))
)]
#[tokio::test]
async fn credentials_resolve(
#[case] config: AwsAuthConfig,
#[case] env: fn(&str) -> Option<String>,
#[case] expected: (&str, &str, Option<&str>),
) {
let credentials = S3Credentials::with_env(config, env)
.provide_credentials()
.await
.unwrap();
assert_eq!(
(
credentials.access_key_id(),
credentials.secret_access_key(),
credentials.session_token(),
),
expected
);
}

18
litellm-rust/crates/cache-s3/tests/support/mod.rs
View file

@ -5,7 +5,7 @@ use std::{
sync::{Arc, Mutex},
};
use litellm_auth_aws::AwsAuthConfig;
use litellm_auth_aws::{AwsAuthConfig, AwsAuthService, AwsCredentialsProvider};
use litellm_cache::JsonCodec;
use litellm_cache_s3::{S3Cache, S3CacheConfig, S3Endpoint};
use serde_json::Value;
@ -22,12 +22,16 @@ pub fn config(endpoint: &str) -> S3CacheConfig {
endpoint: Some(S3Endpoint {
url: endpoint.to_string(),
}),
auth: AwsAuthConfig {
access_key_id: Some("key".into()),
secret_access_key: Some("secret".into()),
region_name: Some("us-east-1".to_string()),
..Default::default()
},
credentials: AwsCredentialsProvider::configured(
AwsAuthService::default(),
AwsAuthConfig {
access_key_id: Some("key".into()),
secret_access_key: Some("secret".into()),
region_name: Some("us-east-1".to_string()),
..Default::default()
},
|_| None,
),
}
}

18
litellm-rust/crates/callbacks-legacy-python/python_contract.json
View file

@ -27,6 +27,10 @@
"api_key",
"additional_args"
],
"messages_input": [
"body",
"kwargs"
],
"post_call": [
"logger",
"original_response",
@ -47,9 +51,18 @@
"defers_async_logging": [
"logger"
],
"claim_async_success_log": [
"logger"
],
"mark_async_success_enqueued": [
"logger"
],
"defer_success": [
"logger",
"pending"
"pending",
"response",
"start",
"end"
],
"sync_success_for_async_call": [
"logger",
@ -117,7 +130,8 @@
"chunks",
"start",
"end",
"first_chunk"
"first_chunk",
"stream_teardown"
],
"stream_failure": [
"logger",

346
litellm-rust/crates/callbacks-legacy-python/src/adapter.rs
View file

@ -14,7 +14,7 @@ use pyo3::{
exceptions::{PyBaseException, PyException},
gc::{PyTraverseError, PyVisit},
prelude::*,
types::{PyDateTime, PyDict, PyList},
types::{PyDateTime, PyDict, PyList, PyString},
};
use serde_json::Value;
@ -22,7 +22,7 @@ use crate::{
DeploymentHooks, LegacyCallbacks, PublicCall, PythonLogger,
deferred::{PendingLogging, PendingSuccess},
finalize, is_internal_call,
python::Streaming,
python::{Logging, Streaming},
setup,
};
@ -137,6 +137,40 @@ impl LegacyLogging {
Some(params)
}
// Compared as Python compares them, so a caller's `0` is kept over a typed float's `0.0`.
fn caller_value<'py>(
&self,
py: Python<'py>,
name: &str,
sent: &Value,
) -> PyResult<Option<Bound<'py, PyAny>>> {
match self.call.lookup(py, name)? {
Some(value) if from_py::<Value>(&value).is_ok() && value.eq(to_py(py, sent)?)? => {
Ok(Some(value))
}
_ => Ok(None),
}
}
// Python's Messages handler logs the caller's own optional params, not their typed copies.
fn as_called(&self, py: Python<'_>, context: &RequestContext) -> PyResult<RequestContext> {
let mut context = context.clone();
if self.operation != LoggingOperation::Messages {
return Ok(context);
}
for (name, sent) in context
.optional_params
.as_object_mut()
.into_iter()
.flatten()
{
if let Some(value) = self.caller_value(py, name, sent)? {
*sent = from_py(&value)?;
}
}
Ok(context)
}
/// Python times the provider exchange inside the async `AsyncHTTPHandler.post`; its sync
/// client is untimed.
fn close_api_window(&mut self, py: Python<'_>, end: f64) -> PyResult<()> {
@ -158,13 +192,20 @@ impl LegacyLogging {
}
}
fn input_description(&self) -> &'static str {
match self.operation {
fn logged_input<'py>(
&self,
py: Python<'py>,
body: &Bound<'py, PyDict>,
) -> PyResult<Bound<'py, PyAny>> {
let description = match self.operation {
LoggingOperation::Completion => "Chat completions",
LoggingOperation::Responses => "Responses",
LoggingOperation::Messages => "Messages",
LoggingOperation::Messages => {
return Logging::MessagesInput.call(py, (body, self.call.kwargs().bind(py)));
}
LoggingOperation::Ocr => "OCR document processing",
}
};
Ok(PyString::new(py, description).into_any())
}
/// The key `pre_call` logs. Python's Messages handler logs an empty string, never the key.
@ -251,30 +292,49 @@ impl LegacyLogging {
if !self.asynchronous {
return pending().sync(py);
}
if !self.internal
&& self
if self.internal {
return Ok(());
}
let defers = logger.defers_async_logging(py);
if logger.claim_async_success_log(py)? {
let enqueues = self
.call
.kwargs()
.bind(py)
.get_item("fallbacks")?
.is_none_or(|value| value.is_none())
{
if logger.defers_async_logging(py) {
.is_none_or(|value| value.is_none());
if defers {
let pending = Py::new(
py,
PendingLogging {
pending: Some(pending()),
enqueues,
},
)?;
logger.defer_success(py, pending.bind(py).as_any())?;
} else {
return logger.defer_success(
py,
pending.bind(py).as_any(),
&self.response,
&self.start,
&self.end,
);
}
if enqueues {
pending().asynchronous(py)?;
}
}
if defers {
return Ok(());
}
logger.sync_success_for_async_call(py, &self.response, &self.start, &self.end)
}
fn stream_success(&self, py: Python<'_>, stream: &DeliveredStream) -> PyResult<()> {
fn stream_success(
&self,
py: Python<'_>,
stream: &DeliveredStream,
teardown: bool,
) -> PyResult<()> {
let logger = self.logger()?;
let billing = self.stream_billing().ok_or(HostError::MissingState)?;
let billed = Streaming::Success.call(
@ -291,6 +351,7 @@ impl LegacyLogging {
&self.start,
&self.end,
&stream.first_chunk,
teardown,
),
);
match billed {
@ -442,6 +503,7 @@ impl LegacyLogging {
wire: Box<WireRequest>,
context: &RequestContext,
) -> PyResult<HookStep<Self, Box<WireRequest>>> {
let context = &self.as_called(py, context)?;
let logger = self.logger()?;
logger.update_from_kwargs(
py,
@ -454,12 +516,17 @@ impl LegacyLogging {
.into_bound(py)
.cast_into::<PyDict>()?;
for (name, sent) in wire.body.as_object().into_iter().flatten() {
if let Some(value) = self.call.lookup(py, name)?
&& from_py::<Value>(&value).is_ok_and(|caller| caller == *sent)
{
if let Some(value) = self.caller_value(py, name, sent)? {
body.set_item(name, value)?;
}
}
if self.operation == LoggingOperation::Messages {
// Python's Messages handler merges its transformed request body into the call details.
logger
.object(py)
.getattr("model_call_details")?
.call_method1("update", (&body,))?;
}
let headers = PyDict::new(py);
for (name, value) in &wire.headers {
headers.set_item(name, value)?;
@ -469,9 +536,10 @@ impl LegacyLogging {
headers: headers.clone().unbind(),
context: context.clone(),
});
let input = self.logged_input(py, &body)?;
self.logger()?.pre_call(
py,
self.input_description(),
&input,
self.logged_api_key(context),
&body,
&headers,
@ -525,7 +593,11 @@ impl LegacyLogging {
.getattr("litellm_params")?
.cast_into::<PyDict>()?
.copy()?;
params.set_item("custom_llm_provider", &facts.provider.provider)?;
if self.operation != LoggingOperation::Messages {
// Python's Messages handler hands the provider to the logger beside its params, never
// in them, and the proxy's logger carries over to the next fallback hop.
params.set_item("custom_llm_provider", &facts.provider.provider)?;
}
crate::python::Logging::Update.call(
py,
(
@ -548,8 +620,17 @@ impl LegacyLogging {
ResultSource::Provider => None,
ResultSource::Cache { key } => Some(key.clone()),
};
details.set_item("cache_hit", self.cache_key.is_some())?;
if self.cache_key.is_some() {
details.set_item("cache_hit", true)?;
} else {
details.call_method1("pop", ("cache_hit", py.None()))?;
}
details.set_item("cache_key", self.cache_key.as_deref())?;
if let (LoggingOperation::Messages, ResultSource::Provider, Some(request)) =
(self.operation, &facts.source, &self.request)
{
details.call_method1("update", (request.body.bind(py),))?;
}
Ok(HookStep::Ready(()))
}
@ -592,7 +673,7 @@ impl LegacyLogging {
self.end = Some(datetime(py, timing.end_time)?);
self.response = Some(response.clone_ref(py));
match &self.stream {
Some(stream) => self.stream_success(py, stream)?,
Some(stream) => self.stream_success(py, stream, false)?,
None => self.dispatch_success(py)?,
}
Ok(HookStep::Ready(()))
@ -662,8 +743,28 @@ impl LegacyLogging {
}
}
impl LegacyLogging {
// A caller that goes away mid-stream cancels the call before it ends: Python's Messages
// stream iterator still bills the chunks it read, straight away rather than deferred.
fn bill_abandoned_stream(&mut self, py: Python<'_>) -> PyResult<()> {
if !self.asynchronous || self.end.is_some() {
return Ok(());
}
let Some(stream) = self.stream.take() else {
return Ok(());
};
self.end = Some(datetime(py, epoch_seconds())?);
self.stream_success(py, &stream, true)
}
}
impl PythonOwned for LegacyLogging {
fn close(&mut self, py: Python<'_>) {
if self.logger.is_some()
&& let Err(error) = self.bill_abandoned_stream(py)
{
error.write_unraisable(py, None);
}
if let Some(logger) = self.logger.take()
&& let Err(error) = logger.restore_context(py)
{
@ -1114,10 +1215,10 @@ logger._provider_reported_cost = None
};
let previous = litellm_host::interceptors::ExecutionFacts {
provider: facts.provider.clone(),
cost: litellm_host::interceptors::Cost::Reported {
amount: serde_json::Number::from_f64(9.0).unwrap(),
cost: litellm_host::interceptors::Cost::CacheHit,
source: litellm_host::interceptors::ResultSource::Cache {
key: "earlier-hit".into(),
},
source: litellm_host::interceptors::ResultSource::Provider,
};
logging.result_ready(py, &previous).unwrap();
local(&locals, "logger")
@ -1135,13 +1236,59 @@ logger._provider_reported_cost = None
c"
assert logger.model_call_details.get('response_cost') is None, logger.model_call_details
assert logger._provider_reported_cost == reported, logger._provider_reported_cost
assert logger.model_call_details['cache_hit'] is cache_hit, logger.model_call_details
assert logger.model_call_details.get('cache_hit') is (True if cache_hit else None), logger.model_call_details
assert logger.update['custom_llm_provider'] == 'edenai', logger.update
",
);
});
}
#[rstest]
#[case::messages(crate::LoggingOperation::Messages, false)]
#[case::ocr(crate::LoggingOperation::Ocr, true)]
fn result_ready_puts_the_provider_in_the_logged_params_except_for_messages(
#[case] operation: crate::LoggingOperation,
#[case] in_params: bool,
) {
Python::initialize();
Python::attach(|py| {
let locals = namespace(py, PAYLOAD_LOGGER);
run(
py,
&locals,
c"
logger.litellm_params = {}
logger.optional_params = {}
logger.model_call_details = {}
logger._provider_reported_cost = None
",
);
let mut logging = LegacyLogging {
logger: Some(PythonLogger::new(local(&locals, "logger").unbind())),
operation,
..legacy_call(py, &locals, false)
};
let facts = litellm_host::interceptors::ExecutionFacts {
provider: litellm_host::interceptors::ProviderIdentity {
model: "native-test".into(),
provider: "anthropic".into(),
},
cost: litellm_host::interceptors::Cost::Deferred,
source: litellm_host::interceptors::ResultSource::Provider,
};
logging.result_ready(py, &facts).unwrap();
locals.set_item("in_params", in_params).unwrap();
run(
py,
&locals,
c"
assert logger.update['custom_llm_provider'] == 'anthropic', logger.update
assert ('custom_llm_provider' in logger.update['litellm_params']) is in_params, logger.update
",
);
});
}
const DOCUMENT: &str = "data:application/pdf;base64,YWJj";
const EDITED: &str = "data:application/pdf;base64,ZWRpdGVk";
@ -1276,7 +1423,10 @@ assert logger.update['custom_llm_provider'] == 'edenai', logger.update
#[rstest]
#[case::completion(crate::LoggingOperation::Completion, "Chat completions")]
#[case::responses(crate::LoggingOperation::Responses, "Responses")]
#[case::messages(crate::LoggingOperation::Messages, "Messages")]
#[case::messages(
crate::LoggingOperation::Messages,
r#"[{'role': 'user', 'content': '{"pages": [1]}'}]"#
)]
#[case::ocr(crate::LoggingOperation::Ocr, "OCR document processing")]
fn prepared_arguments_replace_the_legacy_view_without_losing_callback_aliases(
#[case] operation: crate::LoggingOperation,
@ -1317,7 +1467,7 @@ prepared = {'pages': replacement}
&locals,
c"
assert logger.pre['complete_input_dict']['pages'] is replacement
assert logger.pre_input == description
assert str(logger.pre_input) == description
assert original == [0]
",
);
@ -1873,10 +2023,15 @@ mod terminal_tests {
#[case::async_listened(
true,
c"",
&["async_success_handler", "enqueued", "sync_success_for_async_call"]
&["async_success_handler", "mark_async_success_enqueued", "enqueued", "sync_success_for_async_call"]
)]
#[case::async_deferred(true, c"logger._defer_async_logging = True", &["sync_success_for_async_call"])]
#[case::async_deferred(true, c"logger._defer_async_logging = True", &[])]
#[case::async_with_fallbacks(true, c"kwargs = {'fallbacks': ['other']}", &["sync_success_for_async_call"])]
#[case::async_deferred_with_fallbacks(
true,
c"logger._defer_async_logging = True\nkwargs = {'fallbacks': ['other']}",
&[]
)]
fn success_reaches_the_logging_handlers(
#[case] asynchronous: bool,
#[case] script: &CStr,
@ -1902,12 +2057,88 @@ mod terminal_tests {
&locals,
c"
assert all(value is response for name, value in logger.calls if name.endswith('_handler'))
assert hasattr(logger, '_native_pending_logging') == getattr(logger, '_defer_async_logging', False)
assert hasattr(logger, 'enqueue_deferred_logging') == getattr(logger, '_defer_async_logging', False)
assert getattr(logger, 'deferred_success_result', (response,))[0] is response
",
);
});
}
#[rstest]
#[case::accepted(c"", &["sync_success_for_async_call", "async_success_handler", "mark_async_success_enqueued", "enqueued"])]
#[case::accepted_with_fallbacks(c"kwargs = {'fallbacks': ['other']}", &["sync_success_for_async_call"])]
fn a_deferred_success_runs_sync_callbacks_only_once_released(
#[case] script: &CStr,
#[case] expected: &[&str],
) {
Python::initialize();
Python::attach(|py| {
let locals = namespace(
py,
c"response = object()\nlogger._defer_async_logging = True",
);
run(py, &locals, script);
let mut logging = logged(py, &locals, true);
succeed(py, &locals, &mut logging);
run(
py,
&locals,
c"
assert logger.names() == [], logger.calls
logger.enqueue_deferred_logging()
logger.enqueue_deferred_logging()
",
);
let names: Vec<String> = local(&locals, "logger")
.call_method0("names")
.unwrap()
.extract()
.unwrap();
assert_eq!(names, expected);
});
}
#[rstest]
#[case::deferred(
c"logger._defer_async_logging = True",
c"logger.enqueue_deferred_logging()",
&["sync_success_for_async_call", "async_success_handler", "mark_async_success_enqueued", "enqueued"]
)]
#[case::deferred_and_blocked(c"logger._defer_async_logging = True", c"", &[])]
#[case::immediate(
c"",
c"",
&["async_success_handler", "mark_async_success_enqueued", "enqueued", "sync_success_for_async_call", "sync_success_for_async_call"]
)]
fn a_fallback_attempt_on_the_same_logger_leaves_the_first_success_claimed(
#[case] script: &CStr,
#[case] release: &CStr,
#[case] expected: &[&str],
) {
Python::initialize();
Python::attach(|py| {
let locals = namespace(py, c"refused = object()\nresponse = refused");
run(py, &locals, script);
let mut first = logged(py, &locals, true);
succeed(py, &locals, &mut first);
run(py, &locals, c"response = object()");
let mut fallback = logged(py, &locals, true);
succeed(py, &locals, &mut fallback);
run(py, &locals, release);
let names: Vec<String> = local(&locals, "logger")
.call_method0("names")
.unwrap()
.extract()
.unwrap();
assert_eq!(names, expected);
run(
py,
&locals,
c"assert all(value is refused for name, value in logger.calls if name == 'async_success_handler')",
);
});
}
#[rstest]
fn dropped_observations_preserve_deferred_success_and_response_identity() {
Python::initialize();
@ -1933,10 +2164,10 @@ assert hasattr(logger, '_native_pending_logging') == getattr(logger, '_defer_asy
));
assert_eq!(sender.dropped_events(), 1);
run(py, &locals, c"
assert logger.names() == ['sync_success_for_async_call'], logger.calls
logger._native_pending_logging.release(True)
logger._native_pending_logging.release(True)
assert logger.names() == ['sync_success_for_async_call', 'async_success_handler', 'enqueued'], logger.calls
assert logger.names() == [], logger.calls
logger.enqueue_deferred_logging()
logger.enqueue_deferred_logging()
assert logger.names() == ['sync_success_for_async_call', 'async_success_handler', 'mark_async_success_enqueued', 'enqueued'], logger.calls
assert logger.calls[0][1] is response
assert logger.calls[1][1] is response
");
@ -1981,6 +2212,45 @@ assert chunks[1] is last
});
}
#[rstest]
#[case::cancelled_mid_stream(None, &["stamp_stream_head", "stream_opened", "stream_success", "restore"])]
#[case::after_success(Some(true), &["stamp_stream_head", "stream_opened", "stream_success", "restore"])]
#[case::after_failure(Some(false), &["stamp_stream_head", "stream_opened", "stream_failure", "restore"])]
fn closing_an_unfinished_stream_bills_its_delivered_chunks_once(
#[case] ended: Option<bool>,
#[case] expected: &[&str],
) {
Python::initialize();
Python::attach(|py| {
let locals = namespace(
py,
c"first = b'first'\nresponse = None\nfailure = ValueError('provider')\nhead = {}",
);
let mut logging = LegacyLogging {
operation: crate::LoggingOperation::Messages,
..logged(py, &locals, true)
};
logging
.on_stream_open(py, &local(&locals, "head").unbind())
.unwrap();
logging
.on_stream_chunk(py, &local(&locals, "first").unbind())
.unwrap();
match ended {
Some(true) => drop(succeed(py, &locals, &mut logging)),
Some(false) => drop(fail(py, &locals, &mut logging)),
None => {}
}
litellm_host_python::PythonOwned::close(&mut logging, py);
let names: Vec<String> = local(&locals, "logger")
.call_method0("names")
.unwrap()
.extract()
.unwrap();
assert_eq!(names, expected);
});
}
#[rstest]
#[case::synchronous(false, &["failure_handler"])]
#[case::asynchronous(true, &[])]
@ -2009,7 +2279,7 @@ assert chunks[1] is last
}
#[rstest::rstest]
fn internal_async_calls_skip_the_async_success_fan_out() {
fn internal_async_calls_skip_success_logging() {
Python::initialize();
Python::attach(|py| {
let locals = namespace(py, c"response = object()");
@ -2018,11 +2288,7 @@ assert chunks[1] is last
..logged(py, &locals, true)
};
succeed(py, &locals, &mut logging);
run(
py,
&locals,
c"assert logger.names() == ['sync_success_for_async_call'], logger.calls",
);
run(py, &locals, c"assert logger.calls == [], logger.calls");
});
}

33
litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs
View file

@ -26,7 +26,7 @@ pub trait LegacyCallbacks {
fn pre_call(
&self,
py: Python<'_>,
input: &str,
input: &Bound<'_, PyAny>,
api_key: Option<&str>,
body: &Bound<'_, PyDict>,
headers: &Bound<'_, PyDict>,
@ -59,7 +59,16 @@ pub trait LegacyCallbacks {
fn defers_async_logging(&self, py: Python<'_>) -> bool;
fn defer_success(&self, py: Python<'_>, pending: &Bound<'_, PyAny>) -> PyResult<()>;
fn claim_async_success_log(&self, py: Python<'_>) -> PyResult<bool>;
fn defer_success(
&self,
py: Python<'_>,
pending: &Bound<'_, PyAny>,
response: &Option<Py<PyAny>>,
start: &Py<PyAny>,
end: &Option<Py<PyAny>>,
) -> PyResult<()>;
fn sync_success_for_async_call(
&self,
@ -148,7 +157,7 @@ impl LegacyCallbacks for PythonLogger {
fn pre_call(
&self,
py: Python<'_>,
input: &str,
input: &Bound<'_, PyAny>,
api_key: Option<&str>,
body: &Bound<'_, PyDict>,
headers: &Bound<'_, PyDict>,
@ -203,8 +212,21 @@ impl LegacyCallbacks for PythonLogger {
.unwrap_or(false)
}
fn defer_success(&self, py: Python<'_>, pending: &Bound<'_, PyAny>) -> PyResult<()> {
Logging::DeferSuccess.call(py, (self.object(py), pending))?;
fn claim_async_success_log(&self, py: Python<'_>) -> PyResult<bool> {
Logging::ClaimAsyncSuccess
.call(py, (self.object(py),))?
.extract()
}
fn defer_success(
&self,
py: Python<'_>,
pending: &Bound<'_, PyAny>,
response: &Option<Py<PyAny>>,
start: &Py<PyAny>,
end: &Option<Py<PyAny>>,
) -> PyResult<()> {
Logging::DeferSuccess.call(py, (self.object(py), pending, response, start, end))?;
Ok(())
}
@ -252,6 +274,7 @@ impl LegacyCallbacks for PythonLogger {
) -> PyResult<()> {
let coroutine =
Logging::AsyncSuccessHandler.call(py, (self.object(py), response, start, end))?;
Logging::MarkAsyncSuccessEnqueued.call(py, (self.object(py),))?;
let enqueue = Logging::Enqueue.call(py, (&coroutine,));
if enqueue.is_err()
&& let Err(error) = coroutine.call_method0("close")

59
litellm-rust/crates/callbacks-legacy-python/src/deferred.rs
View file

@ -22,21 +22,35 @@ impl PendingSuccess {
self.logger
.enqueue_success(py, &self.response, &self.start, &self.end)
}
pub(crate) fn sync_for_async_call(&self, py: Python<'_>) -> PyResult<()> {
self.logger
.sync_success_for_async_call(py, &self.response, &self.start, &self.end)
}
}
#[pyclass]
pub(crate) struct PendingLogging {
pub(crate) pending: Option<PendingSuccess>,
pub(crate) enqueues: bool,
}
#[pymethods]
impl PendingLogging {
fn release(slf: &Bound<'_, Self>, py: Python<'_>, success: bool) -> PyResult<()> {
let pending = slf.borrow_mut().pending.take();
if let Some(pending) = pending
&& success
{
match pending.asynchronous(py) {
fn __call__(slf: &Bound<'_, Self>, py: Python<'_>) -> PyResult<()> {
let (pending, enqueues) = {
let mut slot = slf.borrow_mut();
(slot.pending.take(), slot.enqueues)
};
if let Some(pending) = pending {
let released = pending.sync_for_async_call(py).and_then(|()| {
if enqueues {
pending.asynchronous(py)
} else {
Ok(())
}
});
match released {
Err(error) if error.is_instance_of::<PyException>(py) => {
error.write_unraisable(py, Some(pending.logger.object(py)));
}
@ -87,6 +101,7 @@ mod tests {
start: py.None(),
end: Some(py.None()),
}),
enqueues: true,
},
)
.unwrap();
@ -108,7 +123,7 @@ observed = []
def on_enqueue(coroutine):
observed.append(marker.get())
pending.release(True)
pending()
logger.on_enqueue = on_enqueue
",
@ -118,28 +133,12 @@ logger.on_enqueue = on_enqueue
&locals,
c"
marker.set('release')
pending.release(True)
pending.release(True)
pending()
pending()
assert observed == ['release'], observed
assert logger.names() == ['async_success_handler', 'enqueued'], logger.calls
assert logger.names() == ['sync_success_for_async_call', 'async_success_handler', 'mark_async_success_enqueued', 'enqueued'], logger.calls
assert logger.calls[0][1] is response
",
);
});
}
#[test]
fn a_blocked_release_drops_the_success_for_good() {
Python::initialize();
Python::attach(|py| {
let locals = defer(py, c"");
run(
py,
&locals,
c"
pending.release(False)
pending.release(True)
assert logger.calls == [], logger.calls
assert logger.calls[1][1] is response
",
);
});
@ -168,7 +167,7 @@ logger.on_enqueue = on_enqueue
locals
.set_item("failure", py.eval(failure, None, Some(&locals)).unwrap())
.unwrap();
let released = local(&locals, "pending").call_method1("release", (true,));
let released = local(&locals, "pending").call0();
match released {
Ok(_) => assert!(!propagates),
Err(error) => {
@ -181,8 +180,8 @@ logger.on_enqueue = on_enqueue
py,
&locals,
c"
pending.release(True)
assert logger.names() == ['async_success_handler', 'enqueued', 'closed'], logger.calls
pending()
assert logger.names() == ['sync_success_for_async_call', 'async_success_handler', 'mark_async_success_enqueued', 'enqueued', 'closed'], logger.calls
assert unraisable_from(logger) == ([] if propagates else [failure])
",
);

6
litellm-rust/crates/callbacks-legacy-python/src/python.rs
View file

@ -42,6 +42,8 @@ pub(crate) enum Logging {
Update,
#[strum(serialize = "pre_call")]
PreCall,
#[strum(serialize = "messages_input")]
MessagesInput,
#[strum(serialize = "post_call")]
PostCall,
#[strum(serialize = "record_response")]
@ -50,6 +52,8 @@ pub(crate) enum Logging {
RecordApiWindow,
#[strum(serialize = "defers_async_logging")]
DefersAsync,
#[strum(serialize = "claim_async_success_log")]
ClaimAsyncSuccess,
#[strum(serialize = "defer_success")]
DeferSuccess,
#[strum(serialize = "sync_success_for_async_call")]
@ -58,6 +62,8 @@ pub(crate) enum Logging {
SubmitSuccess,
#[strum(serialize = "async_success_handler")]
AsyncSuccessHandler,
#[strum(serialize = "mark_async_success_enqueued")]
MarkAsyncSuccessEnqueued,
#[strum(serialize = "enqueue_logging")]
Enqueue,
#[strum(serialize = "failure_handler")]

15
litellm-rust/crates/callbacks-legacy-python/src/test_support.rs
View file

@ -62,6 +62,7 @@ FAKES = {
custom_llm_provider=provider,
),
'pre_call': lambda logger, input, api_key, additional_args: logger.pre_call(input, api_key, additional_args),
'messages_input': lambda body, kwargs: [{'role': 'user', 'content': json.dumps(body)}],
'post_call': lambda logger, original_response, api_key, additional_args: logger.post_call(
original_response, api_key, additional_args
),
@ -72,7 +73,11 @@ FAKES = {
llm_api_duration_ms=(end - start) * 1000
),
'defers_async_logging': lambda logger: bool(getattr(logger, '_defer_async_logging', False)),
'defer_success': lambda logger, pending: setattr(logger, '_native_pending_logging', pending),
'claim_async_success_log': lambda logger: logger.claim_async_success_log(),
'defer_success': lambda logger, pending, response, start, end: (
setattr(logger, 'deferred_success_result', (response, start, end)),
setattr(logger, 'enqueue_deferred_logging', pending),
),
'sync_success_for_async_call': lambda logger, response, start, end: logger.handle_sync_success_callbacks_for_async_calls(
response, start, end
),
@ -81,6 +86,7 @@ FAKES = {
)(error, ''.join(traceback.format_exception(error)), start, end),
'submit_success': lambda logger, response, start, end: logger.record('submit', (response, start, end)),
'async_success_handler': lambda logger, response, start, end: logger.async_success_handler(response, start, end),
'mark_async_success_enqueued': lambda logger: logger.record('mark_async_success_enqueued', None),
'enqueue_logging': lambda coroutine: coroutine.enqueue(),
'restore_context': lambda logger: logger.record('restore', None),
'custom_pricing_fields': lambda: ('ocr_cost_per_page',),
@ -92,7 +98,7 @@ FAKES = {
'after_deployment_failure': lambda kwargs, error, call_type: kwargs['logger'].hook('failure', error, call_type),
'stream_opened': lambda logger, head: logger.record('stream_opened', head),
'stamp_stream_head': lambda head, logger, kwargs, start, end: logger.record('stamp_stream_head', head),
'stream_success': lambda logger, url_route, endpoint_type, request_body, chunks, start, end, first_chunk: logger.record(
'stream_success': lambda logger, url_route, endpoint_type, request_body, chunks, start, end, first_chunk, stream_teardown: logger.record(
'stream_success', list(chunks)
),
'stream_failure': lambda logger, endpoint_type, request_body, chunks, error: logger.record('stream_failure', error),
@ -129,6 +135,7 @@ class StubCoroutine:
class StubLogger:
def __init__(self):
self.calls = []
self.model_call_details = {}
self.hooks = {}
self.on_enqueue = lambda coroutine: None
@ -156,6 +163,10 @@ class StubLogger:
self.record('async_success_handler', response)
return StubCoroutine(self)
def claim_async_success_log(self):
claimed, self.claimed = getattr(self, 'claimed', False), True
return not claimed
def handle_sync_success_callbacks_for_async_calls(self, response, start, end):
self.record('sync_success_for_async_call', response)

9
litellm-rust/crates/circuit-breaker/AGENTS.md
View file

@ -1,13 +1,14 @@
- `fsm` is a generic state machine runtime and knows nothing about breakers. `breaker` is a generic circuit breaker and knows nothing about Redis. Deciding which errors count as failures belongs to the consumer
- `Machine::step` stays pure. `&self` is config and time comes in on the event, never read inside the machine. The only clock read is `Admission`'s drop, at the call-site boundary
- Transition rules are one `match (state, event)` per machine, with no macro DSL and no `_` arm that would hide a new state from the compiler
- `Shared` counts a transition, bumps the generation and calls the observer only when the state's enum variant changes, matching Python bumping `generation` only when `_state` changes. Data inside a variant (failure counters, `since`) changes without a transition
- `Shared` counts a transition, bumps the generation and calls the observer only when the state's enum variant changes. Data inside a variant (failure counters, `since`) changes without a transition
- The observer runs under the lock so transitions are seen in generation order. Keep it to logging or a metric
- A disabled breaker is `Option<Breaker>::None` at the consumer, not a flag inside the machine. `Config` comes in already resolved, reading `REDIS_CIRCUIT_BREAKER_*` env is the consumer's job
## Parity with Python `RedisCircuitBreaker` (`litellm/caching/redis_cache.py`)
- The rstest table `matches_the_python_redis_circuit_breaker` was checked case by case against the Python breaker. A behavior change here needs the same check, by driving the Python class with the same steps
- Divergences, each on purpose
## The only Redis breaker
- Python's `RedisCache` admits through this breaker via `NativeRedisBreaker` (`python-bridge/src/redis/breaker.rs`), sharing one per Redis server with the native client through `redis-connect`'s `RedisHosts`, so the Python `RedisCircuitBreaker` it replaced is gone
- The rstest table `matches_the_python_redis_circuit_breaker` was checked case by case against that Python class and now pins its behavior. A behavior change here changes what every Redis caller in the proxy sees
- Divergences from the old Python class, each on purpose
- monotonic `Instant` instead of `time.time()`, so a wall clock jump can't open or close the breaker
- a probe `Admission` dropped without an outcome records a hard failure and reopens. Python's probe that dies on `CancelledError` (a `BaseException`) records nothing and leaves the breaker HALF_OPEN for good, and a dropped Rust future is common
- `_failure_count` is dropped, it only fed a log line

15
litellm-rust/crates/circuit-breaker/src/breaker.rs
View file

@ -1,5 +1,5 @@
use std::{
sync::Arc,
sync::{Arc, Weak},
time::{Duration, Instant},
};
@ -128,6 +128,9 @@ impl Machine for Config {
#[derive(Clone)]
pub struct Breaker(Arc<Shared<Config>>);
#[derive(Clone)]
pub struct WeakBreaker(Weak<Shared<Config>>);
#[derive(Debug, PartialEq, Eq)]
pub struct Rejected;
@ -165,6 +168,16 @@ impl Breaker {
pub fn state(&self) -> State {
self.0.state()
}
pub fn downgrade(&self) -> WeakBreaker {
WeakBreaker(Arc::downgrade(&self.0))
}
}
impl WeakBreaker {
pub fn upgrade(&self) -> Option<Breaker> {
self.0.upgrade().map(Breaker)
}
}
impl Admission {

4
litellm-rust/crates/circuit-breaker/src/lib.rs
View file

@ -1,5 +1,7 @@
mod breaker;
mod fsm;
pub use breaker::{Admission, Breaker, Config, Event, Failure, Output, Rejected, State};
pub use breaker::{
Admission, Breaker, Config, Event, Failure, Output, Rejected, State, WeakBreaker,
};
pub use fsm::{Machine, Observer, Shared, Transition};

3
litellm-rust/crates/config/AGENTS.md Normal file
View file

@ -0,0 +1,3 @@
## Redis settings
- This crate only deserializes YAML. `RedisSettings` is the raw `coordination_redis` or `cache_params` map and `CacheSettings` tells a `type: redis` cache apart from the rest. Reading those maps the way `litellm/_redis.py` does belongs to `gateway/src/redis`, and turning merged kwargs into a `Resolution` belongs to `redis-resolve`

7
litellm-rust/crates/config/Cargo.toml
View file

@ -6,17 +6,16 @@ license.workspace = true
repository.workspace = true
[dependencies]
litellm-secret-value.workspace = true
litellm-router-types.workspace = true
litellm-secret-value.workspace = true
serde.workspace = true
serde_json.workspace = true
serde_yaml_ng = "0.10.0"
strum.workspace = true
thiserror.workspace = true
tracing.workspace = true
veil.workspace = true
[dev-dependencies]
litellm-error.workspace = true
rstest.workspace = true
serde_json.workspace = true
tempfile.workspace = true
tracing-subscriber = { version = "0.3", default-features = false, features = ["registry", "std"] }

7
litellm-rust/crates/config/src/error.rs
View file

@ -11,11 +11,4 @@ pub enum Error {
path: PathBuf,
source: serde_yaml_ng::Error,
},
#[error("Redis host does not yet support configured setting {0}")]
RedisSetting(String),
#[error("invalid Redis setting {setting}: {reason}")]
RedisInvalid {
setting: &'static str,
reason: &'static str,
},
}

5
litellm-rust/crates/config/src/lib.rs
View file

@ -14,10 +14,7 @@ pub use error::Error;
pub use litellm_router_types::{LitellmParams, Spelled};
pub use mcp::{McpAuth, McpServer, McpTransport};
pub use model::Model;
pub use redis::{
CacheSettings, Lookup, RedisCluster, RedisEndpoint, RedisHost, RedisNode, RedisSettings,
ResolvedRedis, environment_has_target,
};
pub use redis::{CacheSettings, RedisSettings};
pub use settings::{
ClickHouseStoreSettings, GeneralSettings, LiteLlmSettings, RouterSettings, TracingSettings,
TracingStoreSettings,

658
litellm-rust/crates/config/src/redis.rs
View file

@ -1,195 +1,21 @@
use std::{collections::BTreeMap, fmt, time::Duration};
use std::collections::BTreeMap;
use litellm_router_types::Spelled;
use litellm_secret_value::SecretValue;
use serde::{Deserialize, Deserializer, de::Error as _};
use veil::Redact;
use crate::{AdditionalFields, Error, Object, Value};
use crate::{Object, Value};
pub type Lookup<'a> = &'a dyn Fn(&str) -> Option<String>;
const SUPPORTED: &[&str] = &[
"url",
"host",
"port",
"username",
"password",
"ssl",
"db",
"socket_timeout",
"startup_nodes",
"namespace",
];
const IGNORED: &[&str] = &[
"max_connections",
"single_connection_client",
"health_check_interval",
"socket_keepalive",
"socket_keepalive_options",
"decode_responses",
"encoding",
"encoding_errors",
"charset",
"errors",
"client_name",
"lib_name",
"lib_version",
"sentinel_password",
"gcp_ssl_ca_certs",
"azure_client_id",
"azure_tenant_id",
"azure_client_secret",
"aws_iam_user_name",
"aws_iam_cache_name",
"aws_iam_region",
"aws_iam_serverless",
];
const WARN_IGNORED: &[&str] = &[
"max_connections",
"health_check_interval",
"socket_keepalive",
"socket_keepalive_options",
];
const REDIS_PY_KWARGS: &[&str] = &[
"aws_iam_auth",
"aws_iam_cache_name",
"aws_iam_region",
"aws_iam_serverless",
"aws_iam_user_name",
"azure_client_id",
"azure_client_secret",
"azure_redis_ad_token",
"azure_tenant_id",
"cache",
"cache_config",
"charset",
"client_name",
"db",
"decode_responses",
"encoding",
"encoding_errors",
"errors",
"event_dispatcher",
"gcp_service_account",
"gcp_ssl_ca_certs",
"health_check_interval",
"host",
"lib_name",
"lib_version",
"max_connections",
"password",
"port",
"protocol",
"redis_connect_func",
"retry_on_error",
"retry_on_timeout",
"single_connection_client",
"socket_connect_timeout",
"socket_keepalive",
"socket_keepalive_options",
"socket_timeout",
"ssl",
"ssl_ca_certs",
"ssl_ca_data",
"ssl_ca_path",
"ssl_cert_reqs",
"ssl_certfile",
"ssl_check_hostname",
"ssl_ciphers",
"ssl_keyfile",
"ssl_min_version",
"ssl_ocsp_context",
"ssl_ocsp_expected_cert",
"ssl_password",
"ssl_validate_ocsp",
"ssl_validate_ocsp_stapled",
"unix_socket_path",
"url",
"username",
];
const UNSUPPORTED_EXTRA: &[&str] = &[
"credential_provider",
"sentinel_nodes",
"retry",
"connection_pool",
];
const CACHE_ONLY: &[&str] = &[
"type",
"mode",
"ttl",
"default_in_memory_ttl",
"default_in_redis_ttl",
"supported_call_types",
"similarity_threshold",
"max_messages",
"_backend",
];
#[derive(Clone, Default)]
/// A Redis section as `config.yaml` spells it, before any surface reads it.
#[derive(Clone, Default, Deserialize, Redact)]
#[serde(transparent)]
pub struct RedisSettings {
pub url: Option<SecretValue>,
pub host: Option<Spelled<String>>,
pub port: Option<Spelled<u16>>,
pub username: Option<SecretValue>,
pub password: Option<SecretValue>,
pub ssl: Option<Spelled<bool>>,
pub db: Option<Spelled<i64>>,
pub socket_timeout: Option<Spelled<f64>>,
pub startup_nodes: Option<Value>,
pub namespace: Option<Spelled<String>>,
pub additional_fields: AdditionalFields,
unparsed: DeferredInputs,
#[redact(with = "[REDACTED]")]
fields: BTreeMap<String, Value>,
}
impl<'de> Deserialize<'de> for RedisSettings {
fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
let fields = BTreeMap::<String, Value>::deserialize(deserializer)?;
let scalar = |key: &str| fields.get(key).and_then(as_text);
Ok(Self {
url: scalar("url").map(SecretValue::new),
host: scalar("host").map(Spelled::Value),
port: spelled(scalar("port")),
username: scalar("username").map(SecretValue::new),
password: scalar("password").map(SecretValue::new),
ssl: spelled(scalar("ssl")),
db: spelled(scalar("db")),
socket_timeout: spelled(scalar("socket_timeout")),
startup_nodes: fields
.get("startup_nodes")
.filter(|value| !value.is_null())
.cloned(),
namespace: scalar("namespace").map(Spelled::Value),
additional_fields: cloned(&fields, |key, _| !SUPPORTED.contains(&key)),
unparsed: deferred_inputs(&fields),
})
}
}
impl fmt::Debug for RedisSettings {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("RedisSettings")
.field("url", &self.url)
.field("host", &self.host)
.field("port", &self.port)
.field("username", &self.username)
.field("password", &self.password)
.field("ssl", &self.ssl)
.field("db", &self.db)
.field("socket_timeout", &self.socket_timeout)
.field(
"startup_nodes",
&self.startup_nodes.as_ref().map(|_| "<configured>"),
)
.field("namespace", &self.namespace)
.field("additional_fields", &self.additional_fields.keys())
.field("unparsed", &self.unparsed.keys())
.finish()
impl RedisSettings {
pub fn fields(&self) -> &BTreeMap<String, Value> {
&self.fields
}
}
@ -216,452 +42,22 @@ impl<'de> Deserialize<'de> for CacheSettings {
}
}
#[derive(Clone, PartialEq)]
pub struct ResolvedRedis {
pub endpoint: RedisEndpoint,
pub namespace: Option<String>,
pub timeout: Option<Duration>,
}
#[cfg(test)]
mod tests {
use super::RedisSettings;
use rstest::rstest;
impl fmt::Debug for ResolvedRedis {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("ResolvedRedis")
.field("endpoint", &self.endpoint)
.field("namespace", &self.namespace)
.field("timeout", &self.timeout)
.finish()
#[rstest]
#[case::scalar(serde_json::json!("credential-value"))]
#[case::object(serde_json::json!({"secret": "credential-value"}))]
#[case::array(serde_json::json!(["credential-value"]))]
fn credentials_are_redacted(#[case] password: serde_json::Value) {
let settings: RedisSettings = serde_json::from_value(serde_json::json!({
"host": "redis", "password": password,
}))
.unwrap();
let debug = format!("{settings:?}");
assert!(!debug.contains("credential-value"));
assert!(debug.contains("[REDACTED]"));
}
}
#[derive(Clone, PartialEq)]
pub enum RedisEndpoint {
Url(SecretValue),
Host(RedisHost),
Cluster(RedisCluster),
}
impl fmt::Debug for RedisEndpoint {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Url(_) => formatter.debug_tuple("Url").field(&"<redacted>").finish(),
Self::Host(host) => host.fmt(formatter),
Self::Cluster(cluster) => cluster.fmt(formatter),
}
}
}
#[derive(Clone, Debug, PartialEq)]
pub struct RedisHost {
pub host: String,
pub port: u16,
pub username: Option<SecretValue>,
pub password: Option<SecretValue>,
pub tls: bool,
pub db: i64,
}
#[derive(Clone, Debug, PartialEq)]
pub struct RedisCluster {
pub nodes: Vec<RedisNode>,
pub username: Option<SecretValue>,
pub password: Option<SecretValue>,
pub tls: bool,
}
#[derive(Clone, Debug, PartialEq)]
pub struct RedisNode {
pub host: String,
pub port: u16,
}
#[derive(Clone)]
struct Entry {
value: Value,
source: String,
}
type Kwargs = BTreeMap<String, Entry>;
impl RedisSettings {
pub fn resolve(&self, lookup: Lookup<'_>) -> Result<ResolvedRedis, Error> {
self.resolve_with(lookup, false)
}
pub fn resolve_cache(&self, lookup: Lookup<'_>) -> Result<ResolvedRedis, Error> {
self.resolve_with(lookup, true)
}
fn resolve_with(&self, lookup: Lookup<'_>, cache: bool) -> Result<ResolvedRedis, Error> {
let block = self.kwargs(cache);
let block: Kwargs = block
.into_iter()
.filter_map(|(key, entry)| {
let value = environment_reference(&entry.value, lookup)?;
Some((key, Entry { value, ..entry }))
})
.collect();
let environment = environment_kwargs(lookup);
let names_a_target = ["host", "startup_nodes", "sentinel_nodes"]
.iter()
.any(|key| block.contains_key(*key));
let drop_environment_url = names_a_target && !block.contains_key("url");
let merged: Kwargs = environment
.into_iter()
.filter(|(key, _)| !(key == "url" && drop_environment_url))
.chain(block)
.collect();
reject_unsupported(&merged, lookup)?;
warn_ignored(&merged);
let startup_nodes = match merged.get("startup_nodes") {
Some(entry) => Some(cluster_nodes(&entry.value)?),
None => lookup("REDIS_CLUSTER_NODES")
.map(|nodes| cluster_nodes(&Value::String(nodes)))
.transpose()?,
};
let username = secret_value(&merged, "username");
let password = secret_value(&merged, "password");
let tls = bool_value(&merged, "ssl");
let endpoint = match (startup_nodes, text_value(&merged, "url")) {
(Some(nodes), _) if !nodes.is_empty() => RedisEndpoint::Cluster(RedisCluster {
nodes,
username,
password,
tls,
}),
(_, Some(url)) => RedisEndpoint::Url(SecretValue::new(url)),
_ => RedisEndpoint::Host(RedisHost {
host: text_value(&merged, "host").ok_or(Error::RedisInvalid {
setting: "host",
reason: "either host or url must be specified for redis",
})?,
port: integer(&merged, "port")?.unwrap_or(6379),
username,
password,
tls,
db: integer(&merged, "db")?.unwrap_or(0),
}),
};
let timeout = merged
.get("socket_timeout")
.and_then(|entry| as_text(&entry.value))
.map(|seconds| {
seconds
.parse::<f64>()
.ok()
.and_then(|seconds| Duration::try_from_secs_f64(seconds).ok())
.ok_or(Error::RedisInvalid {
setting: "socket_timeout",
reason: "must be a non-negative number of seconds",
})
})
.transpose()?;
Ok(ResolvedRedis {
endpoint,
namespace: text_value(&merged, "namespace"),
timeout,
})
}
fn kwargs(&self, cache: bool) -> Kwargs {
let known = [
(
"url",
self.url.as_ref().map(|value| value.expose().to_owned()),
),
("host", spelled_text(self.host.as_ref())),
("port", spelled_text(self.port.as_ref())),
(
"username",
self.username
.as_ref()
.map(|value| value.expose().to_owned()),
),
(
"password",
self.password
.as_ref()
.map(|value| value.expose().to_owned()),
),
("ssl", spelled_text(self.ssl.as_ref())),
("db", spelled_text(self.db.as_ref())),
("socket_timeout", spelled_text(self.socket_timeout.as_ref())),
("namespace", spelled_text(self.namespace.as_ref())),
];
let known = known.into_iter().filter_map(|(key, value)| {
value.map(|value| {
(
key.to_owned(),
Entry {
value: Value::String(value),
source: key.to_owned(),
},
)
})
});
let startup_nodes = self.startup_nodes.clone().map(|value| {
(
"startup_nodes".to_owned(),
Entry {
value,
source: "startup_nodes".to_owned(),
},
)
});
let extras = self.additional_fields.iter().filter_map(|(key, value)| {
if cache && CACHE_ONLY.contains(&key.as_str()) {
return None;
}
let key = match (cache, key.as_str()) {
(true, "redis_startup_nodes") => "startup_nodes".to_owned(),
_ => key.clone(),
};
Some((
key.clone(),
Entry {
value: value.clone(),
source: key,
},
))
});
let unparsed = self.unparsed.iter().map(|(key, value)| {
(
key.clone(),
Entry {
value: value.0.clone(),
source: key.clone(),
},
)
});
extras
.chain(unparsed)
.chain(known)
.chain(startup_nodes)
.collect()
}
}
pub fn environment_has_target(lookup: Lookup<'_>) -> bool {
[
"REDIS_HOST",
"REDIS_URL",
"REDIS_CLUSTER_NODES",
"REDIS_SENTINEL_NODES",
]
.iter()
.any(|name| lookup(name).is_some())
}
fn reject_unsupported(merged: &Kwargs, lookup: Lookup<'_>) -> Result<(), Error> {
let active_auth = [
merged
.get("gcp_service_account")
.map(|entry| entry.source.as_str())
.or_else(|| lookup("REDIS_GCP_SERVICE_ACCOUNT").map(|_| "REDIS_GCP_SERVICE_ACCOUNT")),
text_value(merged, "azure_redis_ad_token")
.or_else(|| lookup("REDIS_AZURE_AD_TOKEN"))
.filter(|value| value.eq_ignore_ascii_case("true"))
.map(|_| {
merged
.get("azure_redis_ad_token")
.map_or("REDIS_AZURE_AD_TOKEN", |entry| entry.source.as_str())
}),
bool_value(merged, "aws_iam_auth").then(|| merged["aws_iam_auth"].source.as_str()),
merged
.get("sentinel_nodes")
.map(|entry| entry.source.as_str())
.or_else(|| lookup("REDIS_SENTINEL_NODES").map(|_| "REDIS_SENTINEL_NODES")),
];
if let Some(setting) = active_auth.into_iter().flatten().next() {
return Err(Error::RedisSetting(setting.to_owned()));
}
match merged
.iter()
.find(|(key, _)| is_unsupported(key))
.map(|(_, entry)| entry.source.clone())
{
Some(setting) => Err(Error::RedisSetting(setting)),
None => Ok(()),
}
}
fn warn_ignored(merged: &Kwargs) {
for setting in WARN_IGNORED {
if merged.contains_key(*setting) {
tracing::warn!(setting, "Redis setting is accepted but ignored");
}
}
}
fn is_unsupported(key: &str) -> bool {
let known = REDIS_PY_KWARGS.contains(&key) || UNSUPPORTED_EXTRA.contains(&key);
let handled = SUPPORTED.contains(&key)
|| IGNORED.contains(&key)
|| matches!(
key,
"gcp_service_account" | "azure_redis_ad_token" | "aws_iam_auth"
);
known && !handled
}
fn environment_kwargs(lookup: Lookup<'_>) -> Kwargs {
REDIS_PY_KWARGS
.iter()
.filter_map(|key| {
let source = format!("REDIS_{}", key.to_ascii_uppercase());
let value = lookup(&source)?;
Some((
(*key).to_owned(),
Entry {
value: Value::String(value),
source,
},
))
})
.collect()
}
fn environment_reference(value: &Value, lookup: Lookup<'_>) -> Option<Value> {
match value {
Value::String(text) => match text.strip_prefix("os.environ/") {
Some(name) => lookup(name).map(Value::String),
None => Some(value.clone()),
},
Value::Null => None,
_ => Some(value.clone()),
}
}
fn as_text(value: &Value) -> Option<String> {
match value {
Value::String(text) => Some(text.clone()),
Value::Number(number) => Some(number.to_string()),
Value::Bool(flag) => Some(flag.to_string()),
_ => None,
}
}
fn text_value(kwargs: &Kwargs, key: &str) -> Option<String> {
kwargs.get(key).and_then(|entry| as_text(&entry.value))
}
fn secret_value(kwargs: &Kwargs, key: &str) -> Option<SecretValue> {
text_value(kwargs, key).map(SecretValue::new)
}
fn bool_value(kwargs: &Kwargs, key: &str) -> bool {
match kwargs.get(key).map(|entry| &entry.value) {
Some(Value::Bool(flag)) => *flag,
Some(value) => as_text(value)
.is_some_and(|text| matches!(text.to_ascii_lowercase().as_str(), "true" | "1" | "yes")),
None => false,
}
}
fn integer<T: std::str::FromStr>(kwargs: &Kwargs, key: &'static str) -> Result<Option<T>, Error> {
kwargs
.get(key)
.and_then(|entry| as_text(&entry.value))
.map(|text| {
text.parse().map_err(|_| Error::RedisInvalid {
setting: key,
reason: "must be an integer",
})
})
.transpose()
}
const INVALID_NODES: Error = Error::RedisInvalid {
setting: "startup_nodes",
reason: "must be a list of {host, port} objects",
};
fn cluster_nodes(value: &Value) -> Result<Vec<RedisNode>, Error> {
match value {
Value::String(json) => {
let nodes: Vec<serde_json::Map<String, serde_json::Value>> =
serde_json::from_str(json).map_err(|_| INVALID_NODES)?;
nodes
.iter()
.map(|node| {
cluster_node(
node.get("host").and_then(json_text),
node.get("port").and_then(json_text),
)
})
.collect()
}
Value::Sequence(nodes) => nodes
.iter()
.map(|node| {
cluster_node(
node.get("host").and_then(as_text),
node.get("port").and_then(as_text),
)
})
.collect(),
_ => Err(INVALID_NODES),
}
}
fn cluster_node(host: Option<String>, port: Option<String>) -> Result<RedisNode, Error> {
match (host, port.and_then(|port| port.parse().ok())) {
(Some(host), Some(port)) => Ok(RedisNode { host, port }),
_ => Err(INVALID_NODES),
}
}
fn json_text(value: &serde_json::Value) -> Option<String> {
match value {
serde_json::Value::String(text) => Some(text.clone()),
serde_json::Value::Number(number) => Some(number.to_string()),
_ => None,
}
}
fn spelled_text<T: ToString>(value: Option<&Spelled<T>>) -> Option<String> {
value.map(|value| match value {
Spelled::Value(value) => value.to_string(),
Spelled::Text(value) => value.clone(),
})
}
fn spelled<T: std::str::FromStr>(text: Option<String>) -> Option<Spelled<T>> {
text.map(|text| match text.parse() {
Ok(value) => Spelled::Value(value),
Err(_) => Spelled::Text(text),
})
}
fn cloned(
fields: &BTreeMap<String, Value>,
keep: impl Fn(&str, &Value) -> bool,
) -> AdditionalFields {
fields
.iter()
.filter(|(key, value)| keep(key, value))
.map(|(key, value)| (key.clone(), value.clone()))
.collect()
}
#[derive(Clone)]
struct DeferredInput(Value);
type DeferredInputs = BTreeMap<String, DeferredInput>;
fn deferred_inputs(fields: &BTreeMap<String, Value>) -> DeferredInputs {
fields
.iter()
.filter(|(key, value)| is_unparsed(key, value))
.map(|(key, value)| (key.clone(), DeferredInput(value.clone())))
.collect()
}
fn is_unparsed(key: &str, value: &Value) -> bool {
SUPPORTED.contains(&key)
&& match key {
"startup_nodes" => value.is_null(),
_ => as_text(value).is_none(),
}
}

26
litellm-rust/crates/config/tests/config.rs
View file

@ -440,3 +440,29 @@ fn mcp_config_redacts_nested_credentials_and_preserves_policy_for_validation() {
fn rejects_invalid_typed_mcp_settings(#[case] setting: &str) {
assert!(Config::from_yaml(&format!("mcp_servers:\n docs:\n {setting}\n")).is_err());
}
#[rstest]
fn inactive_redis_cache_keeps_malformed_known_inputs_deferred() {
let config = Config::from_yaml(
"litellm_settings:\n cache: false\n cache_params:\n type: redis\n host: []\n port: invalid\n socket_timeout: invalid",
)
.unwrap();
assert!(matches!(
config.litellm_settings.cache_params,
Some(litellm_config::CacheSettings::Redis(_))
));
}
#[rstest]
fn non_redis_cache_settings_remain_opaque() {
let config = Config::from_yaml(
"litellm_settings:\n cache_params:\n type: local\n custom_option:\n nested: true",
)
.unwrap();
assert!(matches!(
config.litellm_settings.cache_params,
Some(litellm_config::CacheSettings::Other(_))
));
}

405
litellm-rust/crates/config/tests/redis.rs
View file

@ -1,405 +0,0 @@
use std::{
collections::HashMap,
fmt::Write,
sync::{Arc, Mutex},
};
use litellm_config::{CacheSettings, Config, Error, RedisEndpoint, RedisSettings, ResolvedRedis};
use rstest::rstest;
use tracing::{Event, Subscriber, field::Visit};
use tracing_subscriber::{Layer, layer::Context, prelude::*};
fn resolve(
yaml: &str,
environment: &[(&str, &str)],
) -> Result<litellm_config::ResolvedRedis, Error> {
let environment: HashMap<String, String> = environment
.iter()
.map(|(key, value)| ((*key).to_owned(), (*value).to_owned()))
.collect();
let config = Config::from_yaml(yaml).unwrap();
config
.general_settings
.coordination_redis
.as_ref()
.unwrap_or(&RedisSettings::default())
.resolve(&|name| environment.get(name).cloned())
}
fn resolve_cache(yaml: &str) -> Result<ResolvedRedis, Error> {
let config = Config::from_yaml(yaml).unwrap();
match config.litellm_settings.cache_params.unwrap() {
CacheSettings::Redis(settings) => settings.resolve_cache(&|_| None),
CacheSettings::Other(_) => panic!("expected a Redis cache"),
}
}
fn cluster_ports(resolved: &ResolvedRedis) -> Vec<u16> {
match &resolved.endpoint {
RedisEndpoint::Cluster(cluster) => cluster.nodes.iter().map(|node| node.port).collect(),
endpoint => panic!("expected a cluster, got {endpoint:?}"),
}
}
#[rstest]
#[case::host(
"general_settings:\n coordination_redis:\n host: redis\n port: 6380\n db: 2\n ssl: true\n namespace: app\n socket_timeout: 1.5",
&[],
"host:redis:6380:2:true:app:1500"
)]
#[case::environment_url(
"model_list: []",
&[("REDIS_URL", "redis://environment:6379")],
"url:redis://environment:6379"
)]
#[case::explicit_host_drops_environment_url(
"general_settings:\n coordination_redis:\n host: configured",
&[("REDIS_URL", "redis://environment:6379")],
"host:configured:6379:0:false::none"
)]
#[case::yaml_cluster(
"general_settings:\n coordination_redis:\n startup_nodes:\n - host: node-a\n port: 7000\n - host: node-b\n port: \"7001\"",
&[],
"cluster:node-a:7000,node-b:7001"
)]
#[case::json_environment_cluster(
"model_list: []",
&[("REDIS_CLUSTER_NODES", r#"[{"host":"node-a","port":"7000"},{"host":"node-b","port":7001}]"#)],
"cluster:node-a:7000,node-b:7001"
)]
#[case::missing_reference_keeps_environment_fallback(
"general_settings:\n coordination_redis:\n host: os.environ/MISSING_HOST",
&[("REDIS_HOST", "environment")],
"host:environment:6379:0:false::none"
)]
#[case::configured_values_beat_the_same_environment_values(
"general_settings:\n coordination_redis:\n host: configured\n port: 6380\n db: 3",
&[("REDIS_HOST", "environment"), ("REDIS_PORT", "6390"), ("REDIS_DB", "9")],
"host:configured:6380:3:false::none"
)]
#[case::environment_url_beats_the_environment_host(
"model_list: []",
&[("REDIS_URL", "redis://environment:6379"), ("REDIS_HOST", "ignored")],
"url:redis://environment:6379"
)]
#[case::startup_nodes_beat_a_url(
"general_settings:\n coordination_redis:\n url: redis://single:6379\n startup_nodes:\n - host: node-a\n port: 7000",
&[],
"cluster:node-a:7000"
)]
#[case::environment_references_resolve_the_host_and_port(
"general_settings:\n coordination_redis:\n host: os.environ/COORD_HOST\n port: os.environ/COORD_PORT",
&[("COORD_HOST", "referenced"), ("COORD_PORT", "6400")],
"host:referenced:6400:0:false::none"
)]
#[case::a_null_host_falls_back_to_the_environment(
"general_settings:\n coordination_redis:\n host: null",
&[("REDIS_HOST", "environment")],
"host:environment:6379:0:false::none"
)]
fn resolves_redis_through_the_parsed_config(
#[case] yaml: &str,
#[case] environment: &[(&str, &str)],
#[case] expected: &str,
) {
let resolved = resolve(yaml, environment).unwrap();
let description = match resolved.endpoint {
RedisEndpoint::Url(url) => format!("url:{}", url.expose()),
RedisEndpoint::Host(host) => format!(
"host:{}:{}:{}:{}:{}:{}",
host.host,
host.port,
host.db,
host.tls,
resolved.namespace.unwrap_or_default(),
resolved
.timeout
.map(|timeout| timeout.as_millis().to_string())
.unwrap_or_else(|| "none".to_owned())
),
RedisEndpoint::Cluster(cluster) => format!(
"cluster:{}",
cluster
.nodes
.iter()
.map(|node| format!("{}:{}", node.host, node.port))
.collect::<Vec<_>>()
.join(",")
),
};
assert_eq!(description, expected);
}
#[rstest]
#[case::custom_tls_file(
"general_settings:\n coordination_redis:\n host: redis\n ssl_ca_certs: /ca.pem",
&[],
"ssl_ca_certs"
)]
#[case::gcp_auth_from_config(
"general_settings:\n coordination_redis:\n host: redis\n gcp_service_account: account",
&[],
"gcp_service_account"
)]
#[case::gcp_auth_from_environment(
"model_list: []",
&[("REDIS_HOST", "redis"), ("REDIS_GCP_SERVICE_ACCOUNT", "account")],
"REDIS_GCP_SERVICE_ACCOUNT"
)]
#[case::azure_auth_from_config(
"general_settings:\n coordination_redis:\n host: redis\n azure_redis_ad_token: true",
&[],
"azure_redis_ad_token"
)]
#[case::azure_auth_from_environment(
"model_list: []",
&[("REDIS_HOST", "redis"), ("REDIS_AZURE_AD_TOKEN", "true")],
"REDIS_AZURE_AD_TOKEN"
)]
#[case::aws_auth_from_config(
"general_settings:\n coordination_redis:\n host: redis\n aws_iam_auth: true",
&[],
"aws_iam_auth"
)]
#[case::aws_auth_from_environment(
"model_list: []",
&[("REDIS_HOST", "redis"), ("REDIS_AWS_IAM_AUTH", "true")],
"REDIS_AWS_IAM_AUTH"
)]
#[case::sentinel_from_config(
"general_settings:\n coordination_redis:\n sentinel_nodes: [[redis, 26379]]",
&[],
"sentinel_nodes"
)]
#[case::sentinel_from_environment(
"model_list: []",
&[("REDIS_SENTINEL_NODES", "[[\"redis\",26379]]")],
"REDIS_SENTINEL_NODES"
)]
fn rejected_settings_name_their_source(
#[case] yaml: &str,
#[case] environment: &[(&str, &str)],
#[case] expected: &str,
) {
let error = resolve(yaml, environment).unwrap_err();
assert!(matches!(error, Error::RedisSetting(setting) if setting == expected));
}
#[rstest]
fn an_explicit_host_keeps_environment_credentials() {
let resolved = resolve(
"general_settings:\n coordination_redis:\n host: configured",
&[
("REDIS_PASSWORD", "environment-secret"),
("REDIS_SSL", "True"),
],
)
.unwrap();
let RedisEndpoint::Host(host) = resolved.endpoint else {
panic!("expected a host endpoint");
};
assert_eq!(
host.password.as_ref().map(|password| password.expose()),
Some("environment-secret")
);
assert!(host.tls);
}
#[rstest]
#[case::only_a_port("general_settings:\n coordination_redis:\n port: 6379", &[])]
#[case::a_structured_host_with_an_environment_target(
"general_settings:\n coordination_redis:\n host: []",
&[("REDIS_HOST", "environment"), ("REDIS_URL", "redis://environment:6379")]
)]
fn a_block_without_a_usable_target_fails(#[case] yaml: &str, #[case] environment: &[(&str, &str)]) {
let error = resolve(yaml, environment).unwrap_err();
assert!(matches!(
error,
Error::RedisInvalid {
setting: "host",
..
}
));
}
#[rstest]
fn the_cache_alias_names_cluster_nodes() {
let resolved = resolve_cache(
"litellm_settings:\n cache_params:\n type: redis\n redis_startup_nodes:\n - host: node-a\n port: 7000",
)
.unwrap();
assert_eq!(cluster_ports(&resolved), vec![7000]);
}
#[rstest]
fn canonical_startup_nodes_beat_the_cache_alias() {
let resolved = resolve_cache(
"litellm_settings:\n cache_params:\n type: redis\n startup_nodes:\n - host: canonical\n port: 7000\n redis_startup_nodes:\n - host: alias\n port: 7100\n - host: alias\n port: 7101",
)
.unwrap();
assert_eq!(cluster_ports(&resolved), vec![7000]);
}
#[rstest]
fn a_null_canonical_startup_nodes_setting_shadows_the_cache_alias() {
let resolved = resolve_cache(
"litellm_settings:\n cache_params:\n type: redis\n host: cache\n startup_nodes: null\n redis_startup_nodes:\n - host: alias\n port: 7100",
)
.unwrap();
assert!(matches!(resolved.endpoint, RedisEndpoint::Host(_)));
}
#[rstest]
fn a_missing_canonical_startup_nodes_reference_shadows_the_cache_alias() {
let resolved = resolve_cache(
"litellm_settings:\n cache_params:\n type: redis\n host: cache\n startup_nodes: os.environ/MISSING_NODES\n redis_startup_nodes:\n - host: alias\n port: 7100",
)
.unwrap();
assert!(matches!(resolved.endpoint, RedisEndpoint::Host(_)));
}
#[rstest]
fn rejected_environment_settings_never_print_the_value() {
let error = resolve(
"model_list: []",
&[
("REDIS_HOST", "redis"),
("REDIS_SSL_CERT_REQS", "private-option-value"),
],
)
.unwrap_err();
let message = error.to_string();
assert!(message.contains("REDIS_SSL_CERT_REQS"));
assert!(!message.contains("private-option-value"));
}
#[rstest]
#[case::negative("-1")]
#[case::not_a_number("slow")]
fn rejects_bad_socket_timeout(#[case] timeout: &str) {
let yaml = format!(
"general_settings:\n coordination_redis:\n host: redis\n socket_timeout: {timeout}"
);
let error = resolve(&yaml, &[]).unwrap_err();
assert!(matches!(
error,
Error::RedisInvalid {
setting: "socket_timeout",
..
}
));
}
#[rstest]
fn defaults_socket_timeout_for_the_gateway() {
let resolved = resolve(
"general_settings:\n coordination_redis:\n host: redis",
&[],
)
.unwrap();
assert_eq!(resolved.timeout, None);
}
#[rstest]
fn debug_redacts_redis_secrets_and_url() {
let config = Config::from_yaml(
"general_settings:\n coordination_redis:\n url: redis://user:password@redis:6379\n username: credential-user\n password: credential-password",
)
.unwrap();
let settings = config.general_settings.coordination_redis.unwrap();
let resolved = settings.resolve(&|_| None).unwrap();
let debug = format!("{settings:?} {resolved:?}");
assert!(!debug.contains("redis://user:password"));
assert!(!debug.contains("credential-user"));
assert!(!debug.contains("credential-password"));
assert!(debug.contains("redacted") || debug.contains("REDACTED"));
}
#[rstest]
fn inactive_redis_cache_keeps_malformed_known_inputs_deferred() {
let config = Config::from_yaml(
"litellm_settings:\n cache: false\n cache_params:\n type: redis\n host: []\n port: invalid\n socket_timeout: invalid",
)
.unwrap();
assert!(matches!(
config.litellm_settings.cache_params,
Some(litellm_config::CacheSettings::Redis(_))
));
}
#[rstest]
fn non_redis_cache_settings_remain_opaque() {
let config = Config::from_yaml(
"litellm_settings:\n cache_params:\n type: local\n custom_option:\n nested: true",
)
.unwrap();
assert!(matches!(
config.litellm_settings.cache_params,
Some(litellm_config::CacheSettings::Other(_))
));
}
#[rstest]
fn legacy_unknown_settings_do_not_change_resolution() {
let resolved = resolve(
"general_settings:\n coordination_redis:\n host: redis\n redis_flush_size: 42\n service_name: primary",
&[],
)
.unwrap();
assert!(matches!(resolved.endpoint, RedisEndpoint::Host(_)));
}
struct WarningCapture(Arc<Mutex<Vec<String>>>);
impl<S: Subscriber> Layer<S> for WarningCapture {
fn on_event(&self, event: &Event<'_>, _context: Context<'_, S>) {
let mut fields = Fields::default();
event.record(&mut fields);
self.0.lock().unwrap().push(fields.0);
}
}
#[derive(Default)]
struct Fields(String);
impl Visit for Fields {
fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
write!(self.0, "{}={value:?} ", field.name()).unwrap();
}
fn record_str(&mut self, field: &tracing::field::Field, value: &str) {
write!(self.0, "{}={value} ", field.name()).unwrap();
}
}
#[rstest]
fn ignored_setting_warning_names_the_setting_without_its_value() {
let events = Arc::new(Mutex::new(Vec::new()));
let subscriber = tracing_subscriber::registry().with(WarningCapture(events.clone()));
tracing::subscriber::with_default(subscriber, || {
resolve(
"general_settings:\n coordination_redis:\n host: redis\n max_connections: 47",
&[],
)
.unwrap();
});
let warnings = events.lock().unwrap().join(" ");
assert!(warnings.contains("max_connections"));
assert!(!warnings.contains("47"));
}

7
litellm-rust/crates/gateway/AGENTS.md
View file

@ -5,8 +5,11 @@
- Test mounting and server lifecycle through public HTTP behavior. Test endpoint semantics in the owning crate
## Redis
- `build_redis` builds the Redis clients once at boot (`src/redis.rs`). Components get a clone as a field, never a per-request extractor or a pool. redis-rs multiplexes one connection per node, so there is no connection pool and pool settings like `max_connections` are accepted and ignored
- `build_redis` builds the Redis clients once at boot (`src/redis/mod.rs`). Components get a clone as a field, never a per-request extractor or a pool. redis-rs multiplexes one connection per node, so there is no connection pool and pool settings like `max_connections` are accepted and ignored
- Which Redis backs which role mirrors the proxy: coordination is `general_settings.coordination_redis`, else the `cache_params` Redis when `cache: true` and `type: redis`, else `REDIS_*` env. The transaction buffer reuses coordination, else env when `use_redis_transaction_buffer` is on. The response cache gets its own client so large values never queue in front of coordination commands
- `litellm-config` follows `litellm/_redis.py`'s merge of config blocks and `REDIS_*` env. A redis-py setting we can't honor (sentinel, IAM, custom TLS files, ...) fails boot with its name, never silently dropped
- `src/redis/layers.rs` is the Rust copy of `_merge_redis_kwargs` in `litellm/_redis.py`, which only this host needs because the Python host merges in Python. It puts the configured kwargs over their `REDIS_*` values and fills the `get_secret` fallbacks from the `SETTINGS` rows in `redis-resolve`, under the same conditions Python reads them, then hands the result to `litellm_redis_resolve::interpret`. Values come from a `SecretSnapshot` filled once at boot, read like `get_secret` (secret manager, then environment, with `true` and `false` read as bools), and the only straight environment read is Azure's `REDIS_USERNAME`
- A surface is a `ConfiguredRedis` constructor: `coordination` for `general_settings.coordination_redis`, which rejects a malformed `port` like the proxy's `CoordinationRedisParams` pydantic model, `cache_params` for `litellm_settings.cache_params`, which drops the cache-only keys and reads the `redis_startup_nodes` alias, and `redis_cache_kwargs` for a `RedisCache.redis_kwargs` dict in the contract tests. A new surface is a new constructor, never a flag
- Nulls depend on the surface. `coordination_redis` drops every null like `model_dump(exclude_none=True)`, `cache_params` keeps a null as an explicit `None` that hides the `REDIS_*` value unless `Cache` or `RedisCache` names that parameter, and `RedisCache.redis_kwargs` keeps every null. A missing `os.environ/` reference is a null, as `load_config` resolves it
- A setting the Rust client cannot honor warns and is dropped, and only a security check it cannot enforce fails boot, per `redis-resolve/AGENTS.md`
- Boot never waits for a configured Redis (clients connect lazily). Only the env-only coordination fallback is pinged, for 2s, and treated as absent when unreachable, as in Python
- `/health/readiness` never touches Redis, matching the proxy

7
litellm-rust/crates/gateway/Cargo.toml
View file

@ -20,9 +20,14 @@ futures-util.workspace = true
litellm-gateway-auth.workspace = true
litellm-gateway-ui.workspace = true
litellm-config.workspace = true
litellm-core-utils.workspace = true
litellm-error.workspace = true
litellm-circuit-breaker.workspace = true
litellm-redis-client.workspace = true
litellm-redis-connect.workspace = true
litellm-redis-resolve.workspace = true
litellm-redis-types.workspace = true
rustls.workspace = true
litellm-http.workspace = true
litellm-llms.workspace = true
litellm-secrets.workspace = true
@ -32,10 +37,12 @@ serde.workspace = true
tracing.workspace = true
tokio = { workspace = true, features = ["signal"] }
uuid.workspace = true
veil.workspace = true
tower-sessions-moka-store = "0.15.0"
[dev-dependencies]
litellm-redis-testing.workspace = true
litellm-http = { workspace = true, features = ["test-support"] }
base64.workspace = true
rstest.workspace = true
tempfile.workspace = true

11
litellm-rust/crates/gateway/src/error.rs
View file

@ -12,13 +12,12 @@ pub enum Error {
Mcp(#[from] litellm_gateway_mcp::ConnectError),
#[error(transparent)]
Config(#[from] litellm_config::Error),
#[error("invalid Redis setting {setting}: {reason}")]
RedisInvalid {
setting: &'static str,
reason: &'static str,
},
#[error(transparent)]
RedisConnect(#[from] litellm_redis_client::ConnectError),
Redis(#[from] litellm_redis_resolve::Error),
#[error("Redis setting {setting} is not plain data")]
RedisNotPlainData { setting: String },
#[error(transparent)]
RedisConnect(#[from] litellm_redis_connect::Error),
#[error("MCP requires working gateway authentication")]
Auth(#[from] litellm_gateway_auth::Error),
}

22
litellm-rust/crates/gateway/src/lib.rs
View file

@ -68,15 +68,29 @@ pub fn build_inference(config: &Config) -> Result<Arc<Gateway>, Error> {
Ok(Arc::new(Gateway::new(resources, http, secrets, router)?))
}
pub async fn build_redis(config: &Config) -> Result<RedisClients, Error> {
let environment = secrets::environment_values(&config.environment_variables)?;
let lookup = |name: &str| {
pub fn install_crypto_provider() {
if rustls::crypto::CryptoProvider::get_default().is_none() {
let _ = rustls::crypto::ring::default_provider().install_default();
}
}
pub async fn build_redis(
config: &Config,
pool: &HttpClientPool,
http: &litellm_http::HttpClientConfig,
) -> Result<RedisClients, Error> {
let environment = Arc::new(secrets::environment_values(&config.environment_variables)?);
let lookup: litellm_redis_connect::SharedEnvironment = Arc::new(move |name: &str| {
environment
.get(name)
.map(|value| value.expose().to_owned())
.or_else(|| std::env::var(name).ok())
});
let resources = redis::Resources {
http: pool.client(http, ClientVariant::Provider)?,
lookup: lookup.clone(),
};
redis::connect(config, &lookup).await
redis::connect(config, &*lookup, &resources).await
}
pub async fn build_mcp(

4
litellm-rust/crates/gateway/src/main.rs
View file

@ -41,6 +41,7 @@ impl Sink for StderrSink {
#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
litellm_gateway::install_crypto_provider();
if let Some(forward) = litellm_gateway_forward::Settings::from_env()? {
let log: LogSettings = settings::from_iter(std::env::vars_os())?;
Logger::new(StderrSink {
@ -55,7 +56,8 @@ async fn main() -> Result<(), Box<dyn Error>> {
Logger::new(StderrSink { level }).install_global()?;
let config = Config::load(settings.litellm_config)?;
let inference = litellm_gateway::build_inference(&config)?;
let _redis = litellm_gateway::build_redis(&config).await?;
let _redis =
litellm_gateway::build_redis(&config, &inference.resources.pool, &inference.http).await?;
let ui = match std::env::var_os("LITELLM_UI_PATH") {
Some(directory) => {
let ui_settings: UiSettings = settings::from_iter(std::env::vars_os())?;

144
litellm-rust/crates/gateway/src/redis/configured.rs Normal file
View file

@ -0,0 +1,144 @@
use std::collections::BTreeMap;
use litellm_config::{RedisSettings, Value};
use litellm_core_utils::settings::Lookup;
use litellm_redis_resolve::{Expected, Kwargs};
use litellm_redis_types::{ProviderHandle, Resolution};
use veil::Redact;
use super::{layers, secrets::SecretSnapshot};
use crate::Error;
const CACHE_ONLY: &[&str] = &[
"type",
"mode",
"ttl",
"default_in_memory_ttl",
"default_in_redis_ttl",
"supported_call_types",
"similarity_threshold",
"max_messages",
"_backend",
];
/// The redis-py kwargs one surface hands `_get_redis_client_logic`, before `REDIS_*` and secrets.
#[derive(Clone, Redact)]
pub(crate) struct ConfiguredRedis {
#[redact(with = "[REDACTED]")]
pub(super) fields: BTreeMap<String, Value>,
pub(super) nulls: Nulls,
pub(super) credential_provider: Option<ProviderHandle>,
}
/// Whether a null setting reaches redis-py as an explicit `None`, which outranks its `REDIS_*` value.
#[derive(Clone, Copy, Debug)]
pub(super) enum Nulls {
/// `coordination_redis` reaches redis-py through `model_dump(exclude_none=True)`.
Dropped,
/// `Cache` drops a null only for the parameters it names and passes any other key on.
KeptUnlessNamedByCache,
/// `RedisCache.redis_kwargs` already holds only the nulls redis-py receives.
#[cfg(test)]
Kept,
}
impl ConfiguredRedis {
/// `general_settings.coordination_redis`, which the proxy validates as a pydantic model first.
/// An `os.environ/` port is only known after merging, where it is read like any other port.
pub(crate) fn coordination(
settings: &RedisSettings,
) -> Result<Self, litellm_redis_resolve::Error> {
if let Some(port) = settings.fields().get("port").and_then(text)
&& !port.starts_with("os.environ/")
&& port.parse::<u16>().is_err()
{
return Err(litellm_redis_resolve::Error::Invalid {
setting: "port".to_owned(),
expected: Expected::Integer,
});
}
Ok(Self {
fields: settings.fields().clone(),
nulls: Nulls::Dropped,
credential_provider: None,
})
}
/// `litellm_settings.cache_params` with `type: redis`, as `litellm.Cache` hands it to `RedisCache`.
pub(crate) fn cache_params(settings: &RedisSettings) -> Self {
let fields = settings.fields();
let alias = fields
.get("redis_startup_nodes")
.filter(|_| !fields.contains_key("startup_nodes"))
.map(|nodes| ("startup_nodes".to_owned(), nodes.clone()));
let fields = fields
.iter()
.filter(|(key, _)| {
key.as_str() != "redis_startup_nodes" && !CACHE_ONLY.contains(&key.as_str())
})
.map(|(key, value)| (key.clone(), value.clone()))
.chain(alias)
.collect();
Self {
fields,
nulls: Nulls::KeptUnlessNamedByCache,
credential_provider: None,
}
}
/// `RedisCache.redis_kwargs`, the dict every Python surface ends in.
#[cfg(test)]
pub(crate) fn redis_cache_kwargs(fields: BTreeMap<String, Value>) -> Self {
Self {
fields,
nulls: Nulls::Kept,
credential_provider: None,
}
}
/// The handle the host holds for the `credential_provider` object in these settings.
#[cfg(test)]
pub(crate) fn with_credential_provider(self, handle: ProviderHandle) -> Self {
Self {
credential_provider: Some(handle),
..self
}
}
/// No configured settings, so only `REDIS_*` and the other environment fallbacks apply.
pub(crate) fn environment_only() -> Self {
Self {
fields: BTreeMap::new(),
nulls: Nulls::Dropped,
credential_provider: None,
}
}
/// `_merge_redis_kwargs` in `litellm/_redis.py`: the kwargs over their `REDIS_*` values with every
/// `get_secret` fallback filled.
pub(crate) fn merged(&self, secrets: &SecretSnapshot) -> Result<Kwargs, Error> {
layers::merged(self, secrets)
}
/// The environment is only for what Python reads straight from `os.environ`, `REDIS_USERNAME` for Azure.
pub(crate) fn resolve(
&self,
secrets: &SecretSnapshot,
environment: &dyn Lookup,
) -> Result<Resolution, Error> {
Ok(litellm_redis_resolve::interpret(
&self.merged(secrets)?,
self.credential_provider,
environment.get("REDIS_USERNAME").as_deref(),
)?)
}
}
pub(super) fn text(value: &Value) -> Option<String> {
match value {
Value::String(text) => Some(text.clone()),
Value::Number(number) => Some(number.to_string()),
Value::Bool(flag) => Some(flag.to_string()),
_ => None,
}
}

226
litellm-rust/crates/gateway/src/redis/layers.rs Normal file
View file

@ -0,0 +1,226 @@
use litellm_config::Value;
use litellm_core_utils::settings::Lookup;
use litellm_redis_resolve::{Entry, Kwargs, SETTINGS, is_present, setting};
use super::configured::{ConfiguredRedis, Nulls};
use crate::Error;
const NAMED_BY_CACHE: &[&str] = &[
"host",
"port",
"password",
"socket_timeout",
"startup_nodes",
"namespace",
"gcp_service_account",
"gcp_ssl_ca_certs",
];
/// How `_merge_redis_kwargs` reads a fallback: `get_secret` turns `true` and `false` into a bool, and
/// `get_secret_str` treats that bool as unset.
#[derive(Clone, Copy)]
enum Read {
Secret,
SecretStr,
}
const TARGETS: &[(&str, Read)] = &[
("startup_nodes", Read::Secret),
("sentinel_nodes", Read::Secret),
("sentinel_password", Read::SecretStr),
("service_name", Read::Secret),
];
const IDENTITIES: &[(&str, Read)] = &[
("gcp_service_account", Read::SecretStr),
("gcp_ssl_ca_certs", Read::SecretStr),
("azure_redis_ad_token", Read::Secret),
];
const AZURE: &[(&str, Read)] = &[
("azure_client_id", Read::SecretStr),
("azure_tenant_id", Read::SecretStr),
("azure_client_secret", Read::SecretStr),
];
const AWS: &[(&str, Read)] = &[("aws_iam_region", Read::SecretStr)];
/// The configured settings over their `REDIS_*` values, `{**environment_kwargs, **env_overrides}`
/// in `_merge_redis_kwargs`, with the `get_secret` fallbacks it reads filled in. A null left in the result is
/// an explicit `None` that hid its environment value.
pub(super) fn merged(configured: &ConfiguredRedis, lookup: &dyn Lookup) -> Result<Kwargs, Error> {
let block: Kwargs = configured
.fields
.iter()
.map(|(key, value)| (key, environment_reference(value, lookup)))
.filter(|(key, value)| !value.is_null() || keeps_null(configured.nulls, key))
.map(|(key, value)| {
let value = serde_json::to_value(&value).map_err(|_| Error::RedisNotPlainData {
setting: key.clone(),
})?;
let entry = Entry {
value,
source: key.clone(),
};
Ok((key.clone(), entry))
})
.collect::<Result<_, Error>>()?;
let names_a_target = ["host", "startup_nodes", "sentinel_nodes"]
.iter()
.any(|key| is_set(&block, key));
let drop_environment_url = names_a_target && !is_set(&block, "url");
let layered: Kwargs = environment(lookup)
.filter(|(key, _)| !(key == "url" && drop_environment_url))
.chain(block)
.collect();
let fallbacks = fallbacks(&layered, configured.credential_provider.is_some(), lookup);
Ok(layered.into_iter().chain(fallbacks).collect())
}
// The fallbacks are read under the same conditions as _merge_redis_kwargs, so a secret Python never asks
// for is never filled in
fn fallbacks(layered: &Kwargs, has_provider: bool, lookup: &dyn Lookup) -> Kwargs {
let fill = |kwargs: &[(&str, Read)]| -> Kwargs {
kwargs
.iter()
.filter_map(|(kwarg, read)| fallback(layered, kwarg, *read, lookup))
.collect()
};
let targets = fill(TARGETS);
if has_provider || is_set(layered, "credential_provider") {
return targets;
}
let identities = fill(IDENTITIES);
let value = |kwarg: &str| {
layered
.get(kwarg)
.filter(|entry| is_present(&entry.value))
.or_else(|| identities.get(kwarg))
.map(|entry| &entry.value)
};
let gcp = value("gcp_service_account").is_some();
let azure = value("azure_redis_ad_token")
.is_some_and(|token| text(token).is_some_and(|token| token.eq_ignore_ascii_case("true")));
let aws = layered
.get("aws_iam_auth")
.is_some_and(|entry| coerces_to_true(&entry.value));
let azure_settings = match azure && !gcp {
true => fill(AZURE),
false => Kwargs::new(),
};
let aws_settings = match aws && !gcp && !azure {
true => fill(AWS),
false => Kwargs::new(),
};
targets
.into_iter()
.chain(identities)
.chain(azure_settings)
.chain(aws_settings)
.collect()
}
fn fallback(
layered: &Kwargs,
kwarg: &str,
read: Read,
lookup: &dyn Lookup,
) -> Option<(String, Entry)> {
if layered
.get(kwarg)
.is_some_and(|entry| is_present(&entry.value))
{
return None;
}
let names = setting(kwarg).map_or(&[][..], |setting| setting.fallback);
names.iter().find_map(|name| {
let value = lookup
.get(name)
.filter(|value| !value.is_empty())
.map(secret)?;
match (read, &value) {
(Read::SecretStr, serde_json::Value::Bool(_)) => None,
_ => {
let entry = Entry {
value,
source: (*name).to_owned(),
};
Some((kwarg.to_owned(), entry))
}
}
})
}
// get_secret reads an environment value of true or false, in any case and padding, as a bool
fn secret_flag(value: &str) -> Option<bool> {
match value.trim().to_ascii_lowercase().as_str() {
"true" => Some(true),
"false" => Some(false),
_ => None,
}
}
fn secret(value: String) -> serde_json::Value {
secret_flag(&value).map_or(serde_json::Value::String(value), serde_json::Value::Bool)
}
// _coerces_to_true: text is read like _str_to_bool, anything else by its truthiness
fn coerces_to_true(value: &serde_json::Value) -> bool {
match value {
serde_json::Value::String(text) => {
matches!(text.to_ascii_lowercase().as_str(), "true" | "1" | "yes")
}
serde_json::Value::Null => false,
serde_json::Value::Bool(flag) => *flag,
serde_json::Value::Number(number) => number.as_f64().is_some_and(|number| number != 0.0),
serde_json::Value::Array(items) => !items.is_empty(),
serde_json::Value::Object(entries) => !entries.is_empty(),
}
}
fn text(value: &serde_json::Value) -> Option<String> {
match value {
serde_json::Value::String(text) => Some(text.clone()),
serde_json::Value::Number(number) => Some(number.to_string()),
serde_json::Value::Bool(flag) => Some(flag.to_string()),
_ => None,
}
}
fn keeps_null(nulls: Nulls, key: &str) -> bool {
match nulls {
Nulls::Dropped => false,
Nulls::KeptUnlessNamedByCache => !NAMED_BY_CACHE.contains(&key),
#[cfg(test)]
Nulls::Kept => true,
}
}
fn environment(lookup: &dyn Lookup) -> impl Iterator<Item = (String, Entry)> + '_ {
SETTINGS.iter().filter_map(|setting| {
let source = setting.environment_name()?;
let value = lookup.get(&source)?;
let entry = Entry {
value: secret(value),
source,
};
Some((setting.kwarg.to_owned(), entry))
})
}
// load_config resolves os.environ/ to None when unset, so a missing reference is a null like any other
fn environment_reference(value: &Value, lookup: &dyn Lookup) -> Value {
match value {
Value::String(text) => match text.strip_prefix("os.environ/") {
Some(name) => lookup.get(name).map_or(Value::Null, |value| {
secret_flag(&value).map_or(Value::String(value), Value::Bool)
}),
None => value.clone(),
},
_ => value.clone(),
}
}
fn is_set(kwargs: &Kwargs, key: &str) -> bool {
kwargs.get(key).is_some_and(|entry| !entry.value.is_null())
}

410
litellm-rust/crates/gateway/src/redis.rs → litellm-rust/crates/gateway/src/redis/mod.rs
View file

@ -1,19 +1,31 @@
use std::time::Duration;
use std::{sync::Arc, time::Duration};
use litellm_circuit_breaker::{Breaker, Config as BreakerConfig, State};
use litellm_config::{
CacheSettings, Config, Lookup, RedisEndpoint, RedisSettings, ResolvedRedis, Value,
environment_has_target,
};
use litellm_redis_client::{
ClusterEndpoint, ClusterNode, DEFAULT_TIMEOUT, Endpoint, HostEndpoint, Namespace, RedisClient,
RedisConfig,
};
use litellm_circuit_breaker::Config as BreakerConfig;
use litellm_config::{CacheSettings, Config, RedisSettings, Value};
use litellm_core_utils::settings::Lookup;
use litellm_redis_client::{RedisClient, RedisEvents};
use litellm_redis_connect::{HostServices, RedisHosts, SharedEnvironment, TracingEvents};
use litellm_redis_resolve::Expected;
use litellm_redis_types::Resolution;
use crate::Error;
mod configured;
mod layers;
mod secrets;
#[cfg(test)]
mod surface_tests;
use configured::ConfiguredRedis;
use secrets::SecretSnapshot;
pub(crate) struct Resources {
pub(crate) http: litellm_http::Client,
pub(crate) lookup: SharedEnvironment,
}
const ENVIRONMENT_PING_TIMEOUT: Duration = Duration::from_secs(2);
type Params = ResolvedRedis;
type Params = Resolution;
pub struct RedisClients {
pub coordination: Option<RedisClient>,
@ -34,12 +46,15 @@ enum Target {
Environment(Params),
}
pub(crate) fn plan(config: &Config, lookup: Lookup<'_>) -> Result<Plan, Error> {
pub(crate) fn plan(config: &Config, lookup: &dyn Lookup) -> Result<Plan, Error> {
let response_cache = redis_cache_params(config)
.map(|settings| settings.resolve_cache(lookup))
.map(|settings| resolve(&ConfiguredRedis::cache_params(settings), lookup))
.transpose()?;
let coordination = match &config.general_settings.coordination_redis {
Some(settings) => Some(Target::Configured(settings.resolve(lookup)?)),
Some(settings) => Some(Target::Configured(resolve(
&ConfiguredRedis::coordination(settings)?,
lookup,
)?)),
None => match (&response_cache, config.litellm_settings.cache == Some(true)) {
(Some(cache), _) => Some(Target::Configured(cache.clone())),
(None, true) => environment(lookup)?.map(Target::Configured),
@ -57,12 +72,30 @@ pub(crate) fn plan(config: &Config, lookup: Lookup<'_>) -> Result<Plan, Error> {
})
}
pub async fn connect(config: &Config, lookup: Lookup<'_>) -> Result<RedisClients, Error> {
pub(crate) async fn connect(
config: &Config,
lookup: &dyn Lookup,
resources: &Resources,
) -> Result<RedisClients, Error> {
let plan = plan(config, lookup)?;
let breaker = breaker_config(lookup)?;
let events: Arc<dyn RedisEvents> = Arc::new(TracingEvents::new(lookup));
let hosts = RedisHosts::default();
let client = |params: Params| -> Result<RedisClient, Error> {
Ok(hosts.client(
params,
HostServices {
events: Some(events.clone()),
http: resources.http.clone(),
environment: resources.lookup.clone(),
credential_provider: None,
},
breaker,
)?)
};
let coordination = match plan.coordination {
Some(Target::Configured(params)) => Some(client(params, breaker)?),
Some(Target::Environment(params)) => match client(params, breaker) {
Some(Target::Configured(params)) => Some(client(params)?),
Some(Target::Environment(params)) => match client(params) {
Ok(client) => reachable(client).await,
Err(_) => {
tracing::warn!(
@ -75,13 +108,10 @@ pub async fn connect(config: &Config, lookup: Lookup<'_>) -> Result<RedisClients
};
let transaction_buffer = match (&coordination, plan.transaction_buffer_fallback) {
(Some(coordination), _) => Some(coordination.clone()),
(None, Some(params)) => Some(client(params, breaker)?),
(None, Some(params)) => Some(client(params)?),
(None, None) => None,
};
let response_cache = plan
.response_cache
.map(|params| client(params, breaker))
.transpose()?;
let response_cache = plan.response_cache.map(client).transpose()?;
Ok(RedisClients {
coordination,
transaction_buffer,
@ -98,17 +128,38 @@ fn redis_cache_params(config: &Config) -> Option<&RedisSettings> {
}
}
fn environment(lookup: Lookup<'_>) -> Result<Option<Params>, Error> {
fn resolve(configured: &ConfiguredRedis, lookup: &dyn Lookup) -> Result<Params, Error> {
let secrets = SecretSnapshot::read(configured.secret_names(), lookup);
let resolution = configured.resolve(&secrets, lookup)?;
for ignored in &resolution.ignored {
tracing::warn!(
"Redis setting {} is not supported by the Rust client and is ignored: {}",
ignored.setting,
ignored.consequence
);
}
Ok(resolution)
}
fn environment(lookup: &dyn Lookup) -> Result<Option<Params>, Error> {
match environment_has_target(lookup) {
true => RedisSettings::default()
.resolve(lookup)
.map(Some)
.map_err(Into::into),
true => Ok(Some(resolve(&ConfiguredRedis::environment_only(), lookup)?)),
false => Ok(None),
}
}
fn optional_environment(lookup: Lookup<'_>) -> Option<Params> {
fn environment_has_target(lookup: &dyn Lookup) -> bool {
[
"REDIS_HOST",
"REDIS_URL",
"REDIS_CLUSTER_NODES",
"REDIS_SENTINEL_NODES",
]
.iter()
.any(|name| lookup.get(name).is_some())
}
fn optional_environment(lookup: &dyn Lookup) -> Option<Params> {
match environment(lookup) {
Ok(params) => params,
Err(error) => {
@ -149,122 +200,101 @@ async fn reachable(client: RedisClient) -> Option<RedisClient> {
}
}
fn client(params: Params, breaker: Option<BreakerConfig>) -> Result<RedisClient, Error> {
tracing::info!(endpoint = ?params.endpoint, namespace = ?params.namespace, "redis client configured");
let endpoint = match params.endpoint {
RedisEndpoint::Url(url) => Endpoint::Url(url),
RedisEndpoint::Host(host) => Endpoint::Host(HostEndpoint {
host: host.host,
port: host.port,
username: host.username,
password: host.password,
tls: host.tls,
db: host.db,
}),
RedisEndpoint::Cluster(cluster) => Endpoint::Cluster(ClusterEndpoint {
nodes: cluster
.nodes
.into_iter()
.map(|node| ClusterNode {
host: node.host,
port: node.port,
})
.collect(),
username: cluster.username,
password: cluster.password,
tls: cluster.tls,
}),
};
Ok(RedisClient::new(RedisConfig {
endpoint,
namespace: params.namespace.map(Namespace::new).unwrap_or_default(),
timeout: params.timeout.unwrap_or(DEFAULT_TIMEOUT),
breaker: breaker.map(|config| Breaker::new(config, Box::new(log_transition))),
})?)
fn parse_seconds(text: &str) -> Option<Duration> {
text.parse::<f64>()
.ok()
.and_then(|seconds| Duration::try_from_secs_f64(seconds).ok())
}
fn log_transition(transition: &litellm_circuit_breaker::Transition<State>) {
match transition.to {
State::Open { .. } => tracing::warn!(
generation = transition.generation,
"redis circuit breaker opened, fast-failing redis calls"
),
State::HalfOpen => tracing::info!(
generation = transition.generation,
"redis circuit breaker half open, probing redis"
),
State::Closed { .. } => tracing::info!(
generation = transition.generation,
"redis circuit breaker closed, redis recovered"
),
}
}
fn breaker_config(lookup: Lookup<'_>) -> Result<Option<BreakerConfig>, Error> {
let enabled = lookup("REDIS_CIRCUIT_BREAKER_ENABLED")
fn breaker_config(lookup: &dyn Lookup) -> Result<Option<BreakerConfig>, Error> {
let enabled = lookup
.get("REDIS_CIRCUIT_BREAKER_ENABLED")
.is_none_or(|value| value.eq_ignore_ascii_case("true"));
if !enabled {
return Ok(None);
}
let defaults = BreakerConfig::default();
let seconds = |name: &'static str, default: Duration| -> Result<Duration, Error> {
lookup(name)
.map(|value| {
value
.parse::<f64>()
.ok()
.and_then(|seconds| Duration::try_from_secs_f64(seconds).ok())
.ok_or(Error::RedisInvalid {
setting: name,
reason: "must be a non-negative number of seconds",
})
})
.unwrap_or(Ok(default))
};
let failure_threshold = lookup("REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD")
Ok(Some(BreakerConfig {
failure_threshold: environment_value(
lookup,
"REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD",
Expected::Integer,
|text| text.parse().ok(),
)?
.unwrap_or(defaults.failure_threshold),
recovery_timeout: environment_value(
lookup,
"REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT",
Expected::Seconds,
parse_seconds,
)?
.unwrap_or(defaults.recovery_timeout),
timeout_min_duration: environment_value(
lookup,
"REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION",
Expected::Seconds,
parse_seconds,
)?
.unwrap_or(defaults.timeout_min_duration),
}))
}
fn environment_value<T>(
lookup: &dyn Lookup,
name: &str,
expected: Expected,
parse: impl Fn(&str) -> Option<T>,
) -> Result<Option<T>, litellm_redis_resolve::Error> {
lookup
.get(name)
.map(|value| {
value.parse().map_err(|_| Error::RedisInvalid {
setting: "REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD",
reason: "must be an integer",
parse(&value).ok_or_else(|| litellm_redis_resolve::Error::Invalid {
setting: name.to_owned(),
expected,
})
})
.unwrap_or(Ok(defaults.failure_threshold))?;
Ok(Some(BreakerConfig {
failure_threshold,
recovery_timeout: seconds(
"REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT",
defaults.recovery_timeout,
)?,
timeout_min_duration: seconds(
"REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION",
defaults.timeout_min_duration,
)?,
}))
.transpose()
}
#[cfg(test)]
mod tests {
use litellm_core_utils::settings::Lookup;
use std::{collections::HashMap, time::Duration};
use litellm_config::{Config, RedisEndpoint, RedisHost};
use litellm_config::Config;
use litellm_redis_client::{Standalone, StaticAuth, commands};
use litellm_redis_testing::RedisServer;
use litellm_redis_types::{Auth, HealthCheck};
use rstest::rstest;
use super::{Params, Plan, Target, connect, plan};
use std::sync::Arc;
use super::{Params, Plan, Resources, Target, plan};
use crate::Error;
async fn connect(config: &Config, lookup: &dyn Lookup) -> Result<super::RedisClients, Error> {
let resources = Resources {
http: litellm_http::Client::for_test(),
lookup: Arc::new(|_: &str| None),
};
super::connect(config, lookup, &resources).await
}
fn host(name: &str, port: u16) -> Params {
Params {
endpoint: RedisEndpoint::Host(RedisHost {
target: litellm_redis_client::Target::Standalone(Standalone {
host: name.to_owned(),
port,
username: None,
password: None,
tls: false,
tls: None,
db: 0,
}),
auth: Auth::Static(StaticAuth::default()),
namespace: None,
timeout: None,
connect_timeout: None,
health_check: HealthCheck::Unset,
keepalive: true,
ignored: Vec::new(),
}
}
@ -273,7 +303,7 @@ mod tests {
.iter()
.map(|(key, value)| ((*key).to_owned(), (*value).to_owned()))
.collect();
plan(&Config::from_yaml(yaml).unwrap(), &|name| {
plan(&Config::from_yaml(yaml).unwrap(), &|name: &str| {
environment.get(name).cloned()
})
}
@ -350,36 +380,44 @@ mod tests {
port_of(&server)
))
.unwrap();
let clients = connect(&config, &|_| None).await.unwrap();
let clients = connect(&config, &|_: &str| None).await.unwrap();
let coordination = clients.coordination.unwrap();
let buffer = clients.transaction_buffer.unwrap();
let cache = clients.response_cache.unwrap();
let unnamespaced = connect(&coordination_config(&server), &|_| None)
let unnamespaced = connect(&coordination_config(&server), &|_: &str| None)
.await
.unwrap()
.coordination
.unwrap();
coordination
.set("rate-limit", "1", Duration::from_secs(60))
.send(commands::set("rate-limit", "1", Duration::from_secs(60)))
.await
.unwrap();
buffer
.set("spend", "2", Duration::from_secs(60))
.send(commands::set("spend", "2", Duration::from_secs(60)))
.await
.unwrap();
assert_eq!(
cache.get::<String>("rate-limit").await.unwrap().as_deref(),
cache
.send(commands::get::<String>("rate-limit"))
.await
.unwrap()
.as_deref(),
Some("1")
);
assert_eq!(
cache.get::<String>("spend").await.unwrap().as_deref(),
cache
.send(commands::get::<String>("spend"))
.await
.unwrap()
.as_deref(),
Some("2")
);
assert_eq!(
unnamespaced
.get::<String>("responses:rate-limit")
.send(commands::get::<String>("responses:rate-limit"))
.await
.unwrap()
.as_deref(),
@ -387,7 +425,7 @@ mod tests {
);
assert_eq!(
unnamespaced
.get::<String>("responses:spend")
.send(commands::get::<String>("responses:spend"))
.await
.unwrap()
.as_deref(),
@ -395,16 +433,21 @@ mod tests {
);
assert!(
unnamespaced
.get::<String>("rate-limit")
.send(commands::get::<String>("rate-limit"))
.await
.unwrap()
.is_none()
);
assert!(
unnamespaced
.send(commands::get::<String>("spend"))
.await
.unwrap()
.is_none()
);
assert!(unnamespaced.get::<String>("spend").await.unwrap().is_none());
}
#[rstest]
#[case::malformed_port(&[("REDIS_HOST", "env"), ("REDIS_PORT", "not-a-port")])]
#[case::malformed_cluster_nodes(&[("REDIS_CLUSTER_NODES", "[{\"host\": ")])]
fn redis_migration_regression_malformed_optional_environment_is_absent(
#[case] environment: &[(&str, &str)],
@ -413,54 +456,45 @@ mod tests {
}
#[rstest]
fn redis_migration_regression_explicit_malformed_setting_still_fails() {
#[case::optional_environment("model_list: []", &[("REDIS_HOST", "env"), ("REDIS_PORT", "not-a-port")])]
#[case::transaction_buffer_environment(
"general_settings:\n use_redis_transaction_buffer: true",
&[("REDIS_HOST", "env"), ("REDIS_PORT", "not-a-port")]
)]
#[case::response_cache(
"litellm_settings:\n cache: true\n cache_params:\n type: redis\n host: env\n port: not-a-port",
&[]
)]
fn a_malformed_port_outside_the_coordination_block_falls_back_to_6379_like_redis_py(
#[case] yaml: &str,
#[case] environment: &[(&str, &str)],
) {
let plan = resolved(yaml, environment).unwrap();
let target = match (plan.coordination, plan.transaction_buffer_fallback) {
(Some(Target::Configured(params) | Target::Environment(params)), _)
| (None, Some(params)) => params,
(None, None) => panic!("expected a Redis target"),
};
assert_eq!(target.target, host("env", 6379).target);
}
#[rstest]
#[case::port("port: not-a-port", "port")]
#[case::ssl("ssl: maybe", "ssl")]
fn a_malformed_coordination_block_fails_like_the_pydantic_model(
#[case] setting: &str,
#[case] expected: &str,
) {
let error = resolved(
"general_settings:\n coordination_redis:\n host: coord\n port: not-a-port",
&format!("general_settings:\n coordination_redis:\n host: coord\n {setting}"),
&[],
)
.unwrap_err();
assert!(matches!(
error,
Error::Config(litellm_config::Error::RedisInvalid {
setting: "port",
..
})
));
}
#[rstest]
fn redis_migration_regression_enabled_transaction_buffer_rejects_invalid_environment() {
let error = resolved(
"general_settings:\n use_redis_transaction_buffer: true",
&[("REDIS_HOST", "env"), ("REDIS_PORT", "not-a-port")],
)
.unwrap_err();
assert!(matches!(
error,
Error::Config(litellm_config::Error::RedisInvalid {
setting: "port",
..
})
));
}
#[rstest]
fn redis_migration_regression_invalid_response_cache_still_fails() {
let error = resolved(
"litellm_settings:\n cache: true\n cache_params:\n type: redis\n host: cache\n port: not-a-port",
&[],
)
.unwrap_err();
assert!(matches!(
error,
Error::Config(litellm_config::Error::RedisInvalid {
setting: "port",
..
})
));
assert!(
matches!(error, Error::Redis(litellm_redis_resolve::Error::Invalid { setting, .. }) if setting == expected)
);
}
#[rstest]
@ -468,7 +502,7 @@ mod tests {
async fn redis_migration_regression_invalid_optional_environment_url_is_absent() {
let config = Config::from_yaml("model_list: []").unwrap();
let clients = connect(&config, &|name| {
let clients = connect(&config, &|name: &str| {
(name == "REDIS_URL").then(|| "not a redis URL".to_owned())
})
.await
@ -506,7 +540,7 @@ mod tests {
let server = RedisServer::start().await.unwrap();
server.stop().await.unwrap();
let clients = connect(&coordination_config(&server), &|_| None)
let clients = connect(&coordination_config(&server), &|_: &str| None)
.await
.unwrap();
let coordination = clients.coordination.unwrap();
@ -521,12 +555,12 @@ mod tests {
#[tokio::test]
async fn the_coordination_client_survives_a_redis_restart_and_backs_the_transaction_buffer() {
let server = RedisServer::start().await.unwrap();
let clients = connect(&coordination_config(&server), &|_| None)
let clients = connect(&coordination_config(&server), &|_: &str| None)
.await
.unwrap();
let coordination = clients.coordination.unwrap();
coordination
.set("written", "1", Duration::from_secs(60))
.send(commands::set("written", "1", Duration::from_secs(60)))
.await
.unwrap();
@ -537,6 +571,30 @@ mod tests {
assert!(buffer.ping().await.is_ok());
}
#[rstest]
#[tokio::test]
async fn an_environment_only_tls_redis_is_reached_with_its_ca() {
crate::install_crypto_provider();
let server = litellm_redis_testing::RedisTlsServer::start(false)
.await
.unwrap();
let ca = server.files().ca.display().to_string();
let port = server.port().to_string();
let config = Config::from_yaml("model_list: []").unwrap();
let clients = connect(&config, &|name: &str| match name {
"REDIS_HOST" => Some(server.host().to_owned()),
"REDIS_PORT" => Some(port.clone()),
"REDIS_SSL" => Some("true".to_owned()),
"REDIS_SSL_CA_CERTS" => Some(ca.clone()),
_ => None,
})
.await
.unwrap();
assert!(clients.coordination.is_some());
}
#[rstest]
#[tokio::test]
async fn an_unreachable_environment_only_redis_leaves_coordination_in_memory() {
@ -545,7 +603,7 @@ mod tests {
server.stop().await.unwrap();
let config = Config::from_yaml("model_list: []").unwrap();
let clients = connect(&config, &|name| match name {
let clients = connect(&config, &|name: &str| match name {
"REDIS_HOST" => Some("127.0.0.1".to_owned()),
"REDIS_PORT" => Some(port.clone()),
_ => None,
@ -564,7 +622,7 @@ mod tests {
let port = port_of(&server);
let config = Config::from_yaml("model_list: []").unwrap();
let clients = connect(&config, &|name| match name {
let clients = connect(&config, &|name: &str| match name {
"REDIS_HOST" => Some("127.0.0.1".to_owned()),
"REDIS_PORT" => Some(port.clone()),
_ => None,

61
litellm-rust/crates/gateway/src/redis/secrets.rs Normal file
View file

@ -0,0 +1,61 @@
use std::collections::{BTreeMap, BTreeSet};
use litellm_core_utils::settings::Lookup;
use litellm_secret_value::SecretValue;
use litellm_config::Value;
use litellm_redis_resolve::SETTINGS;
use super::configured::ConfiguredRedis;
#[derive(Clone, Debug, Default)]
pub(crate) struct SecretSnapshot(BTreeMap<String, Option<SecretValue>>);
impl SecretSnapshot {
pub(crate) fn read(names: BTreeSet<String>, lookup: &dyn Lookup) -> Self {
Self(
names
.into_iter()
.map(|name| {
let value = lookup.get(&name).map(SecretValue::new);
(name, value)
})
.collect(),
)
}
}
impl Lookup for SecretSnapshot {
fn get(&self, name: &str) -> Option<String> {
match self.0.get(name) {
Some(value) => value.as_ref().map(|value| value.expose().to_owned()),
None => {
debug_assert!(
false,
"{name} is missing from ConfiguredRedis::secret_names"
);
tracing::error!(
name,
"a Redis secret was read without being in the snapshot"
);
None
}
}
}
}
impl ConfiguredRedis {
pub(crate) fn secret_names(&self) -> BTreeSet<String> {
let references = self.fields.values().filter_map(|value| match value {
Value::String(text) => text.strip_prefix("os.environ/").map(str::to_owned),
_ => None,
});
let declared = SETTINGS.iter().flat_map(|setting| {
setting
.environment_name()
.into_iter()
.chain(setting.fallback.iter().map(|name| (*name).to_owned()))
});
declared.chain(references).collect()
}
}

964
litellm-rust/crates/gateway/src/redis/surface_tests.rs Normal file
View file

@ -0,0 +1,964 @@
use std::collections::{BTreeMap, HashMap};
use litellm_config::{CacheSettings, Config, Value};
use litellm_redis_resolve::Expected;
use litellm_redis_types::{
Auth, AwsIamAuth, AzureAdAuth, Node, ProviderHandle, Resolution, StaticAuth, Target, TlsOptions,
};
use litellm_secret_value::SecretValue;
use rstest::rstest;
use serde_json::json;
use super::{configured::ConfiguredRedis, secrets::SecretSnapshot};
use crate::Error;
fn resolution(
configured: &ConfiguredRedis,
environment: &HashMap<String, String>,
) -> Result<Resolution, Error> {
let lookup = |name: &str| environment.get(name).cloned();
configured.resolve(
&SecretSnapshot::read(configured.secret_names(), &lookup),
&lookup,
)
}
const PROVIDER: ProviderHandle = ProviderHandle::new(7);
fn redis_cache_kwargs(kwargs: serde_json::Value) -> ConfiguredRedis {
let has_provider = kwargs
.get("credential_provider")
.is_some_and(|provider| !provider.is_null());
let configured = ConfiguredRedis::redis_cache_kwargs(
serde_json::from_value::<BTreeMap<String, Value>>(kwargs).unwrap(),
);
match has_provider {
true => configured.with_credential_provider(PROVIDER),
false => configured,
}
}
fn coordination_resolution(yaml: &str, environment: &[(&str, &str)]) -> Result<Resolution, Error> {
let environment: HashMap<String, String> = environment
.iter()
.map(|(key, value)| ((*key).to_owned(), (*value).to_owned()))
.collect();
let config = Config::from_yaml(yaml).unwrap();
let configured = config
.general_settings
.coordination_redis
.unwrap_or_default();
let configured = ConfiguredRedis::coordination(&configured)?;
resolution(&configured, &environment)
}
fn resolve(yaml: &str, environment: &[(&str, &str)]) -> Result<Resolution, Error> {
coordination_resolution(yaml, environment)
}
fn resolve_cache(yaml: &str) -> Result<Resolution, Error> {
resolve_cache_in(yaml, &[])
}
fn resolve_cache_in(yaml: &str, environment: &[(&str, &str)]) -> Result<Resolution, Error> {
let lookup = |name: &str| {
environment
.iter()
.find(|(key, _)| *key == name)
.map(|(_, value)| (*value).to_owned())
};
let config = Config::from_yaml(yaml).unwrap();
let configured = match config.litellm_settings.cache_params.unwrap() {
CacheSettings::Redis(settings) => ConfiguredRedis::cache_params(&settings),
CacheSettings::Other(_) => panic!("expected a Redis cache"),
};
configured.resolve(
&SecretSnapshot::read(configured.secret_names(), &lookup),
&lookup,
)
}
fn cluster_ports(resolved: &Resolution) -> Vec<u16> {
match &resolved.target {
Target::Cluster(cluster) => cluster.nodes.iter().map(|node| node.port).collect(),
target => panic!("expected a cluster, got {target:?}"),
}
}
#[rstest]
#[case::host(
"general_settings:\n coordination_redis:\n host: redis\n port: 6380\n db: 2\n ssl: true\n namespace: app\n socket_timeout: 1.5",
&[],
"host:redis:6380:2:true:app:1500"
)]
#[case::environment_references_resolve_the_host_and_port(
"general_settings:\n coordination_redis:\n host: os.environ/COORD_HOST\n port: os.environ/COORD_PORT",
&[("COORD_HOST", "referenced"), ("COORD_PORT", "6400")],
"host:referenced:6400:0:false::none"
)]
#[case::a_null_host_falls_back_to_the_environment(
"general_settings:\n coordination_redis:\n host: null",
&[("REDIS_HOST", "environment")],
"host:environment:6379:0:false::none"
)]
fn resolves_redis_through_the_parsed_config(
#[case] yaml: &str,
#[case] environment: &[(&str, &str)],
#[case] expected: &str,
) {
let resolved = resolve(yaml, environment).unwrap();
let description = match resolved.target {
Target::Unix(socket) => format!("unix:{}", socket.path.display()),
Target::Url(url) => format!("url:{}", url.url.expose()),
Target::Sentinel(sentinel) => format!("sentinel:{}", sentinel.service_name),
Target::Standalone(host) => format!(
"host:{}:{}:{}:{}:{}:{}",
host.host,
host.port,
host.db,
host.tls.is_some(),
resolved.namespace.unwrap_or_default(),
resolved
.timeout
.map(|timeout| timeout.as_millis().to_string())
.unwrap_or_else(|| "none".to_owned())
),
Target::Cluster(cluster) => format!(
"cluster:{}",
cluster
.nodes
.iter()
.map(|node| format!("{}:{}", node.host, node.port))
.collect::<Vec<_>>()
.join(",")
),
};
assert_eq!(description, expected);
}
#[rstest]
#[case::tls_ciphers(
"general_settings:\n coordination_redis:\n host: redis\n ssl_ciphers: HIGH",
&[],
"ssl_ciphers"
)]
#[case::ssl_ciphers_from_environment(
"model_list: []",
&[("REDIS_HOST", "redis"), ("REDIS_SSL_CIPHERS", "HIGH")],
"REDIS_SSL_CIPHERS"
)]
#[case::connection_pool(
"general_settings:\n coordination_redis:\n host: redis\n connection_pool: shared",
&[],
"connection_pool"
)]
fn unsupported_settings_resolve_and_are_named_by_their_source(
#[case] yaml: &str,
#[case] environment: &[(&str, &str)],
#[case] expected: &str,
) {
let resolution = coordination_resolution(yaml, environment).unwrap();
assert_eq!(
resolution.target,
Target::Standalone(litellm_redis_types::Standalone {
host: "redis".into(),
port: 6379,
db: 0,
tls: None,
})
);
let [ignored] = resolution.ignored.as_slice() else {
panic!("expected one ignored setting, got {:?}", resolution.ignored);
};
assert_eq!(ignored.setting, expected);
assert!(!ignored.consequence.is_empty());
}
#[rstest]
#[case::kwarg(
"general_settings:\n coordination_redis:\n host: redis\n ssl: true\n ssl_validate_ocsp: true",
&[],
"ssl_validate_ocsp"
)]
#[case::stapled_from_environment(
"model_list: []",
&[("REDIS_HOST", "redis"), ("REDIS_SSL", "true"), ("REDIS_SSL_VALIDATE_OCSP_STAPLED", "yes")],
"REDIS_SSL_VALIDATE_OCSP_STAPLED"
)]
#[case::rediss_url_query(
"general_settings:\n coordination_redis:\n url: rediss://cache:6380/0?ssl_validate_ocsp=true",
&[],
"url ?ssl_validate_ocsp"
)]
fn an_ocsp_check_fails_instead_of_connecting_without_it(
#[case] yaml: &str,
#[case] environment: &[(&str, &str)],
#[case] setting: &str,
) {
let error = coordination_resolution(yaml, environment).unwrap_err();
assert!(
matches!(&error, Error::Redis(litellm_redis_resolve::Error::Unenforceable { setting: named, .. }) if named == setting),
"{error:?}"
);
}
#[rstest]
#[case::kwarg(" host: redis\n ssl: true\n ssl_validate_ocsp: false")]
#[case::rediss_url_query(" url: rediss://cache:6380/0?ssl_validate_ocsp=false")]
fn an_ocsp_flag_set_to_false_asks_for_nothing(#[case] block: &str) {
let yaml = format!("general_settings:\n coordination_redis:\n{block}");
let resolution = coordination_resolution(&yaml, &[]).unwrap();
assert!(resolution.ignored.is_empty(), "{:?}", resolution.ignored);
}
fn static_auth(username: Option<&str>, password: Option<&str>) -> Auth {
Auth::Static(StaticAuth {
username: username.map(litellm_secret_value::SecretValue::new),
password: password.map(litellm_secret_value::SecretValue::new),
})
}
#[rstest]
fn an_explicit_host_keeps_environment_credentials() {
let resolved = resolve(
"general_settings:\n coordination_redis:\n host: configured",
&[
("REDIS_PASSWORD", "environment-secret"),
("REDIS_SSL", "True"),
],
)
.unwrap();
assert_eq!(resolved.auth, static_auth(None, Some("environment-secret")));
assert!(matches!(resolved.target, Target::Standalone(host) if host.tls.is_some()));
}
fn azure(username: Option<&str>) -> Auth {
Auth::AzureAd(AzureAdAuth {
client_id: Some("client".to_owned()),
tenant_id: None,
client_secret: None,
username: username.map(str::to_owned),
})
}
fn aws(serverless: bool) -> Auth {
Auth::AwsIam(AwsIamAuth {
user_name: "iam-user".to_owned(),
cache_name: "cache-1".to_owned(),
region: "us-east-1".to_owned(),
serverless,
})
}
#[rstest]
#[case::static_credentials(
"general_settings:\n coordination_redis:\n host: redis\n username: user\n password: pass",
&[],
static_auth(Some("user"), Some("pass"))
)]
#[case::gcp_beats_azure_and_aws(
"general_settings:\n coordination_redis:\n host: redis\n password: pass\n azure_redis_ad_token: \"True\"\n aws_iam_auth: true",
&[("REDIS_GCP_SERVICE_ACCOUNT", "projects/-/serviceAccounts/sa@p.iam.gserviceaccount.com")],
Auth::GcpIam { service_account: "projects/-/serviceAccounts/sa@p.iam.gserviceaccount.com".to_owned() }
)]
#[case::azure_beats_aws_and_takes_the_environment_username(
"general_settings:\n coordination_redis:\n host: redis\n username: configured-user\n azure_redis_ad_token: true\n azure_client_id: client\n aws_iam_auth: true",
&[("REDIS_USERNAME", "environment-user")],
azure(Some("environment-user"))
)]
#[case::azure_without_an_environment_username_sends_only_the_token(
"general_settings:\n coordination_redis:\n host: redis\n username: configured-user\n azure_redis_ad_token: true\n azure_client_id: client",
&[],
azure(None)
)]
#[case::azure_needs_the_literal_true(
"general_settings:\n coordination_redis:\n host: redis\n azure_redis_ad_token: \"1\"",
&[],
static_auth(None, None)
)]
#[case::aws_with_a_region_from_the_environment(
"general_settings:\n coordination_redis:\n host: redis\n ssl: true\n password: dropped\n aws_iam_auth: \"yes\"\n aws_iam_user_name: iam-user\n aws_iam_cache_name: cache-1",
&[("AWS_DEFAULT_REGION", "us-east-1")],
aws(false)
)]
#[case::aws_serverless_over_a_rediss_url(
"general_settings:\n coordination_redis:\n url: rediss://cache-1:6379\n aws_iam_auth: true\n aws_iam_user_name: iam-user\n aws_iam_cache_name: cache-1\n aws_iam_region: us-east-1\n aws_iam_serverless: true",
&[],
aws(true)
)]
fn auth_follows_python_precedence(
#[case] yaml: &str,
#[case] environment: &[(&str, &str)],
#[case] expected: Auth,
) {
assert_eq!(resolve(yaml, environment).unwrap().auth, expected);
}
#[rstest]
#[case::coordination(
"general_settings:\n coordination_redis:\n host: redis\n credential_provider: host-object",
&[]
)]
#[case::without_a_target(
"general_settings:\n coordination_redis:\n credential_provider: host-object",
&[("REDIS_HOST", "redis")]
)]
fn a_credential_provider_without_a_host_handle_fails(
#[case] yaml: &str,
#[case] environment: &[(&str, &str)],
) {
let error = resolve(yaml, environment).unwrap_err();
assert!(
matches!(
error,
Error::Redis(litellm_redis_resolve::Error::ProviderWithoutHandle)
),
"{error}"
);
}
#[rstest]
#[case::without_tls(
"general_settings:\n coordination_redis:\n host: redis\n aws_iam_auth: true\n aws_iam_user_name: u\n aws_iam_cache_name: c\n aws_iam_region: r",
"requires TLS"
)]
#[case::a_plain_url_ignores_ssl(
"general_settings:\n coordination_redis:\n url: redis://cache:6379\n ssl: true\n aws_iam_auth: true\n aws_iam_user_name: u\n aws_iam_cache_name: c\n aws_iam_region: r",
"requires TLS"
)]
#[case::missing_settings(
"general_settings:\n coordination_redis:\n host: redis\n ssl: true\n aws_iam_auth: true\n aws_iam_cache_name: c",
"requires: aws_iam_user_name, aws_iam_region"
)]
fn incomplete_aws_iam_auth_fails(#[case] yaml: &str, #[case] message: &str) {
let error = resolve(yaml, &[]).unwrap_err();
assert!(error.to_string().contains(message), "{error}");
}
#[rstest]
#[case::defaults_verify_the_certificate_but_not_the_hostname(
" ssl: true",
&[],
Some(TlsOptions { ca_certs: None, ca_data: None, ca_path: None, certfile: None, keyfile: None, key_password: None, verify_certificate: true, check_hostname: false })
)]
#[case::every_option(
" ssl: true\n ssl_ca_certs: /ca.pem\n ssl_certfile: /client.pem\n ssl_keyfile: /client.key\n ssl_password: key-secret\n ssl_cert_reqs: none\n ssl_check_hostname: true",
&[],
Some(TlsOptions { ca_certs: Some("/ca.pem".into()), ca_data: None, ca_path: None, certfile: Some("/client.pem".into()), keyfile: Some("/client.key".into()), key_password: Some(SecretValue::new("key-secret")), verify_certificate: false, check_hostname: true })
)]
#[case::options_from_the_environment(
" ssl: true",
&[("REDIS_SSL_CA_CERTS", "/env-ca.pem"), ("REDIS_SSL_CERT_REQS", "optional"), ("REDIS_SSL_PASSWORD", "env-key-secret")],
Some(TlsOptions { ca_certs: Some("/env-ca.pem".into()), ca_data: None, ca_path: None, certfile: None, keyfile: None, key_password: Some(SecretValue::new("env-key-secret")), verify_certificate: true, check_hostname: false })
)]
#[case::gcp_ca_certs_replace_the_configured_ones(
" ssl: true\n ssl_ca_certs: /ca.pem\n gcp_service_account: account\n gcp_ssl_ca_certs: /gcp-ca.pem",
&[],
Some(TlsOptions { ca_certs: Some("/gcp-ca.pem".into()), ca_data: None, ca_path: None, certfile: None, keyfile: None, key_password: None, verify_certificate: true, check_hostname: false })
)]
#[case::options_without_ssl_stay_plaintext(" ssl_ca_certs: /ca.pem", &[], None)]
fn tls_options_resolve_like_redis_py(
#[case] block: &str,
#[case] environment: &[(&str, &str)],
#[case] expected: Option<TlsOptions>,
) {
let yaml = format!("general_settings:\n coordination_redis:\n host: redis\n{block}");
let resolved = resolve(&yaml, environment).unwrap();
assert!(matches!(resolved.target, Target::Standalone(host) if host.tls == expected));
}
#[rstest]
fn an_unknown_cert_requirement_fails() {
let error = resolve(
"general_settings:\n coordination_redis:\n host: redis\n ssl: true\n ssl_cert_reqs: strict",
&[],
)
.unwrap_err();
assert!(matches!(
error,
Error::Redis(litellm_redis_resolve::Error::Invalid { setting, expected: Expected::CertRequirement }) if setting == "ssl_cert_reqs"
));
}
#[rstest]
#[case::from_config(
"general_settings:\n coordination_redis:\n sentinel_nodes: [[sentinel-a, 26379], [sentinel-b, \"26380\"]]\n service_name: primary\n sentinel_password: sentinel-secret\n db: 2",
&[]
)]
#[case::from_the_environment(
"model_list: []",
&[
("REDIS_SENTINEL_NODES", r#"[["sentinel-a", 26379], ["sentinel-b", "26380"]]"#),
("REDIS_SERVICE_NAME", "primary"),
("REDIS_SENTINEL_PASSWORD", "sentinel-secret"),
("REDIS_DB", "2"),
]
)]
fn sentinel_nodes_resolve_to_a_sentinel_endpoint(
#[case] yaml: &str,
#[case] environment: &[(&str, &str)],
) {
let Target::Sentinel(sentinel) = resolve(yaml, environment).unwrap().target else {
panic!("expected a sentinel endpoint");
};
assert_eq!(
sentinel.monitors,
vec![
Node {
host: "sentinel-a".to_owned(),
port: 26379
},
Node {
host: "sentinel-b".to_owned(),
port: 26380
},
]
);
assert_eq!(sentinel.service_name, "primary");
assert_eq!(
sentinel.monitor_auth,
Some(StaticAuth {
username: None,
password: Some(litellm_secret_value::SecretValue::new("sentinel-secret")),
})
);
assert_eq!(sentinel.db, 2);
}
#[rstest]
#[case::a_url_beats_sentinels(
"general_settings:\n coordination_redis:\n url: redis://single:6379\n sentinel_nodes: [[s, 26379]]\n service_name: primary",
"url:redis://single:6379"
)]
fn sentinel_precedence(#[case] yaml: &str, #[case] expected: &str) {
let resolved = resolve(yaml, &[]).unwrap();
let description = match resolved.target {
Target::Url(url) => format!("url:{}", url.url.expose()),
Target::Sentinel(sentinel) => format!("sentinel:{}", sentinel.service_name),
other => panic!("unexpected {other:?}"),
};
assert_eq!(description, expected);
}
#[rstest]
fn sentinels_without_a_service_name_fail() {
let error = resolve(
"general_settings:\n coordination_redis:\n sentinel_nodes: [[s, 26379]]",
&[],
)
.unwrap_err();
assert!(
matches!(
error,
Error::Redis(litellm_redis_resolve::Error::NoSentinelService)
),
"{error:?}"
);
}
#[rstest]
#[case::configured_host(
"general_settings:\n coordination_redis:\n host: ignored\n sentinel_nodes: [[s, 26379]]\n service_name: primary",
&[],
"host"
)]
#[case::environment_port(
"general_settings:\n coordination_redis:\n sentinel_nodes: [[s, 26379]]\n service_name: primary",
&[("REDIS_PORT", "6380")],
"REDIS_PORT"
)]
fn a_host_or_port_cannot_reach_sentinel(
#[case] yaml: &str,
#[case] environment: &[(&str, &str)],
#[case] source: &str,
) {
let error = resolve(yaml, environment).unwrap_err();
assert!(
matches!(&error, Error::Redis(litellm_redis_resolve::Error::SentinelWithHost { setting }) if setting == source),
"{error:?}"
);
}
#[rstest]
#[case::not_pairs(
"general_settings:\n coordination_redis:\n sentinel_nodes: [{host: s, port: 26379}]\n service_name: primary",
&[],
"sentinel_nodes"
)]
#[case::from_the_environment(
"model_list: []",
&[("REDIS_SENTINEL_NODES", "not-json"), ("REDIS_SERVICE_NAME", "primary")],
"REDIS_SENTINEL_NODES"
)]
fn invalid_sentinel_nodes_name_their_source(
#[case] yaml: &str,
#[case] environment: &[(&str, &str)],
#[case] source: &str,
) {
let error = resolve(yaml, environment).unwrap_err();
assert!(
matches!(&error, Error::Redis(litellm_redis_resolve::Error::Invalid { setting, expected: Expected::SentinelNodes }) if setting == source),
"{error:?}"
);
}
#[rstest]
#[case::from_config("general_settings:\n coordination_redis:\n startup_nodes: [{host: a}]", &[], "startup_nodes")]
#[case::from_the_environment("model_list: []", &[("REDIS_CLUSTER_NODES", "not-json")], "REDIS_CLUSTER_NODES")]
#[case::out_of_range_port_from_the_environment("model_list: []", &[("REDIS_HOST", "redis"), ("REDIS_PORT", "70000")], "REDIS_PORT")]
fn invalid_values_name_their_source(
#[case] yaml: &str,
#[case] environment: &[(&str, &str)],
#[case] source: &str,
) {
let error = resolve(yaml, environment).unwrap_err();
assert!(
matches!(&error, Error::Redis(litellm_redis_resolve::Error::Invalid { setting, .. }) if setting == source),
"{error:?}"
);
}
#[rstest]
#[case::falls_back_to_the_socket_timeout(" socket_timeout: 1.5", Some(1500), None)]
#[case::explicit_connect_timeout(
" socket_timeout: 1.5\n socket_connect_timeout: 0.25",
Some(1500),
Some(250)
)]
fn connect_timeout_is_resolved_separately(
#[case] block: &str,
#[case] timeout_ms: Option<u128>,
#[case] connect_ms: Option<u128>,
) {
let resolved = resolve(
&format!("general_settings:\n coordination_redis:\n host: redis\n{block}"),
&[],
)
.unwrap();
assert_eq!(
resolved.timeout.map(|timeout| timeout.as_millis()),
timeout_ms
);
assert_eq!(
resolved.connect_timeout.map(|timeout| timeout.as_millis()),
connect_ms
);
}
#[rstest]
#[case::only_a_port("general_settings:\n coordination_redis:\n port: 6379", &[])]
#[case::a_structured_host_with_an_environment_target(
"general_settings:\n coordination_redis:\n host: []",
&[("REDIS_HOST", "environment"), ("REDIS_URL", "redis://environment:6379")]
)]
fn a_block_without_a_usable_target_fails(#[case] yaml: &str, #[case] environment: &[(&str, &str)]) {
let error = resolve(yaml, environment).unwrap_err();
assert!(
matches!(error, Error::Redis(litellm_redis_resolve::Error::NoTarget)),
"{error:?}"
);
}
#[rstest]
fn the_cache_alias_names_cluster_nodes() {
let resolved = resolve_cache(
"litellm_settings:\n cache_params:\n type: redis\n redis_startup_nodes:\n - host: node-a\n port: 7000",
)
.unwrap();
assert_eq!(cluster_ports(&resolved), vec![7000]);
}
#[rstest]
fn canonical_startup_nodes_beat_the_cache_alias() {
let resolved = resolve_cache(
"litellm_settings:\n cache_params:\n type: redis\n startup_nodes:\n - host: canonical\n port: 7000\n redis_startup_nodes:\n - host: alias\n port: 7100\n - host: alias\n port: 7101",
)
.unwrap();
assert_eq!(cluster_ports(&resolved), vec![7000]);
}
#[rstest]
fn a_null_canonical_startup_nodes_setting_shadows_the_cache_alias() {
let resolved = resolve_cache(
"litellm_settings:\n cache_params:\n type: redis\n host: cache\n startup_nodes: null\n redis_startup_nodes:\n - host: alias\n port: 7100",
)
.unwrap();
assert!(matches!(resolved.target, Target::Standalone(_)));
}
#[rstest]
fn a_missing_canonical_startup_nodes_reference_shadows_the_cache_alias() {
let resolved = resolve_cache(
"litellm_settings:\n cache_params:\n type: redis\n host: cache\n startup_nodes: os.environ/MISSING_NODES\n redis_startup_nodes:\n - host: alias\n port: 7100",
)
.unwrap();
assert!(matches!(resolved.target, Target::Standalone(_)));
}
#[rstest]
fn unsupported_environment_settings_never_carry_the_value() {
let resolution = coordination_resolution(
"model_list: []",
&[
("REDIS_HOST", "redis"),
("REDIS_SSL_CIPHERS", "private-option-value"),
],
)
.unwrap();
let ignored = format!("{:?}", resolution.ignored);
assert!(ignored.contains("REDIS_SSL_CIPHERS"), "{ignored}");
assert!(!ignored.contains("private-option-value"), "{ignored}");
}
#[rstest]
#[case::negative("-1")]
#[case::infinite("inf")]
fn rejects_bad_socket_timeout(#[case] timeout: &str) {
let yaml = format!(
"general_settings:\n coordination_redis:\n host: redis\n socket_timeout: {timeout}"
);
let error = resolve(&yaml, &[]).unwrap_err();
assert!(matches!(
error,
Error::Redis(litellm_redis_resolve::Error::Invalid { setting, expected: Expected::Seconds }) if setting == "socket_timeout"
));
}
#[rstest]
fn debug_redacts_redis_secrets_and_url() {
let config = Config::from_yaml(
"general_settings:\n coordination_redis:\n url: redis://user:password@redis:6379\n username: credential-user\n password: credential-password\n sentinel_kwargs: {password: credential-passthrough}",
)
.unwrap();
let settings = config.general_settings.coordination_redis.unwrap();
let configured = ConfiguredRedis::coordination(&settings).unwrap();
let resolved = resolution(&configured, &HashMap::new()).unwrap();
let debug = format!("{settings:?} {configured:?} {resolved:?}");
assert!(!debug.contains("redis://user:password"));
assert!(!debug.contains("credential-user"));
assert!(!debug.contains("credential-password"));
assert!(!debug.contains("credential-passthrough"));
assert!(debug.contains("redacted") || debug.contains("REDACTED"));
}
#[rstest]
fn legacy_unknown_settings_do_not_change_resolution() {
let resolved = resolve(
"general_settings:\n coordination_redis:\n host: redis\n redis_flush_size: 42\n service_name: primary",
&[],
)
.unwrap();
assert!(matches!(resolved.target, Target::Standalone(_)));
}
#[rstest]
#[case::configured(
"general_settings:\n coordination_redis:\n host: redis\n max_connections: 47",
&[],
"max_connections"
)]
#[case::from_the_environment(
"general_settings:\n coordination_redis:\n host: redis",
&[("REDIS_MAX_CONNECTIONS", "47")],
"REDIS_MAX_CONNECTIONS"
)]
fn ignored_settings_are_named_by_their_source(
#[case] yaml: &str,
#[case] environment: &[(&str, &str)],
#[case] source: &str,
) {
let resolution = coordination_resolution(yaml, environment).unwrap();
let settings: Vec<&str> = resolution
.ignored
.iter()
.map(|ignored| ignored.setting.as_str())
.collect();
assert_eq!(settings, vec![source]);
}
fn snapshot(values: &[(&str, &str)]) -> SecretSnapshot {
let values: HashMap<&str, &str> = values.iter().copied().collect();
SecretSnapshot::read(
ConfiguredRedis::environment_only().secret_names(),
&|name: &str| values.get(name).map(|value| (*value).to_owned()),
)
}
#[rstest]
fn secret_manager_values_reach_the_resolver_through_the_snapshot() {
let settings = redis_cache_kwargs(json!({"host": "os.environ/COORD_HOST"}));
let secrets = SecretSnapshot::read(settings.secret_names(), &|name: &str| match name {
"COORD_HOST" => Some("from-the-manager".to_owned()),
"REDIS_PASSWORD" => Some("manager-secret".to_owned()),
_ => None,
});
let resolved = settings.resolve(&secrets, &|_: &str| None).unwrap();
assert!(
matches!(&resolved.target, Target::Standalone(host) if host.host == "from-the-manager")
);
assert_eq!(resolved.auth, static_auth(None, Some("manager-secret")));
}
#[rstest]
#[case::from_the_environment(&[], &[("REDIS_USERNAME", "environment-user")], Some("environment-user"))]
#[case::never_from_the_secret_manager(&[("REDIS_USERNAME", "manager-user")], &[], None)]
fn the_azure_username_comes_from_the_process_environment_like_python(
#[case] secrets: &[(&str, &str)],
#[case] environment: &[(&str, &str)],
#[case] expected: Option<&str>,
) {
let settings = redis_cache_kwargs(json!({"host": "redis", "azure_redis_ad_token": "true"}));
let environment: HashMap<&str, &str> = environment.iter().copied().collect();
let resolved = settings
.resolve(&snapshot(secrets), &|name: &str| {
environment.get(name).map(|value| (*value).to_owned())
})
.unwrap();
assert!(
matches!(&resolved.auth, Auth::AzureAd(azure) if azure.username.as_deref() == expected),
"{:?}",
resolved.auth
);
}
#[rstest]
#[case::coordination_drops_the_null(
"general_settings:\n coordination_redis:\n host: redis\n ssl: null",
false,
true
)]
#[case::cache_passes_an_unnamed_null_on(
"litellm_settings:\n cache_params:\n type: redis\n host: redis\n ssl: null",
true,
false
)]
#[case::cache_drops_a_null_it_names(
"litellm_settings:\n cache_params:\n type: redis\n host: null\n ssl: true",
true,
true
)]
fn a_null_overrides_the_environment_only_where_python_passes_it_on(
#[case] yaml: &str,
#[case] cache: bool,
#[case] tls: bool,
) {
let environment = [("REDIS_HOST", "environment"), ("REDIS_SSL", "true")];
let resolved = match cache {
true => resolve_cache_in(yaml, &environment),
false => resolve(yaml, &environment),
}
.unwrap();
let Target::Standalone(host) = resolved.target else {
panic!("expected a standalone target, got {:?}", resolved.target);
};
assert_eq!(host.tls.is_some(), tls);
}
#[rstest]
#[case::empty_cluster_nodes_use_the_environment(
json!({"startup_nodes": []}),
&[("REDIS_CLUSTER_NODES", r#"[{"host": "node-a", "port": 7000}]"#)],
json!({"startup_nodes": "[{\"host\": \"node-a\", \"port\": 7000}]"}),
&[("startup_nodes", "REDIS_CLUSTER_NODES")],
)]
#[case::empty_sentinel_nodes_use_the_environment(
json!({"sentinel_nodes": [], "service_name": "primary"}),
&[("REDIS_SENTINEL_NODES", r#"[["sentinel-a", 26379]]"#)],
json!({"sentinel_nodes": "[[\"sentinel-a\", 26379]]", "service_name": "primary"}),
&[("sentinel_nodes", "REDIS_SENTINEL_NODES")],
)]
#[case::empty_service_name_uses_the_environment(
json!({"sentinel_nodes": [["sentinel-a", 26379]], "service_name": ""}),
&[("REDIS_SERVICE_NAME", "primary")],
json!({"sentinel_nodes": [["sentinel-a", 26379]], "service_name": "primary"}),
&[("service_name", "REDIS_SERVICE_NAME")],
)]
#[case::environment_host_port_and_password(
json!({}),
&[("REDIS_HOST", "environment"), ("REDIS_PASSWORD", "environment-secret"), ("REDIS_PORT", "6390"), ("REDIS_SSL", "True")],
json!({"port": "6390", "ssl": true, "host": "environment", "password": "environment-secret"}),
&[("host", "REDIS_HOST"), ("password", "REDIS_PASSWORD"), ("port", "REDIS_PORT"), ("ssl", "REDIS_SSL")],
)]
#[case::padded_environment_numbers_parse_like_python_int(
json!({}),
&[("REDIS_HOST", "redis"), ("REDIS_PORT", "6_380"), ("REDIS_DB", "3\n")],
json!({"port": "6_380", "host": "redis", "db": "3\n"}),
&[("db", "REDIS_DB"), ("host", "REDIS_HOST"), ("port", "REDIS_PORT")],
)]
#[case::socket_timeout_from_the_environment(
json!({"host": "redis"}),
&[("REDIS_SOCKET_TIMEOUT", "0.75")],
json!({"socket_timeout": "0.75", "host": "redis"}),
&[("socket_timeout", "REDIS_SOCKET_TIMEOUT")],
)]
#[case::environment_url(
json!({}),
&[("REDIS_URL", "redis://environment:6379/1")],
json!({"url": "redis://environment:6379/1"}),
&[("url", "REDIS_URL")],
)]
#[case::environment_url_beats_the_environment_host(
json!({}),
&[("REDIS_HOST", "ignored"), ("REDIS_URL", "redis://from-url:6379")],
json!({"host": "ignored", "url": "redis://from-url:6379"}),
&[("host", "REDIS_HOST"), ("url", "REDIS_URL")],
)]
#[case::an_explicit_host_drops_the_environment_url(
json!({"host": "configured"}),
&[("REDIS_URL", "redis://environment:6379")],
json!({"host": "configured"}),
&[],
)]
#[case::environment_cluster_nodes(
json!({}),
&[("REDIS_CLUSTER_NODES", r#"[{"host": "node-a", "port": "7000"}, {"host": "node-b", "port": 7001}]"#)],
json!({"startup_nodes": "[{\"host\": \"node-a\", \"port\": \"7000\"}, {\"host\": \"node-b\", \"port\": 7001}]"}),
&[("startup_nodes", "REDIS_CLUSTER_NODES")],
)]
#[case::sentinel_from_the_environment(
json!({}),
&[("REDIS_SENTINEL_NODES", r#"[["sentinel-a", 26379]]"#), ("REDIS_SENTINEL_PASSWORD", "sentinel-secret"), ("REDIS_SERVICE_NAME", "primary")],
json!({"sentinel_nodes": "[[\"sentinel-a\", 26379]]", "sentinel_password": "sentinel-secret", "service_name": "primary"}),
&[("sentinel_nodes", "REDIS_SENTINEL_NODES"), ("sentinel_password", "REDIS_SENTINEL_PASSWORD"), ("service_name", "REDIS_SERVICE_NAME")],
)]
#[case::an_explicit_credential_provider_beats_cloud_auth(
json!({"aws_iam_auth": true, "credential_provider": true, "gcp_service_account": "account", "host": "redis", "password": "dropped"}),
&[("REDIS_AZURE_AD_TOKEN", "true")],
json!({"aws_iam_auth": true, "credential_provider": true, "gcp_service_account": "account", "host": "redis", "password": "dropped"}),
&[],
)]
#[case::gcp_beats_azure_and_aws(
json!({"aws_iam_auth": true, "azure_redis_ad_token": "True", "host": "redis", "password": "dropped"}),
&[("REDIS_GCP_SERVICE_ACCOUNT", "projects/-/serviceAccounts/sa@project.iam.gserviceaccount.com")],
json!({"gcp_service_account": "projects/-/serviceAccounts/sa@project.iam.gserviceaccount.com", "aws_iam_auth": true, "azure_redis_ad_token": "True", "host": "redis", "password": "dropped"}),
&[("gcp_service_account", "REDIS_GCP_SERVICE_ACCOUNT")],
)]
#[case::azure_beats_aws_and_takes_the_environment_username(
json!({"aws_iam_auth": true, "azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": true, "azure_tenant_id": "tenant", "host": "redis", "username": "dropped"}),
&[("REDIS_USERNAME", "environment-user")],
json!({"aws_iam_auth": true, "azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": true, "azure_tenant_id": "tenant", "host": "redis", "username": "dropped"}),
&[],
)]
#[case::an_empty_environment_username_sends_only_the_token(
json!({"azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": true, "azure_tenant_id": "tenant", "host": "redis"}),
&[("REDIS_USERNAME", "")],
json!({"azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": true, "azure_tenant_id": "tenant", "host": "redis", "username": ""}),
&[("username", "REDIS_USERNAME")],
)]
#[case::aws_iam_over_tls_with_the_default_region(
json!({"aws_iam_auth": "yes", "aws_iam_cache_name": "cache-1", "aws_iam_user_name": "iam-user", "host": "cache-1", "password": "dropped", "ssl": true}),
&[("AWS_DEFAULT_REGION", "us-east-1")],
json!({"aws_iam_auth": "yes", "aws_iam_cache_name": "cache-1", "aws_iam_user_name": "iam-user", "host": "cache-1", "password": "dropped", "ssl": true, "aws_iam_region": "us-east-1"}),
&[("aws_iam_region", "AWS_DEFAULT_REGION")],
)]
#[case::aws_region_beats_the_default_region(
json!({"aws_iam_auth": true, "aws_iam_cache_name": "cache-1", "aws_iam_user_name": "iam-user", "host": "cache-1", "ssl": true}),
&[("AWS_DEFAULT_REGION", "us-east-1"), ("AWS_REGION", "eu-central-1")],
json!({"aws_iam_auth": true, "aws_iam_cache_name": "cache-1", "aws_iam_user_name": "iam-user", "host": "cache-1", "ssl": true, "aws_iam_region": "eu-central-1"}),
&[("aws_iam_region", "AWS_REGION")],
)]
#[case::a_null_cert_requirement_skips_verification(
json!({"host": "redis", "ssl": true, "ssl_cert_reqs": null}),
&[("REDIS_SSL_CERT_REQS", "required")],
json!({"host": "redis", "ssl": true, "ssl_cert_reqs": null}),
&[],
)]
#[case::a_cert_requirement_spelled_like_python_none_fails(
json!({"host": "redis"}),
&[("REDIS_SSL", "true"), ("REDIS_SSL_CERT_REQS", "None")],
json!({"ssl": true, "ssl_cert_reqs": "None", "host": "redis"}),
&[("ssl", "REDIS_SSL"), ("ssl_cert_reqs", "REDIS_SSL_CERT_REQS")],
)]
#[case::tls_options_from_the_environment(
json!({"host": "redis"}),
&[("REDIS_SSL", "true"), ("REDIS_SSL_CA_CERTS", "/env-ca.pem"), ("REDIS_SSL_CERT_REQS", "optional")],
json!({"ssl_ca_certs": "/env-ca.pem", "ssl": true, "ssl_cert_reqs": "optional", "host": "redis"}),
&[("ssl", "REDIS_SSL"), ("ssl_ca_certs", "REDIS_SSL_CA_CERTS"), ("ssl_cert_reqs", "REDIS_SSL_CERT_REQS")],
)]
#[case::environment_references_resolve(
json!({"host": "os.environ/COORD_HOST", "password": "os.environ/COORD_PASSWORD"}),
&[("COORD_HOST", "referenced"), ("COORD_PASSWORD", "referenced-secret")],
json!({"host": "referenced", "password": "referenced-secret"}),
&[],
)]
#[case::malformed_port_from_the_environment(
json!({}),
&[("REDIS_HOST", "redis"), ("REDIS_PORT", "abc")],
json!({"port": "abc", "host": "redis"}),
&[("host", "REDIS_HOST"), ("port", "REDIS_PORT")],
)]
#[case::retry_on_error_from_the_environment(
json!({"host": "redis"}),
&[("REDIS_RETRY_ON_ERROR", "true")],
json!({"retry_on_error": true, "host": "redis"}),
&[("retry_on_error", "REDIS_RETRY_ON_ERROR")],
)]
#[case::ssl_ca_path_from_the_environment(
json!({"host": "redis"}),
&[("REDIS_SSL", "true"), ("REDIS_SSL_CA_PATH", "/etc/ssl/private-cas")],
json!({"ssl": true, "ssl_ca_path": "/etc/ssl/private-cas", "host": "redis"}),
&[("ssl", "REDIS_SSL"), ("ssl_ca_path", "REDIS_SSL_CA_PATH")],
)]
fn merges_like_python_redis(
#[case] kwargs: serde_json::Value,
#[case] environment: &[(&str, &str)],
#[case] merged: serde_json::Value,
#[case] sources: &[(&str, &str)],
) {
let configured = redis_cache_kwargs(kwargs);
let lookup = |name: &str| {
environment
.iter()
.find(|(key, _)| *key == name)
.map(|(_, value)| (*value).to_owned())
};
let kwargs = configured
.merged(&SecretSnapshot::read(configured.secret_names(), &lookup))
.unwrap();
let values: serde_json::Map<String, serde_json::Value> = kwargs
.iter()
.map(|(kwarg, entry)| (kwarg.clone(), entry.value.clone()))
.collect();
let renamed: Vec<(&str, &str)> = kwargs
.iter()
.filter(|(kwarg, entry)| *kwarg != &entry.source)
.map(|(kwarg, entry)| (kwarg.as_str(), entry.source.as_str()))
.collect();
assert_eq!(serde_json::Value::Object(values), merged);
assert_eq!(renamed, sources);
}

5
litellm-rust/crates/host-python/src/binding.rs
View file

@ -28,6 +28,11 @@ pub trait PythonBinding: PythonOwned {
) -> PyResult<Py<PyAny>>;
/// What the stream carries at hand-off, as the caller's stream receives it.
/// The unary response after the deployment transform, before success hooks run.
fn deliver_response(&mut self, _py: Python<'_>, _response: &Bound<'_, PyAny>) -> PyResult<()> {
Ok(())
}
fn encode_stream_head(
&mut self,
py: Python<'_>,

12
litellm-rust/crates/host-python/src/driver.rs
View file

@ -399,8 +399,16 @@ where
Ok(step) => step,
Err(error) => return self.failure(error, FailureOrigin::Host).await,
};
match self.settle(step).await {
Ok(response) => self.succeeded(response).await,
let response = match self.settle(step).await {
Ok(response) => response,
Err(error) => return self.failure(error, FailureOrigin::Host).await,
};
let delivered = {
let mut parts = self.parts();
Python::attach(|py| parts.binding.deliver_response(py, response.bind(py)))
};
match delivered {
Ok(()) => self.succeeded(response).await,
Err(error) => self.failure(error, FailureOrigin::Host).await,
}
}

4
litellm-rust/crates/host-python/src/lib.rs
View file

@ -36,8 +36,8 @@ pub use marshal::{
pub use owned::PythonOwned;
pub use runtime::{
ForkedAfterNativeRuntimeStarted, ProcessReservedForForking, enter_native, poll_async_value,
ready_future, reserve_process_for_forking, run_async, run_async_value, run_sync,
run_sync_value, runtime_started,
ready_future, reserve_process_for_forking, run_async, run_async_value, run_isolated, run_sync,
run_sync_value, runtime_handle, runtime_started,
};
pub use services::PythonHostCalls;

34
litellm-rust/crates/host-python/src/runtime.rs
View file

@ -68,6 +68,23 @@ fn runtime() -> PyResult<&'static Runtime> {
Ok(pyo3_async_runtimes::tokio::get_runtime())
}
/// Runs `future` on a current-thread runtime built and shut down inside the call, its blocking pool joined before
/// it returns. No thread outlives the call, so unlike every other door this one never claims the process in the fork
/// gate: a process that only used it can still fork. The future must not hold resources tied to another runtime
pub fn run_isolated<F: Future>(future: F) -> std::io::Result<F::Output> {
let runtime = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()?;
let output = runtime.block_on(future);
runtime.shutdown_timeout(Duration::from_secs(5));
Ok(output)
}
/// The native runtime, for work that must outlive the call that schedules it.
pub fn runtime_handle() -> PyResult<Handle> {
Ok(runtime()?.handle().clone())
}
#[expect(clippy::disallowed_methods, reason = "this is the gated door")]
fn future_into_py<F, T>(py: Python<'_>, future: F) -> PyResult<Bound<'_, PyAny>>
where
@ -813,4 +830,21 @@ asyncio.run(exercise())
.expect("async route contract should hold");
});
}
#[rstest]
fn an_isolated_run_joins_blocking_work_it_left_running_before_returning() {
let finished = Arc::new(AtomicBool::new(false));
let flag = Arc::clone(&finished);
let output = run_isolated(async move {
drop(tokio::task::spawn_blocking(move || {
thread::sleep(Duration::from_millis(50));
flag.store(true, Ordering::SeqCst);
}));
7
})
.unwrap();
assert_eq!(output, 7);
assert!(finished.load(Ordering::SeqCst));
}
}

9
litellm-rust/crates/host-python/tests/isolated_run.rs Normal file
View file

@ -0,0 +1,9 @@
use litellm_host_python::{reserve_process_for_forking, run_isolated, runtime_started};
#[test]
fn an_isolated_run_leaves_the_process_free_to_reserve_for_forking() {
run_isolated(async { tokio::task::spawn_blocking(|| ()).await.unwrap() }).unwrap();
assert!(!runtime_started());
assert!(reserve_process_for_forking().is_ok());
}

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

@ -16,6 +16,15 @@ pub enum TlsSource {
#[error(transparent)]
pub struct Error(Arc<Repr>);
// A streamed body the connection ended early, after the response head arrived; `expected` is
// the declared `content-length`, `None` for a chunked body.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct BodyCut {
pub received: u64,
pub expected: Option<u64>,
pub reset: bool,
}
#[derive(Debug, thiserror::Error)]
pub(crate) enum Repr {
#[error("could not read {}", path.display())]
@ -56,6 +65,12 @@ pub(crate) enum Repr {
Connect(#[source] reqwest::Error),
#[error("upstream network error")]
Network(#[source] reqwest::Error),
#[error("upstream network error")]
BodyCut {
#[source]
source: reqwest::Error,
cut: BodyCut,
},
#[error("upstream request timed out")]
Timeout(#[source] reqwest::Error),
/// No bytes arrived within the read timeout, which the caller enforces per read; `body` says
@ -201,6 +216,37 @@ impl Error {
matches!(&*self.0, Repr::ReadTimeout { body: true, .. })
}
pub fn body_read(error: reqwest::Error, received: u64, expected: Option<u64>) -> Self {
let mut source: Option<&(dyn std::error::Error + 'static)> = Some(&error);
let kind = std::iter::from_fn(|| {
let current = source?;
source = current.source();
Some(current)
})
.find_map(|cause| cause.downcast_ref::<io::Error>().map(io::Error::kind));
let reset = match kind {
Some(io::ErrorKind::UnexpectedEof) => false,
Some(io::ErrorKind::ConnectionReset) => true,
_ => return error.into(),
};
Repr::BodyCut {
source: error.without_url(),
cut: BodyCut {
received,
expected,
reset,
},
}
.into()
}
pub fn body_cut(&self) -> Option<BodyCut> {
match &*self.0 {
Repr::BodyCut { cut, .. } => Some(*cut),
_ => None,
}
}
pub fn computed_header(name: impl Into<String>) -> Self {
Repr::ComputedHeader(name.into()).into()
}
@ -226,6 +272,7 @@ impl Error {
pub fn dispatched(&self) -> bool {
match &*self.0 {
Repr::Network(_)
| Repr::BodyCut { .. }
| Repr::Timeout(_)
| Repr::ReadTimeout { .. }
| Repr::Status { .. }
@ -285,6 +332,7 @@ impl Error {
| Repr::Signature(_)
| Repr::Connect(_)
| Repr::Network(_)
| Repr::BodyCut { .. }
| Repr::Timeout(_)
| Repr::ReadTimeout { .. }
| Repr::Status { .. }
@ -318,6 +366,7 @@ impl Error {
| Repr::Signature(_)
| Repr::Connect(_)
| Repr::Network(_)
| Repr::BodyCut { .. }
| Repr::Timeout(_)
| Repr::ReadTimeout { .. }
| Repr::Status { .. }
@ -355,6 +404,7 @@ impl Error {
| Repr::Signature(_)
| Repr::Connect(_)
| Repr::Network(_)
| Repr::BodyCut { .. }
| Repr::Timeout(_)
| Repr::ReadTimeout { .. }
| Repr::Status { .. }
@ -394,6 +444,7 @@ impl Error {
| Repr::UnsignableHeader(_)
| Repr::Signature(_)
| Repr::Network(_)
| Repr::BodyCut { .. }
| Repr::Timeout(_)
| Repr::ReadTimeout { .. }
| Repr::Status { .. }
@ -437,6 +488,7 @@ impl Error {
| Repr::Signature(_)
| Repr::Connect(_)
| Repr::Network(_)
| Repr::BodyCut { .. }
| Repr::Status { .. }
| Repr::WebSocketUrl(_)
| Repr::WebSocketScheme(_)
@ -473,6 +525,7 @@ impl Error {
| Repr::Signature(_)
| Repr::Connect(_)
| Repr::Network(_)
| Repr::BodyCut { .. }
| Repr::Timeout(_)
| Repr::ReadTimeout { .. }
| Repr::WebSocketUrl(_)
@ -513,6 +566,7 @@ impl Error {
| Repr::Signature(_)
| Repr::Connect(_)
| Repr::Network(_)
| Repr::BodyCut { .. }
| Repr::WebSocketUrl(_)
| Repr::WebSocketScheme(_)
| Repr::WebSocketSchemeRewrite(_)
@ -548,6 +602,7 @@ impl Error {
| Repr::Signature(_)
| Repr::Connect(_)
| Repr::Network(_)
| Repr::BodyCut { .. }
| Repr::Timeout(_)
| Repr::ReadTimeout { .. }
| Repr::Status { .. }
@ -742,4 +797,53 @@ mod tests {
assert_eq!(error.upstream_response(), Some((429, body.as_str())));
assert_eq!(error.status(), Some(429));
}
#[rstest]
#[case::chunked("transfer-encoding: chunked\r\n\r\n5\r\nhello\r\n", None)]
#[case::content_length("content-length: 100\r\n\r\nhello", Some(100))]
#[tokio::test]
async fn a_connection_that_ends_mid_body_is_a_body_cut(
#[case] tail: &'static str,
#[case] expected: Option<u64>,
) {
use std::io::{Read as _, Write as _};
let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind");
let address = listener.local_addr().expect("address");
let server = std::thread::spawn(move || {
let (mut stream, _) = listener.accept().expect("accept");
let mut request = [0_u8; 1024];
let _ = stream.read(&mut request).expect("request");
stream
.write_all(format!("HTTP/1.1 200 OK\r\n{tail}").as_bytes())
.expect("response");
});
let mut response = reqwest::Client::builder()
.no_proxy()
.build()
.expect("client")
.get(format!("http://{address}/"))
.send()
.await
.expect("head");
server.join().expect("server");
let mut received = 0_u64;
let error = loop {
match response.chunk().await {
Ok(Some(chunk)) => received += chunk.len() as u64,
Ok(None) => panic!("the body must end with an error"),
Err(error) => break crate::Error::body_read(error, received, expected),
}
};
assert_eq!(
error.body_cut(),
Some(super::BodyCut {
received: 5,
expected,
reset: false,
})
);
assert!(error.dispatched());
assert_eq!(error.to_string(), "upstream network error");
}
}

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

@ -18,7 +18,7 @@ pub mod websocket;
pub use client::Client;
pub use config::{ClientIdentity, HttpClientConfig, Resolution, Verify};
pub use error::{Error, TlsSource};
pub use error::{BodyCut, Error, TlsSource};
pub use pool::{ClientVariant, HttpClientPool};
pub use proxy::EnvironmentProxies;
pub use settings::{HttpSettings, HttpSettingsLayer, SslVerify, TcpKeepalive};

70
litellm-rust/crates/http/src/pool.rs
View file

@ -55,36 +55,13 @@ impl HttpClientPool {
config: &HttpClientConfig,
variant: ClientVariant,
) -> Result<Client, Error> {
let effective = match variant {
ClientVariant::Media => HttpClientConfig {
client_certificate: None,
proxies: EnvironmentProxies::default(),
..config.clone()
},
ClientVariant::UnpinnedMedia => HttpClientConfig {
client_certificate: None,
..config.clone()
},
ClientVariant::WebSocket => HttpClientConfig {
http2: false,
verify: match &config.verify {
Verify::BuiltInRoots => Verify::NativeRoots,
verify => verify.clone(),
},
..config.clone()
},
ClientVariant::Provider | ClientVariant::NoRedirect => config.clone(),
};
let key = (effective, variant);
let key = (effective(config, variant), variant);
if let Some(pooled) = self.lock().get(&key)
&& pooled.built_at.elapsed() < self.ttl
{
return Ok(Client::new(pooled.client.clone()));
}
let client = self
.apply(variant, reqwest::ClientBuilder::try_from(&key.0)?)
.build()
.map_err(Repr::Client)?;
let client = self.build(&key.0, variant)?;
self.lock().insert(
key,
PooledClient {
@ -95,6 +72,26 @@ impl HttpClientPool {
Ok(Client::new(client))
}
/// A client nobody else holds, for a caller that drives it on a runtime of its own and drops both together
pub fn unpooled(
&self,
config: &HttpClientConfig,
variant: ClientVariant,
) -> Result<Client, Error> {
self.build(&effective(config, variant), variant)
.map(Client::new)
}
fn build(
&self,
config: &HttpClientConfig,
variant: ClientVariant,
) -> Result<reqwest::Client, Error> {
self.apply(variant, reqwest::ClientBuilder::try_from(config)?)
.build()
.map_err(|error| Error::from(Repr::Client(error)))
}
fn lock(&self) -> MutexGuard<'_, Clients> {
self.clients.lock().unwrap_or_else(PoisonError::into_inner)
}
@ -116,6 +113,29 @@ impl HttpClientPool {
}
}
fn effective(config: &HttpClientConfig, variant: ClientVariant) -> HttpClientConfig {
match variant {
ClientVariant::Media => HttpClientConfig {
client_certificate: None,
proxies: EnvironmentProxies::default(),
..config.clone()
},
ClientVariant::UnpinnedMedia => HttpClientConfig {
client_certificate: None,
..config.clone()
},
ClientVariant::WebSocket => HttpClientConfig {
http2: false,
verify: match &config.verify {
Verify::BuiltInRoots => Verify::NativeRoots,
verify => verify.clone(),
},
..config.clone()
},
ClientVariant::Provider | ClientVariant::NoRedirect => config.clone(),
}
}
#[cfg(test)]
mod tests {
use std::{

83
litellm-rust/crates/inference-messages/src/common_utils.rs
View file

@ -1,5 +1,5 @@
use litellm_core_utils::get_llm_provider_logic::LlmProviders;
use litellm_http::request::string_headers as shared_string_headers;
use litellm_llms::ErrorDetail;
use litellm_llms::tencent::messages::transformation::TENCENT_MESSAGES_CONFIG;
use litellm_llms::{
anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG,
@ -19,8 +19,6 @@ use serde_json::{Map, Value};
use super::Error;
const HEADER_CONTEXT: &str = "messages";
#[derive(Clone, Copy)]
pub(crate) struct MessagesProvider {
provider: LlmProviders,
@ -75,7 +73,52 @@ pub(crate) fn messages_provider(provider: LlmProviders, model: &str) -> Option<M
pub(super) fn string_headers(
extra_headers: Option<Map<String, Value>>,
) -> Result<Vec<(String, String)>, Error> {
shared_string_headers(HEADER_CONTEXT, extra_headers).map_err(Error::from)
extra_headers
.unwrap_or_default()
.into_iter()
.map(|(name, value)| match value {
Value::String(value) => Ok((name, value)),
value => Err(Error::invalid_request(ErrorDetail::invalid(
"extra_headers",
NonStringHeader {
actual: python_type_name(&value),
name,
},
))),
})
.collect()
}
/// A header value the caller set to something other than a string, named by its Python type
/// since the caller wrote it in Python.
#[derive(Debug)]
struct NonStringHeader {
name: String,
actual: &'static str,
}
impl std::fmt::Display for NonStringHeader {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(
formatter,
"header {} must be a string, got {}",
self.name, self.actual
)
}
}
impl std::error::Error for NonStringHeader {}
fn python_type_name(value: &Value) -> &'static str {
match value {
Value::Null => "NoneType",
Value::Bool(_) => "bool",
Value::Number(number) if number.is_f64() => "float",
Value::Number(_) => "int",
Value::String(_) => "str",
Value::Array(_) => "list",
Value::Object(_) => "dict",
}
}
#[cfg(test)]
@ -85,7 +128,6 @@ mod tests {
use rstest::rstest;
use super::{messages_provider, string_headers};
use crate::Error;
use litellm_core_utils::get_llm_provider_logic::LlmProviders;
#[rstest]
@ -113,17 +155,28 @@ mod tests {
assert!(messages_provider(provider, model).is_none());
}
#[test]
fn string_headers_rejects_non_string_values() {
let headers = json!({"x-count": 3}).as_object().unwrap().clone();
let err = string_headers(Some(headers)).expect_err("non-string header rejected");
#[rstest]
#[case::int(json!(12345), "int")]
#[case::float(json!(1.5), "float")]
#[case::bool(json!(true), "bool")]
#[case::null(json!(null), "NoneType")]
#[case::list(json!([1]), "list")]
#[case::dict(json!({"a": 1}), "dict")]
fn non_string_header_is_a_bad_request_naming_the_header_and_its_type(
#[case] value: serde_json::Value,
#[case] python_type: &str,
) {
let headers = json!({"x-kept": "v", "x-count": value})
.as_object()
.unwrap()
.clone();
let error = string_headers(Some(headers)).expect_err("non-string header refused");
assert!(error.is_request(), "{error:?}");
assert_eq!(error.upstream_response(), None);
let rejection = error.rejected_request().expect("a rejected request");
assert_eq!(
err,
Error::Headers(litellm_http::request::HeaderError {
context: "messages",
name: "x-count".to_string(),
actual: "number",
})
litellm_error::Report::new(rejection).to_string(),
format!("invalid extra_headers: header x-count must be a string, got {python_type}")
);
}
}

64
litellm-rust/crates/inference-messages/src/handler.rs
View file

@ -26,10 +26,10 @@ use litellm_llms_types::formats::messages::{
};
use litellm_secrets::source::Secrets;
use litellm_tracing::ByteChunk;
use serde_json::Value;
use serde_json::{Map, Value};
use super::{
Error, MessagesCall, MessagesCallResponse, MessagesRoute,
Error, MessagesBody, MessagesCall, MessagesCallResponse, MessagesRoute,
common_utils::{MessagesProvider, merge_messages_header_sources, string_headers},
constants::MESSAGES_TIMEOUT_SECS,
route::Messages,
@ -85,7 +85,10 @@ pub(super) async fn execute(
),
};
let reported = match &output {
MessagesCallResponse::Complete(response) => provider.config().reported_cost(response),
MessagesCallResponse::Complete(response) => match response {
MessagesBody::Message(message) => provider.config().reported_cost(message),
MessagesBody::Unchecked(_) => None,
},
MessagesCallResponse::Stream { .. } => None,
};
context
@ -130,9 +133,10 @@ async fn provider_call(
)?;
let transform_context =
MessagesTransformContext::new(shaping.capabilities, shaping.settings.drop_params);
let filtered = requested_optional_params(sanitized, provider, &transform_context)?;
let trimmed =
without_additional_drop_params(filtered, &shaping.settings.additional_drop_params)?;
let without_listed =
without_additional_drop_params(sanitized, &shaping.settings.additional_drop_params)?;
let trimmed = requested_optional_params(without_listed, provider, &transform_context)?;
let requested_params = serde_json::to_value(&trimmed.params).map_err(serialize_failure)?;
let scoped =
get_provider_specific_headers(provider_specific_header.as_ref(), provider.as_str());
@ -173,7 +177,7 @@ async fn provider_call(
let request_context = RequestContext {
model: body.model.clone(),
custom_llm_provider: provider.as_str().to_string(),
optional_params: serde_json::to_value(&body.params).map_err(serialize_failure)?,
optional_params: requested_params,
secret_fields: Vec::new(),
api_key: api_key.map(SecretValue::new),
};
@ -346,8 +350,7 @@ async fn call_provider(
// body it cannot parse, and the public mapping classifies it by that status.
return Err(litellm_http::Error::upstream(status, text).into());
}
decode_response(config, &identity.model, &text)
.map(|message| MessagesCallResponse::Complete(Box::new(message)))
decode_response(config, &identity.model, &text).map(MessagesCallResponse::Complete)
}
fn serialize_failure(err: serde_json::Error) -> Error {
@ -383,13 +386,20 @@ fn timed_chunks(
response: reqwest::Response,
timeout: Duration,
) -> impl futures_util::Stream<Item = Result<Bytes, litellm_http::Error>> + Send + 'static {
futures_util::stream::try_unfold(response, move |mut response| async move {
let chunk = tokio::time::timeout(timeout, response.chunk())
.await
.map_err(litellm_http::Error::body_read_timeout)?
.map_err(litellm_http::Error::from)?;
Ok(chunk.map(|chunk| (chunk, response)))
})
let expected = response.content_length();
futures_util::stream::try_unfold(
(response, 0_u64),
move |(mut response, received)| async move {
let chunk = tokio::time::timeout(timeout, response.chunk())
.await
.map_err(litellm_http::Error::body_read_timeout)?
.map_err(|error| litellm_http::Error::body_read(error, received, expected))?;
Ok(chunk.map(|chunk| {
let received = received + chunk.len() as u64;
(chunk, (response, received))
}))
},
)
}
/// The provider's error, its body read under `timeout`: a stream has no overall deadline once
@ -411,15 +421,23 @@ fn decode_response(
config: &dyn BaseMessagesConfig,
model: &str,
text: &str,
) -> Result<MessagesResponse, Error> {
let response = serde_json::from_str(text).map_err(|err| {
Error::invalid_response(litellm_llms::ErrorDetail::invalid(
"messages response JSON",
err,
))
})?;
) -> Result<MessagesBody, Error> {
let response = match serde_json::from_str::<MessagesResponse>(text) {
Ok(response) => response,
Err(err) => {
let Ok(mut body) = serde_json::from_str::<Map<String, Value>>(text) else {
return Err(Error::invalid_response(litellm_llms::ErrorDetail::invalid(
"messages response JSON",
err,
)));
};
body.entry("stop_sequence").or_insert(Value::Null);
return Ok(MessagesBody::Unchecked(body));
}
};
config
.transform_anthropic_messages_response(model, response)
.map(|message| MessagesBody::Message(Box::new(message)))
.map_err(Error::from)
.map_err(Error::after_send)
}

4
litellm-rust/crates/inference-messages/src/lib.rs
View file

@ -17,8 +17,8 @@ use crate::common_utils::messages_provider;
pub use common_utils::merge_messages_header_sources;
pub use litellm_inference::Error;
pub use types::{
MessagesCall, MessagesCallResponse, MessagesSettings, MessagesShaping, litellm_params,
messages_body,
MessagesBody, MessagesCall, MessagesCallResponse, MessagesSettings, MessagesShaping,
litellm_params, messages_body,
};
#[derive(Clone)]

7
litellm-rust/crates/inference-messages/src/route.rs
View file

@ -5,11 +5,10 @@ use litellm_host::{
call::{HostedCompletion, HostedMachine, hosted_call},
protocol::Protocol,
};
use litellm_llms_types::formats::messages::MessagesResponse;
use super::{Error, MessagesCall};
use super::{Error, MessagesBody, MessagesCall};
pub type MessagesOutput = HostedCompletion<Box<MessagesResponse>>;
pub type MessagesOutput = HostedCompletion<MessagesBody>;
/// The upstream response as the caller sees it at stream hand-off, before any chunk.
pub struct MessagesStreamHead {
@ -19,7 +18,7 @@ pub struct MessagesStreamHead {
pub struct Messages;
impl Protocol for Messages {
type Response = Box<MessagesResponse>;
type Response = MessagesBody;
type Error = Error;
type Request = MessagesCall;
type HostCall = Infallible;

12
litellm-rust/crates/inference-messages/src/types.rs
View file

@ -42,7 +42,16 @@ pub(super) fn invalid_request(err: serde_json::Error) -> Error {
}
pub type MessagesCallResponse =
CallOutput<Box<MessagesResponse>, super::route::MessagesStreamHead, Bytes, Error>;
CallOutput<MessagesBody, super::route::MessagesStreamHead, Bytes, Error>;
// Python's `transform_anthropic_messages_response` returns any JSON object the provider sends
// with a 200, message or not, so one that is not a message passes through as sent.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum MessagesBody {
Message(Box<MessagesResponse>),
Unchecked(Map<String, Value>),
}
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct MessagesShaping {
@ -144,6 +153,7 @@ mod tests {
supports_reasoning: true,
supports_adaptive_thinking: true,
thinking_always_on: false,
thinking_on_by_default: false,
supports_legacy_thinking: false,
supports_output_config: true,
supports_sampling_params: false,

9
litellm-rust/crates/inference-messages/tests/bedrock_native/main.rs
View file

@ -1,7 +1,9 @@
use std::{sync::Arc, time::Duration};
use futures_util::TryStreamExt;
use litellm_inference_messages::{Error, MessagesCall, MessagesCallResponse, MessagesShaping};
use litellm_inference_messages::{
Error, MessagesBody, MessagesCall, MessagesCallResponse, MessagesShaping,
};
use litellm_inference_testing::{RecordingSecrets, no_secrets};
use litellm_llms_types::formats::messages::{MessagesRequest, MessagesResponse};
use rstest::{fixture, rstest};
@ -74,7 +76,10 @@ async fn run(call: MessagesCall) -> Result<MessagesCallResponse, Error> {
async fn run_message(call: MessagesCall) -> MessagesResponse {
match run(call).await.unwrap() {
MessagesCallResponse::Complete(message) => *message,
MessagesCallResponse::Complete(message) => match message {
MessagesBody::Message(message) => *message,
MessagesBody::Unchecked(_) => panic!("expected a message"),
},
MessagesCallResponse::Stream { .. } => panic!("expected a completed message"),
}
}

72
litellm-rust/crates/inference-messages/tests/caching.rs
View file

@ -50,42 +50,47 @@ async fn messages_cache_identity_includes_provider_native_parameters(
#[case] original: Value,
#[case] changed: Value,
) {
use litellm_inference_messages::route::Messages;
use litellm_inference_messages::{MessagesBody, route::Messages};
use litellm_llms_types::formats::messages::MessagesResponse;
let calls = AtomicUsize::new(0);
for (value, expected_call) in [(original.clone(), 0), (changed, 1), (original, 0)] {
let response =
execute_unary::<Messages, _, _>(
CacheRequest::from_wire(
ProviderIdentity {
model: "test".into(),
provider: "anthropic".into(),
},
Some(&WireRequest {
url: "https://example.test/v1/messages".into(),
headers: vec![],
body: json!({
"model":"test", "messages":[{"role":"user","content":"hello"}],
"max_tokens":32, (field):value
}),
}),
),
Some(cache.clone()),
Some(CacheOptions::new(CacheScope::Shared)),
&(),
None,
|| async {
let call = calls.fetch_add(1, Ordering::SeqCst);
Ok(Box::new(serde_json::from_value::<MessagesResponse>(json!({
"id":call.to_string(), "type":"message", "role":"assistant", "model":"test",
"content":[{"type":"text","text":format!("answer {call}")}],
"stop_reason":"end_turn", "stop_sequence":null
})).unwrap()))
let response = execute_unary::<Messages, _, _>(
CacheRequest::from_wire(
ProviderIdentity {
model: "test".into(),
provider: "anthropic".into(),
},
)
.await
.unwrap();
Some(&WireRequest {
url: "https://example.test/v1/messages".into(),
headers: vec![],
body: json!({
"model":"test", "messages":[{"role":"user","content":"hello"}],
"max_tokens":32, (field):value
}),
}),
),
Some(cache.clone()),
Some(CacheOptions::new(CacheScope::Shared)),
&(),
None,
|| async {
let call = calls.fetch_add(1, Ordering::SeqCst);
Ok(MessagesBody::Message(Box::new(
serde_json::from_value::<MessagesResponse>(json!({
"id":call.to_string(), "type":"message", "role":"assistant", "model":"test",
"content":[{"type":"text","text":format!("answer {call}")}],
"stop_reason":"end_turn", "stop_sequence":null
}))
.unwrap(),
)))
},
)
.await
.unwrap();
let MessagesBody::Message(response) = response else {
panic!("expected a message");
};
assert_eq!(response.id, expected_call.to_string());
assert_eq!(
response.content[0]["text"],
@ -253,7 +258,7 @@ async fn provider_reported_cost_is_delivered_once_and_never_charged_for_a_cache_
#[case] cost: Value,
#[case] expected: Option<f64>,
) {
use litellm_inference_messages::{MessagesCall, MessagesCallResponse};
use litellm_inference_messages::{MessagesBody, MessagesCall, MessagesCallResponse};
let response = json!({"id": "msg_cost", "type": "message", "role": "assistant", "model": "native-test",
"content": [{"type": "text", "text": "answer"}], "stop_reason": "end_turn", "stop_sequence": null,
"usage": {"input_tokens": 3, "output_tokens": 2}, "cost": cost});
@ -277,6 +282,9 @@ async fn provider_reported_cost_is_delivered_once_and_never_charged_for_a_cache_
panic!("expected Messages responses");
};
assert_eq!(first, second);
let MessagesBody::Message(first) = first else {
panic!("expected a message");
};
assert_eq!(first.extra.get("cost"), Some(&cost));
assert_eq!(support::received(&upstream).await.len(), 1);
let facts = hooks.facts.lock().unwrap();

9
litellm-rust/crates/inference-messages/tests/live/support.rs
View file

@ -4,7 +4,9 @@ use futures_util::{FutureExt, TryStreamExt};
use litellm_host::interceptors::{
ExecutionFacts, Interceptors, RawResponse, RequestContext, ResultSource, WireRequest,
};
use litellm_inference_messages::{Error, MessagesCall, MessagesCallResponse, MessagesRoute};
use litellm_inference_messages::{
Error, MessagesBody, MessagesCall, MessagesCallResponse, MessagesRoute,
};
use litellm_inference_testing::live::{LiveResources, model, within_deadline};
use litellm_llms_types::formats::messages::MessagesResponse;
use serde_json::{Value, json};
@ -111,7 +113,10 @@ pub(super) fn call(provider: &str, payload: Value) -> MessagesCall {
pub(super) fn complete(response: MessagesCallResponse) -> Box<MessagesResponse> {
match response {
MessagesCallResponse::Complete(message) => message,
MessagesCallResponse::Complete(message) => match message {
MessagesBody::Message(message) => message,
MessagesBody::Unchecked(_) => panic!("a complete request returned no message"),
},
MessagesCallResponse::Stream { .. } => panic!("a complete request returned a stream"),
}
}

92
litellm-rust/crates/inference-messages/tests/messages/host.rs
View file

@ -6,9 +6,11 @@ use litellm_host::{
interceptors::{ExecutionFacts, RequestContext, ResultSource, WireRequest},
lifecycle::CallEvent,
};
use litellm_inference_messages::{MessagesCallResponse, route::Messages};
use litellm_inference_messages::{MessagesBody, MessagesCallResponse, route::Messages};
use litellm_inference_testing::{RecordingSecrets, no_secrets};
use litellm_llms::base_llm::messages::context::MessagesModelCapabilities as AnthropicModelCapabilities;
use litellm_llms::base_llm::messages::context::{
MessagesModelCapabilities as AnthropicModelCapabilities, SupportedEffortTiers,
};
use rstest::rstest;
use super::*;
@ -338,7 +340,10 @@ async fn response_mode_follows_the_intercepted_request(
}
MessagesCallResponse::Complete(message) => {
assert_eq!(expected_stream, Some(false));
assert_eq!(*message, serde_json::from_value(message_body()).unwrap());
assert_eq!(
message,
MessagesBody::Message(serde_json::from_value(message_body()).unwrap())
);
}
}
Ok::<_, Error>(())
@ -419,7 +424,7 @@ async fn no_raw_response_is_emitted_for_a_stream_or_a_failure(
assert!(host.raw_responses().is_empty());
}
/// Python logs `optional_params` as what it is about to send, so a dropped param must
/// Python logs `optional_params` after its requested-param gate, so a dropped param must
/// not resurface in callbacks.
#[rstest]
#[tokio::test]
@ -450,3 +455,82 @@ async fn the_request_context_carries_the_shaped_params_without_model_or_messages
.unwrap_or_else(|seen| panic!("before_provider_request runs once, saw {}", seen.len()));
assert_eq!(optional_params, json!({"max_tokens": 16}));
}
#[rstest]
#[tokio::test]
async fn the_request_context_leaves_out_a_listed_param_the_model_refuses(call: MessagesCall) {
let upstream = upstream([message_response()]).await;
let host = RecordingHost::passthrough(authenticated(
MessagesCall {
shaping: MessagesShaping {
settings: MessagesSettings {
additional_drop_params: vec!["temperature".into()],
..MessagesSettings::default()
},
capabilities: AnthropicModelCapabilities {
supports_sampling_params: false,
..AnthropicModelCapabilities::default()
},
},
..with_fields(call, json!({"temperature": 0.2}))
},
upstream.uri(),
));
run_through(&host).await.expect("messages call succeeds");
let [optional_params] = <[Value; 1]>::try_from(host.optional_params.into_inner().unwrap())
.unwrap_or_else(|seen| panic!("before_provider_request runs once, saw {}", seen.len()));
assert_eq!(optional_params, json!({"max_tokens": 16}));
}
#[rstest]
#[case::reasoning_effort_on_an_adaptive_model(
AnthropicModelCapabilities {
supports_reasoning: true,
supports_adaptive_thinking: true,
supports_output_config: true,
effort_tiers: SupportedEffortTiers { low: true, ..SupportedEffortTiers::default() },
..AnthropicModelCapabilities::default()
},
json!({"reasoning_effort": "low"})
)]
#[case::adaptive_payload_on_a_legacy_model(
AnthropicModelCapabilities {
supports_reasoning: true,
..AnthropicModelCapabilities::default()
},
json!({"max_tokens": 3000, "thinking": {"type": "adaptive"}, "output_config": {"effort": "high"}})
)]
#[tokio::test]
async fn the_request_context_carries_the_requested_reasoning_not_its_translation(
call: MessagesCall,
#[case] capabilities: AnthropicModelCapabilities,
#[case] fields: Value,
) {
let upstream = upstream([message_response()]).await;
let host = RecordingHost::passthrough(authenticated(
MessagesCall {
shaping: MessagesShaping {
capabilities,
..MessagesShaping::default()
},
..with_fields(call, fields.clone())
},
upstream.uri(),
));
run_through(&host).await.expect("messages call succeeds");
let [optional_params] = <[Value; 1]>::try_from(host.optional_params.into_inner().unwrap())
.unwrap_or_else(|seen| panic!("before_provider_request runs once, saw {}", seen.len()));
assert_eq!(
optional_params,
Value::Object(
object(json!({"max_tokens": 16}))
.into_iter()
.chain(object(fields))
.collect()
)
);
}

7
litellm-rust/crates/inference-messages/tests/messages/main.rs
View file

@ -5,7 +5,7 @@ use std::{
use litellm_http::{HttpSettings, Resolution};
use litellm_inference_messages::{
Error, MessagesCall, MessagesSettings, MessagesShaping,
Error, MessagesBody, MessagesCall, MessagesSettings, MessagesShaping,
route::{Messages, MessagesMachine, MessagesOutput},
};
use litellm_inference_testing::RecordingSecrets;
@ -119,7 +119,10 @@ async fn run(call: MessagesCall) -> Result<MessagesOutput, Error> {
async fn run_message(call: MessagesCall) -> MessagesResponse {
match run(call).await.expect("messages call succeeds") {
MessagesOutput::Complete(message) => *message,
MessagesOutput::Complete(message) => match message {
MessagesBody::Message(message) => *message,
MessagesBody::Unchecked(_) => panic!("a call returned no message"),
},
MessagesOutput::StreamEnded | MessagesOutput::Detached => {
panic!("a non-streaming call returned a stream")
}

24
litellm-rust/crates/inference-messages/tests/messages/request.rs
View file

@ -530,6 +530,30 @@ async fn unsupported_params_are_dropped_under_drop_params_and_rejected_without_i
assert_eq!(sent["max_tokens"], 16);
}
#[rstest]
#[tokio::test]
async fn additional_drop_params_clear_a_param_the_model_does_not_support(call: MessagesCall) {
let upstream = upstream([message_response()]).await;
run_message(MessagesCall {
api_key: Some("sk".into()),
api_base: Some(upstream.uri()),
shaping: MessagesShaping {
capabilities: sampling_removed(),
settings: MessagesSettings {
additional_drop_params: vec!["top_k".into()],
..MessagesSettings::default()
},
},
..with_fields(call, json!({"top_k": 5}))
})
.await;
let sent = only_request(&upstream).await.json();
assert_eq!(sent.get("top_k"), None);
assert_eq!(sent["max_tokens"], 16);
}
#[rstest]
#[case::adaptive_thinking(json!({"type": "adaptive"}), json!({"type": "adaptive", "display": "summarized"}))]
#[case::disabled_thinking(json!({"type": "disabled"}), json!({"type": "disabled"}))]

48
litellm-rust/crates/inference-messages/tests/messages/response.rs
View file

@ -2,7 +2,7 @@ use litellm_host::{
interceptors::{ExecutionFacts, ResultSource},
lifecycle::ExecutionEvent,
};
use litellm_inference_messages::{MessagesCallResponse, messages_body};
use litellm_inference_messages::{MessagesBody, MessagesCallResponse, messages_body};
use litellm_inference_testing::{
RecordingSecrets, http_config, no_secrets, provider_http, resources,
};
@ -48,6 +48,9 @@ async fn calls_defer_execution_until_polled(
let MessagesCallResponse::Complete(response) = future.await.unwrap() else {
panic!("expected a completed message");
};
let MessagesBody::Message(response) = response else {
panic!("expected a message");
};
assert_eq!(
response.content,
message_body()["content"].as_array().unwrap().as_slice()
@ -240,10 +243,42 @@ async fn an_unparseable_success_body_is_reported_as_the_upstream_answer(call: Me
assert_eq!(error.upstream_response(), Some((200, "not json")));
}
#[rstest]
#[case::error_object(
json!({"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}}),
json!({"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"},
"stop_sequence": null}),
)]
#[case::message_missing_id(
json!({"type": "message", "content": [], "stop_sequence": "###"}),
json!({"type": "message", "content": [], "stop_sequence": "###"}),
)]
#[tokio::test]
async fn a_success_body_that_is_not_a_message_passes_through(
call: MessagesCall,
#[case] body: Value,
#[case] expected: Value,
) {
let upstream = upstream([json_response(body)]).await;
let response = run(MessagesCall {
api_key: Some("sk".into()),
api_base: Some(upstream.uri()),
..call
})
.await
.expect("any JSON object is returned");
let MessagesOutput::Complete(response) = response else {
panic!("expected a completed response");
};
assert_eq!(serde_json::to_value(response).unwrap(), expected);
}
#[rstest]
#[tokio::test]
async fn an_unreadable_success_body_is_an_invalid_response(call: MessagesCall) {
let upstream = upstream([json_response(json!({"unexpected": true}))]).await;
let upstream = upstream([json_response(json!(["unexpected"]))]).await;
let error = run(MessagesCall {
api_key: Some("sk".into()),
@ -320,6 +355,9 @@ async fn the_facade_sends_through_the_injected_http_pool_configuration(call: Mes
let MessagesCallResponse::Complete(message) = response else {
panic!("a non-streaming request returns a message");
};
let MessagesBody::Message(message) = message else {
panic!("expected a message");
};
assert_eq!(message.id, "msg_1");
let sent = only_request(&upstream).await;
assert_eq!(sent.header("x-api-key"), Some("sk-ant"));
@ -412,6 +450,9 @@ async fn route_uses_injected_dependencies_and_optional_cache(
else {
panic!("expected a completed message");
};
let MessagesBody::Message(response) = response else {
panic!("expected a message");
};
assert_eq!(
response.content,
message_body()["content"].as_array().unwrap().as_slice()
@ -484,6 +525,9 @@ async fn cache_overrides_preserve_the_routes_isolated_scope(call: MessagesCall)
else {
panic!("expected a completed message");
};
let MessagesBody::Message(response) = response else {
panic!("expected a message");
};
assert_eq!(response.id, expected["id"].as_str().unwrap());
}
assert_eq!(received(&upstream).await.len(), 2);

8
litellm-rust/crates/inference/src/error.rs
View file

@ -164,6 +164,14 @@ impl Error {
}
}
pub fn body_cut(&self) -> Option<litellm_http::BodyCut> {
match self {
Self::Llm(error) => error.body_cut(),
Self::Http(error) => error.body_cut(),
_ => None,
}
}
/// `Some(connect)` when the call timed out, where `connect` says the connection was never
/// established; `None` for every other error.
pub fn timeout(&self) -> Option<bool> {

1
litellm-rust/crates/llms/src/anthropic/common_utils.rs
View file

@ -1913,6 +1913,7 @@ mod tests {
supports_reasoning: false,
supports_adaptive_thinking: false,
thinking_always_on: false,
thinking_on_by_default: false,
supports_legacy_thinking: false,
supports_output_config: false,
supports_sampling_params: true,

37
litellm-rust/crates/llms/src/anthropic/messages/thinking.rs
View file

@ -2,8 +2,8 @@ use litellm_llms_types::{
formats::{
chat_completions::ReasoningEffort,
messages::{
AdaptiveThinking, EffortLevel, MessagesOptionalParams, MessagesRequest, OutputConfig,
ThinkingConfig, ThinkingDisplay,
AdaptiveThinking, DisabledThinking, EffortLevel, MessagesOptionalParams,
MessagesRequest, OutputConfig, ThinkingConfig, ThinkingDisplay,
},
},
recognized::Recognized,
@ -172,7 +172,11 @@ fn translate_reasoning_effort(
) else {
return Ok(MessagesRequest {
params: MessagesOptionalParams {
thinking: None,
thinking: (context.capabilities.thinking_on_by_default
&& !context.capabilities.thinking_always_on)
.then(|| {
Recognized::Known(ThinkingConfig::Disabled(DisabledThinking::default()))
}),
output_config: None,
..request.params
},
@ -748,6 +752,33 @@ mod tests {
);
}
#[rstest]
#[case::default_thinking(true, false, Some(serde_json::json!({"type": "disabled"}))) ]
#[case::legacy(false, false, None)]
#[case::always_on(true, true, None)]
fn none_disables_only_optional_default_thinking(
#[case] thinking_on_by_default: bool,
#[case] thinking_always_on: bool,
#[case] thinking: Option<Value>,
) {
let capabilities = MessagesModelCapabilities {
thinking_on_by_default,
thinking_always_on,
..opus_4_7()
};
let expected = match thinking {
Some(thinking) => serde_json::json!({"max_tokens": 1024, "thinking": thinking}),
None => serde_json::json!({"max_tokens": 1024}),
};
assert_eq!(
translate(
capabilities,
serde_json::json!({"max_tokens": 1024, "reasoning_effort": "none"})
),
Ok(request(expected))
);
}
#[rstest]
#[case::adaptive_model(opus_4_7())]
#[case::effort_capable_model(opus_4_5())]

3
litellm-rust/crates/llms/src/base_llm/messages/context.rs
View file

@ -32,6 +32,8 @@ pub struct MessagesModelCapabilities {
#[serde(default)]
pub thinking_always_on: bool,
#[serde(default)]
pub thinking_on_by_default: bool,
#[serde(default)]
pub supports_legacy_thinking: bool,
#[serde(default)]
pub supports_output_config: bool,
@ -55,6 +57,7 @@ impl Default for MessagesModelCapabilities {
supports_reasoning: false,
supports_adaptive_thinking: false,
thinking_always_on: false,
thinking_on_by_default: false,
supports_legacy_thinking: false,
supports_output_config: false,
supports_sampling_params: true,

7
litellm-rust/crates/llms/src/error.rs
View file

@ -103,6 +103,13 @@ impl Error {
matches!(self, Self::Http(error) if error.is_body_read_timeout())
}
pub fn body_cut(&self) -> Option<litellm_http::BodyCut> {
match self {
Self::Http(error) => error.body_cut(),
_ => None,
}
}
pub fn is_timeout(&self) -> Option<bool> {
match self {
Self::Http(error) if error.is_timeout() => Some(error.is_connect_timeout()),

9
litellm-rust/crates/model-catalog-refresh/AGENTS.md Normal file
View file

@ -0,0 +1,9 @@
# Model catalog refresh
- `fetch` fetches one document over http(s) or `s3://` and never parses it
- Only 429, 500, 502, 503, 504 and transport errors are retried. S3 errors are final
- `Failed.retryable` means only the attempt limit ended the fetch
- A server's Retry-After overrides the backoff, capped at `max_delay`
- `Start::After(failed)` resumes the backoff schedule where another caller's first attempt left off
- Error reasons use the URL with userinfo, query and fragment stripped, since presigned URLs carry credentials
- Bodies over `MAX_BODY_BYTES` are rejected

13
litellm-rust/crates/model-catalog-refresh/Cargo.toml
View file

@ -6,14 +6,19 @@ license.workspace = true
repository.workspace = true
[dependencies]
backon.workspace = true
bytes.workspace = true
litellm-error.workspace = true
litellm-http.workspace = true
litellm-model-catalog.workspace = true
litellm-s3.workspace = true
litellm-tracing.workspace = true
reqwest.workspace = true
thiserror.workspace = true
tokio = { workspace = true, features = ["time"] }
url.workspace = true
[dev-dependencies]
litellm-error.workspace = true
aws-credential-types = "1.3.0"
litellm-http = { workspace = true, features = ["test-support"] }
reqwest.workspace = true
rstest.workspace = true
tokio.workspace = true
tokio = { workspace = true, features = ["macros", "rt-multi-thread", "time"] }

153
litellm-rust/crates/model-catalog-refresh/src/attempt.rs Normal file
View file

@ -0,0 +1,153 @@
use std::time::Duration;
use bytes::Bytes;
use litellm_http::Client;
use litellm_s3::S3Uri;
use litellm_s3::sdk::config::retry::RetryConfig;
use litellm_s3::sdk::config::timeout::TimeoutConfig;
use reqwest::header::{ETAG, HeaderMap, RETRY_AFTER};
use url::Url;
use crate::{Error, RETRYABLE_STATUSES, S3Access};
pub(crate) const MAX_BODY_BYTES: u64 = 100 * 1024 * 1024;
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Fetched {
pub body: Vec<u8>,
pub etag: Option<String>,
}
#[derive(Debug)]
pub(crate) enum Cause {
Status(u16),
Transport(Error),
}
#[derive(Debug)]
pub(crate) struct Rejected {
pub(crate) cause: Cause,
pub(crate) retryable: bool,
pub(crate) retry_after: Option<Duration>,
}
impl Rejected {
pub(crate) fn fatal(error: Error) -> Self {
Self {
cause: Cause::Transport(error),
retryable: false,
retry_after: None,
}
}
}
/// Any status below 400 is the document
pub(crate) async fn http(
client: &Client,
url: Url,
timeout: Duration,
) -> Result<Fetched, Rejected> {
let transport = |error: litellm_http::Error| Rejected {
cause: Cause::Transport(Error::from(error)),
retryable: true,
retry_after: None,
};
let mut response = tokio::time::timeout(timeout, client.get(url).send())
.await
.map_err(litellm_http::Error::read_timeout)
.and_then(|sent| sent.map_err(litellm_http::Error::from))
.map_err(transport)?;
let status = response.status().as_u16();
let etag = etag(response.headers());
let retry_after = response
.headers()
.get(RETRY_AFTER)
.and_then(|value| value.to_str().ok())
.and_then(|value| value.trim().parse().ok())
.map(Duration::from_secs);
if status >= 400 {
return Err(Rejected {
cause: Cause::Status(status),
retryable: RETRYABLE_STATUSES.contains(&status),
retry_after,
});
}
if response
.content_length()
.is_some_and(|length| length > MAX_BODY_BYTES)
{
return Err(Rejected::fatal(Error::body_too_large(MAX_BODY_BYTES)));
}
let mut body = Vec::new();
while let Some(chunk) = tokio::time::timeout(timeout, response.chunk())
.await
.map_err(litellm_http::Error::body_read_timeout)
.and_then(|read| read.map_err(litellm_http::Error::from))
.map_err(transport)?
{
append(&mut body, &chunk).map_err(Rejected::fatal)?;
}
Ok(Fetched { body, etag })
}
/// S3 failures are final: the SDK already classifies its own errors and a missing or denied object stays so
pub(crate) async fn s3(
client: &Client,
s3: &S3Access,
object: &S3Uri,
timeout: Duration,
) -> Result<Fetched, Rejected> {
let fetched = async {
let config = match s3 {
S3Access::Environment { region } => {
litellm_s3::environment_config(client.clone(), region.clone()).await
}
S3Access::Config(config) => config.to_builder(),
}
.retry_config(RetryConfig::disabled())
.timeout_config(
TimeoutConfig::builder()
.operation_attempt_timeout(timeout)
.build(),
)
.build();
let output = litellm_s3::sdk::Client::from_conf(config)
.get_object()
.bucket(&object.bucket)
.key(&object.key)
.send()
.await?;
if output
.content_length()
.is_some_and(|length| length.unsigned_abs() > MAX_BODY_BYTES)
{
return Err(Error::body_too_large(MAX_BODY_BYTES));
}
let etag = output.e_tag().map(str::to_owned);
let mut stream = output.body;
let mut body = Vec::new();
while let Some(chunk) = tokio::time::timeout(timeout, stream.try_next()).await?? {
append(&mut body, &chunk)?;
}
Ok::<_, Error>(Fetched { body, etag })
};
fetched.await.map_err(Rejected::fatal)
}
fn append(body: &mut Vec<u8>, chunk: &Bytes) -> Result<(), Error> {
if (body.len() + chunk.len()) as u64 > MAX_BODY_BYTES {
return Err(Error::body_too_large(MAX_BODY_BYTES));
}
body.extend_from_slice(chunk);
Ok(())
}
fn etag(headers: &HeaderMap) -> Option<String> {
let values: Vec<&str> = headers
.get_all(ETAG)
.iter()
.filter_map(|value| value.to_str().ok())
.collect();
(!values.is_empty()).then(|| values.join(", "))
}

70
litellm-rust/crates/model-catalog-refresh/src/error.rs
View file

@ -1,6 +1,9 @@
use std::io;
use litellm_s3::sdk::config::http::HttpResponse;
use litellm_s3::sdk::error::SdkError;
use litellm_s3::sdk::operation::get_object::GetObjectError;
use litellm_s3::sdk::primitives::ByteStreamError;
use thiserror::Error;
use tokio::time::error::Elapsed;
#[derive(Debug, Error)]
#[error(transparent)]
@ -10,18 +13,48 @@ pub struct Error(Box<Repr>);
enum Repr {
#[error(transparent)]
Http(#[from] litellm_http::Error),
#[error("could not read the catalog")]
Read(#[source] io::Error),
#[error("invalid catalog URL")]
InvalidUrl(#[source] url::ParseError),
#[error("unsupported catalog URL scheme {scheme}")]
UnsupportedScheme { scheme: String },
#[error("invalid S3 URL")]
S3Uri(#[source] litellm_s3::Error),
#[error("S3 GetObject failed")]
S3GetObject(#[source] Box<SdkError<GetObjectError, HttpResponse>>),
#[error("S3 GetObject body read failed")]
S3Body(#[source] ByteStreamError),
#[error("S3 GetObject timed out")]
S3Timeout(#[source] Elapsed),
#[error("response body is over the {limit} byte limit")]
BodyTooLarge { limit: u64 },
}
impl Error {
pub fn unsupported_scheme(scheme: String) -> Self {
Self(Box::new(Repr::UnsupportedScheme { scheme }))
}
pub(crate) fn body_too_large(limit: u64) -> Self {
Self(Box::new(Repr::BodyTooLarge { limit }))
}
}
impl From<SdkError<GetObjectError, HttpResponse>> for Error {
fn from(cause: SdkError<GetObjectError, HttpResponse>) -> Self {
Self(Box::new(Repr::S3GetObject(Box::new(cause))))
}
}
impl From<ByteStreamError> for Error {
fn from(cause: ByteStreamError) -> Self {
Self(Box::new(Repr::S3Body(cause)))
}
}
impl From<Elapsed> for Error {
fn from(cause: Elapsed) -> Self {
Self(Box::new(Repr::S3Timeout(cause)))
}
}
impl From<litellm_http::Error> for Error {
@ -30,9 +63,9 @@ impl From<litellm_http::Error> for Error {
}
}
impl From<io::Error> for Error {
fn from(cause: io::Error) -> Self {
Self(Box::new(Repr::Read(cause)))
impl From<litellm_s3::Error> for Error {
fn from(cause: litellm_s3::Error) -> Self {
Self(Box::new(Repr::S3Uri(cause)))
}
}
@ -44,29 +77,18 @@ impl From<url::ParseError> for Error {
#[cfg(test)]
mod tests {
use std::io;
use litellm_error::{Report, find};
use rstest::rstest;
use super::*;
use crate::{DoneReason, Reply, Step};
#[rstest]
fn exhausted_transport_keeps_its_typed_cause() {
let reply = Reply::Transport {
error: io::Error::from(io::ErrorKind::ConnectionReset).into(),
};
let Reply::Transport { error } = reply else {
panic!("transport failure")
};
let step = Step::Done {
reason: DoneReason::Exhausted(error),
};
let Step::Done {
reason: DoneReason::Exhausted(error),
} = step
else {
panic!("exhausted attempts")
};
fn an_s3_failure_keeps_its_typed_cause_once() {
let error = Error::from(SdkError::<GetObjectError, HttpResponse>::timeout_error(
io::Error::from(io::ErrorKind::ConnectionReset),
));
let cause = find::<io::Error>(&error).expect("typed I/O cause");
assert_eq!(cause.kind(), io::ErrorKind::ConnectionReset);
assert_eq!(

122
litellm-rust/crates/model-catalog-refresh/src/fetch.rs Normal file
View file

@ -0,0 +1,122 @@
use std::borrow::Cow;
use std::sync::atomic::{AtomicU32, Ordering};
use std::time::Duration;
use backon::Retryable;
use litellm_error::Report;
use litellm_http::Client;
use url::Url;
use crate::attempt::{self, Cause, Fetched, Rejected};
use crate::{RetryPolicy, Source};
/// How an `s3://` fetch reaches S3
#[derive(Clone, Debug)]
pub enum S3Access {
/// What an AWS SDK finds in the environment and shared config files, loaded on each `s3://` attempt, with
/// `region` winning over the region it finds
Environment { region: Option<String> },
/// A config the caller built with its region, endpoint and credentials
Config(Box<litellm_s3::sdk::Config>),
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Failed {
pub reason: String,
/// Another attempt could succeed. Only the attempt limit ended the fetch
pub retryable: bool,
pub retry_after: Option<Duration>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Start {
Fresh,
/// Continues a fetch whose first attempt failed retryably elsewhere, waiting before its next attempt
After(Failed),
}
pub async fn fetch(
client: &Client,
s3: &S3Access,
url: &str,
policy: RetryPolicy,
start: Start,
) -> Result<Fetched, Failed> {
let shown = redacted(url);
let source =
Source::parse(url).map_err(|error| Failed::of(&shown, &Rejected::fatal(error), 1))?;
let mut backoff = policy.backoff();
let made = AtomicU32::new(0);
if let Start::After(failed) = start {
let Some(wait) = backoff.next() else {
return Err(failed);
};
tokio::time::sleep(policy.delay(wait, failed.retry_after)).await;
made.store(1, Ordering::Relaxed);
}
let attempt = || {
made.fetch_add(1, Ordering::Relaxed);
async {
match &source {
Source::Http(http) => attempt::http(client, http.clone(), policy.timeout).await,
Source::S3(object) => attempt::s3(client, s3, object, policy.timeout).await,
}
}
};
attempt
.retry(backoff)
.sleep(tokio::time::sleep)
.when(|rejected| rejected.retryable)
.adjust(|rejected, wait| wait.map(|wait| policy.delay(wait, rejected.retry_after)))
.notify(|rejected, wait| {
litellm_tracing::warn!(
"LiteLLM: model cost map fetch attempt {}/{} failed ({}); retrying in {:.1}s",
made.load(Ordering::Relaxed),
policy.max_attempts,
reason(&shown, &rejected.cause),
wait.as_secs_f64(),
);
})
.await
.map_err(|rejected| Failed::of(&shown, &rejected, made.load(Ordering::Relaxed)))
}
impl Failed {
fn of(url: &str, rejected: &Rejected, attempts: u32) -> Self {
let reason = reason(url, &rejected.cause);
Self {
reason: match rejected.retryable && attempts > 1 {
true => format!("{reason} (after {attempts} attempts)"),
false => reason,
},
retryable: rejected.retryable,
retry_after: rejected.retry_after,
}
}
}
fn reason(url: &str, cause: &Cause) -> String {
match cause {
Cause::Status(status) => format!("HTTP {status} from {url}"),
Cause::Transport(error) => format!("error fetching {url}: {}", Report::new(error)),
}
}
/// Drops userinfo, query and fragment, where presigned URLs and tokens carry credentials
fn redacted(url: &str) -> Cow<'_, str> {
let Ok(mut parsed) = Url::parse(url) else {
return Cow::Borrowed(url);
};
if parsed.username().is_empty()
&& parsed.password().is_none()
&& parsed.query().is_none()
&& parsed.fragment().is_none()
{
return Cow::Borrowed(url);
}
let _ = parsed.set_username("");
let _ = parsed.set_password(None);
parsed.set_query(None);
parsed.set_fragment(None);
Cow::Owned(parsed.into())
}

6
litellm-rust/crates/model-catalog-refresh/src/lib.rs
View file

@ -1,9 +1,11 @@
mod attempt;
mod error;
mod fetch;
mod policy;
mod source;
mod step;
pub use attempt::Fetched;
pub use error::Error;
pub use fetch::*;
pub use policy::*;
pub use source::*;
pub use step::*;

74
litellm-rust/crates/model-catalog-refresh/src/policy.rs
View file

@ -1,20 +1,76 @@
use std::time::Duration;
use litellm_model_catalog::IntegrityLimits;
use backon::{BackoffBuilder, ExponentialBackoff, ExponentialBuilder};
/// Main's defaults, 3 attempts, 30 s max wait, 5 s timeout, backoff base 2.0, jitter 0.0 to 1.0 s
#[derive(Clone, Copy, Debug, PartialEq)]
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct RetryPolicy {
pub max_attempts: u32,
pub max_wait: Duration,
pub timeout: Duration,
pub backoff_base_seconds: f64,
pub jitter_min_seconds: f64,
pub jitter_max_seconds: f64,
/// Main's 50 model floor and 0.5 shrink ratio
pub limits: IntegrityLimits,
pub min_delay: Duration,
pub max_delay: Duration,
}
pub const RETRYABLE_STATUSES: [u16; 5] = [429, 500, 502, 503, 504];
/// Main's default `LITELLM_MODEL_COST_MAP_URL`
pub const DEFAULT_REGISTRY_URL: &str =
"https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json";
impl RetryPolicy {
pub(crate) fn backoff(&self) -> ExponentialBackoff {
ExponentialBuilder::new()
.with_min_delay(self.min_delay)
.with_max_delay(self.max_delay)
.with_max_times(self.max_attempts.saturating_sub(1) as usize)
.with_jitter()
.build()
}
pub(crate) fn delay(&self, backoff: Duration, retry_after: Option<Duration>) -> Duration {
retry_after.map_or(backoff, |server| server.min(self.max_delay))
}
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use super::*;
const POLICY: RetryPolicy = RetryPolicy {
max_attempts: 3,
timeout: Duration::from_secs(5),
min_delay: Duration::from_secs(1),
max_delay: Duration::from_secs(30),
};
#[rstest]
#[case::backoff_without_retry_after(None, 4)]
#[case::retry_after_wins(Some(7), 7)]
#[case::retry_after_zero(Some(0), 0)]
#[case::retry_after_capped(Some(120), 30)]
fn the_server_retry_after_overrides_the_backoff_up_to_max_delay(
#[case] retry_after: Option<u64>,
#[case] seconds: u64,
) {
assert_eq!(
POLICY.delay(Duration::from_secs(4), retry_after.map(Duration::from_secs)),
Duration::from_secs(seconds)
);
}
#[rstest]
#[case::one_attempt(1, 0)]
#[case::three_attempts(3, 2)]
#[case::zero_attempts(0, 0)]
fn the_backoff_allows_one_retry_fewer_than_the_attempts(
#[case] max_attempts: u32,
#[case] retries: usize,
) {
let policy = RetryPolicy {
max_attempts,
..POLICY
};
assert_eq!(policy.backoff().count(), retries);
}
}

58
litellm-rust/crates/model-catalog-refresh/src/source.rs
View file

@ -1,17 +1,59 @@
use std::path::PathBuf;
use litellm_s3::S3Uri;
use url::Url;
use crate::Error;
/// Where a fetched catalog comes from
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Source {
Http(Url),
File(PathBuf),
S3(S3Object),
S3(S3Uri),
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct S3Object {
pub bucket: String,
pub key: String,
impl Source {
pub fn parse(url: &str) -> Result<Self, Error> {
let parsed = Url::parse(url)?;
match parsed.scheme() {
"http" | "https" => Ok(Self::Http(parsed)),
"s3" => Ok(Self::S3(S3Uri::try_from(url)?)),
scheme => Err(Error::unsupported_scheme(scheme.to_owned())),
}
}
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use super::*;
#[rstest]
#[case::https("https://example.com/cost.json")]
#[case::http("http://127.0.0.1:8000/cost.json")]
fn http_urls_are_fetched_over_http(#[case] url: &str) {
assert_eq!(
Source::parse(url).unwrap(),
Source::Http(Url::parse(url).unwrap())
);
}
#[rstest]
fn s3_urls_name_a_bucket_and_key() {
assert_eq!(
Source::parse("s3://bucket/prices/cost.json").unwrap(),
Source::S3(S3Uri {
bucket: "bucket".into(),
key: "prices/cost.json".into(),
})
);
}
#[rstest]
#[case::s3_without_a_key("s3://bucket/")]
#[case::unsupported("ftp://example.com/cost.json")]
#[case::file("file:///tmp/cost.json")]
#[case::not_a_url("not a url")]
fn other_urls_are_refused(#[case] url: &str) {
assert!(Source::parse(url).is_err());
}
}

66
litellm-rust/crates/model-catalog-refresh/src/step.rs
View file

@ -1,66 +0,0 @@
use std::time::Duration;
use litellm_model_catalog::Catalog;
use crate::Error;
use crate::source::Source;
/// One attempt. Where to fetch from is the `Source`, so the request carries no URL
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct FetchRequest {
pub if_none_match: Option<String>,
pub timeout: Duration,
}
#[derive(Debug)]
pub enum Reply {
/// An HTTP, S3 or file answer. The body is read only on 200
Response {
status: u16,
etag: Option<String>,
retry_after: Option<Duration>,
body: Option<Vec<u8>>,
},
/// No answer reached us
Transport { error: Error },
}
/// What the fetch state machine asks its driver to do next
#[derive(Debug)]
pub enum Step {
Fetch(FetchRequest),
Sleep(Duration),
/// A parsed document to install, with its rejected row errors
Adopt {
catalog: Box<Catalog>,
installed: Installed,
errors: Vec<litellm_model_catalog::Error>,
},
/// 304, a final status, or retries exhausted
Done {
reason: DoneReason,
},
}
#[derive(Debug)]
pub enum DoneReason {
NotModified,
Status(u16),
Exhausted(Error),
}
/// What is installed now, the input to the next revalidation
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Installed {
pub origin: Origin,
pub etag: Option<String>,
/// `git_blob_id` of the installed bytes
pub revision: String,
pub model_count: usize,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Origin {
Bundled,
Remote { source: Source },
}

541
litellm-rust/crates/model-catalog-refresh/tests/fetch.rs Normal file
View file

@ -0,0 +1,541 @@
use std::io::{BufRead, BufReader, Write};
use std::net::TcpListener;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use aws_credential_types::Credentials;
use aws_credential_types::provider::error::CredentialsError;
use aws_credential_types::provider::{ProvideCredentials, future};
use litellm_http::Client;
use litellm_model_catalog_refresh::{Failed, Fetched, RetryPolicy, S3Access, Start, fetch};
use rstest::rstest;
const POLICY: RetryPolicy = RetryPolicy {
max_attempts: 3,
timeout: Duration::from_secs(5),
min_delay: Duration::from_millis(10),
max_delay: Duration::from_millis(50),
};
const OK: &str =
"HTTP/1.1 200 OK\r\nETag: \"v1\"\r\nContent-Length: 2\r\nConnection: close\r\n\r\n{}";
const UNAVAILABLE: &str = "HTTP/1.1 503 Service Unavailable\r\nRetry-After: 0\r\nContent-Length: 0\r\nConnection: close\r\n\r\n";
const NOT_FOUND: &str = "HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\nConnection: close\r\n\r\n";
const SLOW_BODY: &[u8] = b"{ }";
#[derive(Clone, Copy)]
enum Answer {
Raw(&'static str),
Stall,
StallBody,
SlowBody,
Oversized,
}
/// A scripted upstream: the n-th request gets the n-th answer, repeating the last, and every request head is kept
struct Upstream {
base: String,
heads: Arc<Mutex<Vec<String>>>,
}
impl Upstream {
fn serve(answers: Vec<Answer>) -> Self {
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
let base = format!("http://{}", listener.local_addr().unwrap());
let heads = Arc::new(Mutex::new(Vec::new()));
let recorded = Arc::clone(&heads);
std::thread::spawn(move || {
for (index, stream) in listener.incoming().enumerate() {
let answer = answers[index.min(answers.len() - 1)];
let recorded = Arc::clone(&recorded);
std::thread::spawn(move || {
let mut stream = stream.unwrap();
let mut reader = BufReader::new(stream.try_clone().unwrap());
let mut head = String::new();
loop {
let mut line = String::new();
if reader.read_line(&mut line).unwrap() <= 2 {
break;
}
head.push_str(&line);
}
recorded.lock().unwrap().push(head);
match answer {
Answer::Raw(raw) => stream.write_all(raw.as_bytes()).unwrap(),
Answer::Stall => std::thread::sleep(Duration::from_secs(5)),
Answer::Oversized => {
let chunk = vec![b' '; 1024 * 1024];
let written = stream
.write_all(b"HTTP/1.1 200 OK\r\nConnection: close\r\n\r\n")
.and_then(|()| {
(0..=100).try_for_each(|_| stream.write_all(&chunk))
});
drop(written);
}
Answer::StallBody => {
stream
.write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\n{")
.unwrap();
std::thread::sleep(Duration::from_secs(5));
}
Answer::SlowBody => {
stream.set_nodelay(true).unwrap();
stream
.write_all(
b"HTTP/1.1 200 OK\r\nETag: \"v1\"\r\nContent-Length: 8\r\nConnection: close\r\n\r\n",
)
.unwrap();
for byte in SLOW_BODY {
std::thread::sleep(Duration::from_millis(60));
stream.write_all(&[*byte]).unwrap();
}
}
}
});
}
});
Self { base, heads }
}
fn url(&self) -> String {
format!("{}/cost.json", self.base)
}
fn heads(&self) -> Vec<String> {
self.heads.lock().unwrap().clone()
}
}
fn header<'a>(head: &'a str, name: &str) -> Option<&'a str> {
head.lines().find_map(|line| {
let (key, value) = line.split_once(':')?;
key.eq_ignore_ascii_case(name).then(|| value.trim())
})
}
#[derive(Clone, Debug)]
struct NoCredentials;
impl ProvideCredentials for NoCredentials {
fn provide_credentials<'a>(&'a self) -> future::ProvideCredentials<'a>
where
Self: 'a,
{
future::ProvideCredentials::ready(Err(CredentialsError::not_loaded("no credentials")))
}
}
fn keys() -> Credentials {
Credentials::new(
"AKIDMODELCOST",
"model-cost-secret",
None,
None,
"model-cost-test",
)
}
async fn fetched(url: &str, policy: RetryPolicy, start: Start) -> Result<Fetched, Failed> {
fetch(
&Client::for_test(),
&S3Access::Environment { region: None },
url,
policy,
start,
)
.await
}
#[rstest]
#[case::every_etag_joined(
"HTTP/1.1 200 OK\r\nETag: \"a\"\r\nETag: \"b\"\r\nContent-Length: 2\r\nConnection: close\r\n\r\n{}",
Some("\"a\", \"b\"")
)]
#[case::no_etag(
"HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\n{}",
None
)]
#[tokio::test]
async fn a_200_returns_the_body_and_its_etags(
#[case] response: &'static str,
#[case] etag: Option<&str>,
) {
let upstream = Upstream::serve(vec![Answer::Raw(response)]);
assert_eq!(
fetched(&upstream.url(), POLICY, Start::Fresh).await,
Ok(Fetched {
body: b"{}".to_vec(),
etag: etag.map(str::to_owned),
})
);
assert_eq!(upstream.heads().len(), 1);
}
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn a_fetch_runs_as_a_task_on_a_multi_thread_runtime() {
let upstream = Upstream::serve(vec![Answer::Raw(OK)]);
let url = upstream.url();
let fetched = tokio::spawn(async move {
fetch(
&Client::for_test(),
&S3Access::Environment { region: None },
&url,
POLICY,
Start::Fresh,
)
.await
})
.await
.unwrap();
assert!(fetched.is_ok());
}
#[rstest]
#[tokio::test]
async fn a_redirect_is_not_followed_and_its_body_is_the_document() {
let upstream = Upstream::serve(vec![
Answer::Raw(
"HTTP/1.1 301 Moved Permanently\r\nLocation: /elsewhere.json\r\nContent-Length: 7\r\nConnection: close\r\n\r\n{\"a\":1}",
),
Answer::Raw(OK),
]);
assert_eq!(
fetched(&upstream.url(), POLICY, Start::Fresh).await,
Ok(Fetched {
body: b"{\"a\":1}".to_vec(),
etag: None,
})
);
assert_eq!(upstream.heads().len(), 1);
}
#[rstest]
#[tokio::test]
async fn a_retry_after_shorter_than_the_backoff_is_honored() {
let upstream = Upstream::serve(vec![Answer::Raw(UNAVAILABLE), Answer::Raw(OK)]);
let policy = RetryPolicy {
min_delay: Duration::from_secs(10),
max_delay: Duration::from_secs(10),
..POLICY
};
let fetched = tokio::time::timeout(
Duration::from_secs(2),
fetched(&upstream.url(), policy, Start::Fresh),
)
.await
.expect("Retry-After: 0 replaces the 10s backoff");
assert!(fetched.is_ok(), "{fetched:?}");
assert_eq!(upstream.heads().len(), 2);
}
#[rstest]
#[tokio::test]
async fn retryable_failures_stop_at_the_attempt_limit() {
let upstream = Upstream::serve(vec![Answer::Raw(UNAVAILABLE)]);
let url = upstream.url();
assert_eq!(
fetched(&url, POLICY, Start::Fresh).await,
Err(Failed {
reason: format!("HTTP 503 from {url} (after 3 attempts)"),
retryable: true,
retry_after: Some(Duration::ZERO),
})
);
assert_eq!(upstream.heads().len(), 3);
}
#[rstest]
#[tokio::test]
async fn a_404_fails_on_the_first_attempt() {
let upstream = Upstream::serve(vec![Answer::Raw(NOT_FOUND)]);
let url = upstream.url();
assert_eq!(
fetched(&url, POLICY, Start::Fresh).await,
Err(Failed {
reason: format!("HTTP 404 from {url}"),
retryable: false,
retry_after: None,
})
);
assert_eq!(upstream.heads().len(), 1);
}
#[rstest]
#[tokio::test]
async fn a_stalled_response_times_out_and_is_retried() {
let upstream = Upstream::serve(vec![Answer::Stall]);
let url = upstream.url();
let policy = RetryPolicy {
max_attempts: 2,
timeout: Duration::from_millis(100),
..POLICY
};
let failed = fetched(&url, policy, Start::Fresh).await.unwrap_err();
assert!(
failed
.reason
.starts_with(&format!("error fetching {url}: "))
&& failed.reason.ends_with(" (after 2 attempts)"),
"{}",
failed.reason
);
assert!(failed.retryable);
assert_eq!(upstream.heads().len(), 2);
}
#[rstest]
#[tokio::test]
async fn a_resumed_fetch_makes_only_the_attempts_left() {
let upstream = Upstream::serve(vec![Answer::Raw(UNAVAILABLE)]);
let url = upstream.url();
let first = Failed {
reason: "first".into(),
retryable: true,
retry_after: Some(Duration::ZERO),
};
let failed = fetched(&url, POLICY, Start::After(first))
.await
.unwrap_err();
assert_eq!(
failed.reason,
format!("HTTP 503 from {url} (after 3 attempts)")
);
assert_eq!(upstream.heads().len(), 2);
}
#[rstest]
#[tokio::test]
async fn a_resumed_fetch_with_no_attempts_left_returns_the_failure_it_resumed() {
let upstream = Upstream::serve(vec![Answer::Raw(OK)]);
let first = Failed {
reason: "first".into(),
retryable: true,
retry_after: None,
};
let policy = RetryPolicy {
max_attempts: 1,
..POLICY
};
assert_eq!(
fetched(&upstream.url(), policy, Start::After(first.clone())).await,
Err(first)
);
assert!(upstream.heads().is_empty());
}
#[rstest]
#[tokio::test]
async fn a_failure_reason_drops_url_credentials() {
let upstream = Upstream::serve(vec![Answer::Raw(NOT_FOUND)]);
let address = upstream.base.trim_start_matches("http://");
let url = format!("http://user:hunter2@{address}/cost.json?X-Amz-Signature=sig-secret#frag");
let failed = fetched(&url, POLICY, Start::Fresh).await.unwrap_err();
assert_eq!(failed.reason, format!("HTTP 404 from {}", upstream.url()));
}
#[rstest]
#[case::unsupported_scheme("ftp://example.com/cost.json")]
#[case::not_a_url("not a url")]
#[case::s3_without_a_key("s3://prices")]
#[tokio::test]
async fn an_unfetchable_url_fails_without_a_request(#[case] url: &str) {
let failed = fetched(url, POLICY, Start::Fresh).await.unwrap_err();
assert!(
failed
.reason
.starts_with(&format!("error fetching {url}: ")),
"{}",
failed.reason
);
assert!(!failed.retryable);
}
fn s3(upstream: &Upstream, credentials: impl ProvideCredentials + 'static) -> S3Access {
S3Access::Config(Box::new(
litellm_s3::config(
"us-west-2".into(),
Some(&upstream.base),
Client::for_test(),
credentials,
)
.build(),
))
}
async fn fetched_from_s3(
url: &str,
upstream: &Upstream,
credentials: impl ProvideCredentials + 'static,
) -> Result<Fetched, Failed> {
fetch(
&Client::for_test(),
&s3(upstream, credentials),
url,
POLICY,
Start::Fresh,
)
.await
}
#[rstest]
#[tokio::test]
async fn an_s3_fetch_sends_a_signed_get_object_to_the_endpoint() {
let upstream = Upstream::serve(vec![Answer::Raw(OK)]);
assert_eq!(
fetched_from_s3("s3://prices/maps/cost.json", &upstream, keys()).await,
Ok(Fetched {
body: b"{}".to_vec(),
etag: Some("\"v1\"".into()),
})
);
let heads = upstream.heads();
let [head] = heads.as_slice() else {
panic!("expected one request, got {heads:?}");
};
assert!(head.starts_with("GET /prices/maps/cost.json"), "{head}");
let authorization = header(head, "authorization").unwrap_or_default();
assert!(
authorization.starts_with("AWS4-HMAC-SHA256 Credential=AKIDMODELCOST/")
&& authorization.contains("/us-west-2/s3/aws4_request"),
"{authorization}"
);
}
#[rstest]
#[tokio::test]
async fn an_s3_body_that_stalls_after_the_headers_times_out() {
let upstream = Upstream::serve(vec![Answer::StallBody]);
let url = "s3://prices/cost.json";
let policy = RetryPolicy {
timeout: Duration::from_millis(200),
..POLICY
};
let failed = tokio::time::timeout(
Duration::from_secs(2),
fetch(
&Client::for_test(),
&s3(&upstream, keys()),
url,
policy,
Start::Fresh,
),
)
.await
.expect("the attempt timeout bounds the body read")
.unwrap_err();
assert!(
failed
.reason
.starts_with(&format!("error fetching {url}: ")),
"{}",
failed.reason
);
assert!(!failed.retryable);
assert_eq!(upstream.heads().len(), 1);
}
#[rstest]
#[case::http(|upstream: &Upstream| (S3Access::Environment { region: None }, upstream.url()))]
#[case::s3(|upstream: &Upstream| (s3(upstream, keys()), "s3://prices/cost.json".to_owned()))]
#[tokio::test]
async fn a_body_that_keeps_arriving_finishes_after_the_timeout(
#[case] source: fn(&Upstream) -> (S3Access, String),
) {
let upstream = Upstream::serve(vec![Answer::SlowBody]);
let (access, url) = source(&upstream);
let policy = RetryPolicy {
timeout: Duration::from_millis(300),
..POLICY
};
assert_eq!(
fetch(&Client::for_test(), &access, &url, policy, Start::Fresh).await,
Ok(Fetched {
body: SLOW_BODY.to_vec(),
etag: Some("\"v1\"".into()),
})
);
assert_eq!(upstream.heads().len(), 1);
}
#[rstest]
#[case::missing_key(NOT_FOUND)]
#[case::denied("HTTP/1.1 403 Forbidden\r\nContent-Length: 0\r\nConnection: close\r\n\r\n")]
#[tokio::test]
async fn an_s3_error_fails_without_a_retry(#[case] response: &'static str) {
let upstream = Upstream::serve(vec![Answer::Raw(response), Answer::Raw(OK)]);
let url = "s3://prices/cost.json";
let failed = fetched_from_s3(url, &upstream, keys()).await.unwrap_err();
assert!(
failed
.reason
.starts_with(&format!("error fetching {url}: ")),
"{}",
failed.reason
);
assert!(!failed.retryable);
assert_eq!(upstream.heads().len(), 1);
}
#[rstest]
#[tokio::test]
async fn an_s3_fetch_without_credentials_fails_without_a_request() {
let upstream = Upstream::serve(vec![Answer::Raw(OK)]);
let url = "s3://prices/cost.json";
let failed = fetched_from_s3(url, &upstream, NoCredentials)
.await
.unwrap_err();
assert!(
failed
.reason
.starts_with(&format!("error fetching {url}: ")),
"{}",
failed.reason
);
assert!(!failed.retryable);
assert!(upstream.heads().is_empty());
}
const DECLARED_OVERSIZED: &str =
"HTTP/1.1 200 OK\r\nContent-Length: 104857601\r\nConnection: close\r\n\r\n";
#[rstest]
#[case::http_declared(Answer::Raw(DECLARED_OVERSIZED), false)]
#[case::http_streamed(Answer::Oversized, false)]
#[case::s3_declared(Answer::Raw(DECLARED_OVERSIZED), true)]
#[case::s3_streamed(Answer::Oversized, true)]
#[tokio::test]
async fn a_body_over_100_mib_fails_without_a_retry(#[case] answer: Answer, #[case] s3: bool) {
let upstream = Upstream::serve(vec![answer, Answer::Raw(OK)]);
let failed = match s3 {
true => fetched_from_s3("s3://prices/cost.json", &upstream, keys()).await,
false => fetched(&upstream.url(), POLICY, Start::Fresh).await,
}
.unwrap_err();
assert!(
failed
.reason
.contains("response body is over the 104857600 byte limit"),
"{}",
failed.reason
);
assert!(!failed.retryable);
assert_eq!(upstream.heads().len(), 1);
}
#[rstest]
#[tokio::test]
async fn a_body_of_exactly_100_mib_is_fetched() {
let upstream = Upstream::serve(vec![Answer::Raw(Box::leak(
format!(
"HTTP/1.1 200 OK\r\nContent-Length: 104857600\r\nConnection: close\r\n\r\n{}",
" ".repeat(100 * 1024 * 1024)
)
.into_boxed_str(),
))]);
let fetched = fetched(&upstream.url(), POLICY, Start::Fresh)
.await
.unwrap();
assert_eq!(fetched.body.len(), 100 * 1024 * 1024);
}

9
litellm-rust/crates/model-catalog/AGENTS.md
View file

@ -46,7 +46,14 @@ Every price in a row lowers to one `Price { meter, channel, when }` and its `Uni
`store/` is the catalog's in-memory store and the only place that indexes model names. Its rules are in [src/store/AGENTS.md](src/store/AGENTS.md)
`layers.rs`, `resolve/`, `model_info.rs` and `register.rs` hold types only. The functions over them land with their first caller
`cost_map.rs` holds a host's inputs, the installed document and the `register_model` overlay in `register.rs`. It serves the document in the key order and alias sharing Python's `litellm.model_cost` has, without the overlay, since where a registration lands is a lookup in the published table. It never holds the live table, which belongs to the host
- `CostMap` is immutable behind an `Arc`. A new document is a new `CostMap`
- `Registrations<V>` is generic over the host's value type, because Python registrations can hold `datetime` and other values JSON cannot. The host keeps its own objects (`Py<PyDict>` in python-bridge) and does the merging
- Entry order is latest-registration order. `landed_on` records which registered keys landed on each row key
- `write_pass` stops at the first `place` that returns `None` and reports `Pass::Swapped`. Retrying is the host's job
`resolve/` and `model_info.rs` hold types only. The functions over them land with their first caller
Catalog errors use one opaque `Error` with typed causes and JSON Pointer accessors (`path`, `key`, `field`). `Error::rejects` alone decides whether an error fails the document, rejects its row or neither, and `ErrorKind` is the stable name hosts report it by

2
litellm-rust/crates/model-catalog/Cargo.toml
View file

@ -11,6 +11,7 @@ schema = ["dep:schemars", "litellm-cost/schema"]
[dependencies]
litellm-cost = { workspace = true, features = ["serde"] }
litellm-error.workspace = true
litellm-tracing.workspace = true
indexmap = { version = "2.14.0", features = ["serde"] }
jiff = { version = "0.2.37", default-features = false, features = ["std", "tzdb-bundle-always"] }
@ -19,6 +20,7 @@ schemars = { workspace = true, optional = true }
serde.workspace = true
serde_json = { workspace = true, features = ["preserve_order"] }
serde_with.workspace = true
sha1.workspace = true
strum.workspace = true
thiserror.workspace = true
time.workspace = true

109
litellm-rust/crates/model-catalog/src/catalog.rs
View file

@ -12,7 +12,8 @@ use crate::store::{Rows, Store};
use crate::wire::{EntryWire, Fields};
const SAMPLE_SPEC: &str = "sample_spec";
const FALLBACK_GENERALIZATIONS: &str = "fallback_generalizations";
/// The top-level key of the rules block, which is never a model row
pub const FALLBACK_GENERALIZATIONS: &str = "fallback_generalizations";
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct IntegrityLimits {
@ -111,6 +112,11 @@ impl Rows for CatalogRows {
}
impl IntegrityLimits {
pub const UNCHECKED: Self = Self {
min_model_count: 0,
min_reference_ratio: 0.0,
};
pub fn new(min_model_count: usize, min_reference_ratio: f64) -> Option<Self> {
(0.0..=1.0).contains(&min_reference_ratio).then_some(Self {
min_model_count,
@ -142,6 +148,7 @@ impl Catalog {
let mut sections = Vec::new();
let mut entries = IndexMap::with_capacity(document.len());
let mut alias_lists = Vec::new();
let mut pointer_rows = Vec::new();
for (index, (name, value)) in document.into_iter().enumerate() {
let fields = match value {
Value::Object(fields) if !is_section(&name) => fields,
@ -157,16 +164,34 @@ impl Catalog {
continue;
}
};
let (model, aliases, row_problems) = read_row(&fields);
let row = read_row(&fields);
let readable = row.problems.is_empty();
problems.extend(
row_problems
row.problems
.into_iter()
.map(|problem| problem.prepend(&name)),
);
if let Some(aliases) = aliases {
if let Some(target) = row.alias_of {
if readable {
pointer_rows.push((name.clone(), target));
}
sections.push(Section {
index,
key: name,
value: Value::Object(fields),
});
continue;
}
if let Some(aliases) = row.aliases {
alias_lists.push((name.clone(), aliases));
}
entries.insert(name, ModelEntry { fields, model });
entries.insert(
name,
ModelEntry {
fields,
model: row.model,
},
);
}
if entries.is_empty() {
return Err(Error::new("", Repr::Empty));
@ -190,6 +215,21 @@ impl Catalog {
}
}
}
for (alias, target) in pointer_rows {
if !entries.contains_key(&target) {
problems.push(Error::from_key(&alias, Repr::AliasTargetMissing { target }));
} else if let Some(taken_by) = aliases.get(&alias) {
problems.push(Error::from_key(
&alias,
Repr::AliasTaken {
alias: alias.clone(),
taken_by: taken_by.clone(),
},
));
} else {
aliases.insert(alias, target);
}
}
let fallback_generalizations = sections
.iter()
@ -242,8 +282,8 @@ impl Catalog {
}
pub fn row_errors(model: &str, fields: &Fields) -> Vec<Error> {
let (_, _, errors) = read_row(fields);
errors
read_row(fields)
.problems
.into_iter()
.map(|error| error.prepend(model))
.collect()
@ -255,7 +295,13 @@ impl Catalog {
(0..self.rows().entries.len() + self.sections.len()).map_while(move |index| match sections
.next_if(|section| section.index == index)
{
Some(section) => Some((section.key.as_str(), Written::Section(&section.value))),
Some(section) => Some((
section.key.as_str(),
match &section.value {
Value::Object(fields) if !is_section(&section.key) => Written::Entry(fields),
value => Written::Section(value),
},
)),
None => entries
.next()
.map(|(name, entry)| (name.as_str(), Written::Entry(&entry.fields))),
@ -378,6 +424,20 @@ impl Catalog {
})
}
/// The row under an exact model or alias key, an alias reading its owner's row
pub fn get(&self, key: &str) -> Option<&Fields> {
self.rows().fields(key)
}
/// A top-level value that is not a model row, such as `sample_spec`
pub fn section(&self, key: &str) -> Option<&Value> {
self.sections
.iter()
.filter(|section| is_section(&section.key))
.find(|section| section.key == key)
.map(|section| &section.value)
}
pub fn model_count(&self) -> usize {
self.rows().entries.len()
}
@ -406,16 +466,41 @@ impl Catalog {
}
}
fn read_row(fields: &Fields) -> (Option<Model>, Option<Vec<String>>, Vec<Error>) {
struct RowRead {
model: Option<Model>,
aliases: Option<Vec<String>>,
alias_of: Option<String>,
problems: Vec<Error>,
}
fn read_row(fields: &Fields) -> RowRead {
let wire = EntryWire::read(fields);
if let Some(target) = wire.attributes.alias_of.clone() {
let other_fields = fields
.keys()
.filter(|key| *key != "alias_of")
.map(|key| Error::from_key(key, Repr::AliasRowField));
return RowRead {
model: None,
aliases: None,
alias_of: Some(target),
problems: wire.problems.into_iter().chain(other_fields).collect(),
};
}
let aliases = wire.attributes.aliases.clone();
let timezone_error = unknown_timezone(&wire);
match Model::try_from(wire) {
Ok(model) => (Some(model), aliases, timezone_error.into_iter().collect()),
let (model, problems) = match Model::try_from(wire) {
Ok(model) => (Some(model), timezone_error.into_iter().collect()),
Err(mut errors) => {
errors.extend(timezone_error);
(None, aliases, errors)
(None, errors)
}
};
RowRead {
model,
aliases,
alias_of: None,
problems,
}
}

130
litellm-rust/crates/model-catalog/src/cost_map.rs Normal file
View file

@ -0,0 +1,130 @@
use std::sync::{Arc, LazyLock};
use indexmap::IndexMap;
use serde_json::Value;
use sha1::{Digest, Sha1};
use crate::bundled::BUNDLED_REGISTRY;
use crate::catalog::{Catalog, FALLBACK_GENERALIZATIONS, IntegrityLimits};
use crate::error::Error;
/// Which document is installed: the `git_blob_id` of its bytes and the ETag it was served with
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Provenance {
pub revision: String,
pub etag: Option<String>,
}
/// An installed cost map document. A host builds its live table from it
#[derive(Clone, Debug)]
pub struct CostMap {
base: Arc<Base>,
provenance: Provenance,
}
#[derive(Debug)]
struct Base {
catalog: Catalog,
/// Each owner's names, the owner first, for owners with at least one alias
groups: IndexMap<String, Vec<String>>,
}
impl Base {
fn new(catalog: Catalog) -> Self {
let groups = catalog.aliases().iter().fold(
IndexMap::<String, Vec<String>>::new(),
|mut groups, (alias, owner)| {
groups
.entry(owner.clone())
.or_insert_with(|| vec![owner.clone()])
.push(alias.clone());
groups
},
);
Self { catalog, groups }
}
}
static BUNDLED: LazyLock<CostMap> = LazyLock::new(|| CostMap {
base: Arc::new(Base::new(
Catalog::parse(BUNDLED_REGISTRY)
.expect("the bundled registry parses")
.value,
)),
provenance: Provenance {
revision: git_blob_id(BUNDLED_REGISTRY),
etag: None,
},
});
impl CostMap {
pub fn bundled() -> Self {
BUNDLED.clone()
}
/// Reads a fetched document. One that does not parse, or that has fewer models than `limits` allow against the
/// bundled one, is refused. Every alias problem it has is logged
pub fn read(body: &[u8], etag: Option<String>, limits: IntegrityLimits) -> Result<Self, Error> {
let loaded = Catalog::parse(body)?;
loaded
.value
.check_size(BUNDLED.base.catalog.model_count(), limits)?;
for problem in loaded.errors.iter().filter(|error| error.is_alias()) {
litellm_tracing::warn!(
"LiteLLM: model cost map alias problem: {}",
litellm_error::Report::new(problem)
);
}
Ok(Self {
base: Arc::new(Base::new(loaded.value)),
provenance: Provenance {
revision: git_blob_id(body),
etag,
},
})
}
pub fn catalog(&self) -> &Catalog {
&self.base.catalog
}
pub fn provenance(&self) -> &Provenance {
&self.provenance
}
/// Every name the row under `key` is served under, the owner first. Empty for a row without aliases
pub fn names(&self, key: &str) -> &[String] {
let owner = self
.catalog()
.aliases()
.get(key)
.map_or(key, String::as_str);
self.base
.groups
.get(owner)
.map(Vec::as_slice)
.unwrap_or_default()
}
/// Each owner with its names, the owner first, in the order the document declared the aliases
pub fn alias_groups(&self) -> impl Iterator<Item = (&str, &[String])> {
self.base
.groups
.iter()
.map(|(owner, names)| (owner.as_str(), names.as_slice()))
}
/// The `fallback_generalizations` block as the document wrote it
pub fn fallback_generalizations(&self) -> Option<&Value> {
self.catalog().section(FALLBACK_GENERALIZATIONS)
}
}
/// `git hash-object` of the bytes, which main reports as `source_revision`
pub fn git_blob_id(body: &[u8]) -> String {
let digest = Sha1::new()
.chain_update(format!("blob {}\0", body.len()))
.chain_update(body)
.finalize();
format!("{digest:x}")
}

78
litellm-rust/crates/model-catalog/src/error.rs
View file

@ -67,6 +67,12 @@ pub(crate) enum Repr {
#[error("alias {alias:?} already belongs to {taken_by}")]
#[strum_discriminants(strum(serialize = "alias_taken"))]
AliasTaken { alias: String, taken_by: String },
#[error("alias_of {target:?} is not a model name")]
#[strum_discriminants(strum(serialize = "alias_target_missing"))]
AliasTargetMissing { target: String },
#[error("an alias_of row carries no other field")]
#[strum_discriminants(strum(serialize = "alias_row_field"))]
AliasRowField,
#[error("{first} and {second} price the same meter under the same conditions")]
#[strum_discriminants(strum(serialize = "duplicate_price"))]
DuplicatePrice { first: String, second: String },
@ -141,6 +147,8 @@ impl Error {
| Repr::UnknownTimezone { .. }
| Repr::AliasShadowsModel { .. }
| Repr::AliasTaken { .. }
| Repr::AliasTargetMissing { .. }
| Repr::AliasRowField
| Repr::DuplicatePrice { .. }
| Repr::SuffixInsideBlock { .. }
| Repr::TierRange
@ -159,6 +167,74 @@ impl Error {
}
}
impl Error {
/// The document parsed but holds fewer models than the integrity limits trust
pub fn is_integrity(&self) -> bool {
match self.repr() {
Repr::Empty | Repr::BelowMinimum { .. } | Repr::Shrunk { .. } => true,
Repr::Json(_)
| Repr::NotAnObject
| Repr::Invalid(_)
| Repr::UnknownBlockKey
| Repr::UnknownTimezone { .. }
| Repr::AliasShadowsModel { .. }
| Repr::AliasTaken { .. }
| Repr::AliasTargetMissing { .. }
| Repr::AliasRowField
| Repr::DuplicatePrice { .. }
| Repr::SuffixInsideBlock { .. }
| Repr::TierRange
| Repr::TierReversed
| Repr::TierGap { .. }
| Repr::MixedTierBasis
| Repr::NotPricedByResults
| Repr::BandsFromTwoSources
| Repr::NotAnOffPeakRate
| Repr::NoOffPeakPrices
| Repr::NoOffPeakHours
| Repr::EmptyList
| Repr::InvalidDate { .. }
| Repr::MissingProvider
| Repr::UnreadField
| Repr::NoFallbackRules => false,
}
}
/// A problem with a row's aliases, which the served map skips
pub fn is_alias(&self) -> bool {
match self.repr() {
Repr::AliasShadowsModel { .. }
| Repr::AliasTaken { .. }
| Repr::AliasTargetMissing { .. }
| Repr::AliasRowField => true,
Repr::Invalid(_) => matches!(self.field().as_deref(), Some("aliases" | "alias_of")),
Repr::Json(_)
| Repr::Empty
| Repr::BelowMinimum { .. }
| Repr::Shrunk { .. }
| Repr::NotAnObject
| Repr::UnknownBlockKey
| Repr::UnknownTimezone { .. }
| Repr::DuplicatePrice { .. }
| Repr::SuffixInsideBlock { .. }
| Repr::TierRange
| Repr::TierReversed
| Repr::TierGap { .. }
| Repr::MixedTierBasis
| Repr::NotPricedByResults
| Repr::BandsFromTwoSources
| Repr::NotAnOffPeakRate
| Repr::NoOffPeakPrices
| Repr::NoOffPeakHours
| Repr::EmptyList
| Repr::InvalidDate { .. }
| Repr::MissingProvider
| Repr::UnreadField
| Repr::NoFallbackRules => false,
}
}
}
#[derive(Debug)]
pub struct Loaded<T> {
pub value: T,
@ -327,6 +403,8 @@ mod tests {
#[case::unknown_timezone(Repr::UnknownTimezone { zone: "invalid".into() }, Rejects::Row)]
#[case::alias_shadows_model(Repr::AliasShadowsModel { alias: "model".into() }, Rejects::Row)]
#[case::alias_taken(Repr::AliasTaken { alias: "alias".into(), taken_by: "owner".into() }, Rejects::Row)]
#[case::alias_target_missing(Repr::AliasTargetMissing { target: "owner".into() }, Rejects::Row)]
#[case::alias_row_field(Repr::AliasRowField, Rejects::Row)]
#[case::duplicate_price(Repr::DuplicatePrice { first: "a".into(), second: "b".into() }, Rejects::Row)]
#[case::suffix_inside_block(Repr::SuffixInsideBlock { block: "tiered_pricing" }, Rejects::Row)]
#[case::tier_range(Repr::TierRange, Rejects::Row)]

32
litellm-rust/crates/model-catalog/src/layers.rs
View file

@ -1,32 +0,0 @@
use std::sync::Arc;
use indexmap::IndexMap;
use serde_json::Value;
use crate::catalog::Catalog;
use crate::register::Registrations;
use crate::wire::Fields;
/// One field of a fetched row over the layer below, RFC 7386. An absent field inherits, JSON `null` is `Remove`
#[derive(Clone, Debug, PartialEq)]
pub enum Patch {
Set(Value),
Remove,
}
/// The four layers served rows are composed from, lowest first
#[derive(Clone, Debug)]
#[expect(
dead_code,
reason = "02 and 07 add the functions that read these layers"
)]
pub struct Layers {
/// 0, the wheel's embedded map
bundled: Arc<Catalog>,
/// 1, the last installed fetch, rejected rows left out
fetched: Option<Arc<Catalog>>,
/// 2, `register(persist=False)`, cleared by every install and refilled by the router replay
transient: IndexMap<String, Fields>,
/// 3, `register(persist=True)`, applied over every install
persisted: Registrations,
}

7
litellm-rust/crates/model-catalog/src/lib.rs
View file

@ -17,6 +17,9 @@
//! failing the document.
//! - `FallbackGeneralizations` / `FallbackRule`: regex rules that map unknown model ids onto known
//! families. They are data in the file, not entries.
//! - `CostMap`: an installed document with its provenance, which a host builds its live table from.
//! `Registrations` is the ordered `register_model` overlay over the host's own values, and
//! `write_pass` writes registrations into a host's table through `RegistrationTable`.
//! - Validation: `Catalog::check_size` checks a catalog against `IntegrityLimits` (a remote copy
//! that lost most of its models is refused), and `Catalog::row_errors` validates one row.
//!
@ -37,9 +40,9 @@ macro_rules_attribute::attribute_alias! {
mod bundled;
mod catalog;
mod cost_map;
mod error;
mod fallback;
mod layers;
mod lower;
mod model;
mod model_info;
@ -51,9 +54,9 @@ pub mod wire;
pub use bundled::*;
pub use catalog::*;
pub use cost_map::*;
pub use error::*;
pub use fallback::*;
pub use layers::*;
pub use model::*;
pub use model_info::*;
pub use register::*;

266
litellm-rust/crates/model-catalog/src/register.rs
View file

@ -1,26 +1,252 @@
use std::collections::{BTreeSet, HashMap};
use std::collections::BTreeMap;
use indexmap::IndexMap;
use crate::wire::Fields;
/// Where one `register` call writes and which row it inherits from
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct RegistrationTarget {
/// Always the caller's key, never a case variant of it
pub write: String,
/// The first candidate that is an exact key whose row passes `provider_matches`
pub inherit_from: Option<String>,
/// `register_model(persist_across_reloads=True)` calls, replayed above every installed document
#[derive(Clone, Debug, PartialEq)]
pub struct Registrations<V> {
/// Each registered key's merged value, ordered by its latest registration
entries: IndexMap<String, V>,
/// The registered keys whose write landed on each row key, in landing order
landed_on: IndexMap<String, Vec<String>>,
}
/// Python's inheritance candidates in main's order, `get_model_info`'s then `_resolve_builtin_model_cost_entry`'s
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct RegistrationCandidates(pub Vec<String>);
/// Layer 3, and the fields `register` last wrote, which litellm owns again
#[derive(Clone, Debug, Default)]
#[expect(dead_code, reason = "07 adds the functions that read registrations")]
pub struct Registrations {
persisted: IndexMap<String, Fields>,
registered_fields: HashMap<String, BTreeSet<String>>,
impl<V> Default for Registrations<V> {
fn default() -> Self {
Self {
entries: IndexMap::new(),
landed_on: IndexMap::new(),
}
}
}
/// One registration written to a host's table: the key it was registered under and the row key it landed on
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Landed {
pub registered: String,
pub written: String,
}
impl<V> Registrations<V> {
/// Keeps `value` as the key's registration, already merged with `get(key)` by the host, and makes it the latest
pub fn register(&mut self, key: &str, value: V) {
self.entries.shift_remove(key);
self.entries.insert(key.to_owned(), value);
}
pub fn get(&self, key: &str) -> Option<&V> {
self.entries.get(key)
}
pub fn record(&mut self, landed: &[Landed]) {
for Landed {
registered,
written,
} in landed
{
let names = self.landed_on.entry(written.clone()).or_default();
if !names.contains(registered) {
names.push(registered.clone());
}
}
}
/// Every registration, oldest first
pub fn entries(&self) -> impl Iterator<Item = (&str, &V)> {
self.entries
.iter()
.map(|(key, value)| (key.as_str(), value))
}
/// The registrations that apply to the row under `key`, oldest first: `key`'s own and every one that landed on it
pub fn over<'a>(&'a self, key: &str) -> Vec<&'a V> {
let landed = self.landed_on.get(key).into_iter().flatten();
let by_order: BTreeMap<usize, &V> = landed
.map(String::as_str)
.chain([key])
.filter_map(|name| self.entries.get_full(name))
.map(|(order, _, value)| (order, value))
.collect();
by_order.into_values().collect()
}
pub fn map<W>(&self, mut convert: impl FnMut(&V) -> W) -> Registrations<W> {
Registrations {
entries: self
.entries
.iter()
.map(|(key, value)| (key.clone(), convert(value)))
.collect(),
landed_on: self.landed_on.clone(),
}
}
}
/// A host's live table that registrations are written into
pub trait RegistrationTable {
type Registration;
type Error;
/// Writes `registration` where it lands, `None` when this table is no longer the one the host publishes
fn place(
&mut self,
key: &str,
registration: &Self::Registration,
) -> Result<Option<Placed>, Self::Error>;
}
/// Where a registration was written and whether a built-in row matched it
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Placed {
pub key: String,
pub matched_builtin: bool,
}
#[derive(Debug, PartialEq)]
pub enum Pass<'a, R> {
/// Every entry was written. `unmatched` are those no built-in row matched
Finished {
landed: Vec<Landed>,
unmatched: Vec<(&'a str, &'a R)>,
},
/// The table stopped being the published one before every entry was written
Swapped { landed: Vec<Landed> },
}
impl<R> Pass<'_, R> {
pub fn landed(&self) -> &[Landed] {
match self {
Pass::Finished { landed, .. } | Pass::Swapped { landed } => landed,
}
}
}
/// Writes `entries` into `table` in order
pub fn write_pass<'a, T: RegistrationTable>(
table: &mut T,
entries: &'a [(String, T::Registration)],
) -> Result<Pass<'a, T::Registration>, T::Error> {
let mut landed = Vec::with_capacity(entries.len());
let mut unmatched = Vec::new();
for (key, registration) in entries {
let Some(written) = table.place(key, registration)? else {
return Ok(Pass::Swapped { landed });
};
if !written.matched_builtin {
unmatched.push((key.as_str(), registration));
}
landed.push(Landed {
registered: key.clone(),
written: written.key,
});
}
Ok(Pass::Finished { landed, unmatched })
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use super::*;
fn landed(registered: &str, written: &str) -> Landed {
Landed {
registered: registered.to_owned(),
written: written.to_owned(),
}
}
#[rstest]
fn reregistering_a_key_makes_it_the_latest() {
let mut registrations = Registrations::default();
registrations.register("a", 1);
registrations.register("b", 2);
registrations.register("a", 3);
assert_eq!(
registrations.entries().collect::<Vec<_>>(),
[("b", &2), ("a", &3)]
);
}
#[rstest]
fn over_applies_names_landed_on_the_row_in_registration_order() {
let mut registrations = Registrations::default();
registrations.register("openai/o3-pro", 1);
registrations.register("o3-pro", 2);
registrations.register("other", 3);
registrations.record(&[
landed("openai/o3-pro", "o3-pro"),
landed("o3-pro", "o3-pro"),
]);
registrations.register("openai/o3-pro", 4);
assert_eq!(registrations.over("o3-pro"), [&2, &4]);
}
#[rstest]
fn over_applies_a_keys_own_registration_that_landed_on_another_row() {
let mut registrations = Registrations::default();
registrations.register("openai/o3-pro", 1);
registrations.record(&[landed("openai/o3-pro", "o3-pro")]);
assert_eq!(
(
registrations.over("openai/o3-pro"),
registrations.over("o3-pro")
),
(vec![&1], vec![&1])
);
}
/// A table published until `swaps_after` writes, whose rows match a built-in only when the key starts with `b`
struct Table {
writes: Vec<(String, i32)>,
swaps_after: usize,
}
impl RegistrationTable for Table {
type Registration = i32;
type Error = ();
fn place(&mut self, key: &str, registration: &i32) -> Result<Option<Placed>, ()> {
if self.writes.len() >= self.swaps_after {
return Ok(None);
}
let written = key.to_uppercase();
self.writes.push((written.clone(), *registration));
Ok(Some(Placed {
matched_builtin: written.starts_with('B'),
key: written,
}))
}
}
#[rstest]
#[case::published_throughout(
usize::MAX,
Pass::Finished { landed: vec![landed("a", "A"), landed("b", "B")], unmatched: vec![("a", &1)] },
&[("A", 1), ("B", 2)],
)]
#[case::swapped_after_the_first_write(
1,
Pass::Swapped { landed: vec![landed("a", "A")] },
&[("A", 1)],
)]
fn a_pass_writes_entries_in_order_until_the_table_is_swapped(
#[case] swaps_after: usize,
#[case] expected: Pass<'static, i32>,
#[case] writes: &[(&str, i32)],
) {
let entries = [("a".to_owned(), 1), ("b".to_owned(), 2)];
let mut table = Table {
writes: Vec::new(),
swaps_after,
};
assert_eq!(write_pass(&mut table, &entries), Ok(expected));
assert_eq!(
table.writes,
writes
.iter()
.map(|(key, value)| ((*key).to_owned(), *value))
.collect::<Vec<_>>()
);
}
}

12
litellm-rust/crates/model-catalog/src/schema.rs
View file

@ -104,9 +104,18 @@ pub fn registry_json_schema() -> Schema {
"title".into(),
"LiteLLM model prices and context window registry".into(),
);
definitions.insert(
"aliasRow".into(),
json!({
"type": "object",
"required": ["alias_of"],
"properties": {"alias_of": {"type": "string"}},
"additionalProperties": false,
}),
);
root.insert(
"additionalProperties".into(),
json!({"$ref": "#/$defs/modelEntry"}),
json!({"anyOf": [{"$ref": "#/$defs/modelEntry"}, {"$ref": "#/$defs/aliasRow"}]}),
);
root.insert("$defs".into(), Value::Object(definitions));
Schema::from(root.clone())
@ -153,6 +162,7 @@ fn decorate_model_entry(schema: &mut Value) {
let rates = json!({FLAT_KEY_PATTERN: {"type": "number", "minimum": 0}});
let object = schema.as_object_mut().unwrap();
object.insert("required".into(), json!(["litellm_provider"]));
object.insert("not".into(), json!({"required": ["alias_of"]}));
object.insert("additionalProperties".into(), Value::Bool(true));
object.insert(
"patternProperties".into(),

2
litellm-rust/crates/model-catalog/src/store/AGENTS.md
View file

@ -24,7 +24,7 @@
| Source | Stamp | Blind spot |
|---|---|---|
| `Catalog`'s `CatalogRows` | `()`, every change goes through `Store::edit`, which drops the index | none |
| python-bridge's bound `litellm.model_cost` | the dict object and its `len` | a delete plus an add between two reads, or a `litellm_provider` rewrite on an existing key |
| python-bridge's bound `litellm.model_cost` | the dict object, its `len` and its last key | deletes plus adds between two reads that put the old last key back last, or a `litellm_provider` rewrite on an existing key |
A source with a blind spot calls `Store::invalidate` when it knows keys changed. Until then a key added in the blind spot is not found, and a key moved to another provider stays listed under the old one. Removed keys are still never returned

2
litellm-rust/crates/model-catalog/src/wire/attributes.rs
View file

@ -47,6 +47,8 @@ macro_rules! attributes {
}
attributes! {
/// Model this row is another name for; such a row carries no other key.
alias_of: String,
/// Other names this entry is served under.
aliases: Vec<String>,
audio_transcription_config: String,

79
litellm-rust/crates/model-catalog/tests/catalog.rs
View file

@ -139,6 +139,18 @@ fn a_bad_entry_is_reported_by_model_and_key_and_kept_raw() {
json!({"First": {"litellm_provider": "t", "aliases": ["shared"]}, "Second": {"litellm_provider": "t", "aliases": ["shared"]}}),
r#"Second: alias "shared" already belongs to First"#
)]
#[case::alias_of_target_missing(
json!({"Real": {"litellm_provider": "t"}, "ghost-latest": {"alias_of": "Ghost"}}),
r#"ghost-latest: alias_of "Ghost" is not a model name"#
)]
#[case::alias_of_row_with_other_fields(
json!({"Real": {"litellm_provider": "t"}, "real-latest": {"alias_of": "Real", "mode": "chat"}}),
"real-latest.mode: an alias_of row carries no other field"
)]
#[case::alias_of_claimed_by_a_list(
json!({"Real": {"litellm_provider": "t"}, "real-latest": {"alias_of": "Real"}, "Other": {"litellm_provider": "t", "aliases": ["real-latest"]}}),
r#"real-latest: alias "real-latest" already belongs to Other"#
)]
#[case::aliases_not_a_list(
json!({"a": {"litellm_provider": "t", "aliases": "bad"}}),
"a.aliases: invalid value: invalid type: string \"bad\", expected a sequence"
@ -399,3 +411,70 @@ fn a_bad_value_keeps_its_json_error_as_the_source_and_names_its_key_and_field()
assert_eq!(error.key().as_deref(), Some("open/router"));
assert_eq!(error.field().as_deref(), Some("input_cost_per_token"));
}
#[rstest]
fn an_alias_of_row_is_served_as_its_target_and_written_as_itself() {
let catalog = catalog(json!({
"Real": {"litellm_provider": "t", "input_cost_per_token": 1e-6},
"real-latest": {"alias_of": "Real"},
"After": {"litellm_provider": "t"}
}));
let found = catalog.lookup("REAL-LATEST").unwrap();
assert_eq!(
(found.matched_key, found.canonical_key),
("real-latest", "Real")
);
assert_eq!(
found.entry.field("input_cost_per_token"),
Some(&json!(1e-6))
);
assert_eq!(catalog.model_count(), 2);
assert_eq!(
catalog.served().map(|(key, _)| key).collect::<Vec<_>>(),
["Real", "After", "real-latest"]
);
assert_eq!(
written(&catalog),
r#"{"Real":{"litellm_provider":"t","input_cost_per_token":1e-6},"real-latest":{"alias_of":"Real"},"After":{"litellm_provider":"t"}}"#
);
}
#[rstest]
fn a_dangling_alias_of_row_is_rejected_and_dropped_from_the_written_document() {
let parsed = loaded(json!({
"Real": {"litellm_provider": "t"},
"ghost-latest": {"alias_of": "Ghost"}
}));
let (catalog, removed) = parsed.value.without_rejected(&parsed.errors);
assert_eq!(removed.len(), 1);
assert!(catalog.lookup("ghost-latest").is_none());
assert_eq!(written(&catalog), r#"{"Real":{"litellm_provider":"t"}}"#);
}
#[rstest]
fn an_alias_of_row_with_other_fields_is_never_served() {
let parsed = loaded(json!({
"Real": {"litellm_provider": "t"},
"real-latest": {"alias_of": "Real", "mode": "chat"}
}));
assert!(parsed.value.lookup("real-latest").is_none());
let (catalog, removed) = parsed.value.without_rejected(&parsed.errors);
assert_eq!(removed.len(), 1);
assert!(catalog.lookup("real-latest").is_none());
assert_eq!(written(&catalog), r#"{"Real":{"litellm_provider":"t"}}"#);
}
#[rstest]
fn a_pointer_row_that_repeats_a_listed_alias_is_dropped_and_the_model_is_kept() {
let parsed = loaded(json!({
"Real": {"litellm_provider": "t", "aliases": ["real-latest"]},
"real-latest": {"alias_of": "Real"}
}));
let (catalog, removed) = parsed.value.without_rejected(&parsed.errors);
assert_eq!(removed.len(), 1);
assert_eq!(catalog.lookup("real-latest").unwrap().canonical_key, "Real");
assert_eq!(
written(&catalog),
r#"{"Real":{"litellm_provider":"t","aliases":["real-latest"]}}"#
);
}

53
litellm-rust/crates/model-catalog/tests/cost_map.rs Normal file
View file

@ -0,0 +1,53 @@
use litellm_model_catalog::{
BUNDLED_REGISTRY, CostMap, ErrorKind, IntegrityLimits, Provenance, git_blob_id,
};
use rstest::rstest;
use serde_json::json;
fn any_size() -> IntegrityLimits {
IntegrityLimits::new(0, 0.0).unwrap()
}
#[rstest]
#[case::not_json(b"{".as_slice(), false, ErrorKind::Json)]
#[case::empty(b"{}".as_slice(), true, ErrorKind::Empty)]
fn a_refused_document_says_whether_it_failed_the_integrity_limits(
#[case] body: &[u8],
#[case] integrity: bool,
#[case] kind: ErrorKind,
) {
let error = CostMap::read(body, None, any_size()).unwrap_err();
assert_eq!((error.is_integrity(), error.kind()), (integrity, kind));
}
#[rstest]
fn provenance_is_the_read_bytes_and_etag() {
let body = br#"{"m":{"litellm_provider":"p"}}"#;
let cost_map = CostMap::read(body, Some("\"v1\"".into()), any_size()).unwrap();
assert_eq!(
cost_map.provenance(),
&Provenance {
revision: "eaa747f210a1527d5ea7138be0f768c47578ca68".to_owned(),
etag: Some("\"v1\"".into()),
}
);
assert_eq!(
CostMap::bundled().provenance().revision,
git_blob_id(BUNDLED_REGISTRY)
);
}
#[rstest]
fn every_name_of_an_aliased_row_lists_its_owner_first(
#[values("owner", "short", "owner-latest")] key: &str,
) {
let body = json!({
"owner": {"litellm_provider": "p", "aliases": ["short"]},
"owner-latest": {"alias_of": "owner"},
"plain": {"litellm_provider": "p"},
})
.to_string();
let cost_map = CostMap::read(body.as_bytes(), None, any_size()).unwrap();
assert_eq!(cost_map.names(key), ["owner", "short", "owner-latest"]);
assert!(cost_map.names("plain").is_empty());
}

4
litellm-rust/crates/model-catalog/tests/schema.rs
View file

@ -35,6 +35,7 @@ fn generated_registry_schema_validates_checked_in_catalog() {
#[case(json!({"example": {"litellm_provider": "test"}}))]
#[case(json!({"example": {"litellm_provider": "test", "future_field": true}}))]
#[case(json!({"sample_spec": {"litellm_provider": "placeholder"}}))]
#[case(json!({"example": {"litellm_provider": "test"}, "example-latest": {"alias_of": "example"}}))]
fn generated_registry_schema_keeps_reader_compatibility(#[case] document: Value) {
assert!(registry_validator().is_valid(&document));
}
@ -54,6 +55,9 @@ fn generated_registry_schema_keeps_reader_compatibility(#[case] document: Value)
#[case::empty_windows(json!({"litellm_provider": "test", "off_peak_pricing": {"windows": []}}))]
#[case::invalid_weekday(json!({"litellm_provider": "test", "off_peak_pricing": {"windows": [{"hours_utc": "00:00-01:00", "weekdays": [0]}]}}))]
#[case::invalid_aliases(json!({"litellm_provider": "test", "aliases": "wrong"}))]
#[case::alias_of_on_a_model(json!({"litellm_provider": "test", "alias_of": "other"}))]
#[case::alias_row_with_other_fields(json!({"alias_of": "other", "mode": "chat"}))]
#[case::alias_of_not_a_string(json!({"alias_of": 4}))]
#[case::non_object_model(json!(4))]
fn generated_registry_schema_rejects_invalid_entries(#[case] entry: Value) {
assert!(!registry_validator().is_valid(&json!({"example": entry})));

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

@ -19,6 +19,7 @@ huggingface = ["litellm-token-counter/huggingface"]
tiktoken = ["litellm-token-counter/tiktoken"]
[dependencies]
rustls.workspace = true
litellm-secret-value.workspace = true
fancy-regex.workspace = true
litellm-tracing.workspace = true
@ -30,6 +31,14 @@ bytes.workspace = true
futures-util.workspace = true
litellm-cache.workspace = true
litellm-error.workspace = true
litellm-circuit-breaker.workspace = true
litellm-redis-client.workspace = true
litellm-redis-connect.workspace = true
litellm-redis-store.workspace = true
litellm-redis-resolve.workspace = true
litellm-python-compat.workspace = true
litellm-redis-types.workspace = true
tracing.workspace = true
litellm-cache-response.workspace = true
serde.workspace = true
litellm-auth.workspace = true
@ -51,8 +60,10 @@ litellm-router-types.workspace = true
litellm-router-core.workspace = true
litellm-host-python.workspace = true
litellm-model-catalog.workspace = true
litellm-model-catalog-refresh.workspace = true
litellm-token-counter = { path = "../token-counter", default-features = false }
pyo3.workspace = true
indexmap = "2.14.0"
pyo3 = { workspace = true, features = ["indexmap"] }
prost.workspace = true
pyo3-async-runtimes.workspace = true
reqwest.workspace = true

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