scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-24 00:52:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 91

chore: merge main into litellm_gcs_native_cache

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Yujong Lee 2026-09-21 22:51:05 +00:00
commit 4e2047e9c5
276 changed files with 12498 additions and 4894 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

23
.circleci/config.yml
View file

@ -3050,28 +3050,29 @@ jobs:
- run:
name: Run Docker container with bad DATABASE_URL
command: |
set +e
docker run --name my-app \
-p 4000:4000 \
-e LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY=true \
-e DEFAULT_NUM_WORKERS_LITELLM_PROXY=1 \
-e DATABASE_URL="postgresql://wrong:wrong@wrong:5432/wrong" \
myapp:latest \
--port 4000 > docker_output.log 2>&1 || true
--port 4000 > docker_output.log 2>&1
echo "$?" > docker_exit_code
set -e
- run:
name: Display Docker logs
command: cat docker_output.log
- run:
name: Check for expected error
name: Proxy must refuse to serve on an unreachable database
command: |
if grep -q "Error: P1001: Can't reach database server at" docker_output.log && \
(grep -q "Database setup failed after multiple retries" docker_output.log || \
grep -q "ERROR: Application startup failed. Exiting." docker_output.log); then
echo "Expected error found. Test passed."
else
echo "Expected error not found. Test failed."
cat docker_output.log
exit 1
fi
fail() { echo "FAILED: $1"; cat docker_output.log; exit 1; }
exit_code="$(cat docker_exit_code)"
[ "$exit_code" -ne 0 ] || fail "proxy exited 0 with an unreachable database"
grep -q "P1001" docker_output.log || fail "log does not name the unreachable database server"
! grep -q "Application startup complete" docker_output.log || fail "proxy reached serving state"
! docker exec my-app true 2>/dev/null || fail "container is still running"
echo "Proxy refused to serve (exit $exit_code) and never reached startup. Test passed."
provider_replay_harness:
docker:

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

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

1
Makefile
View file

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

11
docker/README.md
View file

@ -2,6 +2,17 @@
This guide provides instructions for building and running the LiteLLM application using Docker and Docker Compose.
> **Just want to run LiteLLM?** This guide builds from source. To run the published
> image instead, use `docker-compose.quickstart.yml` in this directory — the
> two-service stack (gateway + Postgres) that the
> [Docker quickstart](https://docs.litellm.ai/docs/proxy/docker_quick_start) documents:
>
> ```bash
> curl -sSLO https://github.com/BerriAI/litellm/raw/main/docker/docker-compose.quickstart.yml
> printf 'LITELLM_MASTER_KEY=sk-%s\nLITELLM_SALT_KEY=sk-%s\n' "$(openssl rand -hex 32)" "$(openssl rand -hex 32)" > .env
> docker compose -f docker-compose.quickstart.yml up -d
> ```
## Prerequisites
- Docker

41
docker/docker-compose.quickstart.yml Normal file
View file

@ -0,0 +1,41 @@
# LiteLLM quickstart stack: the gateway plus a Postgres database that stores
# models, virtual keys, and spend logs. Used by
# https://docs.litellm.ai/docs/proxy/docker_quick_start
#
# curl -sSLO https://github.com/BerriAI/litellm/raw/main/docker/docker-compose.quickstart.yml
# printf 'LITELLM_MASTER_KEY=sk-%s\nLITELLM_SALT_KEY=sk-%s\n' "$(openssl rand -hex 32)" "$(openssl rand -hex 32)" > .env
# docker compose -f docker-compose.quickstart.yml up -d
#
# Compose reads .env from this directory. Keep it: regenerating LITELLM_SALT_KEY
# makes credentials already stored in the database unreadable. For anything
# beyond local evaluation, pin the image to a specific release tag.
services:
litellm:
image: docker.litellm.ai/berriai/litellm:main-stable
ports:
- "4000:4000"
environment:
LITELLM_MASTER_KEY: ${LITELLM_MASTER_KEY:?set it in .env - see the header of this file}
LITELLM_SALT_KEY: ${LITELLM_SALT_KEY:?set it in .env - see the header of this file}
DATABASE_URL: postgresql://litellm:litellm@db:5432/litellm
STORE_MODEL_IN_DB: "True"
depends_on:
db:
condition: service_healthy
db:
image: postgres:16
environment:
POSTGRES_USER: litellm
POSTGRES_PASSWORD: litellm
POSTGRES_DB: litellm
healthcheck:
test: ["CMD-SHELL", "pg_isready -U litellm"]
interval: 5s
timeout: 5s
retries: 10
volumes:
- postgres_data:/var/lib/postgresql/data
volumes:
postgres_data:

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

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

28
enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py
View file

@ -6,7 +6,7 @@ same route are non-inference and free.
"""
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING, Dict, Optional, cast
from typing import TYPE_CHECKING, Dict, Final, Optional, Protocol, cast
import litellm
from litellm._logging import verbose_proxy_logger
@ -22,11 +22,31 @@ from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.repositories.prisma_protocols import TableActions
from litellm.router import Router
TERMINAL_RESPONSE_STATUSES = frozenset({"completed", "failed", "cancelled", "incomplete"})
class _ManagedObjectRow(Protocol):
@property
def id(self) -> str: ...
@property
def unified_object_id(self) -> str: ...
@property
def created_by(self) -> str | None: ...
@property
def file_object(self) -> object: ...
def _managed_object_table(prisma_client: "PrismaClient") -> "TableActions[_ManagedObjectRow]":
table: Final[TableActions[_ManagedObjectRow]] = prisma_client.db.litellm_managedobjecttable
return table
class CheckResponsesCost:
def __init__(
self,
@ -128,7 +148,7 @@ class CheckResponsesCost:
f"CheckResponsesCost: stale cleanup failed (poll will continue): {cleanup_err}"
)
jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many(
jobs = await _managed_object_table(self.prisma_client).find_many(
where={
"status": {"in": ["queued", "in_progress"]},
"file_purpose": "response",
@ -138,7 +158,7 @@ class CheckResponsesCost:
)
verbose_proxy_logger.debug(f"Found {len(jobs)} response jobs to check")
completed_jobs = []
completed_jobs: Final[list[_ManagedObjectRow]] = []
for job in jobs:
unified_object_id = job.unified_object_id
@ -189,7 +209,7 @@ class CheckResponsesCost:
# Mark completed jobs in the database
if len(completed_jobs) > 0:
await self.prisma_client.db.litellm_managedobjecttable.update_many(
await _managed_object_table(self.prisma_client).update_many(
where={"id": {"in": [job.id for job in completed_jobs]}},
data={"status": "completed"},
)

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

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

3
litellm-proxy-extras/litellm_proxy_extras/migrations/20260921000000_add_password_reset_columns/migration.sql Normal file
View file

@ -0,0 +1,3 @@
ALTER TABLE "LiteLLM_UserTable" ADD COLUMN IF NOT EXISTS "password_reset_required" BOOLEAN;
ALTER TABLE "LiteLLM_UserTable" ADD COLUMN IF NOT EXISTS "last_breach_check_at" TIMESTAMP(3);

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

@ -246,6 +246,8 @@ model LiteLLM_UserTable {
organization_id String?
object_permission_id String?
password String?
password_reset_required Boolean?
last_breach_check_at DateTime?
teams String[] @default([])
user_role String?
max_budget Float?

128
litellm-rust/Cargo.lock generated
View file

@ -1377,6 +1377,18 @@ dependencies = [
"pin-project-lite",
]
[[package]]
name = "fallible-iterator"
version = "0.3.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2acce4a10f12dc2fb14a218589d4f1f62ef011b2d0cc4b3cb1bba8e94da14649"
[[package]]
name = "fallible-streaming-iterator"
version = "0.1.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7360491ce676a36bf9bb3c56c1aa791658183a54d2744120f27285738d90465a"
[[package]]
name = "fancy-regex"
version = "0.17.0"
@ -1428,6 +1440,12 @@ version = "1.0.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1"
[[package]]
name = "foldhash"
version = "0.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb"
[[package]]
name = "form_urlencoded"
version = "1.2.2"
@ -1892,11 +1910,32 @@ version = "0.12.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888"
[[package]]
name = "hashbrown"
version = "0.16.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100"
dependencies = [
"foldhash",
]
[[package]]
name = "hashbrown"
version = "0.17.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a"
dependencies = [
"foldhash",
]
[[package]]
name = "hashlink"
version = "0.12.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a596f1b20ed2cc5ecac41a164aaebc7258057060f06c0cf7a2ba3991ee7990fb"
dependencies = [
"hashbrown 0.17.1",
]
[[package]]
name = "heck"
@ -2277,6 +2316,12 @@ version = "2.12.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d98f6fed1fde3f8c21bc40a1abb88dd75e67924f9cffc3ef95607bad8017f8e2"
[[package]]
name = "iter-read"
version = "1.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "071ed4cc1afd86650602c7b11aa2e1ce30762a1c27193201cb5cee9c6ebb1294"
[[package]]
name = "itertools"
version = "0.13.0"
@ -2435,6 +2480,17 @@ version = "0.2.186"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66"
[[package]]
name = "libsqlite3-sys"
version = "0.38.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f1d20bef17f513b9b3004532233187769cd072d790971f4e4da0e346eb6401e8"
dependencies = [
"cc",
"pkg-config",
"vcpkg",
]
[[package]]
name = "linux-raw-sys"
version = "0.12.1"
@ -2540,6 +2596,21 @@ dependencies = [
"url",
]
[[package]]
name = "litellm-cache-disk"
version = "0.1.0"
dependencies = [
"litellm-cache",
"py_literal",
"rand 0.8.7",
"rstest",
"rusqlite",
"serde-pickle",
"serde_json",
"tempfile",
"tokio",
]
[[package]]
name = "litellm-cache-gcs"
version = "0.1.0"
@ -2758,6 +2829,7 @@ dependencies = [
"litellm-auth-gcp",
"litellm-cache",
"litellm-cache-azure-blob",
"litellm-cache-disk",
"litellm-cache-gcs",
"litellm-cache-memory",
"litellm-cache-redis",
@ -4049,6 +4121,16 @@ dependencies = [
"windows-sys 0.52.0",
]
[[package]]
name = "rsqlite-vfs"
version = "0.1.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c51c9ae4df8a7fba42103df5c621fa3c37eccf3a3c650879e90fc48b11cc192c"
dependencies = [
"hashbrown 0.16.1",
"thiserror 2.0.19",
]
[[package]]
name = "rstest"
version = "0.26.1"
@ -4089,6 +4171,21 @@ dependencies = [
"syn 2.0.119",
]
[[package]]
name = "rusqlite"
version = "0.40.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "23f2a97da3e3873c73cb2a2e71b35c40ff95e0b1eefa8d72d8499a6928c3b5b3"
dependencies = [
"bitflags 2.13.1",
"fallible-iterator",
"fallible-streaming-iterator",
"hashlink",
"libsqlite3-sys",
"smallvec",
"sqlite-wasm-rs",
]
[[package]]
name = "rustc-hash"
version = "2.1.3"
@ -4346,6 +4443,19 @@ dependencies = [
"serde_derive",
]
[[package]]
name = "serde-pickle"
version = "1.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b641fdc8bcf2781ee78b30c599700d64ad4f412976143e4c5d0b9df906bb4843"
dependencies = [
"byteorder",
"iter-read",
"num-bigint 0.4.8",
"num-traits",
"serde",
]
[[package]]
name = "serde_core"
version = "1.0.229"
@ -4573,6 +4683,18 @@ dependencies = [
"unicode-segmentation",
]
[[package]]
name = "sqlite-wasm-rs"
version = "0.5.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "dc3efc0da82635d7e1ced0053bbbfa8c7ab9645d0bf36ceb4f7127bb85315d75"
dependencies = [
"cc",
"js-sys",
"rsqlite-vfs",
"wasm-bindgen",
]
[[package]]
name = "sse-stream"
version = "0.2.6"
@ -5319,6 +5441,12 @@ version = "0.1.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65"
[[package]]
name = "vcpkg"
version = "0.2.15"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "accd4ea62f7bb7a82fe23066fb0957d48ef677f6eeb8215f372f52e48bb32426"
[[package]]
name = "veil"
version = "0.3.0"

1
litellm-rust/Cargo.toml
View file

@ -33,6 +33,7 @@ litellm-cache-azure-blob = { path = "crates/cache-azure-blob" }
litellm-cache-memory = { path = "crates/cache-memory" }
litellm-cache-redis = { path = "crates/cache-redis" }
litellm-cache-gcs = { path = "crates/cache-gcs" }
litellm-cache-disk = { path = "crates/cache-disk" }
litellm-cache-response = { path = "crates/cache-response" }
litellm-token-counter = { path = "crates/token-counter" }
litellm-token-counter-fast = { path = "crates/token-counter-fast" }

19
litellm-rust/crates/cache-disk/Cargo.toml Normal file
View file

@ -0,0 +1,19 @@
[package]
name = "litellm-cache-disk"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
litellm-cache.workspace = true
py_literal = "0.4.0"
rand.workspace = true
rusqlite = { version = "0.40", features = ["bundled"] }
serde-pickle = "1.2"
serde_json.workspace = true
tokio.workspace = true
[dev-dependencies]
rstest.workspace = true
tempfile = "3.27.0"

10
litellm-rust/crates/cache-disk/src/adapter.rs Normal file
View file

@ -0,0 +1,10 @@
use litellm_cache::Error;
use crate::StoredValue;
pub trait ValueAdapter: Send + Sync + 'static {
fn read(&self, value: StoredValue) -> Result<Option<Vec<u8>>, Error>;
fn write(&self, payload: Vec<u8>) -> StoredValue;
fn counter_seed(&self, value: Option<StoredValue>) -> Result<f64, Error>;
fn counter_value(&self, value: f64) -> StoredValue;
}

301
litellm-rust/crates/cache-disk/src/cache.rs Normal file
View file

@ -0,0 +1,301 @@
use std::{
path::Path,
sync::Arc,
time::{Duration, SystemTime, UNIX_EPOCH},
};
use litellm_cache::{
BaseCache, BatchCache, BatchEntry, CacheCodec, CacheConnectionResult, CacheConnectionStatus,
CounterCache, DeleteCache, Error, ExactCacheContext, FlushCache,
};
use crate::{DiskStore, DiskcacheSqliteStore, PythonDiskCacheAdapter, StoredValue, ValueAdapter};
pub struct DiskCache<S, D = DiskcacheSqliteStore, A = PythonDiskCacheAdapter> {
store: Arc<D>,
adapter: Arc<A>,
codec: S,
}
impl<S: CacheCodec> DiskCache<S> {
pub fn open(directory: impl AsRef<Path>, codec: S) -> Result<Self, Error> {
Ok(Self {
store: Arc::new(DiskcacheSqliteStore::open(directory)?),
adapter: Arc::new(PythonDiskCacheAdapter),
codec,
})
}
}
impl<S: CacheCodec, D: DiskStore> DiskCache<S, D, PythonDiskCacheAdapter> {
pub fn with_store(store: D, codec: S) -> Self {
Self {
store: Arc::new(store),
adapter: Arc::new(PythonDiskCacheAdapter),
codec,
}
}
}
impl<S: CacheCodec, D: DiskStore, A: ValueAdapter> DiskCache<S, D, A> {
pub fn with_adapter(store: D, adapter: A, codec: S) -> Self {
Self {
store: Arc::new(store),
adapter: Arc::new(adapter),
codec,
}
}
pub fn directory(&self) -> &Path {
self.store.directory()
}
fn decode_stored(&self, value: StoredValue) -> Result<Option<S::Value>, Error> {
let Some(bytes) = self.adapter.read(value)? else {
return Ok(None);
};
self.codec.decode(&bytes).map(Some)
}
async fn run_blocking<T, F>(store: Arc<D>, operation: F) -> Result<T, Error>
where
T: Send + 'static,
F: FnOnce(&D) -> Result<T, Error> + Send + 'static,
{
tokio::task::spawn_blocking(move || operation(&store))
.await
.map_err(|_| Error::Unavailable)?
}
}
impl<S: CacheCodec, D: DiskStore, A: ValueAdapter> BaseCache for DiskCache<S, D, A> {
type Value = S::Value;
type Context = ExactCacheContext;
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
context.ttl
}
fn set_cache(
&self,
key: &str,
value: Self::Value,
context: &Self::Context,
) -> Result<(), Error> {
let value = self.adapter.write(self.codec.encode(&value)?);
let expire_time = context.ttl.map(|ttl| unix_now() + ttl.as_secs_f64());
self.store.set(key, value, expire_time, unix_now())
}
fn get_cache(&self, key: &str, _: &Self::Context) -> Result<Option<Self::Value>, Error> {
self.store
.get(key, unix_now())?
.map(|value| self.decode_stored(value))
.transpose()
.map(|value| value.flatten())
}
async fn async_set_cache(
&self,
key: &str,
value: Self::Value,
context: ExactCacheContext,
) -> Result<(), Error> {
let value = self.adapter.write(self.codec.encode(&value)?);
let ttl = context.ttl;
let key = key.to_string();
Self::run_blocking(Arc::clone(&self.store), move |store| {
let expire_time = ttl.map(|ttl| unix_now() + ttl.as_secs_f64());
store.set(&key, value, expire_time, unix_now())
})
.await
}
async fn async_get_cache(
&self,
key: &str,
_: &ExactCacheContext,
) -> Result<Option<Self::Value>, Error> {
let key = key.to_string();
let value = Self::run_blocking(Arc::clone(&self.store), move |store| {
store.get(&key, unix_now())
})
.await?;
value
.map(|value| self.decode_stored(value))
.transpose()
.map(|value| value.flatten())
}
async fn async_set_cache_pipeline(
&self,
entries: Vec<(String, Self::Value)>,
context: ExactCacheContext,
) -> Result<(), Error> {
let entries = entries
.into_iter()
.map(|(key, value)| {
self.codec
.encode(&value)
.map(|value| (key, self.adapter.write(value)))
})
.collect::<Result<Vec<_>, _>>()?;
let expire_after = context.ttl;
Self::run_blocking(Arc::clone(&self.store), move |store| {
for (key, value) in entries {
let expire_time = expire_after.map(|ttl| unix_now() + ttl.as_secs_f64());
store.set(&key, value, expire_time, unix_now())?;
}
Ok(())
})
.await
}
async fn disconnect(&self) -> Result<(), Error> {
Ok(())
}
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
let result = Self::run_blocking(Arc::clone(&self.store), |store| {
store.probe().map(|_| CacheConnectionResult {
status: CacheConnectionStatus::Success,
message: "Disk cache connection test successful".into(),
error: None,
})
})
.await;
Ok(match result {
Ok(result) => result,
Err(error) => CacheConnectionResult {
status: CacheConnectionStatus::Failed,
message: format!("Disk cache connection failed: {error}"),
error: Some(error.to_string()),
},
})
}
}
impl<S: CacheCodec, D: DiskStore, A: ValueAdapter> BatchCache for DiskCache<S, D, A> {
fn batch_get_cache(
&self,
keys: &[String],
context: &ExactCacheContext,
) -> Result<Vec<BatchEntry<Self::Value>>, Error> {
keys.iter()
.map(|key| match self.get_cache(key, context) {
Ok(Some(value)) => Ok(BatchEntry::Hit(value)),
Ok(None) => Ok(BatchEntry::Miss),
Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid),
Err(error) => Err(error),
})
.collect()
}
async fn async_batch_get_cache(
&self,
keys: Vec<String>,
_: ExactCacheContext,
) -> Result<Vec<BatchEntry<Self::Value>>, Error> {
let values = Self::run_blocking(Arc::clone(&self.store), move |store| {
keys.into_iter()
.map(|key| store.get(&key, unix_now()).map(|value| (key, value)))
.collect::<Result<Vec<_>, _>>()
})
.await?;
values
.into_iter()
.map(|(_, value)| match value {
None => Ok(BatchEntry::Miss),
Some(value) => match self.decode_stored(value) {
Ok(Some(value)) => Ok(BatchEntry::Hit(value)),
Ok(None) => Ok(BatchEntry::Miss),
Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid),
Err(error) => Err(error),
},
})
.collect()
}
}
impl<S: CacheCodec, D: DiskStore, A: ValueAdapter> DeleteCache for DiskCache<S, D, A> {
fn delete_cache(&self, key: &str) -> Result<(), Error> {
self.store.pop(key, unix_now()).map(|_| ())
}
async fn async_delete_cache(&self, key: &str) -> Result<(), Error> {
let key = key.to_string();
Self::run_blocking(Arc::clone(&self.store), move |store| {
store.pop(&key, unix_now()).map(|_| ())
})
.await
}
}
impl<S: CacheCodec, D: DiskStore, A: ValueAdapter> FlushCache for DiskCache<S, D, A> {
fn flush_cache(&self) -> Result<(), Error> {
self.store.clear()
}
async fn async_flush_cache(&self) -> Result<(), Error> {
Self::run_blocking(Arc::clone(&self.store), |store| store.clear()).await
}
}
impl<S: CacheCodec<Value = f64>, D: DiskStore, A: ValueAdapter> CounterCache
for DiskCache<S, D, A>
{
fn increment_cache(
&self,
key: &str,
amount: f64,
context: ExactCacheContext,
) -> Result<f64, Error> {
increment(
self.adapter.as_ref(),
self.store.as_ref(),
key,
amount,
context.ttl,
)
}
async fn async_increment(
&self,
key: &str,
amount: f64,
context: ExactCacheContext,
) -> Result<f64, Error> {
let key = key.to_string();
let adapter = Arc::clone(&self.adapter);
Self::run_blocking(Arc::clone(&self.store), move |store| {
increment(adapter.as_ref(), store, &key, amount, context.ttl)
})
.await
}
}
fn increment<A: ValueAdapter, D: DiskStore>(
adapter: &A,
store: &D,
key: &str,
amount: f64,
ttl: Option<Duration>,
) -> Result<f64, Error> {
let mut result = None;
let mut apply = |current: Option<StoredValue>| {
let initial = adapter.counter_seed(current)?;
let value = initial + amount;
let stored = adapter.counter_value(value);
result = Some(value);
Ok((stored, ttl.map(|ttl| unix_now() + ttl.as_secs_f64())))
};
store.update(key, unix_now(), &mut apply)?;
result.ok_or(Error::InvalidEntry)
}
fn unix_now() -> f64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_secs_f64()
}

11
litellm-rust/crates/cache-disk/src/lib.rs Normal file
View file

@ -0,0 +1,11 @@
mod adapter;
mod cache;
mod python;
mod sqlite;
mod store;
pub use adapter::ValueAdapter;
pub use cache::DiskCache;
pub use python::PythonDiskCacheAdapter;
pub use sqlite::DiskcacheSqliteStore;
pub use store::{DiskStore, StoredValue};

77
litellm-rust/crates/cache-disk/src/python/mod.rs Normal file
View file

@ -0,0 +1,77 @@
mod value;
use litellm_cache::Error;
use py_literal::Value;
use crate::{StoredValue, ValueAdapter};
#[derive(Clone, Copy, Debug, Default)]
pub struct PythonDiskCacheAdapter;
impl PythonDiskCacheAdapter {
fn python_get_cache(value: StoredValue) -> Result<Option<Value>, Error> {
let value = match value {
StoredValue::Bytes(value) => Value::Bytes(value),
StoredValue::Text(value) => Value::String(value),
StoredValue::Integer(value) => Value::Integer(value.into()),
StoredValue::Float(value) => Value::Float(value),
StoredValue::Pickle(value) => value::from_pickle(&value)?,
};
if !value::is_truthy(&value) {
return Ok(None);
}
match value {
Value::String(text) => Ok(Some(
value::from_json_text(&text).unwrap_or(Value::String(text)),
)),
Value::Bytes(bytes) => match std::str::from_utf8(&bytes) {
Ok(text) => Ok(Some(
value::from_json_text(text).unwrap_or(Value::Bytes(bytes)),
)),
Err(_) => Ok(Some(Value::Bytes(bytes))),
},
value => Ok(Some(value)),
}
}
}
impl ValueAdapter for PythonDiskCacheAdapter {
fn read(&self, value: StoredValue) -> Result<Option<Vec<u8>>, Error> {
match value {
StoredValue::Text(value) => Ok((!value.is_empty()).then(|| value.into_bytes())),
StoredValue::Bytes(value) => Ok((!value.is_empty()).then_some(value)),
value => {
let Some(value) = Self::python_get_cache(value)? else {
return Ok(None);
};
value::to_json(&value).map(Some)
}
}
}
fn write(&self, payload: Vec<u8>) -> StoredValue {
StoredValue::Bytes(payload)
}
fn counter_seed(&self, value: Option<StoredValue>) -> Result<f64, Error> {
let Some(value) = value else {
return Ok(0.0);
};
let Some(value) = Self::python_get_cache(value)? else {
return Ok(0.0);
};
Ok(if value::is_int(&value) {
value::to_f64(&value).unwrap_or(0.0)
} else {
0.0
})
}
fn counter_value(&self, value: f64) -> StoredValue {
if value.fract() == 0.0 && value >= i64::MIN as f64 && value <= i64::MAX as f64 {
StoredValue::Integer(value as i64)
} else {
StoredValue::Float(value)
}
}
}

173
litellm-rust/crates/cache-disk/src/python/value.rs Normal file
View file

@ -0,0 +1,173 @@
use litellm_cache::Error;
use py_literal::Value;
use serde_json::{Map, Number};
pub(crate) fn from_pickle(bytes: &[u8]) -> Result<Value, Error> {
let value = serde_pickle::value_from_slice(bytes, Default::default())
.map_err(|_| Error::InvalidEntry)?;
from_pickle_value(value)
}
fn from_pickle_value(value: serde_pickle::Value) -> Result<Value, Error> {
match value {
serde_pickle::Value::None => Ok(Value::None),
serde_pickle::Value::Bool(value) => Ok(Value::Boolean(value)),
serde_pickle::Value::I64(value) => integer(value.to_string()),
serde_pickle::Value::Int(value) => integer(value.to_string()),
serde_pickle::Value::F64(value) => Ok(Value::Float(value)),
serde_pickle::Value::String(value) => Ok(Value::String(value)),
serde_pickle::Value::Bytes(value) => Ok(Value::Bytes(value)),
serde_pickle::Value::List(values) => values
.into_iter()
.map(from_pickle_value)
.collect::<Result<Vec<_>, _>>()
.map(Value::List),
serde_pickle::Value::Tuple(values) => values
.into_iter()
.map(from_pickle_value)
.collect::<Result<Vec<_>, _>>()
.map(Value::Tuple),
serde_pickle::Value::Set(values) => values
.into_iter()
.map(from_pickle_hashable)
.collect::<Result<Vec<_>, _>>()
.map(Value::Set),
serde_pickle::Value::FrozenSet(values) => values
.into_iter()
.map(from_pickle_hashable)
.collect::<Result<Vec<_>, _>>()
.map(Value::Set),
serde_pickle::Value::Dict(values) => values
.into_iter()
.map(|(key, value)| Ok((from_pickle_hashable(key)?, from_pickle_value(value)?)))
.collect::<Result<Vec<_>, Error>>()
.map(Value::Dict),
}
}
fn from_pickle_hashable(value: serde_pickle::HashableValue) -> Result<Value, Error> {
Ok(match value {
serde_pickle::HashableValue::None => Value::None,
serde_pickle::HashableValue::Bool(value) => Value::Boolean(value),
serde_pickle::HashableValue::I64(value) => integer(value.to_string())?,
serde_pickle::HashableValue::Int(value) => integer(value.to_string())?,
serde_pickle::HashableValue::F64(value) => Value::Float(value),
serde_pickle::HashableValue::Bytes(value) => Value::Bytes(value),
serde_pickle::HashableValue::String(value) => Value::String(value),
serde_pickle::HashableValue::Tuple(values) => Value::Tuple(
values
.into_iter()
.map(from_pickle_hashable)
.collect::<Result<Vec<_>, _>>()?,
),
serde_pickle::HashableValue::FrozenSet(values) => Value::Set(
values
.into_iter()
.map(from_pickle_hashable)
.collect::<Result<Vec<_>, _>>()?,
),
})
}
fn integer(value: String) -> Result<Value, Error> {
value.parse().map_err(|_| Error::InvalidEntry)
}
pub(crate) fn from_json(value: serde_json::Value) -> Value {
match value {
serde_json::Value::Null => Value::None,
serde_json::Value::Bool(value) => Value::Boolean(value),
serde_json::Value::Number(value) => {
if value.is_i64() || value.is_u64() {
integer(value.to_string())
.unwrap_or(Value::Float(value.as_f64().unwrap_or(f64::NAN)))
} else {
Value::Float(value.as_f64().unwrap_or(f64::NAN))
}
}
serde_json::Value::String(value) => Value::String(value),
serde_json::Value::Array(values) => {
Value::List(values.into_iter().map(from_json).collect())
}
serde_json::Value::Object(values) => Value::Dict(
values
.into_iter()
.map(|(key, value)| (Value::String(key), from_json(value)))
.collect(),
),
}
}
pub(crate) fn from_json_text(value: &str) -> Result<Value, Error> {
serde_json::from_str(value)
.map(from_json)
.map_err(|_| Error::InvalidEntry)
}
pub(crate) fn is_truthy(value: &Value) -> bool {
match value {
Value::None => false,
Value::Boolean(value) => *value,
Value::Integer(value) => value.to_string() != "0",
Value::Float(value) => *value != 0.0,
Value::Complex(value) => value.re != 0.0 || value.im != 0.0,
Value::String(value) => !value.is_empty(),
Value::Bytes(value) => !value.is_empty(),
Value::Tuple(value) | Value::List(value) | Value::Set(value) => !value.is_empty(),
Value::Dict(value) => !value.is_empty(),
}
}
pub(crate) fn is_int(value: &Value) -> bool {
matches!(value, Value::Integer(_) | Value::Boolean(_))
}
pub(crate) fn to_f64(value: &Value) -> Option<f64> {
match value {
Value::Integer(value) => value.to_string().parse().ok(),
Value::Boolean(value) => Some(if *value { 1.0 } else { 0.0 }),
_ => None,
}
}
pub(crate) fn to_json(value: &Value) -> Result<Vec<u8>, Error> {
serde_json::to_vec(&to_json_value(value)?).map_err(|_| Error::InvalidEntry)
}
fn to_json_value(value: &Value) -> Result<serde_json::Value, Error> {
Ok(match value {
Value::None => serde_json::Value::Null,
Value::Boolean(value) => serde_json::Value::Bool(*value),
Value::Integer(value) => serde_json::Value::Number(
value
.to_string()
.parse::<Number>()
.map_err(|_| Error::InvalidEntry)?,
),
Value::Float(value) => {
serde_json::Value::Number(Number::from_f64(*value).ok_or(Error::InvalidEntry)?)
}
Value::Complex(_) | Value::Bytes(_) => return Err(Error::InvalidEntry),
Value::String(value) => serde_json::Value::String(value.clone()),
Value::Tuple(values) | Value::List(values) | Value::Set(values) => {
serde_json::Value::Array(
values
.iter()
.map(to_json_value)
.collect::<Result<Vec<_>, _>>()?,
)
}
Value::Dict(values) => {
let values = values
.iter()
.map(|(key, value)| {
let Value::String(key) = key else {
return Err(Error::InvalidEntry);
};
Ok((key.clone(), to_json_value(value)?))
})
.collect::<Result<Map<String, serde_json::Value>, _>>()?;
serde_json::Value::Object(values)
}
})
}

817
litellm-rust/crates/cache-disk/src/sqlite.rs Normal file
View file

@ -0,0 +1,817 @@
use std::{
collections::HashMap,
fs::{self, OpenOptions},
io::Write,
path::{Path, PathBuf},
sync::Mutex,
};
use litellm_cache::Error;
use rand::RngCore;
use rusqlite::{Connection, OptionalExtension, params, types::Value};
use crate::{DiskStore, StoredValue};
const MODE_RAW: i64 = 1;
const MODE_BINARY: i64 = 2;
const MODE_TEXT: i64 = 3;
const MODE_PICKLE: i64 = 4;
const DEFAULT_DISK_MIN_FILE_SIZE: i64 = 2_i64.pow(15);
const DEFAULT_SIZE_LIMIT: i64 = 2_i64.pow(30);
const DEFAULT_CULL_LIMIT: i64 = 10;
pub struct DiskcacheSqliteStore {
directory: PathBuf,
connection: Mutex<Connection>,
min_file_size: usize,
eviction_policy: String,
size_limit: i64,
cull_limit: i64,
statistics: bool,
}
struct StoredColumns {
size: i64,
mode: i64,
filename: Option<String>,
value: Option<Value>,
}
struct Row {
rowid: i64,
mode: i64,
filename: Option<String>,
value: Value,
}
impl DiskcacheSqliteStore {
pub fn open(directory: impl AsRef<Path>) -> Result<Self, Error> {
let directory = directory.as_ref().to_path_buf();
fs::create_dir_all(&directory).map_err(|_| Error::Unavailable)?;
let directory = std::path::absolute(&directory).map_err(|_| Error::Unavailable)?;
let database = directory.join("cache.db");
let connection = Connection::open(database).map_err(|_| Error::Unavailable)?;
connection
.busy_timeout(std::time::Duration::from_secs(60))
.map_err(|_| Error::Unavailable)?;
let mut settings = read_settings(&connection)?;
for (key, value) in default_settings() {
settings.entry(key).or_insert(value);
}
for (key, value) in settings
.iter()
.filter(|(key, _)| key.starts_with("sqlite_"))
{
apply_pragma(&connection, key, value)?;
}
connection
.execute_batch(
"CREATE TABLE IF NOT EXISTS Settings (
key TEXT NOT NULL UNIQUE,
value
)",
)
.map_err(|_| Error::Unavailable)?;
for (key, value) in &settings {
if !matches!(key.as_str(), "count" | "size" | "hits" | "misses") {
connection
.execute(
"INSERT OR REPLACE INTO Settings VALUES (?, ?)",
params![key, value],
)
.map_err(|_| Error::Unavailable)?;
}
}
for (key, value) in [
("count", Value::Integer(0)),
("size", Value::Integer(0)),
("hits", Value::Integer(0)),
("misses", Value::Integer(0)),
] {
connection
.execute(
"INSERT OR IGNORE INTO Settings VALUES (?, ?)",
params![key, value],
)
.map_err(|_| Error::Unavailable)?;
}
connection
.execute_batch(
"CREATE TABLE IF NOT EXISTS Cache (
rowid INTEGER PRIMARY KEY,
key BLOB,
raw INTEGER,
store_time REAL,
expire_time REAL,
access_time REAL,
access_count INTEGER DEFAULT 0,
tag BLOB,
size INTEGER DEFAULT 0,
mode INTEGER DEFAULT 0,
filename TEXT,
value BLOB
);
CREATE UNIQUE INDEX IF NOT EXISTS Cache_key_raw ON Cache(key, raw);
CREATE INDEX IF NOT EXISTS Cache_expire_time ON Cache(expire_time);",
)
.map_err(|_| Error::Unavailable)?;
let eviction_policy = setting_string(&settings, "eviction_policy")
.unwrap_or_else(|| "least-recently-stored".to_string());
match eviction_policy.as_str() {
"none" => {}
"least-recently-stored" => {
connection
.execute_batch(
"CREATE INDEX IF NOT EXISTS Cache_store_time ON Cache(store_time)",
)
.map_err(|_| Error::Unavailable)?;
}
"least-recently-used" => {
connection
.execute_batch(
"CREATE INDEX IF NOT EXISTS Cache_access_time ON Cache(access_time)",
)
.map_err(|_| Error::Unavailable)?;
}
"least-frequently-used" => {
connection
.execute_batch(
"CREATE INDEX IF NOT EXISTS Cache_access_count ON Cache(access_count)",
)
.map_err(|_| Error::Unavailable)?;
}
_ => return Err(Error::Unavailable),
}
connection
.execute_batch(
"CREATE TRIGGER IF NOT EXISTS Settings_count_insert
AFTER INSERT ON Cache FOR EACH ROW BEGIN
UPDATE Settings SET value = value + 1
WHERE key = \"count\"; END;
CREATE TRIGGER IF NOT EXISTS Settings_count_delete
AFTER DELETE ON Cache FOR EACH ROW BEGIN
UPDATE Settings SET value = value - 1
WHERE key = \"count\"; END;
CREATE TRIGGER IF NOT EXISTS Settings_size_insert
AFTER INSERT ON Cache FOR EACH ROW BEGIN
UPDATE Settings SET value = value + NEW.size
WHERE key = \"size\"; END;
CREATE TRIGGER IF NOT EXISTS Settings_size_update
AFTER UPDATE ON Cache FOR EACH ROW BEGIN
UPDATE Settings
SET value = value + NEW.size - OLD.size
WHERE key = \"size\"; END;
CREATE TRIGGER IF NOT EXISTS Settings_size_delete
AFTER DELETE ON Cache FOR EACH ROW BEGIN
UPDATE Settings SET value = value - OLD.size
WHERE key = \"size\"; END;",
)
.map_err(|_| Error::Unavailable)?;
let min_file_size = setting_i64(&settings, "disk_min_file_size")
.unwrap_or(DEFAULT_DISK_MIN_FILE_SIZE)
.try_into()
.map_err(|_| Error::Unavailable)?;
let size_limit = setting_i64(&settings, "size_limit").unwrap_or(DEFAULT_SIZE_LIMIT);
let cull_limit = setting_i64(&settings, "cull_limit").unwrap_or(DEFAULT_CULL_LIMIT);
let statistics = setting_i64(&settings, "statistics").unwrap_or_default() != 0;
Ok(Self {
directory,
connection: Mutex::new(connection),
min_file_size,
eviction_policy,
size_limit,
cull_limit,
statistics,
})
}
fn set_locked(
&self,
connection: &Connection,
key: &str,
columns: StoredColumns,
expire_time: Option<f64>,
now: f64,
) -> Result<Vec<String>, Error> {
let mut cleanup = Vec::new();
if let Some(old_filename) = connection
.query_row(
"SELECT filename FROM Cache WHERE key = ? AND raw = 1",
params![key],
|row| row.get::<_, Option<String>>(0),
)
.optional()
.map_err(|_| Error::Unavailable)?
.flatten()
{
cleanup.push(old_filename);
}
let (size, mode, filename, value) =
(columns.size, columns.mode, columns.filename, columns.value);
let rowid = connection
.query_row(
"SELECT rowid FROM Cache WHERE key = ? AND raw = 1",
params![key],
|row| row.get::<_, i64>(0),
)
.optional()
.map_err(|_| Error::Unavailable)?;
if let Some(rowid) = rowid {
connection
.execute(
"UPDATE Cache SET store_time = ?, expire_time = ?, access_time = ?,
access_count = 0, tag = NULL, size = ?, mode = ?, filename = ?, value = ?
WHERE rowid = ?",
params![now, expire_time, now, size, mode, filename, value, rowid],
)
.map_err(|_| Error::Unavailable)?;
} else {
connection
.execute(
"INSERT INTO Cache(
key, raw, store_time, expire_time, access_time, access_count,
tag, size, mode, filename, value
) VALUES (?, 1, ?, ?, ?, 0, NULL, ?, ?, ?, ?)",
params![key, now, expire_time, now, size, mode, filename, value],
)
.map_err(|_| Error::Unavailable)?;
}
cleanup.extend(self.cull(connection, now)?);
Ok(cleanup)
}
fn cull(&self, connection: &Connection, now: f64) -> Result<Vec<String>, Error> {
if self.cull_limit <= 0 {
return Ok(Vec::new());
}
let mut cleanup = Vec::new();
let expired = connection
.prepare(
"SELECT rowid, filename FROM Cache
WHERE expire_time IS NOT NULL AND expire_time < ?
ORDER BY expire_time LIMIT ?",
)
.map_err(|_| Error::Unavailable)?
.query_map(params![now, self.cull_limit], |row| {
Ok((row.get::<_, i64>(0)?, row.get::<_, Option<String>>(1)?))
})
.map_err(|_| Error::Unavailable)?
.collect::<Result<Vec<_>, _>>()
.map_err(|_| Error::Unavailable)?;
for (_, filename) in &expired {
if let Some(filename) = filename {
cleanup.push(filename.clone());
}
}
for (rowid, _) in &expired {
connection
.execute("DELETE FROM Cache WHERE rowid = ?", params![rowid])
.map_err(|_| Error::Unavailable)?;
}
let remaining = self.cull_limit - i64::try_from(expired.len()).unwrap_or(self.cull_limit);
if remaining <= 0 || self.volume(connection)? < self.size_limit {
return Ok(cleanup);
}
let order = match self.eviction_policy.as_str() {
"none" => return Ok(cleanup),
"least-recently-stored" => "store_time",
"least-recently-used" => "access_time",
"least-frequently-used" => "access_count",
_ => return Err(Error::Unavailable),
};
let rows = connection
.prepare(&format!(
"SELECT rowid, filename FROM Cache ORDER BY {order} LIMIT ?"
))
.map_err(|_| Error::Unavailable)?
.query_map(params![remaining], |row| {
Ok((row.get::<_, i64>(0)?, row.get::<_, Option<String>>(1)?))
})
.map_err(|_| Error::Unavailable)?
.collect::<Result<Vec<_>, _>>()
.map_err(|_| Error::Unavailable)?;
for (_, filename) in &rows {
if let Some(filename) = filename {
cleanup.push(filename.clone());
}
}
for (rowid, _) in rows {
connection
.execute("DELETE FROM Cache WHERE rowid = ?", params![rowid])
.map_err(|_| Error::Unavailable)?;
}
Ok(cleanup)
}
fn volume(&self, connection: &Connection) -> Result<i64, Error> {
let page_count: i64 = connection
.query_row("PRAGMA page_count", [], |row| row.get(0))
.map_err(|_| Error::Unavailable)?;
let page_size: i64 = connection
.query_row("PRAGMA page_size", [], |row| row.get(0))
.map_err(|_| Error::Unavailable)?;
let size: i64 = connection
.query_row("SELECT value FROM Settings WHERE key = 'size'", [], |row| {
row.get(0)
})
.map_err(|_| Error::Unavailable)?;
Ok(page_count.saturating_mul(page_size).saturating_add(size))
}
}
impl DiskStore for DiskcacheSqliteStore {
fn directory(&self) -> &Path {
&self.directory
}
fn get(&self, key: &str, now: f64) -> Result<Option<StoredValue>, Error> {
let connection = self.connection.lock().map_err(|_| Error::Unavailable)?;
let select = "SELECT rowid, expire_time, mode, filename, value FROM Cache
WHERE key = ? AND raw = 1 AND (expire_time IS NULL OR expire_time > ?)";
let row = connection
.query_row(select, params![key, now], row_from_query)
.optional()
.map_err(|_| Error::Unavailable)?;
if !self.statistics && !has_get_update(&self.eviction_policy) {
return row
.map(|row| fetch_row(&self.directory, row))
.transpose()
.map(|value| value.flatten());
}
transactional(&connection, |connection| {
let row = connection
.query_row(select, params![key, now], row_from_query)
.optional()
.map_err(|_| Error::Unavailable)?;
let Some(row) = row else {
if self.statistics {
connection
.execute(
"UPDATE Settings SET value = value + 1 WHERE key = 'misses'",
[],
)
.map_err(|_| Error::Unavailable)?;
}
return Ok(None);
};
let rowid = row.rowid;
let value = fetch_row(&self.directory, row);
let hit = value.as_ref().is_ok_and(Option::is_some);
if hit && self.statistics {
connection
.execute(
"UPDATE Settings SET value = value + 1 WHERE key = 'hits'",
[],
)
.map_err(|_| Error::Unavailable)?;
} else if !hit && self.statistics {
connection
.execute(
"UPDATE Settings SET value = value + 1 WHERE key = 'misses'",
[],
)
.map_err(|_| Error::Unavailable)?;
}
if has_get_update(&self.eviction_policy) && hit {
let update = match self.eviction_policy.as_str() {
"least-recently-used" => "UPDATE Cache SET access_time = ? WHERE rowid = ?",
"least-frequently-used" => {
"UPDATE Cache SET access_count = access_count + 1 WHERE rowid = ?"
}
_ => return Err(Error::Unavailable),
};
if self.eviction_policy == "least-recently-used" {
connection
.execute(update, params![now, rowid])
.map_err(|_| Error::Unavailable)?;
} else {
connection
.execute(update, params![rowid])
.map_err(|_| Error::Unavailable)?;
}
}
value
})
}
fn set(
&self,
key: &str,
value: StoredValue,
expire_time: Option<f64>,
now: f64,
) -> Result<(), Error> {
let columns = store_value(&self.directory, self.min_file_size, value)?;
let new_filename = columns.filename.clone();
let connection = self.connection.lock().map_err(|_| Error::Unavailable)?;
let result = transactional(&connection, |connection| {
self.set_locked(connection, key, columns, expire_time, now)
});
match result {
Ok(cleanup) => {
cleanup_files(&self.directory, cleanup);
Ok(())
}
Err(error) => {
if let Some(filename) = new_filename {
remove_file(&self.directory, &filename);
}
Err(error)
}
}
}
fn pop(&self, key: &str, now: f64) -> Result<Option<StoredValue>, Error> {
let connection = self.connection.lock().map_err(|_| Error::Unavailable)?;
let selected = transactional(&connection, |connection| {
let row = connection
.query_row(
"SELECT rowid, expire_time, mode, filename, value FROM Cache
WHERE key = ? AND raw = 1
AND (expire_time IS NULL OR expire_time > ?)",
params![key, now],
row_from_query,
)
.optional()
.map_err(|_| Error::Unavailable)?;
let Some(row) = row else {
return Ok(None);
};
connection
.execute("DELETE FROM Cache WHERE rowid = ?", params![row.rowid])
.map_err(|_| Error::Unavailable)?;
Ok(Some(row))
})?;
let Some(row) = selected else {
return Ok(None);
};
let filename = row.filename.clone();
let result = fetch_row(&self.directory, row)?;
if let Some(filename) = filename {
remove_file(&self.directory, &filename);
}
Ok(result)
}
fn clear(&self) -> Result<(), Error> {
let connection = self.connection.lock().map_err(|_| Error::Unavailable)?;
let mut last_rowid = 0_i64;
loop {
let batch = transactional(&connection, |connection| {
let rows = connection
.prepare(
"SELECT rowid, filename FROM Cache
WHERE rowid > ? ORDER BY rowid LIMIT 100",
)
.map_err(|_| Error::Unavailable)?
.query_map(params![last_rowid], |row| {
Ok((row.get::<_, i64>(0)?, row.get::<_, Option<String>>(1)?))
})
.map_err(|_| Error::Unavailable)?
.collect::<Result<Vec<_>, _>>()
.map_err(|_| Error::Unavailable)?;
if rows.is_empty() {
return Ok(rows);
}
let ids = rows
.iter()
.map(|(rowid, _)| rowid.to_string())
.collect::<Vec<_>>()
.join(",");
connection
.execute(&format!("DELETE FROM Cache WHERE rowid IN ({ids})"), [])
.map_err(|_| Error::Unavailable)?;
Ok(rows)
})?;
if batch.is_empty() {
return Ok(());
}
last_rowid = batch.last().map(|(rowid, _)| *rowid).unwrap_or(last_rowid);
cleanup_files(
&self.directory,
batch
.into_iter()
.filter_map(|(_, filename)| filename)
.collect(),
);
}
}
fn update(
&self,
key: &str,
now: f64,
apply: &mut dyn FnMut(Option<StoredValue>) -> Result<(StoredValue, Option<f64>), Error>,
) -> Result<(), Error> {
let connection = self.connection.lock().map_err(|_| Error::Unavailable)?;
let mut created_filename = None;
let result = transactional(&connection, |connection| {
let current = connection
.query_row(
"SELECT rowid, expire_time, mode, filename, value FROM Cache
WHERE key = ? AND raw = 1
AND (expire_time IS NULL OR expire_time > ?)",
params![key, now],
row_from_query,
)
.optional()
.map_err(|_| Error::Unavailable)?
.map(|row| fetch_row(&self.directory, row))
.transpose()?
.flatten();
let (value, expire_time) = apply(current)?;
let columns = store_value(&self.directory, self.min_file_size, value)?;
created_filename = columns.filename.clone();
let cleanup = self.set_locked(connection, key, columns, expire_time, now)?;
Ok(cleanup)
});
match result {
Ok(cleanup) => {
cleanup_files(&self.directory, cleanup);
Ok(())
}
Err(error) => {
if let Some(filename) = created_filename {
remove_file(&self.directory, &filename);
}
Err(error)
}
}
}
fn probe(&self) -> Result<(), Error> {
let connection = self.connection.lock().map_err(|_| Error::Unavailable)?;
connection
.query_row(
"SELECT value FROM Settings WHERE key = 'count'",
[],
|row| row.get::<_, i64>(0),
)
.map(|_| ())
.map_err(|_| Error::Unavailable)
}
}
fn default_settings() -> HashMap<String, Value> {
HashMap::from([
("statistics".to_string(), Value::Integer(0)),
("tag_index".to_string(), Value::Integer(0)),
(
"eviction_policy".to_string(),
Value::Text("least-recently-stored".to_string()),
),
("size_limit".to_string(), Value::Integer(DEFAULT_SIZE_LIMIT)),
("cull_limit".to_string(), Value::Integer(DEFAULT_CULL_LIMIT)),
("sqlite_auto_vacuum".to_string(), Value::Integer(1)),
("sqlite_cache_size".to_string(), Value::Integer(8192)),
(
"sqlite_journal_mode".to_string(),
Value::Text("wal".to_string()),
),
(
"sqlite_mmap_size".to_string(),
Value::Integer(2_i64.pow(26)),
),
("sqlite_synchronous".to_string(), Value::Integer(1)),
(
"disk_min_file_size".to_string(),
Value::Integer(DEFAULT_DISK_MIN_FILE_SIZE),
),
("disk_pickle_protocol".to_string(), Value::Integer(5)),
])
}
fn read_settings(connection: &Connection) -> Result<HashMap<String, Value>, Error> {
let mut statement = match connection.prepare("SELECT key, value FROM Settings") {
Ok(statement) => statement,
Err(_) => return Ok(HashMap::new()),
};
statement
.query_map([], |row| Ok((row.get(0)?, row.get(1)?)))
.map_err(|_| Error::Unavailable)?
.collect::<Result<HashMap<_, _>, _>>()
.map_err(|_| Error::Unavailable)
}
fn apply_pragma(connection: &Connection, key: &str, value: &Value) -> Result<(), Error> {
let pragma = key.strip_prefix("sqlite_").ok_or(Error::Unavailable)?;
match value {
Value::Integer(value) => connection
.pragma_update(None, pragma, value)
.map_err(|_| Error::Unavailable),
Value::Text(value) => connection
.pragma_update(None, pragma, value)
.map_err(|_| Error::Unavailable),
_ => Err(Error::Unavailable),
}
}
fn setting_i64(settings: &HashMap<String, Value>, key: &str) -> Option<i64> {
match settings.get(key) {
Some(Value::Integer(value)) => Some(*value),
_ => None,
}
}
fn setting_string(settings: &HashMap<String, Value>, key: &str) -> Option<String> {
match settings.get(key) {
Some(Value::Text(value)) => Some(value.clone()),
_ => None,
}
}
fn has_get_update(policy: &str) -> bool {
matches!(policy, "least-recently-used" | "least-frequently-used")
}
fn row_from_query(row: &rusqlite::Row<'_>) -> rusqlite::Result<Row> {
Ok(Row {
rowid: row.get(0)?,
mode: row.get(2)?,
filename: row.get(3)?,
value: row.get(4)?,
})
}
fn fetch_row(directory: &Path, row: Row) -> Result<Option<StoredValue>, Error> {
match row.mode {
MODE_RAW => match row.value {
Value::Blob(value) => Ok(Some(StoredValue::Bytes(value))),
Value::Text(value) => Ok(Some(StoredValue::Text(value))),
Value::Integer(value) => Ok(Some(StoredValue::Integer(value))),
Value::Real(value) => Ok(Some(StoredValue::Float(value))),
Value::Null => Err(Error::InvalidEntry),
},
MODE_BINARY | MODE_PICKLE => {
let bytes = match row.value {
Value::Blob(value) => value,
Value::Null => {
let Some(value) = read_file(directory, row.filename.as_deref())? else {
return Ok(None);
};
value
}
_ => return Err(Error::InvalidEntry),
};
Ok(Some(if row.mode == MODE_BINARY {
StoredValue::Bytes(bytes)
} else {
StoredValue::Pickle(bytes)
}))
}
MODE_TEXT => {
let bytes = match row.value {
Value::Null => {
let Some(value) = read_file(directory, row.filename.as_deref())? else {
return Ok(None);
};
value
}
Value::Blob(value) => value,
Value::Text(value) => value.into_bytes(),
_ => return Err(Error::InvalidEntry),
};
Ok(Some(StoredValue::Text(
String::from_utf8(bytes).map_err(|_| Error::InvalidEntry)?,
)))
}
_ => Err(Error::InvalidEntry),
}
}
fn read_file(directory: &Path, filename: Option<&str>) -> Result<Option<Vec<u8>>, Error> {
let Some(filename) = filename else {
return Err(Error::InvalidEntry);
};
match fs::read(directory.join(filename)) {
Ok(value) => Ok(Some(value)),
Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(None),
Err(_) => Err(Error::Unavailable),
}
}
fn store_value(
directory: &Path,
min_file_size: usize,
value: StoredValue,
) -> Result<StoredColumns, Error> {
match value {
StoredValue::Integer(value) => Ok(StoredColumns {
size: 0,
mode: MODE_RAW,
filename: None,
value: Some(Value::Integer(value)),
}),
StoredValue::Float(value) => Ok(StoredColumns {
size: 0,
mode: MODE_RAW,
filename: None,
value: Some(Value::Real(value)),
}),
StoredValue::Text(value) if value.chars().count() < min_file_size => Ok(StoredColumns {
size: 0,
mode: MODE_RAW,
filename: None,
value: Some(Value::Text(value)),
}),
StoredValue::Text(value) => {
let bytes = value.into_bytes();
let filename = write_file(directory, &bytes)?;
Ok(StoredColumns {
size: i64::try_from(bytes.len()).map_err(|_| Error::Unavailable)?,
mode: MODE_TEXT,
filename: Some(filename),
value: None,
})
}
StoredValue::Bytes(value) if value.len() < min_file_size => Ok(StoredColumns {
size: 0,
mode: MODE_RAW,
filename: None,
value: Some(Value::Blob(value)),
}),
StoredValue::Bytes(value) => {
let filename = write_file(directory, &value)?;
Ok(StoredColumns {
size: i64::try_from(value.len()).map_err(|_| Error::Unavailable)?,
mode: MODE_BINARY,
filename: Some(filename),
value: None,
})
}
StoredValue::Pickle(value) if value.len() < min_file_size => Ok(StoredColumns {
size: 0,
mode: MODE_PICKLE,
filename: None,
value: Some(Value::Blob(value)),
}),
StoredValue::Pickle(value) => {
let filename = write_file(directory, &value)?;
Ok(StoredColumns {
size: i64::try_from(value.len()).map_err(|_| Error::Unavailable)?,
mode: MODE_PICKLE,
filename: Some(filename),
value: None,
})
}
}
}
fn write_file(directory: &Path, bytes: &[u8]) -> Result<String, Error> {
let mut random = [0_u8; 16];
rand::rngs::OsRng.fill_bytes(&mut random);
let hex = random
.iter()
.map(|byte| format!("{byte:02x}"))
.collect::<String>();
let filename = format!("{}/{}/{}.val", &hex[..2], &hex[2..4], &hex[4..]);
let path = directory.join(&filename);
if let Some(parent) = path.parent() {
fs::create_dir_all(parent).map_err(|_| Error::Unavailable)?;
}
let mut file = OpenOptions::new()
.write(true)
.create_new(true)
.open(path)
.map_err(|_| Error::Unavailable)?;
file.write_all(bytes).map_err(|_| Error::Unavailable)?;
Ok(filename)
}
fn cleanup_files(directory: &Path, filenames: Vec<String>) {
for filename in filenames {
remove_file(directory, &filename);
}
}
fn remove_file(directory: &Path, filename: &str) {
let path = directory.join(filename);
let _ = fs::remove_file(&path);
}
fn transactional<T>(
connection: &Connection,
operation: impl FnOnce(&Connection) -> Result<T, Error>,
) -> Result<T, Error> {
connection
.execute_batch("BEGIN IMMEDIATE")
.map_err(|_| Error::Unavailable)?;
match operation(connection) {
Ok(value) => {
connection
.execute_batch("COMMIT")
.map_err(|_| Error::Unavailable)?;
Ok(value)
}
Err(error) => {
let _ = connection.execute_batch("ROLLBACK");
Err(error)
}
}
}

33
litellm-rust/crates/cache-disk/src/store.rs Normal file
View file

@ -0,0 +1,33 @@
use std::path::Path;
use litellm_cache::Error;
#[derive(Clone, Debug, PartialEq)]
pub enum StoredValue {
Bytes(Vec<u8>),
Text(String),
Integer(i64),
Float(f64),
Pickle(Vec<u8>),
}
pub trait DiskStore: Send + Sync + 'static {
fn directory(&self) -> &Path;
fn get(&self, key: &str, now: f64) -> Result<Option<StoredValue>, Error>;
fn set(
&self,
key: &str,
value: StoredValue,
expire_time: Option<f64>,
now: f64,
) -> Result<(), Error>;
fn pop(&self, key: &str, now: f64) -> Result<Option<StoredValue>, Error>;
fn clear(&self) -> Result<(), Error>;
fn update(
&self,
key: &str,
now: f64,
apply: &mut dyn FnMut(Option<StoredValue>) -> Result<(StoredValue, Option<f64>), Error>,
) -> Result<(), Error>;
fn probe(&self) -> Result<(), Error>;
}

431
litellm-rust/crates/cache-disk/tests/cache.rs Normal file
View file

@ -0,0 +1,431 @@
use std::{
fs,
path::{Path, PathBuf},
sync::Arc,
thread,
time::Duration,
};
use litellm_cache::{
BaseCache, BatchCache, BatchEntry, CacheCodec, CounterCache, DeleteCache, ExactCacheContext,
FlushCache, JsonCodec,
};
use litellm_cache_disk::{DiskCache, DiskStore, DiskcacheSqliteStore, StoredValue, ValueAdapter};
use rstest::{fixture, rstest};
use rusqlite::Connection;
use serde_json::{Value, json};
use tempfile::TempDir;
struct Sandbox {
directory: TempDir,
}
#[fixture]
fn sandbox() -> Sandbox {
Sandbox {
directory: tempfile::tempdir().unwrap(),
}
}
impl Sandbox {
fn store(&self) -> DiskcacheSqliteStore {
DiskcacheSqliteStore::open(self.directory.path()).unwrap()
}
fn cache<V>(&self) -> DiskCache<JsonCodec<V>>
where
JsonCodec<V>: CacheCodec,
{
DiskCache::open(self.directory.path(), JsonCodec::new()).unwrap()
}
fn db(&self) -> Connection {
Connection::open(self.directory.path().join("cache.db")).unwrap()
}
fn value_files(&self) -> Vec<PathBuf> {
fn visit(directory: &Path, files: &mut Vec<PathBuf>) {
for entry in fs::read_dir(directory).unwrap() {
let path = entry.unwrap().path();
if path.is_dir() {
visit(&path, files);
} else if path.extension().is_some_and(|extension| extension == "val") {
files.push(path);
}
}
}
let mut files = Vec::new();
visit(self.directory.path(), &mut files);
files
}
}
#[rstest]
fn relative_store_directory_is_absolutized(sandbox: Sandbox) {
let relative = PathBuf::from(format!(
".litellm-cache-disk-{}",
sandbox
.directory
.path()
.file_name()
.unwrap()
.to_string_lossy()
));
let store = DiskcacheSqliteStore::open(&relative).unwrap();
assert!(store.directory().is_absolute());
assert!(store.directory().ends_with(&relative));
let directory = store.directory().to_path_buf();
drop(store);
fs::remove_dir_all(directory).unwrap();
}
#[derive(Clone, Copy, Debug, Default)]
struct TextAdapter;
impl ValueAdapter for TextAdapter {
fn read(&self, value: StoredValue) -> Result<Option<Vec<u8>>, litellm_cache::Error> {
match value {
StoredValue::Text(value) => Ok(Some(value.into_bytes())),
_ => Ok(None),
}
}
fn write(&self, payload: Vec<u8>) -> StoredValue {
StoredValue::Text(String::from_utf8(payload).unwrap())
}
fn counter_seed(&self, _: Option<StoredValue>) -> Result<f64, litellm_cache::Error> {
Ok(0.0)
}
fn counter_value(&self, value: f64) -> StoredValue {
if value.fract() == 0.0 {
StoredValue::Integer(value as i64)
} else {
StoredValue::Float(value)
}
}
}
#[rstest]
fn roundtrip_persists_and_reopens(sandbox: Sandbox) {
let context = ExactCacheContext::default();
let opened = sandbox.cache::<Value>();
opened
.set_cache("key", json!({"answer": 42}), &context)
.unwrap();
assert_eq!(
opened.get_cache("key", &context).unwrap(),
Some(json!({"answer": 42}))
);
drop(opened);
let reopened = sandbox.cache::<Value>();
assert_eq!(
reopened.get_cache("key", &context).unwrap(),
Some(json!({"answer": 42}))
);
}
#[rstest]
fn ttl_and_expired_culling_match_cache_contract(sandbox: Sandbox) {
let store = sandbox.store();
store
.set(
"expired",
StoredValue::Bytes(b"old".to_vec()),
Some(10.0),
0.0,
)
.unwrap();
assert_eq!(store.get("expired", 10.0).unwrap(), None);
store
.set("new", StoredValue::Bytes(b"new".to_vec()), None, 11.0)
.unwrap();
assert_eq!(
sandbox
.db()
.query_row("SELECT COUNT(*) FROM Cache", [], |row| row.get::<_, i64>(0))
.unwrap(),
1
);
assert_eq!(
sandbox
.db()
.query_row(
"SELECT value FROM Settings WHERE key = 'count'",
[],
|row| row.get::<_, i64>(0)
)
.unwrap(),
1
);
}
#[rstest]
fn batch_preserves_order_and_classifies_misses_and_invalid_values(sandbox: Sandbox) {
let store = sandbox.store();
store
.set(
"hit",
StoredValue::Bytes(br#"{"ok":true}"#.to_vec()),
None,
0.0,
)
.unwrap();
store
.set(
"invalid",
StoredValue::Pickle(vec![0x80, 0x05, 0x2e]),
None,
0.0,
)
.unwrap();
let entries = sandbox
.cache::<Value>()
.batch_get_cache(
&["hit".into(), "missing".into(), "invalid".into()],
&ExactCacheContext::default(),
)
.unwrap();
assert_eq!(
entries,
vec![
BatchEntry::Hit(json!({"ok": true})),
BatchEntry::Miss,
BatchEntry::Invalid
]
);
}
#[rstest]
#[case(StoredValue::Bytes(Vec::new()))]
#[case(StoredValue::Text(String::new()))]
#[case(StoredValue::Integer(0))]
#[case(StoredValue::Float(0.0))]
#[case(StoredValue::Pickle(vec![0x80, 0x05, 0x4e, 0x2e]))]
#[case(StoredValue::Pickle(vec![0x80, 0x05, 0x89, 0x2e]))]
#[case(StoredValue::Pickle(vec![0x80, 0x05, 0x4b, 0x00, 0x2e]))]
#[case(StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x0a, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x47, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x2e]))]
#[case(StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x04, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x8c, 0x00, 0x94, 0x2e]))]
#[case(StoredValue::Pickle(vec![0x80, 0x05, 0x5d, 0x94, 0x2e]))]
#[case(StoredValue::Pickle(vec![0x80, 0x05, 0x7d, 0x94, 0x2e]))]
#[case(StoredValue::Pickle(vec![0x80, 0x05, 0x29, 0x2e]))]
fn falsy_values_are_misses(sandbox: Sandbox, #[case] value: StoredValue) {
sandbox.store().set("key", value, None, 0.0).unwrap();
assert_eq!(
sandbox
.cache::<Value>()
.get_cache("key", &ExactCacheContext::default())
.unwrap(),
None
);
}
#[rstest]
#[case(Some(StoredValue::Integer(2)), 1.5, 3.5, "real")]
#[case(Some(StoredValue::Integer(2)), 1.0, 3.0, "integer")]
#[case(Some(StoredValue::Float(3.5)), 1.0, 1.0, "integer")]
#[case(Some(StoredValue::Text("not a number".into())), 2.0, 2.0, "integer")]
#[case(Some(StoredValue::Text("5".into())), 2.0, 7.0, "integer")]
#[case(Some(StoredValue::Text("3.5".into())), 2.0, 2.0, "integer")]
#[case(Some(StoredValue::Pickle(vec![0x80, 0x05, 0x88, 0x2e])), 1.0, 2.0, "integer")]
#[case(Some(StoredValue::Pickle(vec![0x80, 0x05, 0x4b, 0x02, 0x2e])), 1.0, 3.0, "integer")]
#[case(Some(StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x0a, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x7d, 0x94, 0x8c, 0x01, 0x61, 0x94, 0x4b, 0x01, 0x73, 0x2e])), 1.0, 1.0, "integer")]
#[case(Some(StoredValue::Pickle(vec![0x80, 0x05, 0x4e, 0x2e])), 4.0, 4.0, "integer")]
fn counters_follow_python_initialization(
sandbox: Sandbox,
#[case] initial: Option<StoredValue>,
#[case] amount: f64,
#[case] expected: f64,
#[case] sqlite_type: &str,
) {
if let Some(initial) = initial {
sandbox.store().set("counter", initial, None, 0.0).unwrap();
}
let cache = sandbox.cache::<f64>();
assert_eq!(
cache
.increment_cache("counter", amount, ExactCacheContext::default())
.unwrap(),
expected
);
assert_eq!(
sandbox
.db()
.query_row(
"SELECT typeof(value) FROM Cache WHERE key = 'counter'",
[],
|row| row.get::<_, String>(0)
)
.unwrap(),
sqlite_type
);
}
#[rstest]
fn counters_are_atomic_across_concurrent_callers(sandbox: Sandbox) {
let cache = Arc::new(sandbox.cache::<f64>());
let workers = (0..8)
.map(|_| {
let cache = Arc::clone(&cache);
thread::spawn(move || {
for _ in 0..25 {
cache
.increment_cache("counter", 1.0, ExactCacheContext::default())
.unwrap();
}
})
})
.collect::<Vec<_>>();
for worker in workers {
worker.join().unwrap();
}
assert_eq!(
cache
.increment_cache("counter", 0.0, ExactCacheContext::default())
.unwrap(),
200.0
);
}
#[rstest]
fn fractional_then_integer_increment_follows_python_behavior(sandbox: Sandbox) {
let cache = sandbox.cache::<f64>();
assert_eq!(
cache
.increment_cache("counter", 3.5, ExactCacheContext::default())
.unwrap(),
3.5
);
assert_eq!(
cache
.increment_cache("counter", 1.0, ExactCacheContext::default())
.unwrap(),
1.0
);
}
#[rstest]
fn increment_ttl_replacement_clears_expiry_without_ttl(sandbox: Sandbox) {
let cache = sandbox.cache::<f64>();
cache
.increment_cache(
"counter",
1.0,
ExactCacheContext {
ttl: Some(Duration::from_secs(60)),
},
)
.unwrap();
assert!(
sandbox
.db()
.query_row(
"SELECT expire_time IS NOT NULL FROM Cache WHERE key = 'counter'",
[],
|row| row.get::<_, bool>(0)
)
.unwrap()
);
cache
.increment_cache("counter", 1.0, ExactCacheContext::default())
.unwrap();
assert!(
!sandbox
.db()
.query_row(
"SELECT expire_time IS NOT NULL FROM Cache WHERE key = 'counter'",
[],
|row| row.get::<_, bool>(0)
)
.unwrap()
);
}
#[rstest]
fn custom_adapter_controls_storage_and_reads(sandbox: Sandbox) {
let cache = DiskCache::with_adapter(sandbox.store(), TextAdapter, JsonCodec::<Value>::new());
cache
.set_cache("key", json!({"answer": 42}), &ExactCacheContext::default())
.unwrap();
assert!(matches!(
sandbox.store().get("key", 0.0).unwrap(),
Some(StoredValue::Text(_))
));
assert_eq!(
cache
.get_cache("key", &ExactCacheContext::default())
.unwrap(),
Some(json!({"answer": 42}))
);
}
#[rstest]
fn delete_flush_and_spilled_file_replacement_clean_up_storage(sandbox: Sandbox) {
let large = vec![b'x'; 32 * 1024];
sandbox
.store()
.set("large", StoredValue::Bytes(large.clone()), None, 0.0)
.unwrap();
assert_eq!(sandbox.value_files().len(), 1);
sandbox
.store()
.set(
"large",
StoredValue::Bytes(vec![b'y'; 32 * 1024]),
None,
0.0,
)
.unwrap();
assert_eq!(sandbox.value_files().len(), 1);
sandbox.store().pop("large", 0.0).unwrap();
assert!(sandbox.value_files().is_empty());
sandbox
.store()
.set("a", StoredValue::Bytes(large.clone()), None, 0.0)
.unwrap();
sandbox
.store()
.set("b", StoredValue::Bytes(large), None, 0.0)
.unwrap();
sandbox.store().clear().unwrap();
assert!(sandbox.value_files().is_empty());
}
#[rstest]
#[tokio::test]
async fn async_operations_connection_and_delete_match_sync_operations(sandbox: Sandbox) {
let cache = sandbox.cache::<Value>();
let context = ExactCacheContext {
ttl: Some(Duration::from_secs(60)),
};
cache
.async_set_cache("a", json!(1), context.clone())
.await
.unwrap();
cache
.async_set_cache_pipeline(
vec![("b".into(), json!(2)), ("c".into(), json!(3))],
context.clone(),
)
.await
.unwrap();
assert_eq!(
cache.async_get_cache("a", &context).await.unwrap(),
Some(json!(1))
);
assert_eq!(
cache
.async_batch_get_cache(vec!["c".into(), "missing".into()], context.clone())
.await
.unwrap(),
vec![BatchEntry::Hit(json!(3)), BatchEntry::Miss]
);
cache.async_delete_cache("a").await.unwrap();
cache.async_flush_cache().await.unwrap();
assert_eq!(
cache.test_connection().await.unwrap().status,
litellm_cache::CacheConnectionStatus::Success
);
}

113
litellm-rust/crates/cache-disk/tests/python_compat.rs Normal file
View file

@ -0,0 +1,113 @@
use litellm_cache::Error;
use litellm_cache_disk::{PythonDiskCacheAdapter, StoredValue, ValueAdapter};
use rstest::rstest;
enum ReadExpectation {
Bytes(&'static [u8]),
Miss,
Invalid,
}
#[rstest]
#[case::pickled_dictionary_with_string_keys(
StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x0a, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x7d, 0x94, 0x8c, 0x01, 0x61, 0x94, 0x4b, 0x01, 0x73, 0x2e]),
ReadExpectation::Bytes(br#"{"a":1}"#)
)]
#[case::pickled_list_of_integers(
StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x09, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x5d, 0x94, 0x28, 0x4b, 0x01, 0x4b, 0x02, 0x65, 0x2e]),
ReadExpectation::Bytes(br#"[1,2]"#)
)]
#[case::pickled_tuple_of_integers(
StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x07, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x4b, 0x01, 0x4b, 0x02, 0x86, 0x94, 0x2e]),
ReadExpectation::Bytes(br#"[1,2]"#)
)]
#[case::pickled_set_of_integers(
StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x09, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x8f, 0x94, 0x28, 0x4b, 0x01, 0x4b, 0x02, 0x90, 0x2e]),
ReadExpectation::Bytes(br#"[1,2]"#)
)]
#[case::pickled_response_envelope(
StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x30, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x7d, 0x94, 0x28, 0x8c, 0x09, 0x74, 0x69, 0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d, 0x70, 0x94, 0x47, 0x3f, 0xf8, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x8c, 0x08, 0x72, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x94, 0x8c, 0x08, 0x7b, 0x22, 0x61, 0x22, 0x3a, 0x20, 0x31, 0x7d, 0x94, 0x75, 0x2e]),
ReadExpectation::Bytes(br#"{"response":"{\"a\": 1}","timestamp":1.5}"#)
)]
#[case::non_json_text(
StoredValue::Text("not json".into()),
ReadExpectation::Bytes(b"not json")
)]
#[case::json_text(
StoredValue::Text("{\"a\": 1}".into()),
ReadExpectation::Bytes(br#"{"a": 1}"#)
)]
#[case::non_utf8_bytes(
StoredValue::Bytes(vec![0xff, 0xfe]),
ReadExpectation::Bytes(&[0xff, 0xfe])
)]
#[case::integer_seven(StoredValue::Integer(7), ReadExpectation::Bytes(b"7"))]
#[case::float_one_point_five(StoredValue::Float(1.5), ReadExpectation::Bytes(b"1.5"))]
#[case::pickled_true(
StoredValue::Pickle(vec![0x80, 0x05, 0x88, 0x2e]),
ReadExpectation::Bytes(b"true")
)]
#[case::pickled_negative_integer(
StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x06, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x4a, 0xfd, 0xff, 0xff, 0xff, 0x2e]),
ReadExpectation::Bytes(b"-3")
)]
#[case::pickled_bytes(
StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x09, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x43, 0x05, 0x62, 0x79, 0x74, 0x65, 0x73, 0x94, 0x2e]),
ReadExpectation::Invalid
)]
#[case::pickled_dictionary_with_integer_key(
StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x0a, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x7d, 0x94, 0x4b, 0x01, 0x8c, 0x01, 0x61, 0x94, 0x73, 0x2e]),
ReadExpectation::Invalid
)]
#[case::pickled_complex(
StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x2e, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x8c, 0x08, 0x62, 0x75, 0x69, 0x6c, 0x74, 0x69, 0x6e, 0x73, 0x94, 0x8c, 0x07, 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x94, 0x93, 0x94, 0x47, 0x3f, 0xf0, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x47, 0x40, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x86, 0x94, 0x52, 0x94, 0x2e]),
ReadExpectation::Invalid
)]
#[case::truncated_pickle(
StoredValue::Pickle(vec![0x80, 0x05, 0x2e]),
ReadExpectation::Invalid
)]
#[case::empty_bytes(StoredValue::Bytes(Vec::new()), ReadExpectation::Miss)]
#[case::empty_text(StoredValue::Text(String::new()), ReadExpectation::Miss)]
#[case::zero_integer(StoredValue::Integer(0), ReadExpectation::Miss)]
#[case::zero_float(StoredValue::Float(0.0), ReadExpectation::Miss)]
#[case::pickled_none(StoredValue::Pickle(vec![0x80, 0x05, 0x4e, 0x2e]), ReadExpectation::Miss)]
#[case::pickled_false(StoredValue::Pickle(vec![0x80, 0x05, 0x89, 0x2e]), ReadExpectation::Miss)]
#[case::pickled_zero(StoredValue::Pickle(vec![0x80, 0x05, 0x4b, 0x00, 0x2e]), ReadExpectation::Miss)]
#[case::pickled_zero_float(StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x0a, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x47, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x2e]), ReadExpectation::Miss)]
#[case::pickled_empty_string(StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x04, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x8c, 0x00, 0x94, 0x2e]), ReadExpectation::Miss)]
#[case::pickled_empty_list(StoredValue::Pickle(vec![0x80, 0x05, 0x5d, 0x94, 0x2e]), ReadExpectation::Miss)]
#[case::pickled_empty_dictionary(StoredValue::Pickle(vec![0x80, 0x05, 0x7d, 0x94, 0x2e]), ReadExpectation::Miss)]
#[case::pickled_empty_tuple(StoredValue::Pickle(vec![0x80, 0x05, 0x29, 0x2e]), ReadExpectation::Miss)]
fn python_read_cases(#[case] row: StoredValue, #[case] expected: ReadExpectation) {
let result = PythonDiskCacheAdapter.read(row);
match expected {
ReadExpectation::Bytes(expected) => assert_eq!(result.unwrap().unwrap(), expected),
ReadExpectation::Miss => assert_eq!(result.unwrap(), None),
ReadExpectation::Invalid => assert!(matches!(result, Err(Error::InvalidEntry))),
}
}
#[rstest]
#[case::integer_two(Some(StoredValue::Integer(2)), 2.0)]
#[case::float_three_point_five(Some(StoredValue::Float(3.5)), 0.0)]
#[case::text_not_a_number(Some(StoredValue::Text("not a number".into())), 0.0)]
#[case::text_five(Some(StoredValue::Text("5".into())), 5.0)]
#[case::text_three_point_five(Some(StoredValue::Text("3.5".into())), 0.0)]
#[case::pickled_true(Some(StoredValue::Pickle(vec![0x80, 0x05, 0x88, 0x2e])), 1.0)]
#[case::pickled_two(Some(StoredValue::Pickle(vec![0x80, 0x05, 0x4b, 0x02, 0x2e])), 2.0)]
#[case::pickled_dictionary(Some(StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x0a, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x7d, 0x94, 0x8c, 0x01, 0x61, 0x94, 0x4b, 0x01, 0x73, 0x2e])), 0.0)]
#[case::missing(None, 0.0)]
#[case::pickled_none(Some(StoredValue::Pickle(vec![0x80, 0x05, 0x4e, 0x2e])), 0.0)]
fn python_counter_seed_cases(#[case] row: Option<StoredValue>, #[case] expected: f64) {
assert_eq!(PythonDiskCacheAdapter.counter_seed(row).unwrap(), expected);
}
#[rstest]
#[case::integer_three(3.0, StoredValue::Integer(3))]
#[case::fractional_three_point_five(3.5, StoredValue::Float(3.5))]
#[case::negative_zero(-0.0, StoredValue::Integer(0))]
#[case::large_float(1e300, StoredValue::Float(1e300))]
fn python_counter_value_cases(#[case] value: f64, #[case] expected: StoredValue) {
assert_eq!(PythonDiskCacheAdapter.counter_value(value), expected);
}

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

@ -25,6 +25,7 @@ litellm-cache-azure-blob.workspace = true
litellm-cache-memory.workspace = true
litellm-cache-redis.workspace = true
litellm-cache-gcs.workspace = true
litellm-cache-disk.workspace = true
litellm-cache-response.workspace = true
serde.workspace = true
litellm-auth.workspace = true

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

@ -1,4 +1,4 @@
use std::time::Duration;
use std::{path::PathBuf, time::Duration};
use litellm_cache::CacheType;
use litellm_cache_redis::{RedisNode, RedisTopology};
@ -26,6 +26,10 @@ pub(super) struct MemoryCacheConfig {
pub(super) max_entry_bytes: usize,
}
pub(super) struct DiskCacheConfig {
pub(super) directory: PathBuf,
}
#[derive(Debug, PartialEq)]
pub(super) enum RedisProtocol {
Resp2,
@ -102,6 +106,7 @@ pub(super) enum CacheBackendConfig {
Memory(MemoryCacheConfig),
Redis(Box<RedisCacheConfig>),
Gcs(GcsCacheConfig),
Disk(DiskCacheConfig),
AzureBlob(AzureBlobCacheConfig),
}
@ -118,6 +123,7 @@ pub(super) enum UnsupportedCacheConfig {
RedisConnection,
RedisOption,
GcsBucket,
DiskStore,
}
impl UnsupportedCacheConfig {
@ -129,6 +135,7 @@ impl UnsupportedCacheConfig {
Self::RedisConnection => "native Redis connection type is not implemented",
Self::RedisOption => "native Redis configuration requires Python",
Self::GcsBucket => "native GCS cache requires a configured bucket name",
Self::DiskStore => "native disk cache requires the built-in diskcache store",
}
}
}
@ -178,6 +185,13 @@ impl NativeCacheConfig {
}))),
Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)),
},
Some(CacheType::Disk) => match project_disk(&backend)? {
Ok(backend) => Ok(CacheConfigProjection::Native(Box::new(Self {
policy,
backend: CacheBackendConfig::Disk(backend),
}))),
Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)),
},
Some(CacheType::AzureBlob) => project_azure_blob(&backend).map(|backend| {
CacheConfigProjection::Native(Box::new(Self {
policy,
@ -188,7 +202,6 @@ impl NativeCacheConfig {
CacheType::RedisSemantic
| CacheType::ValkeySemantic
| CacheType::S3
| CacheType::Disk
| CacheType::QdrantSemantic,
)
| None => Ok(CacheConfigProjection::Unsupported(
@ -201,7 +214,9 @@ impl NativeCacheConfig {
let default_ttl = match &self.backend {
CacheBackendConfig::Memory(config) => Some(config.default_ttl),
CacheBackendConfig::Redis(config) => Some(config.default_ttl),
CacheBackendConfig::AzureBlob(_) | CacheBackendConfig::Gcs(_) => None,
CacheBackendConfig::Disk(_)
| CacheBackendConfig::AzureBlob(_)
| CacheBackendConfig::Gcs(_) => None,
};
if service.default_ttl() != default_ttl {
return Some("facade and native backend default TTLs must match");
@ -253,6 +268,17 @@ impl NativeCacheConfig {
Some("facade and native backend credentials must match")
}
CacheBackendConfig::Gcs(_) => None,
CacheBackendConfig::Disk(_) if service.kind() != "disk" => {
Some("facade and native backend types must match")
}
CacheBackendConfig::Disk(config) => {
let Some(directory) = service.directory() else {
return Some("facade and native backend types must match");
};
let native = std::fs::canonicalize(directory).ok();
let facade = std::fs::canonicalize(&config.directory).ok();
(native != facade).then_some("facade and native backend directories must match")
}
CacheBackendConfig::AzureBlob(config) => match service.azure_blob_identity() {
None => Some("facade and native backend types must match"),
Some((account_url, container))
@ -310,6 +336,21 @@ fn project_gcs(
}))
}
#[inline(never)]
fn project_disk(
backend: &Bound<'_, PyAny>,
) -> PyResult<Result<DiskCacheConfig, UnsupportedCacheConfig>> {
let store = backend.getattr("disk_cache")?;
if !instance_class_is(&store, "diskcache.core", "Cache")?
|| !instance_class_is(&store.getattr("_disk")?, "diskcache.core", "Disk")?
{
return Ok(Err(UnsupportedCacheConfig::DiskStore));
}
Ok(Ok(DiskCacheConfig {
directory: PathBuf::from(store.getattr("directory")?.extract::<String>()?),
}))
}
#[inline(never)]
fn project_redis(
backend: &Bound<'_, PyAny>,
@ -697,8 +738,8 @@ mod tests {
use litellm_cache_redis::{RedisNode, RedisTopology};
use super::{
CacheBackendConfig, CacheConfigProjection, CertificateRequirement, GcsCacheConfig,
NativeCacheConfig, RedisProtocol, UnsupportedCacheConfig,
CacheBackendConfig, CacheConfigProjection, CachePolicy, CertificateRequirement,
DiskCacheConfig, GcsCacheConfig, NativeCacheConfig, RedisProtocol, UnsupportedCacheConfig,
};
use crate::cache::native::NativeResponseCache;
@ -898,6 +939,85 @@ mod tests {
assert_eq!(reason.message(), "native Redis credentials require Python");
});
}
#[test]
fn projects_builtin_disk_configuration_and_rejects_custom_stores() {
Python::initialize();
Python::attach(|py| {
let root =
std::env::temp_dir().join(format!("litellm-disk-config-{}", std::process::id()));
let directory = root.to_string_lossy();
let disk_facade = facade(
py,
&format!(
"Cache = type('Cache', (), {{'__module__': 'diskcache.core'}})\n\
Disk = type('Disk', (), {{'__module__': 'diskcache.core'}})\n\
store = Cache()\n\
store._disk = Disk()\n\
store.directory = {directory:?}\n\
backend = SimpleNamespace(disk_cache=store)\n\
facade = SimpleNamespace(type='disk', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)"
),
);
let CacheConfigProjection::Native(config) =
NativeCacheConfig::project(&disk_facade).unwrap()
else {
panic!("disk cache should be supported");
};
let CacheBackendConfig::Disk(disk) = config.backend else {
panic!("expected disk configuration");
};
assert_eq!(disk.directory, root);
let matching = NativeResponseCache::disk(&directory).unwrap();
assert_eq!(
(NativeCacheConfig {
policy: config.policy,
backend: CacheBackendConfig::Disk(disk),
})
.service_mismatch(&matching),
None
);
let other = NativeResponseCache::disk(&root.join("other").to_string_lossy()).unwrap();
let mismatch = NativeCacheConfig {
policy: CachePolicy {
mode: "default-on".into(),
ttl: None,
namespace: None,
supported_call_types: None,
redis_flush_size: None,
semantic_cache_scope: "key".into(),
},
backend: CacheBackendConfig::Disk(DiskCacheConfig {
directory: root.clone(),
}),
};
assert_eq!(
mismatch.service_mismatch(&other),
Some("facade and native backend directories must match")
);
let custom = facade(
py,
&format!(
"CustomCache = type('CustomCache', (), {{'__module__': 'mypkg'}})\n\
CustomDisk = type('CustomDisk', (), {{'__module__': 'mypkg'}})\n\
store = CustomCache()\n\
store._disk = CustomDisk()\n\
store.directory = {directory:?}\n\
backend = SimpleNamespace(disk_cache=store)\n\
facade = SimpleNamespace(type='disk', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)"
),
);
let CacheConfigProjection::Unsupported(reason) =
NativeCacheConfig::project(&custom).unwrap()
else {
panic!("custom disk store must stay on Python");
};
assert_eq!(
reason.message(),
"native disk cache requires the built-in diskcache store"
);
});
}
#[test]
fn projects_cluster_startup_nodes_as_redis_topology() {

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

@ -34,6 +34,11 @@ struct RedisPoolGuard {
attributes: RedisPoolAttributes,
}
struct DiskStoreGuard {
reference: Py<PyAny>,
directory: String,
}
struct AzureBlobClientGuard {
sync_client: Py<PyAny>,
async_client: Py<PyAny>,
@ -46,7 +51,6 @@ enum ConnectionGuard {
RedisPool(RedisPoolGuard),
AzureBlob(AzureBlobClientGuard),
}
struct RedisPoolAttributes {
pool: &'static str,
connection_class: &'static str,
@ -68,6 +72,7 @@ const CLUSTER_POOL: RedisPoolAttributes = RedisPoolAttributes {
pub(super) struct FacadeGuard {
outer: ObjectGuard,
backend: ObjectGuard,
disk_store: Option<DiskStoreGuard>,
connection: ConnectionGuard,
}
@ -218,6 +223,26 @@ impl RedisPoolGuard {
}
}
impl DiskStoreGuard {
fn capture(backend: &Bound<'_, PyAny>) -> PyResult<Self> {
let store = backend.getattr("disk_cache")?;
Ok(Self {
reference: store.clone().unbind(),
directory: store.getattr("directory")?.extract()?,
})
}
fn matches(&self, py: Python<'_>, backend: &Bound<'_, PyAny>) -> PyResult<bool> {
let store = backend.getattr("disk_cache")?;
Ok(self.reference.bind(py).is(&store)
&& self.directory == store.getattr("directory")?.extract::<String>()?)
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&self.reference)
}
}
impl AzureBlobClientGuard {
fn capture(backend: &Bound<'_, PyAny>) -> PyResult<Self> {
let sync_client = backend.getattr("container_client")?;
@ -296,6 +321,7 @@ impl FacadeGuard {
"redis",
),
("gcs", _) => ("litellm.caching.gcs_cache", "GCSCache", "gcs"),
("disk", _) => ("litellm.caching.disk_cache", "DiskCache", "disk"),
("azure-blob", _) => (
"litellm.caching.azure_blob_cache",
"AzureBlobCache",
@ -349,6 +375,9 @@ impl FacadeGuard {
"path_service_account",
],
)?,
disk_store: (kind == "disk")
.then(|| DiskStoreGuard::capture(&backend))
.transpose()?,
connection: ConnectionGuard::capture(kind, cluster, &backend)?,
})
}
@ -361,12 +390,20 @@ impl FacadeGuard {
if !self.backend.matches(py, &backend)? {
return Ok(false);
}
if let Some(guard) = &self.disk_store
&& !guard.matches(py, &backend)?
{
return Ok(false);
}
self.connection.matches(py, &backend)
}
pub(super) fn traverse(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
self.outer.traverse(&visit)?;
self.backend.traverse(&visit)?;
if let Some(guard) = &self.disk_store {
guard.traverse(&visit)?;
}
self.connection.traverse(&visit)
}
}

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

@ -91,6 +91,18 @@ impl CacheTestHandle {
})
}
#[staticmethod]
#[pyo3(signature = (directory))]
fn disk(py: Python<'_>, directory: String) -> PyResult<Self> {
let service =
release_gil(py, move || NativeResponseCache::disk(&directory)).map_err(cache_error)?;
Ok(Self {
service,
guard: None,
pid: std::process::id(),
})
}
#[staticmethod]
#[pyo3(signature = (account_url, container))]
fn azure_blob(py: Python<'_>, account_url: String, container: String) -> PyResult<Self> {

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

@ -1,7 +1,8 @@
use std::{sync::Arc, time::Duration};
use std::{path::Path, sync::Arc, time::Duration};
use litellm_cache::{CacheCodec, CacheConnectionResult, Error};
use litellm_cache_azure_blob::AzureBlobCache;
use litellm_cache_disk::DiskCache;
use litellm_cache_gcs::{GcsCache, GcsConfig, StaticTokenSource};
use litellm_cache_memory::InMemoryCache;
use litellm_cache_redis::{RedisCache, RedisTopology};
@ -18,6 +19,7 @@ pub(super) enum NativeResponseCache {
buffer: Option<Arc<WriteBuffer>>,
},
Gcs(Arc<ResponseCache<GcsCache<ResponseCacheCodec>>>),
Disk(Arc<ResponseCache<DiskCache<ResponseCacheCodec>>>),
AzureBlob(Arc<ResponseCache<AzureBlobCache<ResponseCacheCodec>>>),
}
@ -49,6 +51,10 @@ impl NativeResponseCache {
buffer: None,
})
}
pub fn disk(directory: &str) -> Result<Self, Error> {
let cache = DiskCache::open(directory, ResponseCacheCodec)?;
Ok(Self::Disk(Arc::new(ResponseCache::new(Arc::new(cache)))))
}
pub fn gcs(config: GcsConfig, token: Option<String>) -> Result<Self, Error> {
let backend = match token {
@ -81,7 +87,7 @@ impl NativeResponseCache {
cache.backend().account_url(),
cache.backend().container_name(),
)),
Self::Memory(_) | Self::Redis { .. } | Self::Gcs(_) => None,
Self::Memory(_) | Self::Redis { .. } | Self::Disk(_) | Self::Gcs(_) => None,
}
}
}
@ -92,6 +98,7 @@ impl NativeResponseCache {
Self::Memory(_) => "memory",
Self::Redis { .. } => "redis",
Self::Gcs(_) => "gcs",
Self::Disk(_) => "disk",
Self::AzureBlob(_) => "azure-blob",
}
}
@ -101,13 +108,14 @@ impl NativeResponseCache {
Self::Memory(cache) => cache.default_ttl(),
Self::Redis { cache, .. } => cache.default_ttl(),
Self::Gcs(cache) => cache.default_ttl(),
Self::Disk(cache) => cache.default_ttl(),
Self::AzureBlob(cache) => cache.default_ttl(),
}
}
pub fn namespace(&self) -> Option<&str> {
match self {
Self::Memory(_) | Self::AzureBlob(_) => None,
Self::Memory(_) | Self::Disk(_) | Self::AzureBlob(_) => None,
Self::Redis { cache, .. } => cache.backend().namespace(),
Self::Gcs(_) => None,
}
@ -115,7 +123,7 @@ impl NativeResponseCache {
pub fn topology(&self) -> Option<&RedisTopology> {
match self {
Self::Memory(_) | Self::AzureBlob(_) | Self::Gcs(_) => None,
Self::Memory(_) | Self::Disk(_) | Self::AzureBlob(_) | Self::Gcs(_) => None,
Self::Redis { cache, .. } => Some(cache.backend().topology()),
}
}
@ -123,14 +131,14 @@ impl NativeResponseCache {
pub fn capacity(&self) -> Option<usize> {
match self {
Self::Memory(cache) => Some(cache.backend().max_size_in_memory()),
Self::Redis { .. } | Self::AzureBlob(_) | Self::Gcs(_) => None,
Self::Redis { .. } | Self::Disk(_) | Self::AzureBlob(_) | Self::Gcs(_) => None,
}
}
pub fn max_entry_bytes(&self) -> Option<usize> {
match self {
Self::Memory(cache) => cache.backend().max_entry_bytes(),
Self::Redis { .. } | Self::AzureBlob(_) | Self::Gcs(_) => None,
Self::Redis { .. } | Self::Disk(_) | Self::AzureBlob(_) | Self::Gcs(_) => None,
}
}
@ -144,6 +152,13 @@ impl NativeResponseCache {
}
}
pub fn directory(&self) -> Option<&Path> {
match self {
Self::Disk(cache) => Some(cache.backend().directory()),
Self::Memory(_) | Self::Redis { .. } | Self::AzureBlob(_) | Self::Gcs(_) => None,
}
}
pub fn lookup(
&self,
request: &ResponseCacheRequest,
@ -153,6 +168,7 @@ impl NativeResponseCache {
Self::Memory(cache) => cache.lookup(request, now),
Self::Redis { cache, .. } => cache.lookup(request, now),
Self::Gcs(cache) => cache.lookup(request, now),
Self::Disk(cache) => cache.lookup(request, now),
Self::AzureBlob(cache) => cache.lookup(request, now),
}
}
@ -167,6 +183,7 @@ impl NativeResponseCache {
Self::Memory(cache) => cache.store(request, response, now),
Self::Redis { cache, .. } => cache.store(request, response, now),
Self::Gcs(cache) => cache.store(request, response, now),
Self::Disk(cache) => cache.store(request, response, now),
Self::AzureBlob(cache) => cache.store(request, response, now),
}
}
@ -180,6 +197,7 @@ impl NativeResponseCache {
Self::Memory(cache) => cache.lookup_batch(requests, now),
Self::Redis { cache, .. } => cache.lookup_batch(requests, now),
Self::Gcs(cache) => cache.lookup_batch(requests, now),
Self::Disk(cache) => cache.lookup_batch(requests, now),
Self::AzureBlob(cache) => cache.lookup_batch(requests, now),
}
}
@ -193,6 +211,7 @@ impl NativeResponseCache {
Self::Memory(cache) => cache.async_lookup(request, now).await,
Self::Redis { cache, .. } => cache.async_lookup(request, now).await,
Self::Gcs(cache) => cache.async_lookup(request, now).await,
Self::Disk(cache) => cache.async_lookup(request, now).await,
Self::AzureBlob(cache) => cache.async_lookup(request, now).await,
}
}
@ -214,6 +233,7 @@ impl NativeResponseCache {
buffer: Some(buffer),
} => buffer.async_store(cache, request, response, now).await,
Self::Gcs(cache) => cache.async_store(request, response, now).await,
Self::Disk(cache) => cache.async_store(request, response, now).await,
Self::AzureBlob(cache) => cache.async_store(request, response, now).await,
}
}
@ -227,6 +247,7 @@ impl NativeResponseCache {
Self::Memory(cache) => cache.async_lookup_batch(requests, now).await,
Self::Redis { cache, .. } => cache.async_lookup_batch(requests, now).await,
Self::Gcs(cache) => cache.async_lookup_batch(requests, now).await,
Self::Disk(cache) => cache.async_lookup_batch(requests, now).await,
Self::AzureBlob(cache) => cache.async_lookup_batch(requests, now).await,
}
}
@ -240,6 +261,7 @@ impl NativeResponseCache {
Self::Memory(cache) => cache.async_store_batch(entries, now).await,
Self::Redis { cache, .. } => cache.async_store_batch(entries, now).await,
Self::Gcs(cache) => cache.async_store_batch(entries, now).await,
Self::Disk(cache) => cache.async_store_batch(entries, now).await,
Self::AzureBlob(cache) => cache.async_store_batch(entries, now).await,
}
}
@ -254,6 +276,7 @@ impl NativeResponseCache {
cache.async_flush().await
}
Self::Gcs(cache) => cache.async_flush().await,
Self::Disk(cache) => cache.async_flush().await,
Self::AzureBlob(cache) => cache.async_flush().await,
}
}
@ -263,6 +286,7 @@ impl NativeResponseCache {
Self::Memory(cache) => cache.test_connection().await,
Self::Redis { cache, .. } => cache.test_connection().await,
Self::Gcs(cache) => cache.test_connection().await,
Self::Disk(cache) => cache.test_connection().await,
Self::AzureBlob(cache) => cache.test_connection().await,
}
}

4
litellm/a2a_protocol/providers/pydantic_ai_agents/config.py
View file

@ -23,7 +23,7 @@ class PydanticAIProviderConfig(BaseA2AProviderConfig):
params: dict[str, Any],
api_base: str | None = None,
**kwargs: Any,
) -> dict[str, Any]:
) -> dict[str, object]:
"""Handle non-streaming request to Pydantic AI agent."""
if api_base is None:
raise ValueError("api_base is required for PydanticAIProviderConfig")
@ -41,7 +41,7 @@ class PydanticAIProviderConfig(BaseA2AProviderConfig):
params: dict[str, Any],
api_base: str | None = None,
**kwargs,
) -> AsyncIterator[dict[str, Any]]:
) -> AsyncIterator[dict[str, object]]:
"""Handle streaming request with fake streaming."""
if not api_base:
raise ValueError("api_base is required for Pydantic AI agents")

2
litellm/caching/caching.py
View file

@ -81,7 +81,7 @@ class Cache:
s3_aws_access_key_id: str | None = None,
s3_aws_secret_access_key: str | None = None,
s3_aws_session_token: str | None = None,
s3_config: Any | None = None,
s3_config: object | None = None,
s3_path: str | None = None,
gcs_bucket_name: str | None = None,
gcs_path_service_account: str | None = None,

10
litellm/caching/caching_handler.py
View file

@ -74,7 +74,7 @@ class CachingHandlerResponse(BaseModel):
For embeddings there can be a cache hit for some of the inputs in the list and a cache miss for others
"""
cached_result: Any | None = None
cached_result: object | None = None
final_embedding_cached_response: EmbeddingResponse | None = None
embedding_all_elements_cache_hit: bool = False # this is set to True when all elements in the list have a cache hit in the embedding cache, if true return the final_embedding_cached_response no need to make an API call
@ -722,7 +722,7 @@ class LLMCachingHandler:
async def _retrieve_from_cache(
self, call_type: str, kwargs: dict[str, object], args: tuple[object, ...]
) -> Any | None:
) -> object | None:
"""
Internal method to
- get cache key
@ -968,7 +968,7 @@ class LLMCachingHandler:
def _convert_cached_stream_response(
self,
cached_result: Any,
cached_result: dict[str, object],
call_type: str,
logging_obj: LiteLLMLoggingObj,
model: str,
@ -997,7 +997,7 @@ class LLMCachingHandler:
async def async_set_cache(
self,
result: Any,
result: object,
original_function: Callable,
kwargs: dict[str, Any],
args: tuple[object, ...] | None = None,
@ -1065,7 +1065,7 @@ class LLMCachingHandler:
def sync_set_cache(
self,
result: Any,
result: object,
kwargs: dict[str, object],
args: tuple[object, ...] | None = None,
):

4
litellm/caching/redis_cache.py
View file

@ -981,7 +981,7 @@ class RedisCache(BaseCache):
client: object = None,
) -> object:
async def execute() -> object:
executor: Callable[..., Awaitable[Any]] | None = litellm.in_memory_llm_clients_cache.get_cache(
executor: Callable[..., Awaitable[object]] | None = litellm.in_memory_llm_clients_cache.get_cache(
key=script_cache_key
)
if executor is None:
@ -993,7 +993,7 @@ class RedisCache(BaseCache):
return run_script
def _register_script_for_current_loop(self, script: str) -> Callable[..., Awaitable[Any]]:
def _register_script_for_current_loop(self, script: str) -> Callable[..., Awaitable[object]]:
"""
Register the script against the current event loop's Redis client.

6
litellm/completion_extras/litellm_responses_transformation/handler.py
View file

@ -2,7 +2,7 @@
Handler for transforming /chat/completions api requests to litellm.responses requests
"""
from collections.abc import Coroutine
from collections.abc import AsyncIterable, Coroutine, Iterable
from typing import TYPE_CHECKING, Any, Final, Union
from typing_extensions import TypedDict
@ -74,7 +74,7 @@ class ResponsesToCompletionBridgeHandler:
existing.setdefault(key, value)
return response
def _collect_response_from_stream(self, stream_iter: Any) -> "ResponsesAPIResponse":
def _collect_response_from_stream(self, stream_iter: Iterable[object]) -> "ResponsesAPIResponse":
for _ in stream_iter:
pass
@ -89,7 +89,7 @@ class ResponsesToCompletionBridgeHandler:
raise ValueError("Stream completed response is invalid")
return response
async def _collect_response_from_stream_async(self, stream_iter: Any) -> "ResponsesAPIResponse":
async def _collect_response_from_stream_async(self, stream_iter: AsyncIterable[object]) -> "ResponsesAPIResponse":
async for _ in stream_iter:
pass

30
litellm/completion_extras/litellm_responses_transformation/transformation.py
View file

@ -6,7 +6,7 @@ import json
import os
from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping, Sequence
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict, Union, cast, get_args
from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict, TypeVar, Union, cast, get_args
from openai.types.chat import ChatCompletion
from openai.types.responses import Response
@ -52,7 +52,7 @@ from litellm.types.llms.openai import (
from litellm.types.utils import GenericStreamingChunk, ModelResponseStream
if TYPE_CHECKING:
from openai.types.responses import ResponseInputImageParam
from openai.types.responses import ResponseInputImageParam, ResponseOutputItem
from openai.types.responses.response_text_config_param import (
ResponseTextConfigParam as ResponseText,
)
@ -197,6 +197,9 @@ def _as_chat_reasoning_items(
return cast(list[ChatCompletionReasoningItem], list(reasoning_items))
_ToolChoiceT = TypeVar("_ToolChoiceT")
def _map_incomplete_reason_to_finish_reason(incomplete_reason: str | None) -> Literal["length", "content_filter"]:
if incomplete_reason == "content_filter":
return "content_filter"
@ -291,7 +294,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
def __init__(self):
pass
def _normalize_tool_choice_for_responses_api(self, tool_choice: Any) -> Any:
def _normalize_tool_choice_for_responses_api(
self, tool_choice: _ToolChoiceT
) -> _ToolChoiceT | ToolChoiceFunctionParam | ToolChoiceCustomParam | Literal["auto", "none", "required"]:
"""Chat tool_choice nests the name under function/custom; Responses API expects top-level name."""
if not isinstance(tool_choice, dict):
return tool_choice
@ -497,7 +502,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
responses_api_request["max_output_tokens"] = value
elif key == "tools" and value is not None:
responses_api_request["tools"] = self._convert_tools_to_responses_format(
cast(list[dict[str, Any]], value)
cast(list[dict[str, object]], value)
)
elif key == "response_format":
text_format = self._transform_response_format_to_text_format(value)
@ -828,7 +833,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
response_output: Final = response_payload.get("output")
if not isinstance(response_output, list) or len(response_output) == 0:
return None
return cast(list[dict[str, Any]], response_output)
return cast(list[dict[str, object]], response_output)
@classmethod
def _recover_output_items_from_raw_sse(cls, raw_sse: str | None) -> list[dict[str, object]]:
@ -911,10 +916,12 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
output_items = raw_response.output
if len(output_items) == 0:
recovered_output_items: Final = self._recover_output_items_from_logging(logging_obj)
recovered_output_items: Final[list[ResponseOutputItem | dict[str, object]]] = [
*self._recover_output_items_from_logging(logging_obj)
]
if recovered_output_items:
output_items = cast(Any, recovered_output_items)
raw_response.output = cast(Any, recovered_output_items)
output_items = recovered_output_items
raw_response.output = recovered_output_items
verbose_logger.warning(
"Recovered empty Responses API output from raw SSE for model=%s",
model,
@ -1110,7 +1117,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
verbose_logger.debug("Chat provider: Other content type -> %s", result)
return result
def _convert_tools_to_responses_format(self, tools: list[dict[str, Any]]) -> list["ALL_RESPONSES_API_TOOL_PARAMS"]:
def _convert_tools_to_responses_format(
self, tools: list[dict[str, object]]
) -> list["ALL_RESPONSES_API_TOOL_PARAMS"]:
"""Convert chat completion tools to responses API tools format"""
responses_tools: Final[list[ALL_RESPONSES_API_TOOL_PARAMS]] = []
for tool in tools:
@ -1126,12 +1135,11 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
description=function_tool.get("description"),
)
)
elif tool.get("type") == "custom" and isinstance(tool.get("custom"), dict):
elif tool.get("type") == "custom" and isinstance(custom_payload := tool.get("custom"), dict):
from litellm.litellm_core_utils.prompt_templates.common_utils import (
convert_custom_tool_format_to_responses_shape,
)
custom_payload = tool["custom"]
flat_custom = CustomToolParam(type="custom", name=custom_payload.get("name", ""))
if custom_payload.get("description") is not None:
flat_custom["description"] = custom_payload["description"]

3
litellm/constants.py
View file

@ -2121,3 +2121,6 @@ BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY: Final = "batch_enqueued_token_limit"
# Shared read-only empty mapping, for defaulting optional Mapping parameters without
# constructing a fresh mutable dict at each call site.
EMPTY_MAPPING: Final = MappingProxyType({})
# API endpoint for breached password k-anonymity search
HIBP_RANGE_API_BASE: Final = "https://api.pwnedpasswords.com/range"

6
litellm/cost_calculator.py
View file

@ -353,7 +353,7 @@ def cost_per_token(
data_residency: str | None = None, # for OpenAI regional-processing uplift (e.g. "eu", "us")
### VERTEX LOCATION ###
vertex_location: str | None = None, # for Vertex AI regional-endpoint uplift (e.g. "us-east5", "global")
response: Any | None = None,
response: object | None = None,
### REQUEST MODEL ###
request_model: str | None = None, # original request model for router detection
custom_model_info: OCRPricing | None = None,
@ -609,7 +609,7 @@ def cost_per_token(
model=model,
custom_llm_provider=custom_llm_provider,
number_of_queries=number_of_queries or 1,
optional_params=(response._hidden_params if response and hasattr(response, "_hidden_params") else None),
optional_params=(getattr(response, "_hidden_params", None) if response else None),
)
elif custom_llm_provider == "vertex_ai":
cost_router: Final = google_cost_router(
@ -999,7 +999,7 @@ def _is_known_usage_objects(usage_obj):
)
def _infer_call_type(call_type: CallTypesLiteral | None, completion_response: Any) -> CallTypesLiteral | None:
def _infer_call_type(call_type: CallTypesLiteral | None, completion_response: object) -> CallTypesLiteral | None:
if call_type is not None:
return call_type

7
litellm/integrations/argilla.py
View file

@ -7,7 +7,8 @@ import json
import os
import random
import types
from typing import Any, Final
from collections.abc import Mapping
from typing import Final
import httpx
from pydantic import BaseModel
@ -69,7 +70,7 @@ class ArgillaLogger(CustomBatchLogger):
self.flush_lock = asyncio.Lock()
super().__init__(**kwargs, flush_lock=self.flush_lock)
def validate_argilla_transformation_object(self, argilla_transformation_object: dict[str, Any]):
def validate_argilla_transformation_object(self, argilla_transformation_object: Mapping[str, object]):
if not isinstance(argilla_transformation_object, dict):
raise Exception("'argilla_transformation_object' must be a dictionary, to log your payload to Argilla.")
@ -115,7 +116,7 @@ class ArgillaLogger(CustomBatchLogger):
ARGILLA_DATASET_NAME=_credentials_dataset_name,
)
def get_chat_messages(self, payload: StandardLoggingPayload) -> list[dict[str, Any]]:
def get_chat_messages(self, payload: StandardLoggingPayload) -> list[dict[str, object]]:
payload_messages: Final = payload.get("messages", None)
if payload_messages is None:

8
litellm/integrations/braintrust_logging.py
View file

@ -139,13 +139,13 @@ class BraintrustLogger(CustomLogger):
):
output = None
elif response_obj is not None and isinstance(response_obj, litellm.ModelResponse):
output = response_obj["choices"][0]["message"].json()
output = response_obj.choices[0].message.json()
choices = response_obj["choices"]
elif response_obj is not None and isinstance(response_obj, litellm.TextCompletionResponse):
output = response_obj.choices[0].text
choices = response_obj.choices
elif response_obj is not None and isinstance(response_obj, litellm.ImageResponse):
output = response_obj["data"]
output = response_obj.data
litellm_params: Final = kwargs.get("litellm_params", {}) or {}
dynamic_metadata: Final = litellm_params.get("metadata", {}) or {}
@ -264,13 +264,13 @@ class BraintrustLogger(CustomLogger):
):
output = None
elif response_obj is not None and isinstance(response_obj, litellm.ModelResponse):
output = response_obj["choices"][0]["message"].json()
output = response_obj.choices[0].message.json()
choices = response_obj["choices"]
elif response_obj is not None and isinstance(response_obj, litellm.TextCompletionResponse):
output = response_obj.choices[0].text
choices = response_obj.choices
elif response_obj is not None and isinstance(response_obj, litellm.ImageResponse):
output = response_obj["data"]
output = response_obj.data
litellm_params: Final = kwargs.get("litellm_params", {})
dynamic_metadata: Final = litellm_params.get("metadata", {}) or {}

2
litellm/integrations/custom_guardrail.py
View file

@ -150,7 +150,7 @@ class CustomGuardrail(CustomLogger):
def __init_subclass__(cls, **kwargs: object) -> None: # kwargs-ok: forwarded to cooperative __init_subclass__ hooks
super().__init_subclass__(**kwargs)
own_apply_guardrail: Final = cls.__dict__.get("apply_guardrail")
own_apply_guardrail: Final[object] = cls.__dict__.get("apply_guardrail")
if own_apply_guardrail is None or LOGS_GUARDRAIL_INFORMATION_MARKER in vars(own_apply_guardrail):
return
cls.apply_guardrail = log_guardrail_information(own_apply_guardrail)

42
litellm/integrations/datadog/datadog_llm_obs.py
View file

@ -54,7 +54,7 @@ from litellm.types.utils import (
StandardLoggingPayloadErrorInformation,
)
_EMPTY_MAPPING: Final[Mapping[str, Any]] = MappingProxyType({})
_EMPTY_MAPPING: Final[Mapping[str, object]] = MappingProxyType({})
_EMPTY_MESSAGE: Final[Message] = {"role": "", "content": ""}
_MAX_PARSED_TOOL_ARGUMENT_CHARS: Final = 256 * 1024
_SAFE_REDACTED_MESSAGE_ROLES: Final = frozenset(
@ -154,7 +154,7 @@ def _guardrail_information_without_prompt_carriers(
return tuple(_guardrail_entry_without_prompt_carriers(entry) for entry in _guardrail_entries(guardrail_information))
def _metadata_without_prompt_carriers(standard_logging_metadata: Mapping[str, Any]) -> Mapping[str, Any]:
def _metadata_without_prompt_carriers(standard_logging_metadata: Mapping[str, object]) -> Mapping[str, object]:
"""The metadata minus the records that quote prompts, tool arguments, tool results, or retrieved text."""
return MappingProxyType(
{
@ -237,7 +237,7 @@ def _declared_cost_tags(span_tags: Sequence[str]) -> tuple[str, ...]:
return tuple(dimension for dimension in _COST_DIMENSIONS if dimension in present)
def _reasoning_output_tokens(usage_object: Mapping[str, Any] | None) -> float:
def _reasoning_output_tokens(usage_object: Mapping[str, object] | None) -> float:
"""The provider's reasoning-token count, from either the chat or the responses spelling."""
if usage_object is None:
return 0.0
@ -254,20 +254,24 @@ def _reasoning_output_tokens(usage_object: Mapping[str, Any] | None) -> float:
)
def _mapping_field(source: Mapping[str, Any], key: str) -> Mapping[str, Any]:
def _mapping_field(source: Mapping[str, object], key: str) -> Mapping[str, object]:
"""The value at `key` when it is a mapping, else an empty one."""
value: Final = source.get(key)
return value if isinstance(value, dict) else _EMPTY_MAPPING
def _content_blocks(message: Mapping[str, Any]) -> tuple[Mapping[str, Any], ...]:
def _text_field(source: Mapping[str, object], key: str, default: str = "") -> str:
return _safe_identifier(source.get(key, default))
def _content_blocks(message: Mapping[str, object]) -> tuple[Mapping[str, object], ...]:
content: Final = message.get("content")
if not isinstance(content, list):
return ()
return tuple(block for block in content if isinstance(block, dict))
def _to_dd_arguments(raw_arguments: object) -> dict[str, Any] | str:
def _to_dd_arguments(raw_arguments: object) -> dict[str, object] | str:
"""
Arguments as the object LLM Obs types them as, or the raw string when they are not one.
@ -282,7 +286,7 @@ def _to_dd_arguments(raw_arguments: object) -> dict[str, Any] | str:
return parsed if isinstance(parsed, dict) else raw_arguments
def _to_dd_tool_calls(message: Mapping[str, Any]) -> tuple[ToolCall, ...]:
def _to_dd_tool_calls(message: Mapping[str, object]) -> tuple[ToolCall, ...]:
"""
The tool calls a message carries, in LLM Obs' ToolCall schema, from either dialect.
@ -293,10 +297,10 @@ def _to_dd_tool_calls(message: Mapping[str, Any]) -> tuple[ToolCall, ...]:
raw_tool_calls: Final = message.get("tool_calls")
openai_calls: Final = tuple(
ToolCall(
name=function.get("name", ""),
name=_text_field(function, "name"),
arguments=_to_dd_arguments(function.get("arguments", "")),
tool_id=tool_call.get("id", ""),
type=tool_call.get("type", "function"),
tool_id=_text_field(tool_call, "id"),
type=_text_field(tool_call, "type", "function"),
)
for tool_call in (raw_tool_calls if isinstance(raw_tool_calls, list) else ())
if isinstance(tool_call, dict)
@ -304,9 +308,9 @@ def _to_dd_tool_calls(message: Mapping[str, Any]) -> tuple[ToolCall, ...]:
)
anthropic_calls: Final = tuple(
ToolCall(
name=block.get("name", ""),
name=_text_field(block, "name"),
arguments=_to_dd_arguments(block.get("input") or {}),
tool_id=block.get("id", ""),
tool_id=_text_field(block, "id"),
type="tool_use",
)
for block in _content_blocks(message)
@ -315,7 +319,7 @@ def _to_dd_tool_calls(message: Mapping[str, Any]) -> tuple[ToolCall, ...]:
return openai_calls + anthropic_calls
def _to_dd_tool_results(message: Mapping[str, Any], tool_call_names: Mapping[str, str]) -> tuple[ToolResult, ...]:
def _to_dd_tool_results(message: Mapping[str, object], tool_call_names: Mapping[str, str]) -> tuple[ToolResult, ...]:
"""
The tool results a message carries, linked back to the call each answers.
@ -400,14 +404,14 @@ def _to_dd_messages(messages: object) -> tuple[Message, ...]:
return tuple(_to_dd_message(message, tool_call_names) for message in messages)
def _to_dd_tool_definition(entry: Mapping[str, Any]) -> ToolDefinition | None:
def _to_dd_tool_definition(entry: Mapping[str, object]) -> ToolDefinition | None:
function: Final = entry.get("function")
declared: Final[Mapping[str, Any]] = function if isinstance(function, dict) else entry
name: Final = declared.get("name")
declared: Final[Mapping[str, object]] = function if isinstance(function, dict) else entry
name: Final = _text_field(declared, "name")
if not name:
return None
schema: Final = declared.get("parameters") or declared.get("input_schema")
description: Final = declared.get("description", "")
description: Final = _text_field(declared, "description")
if not isinstance(schema, dict):
return ToolDefinition(name=name, description=description)
return ToolDefinition(name=name, description=description, schema=schema)
@ -683,7 +687,7 @@ class DataDogLLMObsLogger(CustomBatchLogger):
if callable(current_span_fn):
current_span: Final = current_span_fn()
if current_span is not None:
trace_id: Final = getattr(current_span, "trace_id", None)
trace_id: Final[object] = getattr(current_span, "trace_id", None)
if trace_id is not None:
return str(trace_id)
except Exception:
@ -716,7 +720,7 @@ class DataDogLLMObsLogger(CustomBatchLogger):
def redacts_messages_itself(self) -> bool:
return True
def _payload_logging_is_off(self, kwargs: Mapping[str, Any]) -> bool:
def _payload_logging_is_off(self, kwargs: Mapping[str, object]) -> bool:
return (
bool(self.turn_off_message_logging)
or self.message_logging is not True

15
litellm/integrations/dynamodb.py
View file

@ -3,12 +3,21 @@
import os
import traceback
from typing import Any, Final
from collections.abc import Mapping
from typing import Final, Protocol
import litellm
from litellm._uuid import uuid
class _DynamoTable(Protocol):
def put_item(self, *, Item: Mapping[str, object]) -> object: ...
class _DynamoResource(Protocol):
def Table(self, name: str) -> _DynamoTable: ...
class DyanmoDBLogger:
# Class variables or attributes
@ -16,7 +25,7 @@ class DyanmoDBLogger:
# Instance variables
import boto3
self.dynamodb: Any = boto3.resource("dynamodb", region_name=os.environ["AWS_REGION_NAME"])
self.dynamodb: Final[_DynamoResource] = boto3.resource("dynamodb", region_name=os.environ["AWS_REGION_NAME"])
if litellm.dynamodb_table_name is None:
raise ValueError(
"LiteLLM Error, trying to use DynamoDB but not table name passed. Create a table and set `litellm.dynamodb_table_name=<your-table>`"
@ -41,7 +50,7 @@ class DyanmoDBLogger:
id: Final = response_obj.get("id", str(uuid.uuid4()))
# Build the initial payload
payload: Final = {
payload: Final[dict[str, object]] = {
"id": id,
"call_type": call_type,
"startTime": start_time,

6
litellm/integrations/focus/database.py
View file

@ -3,7 +3,7 @@
from __future__ import annotations
from datetime import datetime
from typing import Any, Final
from typing import Final
import polars as pl
@ -32,7 +32,7 @@ class FocusLiteLLMDatabase:
client: Final = self._ensure_prisma_client()
where_clauses: Final[list[str]] = []
query_params: Final[list[Any]] = []
query_params: Final[list[datetime | int]] = []
placeholder_index = 1
if start_time_utc:
where_clauses.append(f"dus.updated_at >= ${placeholder_index}::timestamptz")
@ -112,7 +112,7 @@ class FocusLiteLLMDatabase:
except Exception as exc:
raise RuntimeError(f"Error retrieving usage data: {exc}") from exc
async def get_table_info(self) -> dict[str, Any]:
async def get_table_info(self) -> dict[str, object]:
"""Return metadata about the spend table for diagnostics."""
client: Final = self._ensure_prisma_client()

5
litellm/integrations/focus/destinations/vantage_destination.py
View file

@ -4,7 +4,8 @@ from __future__ import annotations
import csv
import io
from typing import Any, Final
from collections.abc import Mapping
from typing import Final
import httpx # noqa: F401 - used at runtime (AsyncClient, HTTPStatusError)
@ -94,7 +95,7 @@ class FocusVantageDestination(FocusDestination):
self,
*,
prefix: str,
config: dict[str, Any] | None = None,
config: Mapping[str, object] | None = None,
) -> None:
config = config or {}
api_key: Final = config.get("api_key")

11
litellm/integrations/galileo.py
View file

@ -396,12 +396,13 @@ class GalileoObserve(CustomLogger):
)
@staticmethod
def _log_v2_payload_validation(payload: dict[str, Any]) -> None:
def _log_v2_payload_validation(payload: dict[str, object]) -> None:
missing_fields: Final[list[str]] = []
traces: Final[Sequence[object]] = payload.get("traces", [])
if not traces:
traces_value: Final = payload.get("traces", [])
if not traces_value:
missing_fields.append("traces")
traces: Final[Sequence[object]] = traces_value if isinstance(traces_value, list) else []
for trace_index, trace in enumerate(traces):
if not isinstance(trace, dict):
continue
@ -425,8 +426,8 @@ class GalileoObserve(CustomLogger):
missing_fields,
)
def _log_flush_payload(self, url: str, payload: dict[str, Any]) -> None:
traces: Final[Sequence[object]] = payload.get("traces", [])
def _log_flush_payload(self, url: str, payload: dict[str, object]) -> None:
traces: Final = payload.get("traces")
verbose_logger.debug(
"Galileo Logger flush URL: %s trace_count=%s",
url,

8
litellm/integrations/langfuse/langfuse.py
View file

@ -4,7 +4,7 @@ import inspect
import os
import re
import traceback
from collections.abc import Callable, Iterable, Mapping
from collections.abc import Callable, Iterable, Mapping, Sequence
from datetime import datetime
from functools import lru_cache
from types import MappingProxyType
@ -432,7 +432,7 @@ class LangFuseLogger:
prompt: dict,
level: str,
status_message: str | None,
) -> tuple[dict | None, str | dict | list | None]:
) -> tuple[dict | None, str | dict | Sequence[object] | None]:
"""
Get the input and output content for Langfuse logging
@ -448,7 +448,7 @@ class LangFuseLogger:
output: The output content for Langfuse logging
"""
input = None
output: str | dict | list[Any] | None = None
output: str | dict | Sequence[object] | None = None
if level == "ERROR" and status_message is not None and isinstance(status_message, str):
input = prompt
output = status_message
@ -508,7 +508,7 @@ class LangFuseLogger:
user_id: str | None,
metadata: dict[str, object],
litellm_params: dict,
output: str | dict | list | None,
output: str | dict | Sequence[object] | None,
start_time: datetime | None,
end_time: datetime | None,
kwargs: dict,

3
litellm/integrations/langfuse/langfuse_otel_attributes.py
View file

@ -5,6 +5,7 @@ Relevant Issue: https://github.com/BerriAI/litellm/issues/13764
"""
import json
from collections.abc import Mapping
from typing import TYPE_CHECKING, Any, Final
from pydantic import BaseModel
@ -40,7 +41,7 @@ def get_output_content_by_type(
| HttpxBinaryResponseContent
| ResponsesAPIResponse
| list,
kwargs: dict[str, Any] | None = None,
kwargs: Mapping[str, object] | None = None,
) -> str:
"""
Extract output content from response objects based on their type.

8
litellm/integrations/langsmith.py
View file

@ -77,9 +77,9 @@ class LangsmithLogger(CustomBatchLogger):
if _batch_size:
self.batch_size = int(_batch_size)
self.log_queue: list[LangsmithQueueObject] = []
self._flush_task: asyncio.Task[Any] | None = self._start_periodic_flush_task()
self._flush_task: asyncio.Task[None] | None = self._start_periodic_flush_task()
def _start_periodic_flush_task(self) -> asyncio.Task[Any] | None:
def _start_periodic_flush_task(self) -> asyncio.Task[None] | None:
"""Start the periodic flush task only when an event loop is already running."""
try:
loop: Final = asyncio.get_running_loop()
@ -154,9 +154,9 @@ class LangsmithLogger(CustomBatchLogger):
return self._redact_metadata(extra_metadata)
def _build_outputs_with_usage(self, payload: StandardLoggingPayload) -> dict[str, Any]:
def _build_outputs_with_usage(self, payload: StandardLoggingPayload) -> dict[str, object]:
response: Final = payload["response"]
outputs: dict[str, Any]
outputs: dict[str, object]
if isinstance(response, dict):
outputs = {**response}
else:

14
litellm/integrations/otel/model/metadata.py
View file

@ -36,7 +36,7 @@ model. They coincide on the SDK path, which is correct.
from __future__ import annotations
from collections.abc import Iterator, Mapping
from collections.abc import Callable, Iterator, Mapping
from dataclasses import dataclass, field
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, cast
@ -61,7 +61,7 @@ class RequestIdentity:
# The team's free-form metadata, carried raw (empty/missing -> None) and
# filtered to an operator allowlist only at Baggage-promotion time, so an
# unconfigured deployment never promotes any of it.
team_metadata: Mapping[str, Any] | None = None
team_metadata: Mapping[str, object] | None = None
key_hash: str | None = None
end_user: str | None = None
# The model litellm dispatched to the provider. Only known once the call
@ -111,7 +111,7 @@ class RequestIdentity:
snapshot) is flattened to dotted keys so ``requester_metadata.<key>``
resolves too.
"""
get: Final = lambda name: getattr(auth, name, None) # noqa: E731
get: Final[Callable[[str], object]] = lambda name: getattr(auth, name, None) # noqa: E731
auth_meta: Final = tuple(
(meta_key, str(value))
for meta_key, attr in (
@ -228,7 +228,7 @@ class LLMCallEvent:
trace: TraceControls
@classmethod
def from_dict(cls, kwargs: Mapping[str, Any]) -> LLMCallEvent:
def from_dict(cls, kwargs: Mapping[str, object]) -> LLMCallEvent:
raw_payload: Final = kwargs.get("standard_logging_object")
payload: Final = cast("StandardLoggingPayload", raw_payload) if raw_payload else None
operation: Final = resolve_operation(as_str(kwargs.get("call_type")))
@ -251,7 +251,7 @@ def time_to_first_chunk_seconds(kwargs: Mapping[str, Any]) -> float | None:
to the first streamed chunk (``completion_start_time``); ``None`` for
non-streaming calls, where ``completion_start_time`` is backfilled with the
end time and would not measure first-chunk latency."""
optional_params: Final = cast(Mapping[str, Any], kwargs.get("optional_params") or {})
optional_params: Final = cast(Mapping[str, object], kwargs.get("optional_params") or {})
if not optional_params.get("stream"):
return None
api_call_start: Final = to_seconds(kwargs.get("api_call_start_time"))
@ -312,7 +312,7 @@ def _metadata_dicts(
)
def _call_id(payload: StandardLoggingPayload | None, kwargs: Mapping[str, Any]) -> str | None:
def _call_id(payload: StandardLoggingPayload | None, kwargs: Mapping[str, object]) -> str | None:
"""The call id from the payload (when closed) or the bare kwargs (at pre_call)."""
if payload is not None:
call_id: Final = as_str(payload.get("litellm_call_id")) or as_str(payload.get("id"))
@ -385,7 +385,7 @@ def _model_info_id(model_info: object) -> str | None:
return None
def _team_metadata_dict(value: object) -> Mapping[str, Any] | None:
def _team_metadata_dict(value: object) -> Mapping[str, object] | None:
"""The team's free-form metadata as a raw mapping, or ``None`` when missing
or empty.

17
litellm/integrations/otel/mount.py
View file

@ -12,11 +12,14 @@ when the feature gate is off.
"""
import os
from typing import Any, Final
from typing import TYPE_CHECKING, Final, Protocol
from litellm._logging import verbose_logger
from litellm.integrations.otel.model.config import is_otel_v2_enabled
if TYPE_CHECKING:
from fastapi import FastAPI
# Routes excluded from server-span tracing by default: high-frequency pollers and
# static UI/docs assets, none of which are LLM traffic. Entries are substring-matched
# against the request path (unanchored, so they survive a ``server_root_path`` prefix
@ -65,7 +68,15 @@ PASSTHROUGH_PREFIXES: Final = frozenset(
)
def _passthrough_span_name_hook(span: Any, scope: dict) -> None:
class _RenameableSpan(Protocol):
def is_recording(self) -> bool: ...
def update_name(self, name: str) -> None: ...
def set_attribute(self, key: str, value: str) -> None: ...
def _passthrough_span_name_hook(span: "_RenameableSpan | None", scope: dict) -> None:
"""FastAPI ``server_request_hook``: give passthrough server spans a useful name.
The instrumentation matches the route at span creation, so both the span name
@ -88,7 +99,7 @@ def _passthrough_span_name_hook(span: Any, scope: dict) -> None:
pass
def instrument_fastapi_app(app: Any) -> None:
def instrument_fastapi_app(app: "FastAPI") -> None:
"""Attach OTel server-span instrumentation to the proxy FastAPI app.
Safe no-op when the V2 gate is off or ``opentelemetry-instrumentation-fastapi``

6
litellm/litellm_core_utils/coroutine_checker.py
View file

@ -16,7 +16,7 @@ class CoroutineChecker:
"""
def __init__(self):
self._cache = WeakKeyDictionary()
self._cache: WeakKeyDictionary[object, bool] = WeakKeyDictionary()
self._max_size = COROUTINE_CHECKER_MAX_SIZE_IN_MEMORY
def is_async_callable(self, callback: Any) -> bool:
@ -33,10 +33,10 @@ class CoroutineChecker:
pass
# Determine target - optimized path for common cases
target = callback
target: object = callback
if not inspect.isfunction(target) and not inspect.ismethod(target):
try:
call_attr: Final = getattr(target, "__call__", None) # noqa: B004 # value unwrap so iscoroutinefunction sees through functors
call_attr: Final[object] = getattr(target, "__call__", None) # noqa: B004 # value unwrap so iscoroutinefunction sees through functors
if call_attr is not None:
target = call_attr
except Exception:

6
litellm/litellm_core_utils/exception_mapping_utils.py
View file

@ -4,7 +4,7 @@ import re
import traceback
from collections.abc import Mapping
from types import MappingProxyType
from typing import Any, Final, Protocol, cast
from typing import Final, Protocol, cast
import httpx
@ -194,7 +194,7 @@ def _get_response_headers(original_exception: Exception) -> httpx.Headers | None
_response_headers: httpx.Headers | None = None
try:
_response_headers = getattr(original_exception, "headers", None)
error_response: Final = getattr(original_exception, "response", None)
error_response: Final[object] = getattr(original_exception, "response", None)
if not _response_headers and error_response:
_response_headers = getattr(error_response, "headers", None)
if not _response_headers:
@ -211,7 +211,7 @@ def _accepted_init_kwargs(exception_class: type[Exception], candidates: Mapping[
def extract_and_raise_litellm_exception(
response: Any | None,
response: object | None,
error_str: str,
model: str,
custom_llm_provider: str,

15
litellm/litellm_core_utils/llm_cost_calc/utils.py
View file

@ -6,7 +6,7 @@ from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime, timezone, tzinfo
from types import MappingProxyType
from typing import Any, Final, Literal, TypedDict, cast
from typing import Final, Literal, TypedDict, cast
from zoneinfo import ZoneInfo, ZoneInfoNotFoundError
from typing_extensions import ReadOnly
@ -100,7 +100,7 @@ def _requested_image_size(optional_params: Mapping[str, object] | None) -> str |
return value if value is not None and _IMAGE_SIZE_PATTERN.fullmatch(value) else None
def get_web_search_requests(server_tool_use: Any) -> int | None:
def get_web_search_requests(server_tool_use: object) -> int | None:
"""
Tolerantly read ``web_search_requests`` from a ``server_tool_use`` value
that may be ``None``, a ``dict``, a ``ServerToolUse`` pydantic instance,
@ -1653,7 +1653,7 @@ def calculate_image_response_cost_from_usage(
if prompt_tokens == 0 and completion_tokens == 0 and total_tokens == 0:
return None
input_tokens_details: Final = getattr(usage, "input_tokens_details", None)
input_tokens_details: Final[object] = getattr(usage, "input_tokens_details", None)
prompt_tokens_details: PromptTokensDetailsWrapper | None = None
if input_tokens_details is not None:
# input_tokens_details may be a dict (e.g. OpenAI image edit responses)
@ -1666,9 +1666,12 @@ def calculate_image_response_cost_from_usage(
cached_tokens=0,
)
output_tokens_details = getattr(usage, "completion_tokens_details", None)
if output_tokens_details is None:
output_tokens_details = getattr(usage, "output_tokens_details", None)
completion_tokens_details_attr: Final[object] = getattr(usage, "completion_tokens_details", None)
output_tokens_details: Final[object] = (
getattr(usage, "output_tokens_details", None)
if completion_tokens_details_attr is None
else completion_tokens_details_attr
)
if output_tokens_details is None:
completion_tokens_details = CompletionTokensDetailsWrapper(

4
litellm/litellm_core_utils/llm_response_utils/response_metadata.py
View file

@ -1,7 +1,7 @@
import datetime
from collections.abc import Mapping
from functools import reduce
from typing import Any, Final
from typing import Final
import httpx
@ -106,7 +106,7 @@ class ResponseMetadata:
Handles setting and managing `_hidden_params`, `response_time_ms`, and `litellm_overhead_time_ms` for LiteLLM responses
"""
def __init__(self, result: Any):
def __init__(self, result: object):
self.result = result
self._hidden_params: HiddenParams | dict = getattr(result, "_hidden_params", {}) or {}

39
litellm/litellm_core_utils/prompt_templates/common_utils.py
View file

@ -13,14 +13,6 @@ from pathlib import Path
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, TypeVar, cast
from openai.types.chat.chat_completion_custom_tool_param import (
CustomFormatGrammar,
CustomFormatGrammarGrammar,
)
from openai.types.shared_params.custom_tool_input_format import (
Grammar as ResponsesGrammarFormat,
)
import litellm
from litellm import verbose_logger
from litellm.router_utils.batch_utils import InMemoryFile
@ -59,7 +51,7 @@ if TYPE_CHECKING:
def handle_any_messages_to_chat_completion_str_messages_conversion(
messages: Any,
messages: object,
) -> list[dict[str, str]]:
"""
Handles any messages to chat completion str messages conversion
@ -804,7 +796,7 @@ def extract_file_metadata(file_data: FileTypes) -> tuple[str | None, str | None]
"""
filename: str | None = None
content_type: str | None = None
file_content: Any = None
file_content: object = None
if isinstance(file_data, tuple):
if len(file_data) == 2:
@ -1002,7 +994,7 @@ def unpack_defs(
# Use iterative approach with queue to avoid recursion
# Each item in queue is (node, parent_container, key/index, active_defs, ref_chain)
queue: Final[deque[tuple[Any, dict | list | None, str | int | None, dict, set]]] = deque(
queue: Final[deque[tuple[object, dict | list | None, str | int | None, dict, set]]] = deque(
[(schema, None, None, root_defs, set())]
)
inlined_bytes = 0
@ -1624,7 +1616,10 @@ def is_function_call(optional_params: dict) -> bool:
return False
def convert_custom_tool_format_to_chat_shape(format_obj: Mapping[str, Any]) -> Mapping[str, Any]:
_CUSTOM_GRAMMAR_FIELDS: Final = ("definition", "syntax")
def convert_custom_tool_format_to_chat_shape(format_obj: Mapping[str, object]) -> Mapping[str, object]:
"""
Responses API grammar formats are flat ({"type": "grammar", "definition", "syntax"});
Chat Completions wraps the same fields in a "grammar" object. Text formats are
@ -1632,15 +1627,11 @@ def convert_custom_tool_format_to_chat_shape(format_obj: Mapping[str, Any]) -> M
"""
if format_obj.get("type") != "grammar" or "grammar" in format_obj:
return format_obj
grammar: Final = CustomFormatGrammarGrammar()
if "definition" in format_obj:
grammar["definition"] = format_obj["definition"]
if "syntax" in format_obj:
grammar["syntax"] = format_obj["syntax"]
return CustomFormatGrammar(type="grammar", grammar=grammar)
grammar: Final[Mapping[str, object]] = {key: format_obj[key] for key in _CUSTOM_GRAMMAR_FIELDS if key in format_obj}
return {"type": "grammar", "grammar": grammar}
def convert_custom_tool_format_to_responses_shape(format_obj: Mapping[str, Any]) -> Mapping[str, Any]:
def convert_custom_tool_format_to_responses_shape(format_obj: Mapping[str, object]) -> Mapping[str, object]:
"""
Inverse of convert_custom_tool_format_to_chat_shape: unwrap the Chat Completions
"grammar" object into the flat Responses API grammar shape.
@ -1648,12 +1639,10 @@ def convert_custom_tool_format_to_responses_shape(format_obj: Mapping[str, Any])
grammar: Final = format_obj.get("grammar")
if format_obj.get("type") != "grammar" or not isinstance(grammar, dict):
return format_obj
flat: Final = ResponsesGrammarFormat(type="grammar")
if "definition" in grammar:
flat["definition"] = grammar["definition"]
if "syntax" in grammar:
flat["syntax"] = grammar["syntax"]
return flat
return {
"type": "grammar",
**{key: grammar[key] for key in _CUSTOM_GRAMMAR_FIELDS if key in grammar},
}
def get_file_ids_from_messages(messages: list[AllMessageValues]) -> list[str]:

26
litellm/litellm_core_utils/prompt_templates/huggingface_template_handler.py
View file

@ -1,6 +1,8 @@
import json
from datetime import datetime
from typing import Any, Final
from typing import Any, Final, Literal
from typing_extensions import NotRequired, ReadOnly, TypedDict
from litellm.llms.custom_httpx.http_handler import (
_get_httpx_client,
@ -9,6 +11,20 @@ from litellm.llms.custom_httpx.http_handler import (
from litellm.types.llms.custom_http import httpxSpecialProvider
class _TokenizerConfigResult(TypedDict):
"""Outcome of a tokenizer_config.json fetch, carrying the parsed document when the fetch succeeded."""
status: ReadOnly[Literal["success", "failure"]]
tokenizer: NotRequired[ReadOnly[object]]
class _ChatTemplateFileResult(TypedDict):
"""Outcome of a chat template file fetch, carrying the template body when the fetch succeeded."""
status: ReadOnly[Literal["success", "failure"]]
chat_template: NotRequired[ReadOnly[str]]
def strftime_now(fmt: str) -> str:
"""
Custom function for templates that need current date/time formatting (e.g., gpt-oss)
@ -22,7 +38,7 @@ def strftime_now(fmt: str) -> str:
return datetime.now().strftime(fmt)
def _get_tokenizer_config(hf_model_name: str) -> dict[str, Any]:
def _get_tokenizer_config(hf_model_name: str) -> _TokenizerConfigResult:
"""
Fetch tokenizer_config.json from HuggingFace (sync)
@ -45,7 +61,7 @@ def _get_tokenizer_config(hf_model_name: str) -> dict[str, Any]:
return {"status": "failure"}
async def _aget_tokenizer_config(hf_model_name: str) -> dict[str, Any]:
async def _aget_tokenizer_config(hf_model_name: str) -> _TokenizerConfigResult:
"""
Fetch tokenizer_config.json from HuggingFace (async)
@ -70,7 +86,7 @@ async def _aget_tokenizer_config(hf_model_name: str) -> dict[str, Any]:
return {"status": "failure"}
def _get_chat_template_file(hf_model_name: str) -> dict[str, Any]:
def _get_chat_template_file(hf_model_name: str) -> _ChatTemplateFileResult:
"""
Fetch chat template from separate .jinja file (sync)
@ -98,7 +114,7 @@ def _get_chat_template_file(hf_model_name: str) -> dict[str, Any]:
return {"status": "failure"}
async def _aget_chat_template_file(hf_model_name: str) -> dict[str, Any]:
async def _aget_chat_template_file(hf_model_name: str) -> _ChatTemplateFileResult:
"""
Fetch chat template from separate .jinja file (async)

10
litellm/litellm_core_utils/sensitive_data_masker.py
View file

@ -93,13 +93,13 @@ class SensitiveDataMasker:
def _mask_sequence(
self,
values: list[Any],
values: Sequence[object],
depth: int,
max_depth: int,
excluded_keys: set[str] | None,
key_is_sensitive: bool,
) -> list[Any]:
masked_items: Final[list[Any]] = []
) -> Sequence[object]:
masked_items: Final[list[object]] = []
if depth >= max_depth:
return values
@ -222,7 +222,7 @@ class _PayloadWalker:
return [self.walk(item, key_is_sensitive, depth + 1) for item in node]
def mask_sensitive_keys(data: dict[str, Any], sensitive_fields: set[str]) -> dict[str, Any]:
def mask_sensitive_keys(data: Mapping[str, object], sensitive_fields: set[str]) -> dict[str, object]:
"""Return a new dict with values masked for keys listed in ``sensitive_fields``.
Unlike :meth:`SensitiveDataMasker.mask_dict`, this does exact key-name
@ -234,7 +234,7 @@ def mask_sensitive_keys(data: dict[str, Any], sensitive_fields: set[str]) -> dic
range and are replaced with a fixed-length all-mask string, so a short
credential is never returned verbatim.
"""
masked: Final[dict[str, Any]] = {}
masked: Final[dict[str, object]] = {}
mask_char: Final = _default_masker.mask_char
min_visible: Final = _default_masker.visible_prefix + _default_masker.visible_suffix
for key, value in data.items():

46
litellm/litellm_core_utils/streaming_chunk_builder_utils.py
View file

@ -839,15 +839,17 @@ class ChunkProcessor:
UsagePerChunk,
)
# # Update usage information if needed
prompt_tokens = 0
completion_tokens = 0
# None means no usage chunk reported the count, which is the only case
# calculate_usage() estimates with the tokenizer. An explicit provider 0
# is a reported count and stays 0; a reported count is never replaced by
# a later chunk's 0 (Ollama sends 0/0 on every chunk before the done one).
prompt_tokens: int | None = None
completion_tokens: int | None = None
# Anthropic's `message_start` SSE event carries usage.output_tokens=1 as a
# cursor/placeholder; the real value only arrives in `message_delta`.
# If a stream is cancelled before `message_delta` lands, the last-wins
# accumulator below leaves completion_tokens stuck at 1 — which then
# bypasses the `completion_tokens or token_counter(...)` fallback in
# calculate_usage() because 1 is truthy. Count the completion-bearing
# If a stream is cancelled before `message_delta` lands, the accumulator
# below leaves completion_tokens stuck at 1, a reported count that
# calculate_usage() would keep. Count the completion-bearing
# usage events so `_reset_anthropic_cursor_completion_tokens` can tell a
# legitimate single-token reply (Anthropic emits 1 in BOTH message_start
# AND message_delta, so >=2 events is positive evidence message_delta
@ -875,10 +877,15 @@ class ChunkProcessor:
if usage_chunk is not None:
usage_chunk_dict = self._usage_chunk_calculation_helper(usage_chunk)
if usage_chunk_dict["prompt_tokens"] is not None and usage_chunk_dict["prompt_tokens"] > 0:
if usage_chunk_dict["prompt_tokens"] is not None and (
usage_chunk_dict["prompt_tokens"] > 0 or prompt_tokens is None
):
prompt_tokens = usage_chunk_dict["prompt_tokens"]
if usage_chunk_dict["completion_tokens"] is not None and usage_chunk_dict["completion_tokens"] > 0:
if usage_chunk_dict["completion_tokens"] is not None and (
usage_chunk_dict["completion_tokens"] > 0 or completion_tokens is None
):
completion_tokens = usage_chunk_dict["completion_tokens"]
if usage_chunk_dict["completion_tokens"] is not None and usage_chunk_dict["completion_tokens"] > 0:
completion_usage_updates += 1
if usage_chunk_dict["cache_creation_input_tokens"] is not None and (
usage_chunk_dict["cache_creation_input_tokens"] > 0 or cache_creation_input_tokens is None
@ -995,10 +1002,10 @@ class ChunkProcessor:
@staticmethod
def _reset_anthropic_cursor_completion_tokens(
chunks: Sequence["_UsageBearingChunk | ModelResponse"],
completion_tokens: int,
completion_tokens: int | None,
completion_usage_updates: int,
) -> int:
"""Reset a stale Anthropic ``message_start`` cursor placeholder to 0.
) -> int | None:
"""Reset a stale Anthropic ``message_start`` cursor placeholder to unreported.
See the ``completion_usage_updates`` comment in
``_calculate_usage_per_chunk``. The accumulated value is NOT a stale
@ -1006,8 +1013,8 @@ class ChunkProcessor:
carried a ``finish_reason`` (positive evidence ``message_delta``
arrived). Otherwise the only completion update we ever saw was the
Anthropic ``message_start`` cursor, a small placeholder whose magnitude
varies per request (1 and 8 both observed live), so reset to 0 and let
``calculate_usage()``'s ``or token_counter(...)`` fallback estimate from
varies per request (1 and 8 both observed live), so reset to None and let
``calculate_usage()``'s ``token_counter(...)`` fallback estimate from
the actually-received text and reasoning instead. Gated on
``custom_llm_provider == "anthropic"`` so the heuristic (which encodes
Anthropic's specific message_start SSE shape) does not silently affect
@ -1028,7 +1035,7 @@ class ChunkProcessor:
custom_llm_provider = hp.get("custom_llm_provider")
if custom_llm_provider == "anthropic":
return 0
return None
return completion_tokens
def calculate_usage(
@ -1063,15 +1070,18 @@ class ChunkProcessor:
cost: Final[float | None] = calculated_usage_per_chunk["cost"]
try:
returned_usage.prompt_tokens = prompt_tokens or (
count_prompt_tokens() if count_prompt_tokens else token_counter(model=model, messages=messages)
returned_usage.prompt_tokens = (
prompt_tokens
if prompt_tokens is not None
else (count_prompt_tokens() if count_prompt_tokens else token_counter(model=model, messages=messages))
)
except Exception: # don't allow this failing to block a complete streaming response from being returned
print_verbose("token_counter failed, assuming prompt tokens is 0")
returned_usage.prompt_tokens = 0
returned_usage.completion_tokens = (
completion_tokens
or (
if completion_tokens is not None
else (
token_counter(
model=model,
text=completion_output,

2
litellm/llms/a2a/chat/guardrail_translation/handler.py
View file

@ -125,7 +125,7 @@ class A2AGuardrailHandler(BaseTranslation):
litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None,
user_api_key_dict: Optional["UserAPIKeyAuth"] = None,
request_data: dict | None = None,
) -> Any:
) -> object:
"""
Process A2A output response by applying guardrails to text content.

24
litellm/llms/anthropic/chat/transformation.py
View file

@ -6,7 +6,7 @@ from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, NoReturn, cast
import httpx
from pydantic import ValidationError
from pydantic import BaseModel, ValidationError
from typing_extensions import ReadOnly, TypedDict
import litellm
@ -151,7 +151,7 @@ class _AnthropicToolResultBlock(TypedDict, total=False):
content: ReadOnly[object]
_ENUM_TYPE_CHECKS: Final[Mapping[str, Callable[[object], bool]]] = MappingProxyType(
_ENUM_TYPE_CHECKS: Final[Mapping[object, Callable[[object], bool]]] = MappingProxyType(
{
"null": lambda v: v is None,
"boolean": lambda v: isinstance(v, bool),
@ -164,7 +164,7 @@ _ENUM_TYPE_CHECKS: Final[Mapping[str, Callable[[object], bool]]] = MappingProxyT
)
def _enum_conflicts_with_declared_type(schema: Mapping[str, Any]) -> bool:
def _enum_conflicts_with_declared_type(schema: Mapping[str, object]) -> bool:
"""Whether ``schema``'s ``enum`` cannot match its declared ``type``."""
enum_values: Final = schema.get("enum")
declared_type: Final = schema.get("type")
@ -659,7 +659,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
return result
def get_json_schema_from_pydantic_object(self, response_format: Any | dict | None) -> dict | None:
def get_json_schema_from_pydantic_object(self, response_format: type[BaseModel] | dict | None) -> dict | None:
return type_to_response_format_param(
response_format, ref_template="/$defs/{model}"
) # Relevant issue: https://github.com/BerriAI/litellm/issues/7755
@ -1072,7 +1072,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
@staticmethod
def _sanitize_tool_names_in_request(
optional_params: dict[str, Any],
optional_params: dict[str, object],
) -> tuple[dict[str, str], dict[str, str]]:
"""Sanitize ``optional_params['tools']`` and ``optional_params['tool_choice']``
in place so every name matches Anthropic's ``^[a-zA-Z0-9_-]{1,128}$``.
@ -1119,7 +1119,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
# so a caller reusing the same tool list/dicts across requests
# doesn't see its inputs permanently rewritten (which would also
# drop the original key from `forward` on the next request).
new_tools: Final[list[Any]] = []
new_tools: Final[list[object]] = []
for t in tools:
if (
isinstance(t, dict)
@ -1442,7 +1442,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
entry_type = entry.get("type")
if entry_type == "compaction":
anthropic_edit: dict[str, Any] = {"type": "compact_20260112"}
anthropic_edit: dict[str, object] = {"type": "compact_20260112"}
compact_threshold = entry.get("compact_threshold")
# Rewrite to 'trigger' with correct nesting if threshold exists
if compact_threshold is not None and isinstance(compact_threshold, (int, float)):
@ -2442,9 +2442,11 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
code_by_id: Final[dict[str, str]] = {}
for tc in tool_calls:
try:
args = json.loads(tc.get("function", {}).get("arguments", "{}"))
args: object = json.loads(tc.get("function", {}).get("arguments", "{}"))
if not isinstance(args, Mapping):
continue
call_id = tc.get("id")
command = args.get("command", "")
command: object = args.get("command", "")
if isinstance(call_id, str):
code_by_id[call_id] = command if isinstance(command, str) else ""
except Exception:
@ -2514,8 +2516,8 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
tool_results: Sequence[_AnthropicToolResultBlock] | None,
compaction_blocks: Sequence[object] | None,
tool_calls: list[ChatCompletionToolCallChunk],
) -> dict[str, Any]:
provider_specific_fields: Final[dict[str, Any]] = {
) -> dict[str, object]:
provider_specific_fields: Final[dict[str, object]] = {
"citations": citations,
"thinking_blocks": thinking_blocks,
}

10
litellm/llms/anthropic/common_utils.py
View file

@ -7,7 +7,7 @@ import re
from collections.abc import Mapping, MutableMapping, Sequence
from datetime import datetime, timezone
from types import MappingProxyType
from typing import Any, Final, Literal
from typing import Any, Final, Literal, TypeVar
import httpx
from pydantic import BaseModel, ConfigDict, StrictBool, TypeAdapter, ValidationError
@ -40,6 +40,8 @@ from litellm.types.llms.anthropic import (
from litellm.types.llms.openai import AllMessageValues
from litellm.types.proxy.model_listing import ModelInfoResponse
_MessageT = TypeVar("_MessageT")
DROP_FORCED_TOOL_CHOICE_WARNING: Final = (
"Downgrading forced tool_choice to 'auto' for model=%s (drop_params=True): this model rejects tool_choice type "
"'any'/'tool' with a 400 because thinking is always on and a forced call would skip it."
@ -1121,7 +1123,7 @@ class AnthropicModelInfo(BaseLLMModelInfo):
return AnthropicTokenCounter()
def strip_advisor_blocks_from_messages(messages: list[Any], replace_with_text: bool = False) -> list[Any]:
def strip_advisor_blocks_from_messages(messages: list[_MessageT], replace_with_text: bool = False) -> list[_MessageT]:
"""
Remove (or replace) server_tool_use (name='advisor') and advisor_tool_result blocks
from assistant message content.
@ -1228,7 +1230,7 @@ def is_anthropic_invalid_thinking_block_error(error_text: str) -> bool:
return "must contain thinking" in lower
def strip_thinking_blocks_from_anthropic_messages(messages: list[Any]) -> list[Any]:
def strip_thinking_blocks_from_anthropic_messages(messages: Sequence[object]) -> list[object]:
"""
Return a new message list with thinking / redacted_thinking content blocks removed
from each message. Used to recover from invalid thinking signatures on retry.
@ -1236,7 +1238,7 @@ def strip_thinking_blocks_from_anthropic_messages(messages: list[Any]) -> list[A
Messages whose content is a list and becomes empty after stripping are omitted,
since Anthropic rejects empty content arrays.
"""
out: Final[list[Any]] = []
out: Final[list[object]] = []
for m in messages:
if not isinstance(m, dict):
out.append(m)

16
litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py
View file

@ -25,6 +25,9 @@ from litellm.constants import STREAM_SSE_KEEPALIVE_PING_BYTES
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.anthropic_messages.transformation import (
BaseAnthropicMessagesConfig,
)
HOLD_BACK_PING_INTERVAL_SECONDS: Final = 15.0
SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES: Final = (
@ -182,7 +185,7 @@ class AgenticAnthropicStreamingIterator:
http_handler: Any,
model: str,
messages: list[dict],
anthropic_messages_provider_config: Any,
anthropic_messages_provider_config: "BaseAnthropicMessagesConfig",
anthropic_messages_optional_request_params: dict,
logging_obj: "LiteLLMLoggingObj",
custom_llm_provider: str,
@ -402,7 +405,7 @@ class AgenticAnthropicStreamingIterator:
@staticmethod
def _rebuild_anthropic_response_from_sse(
raw_bytes: list[bytes],
) -> dict[str, Any] | None:
) -> dict[str, object] | None:
"""
Parse collected SSE bytes into an Anthropic Messages response dict.
@ -416,17 +419,18 @@ class AgenticAnthropicStreamingIterator:
"""
events: Final = _parse_sse_events(b"".join(raw_bytes))
response: Final[dict[str, Any]] = {
content: Final[list[dict[str, object]]] = []
response: Final[dict[str, object]] = {
"id": "",
"type": "message",
"role": "assistant",
"model": "",
"content": [],
"content": content,
"stop_reason": None,
"stop_sequence": None,
"usage": {"input_tokens": 0, "output_tokens": 0},
}
content_blocks: Final[dict[int, dict[str, Any]]] = {}
content_blocks: Final[dict[int, dict[str, object]]] = {}
saw_message_start = False
for event_type, data in events:
@ -448,6 +452,6 @@ class AgenticAnthropicStreamingIterator:
for idx in sorted(content_blocks.keys()):
block = content_blocks[idx]
block.pop("_partial_json", None)
response["content"].append(block)
content.append(block)
return response

6
litellm/llms/anthropic/files/handler.py
View file

@ -185,7 +185,11 @@ class AnthropicFilesHandler:
if not line.strip():
continue
anthropic_result = json.loads(line)
anthropic_result: object = json.loads(line)
if not isinstance(anthropic_result, dict):
raise TypeError(
f"Anthropic batch result line is not a JSON object: {type(anthropic_result).__name__}"
)
custom_id = anthropic_result.get("custom_id", "")
result = anthropic_result.get("result", {})
result_type = result.get("type", "")

9
litellm/llms/azure/assistants.py
View file

@ -1,5 +1,5 @@
from collections.abc import Coroutine, Iterable
from typing import Any, Final, Literal, TypedDict
from typing import Final, Literal, TypedDict
import httpx
from openai import AsyncAzureOpenAI, AzureOpenAI
@ -715,7 +715,8 @@ class AzureAssistantsAPI(BaseAzureLLM):
event_handler: AssistantEventHandler | None,
litellm_params: dict | None = None,
) -> AsyncAssistantStreamManager[AsyncAssistantEventHandler]:
data: Final[dict[str, Any]] = {
stream_fn: Final = client.beta.threads.runs.stream
base_data: Final[_RunThreadStreamData] = {
"thread_id": thread_id,
"assistant_id": assistant_id,
"additional_instructions": additional_instructions,
@ -725,8 +726,8 @@ class AzureAssistantsAPI(BaseAzureLLM):
"tools": tools,
}
if event_handler is not None:
data["event_handler"] = event_handler
return client.beta.threads.runs.stream(**data)
return stream_fn(**base_data, event_handler=event_handler)
return stream_fn(**base_data)
def run_thread_stream(
self,

9
litellm/llms/azure/text_to_speech/transformation.py
View file

@ -19,6 +19,7 @@ from litellm.secret_managers.main import get_secret_str
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.types.llms.openai import HttpxBinaryResponseContent
else:
LiteLLMLoggingObj = Any
@ -67,15 +68,15 @@ class AzureAVATextToSpeechConfig(BaseTextToSpeechConfig):
litellm_params_dict: dict,
logging_obj: "LiteLLMLoggingObj",
timeout: float | httpx.Timeout,
extra_headers: dict[str, Any] | None,
base_llm_http_handler: Any,
extra_headers: dict[str, object] | None,
base_llm_http_handler: "BaseLLMHTTPHandler",
aspeech: bool,
api_base: str | None,
api_key: str | None,
**kwargs: Any,
**kwargs: object,
) -> Union[
"HttpxBinaryResponseContent",
Coroutine[Any, Any, "HttpxBinaryResponseContent"],
Coroutine[object, object, "HttpxBinaryResponseContent"],
]:
"""
Dispatch method to handle Azure AVA TTS requests

2
litellm/llms/azure_ai/anthropic/count_tokens/handler.py
View file

@ -33,7 +33,7 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig):
litellm_params: dict[str, Any] | None = None,
timeout: float | httpx.Timeout | None = None,
tools: list[dict[str, Any]] | None = None,
system: Any | None = None,
system: object = None,
) -> dict[str, Any]:
"""
Handle a CountTokens request using httpx with Azure authentication.

10
litellm/llms/base_llm/videos/transformation.py
View file

@ -180,7 +180,7 @@ class BaseVideoConfig(ABC):
api_base: str,
litellm_params: GenericLiteLLMParams,
headers: dict,
extra_body: dict[str, Any] | None = None,
extra_body: dict[str, object] | None = None,
) -> tuple[str, dict]:
"""
Transform the video remix request into a URL and data
@ -207,7 +207,7 @@ class BaseVideoConfig(ABC):
after: str | None = None,
limit: int | None = None,
order: str | None = None,
extra_query: dict[str, Any] | None = None,
extra_query: dict[str, object] | None = None,
) -> tuple[str, dict]:
"""
Transform the video list request into a URL and params
@ -355,8 +355,8 @@ class BaseVideoConfig(ABC):
litellm_params: GenericLiteLLMParams,
headers: dict,
video_file: FileContent | None = None,
extra_body: dict[str, Any] | None = None,
prefetched_source_data: dict[str, Any] | None = None,
extra_body: dict[str, object] | None = None,
prefetched_source_data: dict[str, object] | None = None,
) -> tuple[str, Mapping[str, object], RequestFiles | None]:
"""
Transform the video edit request into a URL plus either JSON data or
@ -386,7 +386,7 @@ class BaseVideoConfig(ABC):
api_base: str,
litellm_params: GenericLiteLLMParams,
headers: dict,
extra_body: dict[str, Any] | None = None,
extra_body: dict[str, object] | None = None,
) -> tuple[str, dict]:
"""
Transform the video extension request into a URL and JSON data.

2
litellm/llms/bedrock/chat/converse_transformation.py
View file

@ -1126,7 +1126,7 @@ class AmazonConverseConfig(BaseConfig):
return optional_params
def _map_request_metadata_param(self, value: Any, optional_params: dict) -> None:
def _map_request_metadata_param(self, value: object, optional_params: dict) -> None:
if value is not None and isinstance(value, dict):
self._validate_request_metadata(value)
optional_params["requestMetadata"] = value

11
litellm/llms/bedrock/count_tokens/bedrock_token_counter.py
View file

@ -2,6 +2,7 @@
Bedrock Token Counter implementation using the CountTokens API.
"""
from collections.abc import Mapping, Sequence
from typing import Any, Final
from litellm._logging import verbose_logger
@ -26,12 +27,12 @@ class BedrockTokenCounter(BaseTokenCounter):
async def count_tokens(
self,
model_to_use: str,
messages: list[dict[str, Any]] | None,
contents: list[dict[str, Any]] | None,
messages: Sequence[Mapping[str, object]] | None,
contents: Sequence[Mapping[str, object]] | None,
deployment: dict[str, Any] | None = None,
request_model: str = "",
tools: list[dict[str, Any]] | None = None,
system: Any | None = None,
tools: Sequence[Mapping[str, object]] | None = None,
system: object | None = None,
) -> TokenCountResponse | None:
"""
Count tokens using AWS Bedrock's CountTokens API.
@ -56,7 +57,7 @@ class BedrockTokenCounter(BaseTokenCounter):
litellm_params: Final = deployment.get("litellm_params", {})
# Build request data in the format expected by BedrockCountTokensHandler
request_data: Final[dict[str, Any]] = {
request_data: Final[dict[str, object]] = {
"model": model_to_use,
"messages": messages,
}

6
litellm/llms/bedrock/files/transformation.py
View file

@ -375,7 +375,7 @@ def _listed_managed_file(
)
def _uploaded_object_size(litellm_params: Mapping[str, object], raw_response: Response) -> int:
def _uploaded_object_size(litellm_params: Mapping[str, object], response_headers: Mapping[str, str]) -> int:
"""
S3 answers PutObject with an empty body, so the stored object size comes from the
signed request recorded by `transform_create_file_request`, not the response headers.
@ -383,7 +383,7 @@ def _uploaded_object_size(litellm_params: Mapping[str, object], raw_response: Re
uploaded_size: Final = litellm_params.get(UPLOAD_CONTENT_LENGTH_PARAM)
if isinstance(uploaded_size, int):
return uploaded_size
response_content_length: Final = raw_response.headers.get("Content-Length", "0")
response_content_length: Final = response_headers.get("Content-Length", "0")
return int(response_content_length) if response_content_length.isdigit() else 0
@ -1277,7 +1277,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
filename=filename,
created_at=int(time.time()), # Current timestamp
status="uploaded",
bytes=_uploaded_object_size(litellm_params=litellm_params, raw_response=raw_response),
bytes=_uploaded_object_size(litellm_params=litellm_params, response_headers=raw_response.headers),
object="file",
)

4
litellm/llms/bedrock/image_edit/stability_transformation.py
View file

@ -125,7 +125,7 @@ class BedrockStabilityImageEditConfig(BaseImageEditConfig):
}
# Create a copy to not mutate original - convert TypedDict to regular dict
mapped_params: Final[dict[str, Any]] = dict(image_edit_optional_params)
mapped_params: Final[dict[str, object]] = dict(image_edit_optional_params)
for k, v in image_edit_optional_params.items():
if k in param_mapping:
@ -172,7 +172,7 @@ class BedrockStabilityImageEditConfig(BaseImageEditConfig):
Returns the request body dict that will be JSON-encoded by the handler.
"""
# Build Bedrock Stability request
data: Final[dict[str, Any]] = {
data: Final[dict[str, object]] = {
"output_format": "png", # Default to PNG
}

5
litellm/llms/bedrock/passthrough/guardrail_translation/handler.py
View file

@ -14,6 +14,9 @@ from litellm.types.utils import GenericGuardrailAPIInputs
if TYPE_CHECKING:
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.pass_through.guardrail_translation.handler import (
PassThroughEndpointHandler,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.utils import ProxyLogging
@ -27,7 +30,7 @@ def _is_converse_endpoint(endpoint: str) -> bool:
return bool(parts) and parts[-1] in _CONVERSE_ACTIONS
def _generic_passthrough_handler() -> BaseTranslation:
def _generic_passthrough_handler() -> "PassThroughEndpointHandler":
"""
Fallback for non-Converse Bedrock routes (e.g. invoke). The generic
handler scans the full request/response payload so blocking guardrails

8
litellm/llms/bedrock_mantle/responses/transformation.py
View file

@ -16,8 +16,8 @@ BaseAWSLLM._sign_request after the request body is finalized.
"""
import json
from collections.abc import Mapping
from typing import Any, Final, cast # noqa: TID251 # map_openai_params returns the filtered params as a bare dict
from collections.abc import Mapping, Sequence
from typing import Final, cast # noqa: TID251 # map_openai_params returns the filtered params as a bare dict
import httpx
from typing_extensions import ReadOnly, TypedDict
@ -142,9 +142,9 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI
return False
@staticmethod
def _filter_unsupported_tools(tools: list[Any]) -> list[Any]:
def _filter_unsupported_tools(tools: "Sequence[object]") -> "list[object]":
"""Keep only tool types Mantle's Responses API accepts."""
kept: Final[list[Any]] = []
kept: Final[list[object]] = []
dropped_types: Final[list[str]] = []
for tool in tools:
if not isinstance(tool, dict):

4
litellm/llms/chatgpt/common_utils.py
View file

@ -268,7 +268,7 @@ def _normalize_litellm_params(litellm_params: Any | None) -> dict:
return {}
def get_chatgpt_session_id(litellm_params: Any | None) -> str | None:
def get_chatgpt_session_id(litellm_params: object) -> str | None:
params: Final = _normalize_litellm_params(litellm_params)
for key in ("litellm_session_id", "session_id"):
value = params.get(key)
@ -286,5 +286,5 @@ def get_chatgpt_session_id(litellm_params: Any | None) -> str | None:
return None
def ensure_chatgpt_session_id(litellm_params: Any | None) -> str:
def ensure_chatgpt_session_id(litellm_params: object) -> str:
return get_chatgpt_session_id(litellm_params) or str(uuid4())

22
litellm/llms/chatgpt/responses/transformation.py
View file

@ -1,5 +1,8 @@
from collections.abc import Mapping
from typing import TYPE_CHECKING, Any, Final
import httpx
from litellm.exceptions import AuthenticationError
from litellm.litellm_core_utils.core_helpers import process_response_headers
from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
@ -13,6 +16,7 @@ from litellm.responses.sse_output_recovery import (
record_output_text_chunk,
)
from litellm.types.llms.openai import (
ResponseInputParam,
ResponsesAPIResponse,
ResponsesAPIStreamEvents,
)
@ -64,7 +68,7 @@ class ChatGPTResponsesAPIConfig(OpenAIResponsesAPIConfig):
def transform_responses_api_request(
self,
model: str,
input: Any,
input: str | ResponseInputParam,
response_api_optional_request_params: dict,
litellm_params: GenericLiteLLMParams,
headers: dict,
@ -109,9 +113,9 @@ class ChatGPTResponsesAPIConfig(OpenAIResponsesAPIConfig):
def transform_response_api_response(
self,
model: str,
raw_response: Any,
raw_response: httpx.Response,
logging_obj: "LiteLLMLoggingObj",
):
) -> ResponsesAPIResponse:
body_text: Final = raw_response.text or ""
if not self._should_parse_as_sse(raw_response=raw_response, body_text=body_text):
return super().transform_response_api_response(
@ -135,7 +139,7 @@ class ChatGPTResponsesAPIConfig(OpenAIResponsesAPIConfig):
self._attach_response_headers(completed_response=completed_response, raw_response=raw_response)
return completed_response
def _should_parse_as_sse(self, raw_response: Any, body_text: str) -> bool:
def _should_parse_as_sse(self, raw_response: httpx.Response, body_text: str) -> bool:
content_type: Final = (raw_response.headers or {}).get("content-type", "")
if "text/event-stream" in content_type.lower():
return True
@ -150,8 +154,8 @@ class ChatGPTResponsesAPIConfig(OpenAIResponsesAPIConfig):
def _extract_completed_response_from_sse(self, body_text: str) -> tuple[ResponsesAPIResponse | None, str | None]:
completed_response = None
error_message = None
streamed_output_items: Final[dict[int, dict]] = {}
text_only_output_items: Final[dict[int, dict]] = {}
streamed_output_items: Final[dict[int, dict[str, object]]] = {}
text_only_output_items: Final[dict[int, dict[str, object]]] = {}
for chunk in body_text.splitlines():
parsed_chunk = parse_sse_json_chunk(chunk)
if parsed_chunk is None:
@ -178,7 +182,7 @@ class ChatGPTResponsesAPIConfig(OpenAIResponsesAPIConfig):
# output_index, but text-only items at indices without a
# matching OUTPUT_ITEM_DONE must still be preserved (e.g.
# providers that emit only OUTPUT_TEXT_DONE for some indices).
merged_items: dict[int, dict] = {**text_only_output_items}
merged_items: dict[int, dict[str, object]] = {**text_only_output_items}
merged_items.update(streamed_output_items)
completed_response = self._build_completed_response_from_chunk(
parsed_chunk=parsed_chunk,
@ -197,7 +201,7 @@ class ChatGPTResponsesAPIConfig(OpenAIResponsesAPIConfig):
return completed_response, error_message
def _build_completed_response_from_chunk(
self, parsed_chunk: dict[str, Any], streamed_output_items: dict[int, dict]
self, parsed_chunk: Mapping[str, object], streamed_output_items: Mapping[int, dict[str, object]]
) -> ResponsesAPIResponse | None:
response_payload = parsed_chunk.get("response")
if not isinstance(response_payload, dict):
@ -223,7 +227,7 @@ class ChatGPTResponsesAPIConfig(OpenAIResponsesAPIConfig):
def _attach_response_headers(
self,
completed_response: ResponsesAPIResponse,
raw_response: Any,
raw_response: httpx.Response,
) -> None:
raw_headers: Final = dict(raw_response.headers)
processed_headers: Final = process_response_headers(raw_headers)

2
litellm/llms/cohere/chat/transformation.py
View file

@ -110,7 +110,7 @@ class CohereChatConfig(BaseConfig):
tool_results: list | None = None,
seed: int | None = None,
) -> None:
locals_: Final = locals().copy()
locals_: Final[dict[str, object]] = locals().copy()
for key, value in locals_.items():
if key != "self" and value is not None:
setattr(self.__class__, key, value)

15
litellm/llms/cohere/embed/v1_transformation.py
View file

@ -2,7 +2,8 @@
Legacy /v1/embedding transformation logic for Bedrock Cohere.
"""
from typing import Any, Final
from collections.abc import Sized
from typing import Final, Protocol
import httpx
@ -16,6 +17,12 @@ from litellm.types.utils import EmbeddingResponse, PromptTokensDetailsWrapper, U
from litellm.utils import is_base64_encoded
class _SupportsEncode(Protocol):
"""Tokenizer handle: the embedding usage path only encodes text to measure its token length."""
def encode(self, text: str, /) -> Sized: ...
class CohereEmbeddingConfig:
"""
Reference: https://docs.cohere.com/v2/reference/embed
@ -61,7 +68,7 @@ class CohereEmbeddingConfig:
return transformed_request
def _calculate_usage(self, input: list[str], encoding: Any, meta: dict) -> Usage:
def _calculate_usage(self, input: list[str], encoding: _SupportsEncode, meta: dict) -> Usage:
input_tokens = 0
text_tokens: Final[int | None] = meta.get("billed_units", {}).get("input_tokens")
@ -97,7 +104,7 @@ class CohereEmbeddingConfig:
data: dict | CohereEmbeddingRequest,
model_response: EmbeddingResponse,
model: str,
encoding: Any,
encoding: _SupportsEncode,
input: list,
) -> EmbeddingResponse:
response_json: Final = response.json()
@ -121,7 +128,7 @@ class CohereEmbeddingConfig:
response_json: dict,
model_response: EmbeddingResponse,
model: str,
encoding: Any,
encoding: _SupportsEncode,
input: list,
) -> EmbeddingResponse:
"""

5
litellm/llms/custom_httpx/http_handler.py
View file

@ -479,6 +479,11 @@ def _safe_get_response_text(response: httpx.Response) -> str:
return ""
def header_value(headers: Mapping[str, str], name: str) -> str | None:
"""Read one header as ``str | None``; ``httpx.Headers.get`` itself is typed ``Any``."""
return headers.get(name)
async def _safe_aread_response(response: httpx.Response, timeout: float | None = None) -> bytes:
"""Safely read async response body, falling back to empty bytes on errors."""
try:

36
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -4,6 +4,7 @@ import ssl
from collections.abc import AsyncGenerator, AsyncIterator, Coroutine, Iterator, Mapping, Sequence
from contextlib import asynccontextmanager
from functools import lru_cache
from itertools import chain
from types import MappingProxyType, ModuleType
from typing import (
TYPE_CHECKING,
@ -5613,18 +5614,29 @@ class BaseLLMHTTPHandler:
}
internal_keys: Final = {"litellm_logging_obj"}
kwargs_for_followup: Final = {
k: v
for k, v in kwargs.items()
if not is_interception_internal_key(k, prefixes=NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES)
and k != "_code_interpreter_interception_converted_stream"
and k not in internal_keys
and k not in optional_params
}
kwargs_for_followup.update(patch.kwargs)
kwargs_for_followup["_agentic_loop_depth"] = depth + 1
kwargs_for_followup["max_agentic_loops"] = max_loops
kwargs_for_followup["_agentic_loop_fingerprints"] = fingerprints + [fingerprint]
kwargs_for_followup: Final = MappingProxyType(
{
key: value
for key, value in chain(
(
(k, v)
for k, v in kwargs.items()
if not is_interception_internal_key(
k, prefixes=NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES
)
and k != "_code_interpreter_interception_converted_stream"
and k not in internal_keys
and k not in optional_params
),
((k, v) for k, v in patch.kwargs.items() if k not in optional_params),
(
("_agentic_loop_depth", depth + 1),
("max_agentic_loops", max_loops),
("_agentic_loop_fingerprints", fingerprints + [fingerprint]),
),
)
}
)
try:
response: ResponsesAPIResponse | BaseResponsesAPIStreamingIterator = await litellm.aresponses(

10
litellm/llms/databricks/chat/transformation.py
View file

@ -3,7 +3,7 @@ Translates from OpenAI's `/v1/chat/completions` to Databricks' `/chat/completion
"""
import os
from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping
from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final, Literal, cast, overload
import httpx
@ -67,7 +67,7 @@ def _is_bare_assistant_message(message_dict: Mapping[str, object]) -> bool:
)
def _sanitize_empty_content(message_dict: dict[str, Any]) -> None:
def _sanitize_empty_content(message_dict: dict[str, object]) -> None:
"""
Remove or filter content so empty text blocks are not sent.
Databricks Model Serving uses Anthropic Messages API spec and rejects empty text blocks.
@ -430,7 +430,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
@overload
def _transform_messages(
self, messages: list[AllMessageValues], model: str, is_async: Literal[True]
) -> Coroutine[Any, Any, list[AllMessageValues]]: ...
) -> Coroutine[object, object, list[AllMessageValues]]: ...
@overload
def _transform_messages(
@ -442,7 +442,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
def _transform_messages(
self, messages: list[AllMessageValues], model: str, is_async: bool = False
) -> list[AllMessageValues] | Coroutine[Any, Any, list[AllMessageValues]]:
) -> list[AllMessageValues] | Coroutine[object, object, list[AllMessageValues]]:
"""
Databricks does not support:
- 'name' in user message.
@ -564,7 +564,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
@staticmethod
def extract_citations(
content: AllDatabricksContentValues | None,
) -> list[Any] | None:
) -> Sequence[Sequence[Mapping[str, object]]] | None:
if content is None:
return None
citations: Final = []

4
litellm/llms/fireworks_ai/chat/transformation.py
View file

@ -1,6 +1,6 @@
import json
from collections.abc import AsyncIterator, Iterator, Mapping
from typing import TYPE_CHECKING, Any, Final, Literal, cast
from typing import TYPE_CHECKING, Final, Literal, cast
import httpx
@ -759,7 +759,7 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig):
streaming_response: Iterator[str] | AsyncIterator[str] | ModelResponse,
sync_stream: bool,
json_mode: bool | None = False,
) -> Any:
) -> "FireworksAIChatCompletionStreamingHandler":
return FireworksAIChatCompletionStreamingHandler(
streaming_response=streaming_response,
sync_stream=sync_stream,

22
litellm/llms/gdc/chat/transformation.py
View file

@ -7,14 +7,32 @@ import os
import re
import threading
from collections.abc import Callable
from typing import Any, Final, Protocol
from typing import Final, Protocol
from urllib.parse import urlsplit
from typing_extensions import ReadOnly, TypedDict, Unpack
import litellm
from litellm.llms.openai_like.chat.transformation import OpenAILikeChatConfig
from litellm.types.llms.openai import AllMessageValues
class _OpenAIGPTConfigOptions(TypedDict, total=False):
"""The sampling defaults ``OpenAIGPTConfig.__init__`` accepts and stashes on the class."""
frequency_penalty: ReadOnly[int | None]
function_call: ReadOnly[str | dict[str, object] | None]
functions: ReadOnly[list[object] | None]
logit_bias: ReadOnly[dict[str, object] | None]
max_tokens: ReadOnly[int | None]
n: ReadOnly[int | None]
presence_penalty: ReadOnly[int | None]
stop: ReadOnly[str | list[object] | None]
temperature: ReadOnly[int | None]
top_p: ReadOnly[int | None]
response_format: ReadOnly[dict[str, object] | None]
class _GDCHAudienceCredentials(Protocol):
"""A GDCH service account credential already bound to an audience, ready to mint a bearer token."""
@ -32,7 +50,7 @@ class GDCGeminiConfig(OpenAILikeChatConfig):
_GDCH_CREDENTIAL_TYPE: Final[str] = "gdch_service_account"
_PATH_ID_PATTERN: Final[re.Pattern[str]] = re.compile(r"^[a-zA-Z0-9_-]+$")
def __init__(self, **kwargs: Any) -> None:
def __init__(self, **kwargs: Unpack[_OpenAIGPTConfigOptions]) -> None:
super().__init__(**kwargs)
self._creds_lock = threading.Lock()
self._gdch_creds_cache: dict[tuple[str, str], _GDCHAudienceCredentials] = {}

2
litellm/llms/gemini/count_tokens/handler.py
View file

@ -84,7 +84,7 @@ class GoogleAIStudioTokenCounter:
api_key: str | None = None,
api_base: str | None = None,
timeout: float | httpx.Timeout | None = None,
**kwargs,
**kwargs: object,
) -> dict[str, Any]:
"""
Count tokens using Google Gen AI Studio countTokens endpoint.

13
litellm/llms/gemini/image_edit/transformation.py
View file

@ -1,4 +1,5 @@
import base64
from collections.abc import Mapping
from io import BufferedReader, BytesIO
from typing import TYPE_CHECKING, Any, Final, cast
@ -44,7 +45,7 @@ class GeminiImageEditConfig(BaseImageEditConfig):
image_edit_optional_params: ImageEditOptionalRequestParams,
model: str,
drop_params: bool,
) -> dict[str, Any]:
) -> dict[str, object]:
return map_openai_image_params_to_gemini(
params=image_edit_optional_params,
model=model,
@ -87,10 +88,10 @@ class GeminiImageEditConfig(BaseImageEditConfig):
model: str,
prompt: str | None,
image: FileTypes | None,
image_edit_optional_request_params: dict[str, Any],
image_edit_optional_request_params: Mapping[str, object],
litellm_params: GenericLiteLLMParams,
headers: dict,
) -> tuple[dict[str, Any], RequestFiles | None]:
) -> tuple[dict[str, object], RequestFiles | None]:
inline_parts: Final = self._prepare_inline_image_parts(image) if image else []
if not inline_parts:
raise ValueError("Gemini image edit requires at least one image.")
@ -106,7 +107,7 @@ class GeminiImageEditConfig(BaseImageEditConfig):
}
]
request_body: Final[dict[str, Any]] = {"contents": contents}
request_body: Final[dict[str, object]] = {"contents": contents}
request_body["generationConfig"] = get_gemini_image_generation_config(
model=model,
@ -153,14 +154,14 @@ class GeminiImageEditConfig(BaseImageEditConfig):
model_response.usage = transform_gemini_image_usage(response_json["usageMetadata"])
return model_response
def _prepare_inline_image_parts(self, image: FileTypes | list[FileTypes]) -> list[dict[str, Any]]:
def _prepare_inline_image_parts(self, image: FileTypes | list[FileTypes]) -> list[dict[str, object]]:
images: list[FileTypes]
if isinstance(image, list):
images = image
else:
images = [image]
inline_parts: Final[list[dict[str, Any]]] = []
inline_parts: Final[list[dict[str, object]]] = []
for img in images:
if img is None:
continue

14
litellm/llms/gigachat/chat/transformation.py
View file

@ -81,9 +81,17 @@ class GigaChatConfig(BaseConfig):
repetition_penalty: float | None = None,
profanity_check: bool | None = None,
) -> None:
locals_: Final = locals().copy()
for key, value in locals_.items():
if key != "self" and value is not None:
config_params: Final[Mapping[str, float | int | bool | None]] = MappingProxyType(
{
"temperature": temperature,
"top_p": top_p,
"max_tokens": max_tokens,
"repetition_penalty": repetition_penalty,
"profanity_check": profanity_check,
}
)
for key, value in config_params.items():
if value is not None:
setattr(self.__class__, key, value)
# Instance variables for current request context
self._current_credentials: str | None = None

7
litellm/llms/github_copilot/responses/transformation.py
View file

@ -19,6 +19,7 @@ from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfi
from litellm.types.llms.openai import (
ResponseInputParam,
ResponsesAPIOptionalRequestParams,
ResponsesAPIStreamingResponse,
)
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import LlmProviders
@ -129,7 +130,7 @@ class GithubCopilotResponsesAPIConfig(OpenAIResponsesAPIConfig):
model: str,
parsed_chunk: dict,
logging_obj: LiteLLMLoggingObj,
) -> Any:
) -> ResponsesAPIStreamingResponse:
parsed_chunk = self._normalize_stream_item_id(parsed_chunk)
return super().transform_streaming_response(
model=model,
@ -262,7 +263,7 @@ class GithubCopilotResponsesAPIConfig(OpenAIResponsesAPIConfig):
# Return the responses endpoint
return f"{effective_api_base}/responses"
def _handle_reasoning_item(self, item: dict[str, Any]) -> dict[str, Any]:
def _handle_reasoning_item(self, item: dict[str, object]) -> dict[str, object]:
"""
Handle reasoning items for GitHub Copilot, preserving encrypted_content.
@ -280,7 +281,7 @@ class GithubCopilotResponsesAPIConfig(OpenAIResponsesAPIConfig):
# Filter out None values for known problematic fields,
# but preserve encrypted_content even if it exists
filtered_item: Final[dict[str, Any]] = {}
filtered_item: Final[dict[str, object]] = {}
for k, v in item.items():
# Always include encrypted_content if present (even if None)
if k == "encrypted_content":

23
litellm/llms/hosted_vllm/chat/transformation.py
View file

@ -4,7 +4,7 @@ Translate from OpenAI's `/v1/chat/completions` to VLLM's `/v1/chat/completions`
import json
from collections.abc import Coroutine
from typing import Any, Final, Literal, cast, overload
from typing import Final, Literal, cast, overload
from litellm.litellm_core_utils.prompt_templates.common_utils import (
_get_image_mime_type_from_url,
@ -28,12 +28,12 @@ from ...openai.chat.gpt_transformation import OpenAIGPTConfig
class HostedVLLMChatConfig(OpenAIGPTConfig):
def _convert_custom_tools_to_function_tools(self, tools: list[dict[str, Any]]) -> list[dict[str, Any]]:
def _convert_custom_tools_to_function_tools(self, tools: list[dict[str, object]]) -> list[dict[str, object]]:
"""
vLLM chat completions currently accepts only OpenAI function tools.
Convert custom tools into function tools so request validation does not fail.
"""
converted_tools: Final[list[dict[str, Any]]] = []
converted_tools: Final[list[dict[str, object]]] = []
for idx, tool in enumerate(tools):
if not isinstance(tool, dict):
converted_tools.append(tool)
@ -63,17 +63,14 @@ class HostedVLLMChatConfig(OpenAIGPTConfig):
"required": ["input"],
}
function_tool: dict[str, Any] = {
"type": "function",
"function": {
"name": str(tool_name),
"parameters": tool_parameters,
},
function_definition: dict[str, object] = {
"name": str(tool_name),
"parameters": tool_parameters,
}
if isinstance(tool_description, str):
function_tool["function"]["description"] = tool_description
function_definition["description"] = tool_description
converted_tools.append(function_tool)
converted_tools.append({"type": "function", "function": function_definition})
return converted_tools
@ -148,7 +145,7 @@ class HostedVLLMChatConfig(OpenAIGPTConfig):
@overload
def _transform_messages(
self, messages: list[AllMessageValues], model: str, is_async: Literal[True]
) -> Coroutine[Any, Any, list[AllMessageValues]]: ...
) -> Coroutine[object, object, list[AllMessageValues]]: ...
@overload
def _transform_messages(
@ -160,7 +157,7 @@ class HostedVLLMChatConfig(OpenAIGPTConfig):
def _transform_messages(
self, messages: list[AllMessageValues], model: str, is_async: bool = False
) -> list[AllMessageValues] | Coroutine[Any, Any, list[AllMessageValues]]:
) -> list[AllMessageValues] | Coroutine[object, object, list[AllMessageValues]]:
"""
Support translating:
- video files from file_id or file_data to video_url

10
litellm/llms/huggingface/embedding/transformation.py
View file

@ -84,13 +84,13 @@ class HuggingFaceEmbeddingConfig(BaseConfig):
typical_p: float | None = None,
watermark: bool | None = None,
) -> None:
locals_: Final = locals().copy()
locals_: Final[dict[str, object]] = locals().copy()
for key, value in locals_.items():
if key != "self" and value is not None:
setattr(self.__class__, key, value)
@classmethod
def get_config(cls):
def get_config(cls) -> dict[str, object]:
return super().get_config()
def get_special_options_params(self):
@ -352,17 +352,17 @@ class HuggingFaceEmbeddingConfig(BaseConfig):
model: str,
data: dict,
api_key: str | None = None,
) -> list[dict[str, Any]]:
) -> list[dict[str, str]]:
streamed_response: Final = CustomStreamWrapper(
completion_stream=response.iter_lines(),
model=model,
custom_llm_provider="huggingface",
logging_obj=logging_obj,
)
content = ""
content: str = ""
for chunk in streamed_response:
content += chunk["choices"][0]["delta"]["content"]
completion_response: Final[list[dict[str, Any]]] = [{"generated_text": content}]
completion_response: Final[list[dict[str, str]]] = [{"generated_text": content}]
## LOGGING
logging_obj.post_call(
input=data,

37
litellm/llms/nvidia_riva/audio_transcription/handler.py
View file

@ -27,7 +27,6 @@ without the optional STT extras installed.
import asyncio
import inspect
from collections.abc import Callable, Iterable
from types import ModuleType
from typing import TYPE_CHECKING, Any, Final, Protocol
from litellm.litellm_core_utils.audio_utils.utils import (
@ -95,11 +94,37 @@ class _AudioEncoding(Protocol):
def LINEAR_PCM(self) -> object: ...
def _auth_factory(riva_module: ModuleType) -> Callable[..., _RivaAuth]:
class _RivaClientModule(Protocol):
"""The ``riva.client`` entry points this handler calls."""
@property
def Auth(self) -> Callable[..., _RivaAuth]: ...
@property
def ASRService(self) -> Callable[[_RivaAuth], _AsrService]: ...
class _RivaAsrModule(Protocol):
"""The protobuf constructors this handler calls, from whichever module exposes them."""
@property
def AudioEncoding(self) -> _AudioEncoding: ...
@property
def RecognitionConfig(self) -> Callable[..., _RecognitionConfig]: ...
@property
def StreamingRecognitionConfig(self) -> Callable[..., _StreamingRecognitionConfig]: ...
@property
def EndpointingConfig(self) -> Callable[..., _EndpointingConfig]: ...
def _auth_factory(riva_module: _RivaClientModule) -> Callable[..., _RivaAuth]:
return riva_module.Auth
def _audio_encoding(riva_asr_module: ModuleType) -> _AudioEncoding:
def _audio_encoding(riva_asr_module: _RivaAsrModule) -> _AudioEncoding:
return riva_asr_module.AudioEncoding
@ -317,7 +342,7 @@ class NvidiaRivaAudioTranscription:
def _construct_auth(
self,
riva_module: ModuleType,
riva_module: _RivaClientModule,
api_base: str,
api_key: str | None,
optional_params: dict,
@ -349,7 +374,7 @@ class NvidiaRivaAudioTranscription:
return _auth_factory(riva_module)(None, use_ssl, api_base, metadata)
def _build_recognition_config_proto(
self, riva_asr_module: ModuleType, recognition_config_dict: dict[str, Any]
self, riva_asr_module: _RivaAsrModule, recognition_config_dict: dict[str, Any]
) -> _RecognitionConfig:
encoding_name: Final = (recognition_config_dict.get("encoding") or "LINEAR_PCM").upper()
encoding_enum: Final[object] = getattr(
@ -436,7 +461,7 @@ class NvidiaRivaAudioTranscription:
return final_results
def _import_riva() -> tuple[ModuleType, ModuleType]:
def _import_riva() -> tuple[_RivaClientModule, _RivaAsrModule]:
"""
Lazy import of ``riva.client`` and ``riva.client.proto.riva_asr_pb2``.

24
litellm/llms/ollama/chat/transformation.py
View file

@ -1,6 +1,6 @@
import json
import time
from collections.abc import AsyncIterator, Iterator
from collections.abc import AsyncIterator, Iterator, Mapping
from typing import TYPE_CHECKING, Any, Final, cast
from httpx._models import Headers, Response
@ -124,7 +124,7 @@ class OllamaChatConfig(BaseConfig):
setattr(self.__class__, key, value)
@classmethod
def get_config(cls):
def get_config(cls) -> dict[str, object]:
return super().get_config()
def get_supported_openai_params(self, model: str):
@ -420,6 +420,18 @@ class OllamaChatConfig(BaseConfig):
)
def _done_chunk_usage(chunk: Mapping[str, object]) -> ChatCompletionUsageBlock | None:
prompt_eval_count: Final = chunk.get("prompt_eval_count")
eval_count: Final = chunk.get("eval_count")
if chunk.get("done") is not True or not isinstance(prompt_eval_count, int) or not isinstance(eval_count, int):
return None
return ChatCompletionUsageBlock(
prompt_tokens=prompt_eval_count,
completion_tokens=eval_count,
total_tokens=prompt_eval_count + eval_count,
)
class OllamaChatCompletionResponseIterator(BaseModelResponseIterator):
started_reasoning_content: bool = False
finished_reasoning_content: bool = False
@ -528,17 +540,11 @@ class OllamaChatCompletionResponseIterator(BaseModelResponseIterator):
)
]
usage: Final = ChatCompletionUsageBlock(
prompt_tokens=chunk.get("prompt_eval_count", 0),
completion_tokens=chunk.get("eval_count", 0),
total_tokens=chunk.get("prompt_eval_count", 0) + chunk.get("eval_count", 0),
)
return ModelResponseStream(
id=str(uuid.uuid4()),
object="chat.completion.chunk",
created=int(time.time()), # ollama created_at is in UTC
usage=usage,
usage=_done_chunk_usage(chunk),
model=chunk["model"],
choices=choices,
)

2
litellm/llms/ollama/completion/transformation.py
View file

@ -231,7 +231,7 @@ class OllamaConfig(BaseConfig):
model: str,
api_base: str | None = None,
api_key: str | None = None,
) -> Any:
) -> dict[str, object] | None:
"""
curl http://localhost:11434/api/show -d '{
"name": "mistral"

20
litellm/llms/openai/chat/guardrail_translation/handler.py
View file

@ -17,7 +17,7 @@ This pattern can be replicated for other message formats (e.g., Anthropic).
import json
import time
import uuid
from collections.abc import Sequence
from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Union, cast
@ -269,7 +269,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
def _extract_inputs(
self,
message: dict[str, Any],
message: Mapping[str, object],
msg_idx: int,
texts_to_check: list[str],
images_to_check: list[str],
@ -330,7 +330,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
async def _apply_guardrail_responses_to_input_texts(
self,
messages: list[dict[str, Any]],
messages: list[dict[str, object]],
responses: list[str],
task_mappings: list[tuple[int, int | None]],
) -> None:
@ -355,12 +355,12 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
elif isinstance(content, list) and content_idx_optional is not None:
# Replace specific text item in list content
messages[msg_idx]["content"][content_idx_optional]["text"] = guardrail_response
content[content_idx_optional]["text"] = guardrail_response
async def _apply_guardrail_responses_to_input_tool_calls(
self,
messages: list[dict[str, Any]],
tool_calls: list[dict[str, Any]],
messages: Sequence[Mapping[str, object]],
tool_calls: Sequence[Mapping[str, object]],
task_mappings: list[tuple[int, int]],
) -> None:
"""
@ -412,7 +412,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
texts_to_check: Final[list[str]] = []
images_to_check: Final[list[str]] = []
tool_calls_to_check: Final[list[dict[str, Any]]] = []
tool_calls_to_check: Final[list[dict[str, object]]] = []
text_task_mappings: Final[list[tuple[int, int | None]]] = []
tool_call_task_mappings: Final[list[tuple[int, int]]] = []
# text_task_mappings: Track (choice_index, content_index) for each text
@ -461,8 +461,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
guardrailed_texts: Final = guardrailed_inputs.get("texts", [])
returned_tool_calls: Final = guardrailed_inputs.get("tool_calls")
guardrailed_tool_calls: Final[list[dict[str, Any]]] = (
cast(list[dict[str, Any]], returned_tool_calls)
guardrailed_tool_calls: Final[list[dict[str, object]]] = (
cast(list[dict[str, object]], returned_tool_calls)
if isinstance(returned_tool_calls, list) and len(returned_tool_calls) == len(tool_calls_to_check)
else tool_calls_to_check
)
@ -939,7 +939,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
choice_idx: int,
texts_to_check: list[str],
images_to_check: list[str],
tool_calls_to_check: list[dict[str, Any]],
tool_calls_to_check: list[dict[str, object]],
text_task_mappings: list[tuple[int, int | None]],
tool_call_task_mappings: list[tuple[int, int]],
) -> None:

6
litellm/llms/openai/common_utils.py
View file

@ -10,7 +10,7 @@ import ssl
import time
import uuid
from collections.abc import AsyncIterator, Iterator, Mapping
from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, Optional
from typing import TYPE_CHECKING, Final, Literal, NamedTuple, Optional
from urllib.parse import urlsplit
import httpx
@ -88,8 +88,8 @@ class OpenAIError(BaseLLMException):
###################################################################
def drop_params_from_unprocessable_entity_error(
e: openai.UnprocessableEntityError | httpx.HTTPStatusError,
data: dict[str, Any],
) -> dict[str, Any]:
data: Mapping[str, object],
) -> dict[str, object]:
"""
Helper function to read OpenAI UnprocessableEntityError and drop the params that raised an error from the error message.

61
litellm/llms/openai/openai.py
View file

@ -1,7 +1,7 @@
import time
import types
from collections.abc import AsyncIterator, Callable, Coroutine, Iterable, Iterator, Mapping
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast
from typing import TYPE_CHECKING, Final, Literal, Optional, cast
import httpx
@ -2756,7 +2756,12 @@ class OpenAIAssistantsAPI(BaseLLM):
message_thread: Final = await openai_client.beta.threads.create(**data)
return Thread(**message_thread.dict())
return Thread(
id=message_thread.id,
created_at=message_thread.created_at,
metadata=message_thread.metadata,
object=message_thread.object,
)
# fmt: off
@ -2842,7 +2847,12 @@ class OpenAIAssistantsAPI(BaseLLM):
message_thread: Final = openai_client.beta.threads.create(**data)
return Thread(**message_thread.dict())
return Thread(
id=message_thread.id,
created_at=message_thread.created_at,
metadata=message_thread.metadata,
object=message_thread.object,
)
async def async_get_thread(
self,
@ -2865,7 +2875,12 @@ class OpenAIAssistantsAPI(BaseLLM):
response: Final = await openai_client.beta.threads.retrieve(thread_id=thread_id)
return Thread(**response.dict())
return Thread(
id=response.id,
created_at=response.created_at,
metadata=response.metadata,
object=response.object,
)
# fmt: off
@ -2931,7 +2946,12 @@ class OpenAIAssistantsAPI(BaseLLM):
response: Final = openai_client.beta.threads.retrieve(thread_id=thread_id)
return Thread(**response.dict())
return Thread(
id=response.id,
created_at=response.created_at,
metadata=response.metadata,
object=response.object,
)
def delete_thread(self):
pass
@ -2988,18 +3008,27 @@ class OpenAIAssistantsAPI(BaseLLM):
tools: Iterable[AssistantToolParam] | None,
event_handler: AssistantEventHandler | None,
) -> AsyncAssistantStreamManager[AsyncAssistantEventHandler]:
data: Final[dict[str, Any]] = {
"thread_id": thread_id,
"assistant_id": assistant_id,
"additional_instructions": additional_instructions,
"instructions": instructions,
"metadata": metadata,
"model": model,
"tools": tools,
}
runs_stream: Final = client.beta.threads.runs.stream
if event_handler is not None:
data["event_handler"] = event_handler
return client.beta.threads.runs.stream(**data)
return runs_stream(
thread_id=thread_id,
assistant_id=assistant_id,
additional_instructions=additional_instructions,
instructions=instructions,
metadata=metadata,
model=model,
tools=tools,
event_handler=event_handler,
)
return runs_stream(
thread_id=thread_id,
assistant_id=assistant_id,
additional_instructions=additional_instructions,
instructions=instructions,
metadata=metadata,
model=model,
tools=tools,
)
def run_thread_stream(
self,

2
litellm/llms/openai/responses/guardrail_translation/handler.py
View file

@ -238,7 +238,7 @@ _TOOL_CALL_PAYLOAD_EVENT_TYPES: Final = _TOOL_CALL_PAYLOAD_DELTA_EVENT_TYPES | f
_OUTPUT_ITEM_EVENT_TYPES: Final = frozenset({"response.output_item.added", "response.output_item.done"})
_OUTPUT_TEXT_EVENT_TYPES: Final = frozenset({"response.output_text.delta", "response.output_text.done"})
_PATCHABLE_ITEM_FIELDS: Final[Mapping[str, str]] = MappingProxyType(
{"function_call_output": "output", "message": "content"}
{"function_call_output": "output", "custom_tool_call_output": "output", "message": "content"}
)
_EMPTY_RESPONSES_REQUEST: Final[ResponsesAPIOptionalRequestParams] = {}

6
litellm/llms/openai/videos/transformation.py
View file

@ -237,7 +237,7 @@ class OpenAIVideoConfig(BaseVideoConfig):
api_base: str,
litellm_params: GenericLiteLLMParams,
headers: dict,
extra_body: dict[str, Any] | None = None,
extra_body: dict[str, object] | None = None,
) -> tuple[str, dict]:
"""
Transform the video remix request for OpenAI API.
@ -252,7 +252,7 @@ class OpenAIVideoConfig(BaseVideoConfig):
url: Final = f"{api_base.rstrip('/')}/{encoded_video_id}/remix"
# Prepare the request data
data: Final = {"prompt": prompt}
data: Final[dict[str, object]] = {"prompt": prompt}
# Add any extra body parameters
if extra_body:
@ -305,7 +305,7 @@ class OpenAIVideoConfig(BaseVideoConfig):
after: str | None = None,
limit: int | None = None,
order: str | None = None,
extra_query: dict[str, Any] | None = None,
extra_query: dict[str, object] | None = None,
) -> tuple[str, dict]:
"""
Transform the video list request for OpenAI API.

11
litellm/llms/openrouter/image_edit/transformation.py
View file

@ -90,20 +90,21 @@ class OpenRouterImageEditConfig(BaseImageEditConfig):
drop_params: bool,
) -> dict:
supported_params: Final = self.get_supported_openai_params(model)
mapped_params: Final[dict[str, Any]] = {}
mapped_params: Final[dict[str, object]] = {}
image_config: Final[dict[str, str]] = {}
for key, value in image_edit_optional_params.items():
if key in supported_params:
if key == "size":
if "image_config" not in mapped_params:
mapped_params["image_config"] = {}
mapped_params["image_config"]["aspect_ratio"] = self._map_size_to_aspect_ratio(cast(str, value))
mapped_params["image_config"] = image_config
image_config["aspect_ratio"] = self._map_size_to_aspect_ratio(cast(str, value))
elif key == "quality":
image_size = self._map_quality_to_image_size(cast(str, value))
if image_size:
if "image_config" not in mapped_params:
mapped_params["image_config"] = {}
mapped_params["image_config"]["image_size"] = image_size
mapped_params["image_config"] = image_config
image_config["image_size"] = image_size
else:
mapped_params[key] = value

2
litellm/llms/perplexity/embedding/transformation.py
View file

@ -130,7 +130,7 @@ class PerplexityEmbeddingConfig(BaseEmbeddingConfig):
if isinstance(embedding_value, str):
raw_bytes: Final = base64.b64decode(embedding_value)
count: Final = len(raw_bytes)
int8_values: Final = struct.unpack(f"{count}b", raw_bytes)
int8_values: Final[tuple[int, ...]] = struct.unpack(f"{count}b", raw_bytes)
return [float(v) / 127.0 for v in int8_values]
return embedding_value

4
litellm/llms/snowflake/chat/transformation.py
View file

@ -315,7 +315,7 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
for msg in messages:
if isinstance(msg, dict):
role = msg.get("role", "")
content: Any = msg.get("content", "")
content: object = msg.get("content", "")
msg_cache_control: object = msg.get("cache_control")
else:
role = getattr(msg, "role", "")
@ -463,7 +463,7 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
return body
def _transform_tool_choice_to_anthropic(self, tool_choice: Any) -> dict[str, Any]:
def _transform_tool_choice_to_anthropic(self, tool_choice: object) -> Mapping[str, object]:
"""
Convert tool_choice from OpenAI format to Anthropic format.

4
litellm/llms/stability/image_edit/transformations.py
View file

@ -74,7 +74,7 @@ class StabilityImageEditConfig(BaseImageEditConfig):
}
# Create a copy to not mutate original - convert TypedDict to regular dict
mapped_params: Final[dict[str, Any]] = dict(image_edit_optional_params)
mapped_params: Final[dict[str, object]] = dict(image_edit_optional_params)
for k, v in image_edit_optional_params.items():
if k in param_mapping:
@ -182,7 +182,7 @@ class StabilityImageEditConfig(BaseImageEditConfig):
# Build Stability request
# Populate multipart form-data as separate text fields (data) and files.
# Stability expects prompt/output_format/etc. as normal form fields, not file parts.
data: Final[dict[str, Any]] = {
data: Final[dict[str, object]] = {
"output_format": "png", # Default to PNG
}

8
litellm/llms/triton/completion/transformation.py
View file

@ -4,7 +4,7 @@ Translates from OpenAI's `/v1/chat/completions` endpoint to Triton's `/generate`
import json
from collections.abc import AsyncIterator, Iterator
from typing import TYPE_CHECKING, Any, Final, Literal
from typing import TYPE_CHECKING, Final, Literal
from httpx import Headers, Response
@ -172,7 +172,7 @@ class TritonConfig(BaseConfig):
streaming_response: Iterator[str] | AsyncIterator[str] | ModelResponse,
sync_stream: bool,
json_mode: bool | None = False,
) -> Any:
) -> "TritonResponseIterator":
return TritonResponseIterator(
streaming_response=streaming_response,
sync_stream=sync_stream,
@ -195,14 +195,14 @@ class TritonGenerateConfig(TritonConfig):
) -> dict:
inference_params: Final = optional_params.copy()
stream: Final = inference_params.pop("stream", False)
data_for_triton: Final[dict[str, Any]] = {
data_for_triton: Final[dict[str, object]] = {
"text_input": prompt_factory(model=model, messages=messages),
"parameters": {
"max_tokens": int(optional_params.get("max_tokens", DEFAULT_MAX_TOKENS_FOR_TRITON)),
**inference_params,
},
"stream": bool(stream),
}
data_for_triton["parameters"].update(inference_params)
return data_for_triton
def transform_response(

5
litellm/llms/vertex_ai/fine_tuning/handler.py
View file

@ -280,7 +280,7 @@ class VertexFineTuningAPI(VertexLLM):
vertex_location: str,
vertex_credentials: str,
request_route: str,
):
) -> object:
_auth_header, vertex_project = await self._ensure_access_token_async(
credentials=vertex_credentials,
project_id=vertex_project,
@ -341,5 +341,4 @@ class VertexFineTuningAPI(VertexLLM):
f"Error creating fine tuning job. Status code: {response.status_code}. Response: {response.text}"
)
response_json: Final = response.json()
return response_json
return response.json()

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