scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-25 01:02:15 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 80

Merge remote-tracking branch 'origin/main' into litellm_mcp_ui_prompts_resources

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

# Conflicts:
#	litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
This commit is contained in:
joshua 2026-09-21 22:41:11 +00:00
commit 779e4bd433
257 changed files with 9948 additions and 4425 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

52
.circleci/config.yml
View file

@ -6,6 +6,9 @@ parameters:
migration_candidate_image:
type: string
default: ""
migration_baseline_image:
type: string
default: "ghcr.io/berriai/litellm-database:v1.102.0"
migration_source_sha:
type: string
default: ""
@ -2946,7 +2949,10 @@ jobs:
parameters:
suite:
type: enum
enum: [startup, recovery, legacy]
enum: [startup, recovery, legacy, upgrade, shaped]
baseline:
type: boolean
default: false
machine:
image: ubuntu-2204:2024.04.1
resource_class: large
@ -2954,6 +2960,7 @@ jobs:
environment:
LITELLM_MIGRATION_TESTS: "1"
LITELLM_MIGRATION_TEST_IMAGE: litellm-docker-database:ci
LITELLM_MIGRATION_BASELINE_IMAGE: << pipeline.parameters.migration_baseline_image >>
MIGRATION_TEST_ADMIN_URL: postgresql://postgres:postgres@127.0.0.1:5432/postgres
MIGRATION_TEST_CONTAINER_ADMIN_URL: postgresql://postgres:postgres@host.docker.internal:5432/postgres
MIGRATION_TEST_OUTPUT: /tmp/migration-results
@ -2981,6 +2988,16 @@ jobs:
- wait_for_service:
url: tcp://localhost:5432
timeout: "60"
- when:
condition: << parameters.baseline >>
steps:
- run:
name: Pull the baseline release the upgrade starts from
environment:
BASELINE_IMAGE: << pipeline.parameters.migration_baseline_image >>
command: |
[[ "$BASELINE_IMAGE" =~ ^ghcr.io/berriai/[a-z0-9._/-]+(@sha256:[0-9a-f]{64}|:v[0-9][0-9a-z.-]*)$ ]] || exit 1
docker pull "$BASELINE_IMAGE"
- run:
name: Run migration startup regressions
environment:
@ -3033,28 +3050,29 @@ jobs:
- run:
name: Run Docker container with bad DATABASE_URL
command: |
set +e
docker run --name my-app \
-p 4000:4000 \
-e LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY=true \
-e DEFAULT_NUM_WORKERS_LITELLM_PROXY=1 \
-e DATABASE_URL="postgresql://wrong:wrong@wrong:5432/wrong" \
myapp:latest \
--port 4000 > docker_output.log 2>&1 || true
--port 4000 > docker_output.log 2>&1
echo "$?" > docker_exit_code
set -e
- run:
name: Display Docker logs
command: cat docker_output.log
- run:
name: Check for expected error
name: Proxy must refuse to serve on an unreachable database
command: |
if grep -q "Error: P1001: Can't reach database server at" docker_output.log && \
(grep -q "Database setup failed after multiple retries" docker_output.log || \
grep -q "ERROR: Application startup failed. Exiting." docker_output.log); then
echo "Expected error found. Test passed."
else
echo "Expected error not found. Test failed."
cat docker_output.log
exit 1
fi
fail() { echo "FAILED: $1"; cat docker_output.log; exit 1; }
exit_code="$(cat docker_exit_code)"
[ "$exit_code" -ne 0 ] || fail "proxy exited 0 with an unreachable database"
grep -q "P1001" docker_output.log || fail "log does not name the unreachable database server"
! grep -q "Application startup complete" docker_output.log || fail "proxy reached serving state"
! docker exec my-app true 2>/dev/null || fail "container is still running"
echo "Proxy refused to serve (exit $exit_code) and never reached startup. Test passed."
provider_replay_harness:
docker:
@ -3188,6 +3206,16 @@ workflows:
name: migration-legacy-and-pooling
suite: legacy
requires: [build_docker_database_image]
- migration_startup_tests:
name: migration-upgrade
suite: upgrade
baseline: true
requires: [build_docker_database_image]
- migration_startup_tests:
name: migration-upgrade-shaped
suite: shaped
baseline: true
requires: [build_docker_database_image]
migration_startup_scheduled:
triggers:
- schedule:

3
.circleci/scripts/run_migration_tests.py
View file

@ -13,6 +13,8 @@ SUITES: Final = {
"startup": (("test_startup.py",), 12),
"recovery": (("test_recovery.py",), 15),
"legacy": (("test_legacy.py", "test_pooling.py"), 11),
"upgrade": (("test_upgrade.py", "test_rolling_upgrade.py"), 5),
"shaped": (("test_shaped_database.py",), 1),
}
@ -93,6 +95,7 @@ def main() -> int:
{
**metadata,
"suite": suite,
"baseline_image": os.environ.get("LITELLM_MIGRATION_BASELINE_IMAGE", ""),
"expected_cases": expected,
"passed": passed,
"pytest_exit_code": result.returncode,

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: |

2
.github/workflows/test-rust.yml vendored
View file

@ -130,7 +130,7 @@ jobs:
- name: Test secret manager feature combinations
run: |
cargo test -p litellm-auth-gcp --locked --no-default-features
for features in '' aws google cyberark aws,google aws,google,cyberark; do
for features in '' aws google azure cyberark aws,google aws,azure google,azure aws,google,azure aws,google,cyberark aws,google,azure,cyberark; do
cargo test -p litellm-secrets --locked --no-default-features --features "$features"
done

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

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

103
litellm-rust/Cargo.lock generated
View file

@ -115,6 +115,28 @@ dependencies = [
"pin-project-lite",
]
[[package]]
name = "async-stream"
version = "0.3.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0b5a71a6f37880a80d1d7f19efd781e4b5de42c88f0722cc13bcb6cc2cfe8476"
dependencies = [
"async-stream-impl",
"futures-core",
"pin-project-lite",
]
[[package]]
name = "async-stream-impl"
version = "0.3.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c7c24de15d275a1ecfd47a380fb4d5ec9bfe0933f309ed5e705b775596a3574d"
dependencies = [
"proc-macro2",
"quote",
"syn 2.0.119",
]
[[package]]
name = "async-trait"
version = "0.1.91"
@ -599,6 +621,37 @@ dependencies = [
"url",
]
[[package]]
name = "azure_storage_blob"
version = "1.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "17b10207ecf7d666df6940b50051f433b3cd5d2b9b1dd190613208d7a84e7eed"
dependencies = [
"async-stream",
"async-trait",
"azure_core",
"azure_storage_common",
"bytes",
"futures",
"percent-encoding",
"pin-project",
"serde",
"serde_json",
"time",
"tokio",
]
[[package]]
name = "azure_storage_common"
version = "1.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0af2e6aeb8d76b17fc998f453c320913f73787b944e3cc29509d19411fa0321d"
dependencies = [
"azure_core",
"serde",
"time",
]
[[package]]
name = "base64"
version = "0.13.1"
@ -2470,6 +2523,23 @@ dependencies = [
"tokio",
]
[[package]]
name = "litellm-cache-azure-blob"
version = "0.1.0"
dependencies = [
"async-trait",
"azure_core",
"azure_storage_blob",
"futures-util",
"litellm-auth-azure",
"litellm-auth-types",
"litellm-cache",
"litellm-cache-response",
"serde_json",
"tokio",
"url",
]
[[package]]
name = "litellm-cache-memory"
version = "0.1.0"
@ -2672,6 +2742,7 @@ dependencies = [
"litellm-auth",
"litellm-auth-gcp",
"litellm-cache",
"litellm-cache-azure-blob",
"litellm-cache-memory",
"litellm-cache-redis",
"litellm-cache-response",
@ -2704,6 +2775,7 @@ dependencies = [
"jsonwebtoken",
"litellm-core-utils",
"litellm-secrets-aws",
"litellm-secrets-azure",
"litellm-secrets-cyberark",
"litellm-secrets-google",
"litellm-secrets-types",
@ -2739,6 +2811,26 @@ dependencies = [
"wiremock",
]
[[package]]
name = "litellm-secrets-azure"
version = "0.1.0"
dependencies = [
"litellm-auth-azure",
"litellm-auth-types",
"litellm-core-utils",
"litellm-secrets-types",
"percent-encoding",
"reqwest 0.12.28",
"rstest",
"serde",
"serde_json",
"sha2 0.10.9",
"thiserror 2.0.19",
"tokio",
"veil",
"wiremock",
]
[[package]]
name = "litellm-secrets-cyberark"
version = "0.1.0"
@ -3515,6 +3607,16 @@ version = "1.2.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a1d01941d82fa2ab50be1e79e6714289dd7cde78eba4c074bc5a4374f650dfe0"
[[package]]
name = "quick-xml"
version = "0.41.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e660451e55124f798a69a5af3f49ccfbefbd41910eefd25caf2393e1f3473ec1"
dependencies = [
"memchr",
"serde",
]
[[package]]
name = "quinn"
version = "0.11.11"
@ -5060,6 +5162,7 @@ dependencies = [
"base64 0.22.1",
"bytes",
"futures",
"quick-xml",
"serde",
"serde_json",
"url",

2
litellm-rust/Cargo.toml
View file

@ -22,12 +22,14 @@ litellm-secrets = { path = "crates/secrets" }
litellm-secrets-types = { path = "crates/secrets-types" }
litellm-secrets-aws = { path = "crates/secrets-aws" }
litellm-secrets-google = { path = "crates/secrets-google" }
litellm-secrets-azure = { path = "crates/secrets-azure" }
litellm-secrets-cyberark = { path = "crates/secrets-cyberark" }
litellm-http = { path = "crates/http" }
litellm-llms = { path = "crates/llms" }
litellm-types = { path = "crates/types" }
litellm-core-utils = { path = "crates/core-utils" }
litellm-cache = { path = "crates/cache" }
litellm-cache-azure-blob = { path = "crates/cache-azure-blob" }
litellm-cache-memory = { path = "crates/cache-memory" }
litellm-cache-redis = { path = "crates/cache-redis" }
litellm-cache-response = { path = "crates/cache-response" }

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

@ -4,4 +4,4 @@ mod resolve;
mod types;
pub use resolve::AzureAuthService;
pub use types::AzureAuthInputs;
pub use types::{AzureAuthInputs, ConfigValue};

15
litellm-rust/crates/auth-azure/src/types.rs
View file

@ -51,6 +51,21 @@ pub struct AzureAuthInputs {
}
impl AzureAuthInputs {
pub fn default_credential_for_scope(scope: &str) -> Self {
Self {
azure_scope: ConfigValue::Value(Sourced::new(
scope.to_string(),
InputSource::Deployment,
)),
azure_credential: ConfigValue::Value(Sourced::new(
"DefaultAzureCredential".to_string(),
InputSource::Deployment,
)),
enable_azure_ad_token_refresh: Sourced::new(true, InputSource::Deployment),
..Self::default()
}
}
pub fn or_configured_token_refresh(self, enabled: bool) -> Self {
if *self.enable_azure_ad_token_refresh.value() || !enabled {
return self;

22
litellm-rust/crates/cache-azure-blob/Cargo.toml Normal file
View file

@ -0,0 +1,22 @@
[package]
name = "litellm-cache-azure-blob"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
litellm-auth-azure.workspace = true
litellm-auth-types.workspace = true
litellm-cache.workspace = true
async-trait = "0.1"
azure_core = "1.1.0"
azure_storage_blob = "1.1.0"
futures-util.workspace = true
tokio.workspace = true
url.workspace = true
[dev-dependencies]
litellm-cache-response.workspace = true
serde_json.workspace = true

254
litellm-rust/crates/cache-azure-blob/src/cache.rs Normal file
View file

@ -0,0 +1,254 @@
use std::{sync::Arc, time::Duration};
use azure_core::{
credentials::TokenCredential,
error::ErrorKind,
http::{ClientOptions, RequestContent},
};
use azure_storage_blob::{
BlobContainerClient, BlobContainerClientOptions,
models::{BlobClientUploadOptions, StorageErrorCode},
};
use futures_util::{TryStreamExt, future::try_join_all};
use litellm_cache::{
BaseCache, BatchCache, CacheCodec, CacheConnectionResult, CacheConnectionStatus, Error,
ExactCacheContext, FlushCache,
};
use tokio::runtime::Handle;
use url::Url;
use crate::credential::AzureBlobCredential;
pub struct AzureBlobCache<C> {
container: BlobContainerClient,
codec: C,
runtime: Handle,
account_url: String,
container_name: String,
}
impl<C: CacheCodec> AzureBlobCache<C> {
pub async fn connect(
account_url: &str,
container: &str,
codec: C,
runtime: Handle,
) -> Result<Self, Error> {
Self::connect_with_options(
account_url,
container,
Some(Arc::new(AzureBlobCredential::default())),
ClientOptions::default(),
codec,
runtime,
)
.await
}
pub async fn connect_with_options(
account_url: &str,
container: &str,
credential: Option<Arc<dyn TokenCredential>>,
client_options: ClientOptions,
codec: C,
runtime: Handle,
) -> Result<Self, Error> {
let parsed = Url::parse(account_url).map_err(|_| Error::Unavailable)?;
let account_url = parsed.as_str().trim_end_matches('/').to_string();
let container_url = {
let mut url = parsed;
url.path_segments_mut()
.map_err(|()| Error::Unavailable)?
.pop_if_empty()
.push(container);
url
};
let client = BlobContainerClient::new(
container_url,
credential,
Some(BlobContainerClientOptions {
client_options,
..BlobContainerClientOptions::default()
}),
)
.map_err(|_| Error::Unavailable)?;
let cache = Self {
container: client,
codec,
runtime,
account_url,
container_name: container.to_string(),
};
cache.create_container().await?;
Ok(cache)
}
pub fn account_url(&self) -> &str {
&self.account_url
}
pub fn container_name(&self) -> &str {
&self.container_name
}
async fn create_container(&self) -> Result<(), Error> {
match self.container.create(None).await {
Ok(_) => Ok(()),
Err(error) if is_storage_error(&error, StorageErrorCode::ContainerAlreadyExists) => {
Ok(())
}
Err(_) => Err(Error::Unavailable),
}
}
async fn upload(&self, key: &str, value: &C::Value, overwrite: bool) -> Result<(), Error> {
let payload = self.codec.encode(value)?;
let options = (!overwrite).then(|| BlobClientUploadOptions::default().if_not_exists());
match self
.container
.blob_client(key)
.upload(RequestContent::from(payload), options)
.await
{
Ok(_) => Ok(()),
Err(error) if !overwrite && is_already_present(&error) => Ok(()),
Err(_) => Err(Error::Unavailable),
}
}
async fn download(&self, key: &str) -> Result<Option<C::Value>, Error> {
let response = match self.container.blob_client(key).download(None).await {
Ok(response) => response,
Err(error) if is_storage_error(&error, StorageErrorCode::BlobNotFound) => {
return Ok(None);
}
Err(_) => return Err(Error::Unavailable),
};
let bytes = response
.body
.collect()
.await
.map_err(|_| Error::Unavailable)?;
self.codec.decode(&bytes).map(Some)
}
async fn delete_all_blobs(&self) -> Result<(), Error> {
let mut pages = self
.container
.list_blobs(None)
.map_err(|_| Error::Unavailable)?
.into_pages();
while let Some(page) = pages.try_next().await.map_err(|_| Error::Unavailable)? {
let page = page.into_model().map_err(|_| Error::Unavailable)?;
for name in page.blob_items.into_iter().filter_map(|item| item.name) {
self.container
.blob_client(&name)
.delete(None)
.await
.map_err(|_| Error::Unavailable)?;
}
}
Ok(())
}
fn block_on<T>(&self, future: impl Future<Output = T>) -> T {
self.runtime.block_on(future)
}
}
fn is_already_present(error: &azure_core::Error) -> bool {
is_storage_error(error, StorageErrorCode::BlobAlreadyExists)
|| is_storage_error(error, StorageErrorCode::ConditionNotMet)
}
fn is_storage_error(error: &azure_core::Error, code: StorageErrorCode) -> bool {
matches!(
error.kind(),
ErrorKind::HttpResponse {
error_code: Some(error_code),
..
} if error_code == code.as_ref()
)
}
impl<C: CacheCodec> BaseCache for AzureBlobCache<C> {
type Value = C::Value;
type Context = ExactCacheContext;
fn get_ttl(&self, _: &ExactCacheContext) -> Option<Duration> {
None
}
fn set_cache(&self, key: &str, value: C::Value, _: &ExactCacheContext) -> Result<(), Error> {
self.block_on(self.upload(key, &value, false))
}
fn get_cache(&self, key: &str, _: &ExactCacheContext) -> Result<Option<C::Value>, Error> {
self.block_on(self.download(key))
}
async fn async_set_cache(
&self,
key: &str,
value: C::Value,
_: ExactCacheContext,
) -> Result<(), Error> {
self.upload(key, &value, true).await
}
async fn async_get_cache(
&self,
key: &str,
_: &ExactCacheContext,
) -> Result<Option<C::Value>, Error> {
self.download(key).await
}
async fn async_set_cache_pipeline(
&self,
entries: Vec<(String, C::Value)>,
_: ExactCacheContext,
) -> Result<(), Error> {
try_join_all(
entries
.iter()
.map(|(key, value)| self.upload(key, value, true)),
)
.await
.map(drop)
}
async fn disconnect(&self) -> Result<(), Error> {
Ok(())
}
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
Ok(match self.container.get_properties(None).await {
Ok(_) => CacheConnectionResult {
status: CacheConnectionStatus::Success,
message: "Azure Blob cache connection test successful".into(),
error: None,
},
Err(error) => CacheConnectionResult {
status: CacheConnectionStatus::Failed,
message: format!("Azure Blob connection failed: {error}"),
error: Some(error.to_string()),
},
})
}
}
impl<C: CacheCodec> BatchCache for AzureBlobCache<C> {}
impl<C: CacheCodec> FlushCache for AzureBlobCache<C> {
fn flush_cache(&self) -> Result<(), Error> {
self.block_on(self.delete_all_blobs())
}
async fn async_flush_cache(&self) -> Result<(), Error> {
self.delete_all_blobs().await
}
}
#[cfg(test)]
mod tests;

746
litellm-rust/crates/cache-azure-blob/src/cache/tests.rs vendored Normal file
View file

@ -0,0 +1,746 @@
use std::{
collections::BTreeMap,
sync::{Arc, Mutex},
time::Duration,
};
use azure_core::http::{
AsyncRawResponse, Body, ClientOptions, HttpClient, Method, Request, StatusCode, Transport,
headers::{HeaderName, Headers},
};
use litellm_cache::{
BaseCache, BatchCache, BatchEntry, CacheConnectionStatus, Error, ExactCacheContext, FlushCache,
};
use litellm_cache_response::{
CacheEntry, CacheKeyField, CacheKeyInput, ResponseCache, ResponseCacheCodec,
ResponseCacheRequest, cache_key,
};
use serde_json::json;
use tokio::runtime::Runtime;
use super::AzureBlobCache;
const ACCOUNT_URL: &str = "https://example.blob.core.windows.net";
const CONTAINER: &str = "litellm-cache";
const IF_NONE_MATCH: HeaderName = HeaderName::from_static("if-none-match");
const ERROR_CODE: HeaderName = HeaderName::from_static("x-ms-error-code");
#[derive(Clone, Debug, PartialEq, Eq)]
struct RecordedRequest {
method: Method,
path: String,
query: String,
if_none_match: Option<String>,
}
#[derive(Default)]
struct FakeState {
container_exists: bool,
blobs: BTreeMap<String, Vec<u8>>,
requests: Vec<RecordedRequest>,
failing: bool,
precondition_conflicts: bool,
}
#[derive(Clone, Default)]
struct FakeBlobService {
state: Arc<Mutex<FakeState>>,
}
impl std::fmt::Debug for FakeBlobService {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str("FakeBlobService")
}
}
impl FakeBlobService {
fn with_existing_container() -> Self {
let service = Self::default();
service.state.lock().unwrap().container_exists = true;
service
}
fn blob(&self, name: &str) -> Option<Vec<u8>> {
self.state.lock().unwrap().blobs.get(name).cloned()
}
fn blob_names(&self) -> Vec<String> {
self.state.lock().unwrap().blobs.keys().cloned().collect()
}
fn seed_blob(&self, name: &str, bytes: &[u8]) {
self.state
.lock()
.unwrap()
.blobs
.insert(name.to_string(), bytes.to_vec());
}
fn set_failing(&self, failing: bool) {
self.state.lock().unwrap().failing = failing;
}
fn set_precondition_conflicts(&self, enabled: bool) {
self.state.lock().unwrap().precondition_conflicts = enabled;
}
fn requests(&self) -> Vec<RecordedRequest> {
self.state.lock().unwrap().requests.clone()
}
fn container_exists(&self) -> bool {
self.state.lock().unwrap().container_exists
}
fn respond(status: StatusCode, error_code: Option<&str>, body: Vec<u8>) -> AsyncRawResponse {
let mut headers = Headers::new();
if let Some(code) = error_code {
headers.insert(ERROR_CODE, code.to_string());
}
AsyncRawResponse::from_bytes(status, headers, body)
}
fn list_body(state: &FakeState) -> Vec<u8> {
let mut xml = String::from(
r#"<?xml version="1.0" encoding="utf-8"?><EnumerationResults ServiceEndpoint="https://example.blob.core.windows.net/" ContainerName="litellm-cache"><Blobs>"#,
);
for name in state.blobs.keys() {
xml.push_str(&format!(
"<Blob><Name>{name}</Name><Properties><BlobType>BlockBlob</BlobType></Properties></Blob>"
));
}
xml.push_str("</Blobs><NextMarker /></EnumerationResults>");
xml.into_bytes()
}
}
#[async_trait::async_trait]
impl HttpClient for FakeBlobService {
async fn execute_request(&self, request: &Request) -> azure_core::Result<AsyncRawResponse> {
let mut state = self.state.lock().unwrap();
let path = request.url().path().to_string();
let query = request.url().query().unwrap_or_default().to_string();
let if_none_match = request
.headers()
.get_optional_str(&IF_NONE_MATCH)
.map(str::to_owned);
state.requests.push(RecordedRequest {
method: request.method(),
path: path.clone(),
query: query.clone(),
if_none_match: if_none_match.clone(),
});
if state.failing {
return Ok(Self::respond(
StatusCode::Forbidden,
Some("AuthorizationFailure"),
Vec::new(),
));
}
let container_path = format!("/{CONTAINER}");
let blob_name = path
.strip_prefix(&format!("{container_path}/"))
.map(str::to_owned);
let is_container = path == container_path && query.contains("restype=container");
let response = match (request.method(), is_container, blob_name) {
(Method::Put, true, None) if state.container_exists => Self::respond(
StatusCode::Conflict,
Some("ContainerAlreadyExists"),
Vec::new(),
),
(Method::Put, true, None) => {
state.container_exists = true;
Self::respond(StatusCode::Created, None, Vec::new())
}
(Method::Get, true, None) if query.contains("comp=list") => {
Self::respond(StatusCode::Ok, None, Self::list_body(&state))
}
(Method::Get, true, None) if state.container_exists => {
Self::respond(StatusCode::Ok, None, Vec::new())
}
(Method::Get, true, None) => {
Self::respond(StatusCode::NotFound, Some("ContainerNotFound"), Vec::new())
}
(Method::Put, false, Some(name)) => {
if if_none_match.as_deref() == Some("*") && state.blobs.contains_key(&name) {
if state.precondition_conflicts {
Self::respond(
StatusCode::PreconditionFailed,
Some("ConditionNotMet"),
Vec::new(),
)
} else {
Self::respond(StatusCode::Conflict, Some("BlobAlreadyExists"), Vec::new())
}
} else {
let bytes = match request.body() {
Body::Bytes(bytes) => bytes.to_vec(),
Body::SeekableStream(_) => panic!("unexpected streaming upload"),
};
state.blobs.insert(name, bytes);
Self::respond(StatusCode::Created, None, Vec::new())
}
}
(Method::Get, false, Some(name)) => match state.blobs.get(&name) {
Some(bytes) => Self::respond(StatusCode::Ok, None, bytes.clone()),
None => Self::respond(StatusCode::NotFound, Some("BlobNotFound"), Vec::new()),
},
(Method::Delete, false, Some(name)) => match state.blobs.remove(&name) {
Some(_) => Self::respond(StatusCode::Accepted, None, Vec::new()),
None => Self::respond(StatusCode::NotFound, Some("BlobNotFound"), Vec::new()),
},
(method, _, _) => panic!("unexpected request {method:?} {path}?{query}"),
};
Ok(response)
}
}
struct Fixture {
runtime: Runtime,
service: FakeBlobService,
cache: Arc<AzureBlobCache<ResponseCacheCodec>>,
}
impl Fixture {
fn new(service: FakeBlobService) -> Self {
let runtime = Runtime::new().unwrap();
let cache = runtime
.block_on(Self::connect(&service, runtime.handle().clone()))
.unwrap();
Self {
runtime,
service,
cache: Arc::new(cache),
}
}
async fn connect(
service: &FakeBlobService,
handle: tokio::runtime::Handle,
) -> Result<AzureBlobCache<ResponseCacheCodec>, Error> {
AzureBlobCache::connect_with_options(
ACCOUNT_URL,
CONTAINER,
None,
ClientOptions {
transport: Some(Transport::new(Arc::new(service.clone()))),
..ClientOptions::default()
},
ResponseCacheCodec,
handle,
)
.await
}
fn response_cache(&self) -> ResponseCache<AzureBlobCache<ResponseCacheCodec>> {
ResponseCache::new(self.cache.clone())
}
fn stored_json(&self, key: &str) -> serde_json::Value {
serde_json::from_slice(&self.service.blob(key).expect("blob should exist")).unwrap()
}
}
fn request(model: &str) -> ResponseCacheRequest {
ResponseCacheRequest::new(CacheKeyInput {
fields: vec![CacheKeyField {
name: "model".into(),
value: Some(model.into()),
api_parameter: true,
internal_parameter: false,
}],
preset: None,
namespace: None,
include_provider_parameters: false,
})
}
fn now() -> Duration {
Duration::from_secs(1_700_000_000)
}
fn entry(value: serde_json::Value) -> CacheEntry {
CacheEntry {
timestamp: Some(1_700_000_000.5),
response: value,
}
}
fn no_ttl() -> ExactCacheContext {
ExactCacheContext::default()
}
fn with_ttl(seconds: u64) -> ExactCacheContext {
ExactCacheContext {
ttl: Some(Duration::from_secs(seconds)),
}
}
#[test]
fn connect_creates_the_container_once() {
let fixture = Fixture::new(FakeBlobService::default());
assert!(fixture.service.container_exists());
assert_eq!(
fixture.service.requests(),
vec![RecordedRequest {
method: Method::Put,
path: format!("/{CONTAINER}"),
query: "restype=container".into(),
if_none_match: None,
}]
);
assert_eq!(fixture.cache.account_url(), ACCOUNT_URL);
assert_eq!(fixture.cache.container_name(), CONTAINER);
}
#[test]
fn connect_accepts_an_existing_container() {
let fixture = Fixture::new(FakeBlobService::with_existing_container());
assert!(fixture.service.container_exists());
assert_eq!(fixture.service.requests().len(), 1);
}
#[test]
fn connect_accepts_account_urls_with_trailing_slash() {
let runtime = Runtime::new().unwrap();
let service = FakeBlobService::default();
let cache = runtime
.block_on(AzureBlobCache::connect_with_options(
"https://example.blob.core.windows.net/",
CONTAINER,
None,
ClientOptions {
transport: Some(Transport::new(Arc::new(service.clone()))),
..ClientOptions::default()
},
ResponseCacheCodec,
runtime.handle().clone(),
))
.unwrap();
assert_eq!(service.requests()[0].path, format!("/{CONTAINER}"));
assert_eq!(cache.account_url(), "https://example.blob.core.windows.net");
}
#[test]
fn connect_keeps_account_url_query_parameters_on_the_container_path() {
let runtime = Runtime::new().unwrap();
let service = FakeBlobService::default();
runtime
.block_on(AzureBlobCache::connect_with_options(
"https://example.blob.core.windows.net/?sv=2024-01-01&sig=abc",
CONTAINER,
None,
ClientOptions {
transport: Some(Transport::new(Arc::new(service.clone()))),
..ClientOptions::default()
},
ResponseCacheCodec,
runtime.handle().clone(),
))
.unwrap();
let create = &service.requests()[0];
assert_eq!(create.path, format!("/{CONTAINER}"));
assert!(create.query.contains("sig=abc"));
}
#[test]
fn connect_surfaces_service_failures() {
let runtime = Runtime::new().unwrap();
let service = FakeBlobService::default();
service.set_failing(true);
let result = runtime.block_on(Fixture::connect(&service, runtime.handle().clone()));
assert!(matches!(result, Err(Error::Unavailable)));
}
#[test]
fn sync_set_and_get_round_trip_python_json_shape() {
let fixture = Fixture::new(FakeBlobService::default());
let value = entry(json!({"choices": [{"message": {"content": "héllo 🌍"}}]}));
fixture
.cache
.set_cache("key-1", value.clone(), &no_ttl())
.unwrap();
assert_eq!(
fixture.stored_json("key-1"),
json!({
"timestamp": 1_700_000_000.5,
"response": {"choices": [{"message": {"content": "héllo 🌍"}}]}
})
);
assert_eq!(
fixture.cache.get_cache("key-1", &no_ttl()).unwrap(),
Some(value)
);
}
#[test]
fn sync_set_does_not_overwrite_an_existing_blob() {
let fixture = Fixture::new(FakeBlobService::default());
fixture
.cache
.set_cache("key", entry(json!({"v": "first"})), &no_ttl())
.unwrap();
fixture
.cache
.set_cache("key", entry(json!({"v": "second"})), &no_ttl())
.unwrap();
assert_eq!(
fixture.stored_json("key")["response"],
json!({"v": "first"})
);
let uploads: Vec<_> = fixture
.service
.requests()
.into_iter()
.filter(|request| request.method == Method::Put && request.path.ends_with("/key"))
.collect();
assert_eq!(uploads.len(), 2);
assert!(
uploads
.iter()
.all(|request| request.if_none_match.as_deref() == Some("*"))
);
}
#[test]
fn sync_set_treats_a_precondition_conflict_as_an_existing_blob() {
let fixture = Fixture::new(FakeBlobService::default());
fixture.service.set_precondition_conflicts(true);
fixture
.cache
.set_cache("key", entry(json!({"v": "first"})), &no_ttl())
.unwrap();
fixture
.cache
.set_cache("key", entry(json!({"v": "second"})), &no_ttl())
.unwrap();
assert_eq!(
fixture.stored_json("key")["response"],
json!({"v": "first"})
);
}
#[test]
fn async_set_overwrites_an_existing_blob() {
let fixture = Fixture::new(FakeBlobService::default());
fixture.runtime.block_on(async {
fixture
.cache
.async_set_cache("key", entry(json!({"v": "first"})), no_ttl())
.await
.unwrap();
fixture
.cache
.async_set_cache("key", entry(json!({"v": "second"})), no_ttl())
.await
.unwrap();
assert_eq!(
fixture
.cache
.async_get_cache("key", &no_ttl())
.await
.unwrap(),
Some(entry(json!({"v": "second"})))
);
});
assert_eq!(
fixture.stored_json("key")["response"],
json!({"v": "second"})
);
assert!(
fixture
.service
.requests()
.iter()
.filter(|request| request.method == Method::Put && request.path.ends_with("/key"))
.all(|request| request.if_none_match.is_none())
);
}
#[test]
fn missing_blobs_are_misses() {
let fixture = Fixture::new(FakeBlobService::default());
assert_eq!(fixture.cache.get_cache("absent", &no_ttl()).unwrap(), None);
assert_eq!(
fixture
.runtime
.block_on(fixture.cache.async_get_cache("absent", &no_ttl()))
.unwrap(),
None
);
}
#[test]
fn ttl_is_ignored_and_entries_never_expire() {
let fixture = Fixture::new(FakeBlobService::default());
assert_eq!(fixture.cache.get_ttl(&with_ttl(1)), None);
assert_eq!(fixture.cache.get_ttl(&no_ttl()), None);
fixture
.cache
.set_cache("key", entry(json!("value")), &with_ttl(1))
.unwrap();
std::thread::sleep(Duration::from_millis(1100));
assert_eq!(
fixture.cache.get_cache("key", &with_ttl(1)).unwrap(),
Some(entry(json!("value")))
);
assert!(
fixture
.service
.requests()
.iter()
.all(|request| !request.query.contains("expiry"))
);
}
#[test]
fn malformed_blobs_are_invalid_entries_and_response_cache_misses() {
let fixture = Fixture::new(FakeBlobService::default());
fixture.service.seed_blob("broken-json", b"{not json");
fixture
.service
.seed_blob("broken-utf8", &[0xff, 0xfe, 0x22]);
fixture
.service
.seed_blob("wrong-shape", br#"{"timestamp": "yesterday"}"#);
for key in ["broken-json", "broken-utf8", "wrong-shape"] {
assert!(matches!(
fixture.cache.get_cache(key, &no_ttl()),
Err(Error::InvalidEntry)
));
}
let response_cache = fixture.response_cache();
let broken = request("broken");
fixture
.service
.seed_blob(&cache_key(&broken.key), b"{not json");
assert_eq!(response_cache.lookup(&broken, now()).unwrap(), None);
assert_eq!(
fixture
.runtime
.block_on(response_cache.async_lookup(&broken, now()))
.unwrap(),
None
);
}
#[test]
fn batch_get_preserves_order_and_marks_misses_and_invalid_entries() {
let fixture = Fixture::new(FakeBlobService::default());
fixture
.cache
.set_cache("a", entry(json!("A")), &no_ttl())
.unwrap();
fixture
.cache
.set_cache("c", entry(json!("C")), &no_ttl())
.unwrap();
fixture.service.seed_blob("bad", b"nope");
let keys = ["c", "missing", "a", "bad"].map(String::from);
let sync = fixture.cache.batch_get_cache(&keys, &no_ttl()).unwrap();
assert_eq!(
sync,
vec![
BatchEntry::Hit(entry(json!("C"))),
BatchEntry::Miss,
BatchEntry::Hit(entry(json!("A"))),
BatchEntry::Invalid,
]
);
let asynchronous = fixture
.runtime
.block_on(fixture.cache.async_batch_get_cache(keys.to_vec(), no_ttl()))
.unwrap();
assert_eq!(asynchronous, sync);
let response_cache = fixture.response_cache();
let requests = [request("hit"), request("missing"), request("bad")];
response_cache
.store(&requests[0], json!("HIT"), now())
.unwrap();
fixture
.service
.seed_blob(&cache_key(&requests[2].key), b"nope");
let hits = response_cache.lookup_batch(&requests, now()).unwrap();
assert_eq!(hits.values, vec![Some(json!("HIT")), None, None]);
assert_eq!(hits.missing_indices, vec![1, 2]);
let async_hits = fixture
.runtime
.block_on(response_cache.async_lookup_batch(&requests, now()))
.unwrap();
assert_eq!(async_hits.values, hits.values);
}
#[test]
fn async_pipeline_writes_every_entry_with_overwrite() {
let fixture = Fixture::new(FakeBlobService::default());
fixture.service.seed_blob("k2", b"stale");
fixture
.runtime
.block_on(fixture.cache.async_set_cache_pipeline(
vec![
("k1".into(), entry(json!({"n": 1}))),
("k2".into(), entry(json!({"n": 2}))),
("k3".into(), entry(json!({"n": 3}))),
],
with_ttl(30),
))
.unwrap();
assert_eq!(fixture.service.blob_names(), ["k1", "k2", "k3"]);
assert_eq!(fixture.stored_json("k2")["response"], json!({"n": 2}));
}
#[test]
fn flush_deletes_every_blob_in_the_container() {
let fixture = Fixture::new(FakeBlobService::default());
for key in ["x", "y", "z"] {
fixture
.cache
.set_cache(key, entry(json!(key)), &no_ttl())
.unwrap();
}
fixture.cache.flush_cache().unwrap();
assert!(fixture.service.blob_names().is_empty());
assert!(fixture.service.container_exists());
fixture
.cache
.set_cache("again", entry(json!(1)), &no_ttl())
.unwrap();
fixture
.runtime
.block_on(fixture.cache.async_flush_cache())
.unwrap();
assert!(fixture.service.blob_names().is_empty());
}
#[test]
fn service_failures_map_to_unavailable() {
let fixture = Fixture::new(FakeBlobService::default());
fixture.service.set_failing(true);
assert!(matches!(
fixture.cache.get_cache("key", &no_ttl()),
Err(Error::Unavailable)
));
assert!(matches!(
fixture.cache.set_cache("key", entry(json!(1)), &no_ttl()),
Err(Error::Unavailable)
));
assert!(matches!(
fixture.cache.flush_cache(),
Err(Error::Unavailable)
));
assert!(matches!(
fixture.runtime.block_on(
fixture
.cache
.async_set_cache_pipeline(vec![("k".into(), entry(json!(1)))], no_ttl())
),
Err(Error::Unavailable)
));
}
#[test]
fn test_connection_reports_container_reachability() {
let fixture = Fixture::new(FakeBlobService::default());
let ok = fixture
.runtime
.block_on(fixture.cache.test_connection())
.unwrap();
assert_eq!(ok.status, CacheConnectionStatus::Success);
assert!(ok.error.is_none());
fixture.service.set_failing(true);
let failed = fixture
.runtime
.block_on(fixture.cache.test_connection())
.unwrap();
assert_eq!(failed.status, CacheConnectionStatus::Failed);
assert!(failed.error.is_some());
}
#[test]
fn disconnect_is_idempotent_and_keeps_data() {
let fixture = Fixture::new(FakeBlobService::default());
fixture
.cache
.set_cache("key", entry(json!(1)), &no_ttl())
.unwrap();
fixture.runtime.block_on(async {
fixture.cache.disconnect().await.unwrap();
fixture.cache.disconnect().await.unwrap();
});
assert_eq!(
fixture.cache.get_cache("key", &no_ttl()).unwrap(),
Some(entry(json!(1)))
);
}
#[test]
fn response_cache_stores_and_reads_through_the_backend() {
let fixture = Fixture::new(FakeBlobService::default());
let response_cache = fixture.response_cache();
let mut request = request("gpt");
request.context = with_ttl(60);
let response = json!({"id": "chatcmpl-1"});
response_cache
.store(&request, response.clone(), now())
.unwrap();
assert_eq!(
fixture.stored_json(&cache_key(&request.key)),
json!({"timestamp": 1_700_000_000.0, "response": {"id": "chatcmpl-1"}})
);
assert_eq!(
response_cache
.lookup(&request, now() + Duration::from_secs(3600))
.unwrap(),
Some(response.clone())
);
assert_eq!(
fixture
.runtime
.block_on(response_cache.async_lookup(&request, now() + Duration::from_secs(3600)))
.unwrap(),
Some(response.clone())
);
fixture.runtime.block_on(async {
response_cache
.async_store(&request, json!("replaced"), now())
.await
.unwrap();
assert_eq!(
response_cache.async_lookup(&request, now()).await.unwrap(),
Some(json!("replaced"))
);
response_cache.async_flush().await.unwrap();
assert_eq!(
response_cache.async_lookup(&request, now()).await.unwrap(),
None
);
});
}
#[test]
fn non_object_responses_are_written_serialized_like_python() {
let fixture = Fixture::new(FakeBlobService::default());
fixture
.cache
.set_cache("s", entry(json!("plain")), &no_ttl())
.unwrap();
assert_eq!(
fixture.stored_json("s"),
json!({"timestamp": 1_700_000_000.5, "response": "\"plain\""})
);
assert_eq!(
fixture.cache.get_cache("s", &no_ttl()).unwrap(),
Some(entry(json!("plain")))
);
}

84
litellm-rust/crates/cache-azure-blob/src/credential.rs Normal file
View file

@ -0,0 +1,84 @@
use std::{
fmt,
sync::Arc,
time::{Duration, SystemTime},
};
use azure_core::{
credentials::{AccessToken, TokenCredential, TokenRequestOptions},
error::ErrorKind,
time::OffsetDateTime,
};
use litellm_auth_azure::{AzureAuthInputs, AzureAuthService};
use litellm_auth_types::ResolvedCredential;
const STATIC_TOKEN_LIFETIME: Duration = Duration::from_secs(300);
const LLM_TOKEN_ENV: &str = "AZURE_AD_TOKEN";
type EnvLookup = Arc<dyn Fn(&str) -> Option<String> + Send + Sync>;
pub struct AzureBlobCredential {
service: AzureAuthService,
env_lookup: EnvLookup,
}
impl fmt::Debug for AzureBlobCredential {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str("AzureBlobCredential")
}
}
impl Default for AzureBlobCredential {
fn default() -> Self {
Self::new(
AzureAuthService::default(),
Arc::new(|name| std::env::var(name).ok()),
)
}
}
impl AzureBlobCredential {
pub fn new(service: AzureAuthService, env_lookup: EnvLookup) -> Self {
Self {
service,
env_lookup,
}
}
}
#[async_trait::async_trait]
impl TokenCredential for AzureBlobCredential {
async fn get_token(
&self,
scopes: &[&str],
_options: Option<TokenRequestOptions<'_>>,
) -> azure_core::Result<AccessToken> {
let env_lookup = &self.env_lookup;
let lookup = move |name: &str| (name != LLM_TOKEN_ENV).then(|| env_lookup(name)).flatten();
let credential = self
.service
.get_azure_ad_token(
&AzureAuthInputs::default_credential_for_scope(&scopes.join(" ")),
&lookup,
)
.await
.map_err(|error| {
azure_core::Error::with_message(ErrorKind::Credential, error.to_string())
})?
.ok_or_else(|| {
azure_core::Error::with_message(
ErrorKind::Credential,
"no Azure credential is available for blob storage",
)
})?;
let (token, expires_on) = match credential.into_value() {
ResolvedCredential::AccessToken { token, expires_on } => (token, expires_on),
ResolvedCredential::Static(token) => (token, None),
};
let expires_on = expires_on.unwrap_or_else(|| SystemTime::now() + STATIC_TOKEN_LIFETIME);
Ok(AccessToken::new(
token.expose().to_string(),
OffsetDateTime::from(expires_on),
))
}
}

5
litellm-rust/crates/cache-azure-blob/src/lib.rs Normal file
View file

@ -0,0 +1,5 @@
mod cache;
mod credential;
pub use cache::AzureBlobCache;
pub use credential::AzureBlobCredential;

0
litellm-rust/crates/cache-azure-blob/src/tests.rs Normal file
View file

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

@ -21,6 +21,7 @@ tiktoken = ["litellm-token-counter/tiktoken"]
[dependencies]
bytes.workspace = true
litellm-cache.workspace = true
litellm-cache-azure-blob.workspace = true
litellm-cache-memory.workspace = true
litellm-cache-redis.workspace = true
litellm-cache-response.workspace = true

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

@ -75,6 +75,11 @@ pub(super) struct RedisCacheConfig {
pub(super) connection: RedisConnectionConfig,
}
pub(super) struct AzureBlobCacheConfig {
pub(super) account_url: String,
pub(super) container: String,
}
struct RedisClientProjection<'py> {
topology: RedisTopology,
host: String,
@ -89,6 +94,7 @@ const REDIS_PY_DEFAULT_MAX_CONNECTIONS: usize = 1 << 31;
pub(super) enum CacheBackendConfig {
Memory(MemoryCacheConfig),
Redis(Box<RedisCacheConfig>),
AzureBlob(AzureBlobCacheConfig),
}
#[allow(dead_code, reason = "consumed by the cache activation follow-up")]
@ -155,13 +161,18 @@ impl NativeCacheConfig {
}))),
Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)),
},
Some(CacheType::AzureBlob) => project_azure_blob(&backend).map(|backend| {
CacheConfigProjection::Native(Box::new(Self {
policy,
backend: CacheBackendConfig::AzureBlob(backend),
}))
}),
Some(
CacheType::RedisSemantic
| CacheType::ValkeySemantic
| CacheType::S3
| CacheType::Disk
| CacheType::QdrantSemantic
| CacheType::AzureBlob
| CacheType::Gcs,
)
| None => Ok(CacheConfigProjection::Unsupported(
@ -171,12 +182,12 @@ impl NativeCacheConfig {
}
pub(super) fn service_mismatch(&self, service: &NativeResponseCache) -> Option<&'static str> {
if service.default_ttl()
!= Some(match &self.backend {
CacheBackendConfig::Memory(config) => config.default_ttl,
CacheBackendConfig::Redis(config) => config.default_ttl,
})
{
let default_ttl = match &self.backend {
CacheBackendConfig::Memory(config) => Some(config.default_ttl),
CacheBackendConfig::Redis(config) => Some(config.default_ttl),
CacheBackendConfig::AzureBlob(_) => None,
};
if service.default_ttl() != default_ttl {
return Some("facade and native backend default TTLs must match");
}
match &self.backend {
@ -201,10 +212,34 @@ impl NativeCacheConfig {
CacheBackendConfig::Redis(config) => (service.namespace()
!= config.namespace.as_deref())
.then_some("facade and native backend namespaces must match"),
CacheBackendConfig::AzureBlob(config) => match service.azure_blob_identity() {
None => Some("facade and native backend types must match"),
Some((account_url, container))
if account_url != config.account_url || container != config.container =>
{
Some("facade and native backend containers must match")
}
Some(_) => None,
},
}
}
}
#[inline(never)]
fn project_azure_blob(backend: &Bound<'_, PyAny>) -> PyResult<AzureBlobCacheConfig> {
let client = backend.getattr("container_client")?;
let container = client.getattr("container_name")?.extract::<String>()?;
let url = client.getattr("url")?.extract::<String>()?;
let account_url = url
.strip_suffix(container.as_str())
.and_then(|url| url.strip_suffix('/'))
.ok_or_else(|| PyValueError::new_err("Azure Blob container URL is malformed"))?;
Ok(AzureBlobCacheConfig {
account_url: account_url.to_string(),
container,
})
}
#[inline(never)]
fn project_memory(backend: &Bound<'_, PyAny>) -> PyResult<MemoryCacheConfig> {
let max_size_kib = backend.getattr("max_size_per_item")?.extract::<usize>()?;

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

@ -34,6 +34,19 @@ struct RedisPoolGuard {
attributes: RedisPoolAttributes,
}
struct AzureBlobClientGuard {
sync_client: Py<PyAny>,
async_client: Py<PyAny>,
url: String,
container_name: String,
}
enum ConnectionGuard {
None,
RedisPool(RedisPoolGuard),
AzureBlob(AzureBlobClientGuard),
}
struct RedisPoolAttributes {
pool: &'static str,
connection_class: &'static str,
@ -55,7 +68,7 @@ const CLUSTER_POOL: RedisPoolAttributes = RedisPoolAttributes {
pub(super) struct FacadeGuard {
outer: ObjectGuard,
backend: ObjectGuard,
redis_pool: Option<RedisPoolGuard>,
connection: ConnectionGuard,
}
impl ObjectGuard {
@ -205,6 +218,61 @@ impl RedisPoolGuard {
}
}
impl AzureBlobClientGuard {
fn capture(backend: &Bound<'_, PyAny>) -> PyResult<Self> {
let sync_client = backend.getattr("container_client")?;
Ok(Self {
url: sync_client.getattr("url")?.extract::<String>()?,
container_name: sync_client.getattr("container_name")?.extract::<String>()?,
sync_client: sync_client.unbind(),
async_client: backend.getattr("async_container_client")?.unbind(),
})
}
fn matches(&self, py: Python<'_>, backend: &Bound<'_, PyAny>) -> PyResult<bool> {
let sync_client = backend.getattr("container_client")?;
Ok(self.sync_client.bind(py).is(&sync_client)
&& self
.async_client
.bind(py)
.is(&backend.getattr("async_container_client")?)
&& self.url == sync_client.getattr("url")?.extract::<String>()?
&& self.container_name == sync_client.getattr("container_name")?.extract::<String>()?)
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&self.sync_client)?;
visit.call(&self.async_client)
}
}
impl ConnectionGuard {
fn capture(kind: &str, cluster: bool, backend: &Bound<'_, PyAny>) -> PyResult<Self> {
Ok(match (kind, cluster) {
("redis", false) => Self::RedisPool(RedisPoolGuard::capture(backend, STANDALONE_POOL)?),
("redis", true) => Self::RedisPool(RedisPoolGuard::capture(backend, CLUSTER_POOL)?),
("azure-blob", _) => Self::AzureBlob(AzureBlobClientGuard::capture(backend)?),
_ => Self::None,
})
}
fn matches(&self, py: Python<'_>, backend: &Bound<'_, PyAny>) -> PyResult<bool> {
match self {
Self::None => Ok(true),
Self::RedisPool(guard) => guard.matches(py, backend),
Self::AzureBlob(guard) => guard.matches(py, backend),
}
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
match self {
Self::None => Ok(()),
Self::RedisPool(guard) => guard.traverse(visit),
Self::AzureBlob(guard) => guard.traverse(visit),
}
}
}
impl FacadeGuard {
pub(super) fn capture(
py: Python<'_>,
@ -227,6 +295,11 @@ impl FacadeGuard {
"RedisClusterCache",
"redis",
),
("azure-blob", _) => (
"litellm.caching.azure_blob_cache",
"AzureBlobCache",
"azure-blob",
),
_ => unreachable!(),
};
let backend = facade.getattr("cache")?;
@ -272,11 +345,7 @@ impl FacadeGuard {
"redis_flush_size",
],
)?,
redis_pool: match (kind, cluster) {
("redis", false) => Some(RedisPoolGuard::capture(&backend, STANDALONE_POOL)?),
("redis", true) => Some(RedisPoolGuard::capture(&backend, CLUSTER_POOL)?),
_ => None,
},
connection: ConnectionGuard::capture(kind, cluster, &backend)?,
})
}
@ -288,19 +357,13 @@ impl FacadeGuard {
if !self.backend.matches(py, &backend)? {
return Ok(false);
}
match &self.redis_pool {
Some(guard) => guard.matches(py, &backend),
None => Ok(true),
}
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.redis_pool {
guard.traverse(&visit)?;
}
Ok(())
self.connection.traverse(&visit)
}
}

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

@ -1,5 +1,5 @@
use litellm_cache_redis::{RedisNode, RedisTopology};
use litellm_host_python::release_gil;
use litellm_host_python::{release_gil, run_sync_value};
use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*};
use super::{cache_error, facade::FacadeGuard, native::NativeResponseCache, request::duration};
@ -64,6 +64,21 @@ impl CacheTestHandle {
})
}
#[staticmethod]
#[pyo3(signature = (account_url, container))]
fn azure_blob(py: Python<'_>, account_url: String, container: String) -> PyResult<Self> {
let service = run_sync_value(py, async move {
NativeResponseCache::azure_blob(&account_url, &container)
.await
.map_err(cache_error)
})?;
Ok(Self {
service,
guard: None,
pid: std::process::id(),
})
}
#[getter]
fn backend(&self) -> &'static str {
self.service.kind()

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

@ -1,6 +1,7 @@
use std::{sync::Arc, time::Duration};
use litellm_cache::{CacheCodec, CacheConnectionResult, Error};
use litellm_cache_azure_blob::AzureBlobCache;
use litellm_cache_memory::InMemoryCache;
use litellm_cache_redis::{RedisCache, RedisTopology};
use litellm_cache_response::{
@ -15,6 +16,7 @@ pub(super) enum NativeResponseCache {
cache: Arc<ResponseCache<RedisCache<ResponseCacheCodec>>>,
buffer: Option<Arc<WriteBuffer>>,
},
AzureBlob(Arc<ResponseCache<AzureBlobCache<ResponseCacheCodec>>>),
}
impl NativeResponseCache {
@ -45,6 +47,29 @@ impl NativeResponseCache {
buffer: None,
})
}
pub async fn azure_blob(account_url: &str, container: &str) -> Result<Self, Error> {
let backend = AzureBlobCache::connect(
account_url,
container,
ResponseCacheCodec,
tokio::runtime::Handle::current(),
)
.await?;
Ok(Self::AzureBlob(Arc::new(ResponseCache::new(Arc::new(
backend,
)))))
}
pub fn azure_blob_identity(&self) -> Option<(&str, &str)> {
match self {
Self::AzureBlob(cache) => Some((
cache.backend().account_url(),
cache.backend().container_name(),
)),
Self::Memory(_) | Self::Redis { .. } => None,
}
}
}
impl NativeResponseCache {
@ -52,6 +77,7 @@ impl NativeResponseCache {
match self {
Self::Memory(_) => "memory",
Self::Redis { .. } => "redis",
Self::AzureBlob(_) => "azure-blob",
}
}
@ -59,19 +85,20 @@ impl NativeResponseCache {
match self {
Self::Memory(cache) => cache.default_ttl(),
Self::Redis { cache, .. } => cache.default_ttl(),
Self::AzureBlob(cache) => cache.default_ttl(),
}
}
pub fn namespace(&self) -> Option<&str> {
match self {
Self::Memory(_) => None,
Self::Memory(_) | Self::AzureBlob(_) => None,
Self::Redis { cache, .. } => cache.backend().namespace(),
}
}
pub fn topology(&self) -> Option<&RedisTopology> {
match self {
Self::Memory(_) => None,
Self::Memory(_) | Self::AzureBlob(_) => None,
Self::Redis { cache, .. } => Some(cache.backend().topology()),
}
}
@ -79,14 +106,14 @@ impl NativeResponseCache {
pub fn capacity(&self) -> Option<usize> {
match self {
Self::Memory(cache) => Some(cache.backend().max_size_in_memory()),
Self::Redis { .. } => None,
Self::Redis { .. } | Self::AzureBlob(_) => None,
}
}
pub fn max_entry_bytes(&self) -> Option<usize> {
match self {
Self::Memory(cache) => cache.backend().max_entry_bytes(),
Self::Redis { .. } => None,
Self::Redis { .. } | Self::AzureBlob(_) => None,
}
}
@ -96,7 +123,7 @@ impl NativeResponseCache {
cache,
buffer: flush_size.map(|flush_size| Arc::new(WriteBuffer::new(flush_size))),
},
memory => memory,
other => other,
}
}
@ -108,6 +135,7 @@ impl NativeResponseCache {
match self {
Self::Memory(cache) => cache.lookup(request, now),
Self::Redis { cache, .. } => cache.lookup(request, now),
Self::AzureBlob(cache) => cache.lookup(request, now),
}
}
@ -120,6 +148,7 @@ impl NativeResponseCache {
match self {
Self::Memory(cache) => cache.store(request, response, now),
Self::Redis { cache, .. } => cache.store(request, response, now),
Self::AzureBlob(cache) => cache.store(request, response, now),
}
}
@ -131,6 +160,7 @@ impl NativeResponseCache {
match self {
Self::Memory(cache) => cache.lookup_batch(requests, now),
Self::Redis { cache, .. } => cache.lookup_batch(requests, now),
Self::AzureBlob(cache) => cache.lookup_batch(requests, now),
}
}
@ -142,6 +172,7 @@ impl NativeResponseCache {
match self {
Self::Memory(cache) => cache.async_lookup(request, now).await,
Self::Redis { cache, .. } => cache.async_lookup(request, now).await,
Self::AzureBlob(cache) => cache.async_lookup(request, now).await,
}
}
@ -161,6 +192,7 @@ impl NativeResponseCache {
cache,
buffer: Some(buffer),
} => buffer.async_store(cache, request, response, now).await,
Self::AzureBlob(cache) => cache.async_store(request, response, now).await,
}
}
@ -172,6 +204,7 @@ impl NativeResponseCache {
match self {
Self::Memory(cache) => cache.async_lookup_batch(requests, now).await,
Self::Redis { cache, .. } => cache.async_lookup_batch(requests, now).await,
Self::AzureBlob(cache) => cache.async_lookup_batch(requests, now).await,
}
}
@ -183,6 +216,7 @@ impl NativeResponseCache {
match self {
Self::Memory(cache) => cache.async_store_batch(entries, now).await,
Self::Redis { cache, .. } => cache.async_store_batch(entries, now).await,
Self::AzureBlob(cache) => cache.async_store_batch(entries, now).await,
}
}
@ -195,6 +229,7 @@ impl NativeResponseCache {
}
cache.async_flush().await
}
Self::AzureBlob(cache) => cache.async_flush().await,
}
}
@ -202,6 +237,7 @@ impl NativeResponseCache {
match self {
Self::Memory(cache) => cache.test_connection().await,
Self::Redis { cache, .. } => cache.test_connection().await,
Self::AzureBlob(cache) => cache.test_connection().await,
}
}
}

24
litellm-rust/crates/secrets-azure/Cargo.toml Normal file
View file

@ -0,0 +1,24 @@
[package]
name = "litellm-secrets-azure"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
litellm-auth-azure.workspace = true
litellm-auth-types.workspace = true
litellm-secrets-types.workspace = true
litellm-core-utils.workspace = true
reqwest.workspace = true
serde.workspace = true
thiserror.workspace = true
veil.workspace = true
percent-encoding = "2.3"
[dev-dependencies]
tokio.workspace = true
wiremock = "0.6.5"
rstest.workspace = true
serde_json.workspace = true
sha2.workspace = true

25
litellm-rust/crates/secrets-azure/src/error.rs Normal file
View file

@ -0,0 +1,25 @@
#[derive(thiserror::Error, veil::Redact)]
pub enum Error {
#[error("{0} environment variable is missing")]
MissingEnvironment(&'static str),
#[error("AZURE_KEY_VAULT_URI is not a valid https vault URL")]
VaultUri,
#[error("Azure Key Vault credentials are not configured")]
MissingCredentials,
#[error(transparent)]
Auth(
#[from]
#[redact]
litellm_auth_types::Error,
),
#[error("Azure Key Vault request failed")]
Http(
#[source]
#[redact]
reqwest::Error,
),
#[error("Azure Key Vault returned HTTP {0}")]
Status(u16),
#[error("Azure Key Vault response is missing the secret value")]
MissingValue,
}

118
litellm-rust/crates/secrets-azure/src/key_vault.rs Normal file
View file

@ -0,0 +1,118 @@
use std::sync::Arc;
use litellm_auth_azure::{AzureAuthInputs, AzureAuthService, ConfigValue};
use litellm_auth_types::{InputSource, Sourced};
use litellm_core_utils::settings::Lookup;
use litellm_secrets_types::{Secret, SecretValue};
use percent_encoding::{AsciiSet, NON_ALPHANUMERIC};
use serde::Deserialize;
use crate::Error;
const AZURE_KEY_VAULT_URI: &str = "AZURE_KEY_VAULT_URI";
const API_VERSION: &str = "7.4";
const PATH_SEGMENT: &AsciiSet = &NON_ALPHANUMERIC
.remove(b'-')
.remove(b'.')
.remove(b'_')
.remove(b'~');
#[derive(Clone)]
pub struct AzureKeyVault {
client: reqwest::Client,
vault: reqwest::Url,
auth: Arc<AzureAuthService>,
inputs: Arc<AzureAuthInputs>,
environment: Arc<dyn Lookup + Send + Sync>,
}
#[derive(Deserialize)]
struct SecretResponse {
value: Option<String>,
}
impl AzureKeyVault {
pub fn with_client(
client: reqwest::Client,
vault: reqwest::Url,
environment: Arc<dyn Lookup + Send + Sync>,
) -> Result<Self, Error> {
if vault.host_str().is_none() {
return Err(Error::VaultUri);
}
let inputs = AzureAuthInputs {
azure_scope: ConfigValue::Value(Sourced::new(
scope_for(&vault),
InputSource::Deployment,
)),
enable_azure_ad_token_refresh: Sourced::new(true, InputSource::Deployment),
..AzureAuthInputs::default()
};
Ok(Self {
client,
vault,
auth: Arc::new(AzureAuthService::default()),
inputs: Arc::new(inputs),
environment,
})
}
pub fn new(environment: Arc<dyn Lookup + Send + Sync>) -> Result<Self, Error> {
let value = environment
.get(AZURE_KEY_VAULT_URI)
.ok_or(Error::MissingEnvironment(AZURE_KEY_VAULT_URI))?;
let vault = reqwest::Url::parse(&value).map_err(|_| Error::VaultUri)?;
if vault.scheme() != "https" || vault.host_str().is_none() {
return Err(Error::VaultUri);
}
Self::with_client(reqwest::Client::new(), vault, environment)
}
pub fn scope(&self) -> &str {
self.inputs
.azure_scope
.as_value()
.map(|value| value.value().as_str())
.unwrap_or_default()
}
pub async fn get_secret_from_azure_key_vault(
&self,
name: &str,
) -> Result<Option<Secret>, Error> {
let token = self
.auth
.get_azure_ad_token(&self.inputs, &|key| self.environment.get(key))
.await?
.ok_or(Error::MissingCredentials)?;
let encoded_name = percent_encoding::utf8_percent_encode(name, PATH_SEGMENT);
let url = self
.vault
.join(&format!("secrets/{encoded_name}?api-version={API_VERSION}"))
.map_err(|_| Error::VaultUri)?;
let response = self
.client
.get(url)
.bearer_auth(token.value().secret().expose())
.send()
.await
.map_err(Error::Http)?;
if response.status() == reqwest::StatusCode::NOT_FOUND {
return Ok(None);
}
if response.status() != reqwest::StatusCode::OK {
return Err(Error::Status(response.status().as_u16()));
}
let payload: SecretResponse = response.json().await.map_err(Error::Http)?;
let value = payload.value.ok_or(Error::MissingValue)?;
Ok(Some(Secret::String(SecretValue::new(value))))
}
}
fn scope_for(vault: &reqwest::Url) -> String {
let host = vault.host_str().unwrap_or_default();
let resource = host
.split_once('.')
.map_or(host, |(_, remainder)| remainder);
format!("https://{resource}/.default")
}

7
litellm-rust/crates/secrets-azure/src/lib.rs Normal file
View file

@ -0,0 +1,7 @@
#![forbid(unsafe_code)]
mod error;
mod key_vault;
pub use error::Error;
pub use key_vault::AzureKeyVault;

8
litellm-rust/crates/secrets-azure/tests/fixtures/key_vault_parity.json vendored Normal file
View file

@ -0,0 +1,8 @@
{
"cases": [
{"name": "plain_value", "secret_name": "OPENAI-API-KEY", "response": {"status": 200, "body": {"value": "sk-parity-1", "id": "https://example.vault.azure.net/secrets/OPENAI-API-KEY/abc"}}, "expected": {"value": "sk-parity-1"}},
{"name": "json_value_is_kept_as_string", "secret_name": "JSON-SECRET", "response": {"status": 200, "body": {"value": "{\"api_key\": \"nested\"}", "id": "https://example.vault.azure.net/secrets/JSON-SECRET/abc"}}, "expected": {"value": "{\"api_key\": \"nested\"}"}},
{"name": "missing_secret", "secret_name": "MISSING", "response": {"status": 404, "body": {"error": {"code": "SecretNotFound", "message": "not found"}}}, "expected": {"missing": true}},
{"name": "forbidden", "secret_name": "FORBIDDEN", "response": {"status": 403, "body": {"error": {"code": "Forbidden", "message": "denied"}}}, "expected": {"error": true}}
]
}

222
litellm-rust/crates/secrets-azure/tests/key_vault.rs Normal file
View file

@ -0,0 +1,222 @@
use std::sync::Arc;
use litellm_secrets_azure::{AzureKeyVault, Error};
use litellm_secrets_types::{Secret, SecretValue};
use serde::Deserialize;
use wiremock::{
Mock, MockServer, ResponseTemplate,
matchers::{header, path, query_param},
};
fn manager(server: &MockServer) -> AzureKeyVault {
AzureKeyVault::with_client(
reqwest::Client::new(),
server.uri().parse().unwrap(),
Arc::new(|name: &str| (name == "AZURE_AD_TOKEN").then(|| "fake".to_owned())),
)
.unwrap()
}
#[tokio::test]
async fn reads_secret_with_bearer_token_and_api_version() {
let server = MockServer::start().await;
Mock::given(path("/secrets/OPENAI-API-KEY"))
.and(query_param("api-version", "7.4"))
.and(header("authorization", "Bearer fake"))
.respond_with(
ResponseTemplate::new(200)
.set_body_json(serde_json::json!({"value": "s3cret", "id": "secret-id"})),
)
.expect(1)
.mount(&server)
.await;
let secret = manager(&server)
.get_secret_from_azure_key_vault("OPENAI-API-KEY")
.await
.unwrap()
.unwrap();
assert_eq!(secret, Secret::String(SecretValue::new("s3cret")));
}
#[tokio::test]
async fn percent_encodes_secret_name_path_segment() {
let server = MockServer::start().await;
Mock::given(path("/secrets/name%2Fwith%20spaces"))
.and(query_param("api-version", "7.4"))
.respond_with(
ResponseTemplate::new(200).set_body_json(serde_json::json!({"value": "value"})),
)
.expect(1)
.mount(&server)
.await;
let secret = manager(&server)
.get_secret_from_azure_key_vault("name/with spaces")
.await
.unwrap()
.unwrap();
assert_eq!(secret.as_str(), Some("value"));
}
#[rstest::rstest]
#[case::not_found(404, None)]
#[case::forbidden(403, Some(403))]
#[tokio::test]
async fn handles_statuses(#[case] status: u16, #[case] expected_status: Option<u16>) {
let server = MockServer::start().await;
Mock::given(path("/secrets/NAME"))
.respond_with(ResponseTemplate::new(status))
.expect(1)
.mount(&server)
.await;
let result = manager(&server)
.get_secret_from_azure_key_vault("NAME")
.await;
match expected_status {
None => assert_eq!(result.unwrap(), None),
Some(status) => assert!(matches!(result, Err(Error::Status(actual)) if actual == status)),
}
}
#[tokio::test]
async fn missing_value_is_an_error() {
let server = MockServer::start().await;
Mock::given(path("/secrets/NAME"))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({})))
.expect(1)
.mount(&server)
.await;
assert!(matches!(
manager(&server)
.get_secret_from_azure_key_vault("NAME")
.await,
Err(Error::MissingValue)
));
}
#[test]
fn new_validates_vault_environment() {
assert!(matches!(
AzureKeyVault::new(Arc::new(|_: &str| None)),
Err(Error::MissingEnvironment("AZURE_KEY_VAULT_URI"))
));
assert!(matches!(
AzureKeyVault::new(Arc::new(|name: &str| {
(name == "AZURE_KEY_VAULT_URI").then(|| "http://vault.example".to_owned())
})),
Err(Error::VaultUri)
));
assert!(matches!(
AzureKeyVault::new(Arc::new(|name: &str| {
(name == "AZURE_KEY_VAULT_URI").then(|| "vault.example".to_owned())
})),
Err(Error::VaultUri)
));
}
#[rstest::rstest]
#[case("https://myvault.vault.azure.net", "https://vault.azure.net/.default")]
#[case(
"https://v.vault.usgovcloudapi.net/",
"https://vault.usgovcloudapi.net/.default"
)]
#[case("http://localhost:8080", "https://localhost/.default")]
#[test]
fn derives_scope_from_vault_host(#[case] uri: &str, #[case] expected: &str) {
let manager = AzureKeyVault::with_client(
reqwest::Client::new(),
uri.parse().unwrap(),
Arc::new(|_: &str| None),
)
.unwrap();
assert_eq!(manager.scope(), expected);
}
#[tokio::test]
async fn missing_credentials_do_not_request_vault() {
let server = MockServer::start().await;
Mock::given(path("/secrets/NAME"))
.respond_with(ResponseTemplate::new(200))
.expect(0)
.mount(&server)
.await;
assert!(
manager_without_credentials(&server)
.get_secret_from_azure_key_vault("NAME")
.await
.is_err()
);
}
fn manager_without_credentials(server: &MockServer) -> AzureKeyVault {
AzureKeyVault::with_client(
reqwest::Client::new(),
server.uri().parse().unwrap(),
Arc::new(|name: &str| {
(name == "AZURE_CREDENTIAL").then(|| "ClientSecretCredential".to_owned())
}),
)
.unwrap()
}
#[derive(Deserialize)]
struct Fixture {
cases: Vec<FixtureCase>,
}
#[derive(Deserialize)]
struct FixtureCase {
secret_name: String,
response: FixtureResponse,
expected: FixtureExpected,
}
#[derive(Deserialize)]
struct FixtureResponse {
status: u16,
body: serde_json::Value,
}
#[derive(Deserialize)]
struct FixtureExpected {
value: Option<String>,
missing: Option<bool>,
error: Option<bool>,
}
#[tokio::test]
async fn parity_fixture_matches_python_backend_contract() {
let fixture: Fixture =
serde_json::from_str(include_str!("fixtures/key_vault_parity.json")).unwrap();
for case in fixture.cases {
let server = MockServer::start().await;
Mock::given(path(format!("/secrets/{}", case.secret_name)))
.respond_with(
ResponseTemplate::new(case.response.status).set_body_json(case.response.body),
)
.expect(1)
.mount(&server)
.await;
let result = manager(&server)
.get_secret_from_azure_key_vault(&case.secret_name)
.await;
if case.expected.missing == Some(true) {
assert_eq!(result.unwrap(), None);
} else if case.expected.error == Some(true) {
assert!(result.is_err());
} else {
assert_eq!(
result.unwrap().unwrap().as_str(),
case.expected.value.as_deref()
);
}
}
}

30
litellm-rust/crates/secrets-azure/tests/live.rs Normal file
View file

@ -0,0 +1,30 @@
use std::sync::Arc;
use litellm_core_utils::settings::ProcessEnvironment;
use litellm_secrets_azure::AzureKeyVault;
use litellm_secrets_types::Secret;
#[tokio::test]
#[ignore]
async fn reads_a_real_secret() {
let environment = Arc::new(ProcessEnvironment);
let manager = AzureKeyVault::new(environment).unwrap();
let name = std::env::var("AZURE_KEY_VAULT_LIVE_SECRET_NAME").unwrap();
let secret = manager
.get_secret_from_azure_key_vault(&name)
.await
.unwrap()
.unwrap();
assert!(matches!(&secret, Secret::String(_)));
let host = std::env::var("AZURE_KEY_VAULT_URI")
.unwrap()
.parse::<reqwest::Url>()
.unwrap()
.host_str()
.unwrap()
.to_owned();
let value_len = secret.as_str().unwrap().len();
println!(
"native provider=litellm-secrets-azure vault_host={host} secret={name} value_len={value_len}"
);
}

2
litellm-rust/crates/secrets/Cargo.toml
View file

@ -9,12 +9,14 @@ repository.workspace = true
default = []
aws = ["dep:litellm-secrets-aws"]
google = ["dep:litellm-secrets-google"]
azure = ["dep:litellm-secrets-azure"]
cyberark = ["dep:litellm-secrets-cyberark"]
[dependencies]
litellm-secrets-types.workspace = true
litellm-secrets-aws = { workspace = true, optional = true }
litellm-secrets-google = { workspace = true, optional = true }
litellm-secrets-azure = { workspace = true, optional = true }
litellm-secrets-cyberark = { workspace = true, optional = true }
litellm-core-utils.workspace = true
base64.workspace = true

3
litellm-rust/crates/secrets/src/error.rs
View file

@ -30,6 +30,9 @@ pub enum Error {
#[cfg(feature = "google")]
#[error(transparent)]
Google(#[from] litellm_secrets_google::Error),
#[cfg(feature = "azure")]
#[error(transparent)]
Azure(#[from] litellm_secrets_azure::Error),
#[cfg(feature = "cyberark")]
#[error(transparent)]
Cyberark(#[from] litellm_secrets_cyberark::Error),

9
litellm-rust/crates/secrets/src/handler.rs
View file

@ -13,6 +13,8 @@ pub enum SecretManager {
GoogleKms(crate::google::GoogleKms),
#[cfg(feature = "google")]
GoogleSecretManager(crate::google::GoogleSecretManager),
#[cfg(feature = "azure")]
AzureKeyVault(crate::azure::AzureKeyVault),
#[cfg(feature = "cyberark")]
Cyberark(crate::cyberark::CyberArkSecretManager),
}
@ -29,6 +31,8 @@ impl SecretManager {
Self::GoogleKms(_) => KeyManagementSystem::GoogleKms,
#[cfg(feature = "google")]
Self::GoogleSecretManager(_) => KeyManagementSystem::GoogleSecretManager,
#[cfg(feature = "azure")]
Self::AzureKeyVault(_) => KeyManagementSystem::AzureKeyVault,
#[cfg(feature = "cyberark")]
Self::Cyberark(_) => KeyManagementSystem::Cyberark,
}
@ -82,6 +86,11 @@ pub async fn get_secret_from_manager(
.get_secret_from_google_secret_manager(secret_name)
.await
.map_err(Error::from),
#[cfg(feature = "azure")]
SecretManager::AzureKeyVault(client) => client
.get_secret_from_azure_key_vault(secret_name)
.await
.map_err(Error::from),
#[cfg(feature = "cyberark")]
SecretManager::Cyberark(client) => client
.async_read_secret(secret_name)

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

@ -17,6 +17,8 @@ pub use state::{SecretManagerState, secret_manager_would_be_consulted};
#[cfg(feature = "aws")]
pub use litellm_secrets_aws as aws;
#[cfg(feature = "azure")]
pub use litellm_secrets_azure as azure;
#[cfg(feature = "cyberark")]
pub use litellm_secrets_cyberark as cyberark;
#[cfg(feature = "google")]

61
litellm-rust/crates/secrets/tests/handler.rs
View file

@ -106,6 +106,67 @@ async fn google_handler_requires_canonical_base64_and_preserves_plaintext_whites
));
}
#[cfg(feature = "azure")]
#[tokio::test]
async fn azure_handler_reads_missing_and_failed_secrets() {
use litellm_secrets::{
Error, KeyManagementSettings, KeyManagementSystem, SecretManager, azure::AzureKeyVault,
get_secret_from_manager,
};
use wiremock::{
Mock, MockServer, ResponseTemplate,
matchers::{path, query_param},
};
let server = MockServer::start().await;
Mock::given(path("/secrets/KEY"))
.and(query_param("api-version", "7.4"))
.respond_with(
ResponseTemplate::new(200).set_body_json(serde_json::json!({"value": "value"})),
)
.expect(1)
.mount(&server)
.await;
let manager = SecretManager::AzureKeyVault(
AzureKeyVault::with_client(
reqwest::Client::new(),
server.uri().parse().unwrap(),
std::sync::Arc::new(|name: &str| (name == "AZURE_AD_TOKEN").then(|| "fake".to_owned())),
)
.unwrap(),
);
assert_eq!(manager.system(), KeyManagementSystem::AzureKeyVault);
let settings = KeyManagementSettings::default();
let value = get_secret_from_manager(&manager, "KEY", &settings, &|_: &str| None)
.await
.unwrap()
.unwrap();
assert_eq!(value.as_str(), Some("value"));
let not_found = Mock::given(path("/secrets/MISSING"))
.respond_with(ResponseTemplate::new(404))
.expect(1)
.mount_as_scoped(&server)
.await;
assert_eq!(
get_secret_from_manager(&manager, "MISSING", &settings, &|_: &str| None)
.await
.unwrap(),
None
);
drop(not_found);
Mock::given(path("/secrets/FAILED"))
.respond_with(ResponseTemplate::new(500))
.expect(1)
.mount(&server)
.await;
assert!(matches!(
get_secret_from_manager(&manager, "FAILED", &settings, &|_: &str| None).await,
Err(Error::Azure(_))
));
}
#[cfg(feature = "cyberark")]
#[tokio::test]
async fn cyberark_handler_reads_values_and_surfaces_errors() {

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"]

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,

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.

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